393 lines
16 KiB
Python
393 lines
16 KiB
Python
"""Sequential capture workflow for preprocess asset dataset creation."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextlib import suppress
|
|
from dataclasses import dataclass
|
|
import logging
|
|
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
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
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:
|
|
"""Snapshot of sequential capture progress for the GUI/controller."""
|
|
|
|
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
|
|
|
|
logger.info(
|
|
"Sequential capture session created: kind=%s set=%s combos=%d matrix_radar=%s median_sweeps=%d",
|
|
self._kind,
|
|
self._set_name,
|
|
len(self._combos),
|
|
self._is_matrix_radar,
|
|
self._median_sweep_count,
|
|
)
|
|
|
|
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.from_model(config.input_switch)
|
|
self._output_switch = SwitchService.from_model(config.output_switch)
|
|
|
|
@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()
|
|
logger.info("Sequential capture session opened (kind=%s set=%s)", self._kind, self._set_name)
|
|
except Exception:
|
|
logger.exception("Failed to open sequential capture session (kind=%s set=%s)", self._kind, self._set_name)
|
|
self.close()
|
|
raise
|
|
|
|
def close(self) -> None:
|
|
"""Close all opened hardware resources."""
|
|
was_open = self._opened
|
|
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
|
|
if was_open:
|
|
logger.info("Sequential capture session closed (kind=%s set=%s)", self._kind, self._set_name)
|
|
|
|
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
|
|
logger.debug("Captured matrix combo input=%d output=%d", combo.input, combo.output)
|
|
return trace
|
|
|
|
combined_collection = combine_collections_via_median(collections)
|
|
self._traces.extend(combined_collection.traces)
|
|
self._next_index = len(self._combos)
|
|
logger.info("Captured full matrix combo set (%d traces)", len(combined_collection.traces))
|
|
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
|
|
logger.debug(
|
|
"Captured combo input=%d output=%d (%d/%d)",
|
|
combo.input,
|
|
combo.output,
|
|
self._next_index,
|
|
len(self._combos),
|
|
)
|
|
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
|
|
logger.info("Undid matrix combo set capture (kind=%s set=%s)", self._kind, self._set_name)
|
|
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()
|
|
logger.debug(
|
|
"Undid combo capture input=%d output=%d (%d/%d remaining)",
|
|
expected_combo.input,
|
|
expected_combo.output,
|
|
self._next_index,
|
|
len(self._combos),
|
|
)
|
|
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)
|
|
logger.info(
|
|
"Finalized capture set kind=%s set=%s radar_key=%s traces=%d",
|
|
self._kind,
|
|
self._set_name,
|
|
radar_key,
|
|
len(collection.traces),
|
|
)
|
|
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),
|
|
)
|