"""Atomic NumPy collection writers and processing settings for runtime snapshots.""" from __future__ import annotations import json import logging from contextlib import contextmanager from numbers import Integral from pathlib import Path import shutil from uuid import uuid4 from typing import Any, Iterator import numpy as np from python_app.file_io import atomic_write_json, rename_no_replace from python_app.models.dataset_model import ResultCollection, ResultPayload, SweepCollection from python_app.storage.npz.paths import collection_dir_name from python_app.file_io import sanitize_path_component from python_app.storage.npz.serialize import validate_trace_collection logger = logging.getLogger(__name__) @contextmanager def atomic_snapshot_directory(destination: Path) -> Iterator[Path]: """Publish a complete export directory; discard staging files after a failed write.""" destination.parent.mkdir(parents=True, exist_ok=True) if destination.exists(): raise FileExistsError(f"Snapshot directory already exists: {destination}") staging = destination.with_name(f".{destination.name}.{uuid4().hex}") staging.mkdir() try: yield staging rename_no_replace(staging, destination) finally: if staging.exists(): shutil.rmtree(staging) def save_processing_settings(directory: Path, settings: dict[int, dict] | None) -> None: """Merge immutable revision settings before publishing results that refer to them.""" if not settings: return path = directory / "processing_settings.json" saved = json.loads(path.read_text(encoding="utf-8")) if path.exists() else {} if not isinstance(saved, dict): raise ValueError("Saved processing settings must be a JSON object") updated = dict(saved) for revision, value in settings.items(): if isinstance(revision, bool) or not isinstance(revision, Integral) or not 0 <= revision < 1 << 64 or not isinstance(value, dict): raise ValueError("Processing settings require uint64 revision keys and JSON objects") key = str(revision) if key in saved and saved[key] != value: raise ValueError(f"Processing revision {revision} already has different settings") updated[key] = value if updated != saved: atomic_write_json(path, updated, durable=True) def save_trace_history_numpy( stage_dir: Path, history: list[SweepCollection], *, index_offset: int = 0 ) -> None: """Write raw/preprocessed collections as a NumPy directory tree. ``index_offset`` continues the per-collection directory numbering across calls so a streaming recorder can append successive chunks into the same stage directory without colliding or restarting the index. """ for collection in history: validate_trace_collection(collection) stage_dir.mkdir(parents=True, exist_ok=True) logger.debug("Writing %d NumPy trace collection(s) to %s", len(history), stage_dir) for index, collection in enumerate(history, start=index_offset): collection_dir = stage_dir / collection_dir_name(index, collection.collection_id, collection.monotonic_ns) with atomic_snapshot_directory(collection_dir) as collection_dir: traces_meta: list[dict[str, int | str]] = [] for trace in collection.traces: tag = f"i{trace.combo.input}_o{trace.combo.output}" freq = np.asarray(trace.frequency_hz, dtype=np.float32) s21 = np.asarray(trace.s21, dtype=np.complex64) np.save(collection_dir / f"{tag}_freq.npy", freq) np.save(collection_dir / f"{tag}_s21.npy", s21) traces_meta.append( { "input": int(trace.combo.input), "output": int(trace.combo.output), "points": int(freq.size), # When each combo was measured, which in a switched matrix is # spread across the collection window rather than aligned with it. "capture_start_ns": int(trace.capture_start_ns), "capture_end_ns": int(trace.capture_end_ns), "freq_file": f"{tag}_freq.npy", "s21_file": f"{tag}_s21.npy", } ) (collection_dir / "meta.json").write_text( json.dumps( { "collection_id": int(collection.collection_id), "monotonic_ns": int(collection.monotonic_ns), "capture_start_ns": int(collection.capture_start_ns), "capture_end_ns": int(collection.capture_end_ns), "trace_count": len(collection.traces), "traces": traces_meta, }, indent=2, ), encoding="utf-8", ) def _save_result_payload(directory: Path, base_name: str, payload: ResultPayload) -> dict[str, Any]: """Persist every supported payload kind identically at collection and combo scope.""" meta: dict[str, Any] = {"kind": int(payload.kind), "name": payload.processing_name} def save_array(field: str, array: np.ndarray) -> None: filename = f"{base_name}_{field}.npy" np.save(directory / filename, array, allow_pickle=False) meta[f"{field}_file"] = filename if payload.kind == 1: freq = np.asarray(payload.frequency_hz, dtype=np.float32) trace = np.asarray(payload.trace, dtype=np.complex64) if freq.ndim != 1 or trace.ndim != 1 or freq.shape != trace.shape: raise ValueError("Result trace frequency and values must be equally sized 1D arrays") save_array("freq", freq) save_array("trace", trace) meta["points"] = int(freq.size) elif payload.kind == 2: save_array("scalar", np.asarray([payload.scalar_value], dtype=np.float32)) meta["scalar_value"] = float(payload.scalar_value) elif payload.kind == 3: x_axis = np.asarray(payload.image_x_axis, dtype=np.float32) y_axis = np.asarray(payload.image_y_axis, dtype=np.float32) image = np.asarray(payload.image, dtype=np.float32) if x_axis.ndim != 1 or y_axis.ndim != 1 or image.shape != (y_axis.size, x_axis.size): raise ValueError("Result image axis sizes must match its 2D matrix shape") save_array("x_axis", x_axis) save_array("y_axis", y_axis) save_array("image", image) meta.update(x_points=int(x_axis.size), y_points=int(y_axis.size)) elif payload.kind == 4: table = np.asarray(payload.table, dtype=np.float32) if table.ndim != 2: raise ValueError("Result table payload must be a 2D matrix") save_array("table", table) meta.update(rows=int(table.shape[0]), columns=int(table.shape[1])) else: raise ValueError(f"Unsupported payload kind: {payload.kind}") return meta def save_result_history_numpy( stage_dir: Path, history: list[ResultCollection], *, index_offset: int = 0 ) -> None: """Write result payloads, preserving per-stage indices across streaming chunks.""" stage_dir.mkdir(parents=True, exist_ok=True) logger.debug("Writing %d NumPy result collection(s) to %s", len(history), stage_dir) for index, collection in enumerate(history, start=index_offset): collection_dir = stage_dir / collection_dir_name(index, collection.collection_id, collection.monotonic_ns) with atomic_snapshot_directory(collection_dir) as collection_dir: collection_payload_meta = [] for payload_index, payload in enumerate(collection.collection_payloads): name = sanitize_path_component(payload.processing_name or "processor") base_name = f"collection_{payload_index:03d}_{name}_kind{payload.kind}" collection_payload_meta.append(_save_result_payload(collection_dir, base_name, payload)) blocks_meta = [] for block_index, block in enumerate(collection.blocks): block_dir = collection_dir / f"block_{block_index:03d}_i{block.combo.input}_o{block.combo.output}" block_dir.mkdir(parents=True, exist_ok=False) payload_meta = [] for payload_index, payload in enumerate(block.payloads): name = sanitize_path_component(payload.processing_name or "processor") base_name = f"{payload_index:03d}_{name}_kind{payload.kind}" payload_meta.append(_save_result_payload(block_dir, base_name, payload)) blocks_meta.append({ "input": int(block.combo.input), "output": int(block.combo.output), "payload_count": len(block.payloads), "dir": block_dir.name, "payloads": payload_meta, }) (collection_dir / "meta.json").write_text( json.dumps({ "collection_id": int(collection.collection_id), "monotonic_ns": int(collection.monotonic_ns), "processing_duration_ns": int(collection.processing_duration_ns), "settings_revision": int(collection.settings_revision), "capture_start_ns": int(collection.capture_start_ns), "capture_end_ns": int(collection.capture_end_ns), "speed_m_s": float(collection.speed_m_s), "collection_payload_count": len(collection.collection_payloads), "collection_payloads": collection_payload_meta, "block_count": len(collection.blocks), "blocks": blocks_meta, }, indent=2), encoding="utf-8", )