#!/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 json import os from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from pathlib import Path 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. # 1024 MB is the measured knee on a 16-core machine against a 7507-frame x # 2500-sample angle: 512 MB leaves ~20% of the FFT speedup on the table, and # 1536+ MB costs ~0.4 GB more resident for no further gain. Override with # SRAS_MEM_BUDGET_MB on a smaller machine. _TOTAL_BYTES_BUDGET = int(os.environ.get("SRAS_MEM_BUDGET_MB", 1024)) * 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, budget: int | None = None) -> tuple[int, int]: """(chunk_rows, n_workers) such that all concurrent chunks together fit the budget: n_workers x live_multiplier x chunk_rows x n_frames x samples x 4 bytes <= budget. The worker count is chosen *first* and the chunk sized to it. Sizing the chunk first is the trap: _chunk_rows_for spends whatever budget it is given, so a single chunk would always consume the lot and leave room for exactly one worker — no concurrency, on precisely the large scans that need it most. A caller that is itself running several of these concurrently must pass *both* max_workers=1 and its share of the budget. Capping the workers alone is not enough: the chunk would still be sized against the whole budget, and N concurrent callers would each allocate all of it. """ total = _TOTAL_BYTES_BUDGET if budget is None else max(1, budget) cap = max_workers if max_workers is not None else _MAX_WORKERS bytes_per_row = max(1, live_multiplier * n_frames * samples_per_frame * 4) # A chunk is at least one row, so that alone caps how many can be live. n_workers = int(max(1, min(cap, n_rows, total // bytes_per_row))) chunk_rows = _chunk_rows_for(n_frames, samples_per_frame, max(1, total // (n_workers * live_multiplier))) # With chunk_rows known, more workers than chunks buys nothing. n_workers = int(max(1, min(n_workers, -(-n_rows // chunk_rows)))) return chunk_rows, n_workers def _map_row_chunks(n_rows: int, chunk_rows: int, n_workers: int, fn, should_stop=None): """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. *should_stop* is polled per chunk so a cancelled job abandons work within one chunk rather than one whole angle — on a large scan an angle is ~40 s, which is far too long to block application shutdown. """ bounds = [(r0, min(r0 + chunk_rows, n_rows)) for r0 in range(0, n_rows, chunk_rows)] def run(b): if should_stop is not None and should_stop(): return fn(*b) if n_workers <= 1 or len(bounds) == 1: for b in bounds: run(b) return with ThreadPoolExecutor(max_workers=min(n_workers, len(bounds))) as pool: list(pool.map(run, bounds)) # --------------------------------------------------------------------------- # Image computation # --------------------------------------------------------------------------- def plan_angle_level(sras: SrasFile, live_multiplier: int = 1) -> tuple[int, int]: """(n_angle_workers, per_angle_budget) for a caller that parallelises over angles instead of over rows. Callers that parallelise over angles must not also let each angle parallelise over rows: the two levels multiply, in threads and in memory. Each angle is therefore given max_workers=1 and the returned budget share, which together keep total live buffers within _TOTAL_BYTES_BUDGET. """ biggest = max(range(sras.n_angles), key=lambda a: int(sras.n_frames[a])) _, n_workers = _plan_chunks(int(sras.n_rows[biggest]), int(sras.n_frames[biggest]), sras.samples_per_frame, live_multiplier=live_multiplier) n_workers = max(1, min(n_workers, sras.n_angles)) return n_workers, max(1, _TOTAL_BYTES_BUDGET // n_workers) def compute_dc_image(sras: SrasFile, angle_idx: int, ch_idx: int, max_workers: int | None = None, budget: int | None = None, should_stop=None) -> np.ndarray: """Mean of each waveform → (n_rows, n_frames) float32, in ADC counts. Pass max_workers=1 when the caller is already parallelising over angles. If *should_stop* ever returns True the result is incomplete — the caller is expected to be abandoning it. """ 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, max_workers=max_workers, budget=budget) 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, should_stop=should_stop) return img def dc_image_mv(sras: SrasFile, angle_idx: int, ch_idx: int, max_workers: int | None = None, budget: int | None = None, should_stop=None) -> 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, max_workers=max_workers, budget=budget, should_stop=should_stop), *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, max_workers: int | None = None, budget: int | None = None, should_stop=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, max_workers=max_workers, budget=budget) 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, should_stop=should_stop) 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: if mode not in ("dc", "fft"): return f"unknown cache mode {mode!r} (expected 'dc' or 'fft')" 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] @dataclass class ManualAngleParams: """One angle's manual-alignment state, independent of any canvas. rotation_deg/shift_mm are exactly AngleTransform's non-derived fields — the pair a canvas-bound AngleTransform's matrix/offset get built from once a canvas is decided (build_manual_alignment). Defaults to identity (no rotation, no shift): a fresh angle with no prior alignment is shown raw, exactly as scanned — the same "fully unaligned" state Clear Alignment resets back to. """ rotation_deg: float = 0.0 shift_mm: tuple[float, float] = (0.0, 0.0) 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: """CCW rotation, in degrees and in _rotation_matrix's convention, that maps angle_idx's own local mm frame onto ref_idx's. This is the *negative* of the raw angles_deg delta: the GR rotation stage's reported angle increases in the opposite rotational sense from this module's math-positive (CCW, x toward y) convention in scan mm space. Rotating by +(angles_deg[a] - angles_deg[ref]) therefore turns misalignment the wrong way — confirmed empirically (Auto De-rotate made real scans worse, not better, before this negation). """ return -float(sras.angles_deg[angle_idx] - sras.angles_deg[ref_idx]) def _signal_centroid_mm(sras: SrasFile, angle_idx: int, dc4_mv: np.ndarray) -> tuple[float, float]: """Intensity-weighted centroid (mean x, mean y, weighted by CH4 signal after subtracting this angle's own minimum) in angle_idx's own local mm frame — the alignment pivot used in place of the raw scan-window bbox center (see compute_pivot_points_mm). Weighting by the continuous DC signal, rather than a binary >= dc_threshold_mv mask, means the pivot never depends on how well one shared threshold happens to suit this particular angle: real signal levels vary scan to scan, so a threshold tuned for one angle can leave another angle's binary mask empty — and a centroid of an empty mask has nothing to fall back to *except* the raw bbox center, silently reproducing the exact "aligned to the scan window, not the sample" problem this pivot exists to avoid. Falls back to the bbox center only in the fully-degenerate case of a perfectly flat signal (nothing to weight by at all). """ weights = dc4_mv - dc4_mv.min() total = float(weights.sum()) if total <= 0.0: return _bbox_center_mm(sras, angle_idx) x = sras.x_axis_mm(angle_idx) y = sras.y_positions_mm(angle_idx) cx = float((x * weights.sum(axis=0)).sum() / total) cy = float((y * weights.sum(axis=1)).sum() / total) return cx, cy def compute_pivot_points_mm(sras: SrasFile, dc4_mv: dict[int, np.ndarray] | None = None ) -> dict[int, tuple[float, float]]: """Per-angle alignment pivot, in each angle's own local mm frame: the CH4-signal-weighted centroid of its own footprint (see _signal_centroid_mm), rather than the raw scan-window bbox center. Pivoting on each angle's own content — instead of on wherever its scanned window happened to sit in microscope/global XY space — is what makes alignment purely relative *between scans* rather than to a global coordinate system: the rotation stage's true mechanical axis need not coincide with the scan window's geometric center, and the sample need not be perfectly centered on that axis either, so a bbox-center pivot leaves a residual orbital motion between angles that a content-centroid pivot does not. Deliberately independent of dc_threshold_mv (the RF-mask / overlay threshold): that value is a display/masking choice and must never silently change where alignment pivots. *dc4_mv* lets a caller that has already computed each angle's CH4 mV image (compute_angle_alignment's Step 1, or ManualAlignmentDialog's own cache) reuse it instead of recomputing; angles missing from it are computed fresh via dc_image_mv, which prefers a stored v5/v7 cache over recomputing from raw waveforms. """ dc4_mv = dc4_mv or {} pivots: dict[int, tuple[float, float]] = {} for a in range(sras.n_angles): img = dc4_mv.get(a) if img is None: img = dc_image_mv(sras, a, CH4_IDX) pivots[a] = _signal_centroid_mm(sras, a, img) return pivots def _corners_in_ref_frame(sras: SrasFile, angle_idx: int, ref_idx: int, pivot_mm: dict[int, tuple[float, float]], shift_mm=(0.0, 0.0), theta_deg: float | None = None) -> np.ndarray: """Angle *angle_idx*'s bbox corners, rotated about its own alignment pivot (pivot_mm[angle_idx] — see compute_pivot_points_mm) into the reference frame and translated by *shift_mm*. Shape (4, 2). theta_deg overrides the analytic angles_deg-derived rotation used by default — the manual-alignment path (union_canvas_mm) passes a user-chosen rotation here (which may differ from the known scan-angle delta) without needing a parallel code path. """ theta = _theta_deg(sras, angle_idx, ref_idx) if theta_deg is None else theta_deg R = _rotation_matrix(theta) c_a = np.array(pivot_mm[angle_idx]) c_ref = np.array(pivot_mm[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], pivot_mm: dict[int, tuple[float, float]], theta_deg: float | None = None ) -> 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 is _theta_deg(angle_idx, ref_idx) unless *theta_deg* overrides it (see _corners_in_ref_frame's note, used by the manual- alignment path), c_ref/c_a are each angle's own alignment pivot (pivot_mm — see compute_pivot_points_mm; the CH4-signal-weighted centroid of its own footprint, not the raw scan-window bbox center — this keeps the sample itself centered post-rotation, minimizing required canvas padding and, more importantly, keeping alignment relative to the sample rather than to wherever the scan window sat in microscope/global XY space), and A_out/D are the index<->mm scaling matrices for the canvas pitch and this angle's own native pitch respectively. """ theta = _theta_deg(sras, angle_idx, ref_idx) if theta_deg is None else theta_deg Rinv = _rotation_matrix(theta).T cx_a, cy_a = pivot_mm[angle_idx] cx_ref, cy_ref = pivot_mm[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, pivot_mm: dict[int, tuple[float, 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 alignment pivot) 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, pivot_mm)]) 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) # Parallel over angles, serial within each — see plan_angle_level. n_workers, angle_budget = plan_angle_level(sras) # 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: CH4 DC image + binarized mask per angle, native grid ----- def dc4_for(a: int) -> np.ndarray: dc4 = adc_to_mv( compute_dc_image(sras, a, CH4_IDX, max_workers=1, budget=angle_budget), *sras.cal(CH4_IDX)) tick(0, 25) return dc4 dc4_mv = dict(enumerate(_parallel_map(dc4_for, range(n), n_workers))) masks = {a: (img >= dc_threshold_mv).astype(np.float32) for a, img in dc4_mv.items()} # Each angle's own alignment pivot — the CH4-signal-weighted centroid, # not the raw scan-window bbox center (see compute_pivot_points_mm). # Reuses the dc4_mv images just computed above, so this is free, and is # deliberately independent of dc_threshold_mv (see that function's # docstring) so a threshold that happens to leave some angle's binary # mask empty can't silently degrade the pivot back to the bbox center. pivot_mm = compute_pivot_points_mm(sras, dc4_mv=dc4_mv) # ---- 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, pivot_mm, 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, pivot_mm) 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, pivot_mm) 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, pivot_mm, 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, pivot_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 # --------------------------------------------------------------------------- # Manual alignment (Fusion menu -> Manual Alignment... dialog) # # Skips compute_angle_alignment's mask + phase-correlation search entirely: # every angle's rotation_deg/shift_mm is supplied directly by the caller # (nudged by eye against a live multi-angle mask overlay, or pre-seeded from # a previous compute_angle_alignment run or a saved sidecar). Building the # final AlignmentResult from already-known per-angle parameters is pure # closed-form matrix math (union_canvas_mm + _build_affine_canvas_to_raw) # with no per-pixel image work at all, so build_manual_alignment is cheap # enough to call synchronously on the GUI thread on every edit. The only # genuinely expensive per-pixel operation anywhere in this flow is # reproject_mask, and only ManualAlignmentDialog's own downsampled preview # calls that per keystroke — see that class's docstring for how it limits # each nudge to reprojecting only the actively-edited angle. # --------------------------------------------------------------------------- def union_canvas_mm(sras: SrasFile, ref_angle_idx: int, dx: float, dy: float, per_angle_params: dict[int, ManualAngleParams], pivot_mm: dict[int, tuple[float, float]], margin_frac: float = 0.0 ) -> tuple[tuple[float, float], tuple[int, int]]: """Shared-canvas origin (mm) and (n_rows, n_cols) at pitch (dx, dy) that contains every angle's footprint after applying its own rotation+shift — the generalisation of compute_angle_alignment's Step 3 to arbitrary (not just _theta_deg-analytic) per-angle rotation. Angles missing from per_angle_params default to identity (e.g. a sidecar saved before a rescan added more angles). margin_frac pads the box on every side: 0 for a final canvas (this then reproduces compute_angle_alignment's own Step-3 math exactly, when every angle's rotation_deg equals the analytic delta and shift_mm matches); nonzero for ManualAlignmentDialog's downsampled preview canvas, which needs headroom so an ordinary translation nudge of the active angle never has to trigger a full canvas resize (see that class's docstring — an extreme nudge can still, in principle, push content past this padding; accepted as a known edge case, same as _phase_correlate_shift's wraparound risk note). """ n = sras.n_angles corners = np.vstack([ _corners_in_ref_frame( sras, a, ref_angle_idx, pivot_mm, per_angle_params.get(a, ManualAngleParams()).shift_mm, theta_deg=per_angle_params.get(a, ManualAngleParams()).rotation_deg) for a in range(n)]) x_min, y_min = corners.min(axis=0) x_max, y_max = corners.max(axis=0) if margin_frac: 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 = (float(x_min), float(y_min if dy > 0 else y_max)) return origin, (n_rows, n_cols) def reproject_mask(sras: SrasFile, angle_idx: int, ref_angle_idx: int, mask: np.ndarray, rotation_deg: float, shift_mm: tuple[float, float], canvas_dx: float, canvas_dy: float, canvas_origin_mm: tuple[float, float], canvas_shape: tuple[int, int], pivot_mm: dict[int, tuple[float, float]]) -> np.ndarray: """Resample one angle's binary/float mask onto an arbitrary canvas via an explicit rotation+shift — the single building block ManualAlignmentDialog's live preview repeatedly calls (once per keystroke, for only the actively-nudged angle), since it bypasses compute_angle_alignment's phase-correlation search entirely and just takes rotation_deg/shift_mm as given. order=0 (nearest) matches apply_alignment's own reasoning: a binary mask must never be blended with zero-padding. """ matrix, offset = _build_affine_canvas_to_raw( sras, angle_idx, ref_angle_idx, shift_mm, canvas_dx, canvas_dy, canvas_origin_mm, pivot_mm, theta_deg=rotation_deg) return scipy_ndimage.affine_transform( mask.astype(np.float32, copy=False), matrix, offset=offset, output_shape=canvas_shape, order=0, mode="constant", cval=0.0) def build_manual_alignment(sras: SrasFile, ref_angle_idx: int, dc_threshold_mv: float, per_angle_params: dict[int, ManualAngleParams], pivot_mm: dict[int, tuple[float, float]] | None = None ) -> AlignmentResult: """Build a full, full-resolution AlignmentResult from user-supplied per-angle rotation+shift — the Manual Alignment counterpart to compute_angle_alignment, skipping its mask/phase-correlation search entirely (every angle's transform here is exactly what the caller supplied). Pure matrix/bbox math once pivot_mm is known, so it is cheap enough to call synchronously on the GUI thread. The reference angle's params are always forced to identity, regardless of what per_angle_params holds for it — it defines the shared origin and must never be transformed. Pass an already-computed *pivot_mm* (e.g. ManualAlignmentDialog's own cache, built once from its live CH4 images) to skip recomputing every angle's DC image here; otherwise it's computed fresh via compute_pivot_points_mm, which is still fine for a one-off Save or sidecar restore (just not free on a very large, not-yet-cached scan). dc_threshold_mv itself plays no part in the pivot (see that function's docstring) — it's stored on the returned AlignmentResult purely as a record of the RF-mask threshold in effect at the time. """ n = sras.n_angles dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx) params = {a: per_angle_params.get(a, ManualAngleParams()) for a in range(n)} params[ref_angle_idx] = ManualAngleParams() if pivot_mm is None: pivot_mm = compute_pivot_points_mm(sras) canvas_origin_mm, canvas_shape = union_canvas_mm( sras, ref_angle_idx, dx_ref, dy_ref, params, pivot_mm, margin_frac=0.0) per_angle: dict[int, AngleTransform] = {} for a in range(n): p = params[a] matrix, offset = _build_affine_canvas_to_raw( sras, a, ref_angle_idx, p.shift_mm, dx_ref, dy_ref, canvas_origin_mm, pivot_mm, theta_deg=p.rotation_deg) per_angle[a] = AngleTransform(p.rotation_deg, p.shift_mm, matrix, offset) return AlignmentResult(ref_angle_idx, dc_threshold_mv, canvas_shape, dx_ref, dy_ref, canvas_origin_mm, per_angle) # ---- Sidecar persistence (.sras.align.json) ------------------------- # # Lives here, not sras_format.py: sras_format.py is scoped to the versioned # binary .sras spec itself (see scan_format.md); a manual alignment is a # viewer-computed *derived* artifact, analogous in kind to AlignmentResult — # so it belongs with the alignment math it serialises, which already lives # in this module. json + pathlib are both stdlib, so this doesn't add a new # dependency to a module whose only load-bearing constraint is staying free # of Qt/matplotlib for cheap multiprocessing-child imports. @dataclass class ManualAlignmentSidecar: ref_angle_idx: int dc_threshold_mv: float per_angle: dict[int, ManualAngleParams] def sidecar_path(sras_path) -> Path: """.sras.align.json next to the scan file. A thin, independently testable function since save/load/delete and the GUI's status messages all need the identical path.""" p = Path(sras_path) return p.with_name(p.name + ".align.json") # Bumped from 1 -> 2 when the rotation pivot changed from the raw scan- # window bbox center to a content-derived centroid, *and* _theta_deg's sign # convention was corrected — either change alone makes a version-1 file's # stored rotation_deg/shift_mm numbers describe a different (and, for the # pivot bug, actively wrong) transform than they would today. Loading one # unchanged would silently reproduce exactly the "scans show up everywhere" # symptom these fixes address, so version-1 sidecars are treated as absent # rather than migrated. _SIDECAR_SCHEMA_VERSION = 2 def save_manual_alignment(sras: SrasFile, ref_angle_idx: int, dc_threshold_mv: float, per_angle: dict[int, ManualAngleParams]) -> Path: """Write the sidecar JSON for sras.path (overwriting any existing one) and return the path written. Schema (schema_version 2): { "schema_version": 2, "ref_angle_idx": , "dc_threshold_mv": , "per_angle": { "": {"rotation_deg": , "shift_mm": [, ]}, ... } } Angle indices are JSON object keys, so they round-trip as strings — load_manual_alignment converts them back to int. """ path = sidecar_path(sras.path) payload = { "schema_version": _SIDECAR_SCHEMA_VERSION, "ref_angle_idx": ref_angle_idx, "dc_threshold_mv": dc_threshold_mv, "per_angle": { str(a): {"rotation_deg": p.rotation_deg, "shift_mm": list(p.shift_mm)} for a, p in per_angle.items() }, } path.write_text(json.dumps(payload, indent=2)) return path def load_manual_alignment(sras: SrasFile) -> ManualAlignmentSidecar | None: """Read 's sidecar JSON if present, else None — never raises: a missing, corrupt, foreign, or version-mismatched JSON file must not block opening the .sras file itself (see SrasViewerWindow._on_load_done). Per-angle entries for an angle index no longer present in *sras* (e.g. re-scanned with fewer angles) are silently dropped. """ path = sidecar_path(sras.path) if not path.exists(): return None try: raw = json.loads(path.read_text()) if raw.get("schema_version") != _SIDECAR_SCHEMA_VERSION: return None per_angle = { int(a): ManualAngleParams( float(v["rotation_deg"]), (float(v["shift_mm"][0]), float(v["shift_mm"][1]))) for a, v in raw.get("per_angle", {}).items() if int(a) < sras.n_angles } return ManualAlignmentSidecar( ref_angle_idx=int(raw.get("ref_angle_idx", 0)), dc_threshold_mv=float(raw.get("dc_threshold_mv", 0.0)), per_angle=per_angle) except (OSError, ValueError, KeyError, TypeError, IndexError, json.JSONDecodeError): return None def delete_manual_alignment(sras: SrasFile) -> bool: """Delete the sidecar if present. Returns whether a file actually existed to delete, so Clear Alignment's status message can say so. Genuine I/O errors (permission denied, read-only share) propagate — the caller (ManualAlignmentDialog._on_clear) surfaces them rather than silently pretending the destructive action succeeded.""" path = sidecar_path(sras.path) try: path.unlink() return True except FileNotFoundError: return False