scanengine-3/core/saw_check.py

"""Middle-row SAW quality check: acquire one row per angle, then read the
alignment off the frequencies it produces.

Two halves of one test mode, kept together because neither is much use
without the other:

*Acquisition* — ``middle_row_plan`` reduces a full ScanPlan to a single row
per angle, the row-wise middle of the ROI.  ScanEngine runs the result
exactly like any other scan and writes it as a v10 .sras file
(``sras_format.VERSION_SAW_CHECK``), so a check costs one row-time per angle
instead of the hours a full multi-angle scan takes.

*Analysis* — ``frequency_traces`` turns such a file back into one peak-SAW-
frequency trace per angle, and ``alignment_summary`` reduces those to the
numbers the operator is actually asking about.  Both are Qt-free; the plotting
lives in saw_check_viewer.py.

Why the middle row answers an alignment question: ``scan_geometry.build_plan``
centres every angle's rotated bounding box on the same nominal ROI centre, so
each angle's middle row crosses that one point on the sample.  Every angle
therefore measures the same material, and a spread in the per-angle
frequencies is a property of the rig (or of a genuinely anisotropic sample),
not of where each row happened to land.
"""
from __future__ import annotations

from dataclasses import dataclass, field, replace

import numpy as np

from core.scan_geometry import ScanGeometryError, ScanPlan
from core.sras_analysis import ChannelCalibration, compute_rf_image
from core.sras_format import SrasFile

# Rules of thumb for the read-out, not physics.  A well-aligned rig on an
# isotropic sample reads the same frequency at every angle, so the spread of
# the per-angle medians is the alignment signal — but an anisotropic sample
# genuinely varies with angle, so a wide spread is a prompt to look at the
# curves, never a verdict on its own.
SPREAD_GOOD_PCT = 1.0
SPREAD_MARGINAL_PCT = 3.0
# Below this fraction of unmasked pixels a trace is too sparse to read.
VALID_FRACTION_FLOOR = 0.5


# ── Acquisition side ─────────────────────────────────────────────────────────

def middle_row_plan(plan: ScanPlan) -> ScanPlan:
    """Reduce a scan plan to its row-wise middle row at every angle.

    Each angle keeps the geometry the full scan would have used — same
    x_start, x_delta and n_frames from its own rotated bounding box — and
    scans only the middle entry of its row list, so the check samples exactly
    what the scan would along that row.

    An even row count has no exact middle; the upper of the two central rows
    is taken (``n_rows // 2``), which is also the row the viewer picks when it
    reads the middle row out of a full v6 scan.
    """
    if plan.n_angles == 0:
        raise ScanGeometryError("Cannot build a SAW check from a plan with no angles")

    per_angle = []
    for pa in plan.per_angle:
        if not pa.y_positions:
            raise ScanGeometryError(
                f"Angle {pa.angle_deg:.1f}° has no rows, so it has no middle row to check"
            )
        per_angle.append(replace(pa, n_rows=1,
                                 y_positions=[pa.y_positions[middle_row_index(pa.n_rows)]]))
    return replace(plan, per_angle=per_angle)


def middle_row_index(n_rows: int) -> int:
    """The row this check calls the middle one.  One rule, two callers."""
    return max(0, n_rows // 2)


# ── Analysis side ────────────────────────────────────────────────────────────

@dataclass
class AngleTrace:
    """One angle's peak SAW frequency along its middle row.

    ``freq_mhz`` is NaN wherever the pixel was masked out (CH4 DC below the
    threshold), so the gaps stay gaps instead of reading as 0 MHz.
    """
    angle_idx: int
    angle_deg: float
    row_idx: int
    y_mm: float
    x_mm: np.ndarray            # absolute stage X of each frame
    freq_mhz: np.ndarray        # NaN where masked
    _valid: np.ndarray = field(init=False, repr=False)

    def __post_init__(self):
        self._valid = np.isfinite(self.freq_mhz)

    @property
    def offset_mm(self) -> np.ndarray:
        """X relative to the centre of this row.

        Every angle's row is centred on the same ROI centre, so plotting
        against this puts all the angles' curves over the same piece of
        sample — which is the whole point of the comparison.
        """
        if len(self.x_mm) == 0:
            return self.x_mm
        return self.x_mm - 0.5 * (self.x_mm[0] + self.x_mm[-1])

    @property
    def n_valid(self) -> int:
        return int(self._valid.sum())

    @property
    def valid_fraction(self) -> float:
        return self.n_valid / len(self.freq_mhz) if len(self.freq_mhz) else 0.0

    @property
    def median_mhz(self) -> float:
        return float(np.median(self.freq_mhz[self._valid])) if self.n_valid else float("nan")

    @property
    def std_mhz(self) -> float:
        return float(np.std(self.freq_mhz[self._valid])) if self.n_valid > 1 else float("nan")

    @property
    def drift_mhz_per_mm(self) -> float:
        """Least-squares slope of frequency along the row.

        A flat trace means the response did not change across the ROI; a
        sloped one is the signature of a tilt or a defocus the angle spread
        alone would not show.
        """
        if self.n_valid < 2:
            return float("nan")
        x = self.offset_mm[self._valid]
        if np.ptp(x) == 0:
            return float("nan")
        return float(np.polyfit(x, self.freq_mhz[self._valid], 1)[0])


@dataclass
class AlignmentSummary:
    """What the per-angle traces say about the alignment, in scalars."""
    n_angles: int
    median_mhz: float
    spread_mhz: float               # max − min of the per-angle medians
    spread_pct: float               # that spread as a % of the overall median
    best_angle_deg: float           # angle reading the highest median
    worst_angle_deg: float          # angle reading the lowest median
    worst_drift_mhz_per_mm: float
    worst_drift_angle_deg: float
    min_valid_fraction: float

    @property
    def level(self) -> str:
        """"good" / "marginal" / "poor" — see the module's threshold note."""
        if self.n_angles == 0 or not np.isfinite(self.spread_pct):
            return "poor"
        if self.min_valid_fraction < VALID_FRACTION_FLOOR:
            return "poor"
        if self.spread_pct <= SPREAD_GOOD_PCT:
            return "good"
        if self.spread_pct <= SPREAD_MARGINAL_PCT:
            return "marginal"
        return "poor"

    def describe(self) -> str:
        if self.n_angles == 0:
            return "No angle produced a usable frequency trace."
        if self.min_valid_fraction < VALID_FRACTION_FLOOR:
            return (f"Only {self.min_valid_fraction * 100:.0f} % of the worst angle's row "
                    f"is above the DC threshold — check the detection beam and the "
                    f"threshold before reading the spread.")
        return (f"Per-angle medians span {self.spread_mhz:.3f} MHz "
                f"({self.spread_pct:.2f} % of {self.median_mhz:.3f} MHz), "
                f"lowest at {self.worst_angle_deg:.1f}°, highest at {self.best_angle_deg:.1f}°. "
                f"Largest drift along a row: {self.worst_drift_mhz_per_mm:+.3f} MHz/mm "
                f"at {self.worst_drift_angle_deg:.1f}°.")


def frequency_traces(sras: SrasFile, *, dc_threshold_mv: float = 0.0,
                     background: np.ndarray | None = None,
                     gate_start_ns: float | None = None,
                     gate_end_ns: float | None = None,
                     calib: ChannelCalibration | None = None,
                     on_progress=lambda done, total: None) -> list[AngleTrace]:
    """Peak SAW frequency along the middle row of every angle in ``sras``.

    Works on a v10 check (one row per angle, so the middle row is the only
    row) and on a full v6 scan alike — the same middle row the check would
    have acquired is pulled out of the scan, which is what lets a finished
    scan be re-examined with the check's own read-out.

    Angles with nothing on disk (an aborted file) are skipped rather than
    reported as flat zero.
    """
    calib = calib if calib is not None else ChannelCalibration.from_preambles(sras.preambles)
    freq_axis = sras.freq_axis_mhz(sras.header.samples_per_frame)
    time_axis = sras.time_axis_ns()
    statuses = sras.angle_status()

    traces: list[AngleTrace] = []
    for st in statuses:
        on_progress(st.index, len(statuses))
        if st.n_rows_available < 1:
            continue
        pa = sras.per_angle[st.index]
        row = middle_row_index(st.n_rows_available)
        view = sras.load_angle(st.index, n_rows=st.n_rows_available)[row:row + 1]

        img = compute_rf_image(view, calib, freq_axis, dc_threshold_mv,
                               background=background,
                               gate_start_ns=gate_start_ns, gate_end_ns=gate_end_ns,
                               time_axis_ns=time_axis)
        # compute_rf_image zeroes masked pixels and its FFT never peaks in the
        # suppressed DC bin, so 0 MHz means "no reading" and nothing else.
        freq = img[0].astype(np.float64)
        freq[freq <= 0.0] = np.nan

        traces.append(AngleTrace(
            angle_idx=st.index, angle_deg=pa.angle_deg, row_idx=row,
            y_mm=pa.y_positions[row] if row < len(pa.y_positions) else float("nan"),
            x_mm=sras.x_axis_mm(st.index), freq_mhz=freq,
        ))
    on_progress(len(statuses), len(statuses))
    return traces


def alignment_summary(traces: list[AngleTrace]) -> AlignmentSummary:
    """Reduce per-angle traces to the alignment read-out."""
    usable = [t for t in traces if t.n_valid > 0]
    if not usable:
        nan = float("nan")
        return AlignmentSummary(0, nan, nan, nan, nan, nan, nan, nan, 0.0)

    medians = np.array([t.median_mhz for t in usable])
    overall = float(np.median(medians))
    spread = float(medians.max() - medians.min())
    drifts = [(abs(t.drift_mhz_per_mm), t) for t in usable
              if np.isfinite(t.drift_mhz_per_mm)]
    worst_drift = max(drifts, key=lambda d: d[0])[1] if drifts else None

    return AlignmentSummary(
        n_angles=len(usable),
        median_mhz=overall,
        spread_mhz=spread,
        spread_pct=spread / overall * 100.0 if overall else float("nan"),
        best_angle_deg=usable[int(np.argmax(medians))].angle_deg,
        worst_angle_deg=usable[int(np.argmin(medians))].angle_deg,
        worst_drift_mhz_per_mm=worst_drift.drift_mhz_per_mm if worst_drift else float("nan"),
        worst_drift_angle_deg=worst_drift.angle_deg if worst_drift else float("nan"),
        min_valid_fraction=min(t.valid_fraction for t in usable),
    )