Rework radar system

python_app/hardware_full/librevna_driver/sweep/assembler.py

@@ -32,12 +32,12 @@ def _find_vna_value(
 ) -> complex:
     """Find one complex receiver value for stage/port/ref tuple."""
     source_mask = 1 << port_index
-    if reference:
-        source_mask |= _REFERENCE_FLAG
 
     for idx, flags in enumerate(datapoint.flags):
         if _extract_stage(int(flags)) != stage:
             continue
+        if bool(int(flags) & _REFERENCE_FLAG) != reference:
+            continue
         if (int(flags) & source_mask) != source_mask:
             continue
         return complex(float(datapoint.real[idx]), float(datapoint.imag[idx]))
@@ -48,6 +48,14 @@ def _find_vna_value(
     )
 
 
+def _s21(datapoint: VNADatapointPacket) -> complex:
+    reference = _find_vna_value(datapoint, stage=0, port_index=0, reference=True)
+    receiver = _find_vna_value(datapoint, stage=0, port_index=1, reference=False)
+    if not np.isfinite(reference) or abs(reference) <= 1e-12 or not np.isfinite(receiver):
+        raise ParseError(f"Invalid receiver/reference at point {datapoint.point_number}")
+    return receiver / reference
+
+
 def datapoint_to_raw_point(
     datapoint: VNADatapointPacket,
     settings: VNASweepSettings,
@@ -55,19 +63,9 @@ def datapoint_to_raw_point(
     num_ports: int,
 ) -> VNARawPoint:
     """Convert low-level VNADatapoint packet into callback-friendly object."""
-    measurements: dict[str, complex] = {}
-    stage_by_excited_port = {port: stage for stage, port in enumerate(settings.excited_ports)}
-
-    for excited_port, stage in stage_by_excited_port.items():
-        ref = _find_vna_value(datapoint, stage=stage, port_index=excited_port - 1, reference=True)
-        for receiver_port in range(1, num_ports + 1):
-            measured = _find_vna_value(
-                datapoint,
-                stage=stage,
-                port_index=receiver_port - 1,
-                reference=False,
-            )
-            measurements[f"S{receiver_port}{excited_port}"] = measured / ref
+    if num_ports < 2:
+        raise ParseError("S21 acquisition requires two receiver ports")
+    measurements = {"S21": _s21(datapoint)}
 
     zero_span = settings.f_start_hz == settings.f_stop_hz and math.isclose(
         settings.power_start_dbm,
@@ -108,17 +106,9 @@ def assemble_vna_sweep(
     x = np.empty(expected_points, dtype=np.float64)
     seen = np.zeros(expected_points, dtype=bool)
 
-    stage_by_excited_port = {port: stage for stage, port in enumerate(settings.excited_ports)}
-    receiver_limit = min(num_ports, 2)
-    trace_names = [
-        f"s{receiver_port}{excited_port}"
-        for excited_port in settings.excited_ports
-        for receiver_port in range(1, receiver_limit + 1)
-        if receiver_port <= 2 and excited_port <= 2
-    ]
-    traces: dict[str, np.ndarray] = {
-        name: np.empty(expected_points, dtype=np.complex128) for name in trace_names
-    }
+    if num_ports < 2:
+        raise ParseError("S21 acquisition requires two receiver ports")
+    s21 = np.empty(expected_points, dtype=np.complex128)
 
     zero_span = settings.f_start_hz == settings.f_stop_hz and math.isclose(
         settings.power_start_dbm,
@@ -130,7 +120,7 @@ def assemble_vna_sweep(
 
     for datapoint in datapoints:
         idx = datapoint.point_number
-        if idx >= expected_points:
+        if not 0 <= idx < expected_points:
             raise ParseError(f"Received out-of-range point index {idx}, expected < {expected_points}")
 
         if zero_span:
@@ -140,22 +130,7 @@ def assemble_vna_sweep(
         else:
             x[idx] = float(datapoint.frequency_or_time)
 
-        for excited_port, stage in stage_by_excited_port.items():
-            ref = _find_vna_value(datapoint, stage=stage, port_index=excited_port - 1, reference=True)
-            for receiver_port in range(1, num_ports + 1):
-                if receiver_port > 2 or excited_port > 2:
-                    # Public SweepResult intentionally exposes 2-port canonical traces.
-                    continue
-                measured = _find_vna_value(
-                    datapoint,
-                    stage=stage,
-                    port_index=receiver_port - 1,
-                    reference=False,
-                )
-                trace_name = f"s{receiver_port}{excited_port}"
-                trace = traces.get(trace_name)
-                if trace is not None:
-                    trace[idx] = measured / ref
+        s21[idx] = _s21(datapoint)
 
         seen[idx] = True
 
@@ -171,7 +146,7 @@ def assemble_vna_sweep(
     elif power_sweep:
         x_label = "power_dbm"
 
-    result = SweepResult(x=x, traces=traces, x_label=x_label)
+    result = SweepResult(x=x, traces={"s21": s21}, x_label=x_label)
     logger.debug(
         "Assembled VNA sweep: points=%d x_label=%s traces=%s",
         len(result.x),