281 lines
12 KiB
Python
281 lines
12 KiB
Python
"""Snapshot selection and filesystem writers for runtime collection histories."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
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
|
|
|
|
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)
|
|
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,
|
|
"trace_count": len(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)
|
|
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,
|
|
"block_count": len(collection.blocks),
|
|
},
|
|
indent=2,
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
|
|
def save_trace_history_numpy(stage_dir: Path, history: list[SweepCollection]) -> None:
|
|
"""Write raw/preprocessed collections as NumPy directory tree."""
|
|
stage_dir.mkdir(parents=True, exist_ok=True)
|
|
for index, collection in enumerate(history):
|
|
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_pos}_o{trace.combo.output_pos}"
|
|
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_pos),
|
|
"output": int(trace.combo.output_pos),
|
|
"points": int(freq.size),
|
|
"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),
|
|
"trace_count": len(collection.traces),
|
|
"traces": traces_meta,
|
|
},
|
|
indent=2,
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
|
|
def save_result_history_numpy(stage_dir: Path, history: list[ResultCollection]) -> None:
|
|
"""Write processed result collections as NumPy directory tree."""
|
|
stage_dir.mkdir(parents=True, exist_ok=True)
|
|
for index, collection in enumerate(history):
|
|
collection_dir = stage_dir / collection_dir_name(index, collection.collection_id, collection.monotonic_ns)
|
|
collection_dir.mkdir(parents=True, exist_ok=False)
|
|
|
|
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_pos}_o{block.combo.output_pos}"
|
|
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),
|
|
}
|
|
)
|
|
|
|
blocks_meta.append(
|
|
{
|
|
"input": int(block.combo.input_pos),
|
|
"output": int(block.combo.output_pos),
|
|
"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),
|
|
"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
|