#!/usr/bin/env python3 """Image computation and angle alignment for .sras scans. Depends only on numpy/scipy (+ optional pyfftw) and sras_format, so a multiprocessing child can import it without loading Qt or matplotlib — which matters because Python 3.14 on macOS spawns rather than forks. """ import json import os from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from pathlib import Path import numpy as np import scipy.fft as scipy_fft import scipy.ndimage as scipy_ndimage from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, SrasFile, adc_to_mv # --------------------------------------------------------------------------- # FFT backend # --------------------------------------------------------------------------- try: import pyfftw pyfftw.interfaces.cache.enable() PYFFTW_AVAILABLE = True except ImportError: PYFFTW_AVAILABLE = False _fft_backend = "numpy" # "numpy" or "pyfftw"; set via set_fft_backend() def set_fft_backend(name: str): """Select the rfft implementation. Module-level state, so it must be set explicitly inside each multiprocessing child — it does not survive a spawn.""" global _fft_backend _fft_backend = name if (name != "pyfftw" or PYFFTW_AVAILABLE) else "numpy" def get_fft_backend() -> str: return _fft_backend def _do_rfft(x: np.ndarray, n: int | None = None, axis: int = -1, workers: int = 1) -> np.ndarray: """Dispatch rfft to the selected backend with optional multithreading.""" if _fft_backend == "pyfftw" and PYFFTW_AVAILABLE: return pyfftw.interfaces.numpy_fft.rfft(x, n=n, axis=axis, threads=workers) return scipy_fft.rfft(x, n=n, axis=axis, workers=workers) # --------------------------------------------------------------------------- # Chunking / parallel budget # --------------------------------------------------------------------------- # # Rows are batched so the float32 working buffers for one chunk stay under a # memory budget. A fixed row count (the original design) works fine for small # legacy scans but is catastrophic for a v6 scan with a large per-angle # frame/sample count — e.g. a 7500-frame x 2500-sample angle needs ~2.4 GB for # a single 32-row chunk. # # With chunks running concurrently the budget has to cover *all* live chunks at # once. Note that on a large scan chunk_rows is already clamped to its floor of # 1 row (one row alone is ~75 MB of float32 at 7507x2500), so shrinking the # per-chunk size cannot buy more concurrency — the worker count must be derived # from the budget instead. See _plan_chunks. # 1024 MB is the measured knee on a 16-core machine against a 7507-frame x # 2500-sample angle: 512 MB leaves ~20% of the FFT speedup on the table, and # 1536+ MB costs ~0.4 GB more resident for no further gain. Override with # SRAS_MEM_BUDGET_MB on a smaller machine. _TOTAL_BYTES_BUDGET = int(os.environ.get("SRAS_MEM_BUDGET_MB", 1024)) * 1024 * 1024 _CHUNK_ROWS_MAX = 32 # cap for small scans (original behavior) _MAX_WORKERS = int(os.environ.get("SRAS_MAX_WORKERS", 0)) or (os.cpu_count() or 4) # An rfft chunk holds, live at once: the float32 input, the complex64 # transform, and the float32 power spectrum — roughly 3x the input buffer. _FFT_LIVE_MULTIPLIER = 3 def _chunk_rows_for(n_frames: int, samples_per_frame: int, budget: int = _TOTAL_BYTES_BUDGET) -> int: bytes_per_row = max(1, n_frames * samples_per_frame * 4) # float32 return int(max(1, min(_CHUNK_ROWS_MAX, budget // bytes_per_row))) def _plan_chunks(n_rows: int, n_frames: int, samples_per_frame: int, live_multiplier: int = 1, max_workers: int | None = None, budget: int | None = None) -> tuple[int, int]: """(chunk_rows, n_workers) such that all concurrent chunks together fit the budget: n_workers x live_multiplier x chunk_rows x n_frames x samples x 4 bytes <= budget. The worker count is chosen *first* and the chunk sized to it. Sizing the chunk first is the trap: _chunk_rows_for spends whatever budget it is given, so a single chunk would always consume the lot and leave room for exactly one worker — no concurrency, on precisely the large scans that need it most. A caller that is itself running several of these concurrently must pass *both* max_workers=1 and its share of the budget. Capping the workers alone is not enough: the chunk would still be sized against the whole budget, and N concurrent callers would each allocate all of it. """ total = _TOTAL_BYTES_BUDGET if budget is None else max(1, budget) cap = max_workers if max_workers is not None else _MAX_WORKERS bytes_per_row = max(1, live_multiplier * n_frames * samples_per_frame * 4) # A chunk is at least one row, so that alone caps how many can be live. n_workers = int(max(1, min(cap, n_rows, total // bytes_per_row))) chunk_rows = _chunk_rows_for(n_frames, samples_per_frame, max(1, total // (n_workers * live_multiplier))) # With chunk_rows known, more workers than chunks buys nothing. n_workers = int(max(1, min(n_workers, -(-n_rows // chunk_rows)))) return chunk_rows, n_workers def _map_row_chunks(n_rows: int, chunk_rows: int, n_workers: int, fn, should_stop=None): """Apply fn(r0, r1) over row chunks, in parallel when it pays. Chunks write to disjoint output slices, so no locking is needed. numpy ufuncs, scipy's pocketfft, and memmap page faults all release the GIL, so threads give real parallelism here — and on a very large file they also keep many more page-fault requests in flight, which is what the I/O path wants. *should_stop* is polled per chunk so a cancelled job abandons work within one chunk rather than one whole angle — on a large scan an angle is ~40 s, which is far too long to block application shutdown. """ bounds = [(r0, min(r0 + chunk_rows, n_rows)) for r0 in range(0, n_rows, chunk_rows)] def run(b): if should_stop is not None and should_stop(): return fn(*b) if n_workers <= 1 or len(bounds) == 1: for b in bounds: run(b) return with ThreadPoolExecutor(max_workers=min(n_workers, len(bounds))) as pool: list(pool.map(run, bounds)) # --------------------------------------------------------------------------- # Image computation # --------------------------------------------------------------------------- def plan_angle_level(sras: SrasFile, live_multiplier: int = 1) -> tuple[int, int]: """(n_angle_workers, per_angle_budget) for a caller that parallelises over angles instead of over rows. Callers that parallelise over angles must not also let each angle parallelise over rows: the two levels multiply, in threads and in memory. Each angle is therefore given max_workers=1 and the returned budget share, which together keep total live buffers within _TOTAL_BYTES_BUDGET. """ biggest = max(range(sras.n_angles), key=lambda a: int(sras.n_frames[a])) _, n_workers = _plan_chunks(int(sras.n_rows[biggest]), int(sras.n_frames[biggest]), sras.samples_per_frame, live_multiplier=live_multiplier) n_workers = max(1, min(n_workers, sras.n_angles)) return n_workers, max(1, _TOTAL_BYTES_BUDGET // n_workers) def compute_dc_image(sras: SrasFile, angle_idx: int, ch_idx: int, max_workers: int | None = None, budget: int | None = None, should_stop=None) -> np.ndarray: """Mean of each waveform → (n_rows, n_frames) float32, in ADC counts. Pass max_workers=1 when the caller is already parallelising over angles. If *should_stop* ever returns True the result is incomplete — the caller is expected to be abandoning it. """ n_rows, n_frames = sras.image_shape(angle_idx) data = sras.data[angle_idx] chunk_rows, n_workers = _plan_chunks(n_rows, n_frames, sras.samples_per_frame, max_workers=max_workers, budget=budget) img = np.empty((n_rows, n_frames), dtype=np.float32) def chunk(r0: int, r1: int): img[r0:r1] = data[r0:r1, ch_idx, :, :].astype(np.float32).mean(axis=-1) _map_row_chunks(n_rows, chunk_rows, n_workers, chunk, should_stop=should_stop) return img def dc_image_mv(sras: SrasFile, angle_idx: int, ch_idx: int, max_workers: int | None = None, budget: int | None = None, should_stop=None) -> np.ndarray: """DC image for (angle, channel) in mV, preferring a stored v5/v7 cache.""" cached = sras.cached_dc_mv(angle_idx, ch_idx) if cached is not None: return cached return adc_to_mv( compute_dc_image(sras, angle_idx, ch_idx, max_workers=max_workers, budget=budget, should_stop=should_stop), *sras.cal(ch_idx)) def compute_rf_image(sras: SrasFile, angle_idx: int, dc_threshold_mv: float | None, apply_bg_sub: bool = True, n_fft: int | None = None, dc4_mv: np.ndarray | None = None, max_workers: int | None = None, budget: int | None = None, should_stop=None) -> np.ndarray: """FFT of each CH1 waveform; pixel = peak frequency in MHz. Pixels where CH4_dc < dc_threshold_mv are set to 0 — and the FFT is never run for them, since that's the expensive part. The masked-out CH1 samples are also never *read*: the boolean mask is applied to the raw memmap slice before any dtype conversion, so numpy only pages in the bytes for pixels that pass the threshold (an untouched memmap page is never read from disk). *dc_threshold_mv* of None means "mask nothing", and skips reading and averaging CH4 entirely — worth a third of the I/O when caching a whole file, where the mask is applied later at display time. If the DC4 image for this angle is already known (e.g. from the DC-channel precompute cache), pass it as *dc4_mv* (mV, shape (n_rows, n_frames)) to reuse it instead of re-reading CH4 here. Fast path: if the file has a precomputed peak-frequency image for this angle (v5 PREC or v7 CACH), zero-padding is off, and the bg-sub flag matches, the stored image is used directly — no FFT is run. """ n_rows, n_frames = sras.image_shape(angle_idx) data = sras.data[angle_idx] # ---- Fast path: precomputed image (v5 PREC or v7 CACH) ---------------- cached_freq = sras.precomputed_freq_mhz[angle_idx] if (cached_freq is not None and n_fft is None # no custom zero-padding and sras.precomputed_bg_sub == (apply_bg_sub and sras.background is not None)): freq_img = cached_freq.copy() if dc_threshold_mv is not None: # DC4 mask, in priority order: already-cached DC block, caller- # supplied image, or a fresh (cheap — no FFT) recompute. dc4_img = sras.cached_dc_mv(angle_idx, CH4_IDX) if dc4_img is None: dc4_img = dc4_mv if dc4_mv is not None else adc_to_mv( compute_dc_image(sras, angle_idx, CH4_IDX), *sras.cal(CH4_IDX)) freq_img[dc4_img < dc_threshold_mv] = 0.0 return freq_img # ---- Chunked FFT path -------------------------------------------------- freq_axis = sras.freq_axis_mhz(n_fft) img = np.zeros((n_rows, n_frames), dtype=np.float32) n_fft_bins = n_fft if n_fft is not None else sras.samples_per_frame chunk_rows, n_workers = _plan_chunks( n_rows, n_frames, max(sras.samples_per_frame, n_fft_bins), live_multiplier=_FFT_LIVE_MULTIPLIER, max_workers=max_workers, budget=budget) background = sras.background if (apply_bg_sub and sras.background is not None) else None cal4 = sras.cal(CH4_IDX) def chunk(r0: int, r1: int): if dc_threshold_mv is None: valid = None else: if dc4_mv is not None: dc4_chunk = dc4_mv[r0:r1] else: dc4_chunk = adc_to_mv( data[r0:r1, CH4_IDX, :, :].astype(np.float32).mean(axis=-1), *cal4) valid = dc4_chunk >= dc_threshold_mv # True = run the FFT if not valid.any(): return # Index the raw memmap slice with the boolean mask *before* # converting dtype — this is a lazy view until touched, so only the # selected elements are actually read from disk; masked-out pixels' # pages are never paged in at all. raw = data[r0:r1, CH1_IDX, :, :] waves = (raw[valid] if valid is not None else raw).astype(np.float32) if background is not None: waves -= background # background is 1-D (spf,) # Parallelism comes from the outer chunk loop, so keep the inner # transform single-threaded to avoid oversubscribing the machine. spectrum = _do_rfft(waves, n=n_fft, axis=-1, workers=1) del waves # |z|^2 without np.abs()'s extra full-size temporary. power = spectrum.real ** 2 power += spectrum.imag ** 2 del spectrum # [..., 0] not [:, 0]: the unmasked path keeps the (rows, frames, # bins) shape, where [:, 0] would blank a whole frame. power[..., 0] = 0.0 # suppress DC bin peak_bins = np.argmax(power, axis=-1) del power if valid is not None: img[r0:r1][valid] = freq_axis[peak_bins] else: img[r0:r1] = freq_axis[peak_bins] _map_row_chunks(n_rows, chunk_rows, n_workers, chunk, should_stop=should_stop) return img # --------------------------------------------------------------------------- # Batch cache (Convert menu) — process-pool entry point # --------------------------------------------------------------------------- def cache_file(path: str, mode: str, apply_bg_sub: bool, fft_backend: str = "numpy", max_workers: int = 0) -> str: """Compute and store DC or FFT images for every angle of one file, converting v6 → v7 in place. Returns "" on success or an error message. Module-level and picklable so it can run in a ProcessPoolExecutor. The FFT backend and worker cap are passed explicitly because module globals do not survive a spawn. """ global _MAX_WORKERS try: if mode not in ("dc", "fft"): return f"unknown cache mode {mode!r} (expected 'dc' or 'fft')" set_fft_backend(fft_backend) if max_workers: _MAX_WORKERS = max_workers sras = SrasFile(path) if sras.version not in (6, 7): return (f"unsupported version {sras.version} — only v6/v7 files " "can be batch-cached") n = sras.n_angles if mode == "dc": dc3 = [adc_to_mv(compute_dc_image(sras, a, CH3_IDX), *sras.cal(CH3_IDX)) for a in range(n)] dc4 = [adc_to_mv(compute_dc_image(sras, a, CH4_IDX), *sras.cal(CH4_IDX)) for a in range(n)] sras.write_v7_cache(new_dc3_mv=dc3, new_dc4_mv=dc4) else: effective_bg = apply_bg_sub and sras.background is not None # dc_threshold_mv=None: store unmasked images and mask at display # time (same convention as v5's PREC block). Skipping the mask # also skips reading CH4 entirely. freq = [compute_rf_image(sras, a, dc_threshold_mv=None, apply_bg_sub=effective_bg) for a in range(n)] sras.write_v7_cache(new_freq_mhz=freq, new_bg_sub=effective_bg) return "" except Exception as exc: return str(exc) # --------------------------------------------------------------------------- # Angle alignment (Fusion menu) # # Puts every angle's images onto one shared, zero-padded pixel grid using a # rigid transform only — rotation + translation, never scale. # # Angle 0 (the reference) is the sole coordinate authority: it is the only # angle whose stage XY (x_start_mm / y_positions_mm) is ever read, and the # shared canvas is literally an extension of angle 0's own pixel grid, so the # aligned view carries angle 0's real X/Y axes. Every *other* angle is placed # purely by content — its rotation and translation come from cross-correlating # its CH4 image against angle 0's (register_angle_to_reference) — and its own # stage XY is deliberately never consulted. That is not an oversight: the # rotation stage moves the sample relative to the scan window, so where a # window sat in stage coordinates says nothing about where the sample is, and # an earlier design that pivoted each angle on a signal-weighted centroid of # its own window put every angle on a ~20 mm circle around the optical center # instead of stacking them into one shape. # # Only two coordinate frames exist here: # # local mm — one angle's own physical frame: origin at the *center of its own # pixel array*, x along +column, y along +row, scaled by that angle's own # pitches. Carries no stage position whatsoever. # # ref mm — the reference angle's local mm. A registration result # (rotation_deg, shift_mm) is exactly the rigid map from an angle's local # mm to ref mm: q = R(rotation_deg) @ l + shift_mm. Stage coordinates # re-enter once, at the very end, when the canvas origin is converted to # angle 0's stage mm (AlignmentResult.canvas_origin_mm). # # Rotation is done in mm, never on raw pixel indices: the x pitch # (SrasFile.pixel_x_mm, 5 µm on a real scan) and the y/row pitch (50 µm) differ # by 10x, so rotating the raw index grid would shear the image — an unwanted # anisotropic scale. Registration runs on a resampled *isotropic* grid for the # same reason, and every affine here maps shared-grid index -> mm -> undo # rotation/shift -> that angle's own local mm -> that angle's own raw index, # matching the output->input convention scipy.ndimage.affine_transform wants. # --------------------------------------------------------------------------- @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 _rotation_candidates(nominal_deg: float, search_deg: float, step_deg: float) -> list[float]: """Coarse rotation candidates: a window around *both* signs of the stage's reported angle change. Scoring both 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.""" out: list[float] = [] for center in (-nominal_deg, 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) -> 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 both 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. """ 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) thetas = _rotation_candidates(nominal, search_deg, coarse_step_deg) # ---- 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) 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 _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) ------------------------- # # Lives here, not sras_format.py: sras_format.py is scoped to the versioned # binary .sras spec itself (see scan_format.md); a manual alignment is a # viewer-computed *derived* artifact, analogous in kind to AlignmentResult — # so it belongs with the alignment math it serialises, which already lives # in this module. json + pathlib are both stdlib, so this doesn't add a new # dependency to a module whose only load-bearing constraint is staying free # of Qt/matplotlib for cheap multiprocessing-child imports. @dataclass class ManualAlignmentSidecar: ref_angle_idx: int dc_threshold_mv: float per_angle: dict[int, ManualAngleParams] def sidecar_path(sras_path) -> Path: """.sras.align.json next to the scan file. A thin, independently testable function since save/load/delete and the GUI's status messages all need the identical path.""" p = Path(sras_path) return p.with_name(p.name + ".align.json") # The stored rotation_deg/shift_mm are meaningless without the frame they were # measured in, so this is bumped whenever that frame changes. Each bump makes # older files describe a different (and, for the bugs each bump fixed, actively # wrong) transform than the same numbers would today, and loading one unchanged # would silently reproduce the very "scans show up everywhere" symptom the bump # fixed — so older sidecars are treated as absent rather than migrated. # 1 -> 2 pivot moved from the scan-window bbox center to a content-derived # centroid, and the rotation sign convention was corrected. # 2 -> 3 the content centroid was abandoned entirely: rotation is now about # each angle's own array center, mapped onto the reference's array # center, with shift_mm in the reference's local mm frame. No angle # but the reference contributes stage coordinates any more. _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. Schema (schema_version 3): { "schema_version": 3, "ref_angle_idx": , "dc_threshold_mv": , "per_angle": { "": {"rotation_deg": , "shift_mm": [, ]}, ... } } Angle indices are JSON object keys, so they round-trip as strings — load_manual_alignment converts them back to int. """ path = sidecar_path(sras.path) payload = { "schema_version": _SIDECAR_SCHEMA_VERSION, "ref_angle_idx": ref_angle_idx, "dc_threshold_mv": dc_threshold_mv, "per_angle": { str(a): {"rotation_deg": p.rotation_deg, "shift_mm": list(p.shift_mm)} for a, p in per_angle.items() }, } path.write_text(json.dumps(payload, indent=2)) return path def load_manual_alignment(sras: SrasFile) -> ManualAlignmentSidecar | None: """Read 's sidecar JSON if present, else None — never raises: a missing, corrupt, foreign, or version-mismatched JSON file must not block opening the .sras file itself (see SrasViewerWindow._on_load_done). Per-angle entries for an angle index no longer present in *sras* (e.g. re-scanned with fewer angles) are silently dropped. """ path = sidecar_path(sras.path) if not path.exists(): return None try: raw = json.loads(path.read_text()) if raw.get("schema_version") != _SIDECAR_SCHEMA_VERSION: return None per_angle = { int(a): ManualAngleParams( float(v["rotation_deg"]), (float(v["shift_mm"][0]), float(v["shift_mm"][1]))) for a, v in raw.get("per_angle", {}).items() if int(a) < sras.n_angles } return ManualAlignmentSidecar( ref_angle_idx=int(raw.get("ref_angle_idx", 0)), dc_threshold_mv=float(raw.get("dc_threshold_mv", 0.0)), per_angle=per_angle) except (OSError, ValueError, KeyError, TypeError, IndexError, json.JSONDecodeError): return None def delete_manual_alignment(sras: SrasFile) -> bool: """Delete the sidecar if present. Returns whether a file actually existed to delete, so Clear Alignment's status message can say so. Genuine I/O errors (permission denied, read-only share) propagate — the caller (ManualAlignmentDialog._on_clear) surfaces them rather than silently pretending the destructive action succeeded.""" path = sidecar_path(sras.path) try: path.unlink() return True except FileNotFoundError: return False