"""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), )