"""Sequential capture workflow for preprocess asset dataset creation.""" from __future__ import annotations from contextlib import suppress from dataclasses import dataclass import time 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 create_single_radar_service from python_app.hardware_full.switch_service import SwitchService from python_app.models.dataset_model import ComboKey, SweepCollection, TraceData from python_app.models.run_config_model import ComboModel, RunConfigModel from python_app.storage.npz_store import NpzStore, radar_key_from_config MATRIX_RADAR_MANUAL_CAPTURE_KINDS = frozenset({"s21_calibration", "s11_open", "s11_short", "s11_load"}) DEFAULT_CALIBRATION_MEDIAN_SWEEP_COUNT = 5 @dataclass(slots=True) class SequentialCaptureState: """Immutable view of sequential capture progress.""" 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 class SequentialCaptureSession: """Manage hardware and switch stepping for full combo capture sequence.""" def __init__( self, config: RunConfigModel, kind: str, set_name: str, *, median_sweep_count: int = 1, ) -> None: """Create capture session for one preprocess asset set.""" if kind not in {"s21_calibration", "s21_reference", "s11_open", "s11_short", "s11_load", "s11_reference"}: raise RuntimeError(f"Unsupported capture kind: {kind}") if not set_name: raise RuntimeError("Set name is required") if int(median_sweep_count) < 1: raise RuntimeError("median_sweep_count must be >= 1") self._config = config self._kind = kind self._set_name = set_name self._median_sweep_count = int(median_sweep_count) self._is_matrix_radar = config.is_matrix_radar self._manual_matrix_radar_capture = ( self._is_matrix_radar and kind in MATRIX_RADAR_MANUAL_CAPTURE_KINDS ) self._combos = ( RunConfigModel.build_matrix_radar_virtual_combos() if self._is_matrix_radar else RunConfigModel.build_full_combos(config.input_switch.positions, config.output_switch.positions) ) if not self._combos: raise RuntimeError("No switch combinations available for capture") self._traces: list[TraceData] = [] self._next_index = 0 self._opened = False if self._is_matrix_radar: self._radar: MatrixRadarService = create_matrix_radar_service(config) self._input_switch = None self._output_switch = None else: self._radar = create_single_radar_service(config) self._input_switch = SwitchService( name=config.input_switch.name, positions=config.input_switch.positions, default_position=config.input_switch.default_position, mode=config.input_switch.driver_mode, driver=config.input_switch.driver, gpio_chip=config.input_switch.gpio_chip, pin_a=config.input_switch.pin_a, pin_b=config.input_switch.pin_b, invert_logic=config.input_switch.invert_logic, ) self._output_switch = SwitchService( name=config.output_switch.name, positions=config.output_switch.positions, default_position=config.output_switch.default_position, mode=config.output_switch.driver_mode, driver=config.output_switch.driver, gpio_chip=config.output_switch.gpio_chip, pin_a=config.output_switch.pin_a, pin_b=config.output_switch.pin_b, invert_logic=config.output_switch.invert_logic, ) @property def kind(self) -> str: """Return canonical preprocess asset key for this capture session.""" return self._kind @property def set_name(self) -> str: """Return destination set name.""" return self._set_name def open(self) -> None: """Open radar and switch resources.""" if self._opened: return self._opened = True try: self._radar.open() self._radar.configure(self._config.radar.sweep) if self._input_switch is not None: self._input_switch.open() if self._output_switch is not None: self._output_switch.open() except Exception: self.close() raise def close(self) -> None: """Close all opened hardware resources.""" with suppress(Exception): if self._output_switch is not None: self._output_switch.close() with suppress(Exception): if self._input_switch is not None: self._input_switch.close() with suppress(Exception): self._radar.close() self._opened = False def state(self) -> SequentialCaptureState: """Return current progress snapshot.""" current_combo = self._current_combo() return SequentialCaptureState( kind=self._kind, set_name=self._set_name, captured_count=len(self._traces), total_count=len(self._combos), current_combo=current_combo, can_undo=bool(self._traces), is_complete=self.is_complete(), supports_batch_capture=not self._manual_matrix_radar_capture, ) def capture_current_combo(self) -> TraceData: """Capture one trace for current combo and advance sequence cursor.""" if not self._opened: raise RuntimeError("Capture session is not opened") combo = self._current_combo() if combo is None: raise RuntimeError("Capture session is already complete") if self._is_matrix_radar: collections: list[SweepCollection] = [] for _ in range(self._median_sweep_count): collection = self._radar.acquire_collection(collection_id=1) if not collection.traces: raise RuntimeError("Matrix radar capture returned no traces") collections.append(collection) if self._manual_matrix_radar_capture: per_sweep_traces = [select_trace_for_combo(collection, combo) for collection in collections] trace = combine_traces_via_median(per_sweep_traces) self._traces.append(trace) self._next_index += 1 return trace combined_collection = combine_collections_via_median(collections) self._traces.extend(combined_collection.traces) self._next_index = len(self._combos) return combined_collection.traces[-1] if self._input_switch is None or self._output_switch is None: raise RuntimeError("Switches are not initialised for combo capture") self._output_switch.switch_to(combo.output) self._input_switch.switch_to(combo.input) if self._config.runtime.settling_ms > 0: time.sleep(self._config.runtime.settling_ms / 1000.0) sweep_traces: list[TraceData] = [] for _ in range(self._median_sweep_count): sweep = self._radar.acquire() sweep_traces.append( TraceData( combo=ComboKey(input=combo.input, output=combo.output), frequency_hz=np.asarray(sweep.x, dtype=np.float32), s11=np.asarray(sweep.trace("s11"), dtype=np.complex64), s21=np.asarray(sweep.trace("s21"), dtype=np.complex64), ) ) trace = combine_traces_via_median(sweep_traces) self._traces.append(trace) self._next_index += 1 return trace def undo_last_capture(self) -> TraceData: """Remove the most recently captured trace and rewind cursor by one combo.""" if not self._opened: raise RuntimeError("Capture session is not opened") if not self._traces or self._next_index <= 0: raise RuntimeError("No captured combo is available to undo") if self._is_matrix_radar and not self._manual_matrix_radar_capture: if len(self._traces) != len(self._combos): raise RuntimeError("Capture session state is inconsistent; matrix radar trace matrix is incomplete") removed_trace = self._traces[-1] self._traces.clear() self._next_index = 0 return removed_trace expected_combo = self._combos[self._next_index - 1] removed_trace = self._traces[-1] if ( int(removed_trace.combo.input) != int(expected_combo.input) or int(removed_trace.combo.output) != int(expected_combo.output) ): raise RuntimeError("Capture session state is inconsistent; last trace does not match rewind combo") self._next_index -= 1 self._traces.pop() return removed_trace def last_captured_trace(self) -> TraceData | None: """Return the most recently captured trace, if any.""" if not self._traces: return None return self._traces[-1] def captured_traces(self) -> list[TraceData]: """Return captured traces in acquisition order.""" return list(self._traces) def is_complete(self) -> bool: """Return `True` when all combos were captured.""" return self._next_index >= len(self._combos) def finalize(self, store: NpzStore) -> tuple[str, SweepCollection]: """Persist completed capture into store and return radar key + collection.""" if not self.is_complete(): raise RuntimeError("Capture session is not complete") collection = SweepCollection( collection_id=1, monotonic_ns=time.monotonic_ns(), traces=list(self._traces), ) radar_key = radar_key_from_config( model_name=self._config.radar.model, serial=self._config.radar.serial, sweep_start_hz=self._config.radar.sweep.start_hz, sweep_stop_hz=self._config.radar.sweep.stop_hz, sweep_points=self._config.radar.sweep.points, ifbw_hz=self._config.radar.sweep.if_bandwidth_hz, power_dbm=self._config.radar.sweep.power_dbm, extra_serials=self._config.radar_key_extra_parts() or None, ) store.save_set(self._kind, radar_key, self._set_name, collection) return radar_key, collection def _current_combo(self) -> ComboModel | None: """Return next combo to capture, or `None` if session is complete.""" if self._next_index >= len(self._combos): return None return self._combos[self._next_index] def select_trace_for_combo(collection: SweepCollection, combo: ComboModel) -> TraceData: """Return the trace matching a virtual combo from a full multi-device capture.""" for trace in collection.traces: if ( int(trace.combo.input) == int(combo.input) and int(trace.combo.output) == int(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 S11/S21 are 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}" ) s11_stack = np.stack([np.asarray(t.s11, dtype=np.complex64) for t in traces], axis=0) s21_stack = np.stack([np.asarray(t.s21, dtype=np.complex64) for t in traces], axis=0) s11_median = ( np.median(s11_stack.real, axis=0) + 1j * np.median(s11_stack.imag, axis=0) ).astype(np.complex64) s21_median = ( np.median(s21_stack.real, axis=0) + 1j * np.median(s21_stack.imag, axis=0) ).astype(np.complex64) return TraceData( combo=ComboKey(input=int(combo.input), output=int(combo.output)), frequency_hz=np.asarray(first.frequency_hz, dtype=np.float32), s11=s11_median, s21=s21_median, ) 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] 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): if combo_index >= len(collection.traces): raise RuntimeError( f"Median collections have mismatched trace counts at sweep {collection_index}" ) 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=int(reference.capture_start_ns), capture_end_ns=int(reference.capture_end_ns), )