#!/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 atexit import json import os import threading 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_AVAILABLE = True except ImportError: PYFFTW_AVAILABLE = False try: from threadpoolctl import threadpool_limits except ImportError: threadpool_limits = None _fft_backend = "scipy" # "scipy" 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. "numpy" is accepted as a legacy alias for "scipy".""" global _fft_backend if name == "numpy": name = "scipy" _fft_backend = name if (name == "pyfftw" and PYFFTW_AVAILABLE) else "scipy" 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) # --------------------------------------------------------------------------- # FFT worker pool and per-thread pyFFTW plans # --------------------------------------------------------------------------- _FFT_BLOCK = 512 # waveforms per FFT task. The knee on a 16-core machine: # smaller blocks serialise on GIL-held numpy dispatch, # larger ones lose cache residency and task granularity. _ZOOM_MIN_PAD = 4 # zoom refinement engages at n_fft >= _ZOOM_MIN_PAD * spf _FFT_EXACT_ENV = bool(int(os.environ.get("SRAS_FFT_EXACT", "0") or 0)) _pool_lock = threading.Lock() _pool: ThreadPoolExecutor | None = None def _fft_pool() -> ThreadPoolExecutor: """The persistent process-wide pool for FFT block tasks.""" global _pool with _pool_lock: if _pool is None: _pool = ThreadPoolExecutor(max_workers=_MAX_WORKERS, thread_name_prefix="sras-fft") atexit.register(_pool.shutdown, wait=False, cancel_futures=True) return _pool _WISDOM_PATH = Path.home() / ".cache" / "sras-viewer" / "fftw_wisdom" _wisdom_lock = threading.Lock() _wisdom_loaded = False _fftw_local = threading.local() def _load_wisdom_once(): """Import saved FFTW wisdom so FFTW_MEASURE planning is a one-time cost per machine. Purely an optimisation: failures are ignored.""" global _wisdom_loaded with _wisdom_lock: if _wisdom_loaded: return _wisdom_loaded = True try: pyfftw.import_wisdom(_WISDOM_PATH.read_bytes().split(b"\x00\n")) except Exception: pass def _save_wisdom(): with _wisdom_lock: try: _WISDOM_PATH.parent.mkdir(parents=True, exist_ok=True) _WISDOM_PATH.write_bytes(b"\x00\n".join(pyfftw.export_wisdom())) except Exception: pass def _fftw_block_rfft(waves: np.ndarray, n: int) -> np.ndarray: """rfft of a (B, spf) float32 block via a cached per-thread FFTW plan. Plans have a fixed (_FFT_BLOCK, spf) input shape so each worker thread plans once per transform length; a remainder block runs through the same plan with its tail rows ignored. The returned array is the plan's output buffer — consume it before the next call on the same thread. """ n_wf, spf = waves.shape plans = getattr(_fftw_local, "plans", None) if plans is None: plans = _fftw_local.plans = {} key = (_FFT_BLOCK, spf, n) plan = plans.get(key) if plan is None: _load_wisdom_once() buf = pyfftw.empty_aligned((_FFT_BLOCK, spf), dtype="float32") # No overwrite_input: FFTW must not scribble on input_array, whose # zero-padded tail (columns spf..n) is zeroed exactly once here. plan = pyfftw.builders.rfft(buf, n=n, axis=-1, threads=1, planner_effort="FFTW_MEASURE") plan.input_array[:] = 0.0 plans[key] = plan _save_wisdom() inp = plan.input_array inp[:n_wf, :spf] = waves return plan()[:n_wf] def _block_rfft(waves: np.ndarray, n: int) -> np.ndarray: """Single-threaded rfft of one block; outer parallelism comes from the pool, so the transform itself must not spin up threads.""" if _fft_backend == "pyfftw" and PYFFTW_AVAILABLE: return _fftw_block_rfft(waves, n) return scipy_fft.rfft(waves, n=n, axis=-1, workers=1) # --------------------------------------------------------------------------- # Zoom peak search: coarse rfft + local fine DFT around the winning bin # --------------------------------------------------------------------------- _ZOOM_HALFWIDTH = 0.75 # refinement window half-width, in coarse spacings. # Every fine bin lies within 0.5 spacings of its # nearest coarse bin, and that bin is guaranteed to # be a candidate (see _ZOOM_CAND_RATIO), so 0.5 # suffices; 0.75 adds rounding margin. _ZOOM_CAND_RATIO = 0.7 # refine every coarse bin within this power ratio of # its row's coarse maximum. Quarter-natural-bin # scalloping at the 2x-oversampled coarse grid can # understate a peak by at most ~19% in power, so 0.7 # keeps a wide margin — near-contenders are resolved # on the fine grid, never ranked from coarse samples. @dataclass class _ZoomPlan: """Constants for the coarse+refine peak search, built once per compute_rf_image call and shared across worker threads. *phases* is a lazily-filled window-start -> phase-vector cache; a benign duplicate compute under concurrency is cheaper than locking.""" n_fft: int n_coarse: int m: int # fine bins per refinement window n_bins_fine: int E: np.ndarray # (spf, m) complex64 fine-DFT matrix, relative bins j: np.ndarray # (spf,) float64 sample indices phases: dict def _zoom_plan(spf: int, n_fft: int) -> _ZoomPlan: # 2x-oversampled coarse grid: the padded power spectrum is a trig # polynomial of degree spf-1, so at 2x sampling its global max cannot # hide between coarse bins. n_coarse = scipy_fft.next_fast_len(2 * spf, real=True) n_bins_fine = n_fft // 2 + 1 m = int(np.ceil(2 * _ZOOM_HALFWIDTH * n_fft / n_coarse)) + 1 m = min(m, n_bins_fine - 1) j = np.arange(spf, dtype=np.float64) E = np.exp((-2j * np.pi / n_fft) * np.outer(j, np.arange(m))).astype(np.complex64) return _ZoomPlan(n_fft, n_coarse, m, n_bins_fine, E, j, {}) def _window_start(k_c: np.ndarray, zp: _ZoomPlan) -> np.ndarray: """First fine bin of the refinement window around each coarse bin. Clipped to [1, ...] so fine bin 0 stays excluded (DC suppression).""" k0 = np.floor((k_c - _ZOOM_HALFWIDTH) * (zp.n_fft / zp.n_coarse)).astype(np.int64) return np.clip(k0, 1, max(1, zp.n_bins_fine - zp.m)) def _refine_window(waves: np.ndarray, rows: np.ndarray, k0: int, zp: _ZoomPlan, best_pow: np.ndarray, best_bin: np.ndarray): """Evaluate fine bins k0..k0+m-1 for the given rows and fold the result into the per-row best (power, bin), preserving np.argmax's lowest-bin tie-break.""" ph = zp.phases.get(k0) if ph is None: ph = np.exp((-2j * np.pi * k0 / zp.n_fft) * zp.j).astype(np.complex64) zp.phases[k0] = ph F = (waves[rows] * ph) @ zp.E q = F.real ** 2 q += F.imag ** 2 i = np.argmax(q, axis=1) p = q[np.arange(len(rows)), i] b = k0 + i upd = (p > best_pow[rows]) | ((p == best_pow[rows]) & (b < best_bin[rows])) ridx = rows[upd] best_pow[ridx] = p[upd] best_bin[ridx] = b[upd] def _peak_bins_zoom(waves: np.ndarray, zp: _ZoomPlan) -> np.ndarray: """Peak fine-bin per waveform without materialising the padded spectrum: coarse rfft, then a small fine DFT (one gemm per shared window) on the exact n_fft bin grid. Identity with the full padded argmax is enforced by tests/test_compute.py::test_zoom_identity and the golden-hash sweep.""" n_wf = waves.shape[0] S = _block_rfft(waves, zp.n_coarse) P = S.real ** 2 P += S.imag ** 2 P[:, 0] = 0.0 p_c = np.max(P, axis=1) best_pow = np.full(n_wf, -1.0, dtype=np.float32) best_bin = np.full(n_wf, np.iinfo(np.int64).max, dtype=np.int64) # For a clean signal this yields one or two windows; a noise spectrum # (many near-equal peaks) yields a dozen or so — still a tiny fraction # of the padded grid. thr = np.where(p_c > 0, np.float32(_ZOOM_CAND_RATIO) * p_c, np.float32(np.inf)) rows_c, bins_c = np.nonzero(P >= thr[:, None]) k0_c = _window_start(bins_c, zp) # Zeroing the coarse DC bin (suppression) blinds the candidate scan to # fine bins closer to DC than the first coarse sample — where the # DC-leakage skirt of an un-subtracted offset peaks. Always refine the # DC-adjacent window too. rows_c = np.concatenate([rows_c, np.arange(n_wf)]) k0_c = np.concatenate([k0_c, np.ones(n_wf, dtype=np.int64)]) pair = np.unique(np.stack([rows_c, k0_c], axis=1), axis=0) for k0 in np.unique(pair[:, 1]): _refine_window(waves, pair[pair[:, 1] == k0, 0], int(k0), zp, best_pow, best_bin) # An all-zero spectrum must reproduce argmax-of-zeros = bin 0. best_bin[p_c == 0.0] = 0 return best_bin def _peak_bins_direct(waves: np.ndarray, n_len: int) -> np.ndarray: """Peak bin per waveform via the full transform — for pad factors too small for the zoom search to pay.""" S = _block_rfft(waves, n_len) power = S.real ** 2 power += S.imag ** 2 power[:, 0] = 0.0 return np.argmax(power, axis=1) # --------------------------------------------------------------------------- # Chunking / parallel budget # --------------------------------------------------------------------------- # # Row chunks are budgeted so all concurrently-live working buffers fit in # memory; the worker count is derived from the budget, not vice versa. # Rationale and the measured 1024 MB default: docs/design.md ("Memory budget # and row chunking"). _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) 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_fft_rows(n_frames: int, samples_per_frame: int, budget: int) -> int: """Rows per outer chunk for the block-FFT path. Only the float32 waveform buffer scales with the chunk (in-flight block spectra total a few MB across the whole pool), so budget it with 2x slack and let the block fan-out saturate the pool regardless of pad factor.""" bytes_per_row = max(1, n_frames * samples_per_frame * 4 * 2) 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)) # --------------------------------------------------------------------------- # Row-averaged FFT: same-row, distance-weighted CH1 waveform smoothing # # Averages a pixel's CH1 waveform with its same-row neighbors *before* the # FFT peak search, to improve SNR on noisy scans. Never crosses rows: pixel # pitch is strongly anisotropic and varies by scan, but the X pitch within # one row is a single file-wide constant (SrasFile.pixel_x_mm), so a # Gaussian in pixel-index distance along a row and one in true physical mm # distance are the same function up to that constant scale factor — the # kernel itself needs no pitch, only the GUI's physical-width hint label # does. Rationale and the masked-average/background-subtraction proofs: # docs/design.md ("Row-averaged FFT"). # --------------------------------------------------------------------------- _ROW_AVG_SIGMA_FRAC = 0.5 # sigma = n * this; edge weight (distance n) is # exp(-1/(2*frac**2)) ~= 0.135 of the center tap def _row_average_weights(row_avg_n: int) -> np.ndarray: """(2n+1,) float32 Gaussian weights for distance-weighted row averaging, symmetric around the center tap. Not pre-normalized to sum to 1 — _row_average_waveforms renormalizes per output pixel by the actual sum of included, valid, in-window neighbor weights, not a fixed total.""" n = max(0, int(row_avg_n)) if n == 0: return np.ones(1, dtype=np.float32) d = np.arange(-n, n + 1, dtype=np.float64) sigma = n * _ROW_AVG_SIGMA_FRAC return np.exp(-0.5 * (d / sigma) ** 2).astype(np.float32) def _row_average_waveforms(masked_waves: np.ndarray, valid: np.ndarray, weights: np.ndarray) -> np.ndarray: """Distance-weighted mean of each row position's CH1 waveform with its same-row neighbors, counting only neighbors where *valid* is True. masked_waves: (n_frames, spf) float32 — CH1 samples at valid[f] positions; must already be 0.0 elsewhere (the caller must never have read the raw memmap at an invalid position). valid: (n_frames,) bool. weights: (2n+1,) float32 from _row_average_weights. Returns (n_frames, spf) float32, meaningful only where valid is True (the caller compacts by that same mask right after, matching compute_rf_image's existing masked-compaction contract). Two 1-D correlations along the frame axis — the numerator against the raw masked-zero waveforms, the denominator against the validity mask itself — so a masked neighbor contributes *zero weight* rather than a zero-amplitude sample at full weight, and a window truncated at a row's edge renormalizes correctly with no separate edge case: mode="constant", cval=0.0 pads both convolutions with zero beyond the row's own ends. """ num = scipy_ndimage.correlate1d(masked_waves, weights, axis=0, mode="constant", cval=0.0) den = scipy_ndimage.correlate1d(valid.astype(np.float32), weights, axis=0, mode="constant", cval=0.0) den_safe = np.where(valid, den, np.float32(1.0)) return num / den_safe[:, None] def cached_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, row_avg_n: int = 0, allow_dc_recompute: bool = True) -> np.ndarray | None: """The precomputed-cache fast path for compute_rf_image: a ready-to- display peak-frequency image if this file already has one matching every setting the caller cares about, else None (caller must run a real FFT). A stored image stands in only when *all* hold: this angle actually has a stored image (v5 PREC or v7 CACH); no custom zero-padding is requested (n_fft is None) — the store is always natural-resolution; the stored bg-sub flag matches what the caller wants; and sras.precomputed_row_avg_n == row_avg_n exactly (0 == "raw") — this last check is what stops a raw request from ever being silently served a row-averaged image, or vice versa, or a request at one window size being served a cache stored at a different one. If *allow_dc_recompute* is False and no DC4 image is already cached or supplied via *dc4_mv*, applying the mask would mean reading a whole channel on the caller's behalf; this returns None instead so a caller that wants to stay off the I/O path (e.g. a GUI thread) can choose to fall through to a real compute rather than block. """ cached_freq = sras.precomputed_freq_mhz[angle_idx] if (cached_freq is None or n_fft is not None or sras.precomputed_bg_sub != (apply_bg_sub and sras.background is not None) or sras.precomputed_row_avg_n != row_avg_n): return 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: if dc4_mv is not None: dc4_img = dc4_mv elif allow_dc_recompute: dc4_img = adc_to_mv(compute_dc_image(sras, angle_idx, CH4_IDX), *sras.cal(CH4_IDX)) else: return None freq_img[dc4_img < dc_threshold_mv] = 0.0 return freq_img 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, exact: bool = False, row_avg_n: int = 0) -> 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. At pad factors >= _ZOOM_MIN_PAD the padded spectrum is never materialised: a coarse rfft finds each peak and a local fine DFT resolves it on the exact n_fft bin grid (_peak_bins_zoom). *exact* forces the reference full-padded transform instead — it exists for tests, audits, and the SRAS_FFT_EXACT=1 escape hatch, and is slow and memory-hungry at high pad. *row_avg_n* > 0 averages each pixel's CH1 waveform with its up-to-n same-row neighbors (distance-weighted, valid-neighbors-only per the same dc_threshold_mv mask) before the FFT runs — see _row_average_waveforms. 0 (default) is the raw, unaveraged behavior. Fast path: if the file has a precomputed peak-frequency image for this angle (v5 PREC or v7 CACH) matching every one of the caller's settings — including row_avg_n exactly — the stored image is used directly, no FFT is run. See cached_rf_image. """ n_rows, n_frames = sras.image_shape(angle_idx) data = sras.data[angle_idx] fast = cached_rf_image(sras, angle_idx, dc_threshold_mv, apply_bg_sub=apply_bg_sub, n_fft=n_fft, dc4_mv=dc4_mv, row_avg_n=row_avg_n) if fast is not None: return fast # ---- Chunked FFT path -------------------------------------------------- exact = exact or _FFT_EXACT_ENV spf = sras.samples_per_frame n_len = n_fft if n_fft is not None else spf freq32 = sras.freq_axis_mhz(n_fft).astype(np.float32) img = np.zeros((n_rows, n_frames), dtype=np.float32) background = sras.background if (apply_bg_sub and sras.background is not None) else None cal4 = sras.cal(CH4_IDX) row_avg_weights = _row_average_weights(row_avg_n) if row_avg_n > 0 else None zp = (_zoom_plan(spf, n_fft) if not exact and n_fft is not None and n_fft >= _ZOOM_MIN_PAD * spf else None) total = _TOTAL_BYTES_BUDGET if budget is None else max(1, budget) if row_avg_n > 0: # One extra same-sized transient buffer (the pre-averaging full-row # scratch array) is live per in-flight row; halve the budget so # _plan_fft_rows/the exact-path sizing accounts for it rather than # relying on _plan_fft_rows's existing 2x slack to happen to cover it. total = max(1, total // 2) cap = max_workers if max_workers is not None else _MAX_WORKERS if exact: # The reference path materialises the full padded spectrum, so rows # are budgeted against it (complex64 + power + temp per bin) and the # chunk runs as one serial transform. bytes_per_row = max(1, n_frames * (4 * spf + 16 * (n_len // 2 + 1))) chunk_rows = int(max(1, min(_CHUNK_ROWS_MAX, total // bytes_per_row))) cap = 1 else: chunk_rows = _plan_fft_rows(n_frames, spf, total) pool = _fft_pool() if cap > 1 else None def process(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 counts = (valid.sum(axis=1) if valid is not None else np.full(r1 - r0, n_frames, dtype=np.int64)) offs = np.concatenate(([0], np.cumsum(counts))) n_wf = int(offs[-1]) waves = np.empty((n_wf, spf), dtype=np.float32) def read_row(i: int): if should_stop is not None and should_stop(): return raw = data[r0 + i, CH1_IDX] dst = waves[offs[i]:offs[i + 1]] row_valid = valid[i] if valid is not None else None if row_avg_weights is not None: v = row_valid if row_valid is not None else np.ones(n_frames, dtype=bool) # A plain cast + broadcast multiply, not a boolean-indexed # scatter into a zeroed buffer: numpy's fancy indexing holds # the GIL for its whole duration (measured zero speedup # across threads, and net negative once threads outnumber # physical cores), while a cast and a multiply are ordinary # ufuncs that release it — this is what lets read_row actually # parallelize across the pool instead of serializing on the # scatter/gather. Trade-off: every sample in the row is read, # valid or not (row averaging needs broad neighbor context # regardless, unlike the plain path below, which still skips # masked-out pixels entirely). full = raw.astype(np.float32) * v[:, None] avg = _row_average_waveforms(full, v, row_avg_weights) dst[:] = avg[v] if row_valid is not None else avg else: # 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. dst[:] = raw[row_valid] if row_valid is not None else raw if background is not None: dst -= background # background is 1-D (spf,) if pool is None: for i in range(r1 - r0): read_row(i) else: list(pool.map(read_row, range(r1 - r0))) out = np.empty(n_wf, dtype=np.float32) def fft_block(b0: int): if should_stop is not None and should_stop(): return b1 = min(b0 + _FFT_BLOCK, n_wf) w = waves[b0:b1] bins = (_peak_bins_zoom(w, zp) if zp is not None else _peak_bins_direct(w, n_len)) out[b0:b1] = freq32[bins] if exact: spectrum = _do_rfft(waves, n=n_fft, axis=-1, workers=1) power = spectrum.real ** 2 power += spectrum.imag ** 2 power[:, 0] = 0.0 # suppress DC bin out[:] = freq32[np.argmax(power, axis=1)] elif pool is None: for b0 in range(0, n_wf, _FFT_BLOCK): fft_block(b0) else: list(pool.map(fft_block, range(0, n_wf, _FFT_BLOCK))) # Boolean scatter is row-major, matching the read pass's # concatenation order. if valid is not None: img[r0:r1][valid] = out else: img[r0:r1] = out.reshape(r1 - r0, n_frames) # BLAS must not thread under the pool (the fine-DFT gemm would multiply # against the pool's own workers). limiter = (threadpool_limits(limits=1) if pool is not None and threadpool_limits is not None else None) try: for r0 in range(0, n_rows, chunk_rows): if should_stop is not None and should_stop(): break process(r0, min(r0 + chunk_rows, n_rows)) finally: if limiter is not None: limiter.unregister() return img # --------------------------------------------------------------------------- # Batch cache (Convert menu) — process-pool entry point # --------------------------------------------------------------------------- def cache_file(path: str, mode: str, apply_bg_sub: bool, fft_backend: str = "scipy", max_workers: int = 0, dc_threshold_mv: float | None = None, row_avg_n: 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. mode is "dc", "fft" (raw per-pixel FFT, natural-resolution, unmasked — masking applied at display time), or "fft_rowavg" (same-row, distance-weighted CH1 averaging before the FFT — see compute_rf_image's row_avg_n). fft_rowavg needs *dc_threshold_mv* up front, unlike plain "fft": neighbor validity is baked into the stored numbers, so it can't be deferred to display time the way plain masking can. The stored FFT cache is always natural-resolution (pad 1): the v7 SFFT block records no pad factor, and padded views compute live fast enough (see _peak_bins_zoom) that caching them is not worth a format change. """ global _MAX_WORKERS try: if mode not in ("dc", "fft", "fft_rowavg"): return f"unknown cache mode {mode!r} (expected 'dc', 'fft', or 'fft_rowavg')" 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 n_workers, angle_budget = plan_angle_level(sras) if mode == "dc": dc3 = _parallel_map( lambda a: adc_to_mv( compute_dc_image(sras, a, CH3_IDX, max_workers=1, budget=angle_budget), *sras.cal(CH3_IDX)), range(n), n_workers) dc4 = _parallel_map( lambda a: adc_to_mv( compute_dc_image(sras, a, CH4_IDX, max_workers=1, budget=angle_budget), *sras.cal(CH4_IDX)), range(n), n_workers) sras.write_v7_cache(new_dc3_mv=dc3, new_dc4_mv=dc4) elif mode == "fft": 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. The FFT path parallelises # internally over blocks, so angles run one at a time with the # full budget. 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) else: # "fft_rowavg" if row_avg_n <= 0: return "row_avg_n must be a positive neighbor half-width for fft_rowavg mode" if dc_threshold_mv is None: return ("fft_rowavg mode requires a DC threshold " "(neighbor validity depends on it)") effective_bg = apply_bg_sub and sras.background is not None freq = [compute_rf_image(sras, a, dc_threshold_mv=dc_threshold_mv, apply_bg_sub=effective_bg, row_avg_n=row_avg_n) for a in range(n)] sras.write_v7_cache(new_freq_mhz=freq, new_bg_sub=effective_bg, new_row_avg_n=row_avg_n) return "" except Exception as exc: return str(exc) # --------------------------------------------------------------------------- # Angle alignment (Fusion menu) # # Rigid transforms only (rotation + translation, never scale), computed in mm # on two frames: each angle's "local mm" (origin at its own array center) and # the reference angle's local mm. Angle 0 is the sole coordinate authority; # every other angle is placed purely by image content. Why, and the full # frame/affine conventions: docs/design.md ("Angle alignment coordinate # frames"). # --------------------------------------------------------------------------- @dataclass class AngleTransform: rotation_deg: float shift_mm: tuple[float, float] # (dx_mm, dy_mm) in ref mm matrix: np.ndarray # (2,2): canvas (row,col) -> this angle's raw (row,col) offset: np.ndarray # (2,) score: float = 1.0 # registration NCC (1.0 = reference/manual) source: str = "" # image that won registration: "signal"/"mask" @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 rigid map from this angle's local mm to ref mm, the pair a canvas-bound AngleTransform's matrix/offset get built from once a canvas is decided (build_manual_alignment). Defaults to identity: a fresh angle with no prior alignment is shown centered on the reference with no rotation, which is the same "fully unaligned" state Clear Alignment resets back to. """ rotation_deg: float = 0.0 shift_mm: tuple[float, float] = (0.0, 0.0) # Side length of the grid the rotation refinement runs on. The whole cost of # registration scales with it: ~325 ms per candidate rotation at 640, roughly # quadrupling per doubling, against ~200 MB of live masked-correlation buffers # (see _registration_workers). 640 puts a full-size scan at ~0.1 mm/px, which # resolves rotation on a 20 mm sample to well under a tenth of a degree. _DEFAULT_FINE_DIM = 640 @dataclass class RigidFit: """What register_angle_to_reference found for one angle.""" rotation_deg: float shift_mm: tuple[float, float] score: float # zero-mean NCC over the valid overlap source: str # "signal", "mask", or "reference" def _skimage_phase_cross_correlation(): """Import skimage.registration lazily and cache it. Deliberately not a module-level import: this module is imported by every multiprocessing child (see the module docstring), skimage costs ~0.6 s to import, and no child ever registers anything — registration runs in GUI- process threads. """ global _pcc try: return _pcc except NameError: from skimage.registration import phase_cross_correlation as _fn _pcc = _fn return _pcc # ---- Geometry: local mm, ref mm, and the one affine builder --------------- 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). dy keeps its sign, so +row always means the same physical direction as +y.""" 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 _center_idx(sras: SrasFile, angle_idx: int) -> np.ndarray: """(row, col) index of this angle's array center — the origin of its local mm frame. Purely geometric: it depends on the array shape and nothing else, which is what keeps local mm free of stage position.""" n_rows, n_frames = sras.image_shape(angle_idx) return np.array([(n_rows - 1) / 2.0, (n_frames - 1) / 2.0]) def ref_center_mm(sras: SrasFile, ref_angle_idx: int) -> np.ndarray: """Stage mm of the reference angle's array center: ref mm + this == stage mm. The single bridge between ref mm and stage mm, and the only place in the whole alignment path where any angle's stage position is read at all — which is why it takes the *reference* index by name rather than an arbitrary angle. """ x = sras.x_axis_mm(ref_angle_idx) y = sras.y_positions_mm(ref_angle_idx) return np.array([(x[0] + x[-1]) / 2.0, (y[0] + y[-1]) / 2.0], dtype=np.float64) def _local_half_extent_mm(sras: SrasFile, angle_idx: int) -> tuple[float, float]: """(half width, half height) in mm from this angle's array center to the center of its outermost pixel.""" n_rows, n_frames = sras.image_shape(angle_idx) dx, dy = pixel_pitch_mm(sras, angle_idx) return (n_frames - 1) / 2.0 * abs(dx), (n_rows - 1) / 2.0 * abs(dy) 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 nominal_delta_deg(sras: SrasFile, angle_idx: int, ref_idx: int) -> float: """The rotation-stage's own reported angle change between two angles. Used only to *seed* the rotation search, never as the answer: the stage's sign convention relative to this module's math-positive (CCW, x toward y) convention in scan mm is not knowable from the file, so register_angle_to_reference scores both +this and -this and lets the image content decide (see _rotation_candidates). """ return float(sras.angles_deg[angle_idx] - sras.angles_deg[ref_idx]) def _footprint_corners_ref_mm(sras: SrasFile, angle_idx: int, rotation_deg: float, shift_mm: tuple[float, float]) -> np.ndarray: """This angle's 4 footprint corners mapped into ref mm by its rigid transform, shape (4, 2). Zero-padding cost, not content: it is the scan window's corners, which is what the shared canvas has to cover.""" hw, hh = _local_half_extent_mm(sras, angle_idx) corners = np.array([[sx * hw, sy * hh] for sx in (-1.0, 1.0) for sy in (-1.0, 1.0)]) R = _rotation_matrix(rotation_deg) return corners @ R.T + np.asarray(shift_mm, dtype=np.float64) def _affine_out_to_src(*, out_pitch_mm: tuple[float, float], out_origin_ref_mm, src_dx_mm: float, src_dy_mm: float, src_center_idx, src_center_off_mm=(0.0, 0.0), rotation_deg: float = 0.0, shift_mm: tuple[float, float] = (0.0, 0.0) ) -> tuple[np.ndarray, np.ndarray]: """matrix, offset s.t. src_index = matrix @ [row_out, col_out] + offset, matching scipy.ndimage.affine_transform's output->input convention. The single affine builder behind every resampling in this module — the registration grid, the manual-alignment preview and the final canvas all differ only in their arguments. Pipeline (mm unless noted): [X;Y] = A_out @ [row_out;col_out] + out_origin_ref_mm # out idx -> ref mm [lx;ly] = R(rotation)^T @ ([X;Y] - shift) # ref mm -> this angle's local mm [row;col] = D @ ([lx;ly] - src_center_off) + src_center_idx *src_center_off_mm* is the local-mm position of the source array's own center, nonzero only when the source has been block-mean downsampled (its center can land up to half a block off the full-resolution center — see _prepare_reg_image). Everything is expressed relative to array centers, so no per-angle stage coordinate appears anywhere in here. """ Rinv = _rotation_matrix(rotation_deg).T A_out = np.array([[0.0, out_pitch_mm[0]], [out_pitch_mm[1], 0.0]]) # [row,col] -> [x,y] D = np.array([[0.0, 1.0 / src_dy_mm], [1.0 / src_dx_mm, 0.0]]) # [x,y] -> [row,col] b_out = np.asarray(out_origin_ref_mm, dtype=np.float64) shift = np.asarray(shift_mm, dtype=np.float64) off = np.asarray(src_center_off_mm, dtype=np.float64) matrix = D @ Rinv @ A_out offset = D @ (Rinv @ (b_out - shift) - off) + np.asarray(src_center_idx, dtype=np.float64) 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 image) 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_2d(img: np.ndarray, fy: int, fx: int) -> np.ndarray: """Block-mean by independent row/column factors. Independent factors matter because the raw grid is strongly anisotropic (5 µm along x, 50 µm along y): a single square factor would either alias along x or throw away rows.""" fy, fx = max(1, int(fy)), max(1, int(fx)) if fy == 1 and fx == 1: return img h, w = img.shape h2, w2 = (h // fy) * fy, (w // fx) * fx if h2 == 0 or w2 == 0: return img trimmed = img[:h2, :w2] return trimmed.reshape(h2 // fy, fy, w2 // fx, fx).mean(axis=(1, 3)) # ---- Registration: rigid (rotation + translation) fit against the reference # # Everything below works on _RegImage, an angle's image resampled to a shared # *isotropic* grid centered on its own array center. Nothing in here can see a # stage coordinate even in principle, which is the point: the fit is decided by # image content alone. @dataclass class _RegImage: """One angle's image prepared for registration: block-mean downsampled to roughly the registration pitch, carrying the physical pitch it ended up with and the local-mm offset of its own array center from the full-resolution array center (block-mean trims a partial trailing block, so the two centers can differ by up to half a block).""" img: np.ndarray dx_mm: float dy_mm: float center_off_mm: tuple[float, float] def _block_center_off(n_full: int, n_small: int, factor: int, pitch_mm: float) -> float: """Local-mm offset of a block-mean-downsampled array's own center from the full-resolution array's center along one axis. Small pixel k averages full pixels [k*f, k*f + f - 1], so its center sits at full index k*f + (f-1)/2; block-mean also trims a partial trailing block. Both effects together move the small array's center by up to half a block, which has to be accounted for or a downsampled image registers (or previews) at a systematically shifted position. """ return (((n_small - 1) / 2.0) * factor + (factor - 1) / 2.0 - (n_full - 1) / 2.0) * pitch_mm def _prepare_reg_image(sras: SrasFile, angle_idx: int, img: np.ndarray, pitch_mm: float) -> _RegImage: """Block-mean an angle's image down to roughly *pitch_mm* before it is resampled onto the registration grid. Pre-averaging matters: the raw grid is 10x finer along x than along y, so sampling it directly at the (much coarser) isotropic registration pitch would alias badly along x.""" dx, dy = pixel_pitch_mm(sras, angle_idx) fx = max(1, int(pitch_mm / abs(dx))) fy = max(1, int(pitch_mm / abs(dy))) small = block_mean_2d(np.asarray(img, dtype=np.float32), fy, fx) n_rows, n_frames = img.shape return _RegImage( small, dx * fx, dy * fy, (_block_center_off(n_frames, small.shape[1], fx, dx), _block_center_off(n_rows, small.shape[0], fy, dy))) def _embed(reg: _RegImage, pitch_mm: float, n: int, rotation_deg: float, order: int) -> np.ndarray: """Resample a _RegImage onto the shared n x n isotropic registration grid, rotated by *rotation_deg* about the grid center and with no translation (translation is what phase correlation then measures).""" half = (n - 1) / 2.0 * pitch_mm matrix, offset = _affine_out_to_src( out_pitch_mm=(pitch_mm, pitch_mm), out_origin_ref_mm=(-half, -half), src_dx_mm=reg.dx_mm, src_dy_mm=reg.dy_mm, src_center_idx=((reg.img.shape[0] - 1) / 2.0, (reg.img.shape[1] - 1) / 2.0), src_center_off_mm=reg.center_off_mm, rotation_deg=rotation_deg) return scipy_ndimage.affine_transform( reg.img, matrix, offset=offset, output_shape=(n, n), order=order, mode="constant", cval=0.0) def _embed_with_valid(reg: _RegImage, pitch_mm: float, n: int, rotation_deg: float) -> tuple[np.ndarray, np.ndarray]: """Embedded image plus the boolean mask of where that angle actually has data. The valid mask is what lets registration ignore each angle's differently-shaped scan window instead of locking onto its silhouette.""" img = _embed(reg, pitch_mm, n, rotation_deg, order=1) ones = _RegImage(np.ones_like(reg.img), reg.dx_mm, reg.dy_mm, reg.center_off_mm) valid = _embed(ones, pitch_mm, n, rotation_deg, order=0) > 0.5 return img, valid def _shift_into(img: np.ndarray, dr: int, dc: int) -> np.ndarray: """img translated by whole pixels with zero fill (never wrapping, unlike np.roll — wrapped content would score as a spurious match).""" out = np.zeros_like(img) h, w = img.shape sr0, sr1 = max(0, dr), min(h, h + dr) sc0, sc1 = max(0, dc), min(w, w + dc) if sr0 >= sr1 or sc0 >= sc1: return out out[sr0:sr1, sc0:sc1] = img[sr0 - dr:sr1 - dr, sc0 - dc:sc1 - dc] return out def _overlap_ncc(ref: np.ndarray, ref_valid: np.ndarray, mov: np.ndarray, mov_valid: np.ndarray, min_overlap_frac: float = 0.15) -> float: """Zero-mean normalized cross-correlation over the two images' common valid region — the score every rotation candidate is ranked by. Computed on the overlap only, and rejected outright (-1) when the overlap is too small a fraction of the smaller footprint: without that floor a candidate that slides the angles almost entirely apart can win on a handful of coincidentally-similar pixels. """ both = ref_valid & mov_valid n = int(both.sum()) smaller = min(int(ref_valid.sum()), int(mov_valid.sum())) if smaller == 0 or n < min_overlap_frac * smaller or n < 16: return -1.0 a = ref[both].astype(np.float64) b = mov[both].astype(np.float64) a -= a.mean() b -= b.mean() denom = np.sqrt((a * a).sum() * (b * b).sum()) return float((a * b).sum() / denom) if denom > 0 else -1.0 def _masked_shift(ref: np.ndarray, ref_valid: np.ndarray, mov: np.ndarray, mov_valid: np.ndarray) -> tuple[int, int]: """Integer (dr, dc) that best registers *mov* onto *ref*, from skimage's masked FFT phase correlation (Padfield). The masked variant is the whole reason scikit-image is a dependency: plain phase correlation on these images locks onto the scan window's rectangular silhouette, which differs per angle, instead of onto the sample.""" pcc = _skimage_phase_cross_correlation() result = pcc(ref, mov, reference_mask=ref_valid, moving_mask=mov_valid) shift = result[0] if isinstance(result, tuple) else result return int(round(float(shift[0]))), int(round(float(shift[1]))) def _subpixel_residual(ref: np.ndarray, mov: np.ndarray, both: np.ndarray) -> tuple[float, float]: """Sub-pixel leftover shift between two already integer-aligned images, from upsampled phase correlation over their common valid region. A separate pass because skimage's *masked* phase correlation has no upsample_factor; here both images are zeroed outside the shared overlap and mean-subtracted inside it, so the plain upsampled version is well posed. Clamped to ±1 px: this only ever polishes an already-good integer fit, and a larger "residual" means the peak was spurious. normalization=None (plain cross-correlation, not phase correlation) is load-bearing. Whitening the spectrum is what makes phase correlation good at finding a large unknown shift, but here the two images are already aligned to within a pixel and the masked-off surroundings put a hard edge in both: whitened, that edge and the high-frequency noise swamp the true sub-pixel peak and the default returns a flat zero every time. """ if not both.any(): return 0.0, 0.0 a = np.zeros_like(ref) b = np.zeros_like(mov) a[both] = ref[both] - ref[both].mean() b[both] = mov[both] - mov[both].mean() if not (np.any(a) and np.any(b)): return 0.0, 0.0 pcc = _skimage_phase_cross_correlation() result = pcc(a, b, upsample_factor=20, normalization=None) shift = result[0] if isinstance(result, tuple) else result dr, dc = float(shift[0]), float(shift[1]) if abs(dr) > 1.0 or abs(dc) > 1.0: return 0.0, 0.0 return dr, dc def _reg_pitch_and_size(sras: SrasFile, max_dim: int, margin: float = 1.25) -> tuple[float, int]: """Isotropic pitch (mm/px) and side length for the shared registration grid: square, big enough for the largest angle's footprint at any rotation (hence its diagonal) plus *margin* headroom for the translation search. The floor on pitch is the *geometric mean* of the two native pitches, not the coarser of them. The raw grid is anisotropic (5 µm along x, 50 µm along y): flooring at 50 µm would throw away all the extra x detail, and rotation precision depends directly on it — a feature at radius r moves by r·δθ, so at 50 µm a 20 mm-wide sample can only resolve rotation to a few tenths of a degree. The geometric mean interpolates y up by ~3x rather than discarding x, which costs a little memory and buys real angular precision. On a full-size scan max_dim binds first and this floor never applies at all. """ diag = max(float(np.hypot(*(2 * v for v in _local_half_extent_mm(sras, a)))) for a in range(sras.n_angles)) span = diag * margin native = float(np.sqrt( abs(sras.pixel_x_mm) * max(abs(pixel_pitch_mm(sras, a)[1]) for a in range(sras.n_angles)))) pitch = max(span / max_dim, native) n = int(scipy_fft.next_fast_len(max(16, int(np.ceil(span / pitch))))) return pitch, n def _registration_workers(sras: SrasFile, fine_dim: int) -> int: """How many angles may register concurrently. Not plan_angle_level: that budgets for waveform chunks, and registration never touches a waveform — it works on already-computed DC images and a few grid-sized arrays. The real limit is the masked phase correlation, which pads to roughly twice the grid and holds several complex128 arrays of that size live at once, so the count is derived from the same _TOTAL_BYTES_BUDGET the rest of the module honours. The estimate below comes out at 262 MB for the default fine_dim=640, against 199 MB measured — deliberately on the pessimistic side, since overshooting the budget costs swapping while undershooting only costs a little wall time. """ per_worker = 10 * (2 * fine_dim) ** 2 * 16 # ~10 complex128 grids return int(max(1, min(_MAX_WORKERS, sras.n_angles, _TOTAL_BYTES_BUDGET // max(1, per_worker)))) def registration_workers(sras: SrasFile) -> int: """Public: concurrent-registration cap at the default fine grid.""" return _registration_workers(sras, _DEFAULT_FINE_DIM) def default_max_workers() -> int: """Public: the module-wide worker cap (SRAS_MAX_WORKERS or cpu count).""" return _MAX_WORKERS def _rotation_candidates(nominal_deg: float, search_deg: float, step_deg: float, signs: tuple[int, ...] = (-1, 1)) -> list[float]: """Coarse rotation candidates: a window around each requested sign of the stage's reported angle change. Scoring both signs (the default) is what makes the stage's sign convention a non-issue — the images decide which way the stage turns, and a file whose stage reports the opposite sense registers just as well. *signs* exists only so a user who has established which way their own stage turns can halve the coarse sweep from the alignment wizard. It is an override, not an inference: nothing in the file says which sign is right. """ out: list[float] = [] step_deg = max(abs(step_deg), 1e-9) # a 0 step would divide by zero below for sign in signs: center = sign * nominal_deg k = int(np.floor(search_deg / step_deg)) for i in range(-k, k + 1): out.append(center + i * step_deg) # Dedupe (the two windows coincide when nominal_deg is 0) while keeping order. seen: set[float] = set() return [t for t in out if not (round(t, 6) in seen or seen.add(round(t, 6)))] def _score_rotation(ref_img: np.ndarray, ref_valid: np.ndarray, mov: _RegImage, pitch: float, n: int, theta: float, subpixel: bool = False) -> tuple[float, tuple[float, float]]: """Best score this rotation can reach, and the translation that reaches it: rotate, phase-correlate for the shift, score the overlap. *subpixel* also removes the leftover sub-pixel translation before scoring. That matters more than it sounds: without it every candidate is scored at whole-pixel alignment, so rotations that differ by less than one pixel of rim displacement are ranked by quantization noise rather than by fit, and the refinement stalls a degree or so off. Left off for the coarse sweep, which only has to pick a basin, and on for the refinement. """ img, valid = _embed_with_valid(mov, pitch, n, theta) dr, dc = _masked_shift(ref_img, ref_valid, img, valid) shifted = _shift_into(img, dr, dc) shifted_valid = _shift_into(valid.astype(np.float32), dr, dc) > 0.5 if subpixel: sub_dr, sub_dc = _subpixel_residual( ref_img, shifted, ref_valid & shifted_valid) if sub_dr or sub_dc: shifted = scipy_ndimage.shift(shifted, (sub_dr, sub_dc), order=1, mode="constant", cval=0.0) dr, dc = dr + sub_dr, dc + sub_dc return (_overlap_ncc(ref_img, ref_valid, shifted, shifted_valid), (float(dr), float(dc))) def _refine_rotation(ref_img: np.ndarray, ref_valid: np.ndarray, mov: _RegImage, pitch: float, n: int, theta: float, score: float, shift: tuple[float, float], step_deg: float, min_step_deg: float = 0.05, max_evals: int = 80 ) -> tuple[float, float, tuple[float, float]]: """Hill-climb the rotation from the coarse winner: step out while the score improves, halve the step when it doesn't, stop below *min_step_deg*. A walking search rather than a fixed grid around the coarse winner, because the coarse stage ranks rotations at a coarse pitch where a degree can be worth less than the translation quantization — its winner can legitimately land a degree or two off, which a fixed ±½° refinement window could never recover from. Every candidate here is scored with subpixel=True, matching how the incoming *score* was measured. Mixing the two is silently fatal: the sub-pixel-corrected score is strictly the higher of the two, so a subpixel-scored start compared against un-corrected candidates can never be beaten and the search sits still at whatever the coarse stage handed it. """ evals = 0 while step_deg >= min_step_deg and evals < max_evals: trials = [] for cand in (theta - step_deg, theta + step_deg): s, sh = _score_rotation(ref_img, ref_valid, mov, pitch, n, cand, subpixel=True) evals += 1 trials.append((s, cand, sh)) best_s, best_cand, best_sh = max(trials, key=lambda t: t[0]) if best_s > score: theta, score, shift = best_cand, best_s, best_sh else: step_deg /= 2.0 return theta, score, shift def _source_images(sras: SrasFile, angle_idx: int, ref_angle_idx: int, signal_mv: dict[int, np.ndarray], dc_threshold_mv: float, sources) -> list[tuple[str, np.ndarray, np.ndarray]]: """(name, reference image, moving image) per requested registration source. "signal" is each angle's own CH4 image minus its own minimum — subtracting per-angle rather than globally keeps a nonzero DC baseline from reading as a high-contrast edge against the zero padding. "mask" is the binarized >= dc_threshold_mv image, the same silhouette the overlay draws. A mask that is empty or completely full for either angle carries no registration information at all, so that source is dropped rather than scored. """ out = [] for name in sources: pair = [] for a in (ref_angle_idx, angle_idx): img = signal_mv[a] if name == "mask": m = img >= dc_threshold_mv if not m.any() or m.all(): pair = [] break pair.append(m.astype(np.float32)) else: pair.append((img - img.min()).astype(np.float32)) if pair: out.append((name, pair[0], pair[1])) return out def register_angle_to_reference( sras: SrasFile, angle_idx: int, ref_angle_idx: int, signal_mv: dict[int, np.ndarray], *, dc_threshold_mv: float = 0.0, sources: tuple[str, ...] = ("signal", "mask"), coarse_dim: int = 256, fine_dim: int = _DEFAULT_FINE_DIM, search_deg: float = 6.0, coarse_step_deg: float = 2.0, seed_deg: float | None = None, seed_signs: tuple[int, ...] = (-1, 1), refine: bool = True) -> RigidFit: """Rigid (rotation + translation, never scale) fit of *angle_idx* onto *ref_angle_idx*, found entirely by cross-correlating image content. Two stages: 1. Coarse sweep at *coarse_dim* over a ±*search_deg* window around the requested signs of the stage's reported angle change (see _rotation_candidates), for each requested source image, scored by _overlap_ncc. Only has to pick the right basin. 2. Hill-climbing refinement of the winning (source, rotation) at *fine_dim* with sub-pixel translation folded into every score (_refine_rotation, _score_rotation), down to 0.05°. The returned rotation_deg/shift_mm are the rigid map from this angle's local mm to ref mm (q = R @ l + shift); score is the final NCC, which the caller can surface so a bad scan is visible rather than silently fused in. Returns identity for the reference angle itself. The last three arguments only exist to let the alignment wizard expose the rotation search; every default reproduces the search this function has always done. *seed_deg* replaces the stage's reported angle change as the center of the coarse sweep — the pre-rotation the search starts from. None (the default) means the stage angle from the file's Angle Table, which is nearly always what you want: it puts the sweep within a couple of degrees of the answer. Pass 0.0 to search around no rotation at all, which is the honest choice for a file whose stage angles are known to be wrong. *refine=False* stops after the coarse sweep, leaving rotation on the *coarse_step_deg* grid. Combined with search_deg=0.0 and a single seed sign it pins rotation to exactly the seed and searches translation only — for a scan whose stage angles are trusted more than its image content. """ if angle_idx == ref_angle_idx: return RigidFit(0.0, (0.0, 0.0), 1.0, "reference") candidates = _source_images(sras, angle_idx, ref_angle_idx, signal_mv, dc_threshold_mv, sources) if not candidates: return RigidFit(0.0, (0.0, 0.0), -1.0, "none") nominal = (nominal_delta_deg(sras, angle_idx, ref_angle_idx) if seed_deg is None else float(seed_deg)) thetas = _rotation_candidates(nominal, search_deg, coarse_step_deg, seed_signs) # ---- Stage 1: coarse sweep, every source ------------------------------ pitch_c, n_c = _reg_pitch_and_size(sras, coarse_dim) best = (-2.0, 0.0, (0, 0), "none") # score, theta, (dr, dc), source for name, ref_raw, mov_raw in candidates: ref_reg = _prepare_reg_image(sras, ref_angle_idx, ref_raw, pitch_c) mov_reg = _prepare_reg_image(sras, angle_idx, mov_raw, pitch_c) ref_img, ref_valid = _embed_with_valid(ref_reg, pitch_c, n_c, 0.0) for theta in thetas: score, shift = _score_rotation(ref_img, ref_valid, mov_reg, pitch_c, n_c, theta) if score > best[0]: best = (score, theta, shift, name) if best[3] == "none": return RigidFit(0.0, (0.0, 0.0), -1.0, "none") # ---- Stage 2: refine the winner at full registration resolution ------- name = best[3] ref_raw, mov_raw = next((r, m) for n_, r, m in candidates if n_ == name) pitch_f, n_f = _reg_pitch_and_size(sras, fine_dim) ref_reg = _prepare_reg_image(sras, ref_angle_idx, ref_raw, pitch_f) mov_reg = _prepare_reg_image(sras, angle_idx, mov_raw, pitch_f) ref_img, ref_valid = _embed_with_valid(ref_reg, pitch_f, n_f, 0.0) theta = best[1] score, shift = _score_rotation(ref_img, ref_valid, mov_reg, pitch_f, n_f, theta, subpixel=True) if refine: theta, score, shift = _refine_rotation( ref_img, ref_valid, mov_reg, pitch_f, n_f, theta, score, shift, step_deg=coarse_step_deg) dr, dc = shift return RigidFit(float(theta), (float(dc * pitch_f), float(dr * pitch_f)), float(score), name) # ---- Shared canvas: angle 0's own pixel grid, extended -------------------- def canvas_for_params(sras: SrasFile, ref_angle_idx: int, pitch_mm: tuple[float, float], per_angle_params: dict[int, ManualAngleParams], *, margin_frac: float = 0.0, snap: bool = True ) -> tuple[tuple[float, float], tuple[int, int]]: """Shared-canvas origin (stage mm) and (n_rows, n_cols) at *pitch_mm* that contains every angle's footprint after its own rigid transform. Angles missing from per_angle_params default to identity (e.g. a sidecar saved before a rescan added more angles). snap=True aligns the canvas grid with the reference angle's own pixel grid, so the reference lands on integer canvas pixels and is resampled by an exact integer translation — the concrete meaning of "the canvas carries angle 0's X/Y coordinates". It requires pitch_mm to be the reference's own pitch; the manual-alignment preview passes a coarser pitch and snap=False. margin_frac pads the box on every side: 0 for a final canvas, nonzero for ManualAlignmentDialog's preview canvas, which needs headroom so an ordinary translation nudge never has to trigger a full canvas resize (an extreme nudge can still push content past this padding; accepted, and cheap to recover from by re-opening the dialog). """ dx, dy = pitch_mm corners = np.vstack([ _footprint_corners_ref_mm( sras, a, per_angle_params.get(a, ManualAngleParams()).rotation_deg, per_angle_params.get(a, ManualAngleParams()).shift_mm) for a in range(sras.n_angles)]) 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 center = ref_center_mm(sras, ref_angle_idx) if snap: # Express the box in the reference's own pixel indices and grow it # outward to whole pixels, so canvas index k lands exactly where the # reference's own pixel (k + const) does. cy, cx = _center_idx(sras, ref_angle_idx) cols = sorted((x_min / dx + cx, x_max / dx + cx)) rows = sorted((y_min / dy + cy, y_max / dy + cy)) col0, col1 = int(np.floor(cols[0])), int(np.ceil(cols[1])) row0, row1 = int(np.floor(rows[0])), int(np.ceil(rows[1])) origin_ref = np.array([(col0 - cx) * dx, (row0 - cy) * dy]) shape = (row1 - row0 + 1, col1 - col0 + 1) else: origin_ref = np.array([x_min, y_min if dy > 0 else y_max]) shape = (int(np.ceil((y_max - y_min) / abs(dy))) + 1, int(np.ceil((x_max - x_min) / dx)) + 1) origin_stage = origin_ref + center return (float(origin_stage[0]), float(origin_stage[1])), shape def build_canvas_affine(sras: SrasFile, angle_idx: int, ref_angle_idx: int, rotation_deg: float, shift_mm: tuple[float, float], pitch_mm: tuple[float, float], canvas_origin_mm: tuple[float, float], *, src_downsample: tuple[int, int] = (1, 1) ) -> tuple[np.ndarray, np.ndarray]: """canvas index -> this angle's raw index, for a canvas whose origin is given in *stage* mm (the reference's frame). *src_downsample* is the (rows, cols) block-mean factor already applied to the image the caller will resample — 1:1 for the raw image, coarser for the manual-alignment preview's downsampled masks.""" dx_a, dy_a = pixel_pitch_mm(sras, angle_idx) n_rows, n_frames = sras.image_shape(angle_idx) fy, fx = (max(1, int(v)) for v in src_downsample) if (fy, fx) != (1, 1): # Same block-mean bookkeeping _prepare_reg_image does: the downsampled # array's own center can sit up to half a block off the full-resolution # center, and that offset has to be undone here or every preview layer # lands slightly (and inconsistently) off. nr_s, nf_s = n_rows // fy, n_frames // fx src_dx, src_dy = dx_a * fx, dy_a * fy src_center = ((nr_s - 1) / 2.0, (nf_s - 1) / 2.0) src_off = (_block_center_off(n_frames, nf_s, fx, dx_a), _block_center_off(n_rows, nr_s, fy, dy_a)) else: src_dx, src_dy = dx_a, dy_a src_center = _center_idx(sras, angle_idx) src_off = (0.0, 0.0) origin_ref = np.asarray(canvas_origin_mm, dtype=np.float64) \ - ref_center_mm(sras, ref_angle_idx) return _affine_out_to_src( out_pitch_mm=pitch_mm, out_origin_ref_mm=origin_ref, src_dx_mm=src_dx, src_dy_mm=src_dy, src_center_idx=src_center, src_center_off_mm=src_off, rotation_deg=rotation_deg, shift_mm=shift_mm) def _result_from_params(sras: SrasFile, ref_angle_idx: int, dc_threshold_mv: float, params: dict[int, ManualAngleParams], extra: dict[int, tuple[float, str]] | None = None ) -> AlignmentResult: """Assemble the final AlignmentResult from per-angle rigid parameters: pick the shared canvas, then build each angle's canvas->raw affine. Pure matrix and bbox math, so it is cheap enough to call synchronously on the GUI thread on every manual edit.""" pitch = pixel_pitch_mm(sras, ref_angle_idx) canvas_origin_mm, canvas_shape = canvas_for_params( sras, ref_angle_idx, pitch, params, snap=True) extra = extra or {} per_angle: dict[int, AngleTransform] = {} for a in range(sras.n_angles): p = params.get(a, ManualAngleParams()) matrix, offset = build_canvas_affine( sras, a, ref_angle_idx, p.rotation_deg, p.shift_mm, pitch, canvas_origin_mm) score, source = extra.get(a, (1.0, "")) per_angle[a] = AngleTransform(p.rotation_deg, p.shift_mm, matrix, offset, score, source) return AlignmentResult(ref_angle_idx, dc_threshold_mv, canvas_shape, pitch[0], pitch[1], canvas_origin_mm, per_angle) def crop_alignment_result(result: AlignmentResult, row0: int, col0: int, n_rows: int, n_cols: int) -> AlignmentResult: """The same alignment restricted to a rectangular window of its canvas — canvas pixel (row0, col0) becomes the cropped canvas's (0, 0). Cropping folds into each angle's existing affine instead of becoming a second transform, because a canvas crop is *pure index translation*. From _affine_out_to_src, matrix = D @ Rinv @ A_out depends only on pitch and rotation, and A_out @ [row0, col0] is exactly the mm displacement of the new origin, so: matrix @ [r', c'] + (offset + matrix @ [row0, col0]) == matrix @ [r' + row0, c' + col0] + offset i.e. shifting the offset by matrix @ [row0, col0] reproduces the original mapping at the shifted indices, exactly. That equality is what lets apply_alignment, reproject_mask and the aligned .sras exporter all keep working on a cropped result with no special-casing — and it is why the crop preview a user approves is guaranteed to be the same pixels the exporter writes. Bounds are the caller's responsibility (the wizard's ROI page clamps to the canvas): an out-of-range window is geometrically well-defined here and simply resamples padding. """ if n_rows <= 0 or n_cols <= 0: raise ValueError(f"empty crop: {n_rows} x {n_cols}") delta = np.array([float(row0), float(col0)]) per_angle = { a: AngleTransform(t.rotation_deg, t.shift_mm, t.matrix, t.offset + t.matrix @ delta, t.score, t.source) for a, t in result.per_angle.items() } origin = (result.canvas_origin_mm[0] + col0 * result.canvas_dx_mm, result.canvas_origin_mm[1] + row0 * result.canvas_dy_mm) return AlignmentResult(result.ref_angle_idx, result.dc_threshold_mv, (int(n_rows), int(n_cols)), result.canvas_dx_mm, result.canvas_dy_mm, origin, per_angle) def overlap_stats(counts: np.ndarray, n_angles: int) -> dict: """Summarize a per-pixel "how many angles cover this pixel" image — the number the alignment wizard's mask-stack view is colored by. Separated from the drawing code because it is the actual judgement the user makes on that screen ("do the angles land on top of each other?"), and a plain array-in/dict-out function can be tested without Qt. """ counts = np.asarray(counts) union = int(np.count_nonzero(counts)) full = int(np.count_nonzero(counts >= n_angles)) return { "union_px": union, "full_px": full, "full_frac": (full / union) if union else 0.0, "mean_count": float(counts[counts > 0].mean()) if union else 0.0, "max_count": int(counts.max()) if counts.size else 0, "empty": union == 0, } def largest_rect_at_least(counts: np.ndarray, min_count: int ) -> tuple[int, int, int, int] | None: """Largest axis-aligned rectangle whose every pixel has counts >= min_count, as (row0, col0, n_rows, n_cols), or None if no pixel qualifies. Backs the wizard's "fit to full overlap" button. A *bounding box* of the qualifying pixels would be the obvious thing and is wrong: the full-overlap region of several rotated scans is roughly a disc, whose bounding box has corners no angle covers at all. Offering that as the crop would hand the user padding they explicitly asked to avoid, so this finds a rectangle that is entirely inside the region. Standard largest-rectangle-in-a-histogram sweep — O(rows * cols) on a preview-sized array, so it is instant at interactive rates. """ good = np.asarray(counts) >= min_count if not good.any(): return None n_rows, n_cols = good.shape best = (0, 0, 0, 0) # area, row0, col0, ... best_area = 0 heights = np.zeros(n_cols, dtype=np.int64) for r in range(n_rows): heights = np.where(good[r], heights + 1, 0) # Sentinel column of height 0 flushes the stack at the end of the row. stack: list[tuple[int, int]] = [] # (start col, height) for c in range(n_cols + 1): h = int(heights[c]) if c < n_cols else 0 start = c while stack and stack[-1][1] >= h: s, sh = stack.pop() area = sh * (c - s) if area > best_area: best_area = area best = (s, sh, c - s, r) start = s if h: stack.append((start, h)) col0, height, width, row_end = best return (row_end - height + 1, col0, height, width) 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, fine_dim: int = _DEFAULT_FINE_DIM) -> AlignmentResult: """Top-level alignment driver: register every angle onto *ref_angle_idx* by content, then lay them all out on that angle's own coordinate grid. 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 # 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 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))) # ---- Step 2: rigid registration of every angle against the reference -- # Its own worker count: registration is bounded by grid-sized FFT buffers, # not by the waveform chunking plan_angle_level budgets for. _ticks.clear() def fit_for(a: int) -> RigidFit: fit = register_angle_to_reference( sras, a, ref_angle_idx, dc4_mv, dc_threshold_mv=dc_threshold_mv, fine_dim=fine_dim) tick(25, 65) return fit fits = dict(enumerate(_parallel_map( fit_for, range(n), _registration_workers(sras, fine_dim)))) # ---- Step 3/4: shared canvas on the reference's grid, per-angle affines params = {a: ManualAngleParams(f.rotation_deg, f.shift_mm) for a, f in fits.items()} extra = {a: (f.score, f.source) for a, f in fits.items()} result = _result_from_params(sras, ref_angle_idx, dc_threshold_mv, params, extra) if progress_cb: progress_cb(100) return result # --------------------------------------------------------------------------- # Manual alignment (Fusion menu -> Manual Alignment... dialog) # # Skips register_angle_to_reference's 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 registration # run or a saved sidecar). Building the final AlignmentResult from already-known # per-angle parameters is pure closed-form matrix math (_result_from_params) # 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 reproject_mask(sras: SrasFile, angle_idx: int, ref_angle_idx: int, mask: np.ndarray, rotation_deg: float, shift_mm: tuple[float, float], canvas_pitch_mm: tuple[float, float], canvas_origin_mm: tuple[float, float], canvas_shape: tuple[int, int], *, src_downsample: tuple[int, int] = (1, 1)) -> 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). *src_downsample* must match the (rows, cols) block-mean factor already applied to *mask*, or the reprojection lands at the wrong scale. order=0 (nearest) matches apply_alignment's own reasoning: a binary mask must never be blended with zero-padding.""" matrix, offset = build_canvas_affine( sras, angle_idx, ref_angle_idx, rotation_deg, shift_mm, canvas_pitch_mm, canvas_origin_mm, src_downsample=src_downsample) 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] ) -> AlignmentResult: """Build a full, full-resolution AlignmentResult from user-supplied per-angle rotation+shift — the Manual Alignment counterpart to compute_angle_alignment, skipping its registration search entirely (every angle's transform here is exactly what the caller supplied). 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. dc_threshold_mv plays no part in the geometry; it's stored on the returned AlignmentResult purely as a record of the RF-mask threshold in effect at the time. """ params = {a: per_angle_params.get(a, ManualAngleParams()) for a in range(sras.n_angles)} params[ref_angle_idx] = ManualAngleParams() return _result_from_params(sras, ref_angle_idx, dc_threshold_mv, params) # ---- Sidecar persistence (.sras.align.json) ------------------------- # A viewer-computed derived artifact, so it lives with the alignment math # rather than in sras_format (see docs/design.md, "Manual-alignment sidecar"). @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 whenever the frame the stored numbers are measured in changes; older # sidecars are treated as absent, never migrated. Bump history: # docs/design.md ("Schema history"). _SIDECAR_SCHEMA_VERSION = 3 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. Angle indices become JSON object keys, so they round-trip as strings — load_manual_alignment converts them back.""" 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