init commit

python_app/hardware_full/librevna_driver/sweep/assembler.py

@@ -0,0 +1,181 @@
+"""Helpers to assemble high-level sweep results from packet streams."""
+
+from __future__ import annotations
+
+import logging
+import math
+from typing import Iterable
+
+import numpy as np
+
+from ..exceptions import IncompleteSweepError, ParseError
+from ..enums import SweepKind
+from ..models import VNADatapointPacket, VNARawPoint, VNASweepSettings, SweepResult
+
+logger = logging.getLogger(__name__)
+
+
+_REFERENCE_FLAG = 0x10
+
+
+def _extract_stage(flags: int) -> int:
+    """Return stage index encoded inside datapoint flag byte."""
+    return flags >> 5
+
+
+def _find_vna_value(
+    datapoint: VNADatapointPacket,
+    *,
+    stage: int,
+    port_index: int,
+    reference: bool,
+) -> 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 (int(flags) & source_mask) != source_mask:
+            continue
+        return complex(float(datapoint.real[idx]), float(datapoint.imag[idx]))
+
+    kind = "reference" if reference else "receiver"
+    raise ParseError(
+        f"Missing {kind} value for stage={stage}, port_index={port_index}, point={datapoint.point_number}"
+    )
+
+
+def datapoint_to_raw_point(
+    datapoint: VNADatapointPacket,
+    settings: VNASweepSettings,
+    *,
+    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
+
+    zero_span = settings.f_start_hz == settings.f_stop_hz and math.isclose(
+        settings.power_start_dbm,
+        settings.power_stop_dbm,
+        rel_tol=0.0,
+        abs_tol=0.0,
+    )
+
+    if zero_span:
+        return VNARawPoint(
+            point_number=datapoint.point_number,
+            frequency_hz=None,
+            time_s=float(datapoint.frequency_or_time) * 1e-6,
+            power_dbm=None,
+            measurements=measurements,
+        )
+
+    return VNARawPoint(
+        point_number=datapoint.point_number,
+        frequency_hz=float(datapoint.frequency_or_time),
+        time_s=None,
+        power_dbm=float(datapoint.cdbm) / 100.0,
+        measurements=measurements,
+    )
+
+
+def assemble_vna_sweep(
+    datapoints: Iterable[VNADatapointPacket],
+    settings: VNASweepSettings,
+    *,
+    num_ports: int,
+    expected_points: int,
+) -> SweepResult:
+    """Build a complete VNA sweep from raw datapoints."""
+    if expected_points <= 0:
+        raise ValueError("expected_points must be > 0")
+
+    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
+    }
+
+    zero_span = settings.f_start_hz == settings.f_stop_hz and math.isclose(
+        settings.power_start_dbm,
+        settings.power_stop_dbm,
+        rel_tol=0.0,
+        abs_tol=0.0,
+    )
+    power_sweep = settings.kind == SweepKind.POWER
+
+    for datapoint in datapoints:
+        idx = datapoint.point_number
+        if idx >= expected_points:
+            raise ParseError(f"Received out-of-range point index {idx}, expected < {expected_points}")
+
+        if zero_span:
+            x[idx] = float(datapoint.frequency_or_time) * 1e-6
+        elif power_sweep:
+            x[idx] = float(datapoint.cdbm) / 100.0
+        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
+
+        seen[idx] = True
+
+    missing = np.flatnonzero(~seen)
+    if missing.size > 0:
+        raise IncompleteSweepError(
+            f"Sweep is incomplete: received {int(seen.sum())}/{expected_points} points"
+        )
+
+    x_label = "frequency_hz"
+    if zero_span:
+        x_label = "time_s"
+    elif power_sweep:
+        x_label = "power_dbm"
+
+    result = SweepResult(x=x, traces=traces, x_label=x_label)
+    logger.debug(
+        "Assembled VNA sweep: points=%d x_label=%s traces=%s",
+        len(result.x),
+        result.x_label,
+        sorted(result.traces.keys()),
+    )
+    return result