sras-viewer/tools/check_equivalence.py

#!/usr/bin/env python3
"""Golden-output equivalence harness for the sras-viewer refactor.

Computes a battery of DC / FFT / alignment outputs and prints a stable hash
for each. Run it on the pre-refactor commit to capture a baseline, then again
after the refactor and diff the two reports — every line must match.

Imports work against both the pre-refactor monolith (`sras_viewer`) and the
post-refactor split (`sras_format` + `sras_compute`), so the *same* script
produces both sides of the comparison.

Hashes canonicalise to native little-endian float64 before hashing, so a
deliberate dtype/byte-order change that preserves values does not show up as
a false mismatch.

Usage:
    python tools/check_equivalence.py --out baseline.txt
    python tools/check_equivalence.py --out after.txt --real /path/to/big.sras
    diff baseline.txt after.txt
"""

import argparse
import copy
import hashlib
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

# --- Import shim: split modules if present, else the monolith --------------
try:
    from sras_format import SrasFile, CH1_IDX, CH3_IDX, CH4_IDX, adc_to_mv
    import sras_compute as C
    _LAYOUT = "split"
except ImportError:
    import sras_viewer as _V
    from sras_viewer import SrasFile, CH1_IDX, CH3_IDX, CH4_IDX, adc_to_mv
    C = _V
    _LAYOUT = "monolith"

compute_dc_image = C.compute_dc_image
compute_rf_image = C.compute_rf_image
compute_alignment = C._compute_angle_alignment
apply_alignment = C.apply_alignment

import tools.make_test_sras as gen  # noqa: E402


def h(arr) -> str:
    """Stable hash of an array's *values*, independent of dtype/byte order."""
    a = np.ascontiguousarray(np.asarray(arr, dtype=np.float64))
    return hashlib.sha256(a.tobytes()).hexdigest()[:16]


def row_slice(sras: SrasFile, angle_idx: int, n_rows: int) -> SrasFile:
    """A shallow view of *sras* restricted to the first *n_rows* rows of
    *angle_idx*, so the huge real file can be exercised in seconds."""
    view = copy.copy(sras)
    n = min(int(n_rows), int(sras.n_rows[angle_idx]))
    view.n_rows = np.array(sras.n_rows, copy=True)
    view.n_rows[angle_idx] = n
    view.data = list(sras.data)
    view.data[angle_idx] = sras.data[angle_idx][:n]
    view._y_pos_per_angle = list(sras._y_pos_per_angle)
    view._y_pos_per_angle[angle_idx] = sras._y_pos_per_angle[angle_idx][:n]
    return view


def report(lines: list[str], label: str, value: str):
    lines.append(f"{label:<58} {value}")


def check_file(path: Path, lines: list[str], tag: str,
               angles: list[int], n_rows: int | None):
    sras = SrasFile(str(path))
    report(lines, f"[{tag}] version/n_angles",
           f"v{sras.version} n={sras.n_angles}")
    report(lines, f"[{tag}] geometry",
           f"rows={list(map(int, sras.n_rows))} frames={list(map(int, sras.n_frames))}")
    report(lines, f"[{tag}] calibration",
           " ".join(f"{m:.6g}/{o:.6g}/{z:.6g}" for m, o, z in
                    zip(sras.ch_ymult_mv, sras.ch_yoff_adc, sras.ch_yzero_mv)))
    report(lines, f"[{tag}] angles_deg", h(sras.angles_deg))
    if sras.background is not None:
        report(lines, f"[{tag}] background", h(sras.background))

    for a in angles:
        if a >= sras.n_angles:
            continue
        s = row_slice(sras, a, n_rows) if n_rows else sras

        report(lines, f"[{tag}] x_axis_mm(a={a})", h(s.x_axis_mm(a)))
        report(lines, f"[{tag}] y_positions_mm(a={a})", h(s.y_positions_mm(a)))

        for ch, name in ((CH3_IDX, "CH3"), (CH4_IDX, "CH4")):
            raw = compute_dc_image(s, a, ch)
            report(lines, f"[{tag}] dc_adc(a={a},{name})", h(raw))
            mv = adc_to_mv(raw, s.ch_ymult_mv[ch], s.ch_yoff_adc[ch],
                           s.ch_yzero_mv[ch])
            report(lines, f"[{tag}] dc_mv(a={a},{name})", h(mv))

        dc4 = adc_to_mv(compute_dc_image(s, a, CH4_IDX),
                        s.ch_ymult_mv[CH4_IDX], s.ch_yoff_adc[CH4_IDX],
                        s.ch_yzero_mv[CH4_IDX])
        thresholds = [-1e9, float(np.median(dc4))]

        for bg in (False, True):
            if bg and s.background is None:
                continue
            for pad in (1, 2):
                n_fft = s.samples_per_frame * pad if pad > 1 else None
                for ti, thr in enumerate(thresholds):
                    img = compute_rf_image(s, a, dc_threshold_mv=thr,
                                           apply_bg_sub=bg, n_fft=n_fft)
                    report(lines,
                           f"[{tag}] rf(a={a},bg={int(bg)},pad={pad},thr{ti})",
                           h(img))
                    # dc4_mv passthrough must give the identical result
                    img2 = compute_rf_image(s, a, dc_threshold_mv=thr,
                                            apply_bg_sub=bg, n_fft=n_fft,
                                            dc4_mv=dc4)
                    same = "SAME" if h(img) == h(img2) else "DIFFER"
                    report(lines,
                           f"[{tag}] rf-dc4arg(a={a},bg={int(bg)},pad={pad},thr{ti})",
                           same)


def check_alignment(path: Path, lines: list[str], tag: str):
    sras = SrasFile(str(path))
    if sras.n_angles < 2:
        return
    dc4 = adc_to_mv(compute_dc_image(sras, 0, CH4_IDX),
                    sras.ch_ymult_mv[CH4_IDX], sras.ch_yoff_adc[CH4_IDX],
                    sras.ch_yzero_mv[CH4_IDX])
    thr = float(np.median(dc4))
    res = compute_alignment(sras, 0, thr)
    report(lines, f"[{tag}] align canvas_shape", str(res.canvas_shape))
    report(lines, f"[{tag}] align canvas_origin",
           f"{res.canvas_origin_mm[0]:.9g},{res.canvas_origin_mm[1]:.9g}")
    report(lines, f"[{tag}] align canvas_pitch",
           f"{res.canvas_dx_mm:.9g},{res.canvas_dy_mm:.9g}")
    for a in sorted(res.per_angle):
        t = res.per_angle[a]
        report(lines, f"[{tag}] align shift_mm(a={a})",
               f"{t.shift_mm[0]:.9g},{t.shift_mm[1]:.9g}")
        report(lines, f"[{tag}] align rot(a={a})", f"{t.rotation_deg:.9g}")
        report(lines, f"[{tag}] align matrix(a={a})", h(t.matrix))
        report(lines, f"[{tag}] align offset(a={a})", h(t.offset))
        img = compute_dc_image(sras, a, CH4_IDX)
        report(lines, f"[{tag}] align resampled(a={a})",
               h(apply_alignment(res, a, img)))


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--out", required=True, help="report file to write")
    p.add_argument("--real", help="optional path to a real .sras file")
    p.add_argument("--real-rows", type=int, default=2,
                   help="rows per angle to sample from the real file")
    p.add_argument("--real-angles", type=int, default=2,
                   help="how many angles to sample from the real file")
    p.add_argument("--scratch", default=".",
                   help="directory for generated synthetic files")
    args = p.parse_args()

    lines = [f"# layout: {_LAYOUT}", f"# numpy: {np.__version__}"]

    scratch = Path(args.scratch)
    synth = scratch / "equiv_synth.sras"
    gen.write(synth, n_angles=4, seed=0, samples_per_frame=64)
    check_file(synth, lines, "synth", angles=[0, 1, 2, 3], n_rows=None)
    check_alignment(synth, lines, "synth")

    # A second synthetic with an odd sample count, to catch off-by-one in
    # rfft bin handling and chunk-boundary arithmetic.
    synth_odd = scratch / "equiv_synth_odd.sras"
    gen.write(synth_odd, n_angles=2, seed=7, samples_per_frame=37)
    check_file(synth_odd, lines, "odd", angles=[0, 1], n_rows=None)

    if args.real:
        real = Path(args.real)
        if real.exists():
            check_file(real, lines, "real",
                       angles=list(range(args.real_angles)),
                       n_rows=args.real_rows)
        else:
            lines.append(f"# real file not found: {real}")

    Path(args.out).write_text("\n".join(lines) + "\n")
    print("\n".join(lines))
    print(f"\nWrote {args.out} ({len(lines)} lines)")


if __name__ == "__main__":
    main()