radar_system/python_app/hardware_full/kamil_adc/processing.py

"""Signal processing for the Kamil ADC reference-channel acquisition.

Each sweep arrives as two aligned complex arrays — the *main* signal and a
*reference* signal — sampled at the same step indices (see
:mod:`python_app.hardware_full.kamil_adc.protocol`). Turning that into a stable,
comparable S21 trace is a fixed three-stage pipeline:

1. **Frequency axis from the reference phase.** The reference arm has a constant
   electrical delay, so its unwrapped phase is an affine function of frequency.
   Two fixed calibration anchors ``(phase0, freq0)`` and ``(phase1, freq1)`` —
   supplied by config, *never* derived from the live sweep — define that law::

       f(phase) = freq0 + (phase - phase0) * (freq1 - freq0) / (phase1 - phase0)

   Trigger jitter shifts every sample's absolute phase together, so the measured
   band floats from sweep to sweep around the fixed calibration. ``np.unwrap``
   anchors each sweep's ramp to ``np.angle(reference[0])`` on the (-pi, pi] branch,
   so once that float carries the anchor across the +/-pi cut the whole ramp jumps
   one 2*pi turn (~117 MHz on the rig) even though nothing physical moved. The
   genuine float is slow (well under pi between consecutive sweeps) while the wrap
   is a discrete 2*pi step, so the processor *unwraps the anchor across sweeps*:
   each sweep's anchor is snapped onto the branch nearest the previous accepted
   sweep (the first sweep onto the calibration ``phase0`` branch). Validated on
   10 000 live sweeps: 506 branch wraps, yet the largest cross-sweep anchor step
   stayed at 1.6 rad (< pi), and the correction cut badly-distorted pass-through
   sweeps from 580 to 11 while leaving clean sweeps untouched.

2. **Amplitude normalization.** ``S = main / |reference|`` divides out the
   stimulus amplitude. Only the magnitude is removed; the reference phase is used
   solely for the axis above.

3. **Crop + resample onto a fixed grid.** The floating per-sweep axis is resampled
   (linear, on real and imaginary parts) onto a single hardcoded grid
   ``linspace(band_start, band_stop, band_points)``. Every sweep then shares a
   byte-identical frequency axis, so traces can be averaged and subtracted. A
   sweep whose floated range does not fully span the band is *rejected* rather
   than edge-extrapolated.

Everything here is pure and free of I/O so it can be unit-tested in isolation.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from python_app.models.run_config_schema import KamilAdcModel

# A reference sample whose magnitude falls to (effectively) zero carries no usable
# amplitude or phase; such points are dropped before normalization rather than
# producing a division blow-up. The threshold only guards genuine zeros — real
# reference frames are emitted from settled measurements and sit far above it.
_REFERENCE_AMPLITUDE_FLOOR = 1e-9

# Linear interpolation onto the band needs at least two distinct source samples;
# a sweep yielding fewer usable points is malformed and rejected.
_MIN_USABLE_POINTS = 2

# One full turn of reference phase. ``np.unwrap`` anchors each sweep's phase ramp
# to the raw angle of the first sample on the (-pi, pi] branch, so a stray sweep
# whose anchor crossed the +/-pi cut is offset by exactly this; the cross-sweep
# anchor tracking snaps it back (see ``KamilAdcSweepProcessor._align_phase_branch``).
_PHASE_BRANCH_PERIOD_RAD = 2.0 * np.pi


@dataclass(frozen=True, slots=True)
class KamilAdcProcessingParams:
    """Immutable phase→frequency calibration and output-grid definition.

    Built from :class:`~python_app.models.run_config_schema.KamilAdcModel` via
    :meth:`from_kamil_model`, but kept decoupled from the config types so the
    processing can be exercised with plain numbers in tests.
    """

    phase0_rad: float
    freq0_hz: float
    phase1_rad: float
    freq1_hz: float
    band_start_hz: float
    band_stop_hz: float
    band_points: int
    reference_amplitude_floor: float = _REFERENCE_AMPLITUDE_FLOOR

    def __post_init__(self) -> None:
        if not np.isfinite([self.phase0_rad, self.phase1_rad, self.freq0_hz, self.freq1_hz]).all():
            raise ValueError("Kamil ADC phase calibration anchors must be finite")
        if self.phase0_rad == self.phase1_rad:
            raise ValueError("Kamil ADC phase calibration anchors must use distinct phases")
        if self.freq0_hz == self.freq1_hz:
            raise ValueError("Kamil ADC phase calibration anchors must map to distinct frequencies")
        if not (np.isfinite(self.band_start_hz) and np.isfinite(self.band_stop_hz)):
            raise ValueError("Kamil ADC band edges must be finite")
        if self.band_stop_hz <= self.band_start_hz:
            raise ValueError("Kamil ADC band stop_hz must be greater than start_hz")
        if self.band_points < 2:
            raise ValueError("Kamil ADC band points must be >= 2")
        if self.reference_amplitude_floor <= 0.0:
            raise ValueError("Kamil ADC reference amplitude floor must be > 0")

    @classmethod
    def from_kamil_model(cls, kamil_adc: KamilAdcModel) -> KamilAdcProcessingParams:
        """Build parameters from the ``radar.kamil_adc`` config section."""
        calibration = kamil_adc.phase_calibration
        band = kamil_adc.band
        return cls(
            phase0_rad=float(calibration.phase0_rad),
            freq0_hz=float(calibration.freq0_hz),
            phase1_rad=float(calibration.phase1_rad),
            freq1_hz=float(calibration.freq1_hz),
            band_start_hz=float(band.start_hz),
            band_stop_hz=float(band.stop_hz),
            band_points=int(band.points),
        )

    @property
    def hz_per_rad(self) -> float:
        """Frequency change per radian of reference phase (the calibration slope)."""
        return (self.freq1_hz - self.freq0_hz) / (self.phase1_rad - self.phase0_rad)


class KamilAdcSweepProcessor:
    """Turns aligned (main, reference) sweeps into S21 traces on a fixed grid.

    The output grid is computed once from the parameters and reused for every
    sweep, so all traces this processor emits share one identical frequency axis.
    """

    __slots__ = ("_params", "_grid_hz", "_previous_anchor_rad")

    def __init__(self, params: KamilAdcProcessingParams) -> None:
        self._params = params
        self._grid_hz = np.linspace(
            params.band_start_hz, params.band_stop_hz, params.band_points, dtype=np.float64
        )
        # Absolute (branch-tracked) reference phase of the last ACCEPTED sweep's
        # first sample, carried across sweeps so a +/-pi anchor wrap can be undone.
        # ``None`` until the first sweep is accepted; reset by building a new
        # processor (i.e. on reconfigure).
        self._previous_anchor_rad: float | None = None

    @property
    def params(self) -> KamilAdcProcessingParams:
        return self._params

    @property
    def grid_hz(self) -> np.ndarray:
        """The fixed output frequency axis (float32), identical for every sweep."""
        return self._grid_hz.astype(np.float32)

    def reference_frequency_axis(self, reference: np.ndarray) -> np.ndarray:
        """Map a reference signal's absolute unwrapped phase to frequency (Hz).

        Stateless and *uncorrected* (no cross-sweep branch tracking), so it shows
        the raw per-sweep axis — used by diagnostics. The live path
        (:meth:`process`) applies the branch correction. Returns frequencies in
        *step order* (not sorted); see the module docstring for the calibration law.
        """
        return self._phase_to_frequency(np.unwrap(np.angle(np.asarray(reference))))

    def _phase_to_frequency(self, phase: np.ndarray) -> np.ndarray:
        """Apply the affine ``phase -> frequency`` calibration law."""
        return self._params.freq0_hz + (phase - self._params.phase0_rad) * self._params.hz_per_rad

    def _align_phase_branch(self, phase: np.ndarray) -> tuple[np.ndarray, float]:
        """Undo a stray +/-pi anchor wrap by unwrapping the anchor across sweeps.

        ``np.unwrap`` pins the whole ramp to ``phase[0]`` on the (-pi, pi] branch,
        so a slow physical float that drags the anchor over the +/-pi cut flips the
        entire sweep by one :data:`_PHASE_BRANCH_PERIOD_RAD`. We snap this sweep's
        anchor onto the branch nearest the previous accepted sweep's anchor (the
        first sweep onto the calibration ``phase0``), then shift the whole ramp by
        the same whole number of turns.

        Returns ``(branch_aligned_phase, anchor_to_commit)``. The caller commits the
        anchor only once the sweep is accepted, so a rejected/corrupt sweep can
        never latch the tracker onto a wrong branch. Genuine sub-pi sweep-to-sweep
        float rounds to zero turns and is preserved untouched.
        """
        anchor = float(phase[0])
        reference_anchor = (
            self._previous_anchor_rad
            if self._previous_anchor_rad is not None
            else self._params.phase0_rad
        )
        turns = round((reference_anchor - anchor) / _PHASE_BRANCH_PERIOD_RAD)
        if turns:
            shift = turns * _PHASE_BRANCH_PERIOD_RAD
            return phase + shift, anchor + shift
        return phase, anchor

    def process(self, main: np.ndarray, reference: np.ndarray) -> np.ndarray | None:
        """Return the S21 trace resampled onto the fixed grid, or ``None`` to reject.

        ``main`` and ``reference`` are equal-length complex arrays ordered by
        ascending step index. ``None`` is returned when the sweep is malformed
        (too few usable points) or does not fully cover the configured band.
        """
        main = np.asarray(main, dtype=np.complex128)
        reference = np.asarray(reference, dtype=np.complex128)
        if main.size < _MIN_USABLE_POINTS or main.size != reference.size:
            return None

        # Frequency axis from the absolute unwrapped reference phase (step order),
        # with a stray +/-pi anchor wrap undone relative to the last accepted sweep.
        # The aligned anchor is committed only if this sweep is accepted (below).
        phase, candidate_anchor = self._align_phase_branch(np.unwrap(np.angle(reference)))
        freqs = self._phase_to_frequency(phase)

        # Amplitude-only normalization; drop points where the reference vanished.
        reference_amplitude = np.abs(reference)
        with np.errstate(divide="ignore", invalid="ignore"):
            s21 = main / reference_amplitude
        usable = (
            (reference_amplitude > self._params.reference_amplitude_floor)
            & np.isfinite(freqs)
            & np.isfinite(s21.real)
            & np.isfinite(s21.imag)
        )
        if int(np.count_nonzero(usable)) < _MIN_USABLE_POINTS:
            return None

        freqs = freqs[usable]
        s21 = s21[usable]

        # Sort onto a monotonically increasing axis (handles either sweep
        # direction) so interpolation and the coverage check are well defined.
        order = np.argsort(freqs, kind="stable")
        freqs = freqs[order]
        s21 = s21[order]

        # Reject sweeps that do not span the whole band: interpolating past the
        # measured edge would inject flat, non-physical points.
        if freqs[0] > self._params.band_start_hz or freqs[-1] < self._params.band_stop_hz:
            return None

        # The sweep is accepted: commit its branch-aligned anchor so the next
        # sweep is tracked relative to it (and a wrap is measured against a real,
        # in-band reference rather than a rejected one).
        self._previous_anchor_rad = candidate_anchor

        # Linear interpolation on real/imaginary parts. Because the band lies
        # within [freqs[0], freqs[-1]], np.interp never extrapolates here.
        real = np.interp(self._grid_hz, freqs, s21.real)
        imag = np.interp(self._grid_hz, freqs, s21.imag)
        return (real + 1j * imag).astype(np.complex64)