Files
radar_system/python_app/scripts/replay_raw_producer.py
T
2026-10-06 12:32:26 +03:00

258 lines
11 KiB
Python

"""Raw producer that replays a recorded snapshot dataset (no hardware).
Selected by the supervisor instead of the model's normal acquisition producer
when ``radar.replay.dataset_dir`` is non-empty. The dataset is a directory of
snapshot folders named ``NNNN_idK_ns<timestamp>`` (the VNA_AlGo
preprocessed-capture layout), each holding ``i{input}_o{output}_s21.npy`` and
``i{input}_o{output}_freq.npy`` files for every run combo.
The recorded data is already background-subtracted, so the profile must pair
this producer with NEUTRAL preprocess sets (calibration = 1, reference = 0,
no notch) — then the C++ preprocessor passes the traces through bit-exactly
and the processors (e.g. sar3d) see the recorded scene.
Timeline: snapshot timestamps are re-anchored to this process's monotonic
clock, preserving the recorded inter-frame gaps. On loop, the seam advances by
the median inter-frame gap, so motion-model consumers (sar3d aperture offsets)
never see time jump backwards. ``collection_id`` increases forever.
"""
from __future__ import annotations
from contextlib import ExitStack
import json
import logging
from pathlib import Path
import re
import statistics
import time
import numpy as np
from python_app.models.dataset_model import ComboKey, SweepCollection, TraceData
from python_app.orchestration.shm.decoder import RAW_MAGIC
from python_app.scripts.producer_common import RawPublisher, load_producer_config
from python_app.storage.npz.serialize import serialize_trace_collection
logger = logging.getLogger(__name__)
_SNAPSHOT_NAME = re.compile(r"^(\d+)_id\d+_ns(\d+)$")
class ReplayDataset:
"""Recorded snapshots preloaded into memory, with their recorded timeline."""
def __init__(self, snapshots: list[list[TraceData]], timestamps_ns: list[int]) -> None:
self.snapshots = snapshots
self.timestamps_ns = timestamps_ns
gaps = [b - a for a, b in zip(timestamps_ns, timestamps_ns[1:])]
# Median gap doubles as the loop-seam step and the capture-window span;
# a single-snapshot dataset falls back to an arbitrary sane cadence.
self.median_gap_ns = int(statistics.median(gaps)) if gaps else 160_000_000
def __len__(self) -> int:
return len(self.snapshots)
def _load_dataset(dataset_dir: Path, combos: list[ComboKey], expected_points: int) -> ReplayDataset:
"""Preload every snapshot's traces; fail loudly on a malformed dataset."""
if not dataset_dir.is_dir():
raise RuntimeError(f"radar.replay.dataset_dir is not a directory: {dataset_dir}")
named: list[tuple[int, str, int]] = []
for entry in dataset_dir.iterdir():
match = _SNAPSHOT_NAME.match(entry.name)
if entry.is_dir() and match:
named.append((int(match.group(1)), entry.name, int(match.group(2))))
if not named:
raise RuntimeError(
f"No snapshot folders (NNNN_idK_ns<ts>) found in {dataset_dir}"
)
snapshots: list[list[TraceData]] = []
timestamps_ns: list[int] = []
for _, name, timestamp_ns in sorted(named):
folder = dataset_dir / name
# Recorded per-trace capture windows from the snapshot's meta.json: the
# sar3d segmented timing model reconstructs the acquisition segments
# from them. Missing/partial meta leaves the zero "unmeasured" sentinel
# and the consumer falls back to its uniform-segment model.
windows: dict[tuple[int, int], tuple[int, int]] = {}
meta_path = folder / "meta.json"
if meta_path.is_file():
try:
meta = json.loads(meta_path.read_text())
for entry in meta.get("traces", []):
windows[(int(entry["input"]), int(entry["output"]))] = (
int(entry["capture_start_ns"]),
int(entry["capture_end_ns"]),
)
except (OSError, ValueError, KeyError, TypeError):
windows = {}
traces: list[TraceData] = []
for combo in combos:
s21_path = folder / f"i{combo.input}_o{combo.output}_s21.npy"
freq_path = folder / f"i{combo.input}_o{combo.output}_freq.npy"
if not s21_path.is_file() or not freq_path.is_file():
raise RuntimeError(
f"Snapshot {name} is missing files for combo "
f"input={combo.input} output={combo.output}: expected "
f"{s21_path.name} and {freq_path.name}. Check run.combos vs the dataset."
)
s21 = np.load(s21_path, allow_pickle=False).astype(np.complex64, copy=False)
freq = np.load(freq_path, allow_pickle=False).astype(np.float32, copy=False)
if freq.ndim != 1 or s21.shape != freq.shape:
raise RuntimeError(f"Snapshot {name}: frequency and S21 must be equally sized 1D arrays")
if expected_points and s21.size != expected_points:
raise RuntimeError(
f"{s21_path}: {s21.size} points, but radar.sweep.points = "
f"{expected_points}. Align the profile's sweep with the dataset."
)
window = windows.get((combo.input, combo.output), (0, 0))
traces.append(
TraceData(
combo=combo,
frequency_hz=freq,
s21=s21,
capture_start_ns=window[0],
capture_end_ns=window[1],
)
)
snapshots.append(traces)
timestamps_ns.append(timestamp_ns)
return ReplayDataset(snapshots, timestamps_ns)
def _build_collection(
dataset: ReplayDataset,
index: int,
collection_id: int,
frame_ns: int,
) -> SweepCollection:
"""Wrap one preloaded snapshot in a SweepCollection at the given frame time.
Recorded per-trace capture windows are re-anchored by the same shift as the
frame time, preserving the real intra-frame segment layout. Traces without
recorded windows get a synthetic window spanning the median frame gap.
"""
span_ns = dataset.median_gap_ns
start_ns = frame_ns - span_ns // 2
end_ns = frame_ns + span_ns // 2
shift_ns = frame_ns - dataset.timestamps_ns[index]
traces = []
for source in dataset.snapshots[index]:
if source.capture_end_ns > source.capture_start_ns:
capture = (source.capture_start_ns + shift_ns, source.capture_end_ns + shift_ns)
else:
capture = (start_ns, end_ns)
traces.append(
TraceData(
combo=source.combo,
frequency_hz=source.frequency_hz,
s21=source.s21,
capture_start_ns=capture[0],
capture_end_ns=capture[1],
)
)
return SweepCollection(
collection_id=collection_id,
monotonic_ns=frame_ns,
traces=traces,
capture_start_ns=min(t.capture_start_ns for t in traces),
capture_end_ns=max(t.capture_end_ns for t in traces),
)
def main() -> int:
"""Replay the configured dataset into the raw SHM rings until stopped."""
_, config, stop_requested = load_producer_config("Replay a recorded dataset to raw SHM rings")
replay = config.radar.replay
if not replay.dataset_dir:
raise RuntimeError("replay_raw_producer requires a non-empty radar.replay.dataset_dir")
rate_scale = float(replay.rate_scale) if replay.rate_scale > 0 else 1.0
combos = [ComboKey(input=combo.input, output=combo.output) for combo in config.combos]
# expanduser: profiles are copied between machines, so "~/..." dataset paths
# must resolve against the local home instead of failing as a literal "~".
dataset = _load_dataset(
Path(replay.dataset_dir).expanduser(), combos, int(config.radar.sweep.points)
)
logger.info(
"Replay producer starting: dataset=%s, snapshots=%d, combos=%d, "
"median_gap=%.3f s, loop=%s, rate_scale=%.2f",
replay.dataset_dir, len(dataset), len(combos),
dataset.median_gap_ns / 1e9, replay.loop, rate_scale,
)
cleanup = ExitStack()
try:
publisher = RawPublisher(config.rings, cleanup)
payload_size = len(serialize_trace_collection(
_build_collection(dataset, 0, 1, time.monotonic_ns()), RAW_MAGIC, require_capture_window=True,
))
for name in ("raw", "raw_tap", "preprocessed", "preprocessed_tap"):
ring = getattr(config.rings, name)
if name.endswith("_tap") and not ring.name:
continue
if ring.slot_size_bytes < payload_size:
raise ValueError(f"Replay requires {payload_size} bytes per slot; ring {ring.name} has {ring.slot_size_bytes}")
base_ns = time.monotonic_ns()
wall_start_s = time.monotonic()
first_ts = dataset.timestamps_ns[0]
cycle_span_ns = dataset.timestamps_ns[-1] - first_ts + dataset.median_gap_ns
collection_id = 1
cycle = 0
while not stop_requested.is_set():
for index in range(len(dataset)):
if stop_requested.is_set():
break
# Monotonic replay timeline: recorded offsets re-anchored at
# base_ns, each loop shifted by one full cycle span.
frame_ns = (
base_ns
+ (dataset.timestamps_ns[index] - first_ts)
+ cycle * cycle_span_ns
)
# Publish on the recorded cadence (scaled); waiting via the stop
# event keeps shutdown prompt.
deadline_s = wall_start_s + (frame_ns - base_ns) / 1e9 / rate_scale
delay_s = deadline_s - time.monotonic()
if delay_s > 0 and stop_requested.wait(delay_s):
break
collection = _build_collection(dataset, index, collection_id, frame_ns)
# Acquisition identity uses real monotonic time, even at accelerated
# replay rates. Keep the recorded duration for legacy GPR and the
# virtual per-trace timeline for SAR aperture/motion reconstruction.
duration_ns = collection.capture_end_ns - collection.capture_start_ns
collection.capture_start_ns = time.monotonic_ns()
collection.capture_end_ns = collection.capture_start_ns + duration_ns
if not publisher.publish(collection):
raise RuntimeError(f"Replay raw ring rejected collection {collection_id}")
if collection_id == 1 or collection_id % 50 == 0:
logger.info(
"replay collection %d published (snapshot %d/%d, cycle %d)",
collection_id, index + 1, len(dataset), cycle,
)
collection_id += 1
if not config.runtime.continuous:
stop_requested.set()
break
if not replay.loop:
break
cycle += 1
finally:
cleanup.close()
logger.info("replay raw producer stopped")
return 0
if __name__ == "__main__":
raise SystemExit(main())