From d5028db445d5ddbb5ade8fcb00ceaa98893372c7 Mon Sep 17 00:00:00 2001 From: Thomas Ales Date: Thu, 6 Aug 2026 09:35:30 -0500 Subject: [PATCH] Add scan-editing CLI and alignment tests; extend Manual Alignment correlation Continues the Manual Alignment work: refines the FFT cross-correlation and mask handling, adds sras_edit_scans.py (drop/renumber bad angle scans), tools/test_alignment.py (registration ground-truth suite), and a rotating test fixture in make_test_sras.py. Co-Authored-By: Claude Fable 5 --- sras_compute.py | 1164 +++++++++++++++++++++------------- sras_edit_scans.py | 295 +++++++++ sras_viewer.py | 204 +++--- sras_viewer_requirements.txt | 8 + sras_workers.py | 61 +- tools/make_test_sras.py | 132 ++++ tools/test_alignment.py | 223 +++++++ tools/test_gui.py | 129 ++-- 8 files changed, 1619 insertions(+), 597 deletions(-) create mode 100644 sras_edit_scans.py create mode 100644 tools/test_alignment.py diff --git a/sras_compute.py b/sras_compute.py index ee34166..28c8b3d 100644 --- a/sras_compute.py +++ b/sras_compute.py @@ -364,28 +364,50 @@ def cache_file(path: str, mode: str, apply_bg_sub: bool, # Angle alignment (Fusion menu) # # Puts every angle's images onto one shared, zero-padded pixel grid using a -# rigid transform only (rotation + translation, never scale). Rotation for -# angle `a` is the *known* scan-angle delta relative to a reference angle — -# never searched. Only the residual translation is found, via FFT phase -# correlation of each angle's binarized CH4 ("dc-mask") image. +# rigid transform only — rotation + translation, never scale. # -# Rotation is done in physical mm space rather than on raw pixel indices: -# the x-pixel pitch (SrasFile.pixel_x_mm) is file-wide constant but the -# y-pixel pitch (row spacing) can differ from it, and for v6 files can even -# vary per angle. Rotating the raw index grid directly would implicitly -# assume square pixels and shear a non-square-pixel image — an unwanted -# effective anisotropic scale. Instead each angle gets one affine that maps -# shared-canvas pixel index -> mm -> undo rotation/shift -> that angle's own -# local mm -> that angle's own raw pixel index, matching the output->input -# convention scipy.ndimage.affine_transform expects. +# 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) found by phase correlation + 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 @@ -403,36 +425,90 @@ class AlignmentResult: class ManualAngleParams: """One angle's manual-alignment state, independent of any canvas. - rotation_deg/shift_mm are exactly AngleTransform's non-derived fields — - the pair a canvas-bound AngleTransform's matrix/offset get built from - once a canvas is decided (build_manual_alignment). Defaults to identity - (no rotation, no shift): a fresh angle with no prior alignment is shown - raw, exactly as scanned — the same "fully unaligned" state Clear - Alignment resets back to. + rotation_deg/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).""" + 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 _bbox_center_mm(sras: SrasFile, angle_idx: int) -> tuple[float, float]: - x = sras.x_axis_mm(angle_idx) - y = sras.y_positions_mm(angle_idx) - return float((x[0] + x[-1]) / 2.0), float((y[0] + y[-1]) / 2.0) +def _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 _bbox_corners_mm(sras: SrasFile, angle_idx: int) -> np.ndarray: - """4 corners (x, y) of this angle's raw mm bounding box, shape (4, 2).""" - x = sras.x_axis_mm(angle_idx) - y = sras.y_positions_mm(angle_idx) - return np.array([[xx, yy] for xx in (x[0], x[-1]) for yy in (y[0], y[-1])]) +def 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: @@ -441,152 +517,63 @@ def _rotation_matrix(theta_deg: float) -> np.ndarray: return np.array([[c, -s], [s, c]]) # CCW rotation acting on (x, y) -def _theta_deg(sras: SrasFile, angle_idx: int, ref_idx: int) -> float: - """CCW rotation, in degrees and in _rotation_matrix's convention, that - maps angle_idx's own local mm frame onto ref_idx's. +def _nominal_delta_deg(sras: SrasFile, angle_idx: int, ref_idx: int) -> float: + """The rotation-stage's own reported angle change between two angles. - This is the *negative* of the raw angles_deg delta: the GR rotation - stage's reported angle increases in the opposite rotational sense from - this module's math-positive (CCW, x toward y) convention in scan mm - space. Rotating by +(angles_deg[a] - angles_deg[ref]) therefore turns - misalignment the wrong way — confirmed empirically (Auto De-rotate made - real scans worse, not better, before this negation). + 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]) + return float(sras.angles_deg[angle_idx] - sras.angles_deg[ref_idx]) -def _signal_centroid_mm(sras: SrasFile, angle_idx: int, - dc4_mv: np.ndarray) -> tuple[float, float]: - """Intensity-weighted centroid (mean x, mean y, weighted by CH4 signal - after subtracting this angle's own minimum) in angle_idx's own local mm - frame — the alignment pivot used in place of the raw scan-window bbox - center (see compute_pivot_points_mm). - - Weighting by the continuous DC signal, rather than a binary >= - dc_threshold_mv mask, means the pivot never depends on how well one - shared threshold happens to suit this particular angle: real signal - levels vary scan to scan, so a threshold tuned for one angle can leave - another angle's binary mask empty — and a centroid of an empty mask has - nothing to fall back to *except* the raw bbox center, silently - reproducing the exact "aligned to the scan window, not the sample" - problem this pivot exists to avoid. Falls back to the bbox center only - in the fully-degenerate case of a perfectly flat signal (nothing to - weight by at all). - """ - weights = dc4_mv - dc4_mv.min() - total = float(weights.sum()) - if total <= 0.0: - return _bbox_center_mm(sras, angle_idx) - x = sras.x_axis_mm(angle_idx) - y = sras.y_positions_mm(angle_idx) - cx = float((x * weights.sum(axis=0)).sum() / total) - cy = float((y * weights.sum(axis=1)).sum() / total) - return cx, cy +def _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 compute_pivot_points_mm(sras: SrasFile, - dc4_mv: dict[int, np.ndarray] | None = None - ) -> dict[int, tuple[float, float]]: - """Per-angle alignment pivot, in each angle's own local mm frame: the - CH4-signal-weighted centroid of its own footprint (see - _signal_centroid_mm), rather than the raw scan-window bbox center. - - Pivoting on each angle's own content — instead of on wherever its - scanned window happened to sit in microscope/global XY space — is what - makes alignment purely relative *between scans* rather than to a global - coordinate system: the rotation stage's true mechanical axis need not - coincide with the scan window's geometric center, and the sample need - not be perfectly centered on that axis either, so a bbox-center pivot - leaves a residual orbital motion between angles that a content-centroid - pivot does not. - - Deliberately independent of dc_threshold_mv (the RF-mask / overlay - threshold): that value is a display/masking choice and must never - silently change where alignment pivots. - - *dc4_mv* lets a caller that has already computed each angle's CH4 mV - image (compute_angle_alignment's Step 1, or ManualAlignmentDialog's own - cache) reuse it instead of recomputing; angles missing from it are - computed fresh via dc_image_mv, which prefers a stored v5/v7 cache over - recomputing from raw waveforms. - """ - dc4_mv = dc4_mv or {} - pivots: dict[int, tuple[float, float]] = {} - for a in range(sras.n_angles): - img = dc4_mv.get(a) - if img is None: - img = dc_image_mv(sras, a, CH4_IDX) - pivots[a] = _signal_centroid_mm(sras, a, img) - return pivots - - -def _corners_in_ref_frame(sras: SrasFile, angle_idx: int, ref_idx: int, - pivot_mm: dict[int, tuple[float, float]], - shift_mm=(0.0, 0.0), - theta_deg: float | None = None) -> np.ndarray: - """Angle *angle_idx*'s bbox corners, rotated about its own alignment - pivot (pivot_mm[angle_idx] — see compute_pivot_points_mm) into the - reference frame and translated by *shift_mm*. Shape (4, 2). - - theta_deg overrides the analytic angles_deg-derived rotation used by - default — the manual-alignment path (union_canvas_mm) passes a - user-chosen rotation here (which may differ from the known scan-angle - delta) without needing a parallel code path. - """ - theta = _theta_deg(sras, angle_idx, ref_idx) if theta_deg is None else theta_deg - R = _rotation_matrix(theta) - c_a = np.array(pivot_mm[angle_idx]) - c_ref = np.array(pivot_mm[ref_idx]) - shift = np.asarray(shift_mm, dtype=np.float64) - return np.array([R @ (corner - c_a) + c_ref + shift - for corner in _bbox_corners_mm(sras, angle_idx)]) - - -def _build_affine_canvas_to_raw(sras: SrasFile, angle_idx: int, ref_idx: int, - shift_mm: tuple[float, float], - canvas_dx: float, canvas_dy: float, - canvas_origin_mm: tuple[float, float], - pivot_mm: dict[int, tuple[float, float]], - theta_deg: float | None = None - ) -> tuple[np.ndarray, np.ndarray]: - """matrix, offset s.t. raw_index = matrix @ [row_out, col_out] + offset, +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. - Pipeline (all mm unless noted): - [X;Y] = A_out @ [row_out;col_out] + b_out # canvas idx -> ref-frame mm - [lx;ly] = R(theta)^T @ ([X;Y]-c_ref-shift) + c_a # undo rotation+shift -> angle a's local mm - [row;col] = D @ ([lx;ly] - [x_start_a; y0_a]) # local mm -> angle a's raw idx + 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. - where theta is _theta_deg(angle_idx, ref_idx) unless *theta_deg* - overrides it (see _corners_in_ref_frame's note, used by the manual- - alignment path), c_ref/c_a are each angle's own alignment pivot - (pivot_mm — see compute_pivot_points_mm; the CH4-signal-weighted - centroid of its own footprint, not the raw scan-window bbox center — - this keeps the sample itself centered post-rotation, minimizing - required canvas padding and, more importantly, keeping alignment - relative to the sample rather than to wherever the scan window sat in - microscope/global XY space), and A_out/D are the index<->mm scaling - matrices for the canvas pitch and this angle's own native pitch - respectively. + 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. """ - theta = _theta_deg(sras, angle_idx, ref_idx) if theta_deg is None else theta_deg - Rinv = _rotation_matrix(theta).T - cx_a, cy_a = pivot_mm[angle_idx] - cx_ref, cy_ref = pivot_mm[ref_idx] - dx_a, dy_a = _pixel_pitch_mm(sras, angle_idx) - x0_a = float(sras.x_start_mm[angle_idx]) - y0_a = float(sras.y_positions_mm(angle_idx)[0]) - - A_out = np.array([[0.0, canvas_dx], [canvas_dy, 0.0]]) # [row,col] -> [X,Y] - b_out = np.array(canvas_origin_mm, dtype=np.float64) - D = np.array([[0.0, 1.0 / dy_a], [1.0 / dx_a, 0.0]]) # [x,y] -> [row,col] - shift = np.array(shift_mm, dtype=np.float64) - c_ref_v = np.array([cx_ref, cy_ref]) - c_a_v = np.array([cx_a, cy_a]) - origin_a = np.array([x0_a, y0_a]) + 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 - c_ref_v - shift) + D @ (c_a_v - origin_a) + offset = D @ (Rinv @ (b_out - shift) - off) + np.asarray(src_center_idx, dtype=np.float64) return matrix, offset @@ -596,7 +583,7 @@ def apply_alignment(result: AlignmentResult, angle_idx: int, img: np.ndarray, that angle's raw (n_rows, n_frames)) onto the shared alignment canvas. order=0 (nearest) avoids blending real data with zero-padding or with masked-out (0-valued) CH1/velocity pixels at mask edges. Channel- - agnostic: the same per-angle transform (found from the CH4 mask) works + 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( @@ -605,122 +592,533 @@ def apply_alignment(result: AlignmentResult, angle_idx: int, img: np.ndarray, mode="constant", cval=0.0) -def _block_mean_downsample(img: np.ndarray, factor: int) -> np.ndarray: - if factor <= 1: +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 // factor) * factor, (w // factor) * factor + h2, w2 = (h // fy) * fy, (w // fx) * fx + if h2 == 0 or w2 == 0: + return img trimmed = img[:h2, :w2] - return trimmed.reshape(h2 // factor, factor, w2 // factor, factor).mean(axis=(1, 3)) + return trimmed.reshape(h2 // fy, fy, w2 // fx, fx).mean(axis=(1, 3)) -def _phase_correlate_shift(ref_img: np.ndarray, mov_img: np.ndarray) -> tuple[int, int]: - """FFT normalized cross-power-spectrum phase correlation. Returns the - integer (dr, dc) pixel shift of mov_img relative to ref_img; both must - be the same shape. Risk: if the true shift is near +/- half the array - size, wraparound can bias the peak — mitigated by the generous - margin_frac padding in _working_canvas_for_pair, which keeps the true - residual shift small relative to the correlation canvas.""" - F1 = scipy_fft.fft2(ref_img.astype(np.float64), workers=-1) - F2 = scipy_fft.fft2(mov_img.astype(np.float64), workers=-1) - R = F1 * np.conj(F2) - R /= np.maximum(np.abs(R), 1e-12) - corr = scipy_fft.ifft2(R, workers=-1).real - dr, dc = np.unravel_index(np.argmax(corr), corr.shape) - h, w = corr.shape - if dr > h // 2: - dr -= h - if dc > w // 2: - dc -= w - return int(dr), int(dc) +# ---- 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 _working_canvas_for_pair(sras: SrasFile, ref_idx: int, a_idx: int, - dx: float, dy: float, - pivot_mm: dict[int, tuple[float, float]], - margin_frac: float = 0.3 - ) -> tuple[tuple[float, float], tuple[int, int]]: - """Union of the reference's own raw bbox and angle a's raw bbox rotated - (about its own alignment pivot) into the ref frame with zero shift, - padded by margin_frac on each side — sized generously so the true - phase-correlation shift lands well inside the canvas (see - _phase_correlate_shift's wraparound note).""" - pts = np.vstack([_bbox_corners_mm(sras, ref_idx), - _corners_in_ref_frame(sras, a_idx, ref_idx, pivot_mm)]) - x_min, y_min = pts.min(axis=0) - x_max, y_max = pts.max(axis=0) - pad_x, pad_y = (x_max - x_min) * margin_frac, (y_max - y_min) * margin_frac - x_min, x_max = x_min - pad_x, x_max + pad_x - y_min, y_max = y_min - pad_y, y_max + pad_y - n_cols = int(np.ceil((x_max - x_min) / dx)) + 1 - n_rows = int(np.ceil((y_max - y_min) / abs(dy))) + 1 - origin = (x_min, y_min if dy > 0 else y_max) - return origin, (n_rows, n_cols) +def _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 correlate_translation_mm(sras: SrasFile, angle_idx: int, ref_angle_idx: int, - signal_mv: dict[int, np.ndarray], - pivot_mm: dict[int, tuple[float, float]], *, - use_mask: bool = False, - dc_threshold_mv: float = 0.0, - margin_frac: float = 0.3, - max_corr_dim: int = 1024) -> tuple[float, float]: - """FFT phase-correlation translation (shift_x_mm, shift_y_mm) that best - lines up angle_idx's CH4 content onto ref_angle_idx's, given each angle's - own alignment pivot (pivot_mm — see compute_pivot_points_mm) already - rotated about that pivot by the analytic scan-angle delta with zero - shift. Returns (0.0, 0.0) unconditionally for the reference angle. +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))) - By default (use_mask=False) correlates on each angle's own raw CH4 - signal minus its own minimum — the same weighting _signal_centroid_mm - uses — rather than a dc_threshold_mv binary mask: a shared threshold - that doesn't suit every angle's real signal level can empty or distort - one angle's mask (the same failure mode the alignment pivot was fixed - to avoid), and correlating on the raw signal also lets phase - correlation lock onto real internal sample structure rather than just - the scan window's silhouette. Subtracting each angle's own minimum - (rather than using the signal as-is) keeps the zero-padding surrounding - the rotated image from reading as a spurious high-contrast edge against - a nonzero DC baseline. use_mask=True switches to the binary - >= dc_threshold_mv mask instead (compute_angle_alignment's original - behavior), for cases where raw-signal correlation locks onto noise. + +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 (0.0, 0.0) + return RigidFit(0.0, (0.0, 0.0), 1.0, "reference") - def corr_img(a: int) -> np.ndarray: - img = signal_mv[a] - if use_mask: - return (img >= dc_threshold_mv).astype(np.float32) - return (img - img.min()).astype(np.float32) + 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") - img_a = corr_img(angle_idx) - img_ref = corr_img(ref_angle_idx) + nominal = _nominal_delta_deg(sras, angle_idx, ref_angle_idx) + thetas = _rotation_candidates(nominal, search_deg, coarse_step_deg) - dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx) - max_dim = max(img_a.shape + img_ref.shape) - factor = max(1, int(np.ceil(max_dim / max_corr_dim))) - dx_c, dy_c = dx_ref * factor, dy_ref * factor - small_a = _block_mean_downsample(img_a, factor) - small_ref = _block_mean_downsample(img_ref, factor) + # ---- 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) - work_origin, work_shape = _working_canvas_for_pair( - sras, ref_angle_idx, angle_idx, dx_c, dy_c, pivot_mm, margin_frac=margin_frac) + if best[3] == "none": + return RigidFit(0.0, (0.0, 0.0), -1.0, "none") - m_a, o_a = _build_affine_canvas_to_raw( - sras, angle_idx, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin, pivot_mm) - m_ref, o_ref = _build_affine_canvas_to_raw( - sras, ref_angle_idx, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin, pivot_mm) + # ---- 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) - rotated_a = scipy_ndimage.affine_transform( - small_a, m_a / factor, offset=o_a / factor, - output_shape=work_shape, order=0, mode="constant", cval=0.0) - embedded_ref = scipy_ndimage.affine_transform( - small_ref, m_ref / factor, offset=o_ref / factor, - output_shape=work_shape, order=0, mode="constant", cval=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 = _phase_correlate_shift(embedded_ref, rotated_a) - return (dc * dx_c, dr * dy_c) + 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: @@ -733,14 +1131,17 @@ def _parallel_map(fn, items, n_workers: int) -> list: def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int, - dc_threshold_mv: float, - progress_cb=None) -> AlignmentResult: - """Top-level alignment driver. Runs on a background thread (see - AngleAlignmentWorker) — deliberately recomputes CH4 DC images from - scratch rather than reading the GUI-thread _dc_cache dict, since - background-thread workers must not touch GUI-thread-owned caches.""" + 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 - dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx) # Parallel over angles, serial within each — see plan_angle_level. n_workers, angle_budget = plan_angle_level(sras) @@ -753,7 +1154,7 @@ def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int, _ticks.append(1) progress_cb(base + int(min(len(_ticks), n) / n * span)) - # ---- Step 1: CH4 DC image + binarized mask per angle, native grid ----- + # ---- 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), @@ -763,51 +1164,28 @@ def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int, dc4_mv = dict(enumerate(_parallel_map(dc4_for, range(n), n_workers))) - # Each angle's own alignment pivot — the CH4-signal-weighted centroid, - # not the raw scan-window bbox center (see compute_pivot_points_mm). - # Reuses the dc4_mv images just computed above, so this is free, and is - # deliberately independent of dc_threshold_mv (see that function's - # docstring) so a threshold that happens to leave some angle's binary - # mask empty can't silently degrade the pivot back to the bbox center. - pivot_mm = compute_pivot_points_mm(sras, dc4_mv=dc4_mv) + # ---- Step 2: 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() - # ---- Step 2: coarse translation via FFT phase correlation, on each - # angle's binarized CH4 mask against the reference's (see - # correlate_translation_mm — also the engine behind ManualAlignmentDialog's - # Auto Cross-Correlate button, there defaulting to the raw signal instead - # of a mask). - def shift_for(a: int) -> tuple[float, float]: - shift = correlate_translation_mm( - sras, a, ref_angle_idx, dc4_mv, pivot_mm, - use_mask=True, dc_threshold_mv=dc_threshold_mv) - tick(25, 50) - return shift + 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 - shifts_mm = dict(enumerate(_parallel_map(shift_for, range(n), n_workers))) + fits = dict(enumerate(_parallel_map( + fit_for, range(n), _registration_workers(sras, fine_dim)))) - # ---- Step 3: union bounding box over all angles (rotation+shift applied) - corners = np.vstack([ - _corners_in_ref_frame(sras, a, ref_angle_idx, pivot_mm, shifts_mm[a]) - for a in range(n)]) - x_min, y_min = corners.min(axis=0) - x_max, y_max = corners.max(axis=0) - n_cols = int(np.ceil((x_max - x_min) / dx_ref)) + 1 - n_rows = int(np.ceil((y_max - y_min) / abs(dy_ref))) + 1 - canvas_origin_mm = (float(x_min), float(y_min if dy_ref > 0 else y_max)) - - # ---- Step 4: final per-angle full-resolution affine (canvas -> raw idx) - per_angle: dict[int, AngleTransform] = {} - for a in range(n): - matrix, offset = _build_affine_canvas_to_raw( - sras, a, ref_angle_idx, shifts_mm[a], dx_ref, dy_ref, canvas_origin_mm, - pivot_mm) - per_angle[a] = AngleTransform( - _theta_deg(sras, a, ref_angle_idx), shifts_mm[a], matrix, offset) - if progress_cb: - progress_cb(75 + int((a + 1) / n * 25)) - - return AlignmentResult(ref_angle_idx, dc_threshold_mv, (n_rows, n_cols), - dx_ref, dy_ref, canvas_origin_mm, per_angle) + # ---- 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 # Back-compat alias for the pre-split private name (used by tooling). @@ -817,80 +1195,36 @@ _compute_angle_alignment = compute_angle_alignment # --------------------------------------------------------------------------- # Manual alignment (Fusion menu -> Manual Alignment... dialog) # -# Skips compute_angle_alignment's mask + phase-correlation search entirely: -# every angle's rotation_deg/shift_mm is supplied directly by the caller -# (nudged by eye against a live multi-angle mask overlay, or pre-seeded from -# a previous compute_angle_alignment run or a saved sidecar). Building the -# final AlignmentResult from already-known per-angle parameters is pure -# closed-form matrix math (union_canvas_mm + _build_affine_canvas_to_raw) -# with no per-pixel image work at all, so build_manual_alignment is cheap -# enough to call synchronously on the GUI thread on every edit. The only -# genuinely expensive per-pixel operation anywhere in this flow is -# reproject_mask, and only ManualAlignmentDialog's own downsampled preview -# calls that per keystroke — see that class's docstring for how it limits -# each nudge to reprojecting only the actively-edited angle. +# 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 union_canvas_mm(sras: SrasFile, ref_angle_idx: int, dx: float, dy: float, - per_angle_params: dict[int, ManualAngleParams], - pivot_mm: dict[int, tuple[float, float]], - margin_frac: float = 0.0 - ) -> tuple[tuple[float, float], tuple[int, int]]: - """Shared-canvas origin (mm) and (n_rows, n_cols) at pitch (dx, dy) that - contains every angle's footprint after applying its own rotation+shift — - the generalisation of compute_angle_alignment's Step 3 to arbitrary (not - just _theta_deg-analytic) per-angle rotation. Angles missing from - per_angle_params default to identity (e.g. a sidecar saved before a - rescan added more angles). - - margin_frac pads the box on every side: 0 for a final canvas (this then - reproduces compute_angle_alignment's own Step-3 math exactly, when every - angle's rotation_deg equals the analytic delta and shift_mm matches); - nonzero for ManualAlignmentDialog's downsampled preview canvas, which - needs headroom so an ordinary translation nudge of the active angle never - has to trigger a full canvas resize (see that class's docstring — an - extreme nudge can still, in principle, push content past this padding; - accepted as a known edge case, same as _phase_correlate_shift's - wraparound risk note). - """ - n = sras.n_angles - corners = np.vstack([ - _corners_in_ref_frame( - sras, a, ref_angle_idx, pivot_mm, - per_angle_params.get(a, ManualAngleParams()).shift_mm, - theta_deg=per_angle_params.get(a, ManualAngleParams()).rotation_deg) - for a in range(n)]) - x_min, y_min = corners.min(axis=0) - x_max, y_max = corners.max(axis=0) - if margin_frac: - pad_x, pad_y = (x_max - x_min) * margin_frac, (y_max - y_min) * margin_frac - x_min, x_max = x_min - pad_x, x_max + pad_x - y_min, y_max = y_min - pad_y, y_max + pad_y - n_cols = int(np.ceil((x_max - x_min) / dx)) + 1 - n_rows = int(np.ceil((y_max - y_min) / abs(dy))) + 1 - origin = (float(x_min), float(y_min if dy > 0 else y_max)) - return origin, (n_rows, n_cols) - - def reproject_mask(sras: SrasFile, angle_idx: int, ref_angle_idx: int, mask: np.ndarray, rotation_deg: float, shift_mm: tuple[float, float], - canvas_dx: float, canvas_dy: float, + canvas_pitch_mm: tuple[float, float], canvas_origin_mm: tuple[float, float], canvas_shape: tuple[int, int], - pivot_mm: dict[int, tuple[float, float]]) -> np.ndarray: + *, 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), since it bypasses - compute_angle_alignment's phase-correlation search entirely and just - takes rotation_deg/shift_mm as given. order=0 (nearest) matches - apply_alignment's own reasoning: a binary mask must never be blended with - zero-padding. - """ - matrix, offset = _build_affine_canvas_to_raw( - sras, angle_idx, ref_angle_idx, shift_mm, canvas_dx, canvas_dy, - canvas_origin_mm, pivot_mm, theta_deg=rotation_deg) + 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) @@ -898,48 +1232,24 @@ def reproject_mask(sras: SrasFile, angle_idx: int, ref_angle_idx: int, def build_manual_alignment(sras: SrasFile, ref_angle_idx: int, dc_threshold_mv: float, - per_angle_params: dict[int, ManualAngleParams], - pivot_mm: dict[int, tuple[float, float]] | None = None + 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 mask/phase-correlation search - entirely (every angle's transform here is exactly what the caller - supplied). Pure matrix/bbox math once pivot_mm is known, so it is cheap - enough to call synchronously on the GUI thread. The reference angle's - params are always forced to identity, regardless of what - per_angle_params holds for it — it defines the shared origin and must - never be transformed. + 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. - Pass an already-computed *pivot_mm* (e.g. ManualAlignmentDialog's own - cache, built once from its live CH4 images) to skip recomputing every - angle's DC image here; otherwise it's computed fresh via - compute_pivot_points_mm, which is still fine for a one-off Save or - sidecar restore (just not free on a very large, not-yet-cached scan). - dc_threshold_mv itself plays no part in the pivot (see that function's - docstring) — it's stored on the returned AlignmentResult purely as a - record of the RF-mask threshold in effect at the time. + 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. """ - n = sras.n_angles - dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx) - params = {a: per_angle_params.get(a, ManualAngleParams()) for a in range(n)} + params = {a: per_angle_params.get(a, ManualAngleParams()) + for a in range(sras.n_angles)} params[ref_angle_idx] = ManualAngleParams() - if pivot_mm is None: - pivot_mm = compute_pivot_points_mm(sras) - - canvas_origin_mm, canvas_shape = union_canvas_mm( - sras, ref_angle_idx, dx_ref, dy_ref, params, pivot_mm, margin_frac=0.0) - - per_angle: dict[int, AngleTransform] = {} - for a in range(n): - p = params[a] - matrix, offset = _build_affine_canvas_to_raw( - sras, a, ref_angle_idx, p.shift_mm, dx_ref, dy_ref, - canvas_origin_mm, pivot_mm, theta_deg=p.rotation_deg) - per_angle[a] = AngleTransform(p.rotation_deg, p.shift_mm, matrix, offset) - - return AlignmentResult(ref_angle_idx, dc_threshold_mv, canvas_shape, - dx_ref, dy_ref, canvas_origin_mm, per_angle) + return _result_from_params(sras, ref_angle_idx, dc_threshold_mv, params) # ---- Sidecar persistence (.sras.align.json) ------------------------- @@ -967,15 +1277,19 @@ def sidecar_path(sras_path) -> Path: return p.with_name(p.name + ".align.json") -# Bumped from 1 -> 2 when the rotation pivot changed from the raw scan- -# window bbox center to a content-derived centroid, *and* _theta_deg's sign -# convention was corrected — either change alone makes a version-1 file's -# stored rotation_deg/shift_mm numbers describe a different (and, for the -# pivot bug, actively wrong) transform than they would today. Loading one -# unchanged would silently reproduce exactly the "scans show up everywhere" -# symptom these fixes address, so version-1 sidecars are treated as absent -# rather than migrated. -_SIDECAR_SCHEMA_VERSION = 2 +# 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, @@ -984,9 +1298,9 @@ def save_manual_alignment(sras: SrasFile, ref_angle_idx: int, """Write the sidecar JSON for sras.path (overwriting any existing one) and return the path written. - Schema (schema_version 2): + Schema (schema_version 3): { - "schema_version": 2, + "schema_version": 3, "ref_angle_idx": , "dc_threshold_mv": , "per_angle": { diff --git a/sras_edit_scans.py b/sras_edit_scans.py new file mode 100644 index 0000000..8774daf --- /dev/null +++ b/sras_edit_scans.py @@ -0,0 +1,295 @@ +#!/usr/bin/env python3 +""" +sras_edit_scans.py — Remove one or more angle scans from a .sras file. + +A .sras file holds one or more "angles" (rotation positions); the viewer +cross-correlates each non-reference angle against the reference to align +them. If one angle's acquisition went wrong (stage glitch, bad trigger, +laser dropout, ...) it throws off that alignment for the whole file. This +tool drops the bad angle(s) and renumbers the rest, writing a new .sras file +with everything else — waveform samples, calibration preambles, background +waveform, row/geometry tables — carried over byte-for-byte. + +Handles v2-v7. Any precomputed FFT/DC cache (v5 PREC tail, v7 CACH tail) is +dropped on write, since it's indexed by angle and would be stale/misaligned +after renumbering; the viewer just recomputes it next time the file opens. + +Usage: + python sras_edit_scans.py input.sras --list + python sras_edit_scans.py input.sras output.sras --drop 2,5 + python sras_edit_scans.py input.sras output.sras --keep 0,1,3,4,6 +""" + +import argparse +import struct +import sys +from pathlib import Path + +from sras_format import ( + GEO_FMT_V6, GEO_SIZE_V6, HDR_FMT, HDR_FMT_V6, HDR_SIZE, HDR_SIZE_V6, + SrasFile, +) + +_LEGACY_VERSIONS = (2, 3, 4, 5) +_V6_VERSIONS = (6, 7) + + +def parse_args(): + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("input", help="Input .sras file") + p.add_argument("output", nargs="?", help="Output .sras file (omit with --list)") + p.add_argument("--list", action="store_true", + help="Print each angle's index/degrees/geometry and exit") + g = p.add_mutually_exclusive_group() + g.add_argument("--drop", metavar="I,J,...", + help="Comma-separated angle indices to remove") + g.add_argument("--keep", metavar="I,J,...", + help="Comma-separated angle indices to keep (all others dropped)") + return p.parse_args() + + +def _parse_index_list(s: str, n_angles: int) -> set[int]: + out = set() + for piece in s.split(","): + piece = piece.strip() + if not piece: + continue + i = int(piece) + if not (0 <= i < n_angles): + raise ValueError(f"angle index {i} out of range [0, {n_angles - 1}]") + out.add(i) + return out + + +def print_listing(sras: SrasFile): + print(f"\n{'idx':>4} {'angle_deg':>10} {'x_start_mm':>11} {'rows':>6} {'frames':>7}") + for a in range(sras.n_angles): + print(f"{a:>4} {sras.angles_deg[a]:>10.4f} {sras.x_start_mm[a]:>11.4f} " + f"{int(sras.n_rows[a]):>6} {int(sras.n_frames[a]):>7}") + + +def _copy_range(fin, fout, offset: int, nbytes: int, chunk: int = 64 * 1024 * 1024): + """Stream *nbytes* raw bytes from *fin* at *offset* into *fout*, without + ever holding more than one chunk in memory (waveform blocks can be + hundreds of MB to low GB each).""" + fin.seek(offset) + remaining = nbytes + while remaining: + buf = fin.read(min(chunk, remaining)) + if not buf: + raise IOError("unexpected EOF while copying waveform data") + fout.write(buf) + remaining -= len(buf) + + +# --------------------------------------------------------------------------- +# Legacy (v2-v5): uniform geometry across angles, one flat waveform block +# --------------------------------------------------------------------------- + +def _write_legacy(sras: SrasFile, keep: list[int], out_path: Path): + n_rows = int(sras.n_rows[0]) + n_frames = int(sras.n_frames[0]) # uniform across angles for v2-v5 + n_ch = sras.n_channels + spf = sras.samples_per_frame + bps = sras.bytes_per_sample + + header = struct.pack( + HDR_FMT, b"SRAS", sras.version, len(keep), n_rows, + float(sras.x_start_mm[0]), float(sras.x_delta_mm), + sras.velocity_mm_s, sras.laser_freq_hz, + n_frames, spf, sras.sample_rate_hz, bps, n_ch, + ) + + # Row table + preambles + background sit right after the angle table and + # don't vary per angle — copy that whole span through unmodified. + angle_table_size = sras.n_angles * 4 + with open(sras.path, "rb") as f: + f.seek(HDR_SIZE + angle_table_size) + shared_mid = f.read(sras._data_offset - (HDR_SIZE + angle_table_size)) + + angle_bytes = n_rows * n_ch * n_frames * spf * bps + + with open(sras.path, "rb") as fin, open(out_path, "wb") as fout: + fout.write(header) + fout.write(sras.angles_deg[keep].astype(">f4").tobytes()) + fout.write(shared_mid) + for a in keep: + _copy_range(fin, fout, sras._data_offset + a * angle_bytes, angle_bytes) + + +# --------------------------------------------------------------------------- +# v6/v7: per-angle geometry, ragged waveform blocks +# --------------------------------------------------------------------------- + +def _read_v6_sections(path: Path) -> dict: + """Raw, low-level read of everything before the waveform data. + + SrasFile._parse_v6 reads and discards each angle's x_delta (it's not + part of the display-facing geometry it exposes), so a round-trip through + SrasFile would silently drop that field. Re-parsing here keeps every + byte of the Per-Angle Geometry Table intact. + """ + with open(path, "rb") as f: + hdr_raw = f.read(HDR_SIZE_V6) + (magic, ver, n_angles_declared, x_start_nom, y_start_nom, x_delta_nom, + y_delta_nom, row_spacing, vel, freq, spf, sr, bps, + n_ch) = struct.unpack(HDR_FMT_V6, hdr_raw) + + angles_deg = list(struct.unpack(f">{n_angles_declared}f", + f.read(n_angles_declared * 4))) + + geo = [struct.unpack(GEO_FMT_V6, f.read(GEO_SIZE_V6)) + for _ in range(n_angles_declared)] # (x_start, x_delta, n_frames, n_rows) + + row_table = [f.read(geo[a][3] * 4) for a in range(n_angles_declared)] + + preamble_start = f.tell() + for _ in range(n_ch): + (length,) = struct.unpack(">H", f.read(2)) + f.read(length) + preambles_raw = _reread_span(f, preamble_start) + + bg_start = f.tell() + (n_bg,) = struct.unpack(">I", f.read(4)) + f.read(n_bg) + background_raw = _reread_span(f, bg_start) + + data_offset = f.tell() + + return { + "version": ver, "n_angles_declared": n_angles_declared, + "x_start_nom": x_start_nom, "y_start_nom": y_start_nom, + "x_delta_nom": x_delta_nom, "y_delta_nom": y_delta_nom, + "row_spacing": row_spacing, "vel": vel, "freq": freq, + "spf": spf, "sr": sr, "bps": bps, "n_ch": n_ch, + "angles_deg": angles_deg, "geo": geo, "row_table": row_table, + "preambles_raw": preambles_raw, "background_raw": background_raw, + "data_offset": data_offset, + } + + +def _reread_span(f, start: int) -> bytes: + end = f.tell() + f.seek(start) + span = f.read(end - start) + f.seek(end) + return span + + +def _v6_angle_offsets(sections: dict, file_size: int) -> list[tuple[int, int]]: + """(offset, nbytes) of each declared angle's waveform block, stopping at + the first angle whose data isn't fully on disk (aborted scan).""" + n_ch, spf, bps = sections["n_ch"], sections["spf"], sections["bps"] + offsets = [] + offset = sections["data_offset"] + for xs, xd, nf, nr in sections["geo"]: + nbytes = nr * n_ch * nf * spf * bps + if offset + nbytes > file_size: + break + offsets.append((offset, nbytes)) + offset += nbytes + return offsets + + +def _write_v6(in_path: Path, sections: dict, keep: list[int], out_path: Path): + file_size = in_path.stat().st_size + offsets = _v6_angle_offsets(sections, file_size) + geo = sections["geo"] + + header = struct.pack( + HDR_FMT_V6, b"SRAS", sections["version"], len(keep), + sections["x_start_nom"], sections["y_start_nom"], + sections["x_delta_nom"], sections["y_delta_nom"], + sections["row_spacing"], sections["vel"], sections["freq"], + sections["spf"], sections["sr"], sections["bps"], sections["n_ch"], + ) + + with open(in_path, "rb") as fin, open(out_path, "wb") as fout: + fout.write(header) + fout.write(struct.pack(f">{len(keep)}f", + *[sections["angles_deg"][i] for i in keep])) + for i in keep: + fout.write(struct.pack(GEO_FMT_V6, *geo[i])) + for i in keep: + fout.write(sections["row_table"][i]) + fout.write(sections["preambles_raw"]) + fout.write(sections["background_raw"]) + for i in keep: + off, nbytes = offsets[i] + _copy_range(fin, fout, off, nbytes) + + +def main(): + args = parse_args() + in_path = Path(args.input) + if not in_path.exists(): + print(f"Error: input file not found: {in_path}", file=sys.stderr) + sys.exit(1) + + print(f"Reading {in_path} ...", flush=True) + try: + sras = SrasFile(str(in_path)) + except ValueError as e: + print(f"Error: {e}", file=sys.stderr) + sys.exit(1) + + if sras.version not in (*_LEGACY_VERSIONS, *_V6_VERSIONS): + print(f"Error: unsupported .sras version: {sras.version}", file=sys.stderr) + sys.exit(1) + + aborted_note = " (scan aborted; trailing angle(s) already excluded)" if sras.scan_aborted else "" + print(f" Version : v{sras.version}", flush=True) + print(f" Angles : {sras.n_angles}{aborted_note}", flush=True) + + if args.list: + print_listing(sras) + return + + if not args.output: + print("Error: output path required unless --list is given.", file=sys.stderr) + sys.exit(1) + if not (args.drop or args.keep): + print("Error: specify --drop or --keep (see --list for indices).", file=sys.stderr) + sys.exit(1) + + out_path = Path(args.output) + if out_path.resolve() == in_path.resolve(): + print("Error: output path must differ from input path.", file=sys.stderr) + sys.exit(1) + + try: + if args.drop: + drop = _parse_index_list(args.drop, sras.n_angles) + keep = [a for a in range(sras.n_angles) if a not in drop] + else: + keep = sorted(_parse_index_list(args.keep, sras.n_angles)) + except ValueError as e: + print(f"Error: {e}", file=sys.stderr) + sys.exit(1) + + if not keep: + print("Error: at least one angle must remain.", file=sys.stderr) + sys.exit(1) + + dropped = [a for a in range(sras.n_angles) if a not in keep] + print(f"\nDropping angle(s): {dropped}") + print(f"Keeping angle(s) : {keep} ({len(keep)} of {sras.n_angles})") + print(f"\nWriting {out_path} ...", flush=True) + + if sras.version in _LEGACY_VERSIONS: + _write_legacy(sras, keep, out_path) + else: + _write_v6(in_path, _read_v6_sections(in_path), keep, out_path) + + in_mb = in_path.stat().st_size / 1024**2 + out_mb = out_path.stat().st_size / 1024**2 + print(f" Input size : {in_mb:.1f} MB") + print(f" Output size: {out_mb:.1f} MB") + print("Done.") + print("Note: any precomputed FFT/DC cache was dropped (it's indexed by " + "angle); the viewer will recompute it next time this file opens.") + + +if __name__ == "__main__": + main() diff --git a/sras_viewer.py b/sras_viewer.py index fb7e8b5..54272c0 100644 --- a/sras_viewer.py +++ b/sras_viewer.py @@ -821,18 +821,18 @@ class ManualAlignmentDialog(QDialog): misalignment is visible by eye. Reference angle (always index 0) is ground truth and never moves; every other angle is aligned to it. The user picks an "active" angle and nudges its rotation+translation with - the keyboard; Auto De-rotate sets every non-reference angle's rotation to - the known, analytic scan-angle delta without touching any translation; - Auto Cross-Correlate does the same rotation and additionally sets - translation to the FFT-phase-correlation best fit against the reference - (see compute.correlate_translation_mm) — meant to get every angle roughly - stacked on top of each other so keyboard nudging only has to make small - corrections, not find a coarse alignment from scratch. Save writes a - JSON sidecar next to the .sras file and hands a freshly-built, full- - resolution AlignmentResult back to the main window — the exact same - object shape compute_angle_alignment produces, so every existing - Aligned-View code path (apply_alignment, _aligned_canvas_axes, the - pixel-inspector inverse-transform) works completely unmodified. + the keyboard; Auto Cross-Correlate finds every non-reference angle's + rotation *and* translation by registering its image against the + reference's (see compute.register_angle_to_reference) — meant to get every + angle stacked on top of each other so keyboard nudging only has to make + small corrections, not find an alignment from scratch; Auto De-rotate is + the weaker fallback that just seeds rotation from the stage's reported + angle, leaving translation alone. Save writes a JSON sidecar next to the + .sras file and hands a freshly-built, full-resolution AlignmentResult back + to the main window — the exact same object shape compute_angle_alignment + produces, so every existing Aligned-View code path (apply_alignment, + _aligned_canvas_axes, the pixel-inspector inverse-transform) works + completely unmodified. Non-modal by design (shown via .show(), never .exec() or setModal(True)) so the user can still interact with the main window. Talks back to @@ -853,6 +853,16 @@ class ManualAlignmentDialog(QDialog): _ACTIVE_ALPHA = 0.75 _MAX_PREVIEW_DIM = 1024 + # (label, sources passed to compute.register_angle_to_reference). "Both" + # registers on each and keeps whichever scores higher per angle, which + # costs roughly double but removes the failure mode where the single + # chosen source is the one that happens to be uninformative for one angle. + _CORRELATE_SOURCES = ( + ("Both, keep best (recommended)", ("signal", "mask")), + ("Raw signal", ("signal",)), + ("Thresholded mask", ("mask",)), + ) + def __init__(self, parent: "SrasViewerWindow", sras: SrasFile, *, ref_angle_idx: int, dc_threshold_mv: float, seed_per_angle: dict[int, ManualAngleParams] | None, @@ -861,15 +871,16 @@ class ManualAlignmentDialog(QDialog): self._parent = parent self._sras = sras self._ref_angle_idx = ref_angle_idx - self._downsample_factor = 1 + self._downsample = (1, 1) # (rows, cols) block-mean factors self._dc4_mv: dict[int, np.ndarray] = {} self._masks_small: dict[int, np.ndarray] = {} - self._pivot_mm: dict[int, tuple[float, float]] = {} self._preview_layers: dict[int, np.ndarray] = {} self._preview_origin_mm = (0.0, 0.0) self._preview_shape = (1, 1) - self._preview_dx_mm = self._preview_dy_mm = 1.0 + self._preview_pitch_mm = (1.0, 1.0) self._masks_ready = False + self._fit_notes: dict[int, tuple[float, str]] = {} + self._derotate_sign_flipped = False self.setWindowTitle(f"Manual Alignment — {sras.path.name}") self.resize(1150, 760) @@ -1006,25 +1017,27 @@ class ManualAlignmentDialog(QDialog): self.grp_correlate, cl = _group("Cross-Correlate (FFT)") cform = _form() self.combo_correlate_source = QComboBox() - self.combo_correlate_source.addItems( - ["Raw signal (recommended)", "Thresholded mask"]) + for label, sources in self._CORRELATE_SOURCES: + self.combo_correlate_source.addItem(label, sources) cform.addRow("Correlate on:", self.combo_correlate_source) - self.spin_correlate_margin = QDoubleSpinBox() - self.spin_correlate_margin.setRange(0.05, 2.0) - self.spin_correlate_margin.setSingleStep(0.05) - self.spin_correlate_margin.setDecimals(2) - self.spin_correlate_margin.setValue(0.30) - self.spin_correlate_margin.setMinimumWidth(_SPIN_MIN_W) - cform.addRow("Search margin (× extent):", self.spin_correlate_margin) + self.spin_correlate_search_deg = QDoubleSpinBox() + self.spin_correlate_search_deg.setRange(0.0, 180.0) + self.spin_correlate_search_deg.setSingleStep(1.0) + self.spin_correlate_search_deg.setDecimals(1) + self.spin_correlate_search_deg.setSuffix(" °") + self.spin_correlate_search_deg.setValue(6.0) + self.spin_correlate_search_deg.setMinimumWidth(_SPIN_MIN_W) + cform.addRow("Rotation search (±):", self.spin_correlate_search_deg) cl.addLayout(cform) self.btn_auto_correlate = QPushButton("Auto Cross-Correlate (vs Reference)") cl.addWidget(self.btn_auto_correlate) cl.addWidget(_wrap_label( - "Sets rotation to the known scan angle and translation to the " - "FFT-correlated best fit for every non-reference angle. Run this " - "first, then use manual nudging only for small corrections.", - _CSS_HINT)) + "Finds each non-reference angle's rotation *and* translation by " + "cross-correlating its image against the reference's — the stage's " + "reported angle is only the starting point of the search, and both " + "of its signs are tried. Run this first, then nudge only for small " + "corrections.", _CSS_HINT)) panel_l.addWidget(self.grp_correlate) # ---- Actions ------------------------------------------------------ @@ -1091,15 +1104,16 @@ class ManualAlignmentDialog(QDialog): def _finish_mask_prep(self): if len(self._dc4_mv) < self._sras.n_angles: return # a mask-worker error left some angles unfetched - max_dim = max(max(img.shape) for img in self._dc4_mv.values()) - self._downsample_factor = max(1, int(np.ceil(max_dim / self._MAX_PREVIEW_DIM))) + # Rows and columns get their own factor. A real scan is ~7500 frames + # wide but only ~750 rows tall, so one shared factor sized for the + # frames would throw away 8x more row detail than the preview needs and + # leave the overlay too coarse in y to judge alignment by eye. + max_rows = max(img.shape[0] for img in self._dc4_mv.values()) + max_cols = max(img.shape[1] for img in self._dc4_mv.values()) + self._downsample = ( + max(1, int(np.ceil(max_rows / self._MAX_PREVIEW_DIM))), + max(1, int(np.ceil(max_cols / self._MAX_PREVIEW_DIM)))) self._recompute_masks_small() - # Alignment pivot: the CH4-signal-weighted centroid of each angle's - # own footprint (see compute.compute_pivot_points_mm) — computed once - # from the full-res CH4 images and deliberately independent of the - # mask threshold, so it never needs recomputing when that changes - # (unlike _masks_small, which is purely for the overlay's visuals). - self._pivot_mm = compute.compute_pivot_points_mm(self._sras, self._dc4_mv) self._rebuild_preview_canvas() self._set_controls_enabled(True) self.lbl_status.setText("Ready.") @@ -1108,13 +1122,12 @@ class ManualAlignmentDialog(QDialog): """Threshold + downsample every angle's already-in-memory full-res CH4 mV image. Cheap (a compare + block-mean), so this re-runs in full whenever the mask-threshold spin box changes — no re-fetch. - Purely for the overlay's visuals — the alignment pivot does not - depend on this threshold (see _pivot_mm / compute_pivot_points_mm).""" + Purely for the overlay's visuals: no alignment geometry depends on this + threshold, only which pixels the overlay paints.""" threshold = self.spin_mask_threshold_mv.value() - factor = self._downsample_factor + fy, fx = self._downsample self._masks_small = { - a: compute._block_mean_downsample( - (img >= threshold).astype(np.float32), factor) + a: compute._block_mean_2d((img >= threshold).astype(np.float32), fy, fx) for a, img in self._dc4_mv.items() } @@ -1131,32 +1144,35 @@ class ManualAlignmentDialog(QDialog): active angle. NOT triggered by a translation-only nudge — see _refresh_active_preview_layer.""" dx_ref, dy_ref = compute._pixel_pitch_mm(self._sras, self._ref_angle_idx) - factor = self._downsample_factor - dx_c, dy_c = dx_ref * factor, dy_ref * factor - origin, shape = compute.union_canvas_mm( - self._sras, self._ref_angle_idx, dx_c, dy_c, self._angle_params, - self._pivot_mm, margin_frac=self._PREVIEW_MARGIN_FRAC) + fy, fx = self._downsample + pitch = (dx_ref * fx, dy_ref * fy) + origin, shape = compute.canvas_for_params( + self._sras, self._ref_angle_idx, pitch, self._angle_params, + margin_frac=self._PREVIEW_MARGIN_FRAC, snap=False) self._preview_origin_mm, self._preview_shape = origin, shape - self._preview_dx_mm, self._preview_dy_mm = dx_c, dy_c + self._preview_pitch_mm = pitch self._preview_layers = { - a: compute.reproject_mask( - self._sras, a, self._ref_angle_idx, self._masks_small[a], - self._angle_params[a].rotation_deg, self._angle_params[a].shift_mm, - dx_c, dy_c, origin, shape, self._pivot_mm) - for a in range(self._sras.n_angles) + a: self._reproject(a) for a in range(self._sras.n_angles) } self._redraw_overlay() + def _reproject(self, angle_idx: int) -> np.ndarray: + """One angle's downsampled mask on the current preview canvas. + src_downsample must match _masks_small's block-mean factors, or the + layer lands magnified and offset instead of where the alignment + actually puts it.""" + p = self._angle_params[angle_idx] + return compute.reproject_mask( + self._sras, angle_idx, self._ref_angle_idx, + self._masks_small[angle_idx], p.rotation_deg, p.shift_mm, + self._preview_pitch_mm, self._preview_origin_mm, self._preview_shape, + src_downsample=self._downsample) + def _refresh_active_preview_layer(self): """Cheap path for a translation-only nudge/edit of the active angle: reproject just that one angle's downsampled mask onto the *existing* preview canvas — every other angle's cached layer is untouched.""" - a = self._active_angle - self._preview_layers[a] = compute.reproject_mask( - self._sras, a, self._ref_angle_idx, self._masks_small[a], - self._angle_params[a].rotation_deg, self._angle_params[a].shift_mm, - self._preview_dx_mm, self._preview_dy_mm, - self._preview_origin_mm, self._preview_shape, self._pivot_mm) + self._preview_layers[self._active_angle] = self._reproject(self._active_angle) self._redraw_overlay() def _redraw_overlay(self): @@ -1182,7 +1198,7 @@ class ManualAlignmentDialog(QDialog): rgba[..., 3] = fg_a + rgba[..., 3] * (1 - fg_a) x0, y0 = self._preview_origin_mm - dx, dy = self._preview_dx_mm, self._preview_dy_mm + dx, dy = self._preview_pitch_mm x_axis = x0 + np.arange(n_cols) * dx y_axis = y0 + np.arange(n_rows) * dy extent = [x_axis[0] - dx / 2, x_axis[-1] + dx / 2, @@ -1258,18 +1274,29 @@ class ManualAlignmentDialog(QDialog): # ------------------------------------------------------------------ def _on_auto_derotate(self): + """Seed every angle's rotation from the stage's reported angle. + + A starting point for nudging by eye, not an alignment: the stage's + sign convention relative to this module's is not knowable from the + file, so the sign that lines the scans up is whichever of the two looks + right in the overlay. Auto Cross-Correlate decides that from the images + instead, and is the button to reach for first. + """ + sign = -1.0 if self._derotate_sign_flipped else 1.0 + self._derotate_sign_flipped = not self._derotate_sign_flipped n_changed = 0 for a in range(self._sras.n_angles): if a == self._ref_angle_idx: continue - self._angle_params[a].rotation_deg = compute._theta_deg( + self._angle_params[a].rotation_deg = sign * compute._nominal_delta_deg( self._sras, a, self._ref_angle_idx) n_changed += 1 self._sync_active_spinboxes() self._rebuild_preview_canvas() self.lbl_status.setText( - f"Rotation set to the known scan angle for {n_changed} angle(s) " - "(translation left untouched).") + f"Rotation set to the stage angle ({'−' if sign < 0 else '+'}delta) " + f"for {n_changed} angle(s); translation untouched. Click again to " + "try the opposite sign.") def _on_auto_correlate(self): if not self._masks_ready: @@ -1277,13 +1304,14 @@ class ManualAlignmentDialog(QDialog): angles = [a for a in range(self._sras.n_angles) if a != self._ref_angle_idx] if not angles: return - use_mask = self.combo_correlate_source.currentIndex() == 1 worker = CrossCorrelateWorker( - self._sras, self._ref_angle_idx, angles, self._dc4_mv, self._pivot_mm, - use_mask=use_mask, dc_threshold_mv=self.spin_mask_threshold_mv.value(), - margin_frac=self.spin_correlate_margin.value()) + self._sras, self._ref_angle_idx, angles, self._dc4_mv, + sources=self.combo_correlate_source.currentData(), + dc_threshold_mv=self.spin_mask_threshold_mv.value(), + search_deg=self.spin_correlate_search_deg.value()) self._correlate_done_count = 0 self._correlate_total = len(angles) + self._fit_notes = {} self._set_controls_enabled(False) self.lbl_status.setText(f"Cross-correlating: 0/{self._correlate_total} angle(s)…") started = self._parent._run_worker( @@ -1298,8 +1326,10 @@ class ManualAlignmentDialog(QDialog): self.lbl_status.setText("Could not start cross-correlation (busy) — try again.") def _on_correlate_angle_done(self, angle_idx: int, rotation_deg: float, - shift_x_mm: float, shift_y_mm: float): + shift_x_mm: float, shift_y_mm: float, + score: float, source: str): self._angle_params[angle_idx] = ManualAngleParams(rotation_deg, (shift_x_mm, shift_y_mm)) + self._fit_notes[angle_idx] = (score, source) self._correlate_done_count += 1 self.lbl_status.setText( f"Cross-correlating: {self._correlate_done_count}/{self._correlate_total} angle(s)…") @@ -1311,20 +1341,47 @@ class ManualAlignmentDialog(QDialog): self._sync_active_spinboxes() self._rebuild_preview_canvas() self._set_controls_enabled(True) - source = "thresholded mask" if self.combo_correlate_source.currentIndex() == 1 \ - else "raw signal" self.lbl_status.setText( f"Cross-correlated {self._correlate_done_count} angle(s) against " - f"Angle {self._ref_angle_idx} using the {source}. Nudge from here " - "for any remaining fine correction.") + f"Angle {self._ref_angle_idx}.\n" + self._fit_report()) + + def _fit_report(self) -> str: + """Per-angle registration quality, worst first. + + Surfaced rather than buried because a single bad acquisition (stage + glitch, laser dropout) registers poorly and would otherwise be fused in + silently — seeing which angle it is, is what makes dropping it with + sras_edit_scans.py actionable. The deviation from the stage's own + reported angle is shown alongside: a large one means the search and the + stage disagree, which is either a genuine mechanical error or a sign + that this angle's fit is not to be trusted. + """ + if not self._fit_notes: + return "" + rows = sorted(self._fit_notes.items(), key=lambda kv: kv[1][0]) + worst = rows[0] + lines = [f"Worst fit: angle {worst[0]} (score {worst[1][0]:.3f}, " + f"{worst[1][1]})."] + drifted = [] + for a, _note in rows: + nominal = compute._nominal_delta_deg(self._sras, a, self._ref_angle_idx) + got = self._angle_params[a].rotation_deg + dev = min(abs(got - nominal), abs(got + nominal)) + if dev > 1.0: + drifted.append(f"{a} ({dev:.2f}°)") + if drifted: + lines.append("Rotation differs from the stage angle by >1° for " + "angle(s) " + ", ".join(drifted) + ".") + lines.append("Nudge from here for any remaining fine correction.") + return " ".join(lines) def _on_save(self): threshold = self.spin_mask_threshold_mv.value() resolved = dict(self._angle_params) # already concrete floats try: path = save_manual_alignment(self._sras, self._ref_angle_idx, threshold, resolved) - result = build_manual_alignment(self._sras, self._ref_angle_idx, threshold, - resolved, self._pivot_mm) + result = build_manual_alignment(self._sras, self._ref_angle_idx, + threshold, resolved) except OSError as exc: QMessageBox.warning(self, "Save Alignment Failed", str(exc)) return @@ -1348,6 +1405,7 @@ class ManualAlignmentDialog(QDialog): f"Could not delete the saved alignment file: {exc}") return self._angle_params = {a: ManualAngleParams() for a in range(self._sras.n_angles)} + self._fit_notes = {} self._sync_active_spinboxes() self._rebuild_preview_canvas() self.lbl_status.setText( diff --git a/sras_viewer_requirements.txt b/sras_viewer_requirements.txt index 05adf9a..1448342 100644 --- a/sras_viewer_requirements.txt +++ b/sras_viewer_requirements.txt @@ -1,3 +1,11 @@ PyQt6==6.10.2 numpy==2.4.1 matplotlib==3.10.8 +scipy==1.18.0 +# Angle alignment only: masked FFT phase correlation, which registers scans +# whose valid (scanned) regions differ in shape — see sras_compute's +# _masked_shift. +scikit-image==0.26.0 +# Optional: a faster rfft backend for the RF/FFT images (FFT Options -> pyFFTW). +# The viewer falls back to scipy.fft when it is not installed. +pyFFTW==0.15.1 diff --git a/sras_workers.py b/sras_workers.py index 5696b0e..334e655 100644 --- a/sras_workers.py +++ b/sras_workers.py @@ -282,10 +282,10 @@ class BatchCacheWorker(QObject): class AngleAlignmentWorker(QObject): - """Computes rotation+translation alignment for every angle in *sras*, - referenced to *ref_angle_idx*, from each angle's binarized CH4 mask. - Rotation is analytic (from sras.angles_deg); only translation is found by - phase correlation. + """Computes the rigid (rotation + translation, never scale) alignment for + every angle in *sras* against *ref_angle_idx*, by cross-correlating each + angle's CH4 image against the reference's. Both the rotation and the + translation are found from image content — see compute_angle_alignment. """ progress = pyqtSignal(int) # 0–100 finished = pyqtSignal(object, str) # AlignmentResult|None, error ("" = success) @@ -349,52 +349,53 @@ class Ch4MaskWorker(QObject): class CrossCorrelateWorker(QObject): - """FFT phase-correlation translation for each of *angle_indices* against - *ref_angle_idx*, for ManualAlignmentDialog's Auto Cross-Correlate button. + """Rigid registration (rotation + translation, never scale) of each of + *angle_indices* against *ref_angle_idx*, for ManualAlignmentDialog's Auto + Cross-Correlate button. - Runs on a background thread — a real many-angle, high-resolution scan's - correlation (even at its downsampled working resolution) can take long - enough that doing all of them on the GUI thread would visibly freeze the - dialog. Rotation is set to the same analytic scan-angle delta Auto - De-rotate uses alongside the correlated shift, since a translation - search is only meaningful once both angles' content is already oriented - the same way. dc4_mv/pivot_mm are the dialog's own already-in-memory - per-angle images/pivots — this worker does no fetching of its own. + Runs on a background thread — registering a real many-angle, + high-resolution scan takes long enough that doing it on the GUI thread + would visibly freeze the dialog. Rotation is *searched*, not taken from the + stage's reported angle: see compute.register_angle_to_reference, which + seeds from that angle but scores both of its signs and refines from there. + dc4_mv is the dialog's own already-in-memory per-angle CH4 image — this + worker does no fetching of its own. """ - angle_done = pyqtSignal(int, float, float, float) # angle_idx, rotation_deg, shift_x_mm, shift_y_mm + # angle_idx, rotation_deg, shift_x_mm, shift_y_mm, score, source + angle_done = pyqtSignal(int, float, float, float, float, str) finished = pyqtSignal() error = pyqtSignal(str) def __init__(self, sras: SrasFile, ref_angle_idx: int, angle_indices: list[int], - dc4_mv: dict[int, np.ndarray], pivot_mm: dict[int, tuple[float, float]], - *, use_mask: bool, dc_threshold_mv: float, margin_frac: float): + dc4_mv: dict[int, np.ndarray], *, + sources: tuple[str, ...], dc_threshold_mv: float, + search_deg: float): super().__init__() self._sras = sras self._ref = ref_angle_idx self._angles = angle_indices self._dc4_mv = dc4_mv - self._pivot_mm = pivot_mm - self._use_mask = use_mask + self._sources = sources self._threshold = dc_threshold_mv - self._margin = margin_frac + self._search_deg = search_deg - def _one(self, a: int) -> tuple[int, float, float, float]: - theta = compute._theta_deg(self._sras, a, self._ref) - dx, dy = compute.correlate_translation_mm( - self._sras, a, self._ref, self._dc4_mv, self._pivot_mm, - use_mask=self._use_mask, dc_threshold_mv=self._threshold, - margin_frac=self._margin) - return a, theta, dx, dy + def _one(self, a: int) -> tuple[int, compute.RigidFit]: + return a, compute.register_angle_to_reference( + self._sras, a, self._ref, self._dc4_mv, + dc_threshold_mv=self._threshold, sources=self._sources, + search_deg=self._search_deg) def run(self): try: - n_workers, _budget = compute.plan_angle_level(self._sras) + n_workers = compute._registration_workers( + self._sras, compute._DEFAULT_FINE_DIM) pool = ThreadPoolExecutor(max_workers=max(1, n_workers)) try: futures = [pool.submit(self._one, a) for a in self._angles] for fut in as_completed(futures): - a, theta, dx, dy = fut.result() - self.angle_done.emit(a, theta, dx, dy) + a, fit = fut.result() + self.angle_done.emit(a, fit.rotation_deg, fit.shift_mm[0], + fit.shift_mm[1], fit.score, fit.source) finally: pool.shutdown(wait=True) self.finished.emit() diff --git a/tools/make_test_sras.py b/tools/make_test_sras.py index 07386c3..71e93ef 100644 --- a/tools/make_test_sras.py +++ b/tools/make_test_sras.py @@ -125,6 +125,138 @@ def write(path: Path, n_angles: int = 3, seed: int = 0, return meta +# --------------------------------------------------------------------------- +# Rotating-sample scan: one shape, imaged at several known rotations +# --------------------------------------------------------------------------- +# +# The scan the angle-alignment path actually has to solve: every angle images +# the *same* sample at a different known rotation and offset, and a correct +# alignment stacks them all back into one shape. Two properties are +# deliberately hostile: +# +# * every angle gets a different window size and a different, meaningless +# stage x_start / y0 — alignment must ignore per-angle stage coordinates +# entirely, so any code that reads them will visibly fail here; +# * the pixel grid is strongly anisotropic (5 µm along x, 50 µm along y), +# like the real instrument, so any registration that rotates raw indices +# instead of millimetres shears the image and cannot converge. + +_ROT_DX_MM = 0.005 # x pitch, from velocity/laser_freq below +_ROT_DY_MM = 0.05 # row spacing +_ROT_BG_MV = 4.0 +_ROT_FG_MV = 160.0 + + +# How far the sample sits from the rotation axis. Non-zero on purpose: on the +# real instrument every angle's scan window is centred on the rotation axis +# while the sample is not, so each scan sees the sample somewhere else along a +# circle. That offset is exactly what a wrong rotation pivot turns into a ring +# of scans instead of a stack, so a centred test sample would hide the bug. +_ROT_SAMPLE_OFFSET_MM = (0.55, 0.40) + + +def _sample_shape_mv(u: np.ndarray, v: np.ndarray) -> np.ndarray: + """An asymmetric test sample in its own mm frame, chirally distinct at + every rotation (no 180° ambiguity) and with structure at several radii so + rotation is well determined.""" + u = u - _ROT_SAMPLE_OFFSET_MM[0] + v = v - _ROT_SAMPLE_OFFSET_MM[1] + img = np.full(u.shape, _ROT_BG_MV, dtype=np.float32) + img[((u / 0.85) ** 2 + (v / 0.40) ** 2) <= 1.0] = _ROT_FG_MV # bar + img[(np.abs(u - 0.55) <= 0.22) & (np.abs(v - 0.62) <= 0.22)] = _ROT_FG_MV # nub + img[((u + 0.75) ** 2 + (v + 0.30) ** 2) <= 0.20 ** 2] = _ROT_FG_MV # dot + return img + + +def _rot(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]]) + + +def write_rotating(path: Path, n_angles: int = 5, samples_per_frame: int = 4, + seed: int = 0) -> dict: + """Write a v6 file whose CH4 DC image is one sample seen at n_angles known + rotations, and return the ground truth each angle should register to. + + ``truth[a] = (rotation_deg, (shift_x_mm, shift_y_mm))`` is the rigid map + from angle *a*'s local mm (origin at its own array center) to angle 0's — + exactly what ``register_angle_to_reference`` is supposed to recover. + """ + rng = np.random.default_rng(seed) + n_ch, bps = 3, 1 + cal = [(1.5625e-3, -87.04, 0.0), (2.0e-3, -60.0, 1.0e-3), (2.5e-3, -40.0, -2.0e-3)] + ymult_mv, yoff, yzero_mv = cal[2][0] * 1000, cal[2][1], cal[2][2] * 1000 + + stage_angles, geom, x_starts, y_starts, thetas, offsets = [], [], [], [], [], [] + for a in range(n_angles): + stage = -37.0 * a # what the rotation stage reports + stage_angles.append(stage) + # The true image rotation is the negative of the stage's reported + # angle: the stage's positive sense is the opposite of math-positive + # (x toward y) in scan mm. Nothing may depend on knowing that — the + # registration search tries both signs. + thetas.append(-stage) + offsets.append((0.0, 0.0) if a == 0 + else (float(rng.uniform(-0.3, 0.3)), float(rng.uniform(-0.3, 0.3)))) + # A different window per angle, all centred on the same array center — + # the real instrument grows each angle's axis-aligned bounding box to + # cover the rotated ROI. Sized so the off-axis sample stays inside every + # window at every angle, keeping the expected result unambiguous. + geom.append((88 + 8 * a, 780 + 60 * a)) + # Meaningless per-angle stage positions: correct alignment never reads + # them, so scattering them proves it. + x_starts.append(float(20.0 + rng.uniform(-6.0, 6.0))) + y_starts.append(float(30.0 + rng.uniform(-6.0, 6.0))) + + out = bytearray() + out += struct.pack( + HDR_FMT_V6, b"SRAS", 6, n_angles, + x_starts[0], y_starts[0], 1.0, 1.0, _ROT_DY_MM, + _VELOCITY_MM_S, _VELOCITY_MM_S / _ROT_DX_MM, # velocity/freq -> 5 µm pitch + samples_per_frame, _SAMPLE_RATE_HZ, bps, n_ch, + ) + out += np.array(stage_angles, dtype=">f4").tobytes() + for a, (n_rows, n_frames) in enumerate(geom): + out += struct.pack(GEO_FMT_V6, x_starts[a], 1.0, n_frames, n_rows) + for a, (n_rows, _) in enumerate(geom): + out += (y_starts[a] + np.arange(n_rows) * _ROT_DY_MM).astype(">f4").tobytes() + for ymult_v, yoff_a, yzero_v in cal: + p = _preamble(ymult_v, yoff_a, yzero_v) + out += struct.pack(">H", len(p)) + p + background = rng.integers(-8, 9, size=samples_per_frame, dtype=np.int8) + out += struct.pack(">I", samples_per_frame) + background.tobytes() + + truth, dc4_images = {}, [] + for a, (n_rows, n_frames) in enumerate(geom): + # Local mm of every pixel, measured from this angle's own array center. + lx = (np.arange(n_frames) - (n_frames - 1) / 2.0) * _ROT_DX_MM + ly = (np.arange(n_rows) - (n_rows - 1) / 2.0) * _ROT_DY_MM + gx, gy = np.meshgrid(lx, ly) + # local = R(theta) @ sample + offset, so sample = R(theta)^T @ (local - offset) + rel = np.stack([gx - offsets[a][0], gy - offsets[a][1]], axis=-1) + s = rel @ _rot(thetas[a]) # == rel @ R^T.T == R^T @ rel + dc4 = _sample_shape_mv(s[..., 0], s[..., 1]) + dc4_images.append(dc4) + + inv = _rot(-thetas[a]) + truth[a] = (-thetas[a], + tuple(float(v) for v in -(inv @ np.array(offsets[a])))) + + adc4 = np.clip(np.round((dc4 - yzero_mv) / ymult_mv + yoff), -128, 127).astype(np.int8) + block = np.zeros((n_rows, n_ch, n_frames, samples_per_frame), dtype=np.int8) + block[:, 2] = adc4[:, :, None] # CH4 carries the sample + block[:, 1] = 10 # CH3 flat + block[:, 0] = rng.integers(-40, 41, size=(n_rows, n_frames, samples_per_frame), + dtype=np.int8) # CH1 noise + out += block.tobytes() + + path.write_bytes(bytes(out)) + return {"n_angles": n_angles, "geometry": geom, "stage_angles_deg": stage_angles, + "truth": truth, "dc4_mv": dc4_images, "x_starts": x_starts, + "y_starts": y_starts, "dx_mm": _ROT_DX_MM, "dy_mm": _ROT_DY_MM} + + HDR_FMT_LEGACY = ">4sBHHffffIIdBB" diff --git a/tools/test_alignment.py b/tools/test_alignment.py new file mode 100644 index 0000000..f78f9fa --- /dev/null +++ b/tools/test_alignment.py @@ -0,0 +1,223 @@ +#!/usr/bin/env python3 +"""Angle-alignment tests: does registration actually stack the scans? + +Builds a synthetic scan in which one sample is imaged at several *known* +rotations and offsets (tools/make_test_sras.write_rotating) and checks that the +alignment path recovers them, that the shared canvas is angle 0's own pixel +grid extended, and that nothing in the result depends on any other angle's +stage coordinates. + +No Qt — this exercises sras_compute directly. See tools/test_gui.py for the +dialog and Aligned-View plumbing. + +Usage: python tools/test_alignment.py +""" + +import sys +import tempfile +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +import sras_compute as compute # noqa: E402 +from sras_format import CH4_IDX, SrasFile, adc_to_mv # noqa: E402 +import tools.make_test_sras as gen # noqa: E402 + +# Registration is limited by how far a feature moves per degree: with this +# sample's ~1 mm radius and a ~16 µm registration pitch, a quarter degree is +# already sub-pixel, so it is the floor of what any metric can resolve here. +_ROT_TOL_DEG = 0.5 +_SHIFT_TOL_MM = 0.02 +_STACK_IOU_MIN = 0.90 +_THRESHOLD_MV = 80.0 + +_failures: list[str] = [] + + +def check(name: str, ok: bool, detail: str = ""): + print(f" {'PASS' if ok else 'FAIL'} {name}" + (f" — {detail}" if detail else "")) + if not ok: + _failures.append(name) + + +def dc4_images(sras: SrasFile) -> dict[int, np.ndarray]: + return {a: adc_to_mv(compute.compute_dc_image(sras, a, CH4_IDX), *sras.cal(CH4_IDX)) + for a in range(sras.n_angles)} + + +def mm_transform(sras: SrasFile, result, angle_idx: int) -> np.ndarray: + """Recover the pure mm-space rotation from a canvas->raw affine. + + matrix == D @ R^T @ A_out, where A_out and D only carry the canvas and + per-angle pixel pitches; undoing both must leave something orthonormal, or + the transform is smuggling in a scale or a shear. + """ + dx_a, dy_a = compute._pixel_pitch_mm(sras, angle_idx) + A_out = np.array([[0.0, result.canvas_dx_mm], [result.canvas_dy_mm, 0.0]]) + D = np.array([[0.0, 1.0 / dy_a], [1.0 / dx_a, 0.0]]) + return np.linalg.inv(D) @ result.per_angle[angle_idx].matrix @ np.linalg.inv(A_out) + + +def main() -> int: + tmpdir = Path(tempfile.mkdtemp(prefix="sras_align_")) + path = tmpdir / "rotating.sras" + meta = gen.write_rotating(path, n_angles=5) + sras = SrasFile(str(path)) + truth = meta["truth"] + + print(f"\nrotating-sample scan: {sras.n_angles} angles, " + f"shapes {[sras.image_shape(a) for a in range(sras.n_angles)]}") + + print("\nper-angle rigid registration (rotation + translation, no scale)") + dc4 = dc4_images(sras) + fits = {a: compute.register_angle_to_reference( + sras, a, 0, dc4, dc_threshold_mv=_THRESHOLD_MV) + for a in range(sras.n_angles)} + for a, fit in fits.items(): + t_rot, t_shift = truth[a] + rot_err = abs(fit.rotation_deg - t_rot) + shift_err = float(np.hypot(fit.shift_mm[0] - t_shift[0], + fit.shift_mm[1] - t_shift[1])) + check(f"angle {a} rotation within {_ROT_TOL_DEG}° of truth", + rot_err <= _ROT_TOL_DEG, + f"got {fit.rotation_deg:.3f}°, truth {t_rot:.3f}° (err {rot_err:.3f}°)") + check(f"angle {a} translation within {_SHIFT_TOL_MM} mm of truth", + shift_err <= _SHIFT_TOL_MM, f"err {shift_err:.4f} mm") + check("reference angle registers as exact identity", + fits[0] == compute.RigidFit(0.0, (0.0, 0.0), 1.0, "reference")) + + # The stage's rotational sense relative to this module's math-positive + # convention is not knowable from the file, and the old code hardcoded a + # guess. Flipping every reported angle must therefore change nothing: the + # search scores both signs and the images decide. + flipped = SrasFile(str(path)) + flipped.angles_deg = -flipped.angles_deg + flipped_fits = {a: compute.register_angle_to_reference( + flipped, a, 0, dc4, dc_threshold_mv=_THRESHOLD_MV) + for a in range(1, flipped.n_angles)} + check("negating every reported stage angle changes no fit", + all(flipped_fits[a] == fits[a] for a in flipped_fits), + str({a: (flipped_fits[a].rotation_deg, fits[a].rotation_deg) + for a in flipped_fits if flipped_fits[a] != fits[a]})) + + print("\nper-angle stage coordinates are not consulted") + # Move every non-reference angle's scan window somewhere else entirely. + # Only angle 0's coordinates may matter, so every fit must be untouched. + moved = SrasFile(str(path)) + for a in range(1, moved.n_angles): + moved.x_start_mm[a] += 13.5 * a + moved._y_pos_per_angle[a] = moved._y_pos_per_angle[a] - 9.25 * a + moved_dc4 = dc4_images(moved) + moved_fits = {a: compute.register_angle_to_reference( + moved, a, 0, moved_dc4, dc_threshold_mv=_THRESHOLD_MV) + for a in range(1, moved.n_angles)} + check("relocating every other angle's scan window changes no fit", + all(moved_fits[a] == fits[a] for a in moved_fits), + str({a: (round(moved_fits[a].rotation_deg, 4), fits[a].rotation_deg) + for a in moved_fits if moved_fits[a] != fits[a]})) + + print("\nshared canvas is angle 0's own pixel grid, extended") + result = compute.compute_angle_alignment(sras, 0, _THRESHOLD_MV) + t0 = result.per_angle[0] + check("angle 0's transform has no rotation, scale or shear", + np.allclose(t0.matrix, np.eye(2)), str(t0.matrix)) + check("angle 0 lands on whole canvas pixels (no resampling of the reference)", + np.allclose(t0.offset, np.round(t0.offset)), str(t0.offset)) + check("canvas pitch is angle 0's own pitch", + (result.canvas_dx_mm, result.canvas_dy_mm) + == compute._pixel_pitch_mm(sras, 0)) + + n_rows, n_cols = result.canvas_shape + x_axis = result.canvas_origin_mm[0] + np.arange(n_cols) * result.canvas_dx_mm + y_axis = result.canvas_origin_mm[1] + np.arange(n_rows) * result.canvas_dy_mm + row0, col0 = int(round(-t0.offset[0])), int(round(-t0.offset[1])) + a0_rows, a0_cols = sras.image_shape(0) + check("canvas X axis reproduces angle 0's own X coordinates", + np.allclose(x_axis[col0:col0 + a0_cols], sras.x_axis_mm(0))) + check("canvas Y axis reproduces angle 0's own Y coordinates", + np.allclose(y_axis[row0:row0 + a0_rows], sras.y_positions_mm(0))) + check("canvas covers every angle's footprint", + n_rows >= max(int(sras.n_rows[a]) for a in range(sras.n_angles)) + and n_cols >= max(int(sras.n_frames[a]) for a in range(sras.n_angles)), + str(result.canvas_shape)) + + print("\nno scaling anywhere in the per-angle transforms") + for a in range(sras.n_angles): + R = mm_transform(sras, result, a) + check(f"angle {a}'s mm-space transform is a pure rotation", + np.allclose(R @ R.T, np.eye(2), atol=1e-9) + and abs(abs(np.linalg.det(R)) - 1.0) < 1e-9, + f"det={np.linalg.det(R):.6f}") + + print("\nall angles stack into one shape") + aligned = {a: compute.apply_alignment(result, a, dc4[a]) + for a in range(sras.n_angles)} + base = aligned[0] >= _THRESHOLD_MV + for a in range(1, sras.n_angles): + other = aligned[a] >= _THRESHOLD_MV + iou = float((base & other).sum()) / max(1, int((base | other).sum())) + check(f"angle {a}'s aligned sample overlaps angle 0's (IoU >= {_STACK_IOU_MIN})", + iou >= _STACK_IOU_MIN, f"IoU {iou:.4f}") + + print("\ndownsampled preview lands where the full-resolution image does") + # ManualAlignmentDialog reprojects block-mean-downsampled masks, so the + # affine has to account for the factor. When it did not, every preview + # layer came out magnified by that factor and offset — the overlay showed a + # blown-up crop of each mask, which is not something you can align by eye. + pitch = (result.canvas_dx_mm, result.canvas_dy_mm) + a = sras.n_angles - 1 + p = result.per_angle[a] + full_mask = (dc4[a] >= _THRESHOLD_MV).astype(np.float32) + full = compute.reproject_mask( + sras, a, 0, full_mask, p.rotation_deg, p.shift_mm, pitch, + result.canvas_origin_mm, result.canvas_shape) + fy, fx = 4, 16 + small = compute.reproject_mask( + sras, a, 0, compute._block_mean_2d(full_mask, fy, fx), + p.rotation_deg, p.shift_mm, (pitch[0] * fx, pitch[1] * fy), + result.canvas_origin_mm, + (result.canvas_shape[0] // fy, result.canvas_shape[1] // fx), + src_downsample=(fy, fx)) + # Compare in mm, via each layer's own center of mass. + def com_mm(layer, px, py): + rows, cols = np.nonzero(layer > 0.5) + return np.array([cols.mean() * px, rows.mean() * py]) + d = com_mm(small, pitch[0] * fx, pitch[1] * fy) - com_mm(full, *pitch) + check("a downsampled preview layer lands within a preview pixel of the " + "full-resolution one", + abs(d[0]) <= abs(pitch[0] * fx) and abs(d[1]) <= abs(pitch[1] * fy), + f"offset {d[0]:+.4f}, {d[1]:+.4f} mm") + + print("\nmanual path reproduces the same geometry") + params = {a: compute.ManualAngleParams(t.rotation_deg, t.shift_mm) + for a, t in result.per_angle.items()} + manual = compute.build_manual_alignment(sras, 0, _THRESHOLD_MV, params) + check("build_manual_alignment matches compute_angle_alignment for the same params", + manual.canvas_shape == result.canvas_shape + and np.allclose(manual.canvas_origin_mm, result.canvas_origin_mm) + and all(np.allclose(manual.per_angle[a].matrix, result.per_angle[a].matrix) + and np.allclose(manual.per_angle[a].offset, result.per_angle[a].offset) + for a in range(sras.n_angles))) + + print("\nsidecar round-trip") + compute.save_manual_alignment(sras, 0, _THRESHOLD_MV, params) + loaded = compute.load_manual_alignment(sras) + check("sidecar reloads every angle's params", + loaded is not None + and all(np.isclose(loaded.per_angle[a].rotation_deg, params[a].rotation_deg) + and np.allclose(loaded.per_angle[a].shift_mm, params[a].shift_mm) + for a in range(sras.n_angles))) + check("sidecar deletes cleanly", compute.delete_manual_alignment(sras)) + + print() + if _failures: + print(f"{len(_failures)} FAILURE(S): " + ", ".join(_failures)) + return 1 + print("All alignment checks passed.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/test_gui.py b/tools/test_gui.py index 482e11a..36bca72 100644 --- a/tools/test_gui.py +++ b/tools/test_gui.py @@ -27,7 +27,7 @@ from PyQt6.QtWidgets import QApplication, QMessageBox # noqa: E402 sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) import sras_compute as compute # noqa: E402 -from sras_format import CH1_IDX, CH3_IDX, CH4_IDX # noqa: E402 +from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, SrasFile # noqa: E402 from sras_viewer import RoiQuad, SrasViewerWindow, VELOCITY_MODE_IDX # noqa: E402 import tools.make_test_sras as gen # noqa: E402 @@ -240,60 +240,49 @@ def main(): print("\nmanual alignment (Fusion)") check("manual alignment action enabled", win._manual_align_act.isEnabled()) - # --- Alignment pivot is a signal-weighted centroid, not the raw bbox -- - # center, and is independent of any DC threshold (so a threshold that - # happens to leave a real angle's binary mask empty can't silently - # degrade the pivot back to the bbox center). - corner_signal = np.zeros(s.image_shape(0), dtype=np.float32) - corner_signal[0, 0] = 1.0 # single spike -> weighted centroid is exact - expected_corner = (float(s.x_axis_mm(0)[0]), float(s.y_positions_mm(0)[0])) - centroid = compute._signal_centroid_mm(s, 0, corner_signal) - check("signal-weighted centroid of a single spike pixel is that pixel exactly", - np.allclose(centroid, expected_corner), f"{centroid} vs {expected_corner}") - bbox_center = compute._bbox_center_mm(s, 0) - check("signal centroid differs from the raw scan-window bbox center", - not np.allclose(centroid, bbox_center), - f"centroid {centroid} vs bbox center {bbox_center}") + # --- Local mm is anchored on each angle's array center, not its stage -- + # position: that is what makes a scan's placement independent of where its + # window happened to sit. (Registration accuracy itself is covered by + # tools/test_alignment.py, which has a synthetic sample to register.) + n_rows, n_frames = s.image_shape(0) + check("array center is the geometric center of the pixel grid", + np.allclose(compute._center_idx(s, 0), + [(n_rows - 1) / 2, (n_frames - 1) / 2])) + dx0, dy0 = compute._pixel_pitch_mm(s, 0) + check("local half-extent is derived from shape and pitch alone", + np.allclose(compute._local_half_extent_mm(s, 0), + [(n_frames - 1) / 2 * abs(dx0), (n_rows - 1) / 2 * abs(dy0)])) + identity = {a: compute.ManualAngleParams() for a in range(s.n_angles)} + origin_a, shape_a = compute.canvas_for_params(s, 0, (dx0, dy0), identity) + moved = SrasFile(str(path)) + for a in range(1, moved.n_angles): + moved.x_start_mm[a] += 7.5 + moved._y_pos_per_angle[a] = moved._y_pos_per_angle[a] + 3.25 + origin_b, shape_b = compute.canvas_for_params(moved, 0, (dx0, dy0), identity) + check("moving every non-reference angle's scan window leaves the canvas " + "unchanged (only angle 0's coordinates are used)", + shape_a == shape_b and np.allclose(origin_a, origin_b), + f"{origin_a} {shape_a} vs {origin_b} {shape_b}") - # compute_pivot_points_mm should reuse a pre-computed dc4_mv dict rather - # than recomputing from the real DC4 image (which has no such spike and - # would give a different answer if silently recomputed). - reused_pivot = compute.compute_pivot_points_mm(s, dc4_mv={0: corner_signal})[0] - check("compute_pivot_points_mm reuses a pre-computed dc4_mv dict", - np.allclose(reused_pivot, expected_corner)) + # --- Both signs of the stage's reported angle are searched -------------- + cands = compute._rotation_candidates(30.0, 6.0, 2.0) + check("rotation candidates bracket both signs of the stage angle", + min(cands) < -29.0 and max(cands) > 29.0, f"{min(cands)}..{max(cands)}") - # A perfectly flat signal carries no information to weight by, so it - # falls back to the bbox center rather than producing a NaN/degenerate - # centroid. - flat_signal = np.full(s.image_shape(0), 5.0, dtype=np.float32) - flat_centroid = compute._signal_centroid_mm(s, 0, flat_signal) - check("a perfectly flat signal falls back to the bbox center", - np.allclose(flat_centroid, bbox_center)) - - # --- Rotation sign convention: negative of the raw angles_deg delta ---- - check("_theta_deg negates the raw angles_deg delta (GR stage's positive " - "angle is the opposite rotational sense from this module's CCW " - "math convention)", - all(np.isclose(compute._theta_deg(s, a, 0), - -(float(s.angles_deg[a]) - float(s.angles_deg[0]))) - for a in range(s.n_angles))) - - # --- FFT phase correlation recovers a known synthetic pixel shift ------ - rng = np.random.default_rng(0) - corr_ref = np.zeros((40, 50), dtype=np.float32) - corr_ref[10:25, 15:35] = 1.0 - corr_ref += 0.05 * rng.standard_normal(corr_ref.shape).astype(np.float32) - corr_mov = np.roll(corr_ref, shift=(4, -7), axis=(0, 1)) - dr, dc = compute._phase_correlate_shift(corr_ref, corr_mov) - check("phase correlation recovers the shift that aligns mov onto ref", - (dr, dc) == (-4, 7), f"got (dr, dc)={(dr, dc)}") + # --- Whole-pixel translation must not wrap content around the edge ------ + arr = np.zeros((6, 6), dtype=np.float32) + arr[0, 0] = 1.0 + check("_shift_into zero-fills rather than wrapping", + compute._shift_into(arr, -1, -1).sum() == 0.0) + check("_shift_into moves content by exactly the requested offset", + compute._shift_into(arr, 2, 3)[2, 3] == 1.0) # --- Open: must NOT seed from the still-live automatic AlignmentResult -- - # The automatic result's translation comes from FFT phase correlation -- - # the very thing manual mode exists to work around -- so manual mode - # must start from identity (centroids coincide, zero shift) regardless - # of whatever the automatic run last computed. Only a previously *saved - # manual* alignment (sidecar) should ever seed this dialog. + # Manual mode exists to fix up whatever the automatic registration got + # wrong, so it must start from identity (every angle centered on the + # reference, no rotation) regardless of whatever the automatic run last + # computed. Only a previously *saved manual* alignment (sidecar) should + # ever seed this dialog. win._on_manual_alignment() check("dialog opened", win._manual_align_dialog is not None) dlg = win._manual_align_dialog @@ -343,39 +332,41 @@ def main(): check("a real Right-arrow key event nudged shift_x", dlg._angle_params[active].shift_mm[0] > before[0]) - # --- Auto De-rotate: rotation only, translation untouched --------------- + # --- Auto De-rotate: seeds rotation from the stage angle, no translation - shift_before_derotate = dlg._angle_params[active].shift_mm dlg._on_auto_derotate() - expected_theta = compute._theta_deg(s, active, dlg._ref_angle_idx) - check("auto de-rotate set the known analytic angle", - abs(dlg._angle_params[active].rotation_deg - expected_theta) < 1e-6) + nominal = compute._nominal_delta_deg(s, active, dlg._ref_angle_idx) + check("auto de-rotate seeded rotation from the stage's reported angle", + abs(dlg._angle_params[active].rotation_deg - nominal) < 1e-6) check("auto de-rotate left translation untouched", dlg._angle_params[active].shift_mm == shift_before_derotate) check("reference angle stays identity after auto de-rotate", dlg._angle_params[dlg._ref_angle_idx].rotation_deg == 0.0) + # Clicking again offers the other sign, since which one lines the scans up + # is not knowable from the file. + dlg._on_auto_derotate() + check("auto de-rotate offers the opposite sign on a second click", + abs(dlg._angle_params[active].rotation_deg + nominal) < 1e-6) - # --- Auto Cross-Correlate: rotation + FFT-correlated shift, backgrounded - + # --- Auto Cross-Correlate: searches rotation *and* translation ---------- check("cross-correlate action enabled once masks are ready", dlg.btn_auto_correlate.isEnabled()) - dlg._on_auto_correlate() - check("auto cross-correlate completed", wait_until( - lambda: not win._job_running("manual_align_correlate"), timeout_ms=30000)) - check("auto cross-correlate set the known analytic angle for every angle", - all(abs(dlg._angle_params[a].rotation_deg - - compute._theta_deg(s, a, dlg._ref_angle_idx)) < 1e-6 - for a in range(s.n_angles) if a != dlg._ref_angle_idx)) + for label_idx, (label, _sources) in enumerate(dlg._CORRELATE_SOURCES): + dlg.combo_correlate_source.setCurrentIndex(label_idx) + dlg._on_auto_correlate() + check(f"auto cross-correlate completed ({label})", wait_until( + lambda: not win._job_running("manual_align_correlate"), timeout_ms=60000)) + check(f"every non-reference angle got a fit ({label})", + all(a in dlg._fit_notes for a in range(s.n_angles) + if a != dlg._ref_angle_idx)) check("auto cross-correlate reference angle stays identity", dlg._angle_params[dlg._ref_angle_idx] == compute.ManualAngleParams()) check("auto cross-correlate re-enabled controls when done", dlg.grp_correlate.isEnabled() and dlg.btn_save.isEnabled()) check("preview canvas rebuilt after cross-correlate", len(dlg._preview_layers) == s.n_angles) - - # The thresholded-mask option should also work end to end. - dlg.combo_correlate_source.setCurrentIndex(1) # thresholded mask - dlg._on_auto_correlate() - check("auto cross-correlate (thresholded-mask option) completed", wait_until( - lambda: not win._job_running("manual_align_correlate"), timeout_ms=30000)) + check("fit quality is reported per angle", bool(dlg._fit_report()), + dlg._fit_report()) # --- Save ----------------------------------------------------------------- dlg._on_save()