sras-viewer/tests/test_alignment.py

"""Angle-alignment tests: does registration actually stack the scans?

Builds a synthetic scan in which one sample is imaged at several *known*
rotations and offsets (tools/make_test_sras.write_rotating) and checks that the
alignment path recovers them, that the shared canvas is angle 0's own pixel
grid extended, and that nothing in the result depends on any other angle's
stage coordinates.

No Qt — this exercises sras_compute directly. See tests/test_gui.py for the
dialog and Aligned-View plumbing.
"""

from types import SimpleNamespace

import numpy as np
import pytest

import sras_compute as compute
from sras_format import CH4_IDX, SrasFile, adc_to_mv
import tools.make_test_sras as gen

# Registration is limited by how far a feature moves per degree: with this
# sample's ~1 mm radius and a ~16 µm registration pitch, a quarter degree is
# already sub-pixel, so it is the floor of what any metric can resolve here.
_ROT_TOL_DEG = 0.5
_SHIFT_TOL_MM = 0.02
_STACK_IOU_MIN = 0.90
_THRESHOLD_MV = 80.0


def dc4_images(sras: SrasFile) -> dict[int, np.ndarray]:
    return {a: adc_to_mv(compute.compute_dc_image(sras, a, CH4_IDX), *sras.cal(CH4_IDX))
            for a in range(sras.n_angles)}


def mm_transform(sras: SrasFile, result, angle_idx: int) -> np.ndarray:
    """Recover the pure mm-space rotation from a canvas->raw affine.

    matrix == D @ R^T @ A_out, where A_out and D only carry the canvas and
    per-angle pixel pitches; undoing both must leave something orthonormal, or
    the transform is smuggling in a scale or a shear.
    """
    dx_a, dy_a = compute.pixel_pitch_mm(sras, angle_idx)
    A_out = np.array([[0.0, result.canvas_dx_mm], [result.canvas_dy_mm, 0.0]])
    D = np.array([[0.0, 1.0 / dy_a], [1.0 / dx_a, 0.0]])
    return np.linalg.inv(D) @ result.per_angle[angle_idx].matrix @ np.linalg.inv(A_out)


@pytest.fixture(scope="module")
def rig(tmp_path_factory):
    """The rotating-sample scan plus everything computed from it once."""
    tmpdir = tmp_path_factory.mktemp("sras_align")
    path = tmpdir / "rotating.sras"
    meta = gen.write_rotating(path, n_angles=5)
    sras = SrasFile(str(path))
    dc4 = dc4_images(sras)
    fits = {a: compute.register_angle_to_reference(
                sras, a, 0, dc4, dc_threshold_mv=_THRESHOLD_MV)
            for a in range(sras.n_angles)}
    result = compute.compute_angle_alignment(sras, 0, _THRESHOLD_MV)
    return SimpleNamespace(path=path, sras=sras, truth=meta["truth"],
                           dc4=dc4, fits=fits, result=result)


def test_registration_recovers_truth(rig):
    """Per-angle rigid registration (rotation + translation, no scale)."""
    for a, fit in rig.fits.items():
        t_rot, t_shift = rig.truth[a]
        rot_err = abs(fit.rotation_deg - t_rot)
        shift_err = float(np.hypot(fit.shift_mm[0] - t_shift[0],
                                   fit.shift_mm[1] - t_shift[1]))
        assert rot_err <= _ROT_TOL_DEG, \
            (f"angle {a}: got {fit.rotation_deg:.3f}°, truth {t_rot:.3f}° "
             f"(err {rot_err:.3f}°)")
        assert shift_err <= _SHIFT_TOL_MM, f"angle {a}: err {shift_err:.4f} mm"
    assert rig.fits[0] == compute.RigidFit(0.0, (0.0, 0.0), 1.0, "reference"), \
        "reference angle registers as exact identity"


def test_stage_angle_sign_is_not_trusted(rig):
    # The stage's rotational sense relative to this module's math-positive
    # convention is not knowable from the file, and the old code hardcoded a
    # guess. Flipping every reported angle must therefore change nothing: the
    # search scores both signs and the images decide.
    flipped = SrasFile(str(rig.path))
    flipped.angles_deg = -flipped.angles_deg
    flipped_fits = {a: compute.register_angle_to_reference(
                        flipped, a, 0, rig.dc4, dc_threshold_mv=_THRESHOLD_MV)
                    for a in range(1, flipped.n_angles)}
    mismatches = {a: (flipped_fits[a].rotation_deg, rig.fits[a].rotation_deg)
                  for a in flipped_fits if flipped_fits[a] != rig.fits[a]}
    assert not mismatches, \
        f"negating every reported stage angle changed fits: {mismatches}"


def test_stage_coordinates_are_not_consulted(rig):
    # Move every non-reference angle's scan window somewhere else entirely.
    # Only angle 0's coordinates may matter, so every fit must be untouched.
    moved = SrasFile(str(rig.path))
    for a in range(1, moved.n_angles):
        moved.x_start_mm[a] += 13.5 * a
        moved.y_pos_per_angle[a] = moved.y_pos_per_angle[a] - 9.25 * a
    moved_dc4 = dc4_images(moved)
    moved_fits = {a: compute.register_angle_to_reference(
                      moved, a, 0, moved_dc4, dc_threshold_mv=_THRESHOLD_MV)
                  for a in range(1, moved.n_angles)}
    mismatches = {a: (round(moved_fits[a].rotation_deg, 4), rig.fits[a].rotation_deg)
                  for a in moved_fits if moved_fits[a] != rig.fits[a]}
    assert not mismatches, \
        f"relocating every other angle's scan window changed fits: {mismatches}"


def test_canvas_is_reference_grid_extended(rig):
    sras, result = rig.sras, rig.result
    t0 = result.per_angle[0]
    assert np.allclose(t0.matrix, np.eye(2)), \
        f"angle 0's transform has rotation/scale/shear: {t0.matrix}"
    assert np.allclose(t0.offset, np.round(t0.offset)), \
        f"angle 0 does not land on whole canvas pixels: {t0.offset}"
    assert ((result.canvas_dx_mm, result.canvas_dy_mm)
            == compute.pixel_pitch_mm(sras, 0)), \
        "canvas pitch is angle 0's own pitch"

    n_rows, n_cols = result.canvas_shape
    x_axis = result.canvas_origin_mm[0] + np.arange(n_cols) * result.canvas_dx_mm
    y_axis = result.canvas_origin_mm[1] + np.arange(n_rows) * result.canvas_dy_mm
    row0, col0 = int(round(-t0.offset[0])), int(round(-t0.offset[1]))
    a0_rows, a0_cols = sras.image_shape(0)
    assert np.allclose(x_axis[col0:col0 + a0_cols], sras.x_axis_mm(0)), \
        "canvas X axis reproduces angle 0's own X coordinates"
    assert np.allclose(y_axis[row0:row0 + a0_rows], sras.y_positions_mm(0)), \
        "canvas Y axis reproduces angle 0's own Y coordinates"
    assert (n_rows >= max(int(sras.n_rows[a]) for a in range(sras.n_angles))
            and n_cols >= max(int(sras.n_frames[a]) for a in range(sras.n_angles))), \
        f"canvas does not cover every angle's footprint: {result.canvas_shape}"


def test_transforms_are_pure_rotations(rig):
    """No scaling anywhere in the per-angle transforms."""
    for a in range(rig.sras.n_angles):
        R = mm_transform(rig.sras, rig.result, a)
        assert (np.allclose(R @ R.T, np.eye(2), atol=1e-9)
                and abs(abs(np.linalg.det(R)) - 1.0) < 1e-9), \
            f"angle {a}: det={np.linalg.det(R):.6f}"


def test_all_angles_stack(rig):
    aligned = {a: compute.apply_alignment(rig.result, a, rig.dc4[a])
               for a in range(rig.sras.n_angles)}
    base = aligned[0] >= _THRESHOLD_MV
    for a in range(1, rig.sras.n_angles):
        other = aligned[a] >= _THRESHOLD_MV
        iou = float((base & other).sum()) / max(1, int((base | other).sum()))
        assert iou >= _STACK_IOU_MIN, f"angle {a}: IoU {iou:.4f}"


def test_downsampled_preview_lands_with_full_res(rig):
    # The wizard reprojects block-mean-downsampled masks, so the
    # affine has to account for the factor. When it did not, every preview
    # layer came out magnified by that factor and offset — the overlay showed a
    # blown-up crop of each mask, which is not something you can align by eye.
    sras, result = rig.sras, rig.result
    pitch = (result.canvas_dx_mm, result.canvas_dy_mm)
    a = sras.n_angles - 1
    p = result.per_angle[a]
    full_mask = (rig.dc4[a] >= _THRESHOLD_MV).astype(np.float32)
    full = compute.reproject_mask(
        sras, a, 0, full_mask, p.rotation_deg, p.shift_mm, pitch,
        result.canvas_origin_mm, result.canvas_shape)
    fy, fx = 4, 16
    small = compute.reproject_mask(
        sras, a, 0, compute.block_mean_2d(full_mask, fy, fx),
        p.rotation_deg, p.shift_mm, (pitch[0] * fx, pitch[1] * fy),
        result.canvas_origin_mm,
        (result.canvas_shape[0] // fy, result.canvas_shape[1] // fx),
        src_downsample=(fy, fx))

    # Compare in mm, via each layer's own center of mass.
    def com_mm(layer, px, py):
        rows, cols = np.nonzero(layer > 0.5)
        return np.array([cols.mean() * px, rows.mean() * py])

    d = com_mm(small, pitch[0] * fx, pitch[1] * fy) - com_mm(full, *pitch)
    assert (abs(d[0]) <= abs(pitch[0] * fx) and abs(d[1]) <= abs(pitch[1] * fy)), \
        f"downsampled preview offset {d[0]:+.4f}, {d[1]:+.4f} mm"


def test_manual_path_reproduces_geometry(rig):
    sras, result = rig.sras, rig.result
    params = {a: compute.ManualAngleParams(t.rotation_deg, t.shift_mm)
              for a, t in result.per_angle.items()}
    manual = compute.build_manual_alignment(sras, 0, _THRESHOLD_MV, params)
    assert (manual.canvas_shape == result.canvas_shape
            and np.allclose(manual.canvas_origin_mm, result.canvas_origin_mm)
            and all(np.allclose(manual.per_angle[a].matrix, result.per_angle[a].matrix)
                    and np.allclose(manual.per_angle[a].offset, result.per_angle[a].offset)
                    for a in range(sras.n_angles))), \
        "build_manual_alignment matches compute_angle_alignment for the same params"


def test_sidecar_roundtrip(rig):
    sras, result = rig.sras, rig.result
    params = {a: compute.ManualAngleParams(t.rotation_deg, t.shift_mm)
              for a, t in result.per_angle.items()}
    compute.save_manual_alignment(sras, 0, _THRESHOLD_MV, params)
    loaded = compute.load_manual_alignment(sras)
    assert (loaded is not None
            and all(np.isclose(loaded.per_angle[a].rotation_deg, params[a].rotation_deg)
                    and np.allclose(loaded.per_angle[a].shift_mm, params[a].shift_mm)
                    for a in range(sras.n_angles))), \
        "sidecar reloads every angle's params"
    assert compute.delete_manual_alignment(sras), "sidecar deletes cleanly"