sras-viewer/sras_compute.py

#!/usr/bin/env python3
"""Image computation and angle alignment for .sras scans.

Depends only on numpy/scipy (+ optional pyfftw) and sras_format, so a
multiprocessing child can import it without loading Qt or matplotlib —
which matters because Python 3.14 on macOS spawns rather than forks.
"""

import os
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass

import numpy as np
import scipy.fft as scipy_fft
import scipy.ndimage as scipy_ndimage

from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, SrasFile, adc_to_mv

# ---------------------------------------------------------------------------
# FFT backend
# ---------------------------------------------------------------------------

try:
    import pyfftw
    pyfftw.interfaces.cache.enable()
    PYFFTW_AVAILABLE = True
except ImportError:
    PYFFTW_AVAILABLE = False

_fft_backend = "numpy"   # "numpy" or "pyfftw"; set via set_fft_backend()


def set_fft_backend(name: str):
    """Select the rfft implementation. Module-level state, so it must be set
    explicitly inside each multiprocessing child — it does not survive a
    spawn."""
    global _fft_backend
    _fft_backend = name if (name != "pyfftw" or PYFFTW_AVAILABLE) else "numpy"


def get_fft_backend() -> str:
    return _fft_backend


def _do_rfft(x: np.ndarray, n: int | None = None, axis: int = -1,
             workers: int = 1) -> np.ndarray:
    """Dispatch rfft to the selected backend with optional multithreading."""
    if _fft_backend == "pyfftw" and PYFFTW_AVAILABLE:
        return pyfftw.interfaces.numpy_fft.rfft(x, n=n, axis=axis, threads=workers)
    return scipy_fft.rfft(x, n=n, axis=axis, workers=workers)


# ---------------------------------------------------------------------------
# Chunking / parallel budget
# ---------------------------------------------------------------------------
#
# Rows are batched so the float32 working buffers for one chunk stay under a
# memory budget. A fixed row count (the original design) works fine for small
# legacy scans but is catastrophic for a v6 scan with a large per-angle
# frame/sample count — e.g. a 7500-frame x 2500-sample angle needs ~2.4 GB for
# a single 32-row chunk.
#
# With chunks running concurrently the budget has to cover *all* live chunks at
# once. Note that on a large scan chunk_rows is already clamped to its floor of
# 1 row (one row alone is ~75 MB of float32 at 7507x2500), so shrinking the
# per-chunk size cannot buy more concurrency — the worker count must be derived
# from the budget instead. See _plan_chunks.

_TOTAL_BYTES_BUDGET = int(os.environ.get("SRAS_MEM_BUDGET_MB", 1536)) * 1024 * 1024
_CHUNK_ROWS_MAX = 32           # cap for small scans (original behavior)
_MAX_WORKERS = int(os.environ.get("SRAS_MAX_WORKERS", 0)) or (os.cpu_count() or 4)

# An rfft chunk holds, live at once: the float32 input, the complex64
# transform, and the float32 power spectrum — roughly 3x the input buffer.
_FFT_LIVE_MULTIPLIER = 3


def _chunk_rows_for(n_frames: int, samples_per_frame: int,
                    budget: int = _TOTAL_BYTES_BUDGET) -> int:
    bytes_per_row = max(1, n_frames * samples_per_frame * 4)   # float32
    return int(max(1, min(_CHUNK_ROWS_MAX, budget // bytes_per_row)))


def _plan_chunks(n_rows: int, n_frames: int, samples_per_frame: int,
                 live_multiplier: int = 1,
                 max_workers: int | None = None) -> tuple[int, int]:
    """(chunk_rows, n_workers) sized so live_multiplier x chunk_rows x
    n_frames x samples x 4 bytes x n_workers stays within the budget."""
    per_chunk_budget = max(1, _TOTAL_BYTES_BUDGET // live_multiplier)
    chunk_rows = _chunk_rows_for(n_frames, samples_per_frame, per_chunk_budget)

    chunk_bytes = max(1, live_multiplier * chunk_rows * n_frames * samples_per_frame * 4)
    cap = max_workers if max_workers is not None else _MAX_WORKERS
    n_workers = int(max(1, min(cap,
                               _TOTAL_BYTES_BUDGET // chunk_bytes,
                               -(-n_rows // chunk_rows))))
    return chunk_rows, n_workers


def _map_row_chunks(n_rows: int, chunk_rows: int, n_workers: int, fn):
    """Apply fn(r0, r1) over row chunks, in parallel when it pays.

    Chunks write to disjoint output slices, so no locking is needed. numpy
    ufuncs, scipy's pocketfft, and memmap page faults all release the GIL, so
    threads give real parallelism here — and on a very large file they also
    keep many more page-fault requests in flight, which is what the I/O path
    wants.
    """
    bounds = [(r0, min(r0 + chunk_rows, n_rows))
              for r0 in range(0, n_rows, chunk_rows)]
    if n_workers <= 1 or len(bounds) == 1:
        for r0, r1 in bounds:
            fn(r0, r1)
        return
    with ThreadPoolExecutor(max_workers=min(n_workers, len(bounds))) as pool:
        list(pool.map(lambda b: fn(*b), bounds))


# ---------------------------------------------------------------------------
# Image computation
# ---------------------------------------------------------------------------

def compute_dc_image(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray:
    """Mean of each waveform → (n_rows, n_frames) float32, in ADC counts."""
    n_rows, n_frames = sras.image_shape(angle_idx)
    data = sras.data[angle_idx]
    chunk_rows, n_workers = _plan_chunks(n_rows, n_frames, sras.samples_per_frame)
    img = np.empty((n_rows, n_frames), dtype=np.float32)

    def chunk(r0: int, r1: int):
        img[r0:r1] = data[r0:r1, ch_idx, :, :].astype(np.float32).mean(axis=-1)

    _map_row_chunks(n_rows, chunk_rows, n_workers, chunk)
    return img


def dc_image_mv(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray:
    """DC image for (angle, channel) in mV, preferring a stored v5/v7 cache."""
    cached = sras.cached_dc_mv(angle_idx, ch_idx)
    if cached is not None:
        return cached
    return adc_to_mv(compute_dc_image(sras, angle_idx, ch_idx), *sras.cal(ch_idx))


def compute_rf_image(sras: SrasFile, angle_idx: int,
                     dc_threshold_mv: float | None,
                     apply_bg_sub: bool = True,
                     n_fft: int | None = None,
                     dc4_mv: np.ndarray | None = None) -> np.ndarray:
    """FFT of each CH1 waveform; pixel = peak frequency in MHz.

    Pixels where CH4_dc < dc_threshold_mv are set to 0 — and the FFT is
    never run for them, since that's the expensive part. The masked-out
    CH1 samples are also never *read*: the boolean mask is applied to the
    raw memmap slice before any dtype conversion, so numpy only pages in
    the bytes for pixels that pass the threshold (an untouched memmap page
    is never read from disk).

    *dc_threshold_mv* of None means "mask nothing", and skips reading and
    averaging CH4 entirely — worth a third of the I/O when caching a whole
    file, where the mask is applied later at display time.

    If the DC4 image for this angle is already known (e.g. from the
    DC-channel precompute cache), pass it as *dc4_mv* (mV, shape
    (n_rows, n_frames)) to reuse it instead of re-reading CH4 here.

    Fast path: if the file has a precomputed peak-frequency image for this
    angle (v5 PREC or v7 CACH), zero-padding is off, and the bg-sub flag
    matches, the stored image is used directly — no FFT is run.
    """
    n_rows, n_frames = sras.image_shape(angle_idx)
    data = sras.data[angle_idx]

    # ---- Fast path: precomputed image (v5 PREC or v7 CACH) ----------------
    cached_freq = sras.precomputed_freq_mhz[angle_idx]
    if (cached_freq is not None
            and n_fft is None                     # no custom zero-padding
            and sras.precomputed_bg_sub == (apply_bg_sub and sras.background is not None)):
        freq_img = cached_freq.copy()
        if dc_threshold_mv is not None:
            # DC4 mask, in priority order: already-cached DC block, caller-
            # supplied image, or a fresh (cheap — no FFT) recompute.
            dc4_img = sras.cached_dc_mv(angle_idx, CH4_IDX)
            if dc4_img is None:
                dc4_img = dc4_mv if dc4_mv is not None else adc_to_mv(
                    compute_dc_image(sras, angle_idx, CH4_IDX), *sras.cal(CH4_IDX))
            freq_img[dc4_img < dc_threshold_mv] = 0.0
        return freq_img

    # ---- Chunked FFT path --------------------------------------------------
    freq_axis = sras.freq_axis_mhz(n_fft)
    img = np.zeros((n_rows, n_frames), dtype=np.float32)
    n_fft_bins = n_fft if n_fft is not None else sras.samples_per_frame
    chunk_rows, n_workers = _plan_chunks(
        n_rows, n_frames, max(sras.samples_per_frame, n_fft_bins),
        live_multiplier=_FFT_LIVE_MULTIPLIER)
    background = sras.background if (apply_bg_sub and sras.background is not None) else None
    cal4 = sras.cal(CH4_IDX)

    def chunk(r0: int, r1: int):
        if dc_threshold_mv is None:
            valid = None
        else:
            if dc4_mv is not None:
                dc4_chunk = dc4_mv[r0:r1]
            else:
                dc4_chunk = adc_to_mv(
                    data[r0:r1, CH4_IDX, :, :].astype(np.float32).mean(axis=-1), *cal4)
            valid = dc4_chunk >= dc_threshold_mv   # True = run the FFT
            if not valid.any():
                return

        # Index the raw memmap slice with the boolean mask *before*
        # converting dtype — this is a lazy view until touched, so only the
        # selected elements are actually read from disk; masked-out pixels'
        # pages are never paged in at all.
        raw = data[r0:r1, CH1_IDX, :, :]
        waves = (raw[valid] if valid is not None else raw).astype(np.float32)

        if background is not None:
            waves -= background               # background is 1-D (spf,)

        # Parallelism comes from the outer chunk loop, so keep the inner
        # transform single-threaded to avoid oversubscribing the machine.
        spectrum = _do_rfft(waves, n=n_fft, axis=-1, workers=1)
        del waves
        # |z|^2 without np.abs()'s extra full-size temporary.
        power = spectrum.real ** 2
        power += spectrum.imag ** 2
        del spectrum
        # [..., 0] not [:, 0]: the unmasked path keeps the (rows, frames,
        # bins) shape, where [:, 0] would blank a whole frame.
        power[..., 0] = 0.0                   # suppress DC bin
        peak_bins = np.argmax(power, axis=-1)
        del power

        if valid is not None:
            img[r0:r1][valid] = freq_axis[peak_bins]
        else:
            img[r0:r1] = freq_axis[peak_bins]

    _map_row_chunks(n_rows, chunk_rows, n_workers, chunk)
    return img


# ---------------------------------------------------------------------------
# Batch cache (Convert menu) — process-pool entry point
# ---------------------------------------------------------------------------

def cache_file(path: str, mode: str, apply_bg_sub: bool,
               fft_backend: str = "numpy", max_workers: int = 0) -> str:
    """Compute and store DC or FFT images for every angle of one file,
    converting v6 → v7 in place. Returns "" on success or an error message.

    Module-level and picklable so it can run in a ProcessPoolExecutor. The
    FFT backend and worker cap are passed explicitly because module globals
    do not survive a spawn.
    """
    global _MAX_WORKERS
    try:
        set_fft_backend(fft_backend)
        if max_workers:
            _MAX_WORKERS = max_workers

        sras = SrasFile(path)
        if sras.version not in (6, 7):
            return (f"unsupported version {sras.version} — only v6/v7 files "
                    "can be batch-cached")

        n = sras.n_angles
        if mode == "dc":
            dc3 = [adc_to_mv(compute_dc_image(sras, a, CH3_IDX), *sras.cal(CH3_IDX))
                   for a in range(n)]
            dc4 = [adc_to_mv(compute_dc_image(sras, a, CH4_IDX), *sras.cal(CH4_IDX))
                   for a in range(n)]
            sras.write_v7_cache(new_dc3_mv=dc3, new_dc4_mv=dc4)
        else:
            effective_bg = apply_bg_sub and sras.background is not None
            # dc_threshold_mv=None: store unmasked images and mask at display
            # time (same convention as v5's PREC block). Skipping the mask
            # also skips reading CH4 entirely.
            freq = [compute_rf_image(sras, a, dc_threshold_mv=None,
                                     apply_bg_sub=effective_bg)
                    for a in range(n)]
            sras.write_v7_cache(new_freq_mhz=freq, new_bg_sub=effective_bg)
        return ""
    except Exception as exc:
        return str(exc)


# ---------------------------------------------------------------------------
# Angle alignment (Fusion menu)
#
# Puts every angle's images onto one shared, zero-padded pixel grid using a
# rigid transform only (rotation + translation, never scale). Rotation for
# angle `a` is the *known* scan-angle delta relative to a reference angle —
# never searched. Only the residual translation is found, via FFT phase
# correlation of each angle's binarized CH4 ("dc-mask") image.
#
# Rotation is done in physical mm space rather than on raw pixel indices:
# the x-pixel pitch (SrasFile.pixel_x_mm) is file-wide constant but the
# y-pixel pitch (row spacing) can differ from it, and for v6 files can even
# vary per angle. Rotating the raw index grid directly would implicitly
# assume square pixels and shear a non-square-pixel image — an unwanted
# effective anisotropic scale. Instead each angle gets one affine that maps
# shared-canvas pixel index -> mm -> undo rotation/shift -> that angle's own
# local mm -> that angle's own raw pixel index, matching the output->input
# convention scipy.ndimage.affine_transform expects.
# ---------------------------------------------------------------------------

@dataclass
class AngleTransform:
    rotation_deg: float
    shift_mm: tuple[float, float]      # (dx_mm, dy_mm) found by phase correlation
    matrix: np.ndarray                 # (2,2): canvas (row,col) -> this angle's raw (row,col)
    offset: np.ndarray                 # (2,)


@dataclass
class AlignmentResult:
    ref_angle_idx: int
    dc_threshold_mv: float
    canvas_shape: tuple[int, int]          # (n_rows, n_cols)
    canvas_dx_mm: float
    canvas_dy_mm: float
    canvas_origin_mm: tuple[float, float]  # mm at canvas pixel index (0, 0)
    per_angle: dict[int, AngleTransform]


def _pixel_pitch_mm(sras: SrasFile, angle_idx: int) -> tuple[float, float]:
    """(dx, dy) mm/pixel for one angle: dx is the file-wide constant
    pixel_x_mm; dy is this angle's own row spacing (assumed uniform, the same
    assumption _redraw_image makes when it builds the display extent)."""
    y = sras.y_positions_mm(angle_idx)
    return sras.pixel_x_mm, (float(y[1] - y[0]) if len(y) > 1 else 1.0)


def _bbox_center_mm(sras: SrasFile, angle_idx: int) -> tuple[float, float]:
    x = sras.x_axis_mm(angle_idx)
    y = sras.y_positions_mm(angle_idx)
    return float((x[0] + x[-1]) / 2.0), float((y[0] + y[-1]) / 2.0)


def _bbox_corners_mm(sras: SrasFile, angle_idx: int) -> np.ndarray:
    """4 corners (x, y) of this angle's raw mm bounding box, shape (4, 2)."""
    x = sras.x_axis_mm(angle_idx)
    y = sras.y_positions_mm(angle_idx)
    return np.array([[xx, yy] for xx in (x[0], x[-1]) for yy in (y[0], y[-1])])


def _rotation_matrix(theta_deg: float) -> np.ndarray:
    t = np.radians(theta_deg)
    c, s = np.cos(t), np.sin(t)
    return np.array([[c, -s], [s, c]])   # CCW rotation acting on (x, y)


def _theta_deg(sras: SrasFile, angle_idx: int, ref_idx: int) -> float:
    return float(sras.angles_deg[angle_idx] - sras.angles_deg[ref_idx])


def _corners_in_ref_frame(sras: SrasFile, angle_idx: int, ref_idx: int,
                          shift_mm=(0.0, 0.0)) -> np.ndarray:
    """Angle *angle_idx*'s bbox corners, rotated about its own centroid into
    the reference frame and translated by *shift_mm*. Shape (4, 2)."""
    R = _rotation_matrix(_theta_deg(sras, angle_idx, ref_idx))
    c_a = np.array(_bbox_center_mm(sras, angle_idx))
    c_ref = np.array(_bbox_center_mm(sras, ref_idx))
    shift = np.asarray(shift_mm, dtype=np.float64)
    return np.array([R @ (corner - c_a) + c_ref + shift
                     for corner in _bbox_corners_mm(sras, angle_idx)])


def _build_affine_canvas_to_raw(sras: SrasFile, angle_idx: int, ref_idx: int,
                                shift_mm: tuple[float, float],
                                canvas_dx: float, canvas_dy: float,
                                canvas_origin_mm: tuple[float, float]
                                ) -> tuple[np.ndarray, np.ndarray]:
    """matrix, offset s.t. raw_index = matrix @ [row_out, col_out] + offset,
    matching scipy.ndimage.affine_transform's output->input convention.

    Pipeline (all mm unless noted):
      [X;Y]   = A_out @ [row_out;col_out] + b_out        # canvas idx -> ref-frame mm
      [lx;ly] = R(theta)^T @ ([X;Y]-c_ref-shift) + c_a    # undo rotation+shift -> angle a's local mm
      [row;col] = D @ ([lx;ly] - [x_start_a; y0_a])       # local mm -> angle a's raw idx

    where theta = angles_deg[angle_idx] - angles_deg[ref_idx], c_ref/c_a are
    each angle's own raw-bbox mm centroid (the rotation pivot — this keeps
    rotated content centered, minimizing required canvas padding), and
    A_out/D are the index<->mm scaling matrices for the canvas pitch and
    this angle's own native pitch respectively.
    """
    Rinv = _rotation_matrix(_theta_deg(sras, angle_idx, ref_idx)).T
    cx_a, cy_a = _bbox_center_mm(sras, angle_idx)
    cx_ref, cy_ref = _bbox_center_mm(sras, ref_idx)
    dx_a, dy_a = _pixel_pitch_mm(sras, angle_idx)
    x0_a = float(sras.x_start_mm[angle_idx])
    y0_a = float(sras.y_positions_mm(angle_idx)[0])

    A_out = np.array([[0.0, canvas_dx], [canvas_dy, 0.0]])   # [row,col] -> [X,Y]
    b_out = np.array(canvas_origin_mm, dtype=np.float64)
    D = np.array([[0.0, 1.0 / dy_a], [1.0 / dx_a, 0.0]])     # [x,y] -> [row,col]
    shift = np.array(shift_mm, dtype=np.float64)
    c_ref_v = np.array([cx_ref, cy_ref])
    c_a_v = np.array([cx_a, cy_a])
    origin_a = np.array([x0_a, y0_a])

    matrix = D @ Rinv @ A_out
    offset = D @ Rinv @ (b_out - c_ref_v - shift) + D @ (c_a_v - origin_a)
    return matrix, offset


def apply_alignment(result: AlignmentResult, angle_idx: int, img: np.ndarray,
                    order: int = 0) -> np.ndarray:
    """Resample any already-computed 2D image for `angle_idx` (same shape as
    that angle's raw (n_rows, n_frames)) onto the shared alignment canvas.
    order=0 (nearest) avoids blending real data with zero-padding or with
    masked-out (0-valued) CH1/velocity pixels at mask edges. Channel-
    agnostic: the same per-angle transform (found from the CH4 mask) works
    for any channel's image of that angle."""
    t = result.per_angle[angle_idx]
    return scipy_ndimage.affine_transform(
        img.astype(np.float32, copy=False), t.matrix, offset=t.offset,
        output_shape=result.canvas_shape, order=order,
        mode="constant", cval=0.0)


def _block_mean_downsample(img: np.ndarray, factor: int) -> np.ndarray:
    if factor <= 1:
        return img
    h, w = img.shape
    h2, w2 = (h // factor) * factor, (w // factor) * factor
    trimmed = img[:h2, :w2]
    return trimmed.reshape(h2 // factor, factor, w2 // factor, factor).mean(axis=(1, 3))


def _phase_correlate_shift(ref_img: np.ndarray, mov_img: np.ndarray) -> tuple[int, int]:
    """FFT normalized cross-power-spectrum phase correlation. Returns the
    integer (dr, dc) pixel shift of mov_img relative to ref_img; both must
    be the same shape. Risk: if the true shift is near +/- half the array
    size, wraparound can bias the peak — mitigated by the generous
    margin_frac padding in _working_canvas_for_pair, which keeps the true
    residual shift small relative to the correlation canvas."""
    F1 = scipy_fft.fft2(ref_img.astype(np.float64), workers=-1)
    F2 = scipy_fft.fft2(mov_img.astype(np.float64), workers=-1)
    R = F1 * np.conj(F2)
    R /= np.maximum(np.abs(R), 1e-12)
    corr = scipy_fft.ifft2(R, workers=-1).real
    dr, dc = np.unravel_index(np.argmax(corr), corr.shape)
    h, w = corr.shape
    if dr > h // 2:
        dr -= h
    if dc > w // 2:
        dc -= w
    return int(dr), int(dc)


def _working_canvas_for_pair(sras: SrasFile, ref_idx: int, a_idx: int,
                             dx: float, dy: float, margin_frac: float = 0.3
                             ) -> tuple[tuple[float, float], tuple[int, int]]:
    """Union of the reference's own raw bbox and angle a's raw bbox rotated
    (about its own center) into the ref frame with zero shift, padded by
    margin_frac on each side — sized generously so the true phase-
    correlation shift lands well inside the canvas (see
    _phase_correlate_shift's wraparound note)."""
    pts = np.vstack([_bbox_corners_mm(sras, ref_idx),
                     _corners_in_ref_frame(sras, a_idx, ref_idx)])
    x_min, y_min = pts.min(axis=0)
    x_max, y_max = pts.max(axis=0)
    pad_x, pad_y = (x_max - x_min) * margin_frac, (y_max - y_min) * margin_frac
    x_min, x_max = x_min - pad_x, x_max + pad_x
    y_min, y_max = y_min - pad_y, y_max + pad_y
    n_cols = int(np.ceil((x_max - x_min) / dx)) + 1
    n_rows = int(np.ceil((y_max - y_min) / abs(dy))) + 1
    origin = (x_min, y_min if dy > 0 else y_max)
    return origin, (n_rows, n_cols)


def _parallel_map(fn, items, n_workers: int) -> list:
    """fn over items, in order, threaded when it pays."""
    items = list(items)
    if n_workers <= 1 or len(items) <= 1:
        return [fn(x) for x in items]
    with ThreadPoolExecutor(max_workers=min(n_workers, len(items))) as pool:
        return list(pool.map(fn, items))


def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int,
                            dc_threshold_mv: float,
                            progress_cb=None) -> AlignmentResult:
    """Top-level alignment driver. Runs on a background thread (see
    AngleAlignmentWorker) — deliberately recomputes CH4 DC images from
    scratch rather than reading the GUI-thread _dc_cache dict, since
    background-thread workers must not touch GUI-thread-owned caches."""
    n = sras.n_angles   # already the *complete*-angle count for aborted v6 scans
    dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx)
    n_workers = max(1, min(_MAX_WORKERS, n))

    # progress_cb fires from pool threads, so the counter behind it must be
    # atomic. list.append is, and len() of a list is a consistent read.
    _ticks: list[int] = []

    def tick(base: int, span: int):
        if progress_cb:
            _ticks.append(1)
            progress_cb(base + int(min(len(_ticks), n) / n * span))

    # ---- Step 1: binarized CH4 mask per angle, native per-angle grid -----
    def mask_for(a: int) -> np.ndarray:
        dc4 = adc_to_mv(compute_dc_image(sras, a, CH4_IDX), *sras.cal(CH4_IDX))
        m = (dc4 >= dc_threshold_mv).astype(np.float32)
        tick(0, 25)
        return m

    masks = dict(enumerate(_parallel_map(mask_for, range(n), n_workers)))

    # ---- Step 2: coarse correlation stage (downsample first, then rotate)
    # Downsampling before affine_transform (not after) is what keeps this
    # tractable for a v6 scan with thousands of rows/frames per angle.
    max_dim = max(max(m.shape) for m in masks.values())
    factor = max(1, int(np.ceil(max_dim / 1024)))
    dx_c, dy_c = dx_ref * factor, dy_ref * factor
    masks_small = {a: _block_mean_downsample(m, factor) for a, m in masks.items()}

    _ticks.clear()

    def shift_for(a: int) -> tuple[float, float]:
        if a == ref_angle_idx:
            tick(25, 50)
            return (0.0, 0.0)
        work_origin, work_shape = _working_canvas_for_pair(
            sras, ref_angle_idx, a, dx_c, dy_c, margin_frac=0.3)

        # Coarse canvas->raw affine at native pitch, then rescale by /factor
        # so it maps coarse-canvas idx -> coarse (downsampled) raw idx —
        # exact for block-mean downsampling (up to the trimmed remainder).
        m_a, o_a = _build_affine_canvas_to_raw(
            sras, a, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin)
        m_ref, o_ref = _build_affine_canvas_to_raw(
            sras, ref_angle_idx, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin)
        rotated_a = scipy_ndimage.affine_transform(
            masks_small[a], m_a / factor, offset=o_a / factor,
            output_shape=work_shape, order=0, mode="constant", cval=0.0)
        embedded_ref = scipy_ndimage.affine_transform(
            masks_small[ref_angle_idx], m_ref / factor, offset=o_ref / factor,
            output_shape=work_shape, order=0, mode="constant", cval=0.0)

        dr, dc = _phase_correlate_shift(embedded_ref, rotated_a)
        tick(25, 50)
        return (dc * dx_c, dr * dy_c)

    shifts_mm = dict(enumerate(_parallel_map(shift_for, range(n), n_workers)))

    # ---- Step 3: union bounding box over all angles (rotation+shift applied)
    corners = np.vstack([_corners_in_ref_frame(sras, a, ref_angle_idx, shifts_mm[a])
                         for a in range(n)])
    x_min, y_min = corners.min(axis=0)
    x_max, y_max = corners.max(axis=0)
    n_cols = int(np.ceil((x_max - x_min) / dx_ref)) + 1
    n_rows = int(np.ceil((y_max - y_min) / abs(dy_ref))) + 1
    canvas_origin_mm = (float(x_min), float(y_min if dy_ref > 0 else y_max))

    # ---- Step 4: final per-angle full-resolution affine (canvas -> raw idx)
    per_angle: dict[int, AngleTransform] = {}
    for a in range(n):
        matrix, offset = _build_affine_canvas_to_raw(
            sras, a, ref_angle_idx, shifts_mm[a], dx_ref, dy_ref, canvas_origin_mm)
        per_angle[a] = AngleTransform(
            _theta_deg(sras, a, ref_angle_idx), shifts_mm[a], matrix, offset)
        if progress_cb:
            progress_cb(75 + int((a + 1) / n * 25))

    return AlignmentResult(ref_angle_idx, dc_threshold_mv, (n_rows, n_cols),
                           dx_ref, dy_ref, canvas_origin_mm, per_angle)


# Back-compat alias for the pre-split private name (used by tooling).
_compute_angle_alignment = compute_angle_alignment