199 lines
9.7 KiB
Python
199 lines
9.7 KiB
Python
"""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",
|
|
)
|