296 lines
13 KiB
Python
296 lines
13 KiB
Python
"""One capture session for one or more radar sweep variants."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextlib import suppress
|
|
from dataclasses import dataclass
|
|
import time
|
|
from typing import cast
|
|
|
|
import numpy as np
|
|
|
|
from python_app.hardware_full.matrix_radar_service import MatrixRadarService, create_matrix_radar_service
|
|
from python_app.hardware_full.single_radar_service import SingleRadarService, create_single_radar_service
|
|
from python_app.hardware_full.switch_service import SwitchService, select_and_settle
|
|
from python_app.models.dataset_model import ComboKey, SweepCollection, TraceData
|
|
from python_app.models.run_config_schema import ComboModel, RunConfigModel
|
|
from python_app.storage.npz.store import NpzStore
|
|
from python_app.workflows.radar_config_variants import RadarConfigVariant
|
|
|
|
DEFAULT_CALIBRATION_MEDIAN_SWEEP_COUNT = 5
|
|
|
|
|
|
def trace_capture_window(traces: list[TraceData]) -> tuple[int, int]:
|
|
"""Span measured trace windows, retaining zero if timing was not reported."""
|
|
return (min((t.capture_start_ns for t in traces if t.capture_start_ns > 0), default=0),
|
|
max((t.capture_end_ns for t in traces if t.capture_end_ns > 0), default=0))
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class SequentialCaptureState:
|
|
kind: str
|
|
set_name: str
|
|
captured_count: int
|
|
total_count: int
|
|
current_combo: ComboModel | None
|
|
can_undo: bool
|
|
is_complete: bool
|
|
variant_count: int = 1
|
|
supports_batch_capture: bool = True
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class CaptureBatch:
|
|
combo: ComboModel
|
|
traces: tuple[TraceData, ...]
|
|
|
|
@property
|
|
def display_trace(self) -> TraceData:
|
|
return self.traces[-1]
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class SavedCaptureSet:
|
|
display_name: str
|
|
radar_key: str
|
|
collection: SweepCollection
|
|
|
|
|
|
class SequentialCaptureSession:
|
|
"""Capture each antenna pair across variants, committing a complete batch atomically."""
|
|
|
|
def __init__(
|
|
self, config: RunConfigModel, kind: str, set_name: str, *,
|
|
radar_variants: list[RadarConfigVariant] | None = None, median_sweep_count: int = 1,
|
|
) -> None:
|
|
if kind not in {"s21_calibration", "s21_reference"}:
|
|
raise RuntimeError(f"Unsupported capture kind: {kind}")
|
|
if not set_name:
|
|
raise RuntimeError("Set name is required")
|
|
if median_sweep_count < 1:
|
|
raise RuntimeError("median_sweep_count must be >= 1")
|
|
self.kind, self.set_name = kind, set_name
|
|
self._config, self._median_sweep_count = config, median_sweep_count
|
|
self._variants = ([RadarConfigVariant("Current radar", config, config.radar_key())]
|
|
if radar_variants is None else list(radar_variants))
|
|
if not self._variants:
|
|
raise RuntimeError("At least one radar variant is required")
|
|
self._combos = config.build_runtime_combos()
|
|
if not self._combos:
|
|
raise RuntimeError("No switch combinations available for capture")
|
|
self._full_matrix = config.is_multi_device and kind == "s21_reference"
|
|
self._batches: list[CaptureBatch] = []
|
|
self._traces: dict[str, list[TraceData]] = {v.radar_key: [] for v in self._variants}
|
|
self._opened = False
|
|
self._sweep = self._variants[0].config.radar.sweep
|
|
self._radar: MatrixRadarService | SingleRadarService
|
|
if config.is_multi_device:
|
|
self._radar = create_matrix_radar_service(config)
|
|
self._input_switch = self._output_switch = None
|
|
else:
|
|
self._radar = create_single_radar_service(config)
|
|
self._input_switch = SwitchService.from_model(config.input_switch)
|
|
self._output_switch = SwitchService.from_model(config.output_switch)
|
|
|
|
def open(self) -> None:
|
|
if self._opened:
|
|
return
|
|
self._opened = True
|
|
try:
|
|
self._radar.open()
|
|
self._radar.configure(self._sweep)
|
|
for switch in (self._input_switch, self._output_switch):
|
|
if switch is not None:
|
|
switch.open()
|
|
except Exception:
|
|
self.close()
|
|
raise
|
|
|
|
def close(self) -> None:
|
|
for resource in (self._output_switch, self._input_switch, self._radar):
|
|
if resource is not None:
|
|
with suppress(Exception):
|
|
resource.close()
|
|
self._opened = False
|
|
|
|
def state(self) -> SequentialCaptureState:
|
|
complete = self.is_complete()
|
|
return SequentialCaptureState(
|
|
self.kind, self.set_name, len(self._combos) if complete else len(self._batches),
|
|
len(self._combos), None if complete else self._combos[len(self._batches)],
|
|
bool(self._batches), complete, len(self._variants),
|
|
not (self._config.is_multi_device and self.kind == "s21_calibration"),
|
|
)
|
|
|
|
def is_complete(self) -> bool:
|
|
return bool(self._batches) if self._full_matrix else len(self._batches) == len(self._combos)
|
|
|
|
def capture_current_combo(self) -> CaptureBatch:
|
|
if not self._opened:
|
|
raise RuntimeError("Capture session is not opened")
|
|
combo = self.state().current_combo
|
|
if combo is None:
|
|
raise RuntimeError("Capture session is already complete")
|
|
if not self._config.is_multi_device:
|
|
select_and_settle(self._output_switch, self._input_switch, combo.output, combo.input,
|
|
self._config.runtime.settling_ms)
|
|
pending: dict[str, list[TraceData]] = {}
|
|
for variant in self._variants:
|
|
if variant.config.radar.sweep != self._sweep:
|
|
self._radar.configure(variant.config.radar.sweep)
|
|
self._sweep = variant.config.radar.sweep
|
|
if self._config.runtime.settling_ms > 0:
|
|
time.sleep(self._config.runtime.settling_ms / 1000.0)
|
|
if self._config.is_multi_device:
|
|
radar = cast(MatrixRadarService, self._radar)
|
|
if self._full_matrix:
|
|
collections = [radar.acquire_collection() for _ in range(self._median_sweep_count)]
|
|
traces = combine_collections_via_median(collections).traces
|
|
if not traces:
|
|
raise RuntimeError("Matrix radar capture returned no traces")
|
|
else:
|
|
traces = [combine_traces_via_median([
|
|
select_trace_for_combo(radar.acquire_collection(combo=ComboKey(combo.input, combo.output)), combo)
|
|
for _ in range(self._median_sweep_count)])]
|
|
else:
|
|
traces = [capture_single_radar_combo(cast(SingleRadarService, self._radar),
|
|
combo, self._median_sweep_count)]
|
|
pending[variant.radar_key] = traces
|
|
for key, traces in pending.items():
|
|
self._traces[key].extend(traces)
|
|
batch = CaptureBatch(combo, tuple(traces[-1] for traces in pending.values()))
|
|
self._batches.append(batch)
|
|
return batch
|
|
|
|
def undo_last_capture(self) -> None:
|
|
if not self._opened or not self._batches:
|
|
raise RuntimeError("No captured combo is available to undo")
|
|
for traces in self._traces.values():
|
|
del traces[-(len(self._combos) if self._full_matrix else 1):]
|
|
self._batches.pop()
|
|
|
|
def last_captured_trace(self) -> TraceData | None:
|
|
return self._batches[-1].display_trace if self._batches else None
|
|
|
|
def captured_batches(self) -> list[CaptureBatch]:
|
|
return list(self._batches)
|
|
|
|
def finalize(self, store: NpzStore) -> list[SavedCaptureSet]:
|
|
if not self.is_complete():
|
|
raise RuntimeError("Capture session is not complete")
|
|
saved = []
|
|
for variant in self._variants:
|
|
traces = list(self._traces[variant.radar_key])
|
|
start, end = trace_capture_window(traces)
|
|
collection = SweepCollection(1, time.monotonic_ns(), traces, start, end)
|
|
store.save_set(self.kind, variant.radar_key, self.set_name, collection)
|
|
saved.append(SavedCaptureSet(variant.display_name, variant.radar_key, collection))
|
|
return saved
|
|
|
|
|
|
def capture_single_radar_combo(radar: SingleRadarService, combo: ComboModel, count: int) -> TraceData:
|
|
radar.prepare_combo(ComboKey(combo.input, combo.output))
|
|
traces = []
|
|
for _ in range(count):
|
|
start = time.monotonic_ns()
|
|
sweep = radar.acquire()
|
|
traces.append(TraceData(
|
|
ComboKey(combo.input, combo.output), np.asarray(sweep.x, dtype=np.float32),
|
|
np.asarray(sweep.trace("s21"), dtype=np.complex64),
|
|
sweep.capture_start_ns or start, sweep.capture_end_ns or time.monotonic_ns(),
|
|
))
|
|
return combine_traces_via_median(traces)
|
|
|
|
|
|
def select_trace_for_combo(collection: SweepCollection, combo: ComboModel) -> TraceData:
|
|
for trace in collection.traces:
|
|
if trace.combo == ComboKey(combo.input, combo.output):
|
|
return trace
|
|
raise RuntimeError(f"Multi-device capture is missing trace for input={combo.input}, output={combo.output}")
|
|
|
|
|
|
def combine_traces_via_median(traces: list[TraceData]) -> TraceData:
|
|
"""Return one trace whose S21 is the per-point median of the inputs.
|
|
|
|
A single-element input is returned unchanged. With multiple inputs, the real
|
|
and imaginary parts of each complex sample are medianed independently so a
|
|
single bad sweep (e.g. an outlier with random phase) is rejected without
|
|
corrupting the saved calibration trace.
|
|
"""
|
|
if not traces:
|
|
raise RuntimeError("Cannot combine empty sweep list")
|
|
if len(traces) == 1:
|
|
return traces[0]
|
|
|
|
first = traces[0]
|
|
combo = first.combo
|
|
point_count = first.frequency_hz.size
|
|
for index, trace in enumerate(traces[1:], start=1):
|
|
if (
|
|
int(trace.combo.input) != int(combo.input)
|
|
or int(trace.combo.output) != int(combo.output)
|
|
):
|
|
raise RuntimeError(
|
|
f"Median combine combo mismatch at sweep {index}: "
|
|
f"({trace.combo.input},{trace.combo.output}) "
|
|
f"vs ({combo.input},{combo.output})"
|
|
)
|
|
if trace.frequency_hz.size != point_count:
|
|
raise RuntimeError(
|
|
f"Median combine point-count mismatch at sweep {index}: "
|
|
f"{trace.frequency_hz.size} vs {point_count}"
|
|
)
|
|
if not np.array_equal(trace.frequency_hz, first.frequency_hz):
|
|
raise RuntimeError(f"Median combine frequency-axis mismatch at sweep {index}")
|
|
|
|
s21_stack = np.stack([np.asarray(t.s21, dtype=np.complex64) for t in traces], axis=0)
|
|
s21_median = (
|
|
np.median(s21_stack.real, axis=0) + 1j * np.median(s21_stack.imag, axis=0)
|
|
).astype(np.complex64)
|
|
# The median is built from every input sweep, so its window spans all of them.
|
|
capture_start_ns, capture_end_ns = trace_capture_window(traces)
|
|
return TraceData(
|
|
combo=ComboKey(input=int(combo.input), output=int(combo.output)),
|
|
frequency_hz=np.asarray(first.frequency_hz, dtype=np.float32),
|
|
s21=s21_median,
|
|
capture_start_ns=capture_start_ns,
|
|
capture_end_ns=capture_end_ns,
|
|
)
|
|
|
|
|
|
def combine_collections_via_median(collections: list[SweepCollection]) -> SweepCollection:
|
|
"""Combine multi-device matrix captures into one collection with per-combo medians."""
|
|
if not collections:
|
|
raise RuntimeError("Cannot combine empty collection list")
|
|
if len(collections) == 1:
|
|
return collections[0]
|
|
|
|
reference = collections[0]
|
|
for index, collection in enumerate(collections[1:], start=1):
|
|
if len(collection.traces) != len(reference.traces):
|
|
raise RuntimeError(f"Median collections have mismatched trace counts at sweep {index}")
|
|
expected_combos = [(trace.combo.input, trace.combo.output) for trace in reference.traces]
|
|
medianed_traces: list[TraceData] = []
|
|
for combo_index, (input_pos, output_pos) in enumerate(expected_combos):
|
|
per_sweep_traces: list[TraceData] = []
|
|
for collection_index, collection in enumerate(collections):
|
|
trace = collection.traces[combo_index]
|
|
if (
|
|
int(trace.combo.input) != int(input_pos)
|
|
or int(trace.combo.output) != int(output_pos)
|
|
):
|
|
raise RuntimeError(
|
|
f"Median collections trace order mismatch at sweep {collection_index}, "
|
|
f"combo index {combo_index}"
|
|
)
|
|
per_sweep_traces.append(trace)
|
|
medianed_traces.append(combine_traces_via_median(per_sweep_traces))
|
|
|
|
return SweepCollection(
|
|
collection_id=int(reference.collection_id),
|
|
monotonic_ns=int(reference.monotonic_ns),
|
|
traces=medianed_traces,
|
|
capture_start_ns=min((int(item.capture_start_ns) for item in collections if item.capture_start_ns > 0), default=0),
|
|
capture_end_ns=max((int(item.capture_end_ns) for item in collections if item.capture_end_ns > 0), default=0),
|
|
)
|