#!/usr/bin/env python3 """Image computation and angle alignment for .sras scans. Depends only on numpy/scipy/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 pyfftw import scipy.fft as scipy_fft import scipy.ndimage as scipy_ndimage from sras_format import (CH1_IDX, CH3_IDX, CH4_IDX, MAX_PAD_FACTOR, SrasFile, adc_to_mv) # --------------------------------------------------------------------------- # FFT worker pool and per-thread pyFFTW plans # --------------------------------------------------------------------------- _FFT_BLOCK_MAX = 512 # ceiling on waveforms/task: the knee measured on a # 16-core machine at natural resolution (n_len == spf). # Smaller blocks serialise on GIL-held numpy dispatch, # larger ones lose cache residency and task granularity. _FFT_BLOCK_MIN = 32 # floor, so task granularity never collapses at # extreme pad factors — at the cost of exceeding # _FFT_PLAN_BYTES_BUDGET there (see _fft_block_for). _FFT_PLAN_BYTES_BUDGET = int( os.environ.get("SRAS_FFT_PLAN_BUDGET_MB", 16)) * 1024 * 1024 # Per-thread ceiling on one cached pyFFTW plan's resident input+output # buffers (see _fft_block_for). This is PERMANENT memory — the plan # cache is never evicted — so the worst case across the whole pool for # one distinct (spf, n_len) is _MAX_WORKERS * _FFT_PLAN_BYTES_BUDGET. # Deliberately separate from _TOTAL_BYTES_BUDGET/SRAS_MEM_BUDGET_MB, # which bounds transient concurrently-live chunk buffers, freed at the # end of each chunk — this one bounds a permanent per-thread tax that # compounds with worker count, not chunk concurrency. _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 _fft_block_for(spf: int, n_len: int) -> int: """Waveforms per FFT task/plan at transform length n_len. Bounds a cached pyFFTW plan's resident input+output buffers (permanent, one per pool thread, never evicted) to _FFT_PLAN_BYTES_BUDGET regardless of pad factor. Reproduces _FFT_BLOCK_MAX exactly at natural resolution (n_len == spf) — see docs/design.md for the worked numbers at high pad.""" bytes_per_wf = 4 * spf + 8 * (n_len // 2 + 1) block = _FFT_PLAN_BYTES_BUDGET // max(1, bytes_per_wf) return int(min(_FFT_BLOCK_MAX, max(_FFT_BLOCK_MIN, block))) def _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 (block, spf) input shape, block = _fft_block_for(spf, n), so each worker thread plans once per (block size, transform length) pair; a remainder block runs through the same plan with its tail rows ignored. Single-threaded (threads=1): outer parallelism comes from the pool, so the transform itself must not spin up threads. The returned array is the plan's output buffer — consume it before the next call on the same thread. """ n_wf, spf = waves.shape block = _fft_block_for(spf, n) plans = getattr(_fftw_local, "plans", None) if plans is None: plans = _fftw_local.plans = {} key = (block, spf, n) plan = plans.get(key) if plan is None: _load_wisdom_once() buf = pyfftw.empty_aligned((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] # --------------------------------------------------------------------------- # Peak search: full transform, block-parallel # --------------------------------------------------------------------------- def _peak_bins(waves: np.ndarray, n_len: int, min_bin: int = 1) -> np.ndarray: """Peak bin per waveform via the full transform. *min_bin* zeroes every bin below it (always >= 1, so true DC stays suppressed) before the argmax, so the peak search never returns a bin below the caller's frequency floor. Called in _fft_block_for(spf, n_len)-sized blocks, fanned out across _fft_pool() — see docs/design.md.""" S = _block_rfft(waves, n_len) power = S.real ** 2 power += S.imag ** 2 power[:, :min_bin] = 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 memory_budget_bytes() -> int: """The module-wide ceiling on concurrently-live working buffers, for callers outside this module that read the same data (the aligned exporter sizes its source-band reader against it).""" return _TOTAL_BYTES_BUDGET 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 read buffer scales with the chunk; spectrum memory is bounded independently, by _fft_block_for, to _MAX_WORKERS * _FFT_PLAN_BYTES_BUDGET regardless of chunk size or pad factor (see docs/design.md), so this formula budgets the read buffer alone, with 2x slack, and lets the block fan-out saturate the pool.""" 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 pad_factor_for(sras: SrasFile, n_fft: int | None) -> int: """The integer zero-padding factor an *n_fft* request represents, or 0 if it represents none — i.e. it is not a whole multiple of this file's samples_per_frame, so no stored image (which only ever records an integer factor) can answer it. 0 rather than None so callers can compare it straight against ``sras.precomputed_pad_factor``, which is never 0. """ if n_fft is None: return 1 spf = sras.samples_per_frame if spf <= 0 or n_fft % spf: return 0 return max(1, n_fft // spf) def cache_mismatch_reasons(sras: SrasFile, *, n_fft: int | None, apply_bg_sub: bool, row_avg_n: int, min_freq_mhz: float = 0.0) -> list[str]: """Why this file's stored FFT images can't answer a request, as human- readable phrases; empty means they can. Padding, background subtraction and row-averaging are all baked into the stored numbers, so a request that differs in any of them has to go back through a real FFT. This is the single accept rule: cached_rf_image serves a stored image iff this returns nothing, and callers that want to *explain* the miss (rather than silently recompute) format these same strings, so the two can't drift apart. The min peak frequency floor is the asymmetric case: a stored floor *above* the request is a genuine mismatch (bins below the stored floor were never searched, so the stored numbers can't answer a request that wants them considered), but a stored floor at or *below* the request is servable — the difference is re-applied as a display-time mask by cached_rf_image (pixels whose stored peak falls below the requested floor are marked invalid rather than re-resolved). Compared in whole kHz, the header field's own fixed-point grid, so a value that round-trips through the file can never miscompare against itself. """ reasons = [] if round(sras.precomputed_min_freq_mhz * 1000) > round(min_freq_mhz * 1000): reasons.append( f"stored with a {sras.precomputed_min_freq_mhz:g} MHz " f"min-peak-freq floor, requested {min_freq_mhz:g} MHz " f"(bins below the stored floor were never searched)") pad = pad_factor_for(sras, n_fft) if pad != sras.precomputed_pad_factor: want = f"pad {pad}x" if pad else "a ragged n_fft" reasons.append(f"stored at pad {sras.precomputed_pad_factor}x, " f"requested at {want}") if sras.precomputed_bg_sub != (apply_bg_sub and sras.background is not None): reasons.append("stored with background subtraction " f"{'on' if sras.precomputed_bg_sub else 'off'}") if sras.precomputed_row_avg_n != row_avg_n: have = (f"row-averaged (n={sras.precomputed_row_avg_n})" if sras.precomputed_row_avg_n else "raw per-pixel") want = (f"row-averaged (n={row_avg_n})" if row_avg_n else "raw per-pixel") reasons.append(f"stored {have}, requested {want}") return reasons 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, min_freq_mhz: float = 0.0) -> 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); the requested zero-padding matches what the store was computed at (see below); 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. Padding is checked the same way, and for the same reason: a padded FFT interpolates between the natural bins, so it resolves genuinely different peak frequencies. *n_fft* of None means natural resolution, i.e. pad 1; anything else must be an exact whole multiple of samples_per_frame equal to sras.precomputed_pad_factor. A ragged n_fft that is not such a multiple can never match a stored image, since the store only ever records an integer pad factor. 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. *min_freq_mhz* above the stored floor is re-imposed here as a mask: pixels whose stored peak is below it are set to 0.0 — the same sentinel the DC mask uses, and unambiguous, since a genuine peak can never be 0.0 (bin 0 is always excluded from the search). Masked-below-floor pixels are marked invalid, not re-resolved; only a real recompute can recover the strongest peak *above* the floor for them. A request whose floor is *below* the stored one can't be answered at all (see cache_mismatch_reasons) and returns None like any other mismatch. """ cached_freq = sras.precomputed_freq_mhz[angle_idx] if cached_freq is None or cache_mismatch_reasons( sras, n_fft=n_fft, apply_bg_sub=apply_bg_sub, row_avg_n=row_avg_n, min_freq_mhz=min_freq_mhz): return None dc4_img = None 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. Resolved # before the copy below so the allow_dc_recompute bail-out doesn't # allocate a full image it is about to throw away. 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 = cached_freq.copy() if dc4_img is not None: freq_img[dc4_img < dc_threshold_mv] = 0.0 if min_freq_mhz > 0.0: # Strict <, matching the compute path: the first bin at or above # the floor survives searchsorted there, so it must survive here. freq_img[freq_img < min_freq_mhz] = 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, row_avg_n: int = 0, min_freq_mhz: float = 0.0, use_stored: bool = True) -> 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. The peak search runs the full transform (_peak_bins) in _fft_block_for(spf, n_len)-sized blocks, fanned out across the pyFFTW-plan-caching thread pool — see docs/design.md for how the block size is bounded so this stays memory-safe at high pad factors. *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. *min_freq_mhz* excludes every bin below it (true DC, bin 0, is always excluded regardless) from the peak search, on top of — not instead of — the dc_threshold_mv mask. Without it, a pixel that passes the DC-bias threshold but carries only weak real signal can still resolve to a near-zero frequency: the un-subtracted background's DC-leakage skirt then has more power than the genuine (but weak) signal peak. Raising the floor above that skirt forces the search to report the strongest peak that is plausibly real signal instead. 0.0 (default) disables it (only true DC is excluded, the long-standing 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. A *min_freq_mhz* at or above the stored floor is part of that service: pixels whose stored peak falls below it come back masked to 0.0 rather than re-resolved. A request *below* the stored floor is a genuine mismatch and falls through to the real FFT here. *use_stored=False* skips the fast path entirely and always runs the real FFT. Batch recompute (cache_file) needs this: served-then-masked pixels written back into the store would permanently replace peaks a real recompute re-resolves above the floor. """ n_rows, n_frames = sras.image_shape(angle_idx) data = sras.data[angle_idx] if use_stored: 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, min_freq_mhz=min_freq_mhz) if fast is not None: return fast # ---- Chunked FFT path -------------------------------------------------- spf = sras.samples_per_frame n_len = n_fft if n_fft is not None else spf block = _fft_block_for(spf, n_len) freq32 = sras.freq_axis_mhz(n_fft).astype(np.float32) # First fine bin at or above the floor; searchsorted is exact here since # freq32 is the very axis the floor is expressed against. min_bin = max(1, int(np.searchsorted(freq32, min_freq_mhz))) 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 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 accounts for it rather than relying on its 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 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 + block, n_wf) out[b0:b1] = freq32[_peak_bins(waves[b0:b1], n_len, min_bin)] if pool is None: for b0 in range(0, n_wf, block): fft_block(b0) else: list(pool.map(fft_block, range(0, n_wf, 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) 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)) return img # --------------------------------------------------------------------------- # Batch cache (Convert menu) — process-pool entry point # --------------------------------------------------------------------------- def cache_file(path: str, mode: str, apply_bg_sub: bool, max_workers: int = 0, pad_factor: int = 1, dc_threshold_mv: float | None = None, row_avg_n: int = 0, min_freq_mhz: float = 0.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 worker cap is passed explicitly because module globals do not survive a spawn. mode is "dc", "fft" (raw per-pixel FFT, 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. *pad_factor* is the zero-padding factor to resolve the peaks at: 1 (the default) is natural resolution, n_fft == samples_per_frame. It is recorded in the SFFT block so a reader knows which views the stored numbers answer for — and it has to be the pad the viewer is *actually* using, since a pad-1 cache is dead weight to a padded view and vice versa (cached_rf_image refuses the mismatch rather than showing peaks resolved at the wrong resolution). *min_freq_mhz* is the peak-search floor (see compute_rf_image), also recorded in the SFFT block; it is quantized up front to the header field's 0.001 MHz grid so the stored images and the recorded floor can never disagree. Both FFT modes force use_stored=False on compute_rf_image: a batch recompute must always run the real FFT, never serve this file's own existing cache back to itself — with a floor (or a changed DC threshold in fft_rowavg mode) that would silently bake a masked copy of the old image into the store in place of a genuine recompute. """ 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')" # write_v7_cache enforces this bound too, but only once every angle's # FFT has already been computed. Checking up front is the difference # between a bad argument costing nothing and costing the whole run. if not (1 <= pad_factor <= MAX_PAD_FACTOR): return f"pad_factor must be 1-{MAX_PAD_FACTOR}, got {pad_factor}" if not (np.isfinite(min_freq_mhz) and min_freq_mhz >= 0): return f"min_freq_mhz must be a finite value >= 0, got {min_freq_mhz}" # Quantize to the SFFT header's fixed-point grid (whole kHz) before # computing, so the floor the FFT actually ran with is exactly the # floor the header records. min_freq_mhz = round(min_freq_mhz * 1000) / 1000.0 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_fft = sras.samples_per_frame * pad_factor if pad_factor > 1 else None 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, n_fft=n_fft, min_freq_mhz=min_freq_mhz, use_stored=False) for a in range(n)] # new_row_avg_n=0 explicitly: these images are raw per-pixel FFTs, # and carrying forward a row_avg_n left by an earlier fft_rowavg # write would label them as something they are not. sras.write_v7_cache(new_freq_mhz=freq, new_bg_sub=effective_bg, new_row_avg_n=0, new_pad_factor=pad_factor, new_min_freq_mhz=min_freq_mhz) 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, n_fft=n_fft, row_avg_n=row_avg_n, min_freq_mhz=min_freq_mhz, use_stored=False) 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, new_pad_factor=pad_factor, new_min_freq_mhz=min_freq_mhz) 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], ) -> 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). The canvas grid is aligned 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". This requires pitch_mm to be the reference's own pitch. """ 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) center = ref_center_mm(sras, ref_angle_idx) # 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) 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) 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 coarse_rect_to_canvas(rect: tuple[int, int, int, int], downsample: tuple[int, int], canvas_shape: tuple[int, int], *, inset_blocks: int = 0 ) -> tuple[int, int, int, int]: """A rectangle in coarse (block-mean preview) indices as canvas pixels, both as (row0, col0, n_rows, n_cols) and clamped to the canvas. The single place the preview grid's relation to the real canvas is written down: the coarse grid samples canvas pixels 0, f, 2f, …, so a coarse pixel stands for a whole block and every caller has to agree on which canvas pixels that block means, or a crop the user approved on the preview lands a few pixels off in the export. *inset_blocks* shrinks the rectangle by that many coarse blocks on each side. A coarse pixel reported as fully covered stands for a block whose far edge may not be, so "fit to full overlap" insets by 1 to stay honestly inside the overlap region; a union rectangle insets by 0 because it wants to contain the region rather than fit inside it. """ row0, col0, nr, nc = rect fy, fx = downsample n_rows, n_cols = canvas_shape r0 = min(n_rows - 1, (row0 + inset_blocks) * fy) c0 = min(n_cols - 1, (col0 + inset_blocks) * fx) r1 = min(n_rows, (row0 + nr - inset_blocks) * fy) c1 = min(n_cols, (col0 + nc - inset_blocks) * fx) return r0, c0, max(1, r1 - r0), max(1, c1 - c0) 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. Meant for a background thread — 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 the wizard's own downsampled mask stack calls that per keystroke # — see AlignmentWizard.rebuild_stack for how it limits a 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 the alignment wizard's live mask stack repeatedly calls (once per angle per edit). *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 (the wizard's Clear path) surfaces them rather than silently pretending the destructive action succeeded.""" path = sidecar_path(sras.path) try: path.unlink() return True except FileNotFoundError: return False