Files
radar_system/python_app/storage/npz/snapshot_numpy.py

416 lines
19 KiB
Python

"""Snapshot selection and filesystem writers for runtime collection histories."""
from __future__ import annotations
import json
import logging
from pathlib import Path
from typing import Any, TypeVar
import numpy as np
from python_app.models.dataset_model import ResultCollection, SweepCollection
from python_app.storage.npz.paths import collection_dir_name, sanitize_path_component
from python_app.storage.npz.serialize import serialize_result_collection, serialize_trace_collection
logger = logging.getLogger(__name__)
TCollection = TypeVar("TCollection")
def select_aligned_histories(
raw_history: list[SweepCollection],
preprocessed_history: list[SweepCollection],
result_history: list[ResultCollection],
last_n: int,
) -> tuple[list[SweepCollection], list[SweepCollection], list[ResultCollection], dict[str, Any]]:
"""Select aligned tails with `results` as preferred anchor stage."""
raw_index, raw_pos = _index_by_collection_sequence(raw_history)
pre_index, pre_pos = _index_by_collection_sequence(preprocessed_history)
result_index, result_pos = _index_by_collection_sequence(result_history)
if result_index:
ordered_result_keys = sorted(result_index, key=lambda key: result_pos[key])
selected_keys = ordered_result_keys[-last_n:]
selected_raw = [raw_index[key] for key in selected_keys if key in raw_index]
selected_pre = [pre_index[key] for key in selected_keys if key in pre_index]
selected_results = [result_index[key] for key in selected_keys]
raw_missing = len(selected_keys) - len(selected_raw)
pre_missing = len(selected_keys) - len(selected_pre)
return (
selected_raw,
selected_pre,
selected_results,
{
"selection_mode": "result_tail_with_optional_alignment",
"anchor_stage": "results",
"selected_collection_ids": [int(key[0]) for key in selected_keys],
"aligned_key_count": len(selected_keys),
"raw_missing_count": raw_missing,
"preprocessed_missing_count": pre_missing,
},
)
if raw_index:
ordered_raw_keys = sorted(raw_index, key=lambda key: raw_pos[key])
selected_keys = ordered_raw_keys[-last_n:]
selected_raw = [raw_index[key] for key in selected_keys]
selected_pre = [pre_index[key] for key in selected_keys if key in pre_index]
selected_results = [result_index[key] for key in selected_keys if key in result_index]
pre_missing = len(selected_keys) - len(selected_pre)
result_missing = len(selected_keys) - len(selected_results)
return (
selected_raw,
selected_pre,
selected_results,
{
"selection_mode": "raw_tail_with_optional_alignment",
"anchor_stage": "raw",
"selected_collection_ids": [int(key[0]) for key in selected_keys],
"aligned_key_count": len(selected_keys),
"preprocessed_missing_count": pre_missing,
"result_missing_count": result_missing,
},
)
if pre_index:
ordered_pre_keys = sorted(pre_index, key=lambda key: pre_pos[key])
selected_keys = ordered_pre_keys[-last_n:]
selected_raw = [raw_index[key] for key in selected_keys if key in raw_index]
selected_pre = [pre_index[key] for key in selected_keys]
selected_results = [result_index[key] for key in selected_keys if key in result_index]
raw_missing = len(selected_keys) - len(selected_raw)
result_missing = len(selected_keys) - len(selected_results)
return (
selected_raw,
selected_pre,
selected_results,
{
"selection_mode": "preprocessed_tail_with_optional_alignment",
"anchor_stage": "preprocessed",
"selected_collection_ids": [int(key[0]) for key in selected_keys],
"aligned_key_count": len(selected_keys),
"raw_missing_count": raw_missing,
"result_missing_count": result_missing,
},
)
return (
raw_history[-last_n:],
preprocessed_history[-last_n:],
result_history[-last_n:],
{
"selection_mode": "independent_tail",
"anchor_stage": "none",
"selected_collection_ids": [],
"aligned_key_count": 0,
},
)
def save_trace_history_binary(stage_dir: Path, history: list[SweepCollection], magic: int) -> None:
"""Write binary trace history with lightweight metadata sidecars."""
stage_dir.mkdir(parents=True, exist_ok=True)
logger.debug("Writing %d binary trace collection(s) to %s", len(history), stage_dir)
for index, collection in enumerate(history):
binary_path = stage_dir / f"{index:04d}.bin"
metadata_path = stage_dir / f"{index:04d}.json"
binary_path.write_bytes(serialize_trace_collection(collection, magic))
metadata_path.write_text(
json.dumps(
{
"collection_id": collection.collection_id,
"monotonic_ns": collection.monotonic_ns,
"capture_start_ns": int(collection.capture_start_ns),
"capture_end_ns": int(collection.capture_end_ns),
"trace_count": len(collection.traces),
# Also in the .bin trailer; repeated here so per-combo timing is
# readable without decoding the binary payload.
"traces": [
{
"input": int(trace.combo.input),
"output": int(trace.combo.output),
"capture_start_ns": int(trace.capture_start_ns),
"capture_end_ns": int(trace.capture_end_ns),
}
for trace in collection.traces
],
},
indent=2,
),
encoding="utf-8",
)
def save_result_history_binary(stage_dir: Path, history: list[ResultCollection]) -> None:
"""Write binary processed-result history with metadata sidecars."""
stage_dir.mkdir(parents=True, exist_ok=True)
logger.debug("Writing %d binary result collection(s) to %s", len(history), stage_dir)
for index, collection in enumerate(history):
binary_path = stage_dir / f"{index:04d}.bin"
metadata_path = stage_dir / f"{index:04d}.json"
binary_path.write_bytes(serialize_result_collection(collection))
metadata_path.write_text(
json.dumps(
{
"collection_id": collection.collection_id,
"monotonic_ns": collection.monotonic_ns,
"collection_payload_count": len(collection.collection_payloads),
"block_count": len(collection.blocks),
},
indent=2,
),
encoding="utf-8",
)
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.
"""
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)
collection_dir.mkdir(parents=True, exist_ok=False)
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)
s11 = np.asarray(trace.s11, dtype=np.complex64)
s21 = np.asarray(trace.s21, dtype=np.complex64)
np.save(collection_dir / f"{tag}_freq.npy", freq)
np.save(collection_dir / f"{tag}_s11.npy", s11)
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",
"s11_file": f"{tag}_s11.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_history_numpy(
stage_dir: Path, history: list[ResultCollection], *, index_offset: int = 0
) -> None:
"""Write processed result collections as a NumPy directory tree.
``index_offset`` continues the per-collection directory numbering across calls (see
:func:`save_trace_history_numpy`) so streamed chunks append cleanly.
"""
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)
collection_dir.mkdir(parents=True, exist_ok=False)
collection_payload_meta: list[dict[str, int | str | float]] = []
for payload_index, payload in enumerate(collection.collection_payloads):
safe_name = sanitize_path_component(payload.processing_name or "processor")
base_name = f"collection_{payload_index:03d}_{safe_name}_kind{payload.kind}"
if payload.kind == 1:
freq = np.asarray(payload.frequency_hz, dtype=np.float32)
trace = np.asarray(payload.trace, dtype=np.complex64)
np.save(collection_dir / f"{base_name}_freq.npy", freq)
np.save(collection_dir / f"{base_name}_trace.npy", trace)
collection_payload_meta.append(
{
"kind": int(payload.kind),
"name": payload.processing_name,
"points": int(freq.size),
"freq_file": f"{base_name}_freq.npy",
"trace_file": f"{base_name}_trace.npy",
}
)
elif payload.kind == 2:
scalar = np.asarray([float(payload.scalar_value)], dtype=np.float32)
np.save(collection_dir / f"{base_name}_scalar.npy", scalar)
collection_payload_meta.append(
{
"kind": int(payload.kind),
"name": payload.processing_name,
"scalar_file": f"{base_name}_scalar.npy",
"scalar_value": float(payload.scalar_value),
}
)
elif payload.kind == 3:
image_x_axis = np.asarray(payload.image_x_axis, dtype=np.float32)
image_y_axis = np.asarray(payload.image_y_axis, dtype=np.float32)
image = np.asarray(payload.image, dtype=np.float32)
np.save(collection_dir / f"{base_name}_x_axis.npy", image_x_axis)
np.save(collection_dir / f"{base_name}_y_axis.npy", image_y_axis)
np.save(collection_dir / f"{base_name}_image.npy", image)
collection_payload_meta.append(
{
"kind": int(payload.kind),
"name": payload.processing_name,
"x_points": int(image_x_axis.size),
"y_points": int(image_y_axis.size),
"x_axis_file": f"{base_name}_x_axis.npy",
"y_axis_file": f"{base_name}_y_axis.npy",
"image_file": f"{base_name}_image.npy",
}
)
elif payload.kind == 4:
table = np.asarray(payload.table, dtype=np.float32)
np.save(collection_dir / f"{base_name}_table.npy", table)
collection_payload_meta.append(
{
"kind": int(payload.kind),
"name": payload.processing_name,
"rows": int(table.shape[0]) if table.ndim == 2 else 0,
"columns": int(table.shape[1]) if table.ndim == 2 else 0,
"table_file": f"{base_name}_table.npy",
}
)
blocks_meta: list[dict[str, int | str | list[dict[str, int | str | float]]]] = []
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: list[dict[str, int | str | float]] = []
for payload_index, payload in enumerate(block.payloads):
safe_name = sanitize_path_component(payload.processing_name or "processor")
base_name = f"{payload_index:03d}_{safe_name}_kind{payload.kind}"
if payload.kind == 1:
freq = np.asarray(payload.frequency_hz, dtype=np.float32)
trace = np.asarray(payload.trace, dtype=np.complex64)
np.save(block_dir / f"{base_name}_freq.npy", freq)
np.save(block_dir / f"{base_name}_trace.npy", trace)
payload_meta.append(
{
"kind": int(payload.kind),
"name": payload.processing_name,
"points": int(freq.size),
"freq_file": f"{base_name}_freq.npy",
"trace_file": f"{base_name}_trace.npy",
}
)
elif payload.kind == 2:
scalar = np.asarray([float(payload.scalar_value)], dtype=np.float32)
np.save(block_dir / f"{base_name}_scalar.npy", scalar)
payload_meta.append(
{
"kind": int(payload.kind),
"name": payload.processing_name,
"scalar_file": f"{base_name}_scalar.npy",
"scalar_value": float(payload.scalar_value),
}
)
elif payload.kind == 3:
image_x_axis = np.asarray(payload.image_x_axis, dtype=np.float32)
image_y_axis = np.asarray(payload.image_y_axis, dtype=np.float32)
image = np.asarray(payload.image, dtype=np.float32)
np.save(block_dir / f"{base_name}_x_axis.npy", image_x_axis)
np.save(block_dir / f"{base_name}_y_axis.npy", image_y_axis)
np.save(block_dir / f"{base_name}_image.npy", image)
payload_meta.append(
{
"kind": int(payload.kind),
"name": payload.processing_name,
"x_points": int(image_x_axis.size),
"y_points": int(image_y_axis.size),
"x_axis_file": f"{base_name}_x_axis.npy",
"y_axis_file": f"{base_name}_y_axis.npy",
"image_file": f"{base_name}_image.npy",
}
)
elif payload.kind == 4:
table = np.asarray(payload.table, dtype=np.float32)
np.save(block_dir / f"{base_name}_table.npy", table)
payload_meta.append(
{
"kind": int(payload.kind),
"name": payload.processing_name,
"rows": int(table.shape[0]) if table.ndim == 2 else 0,
"columns": int(table.shape[1]) if table.ndim == 2 else 0,
"table_file": f"{base_name}_table.npy",
}
)
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),
"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",
)
def validate_snapshot_name(snapshot_name: str) -> str:
"""Normalize snapshot directory stem."""
return sanitize_path_component(snapshot_name)
def _index_by_collection_sequence(
history: list[TCollection],
) -> tuple[dict[tuple[int, int], TCollection], dict[tuple[int, int], int]]:
"""Index collections by latest-first `(collection_id, occurrence)` preserving history position.
Occurrence is counted from the tail (newest item is occurrence `0` for its
collection id). This avoids cross-run misalignment when collection ids are
reused after restarts and one stage keeps longer history than another.
"""
indexed: dict[tuple[int, int], TCollection] = {}
positions: dict[tuple[int, int], int] = {}
seen_count_from_tail_by_id: dict[int, int] = {}
# Walk from newest to oldest so occurrence=0 always means the most recent
# instance for the given collection id.
for index in range(len(history) - 1, -1, -1):
item = history[index]
collection_id = int(getattr(item, "collection_id"))
occurrence = seen_count_from_tail_by_id.get(collection_id, 0)
key = (collection_id, occurrence)
indexed[key] = item
positions[key] = index
seen_count_from_tail_by_id[collection_id] = occurrence + 1
return indexed, positions