diff --git a/docs/design.md b/docs/design.md index 9c1ef27..74cec99 100644 --- a/docs/design.md +++ b/docs/design.md @@ -171,6 +171,140 @@ for the same reason, and every affine 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. +### Cropping the canvas is index translation, not a second transform + +`crop_alignment_result` restricts an `AlignmentResult` to a rectangular window +of its canvas by folding the crop into each angle's existing affine rather than +composing a new one. From `_affine_out_to_src`, `matrix = D @ Rinv @ A_out` +depends only on the pitches and the rotation, and `A_out @ [row0, col0]` is +exactly the mm displacement of the new origin, so + +``` +matrix @ [r', c'] + (offset + matrix @ [row0, col0]) + == matrix @ [r' + row0, c' + col0] + offset +``` + +identically. `matrix` is untouched and `offset` — which already absorbs the +origin — absorbs the crop too. + +Two things follow, and both are relied on. `apply_alignment`, `reproject_mask` +and the aligned exporter all work on a cropped result with no special-casing: +resampling a cropped result is *exactly* a slice of resampling the full one +(`tests/test_align_export.py::test_crop_is_a_window_of_the_full_canvas` asserts +bit equality). And because the crop offset is a whole number of canvas pixels, +`canvas_for_params`' snap invariant — the reference angle lands on integer +canvas pixels — survives the crop, which is what keeps the reference exportable +as a verbatim block. + +## Aligned export (`sras_align_export.py`) + +`write_aligned_sras` bakes an alignment into a new v6 file: every angle +resampled onto the cropped shared canvas, so all of them end up with identical +geometry and the file opens already aligned. It is the only place in the +codebase that *resamples* waveform data — `sras_edit_scans` and `sras_average` +copy waveform bytes verbatim — which is why it is its own top-level module +rather than part of `sras_format` (scoped to the versioned binary spec, per the +sidecar section's own rule) or `sras_compute` (imported by every +multiprocessing child). + +**Nearest neighbour, never interpolation.** Each output pixel gets exactly one +source pixel's three waveforms, verbatim. Averaging two neighbouring CH1 +packets would synthesise a waveform the instrument never measured, whose FFT +peak is the peak of neither — meaningless for a technique whose entire output is +that peak frequency. The cost is that some source pixels are duplicated and +others dropped, which is the same trade `apply_alignment`'s `order=0` already +makes for the display. + +**The rounding rule is `floor(x + 0.5)`, not `np.rint`.** `scipy.ndimage`'s +`order=0` rounds halves away from zero while `np.rint` rounds them to even. The +canvas is snapped to the reference's own pixel grid, so an angle whose row pitch +differs from the reference's lands on exact half-integers across whole rows — +this is the common case, not a corner case. Getting it wrong shifts those rows +by one source pixel relative to what the Aligned View drew. + +**Out-of-bounds is tested on the fractional coordinate, not the rounded index.** +`scipy`'s `mode="constant"` writes `cval` wherever the coordinate leaves the +range of sample *centres*, `[0, n-1]` — a coordinate of −0.4 rounds to a +perfectly valid index 0 and is still padding. Testing the rounded index instead +puts a one-pixel rim of real data everywhere the preview shows padding. + +**...but with a tolerance (`_EDGE_TOL`).** The affine is built from a chain of +mm-space multiplications, so an exactly-integer transform comes out a few times +1e-13 off: the reference angle's offset is `-20 - 7e-15`, not `-20`. A bare +`>= 0.0` therefore rejects that angle's entire first row, and `<= n-1` its last +column — for the *reference* angle, whose whole job is to pass through as an +exact integer crop. The tolerance is ~7 orders of magnitude above that noise and +~7 below the half-pixel scale at which a rounding decision means anything, so it +can only ever change pixels whose scipy answer was itself decided by noise. + +**Padding is the per-channel ADC code nearest 0 mV, not 0.** Zero ADC decodes to +`(0 - yoff) * ymult + yzero`, which on real calibration is around +100 mV — +above any sensible CH4 mask threshold, so a zero fill would paint a solid +rectangle of "valid" pixels around the sample and corrupt every DC image and ROI +statistic downstream. + +**Source rows are served from sliding in-RAM bands** (`_SourceReader`). A +rotated angle maps one output row to a *diagonal* across the source array, so +the pixels of a single output row come from hundreds of different source rows — +~1.4 MB each on a full-size scan. Indexing a memmap pixel-by-pixel in output +order re-faults nearly the whole angle per output row: terabytes of paging for a +gigabyte of data. Reading a contiguous band per output chunk, with the band +advancing monotonically, costs roughly 2× the source size in total reads. + +**Writes go to `.part` and are `os.replace`d into position.** Not politeness: a +truncated .sras is not detectably broken, because `_parse_v6` drops incomplete +trailing angle blocks and opens what is left as an aborted scan. A half-written +export left in place would silently look like a real file with fewer angles. + +**The Angle Table is carried over unchanged.** Alignment removes the *spatial* +rotation of the sample; it does not change which acoustic propagation direction +each angle measured, and that direction is the scientific content of a +multi-angle scan. Zeroing the table would make the export self-consistent for +re-registration and useless for anisotropy work. The consequence is that +re-registering an export needs `seed_deg=0.0` to put 0° inside the coarse sweep, +since `nominal_delta_deg` is still non-zero — which is exactly what the seed +parameter exists for. + +## Alignment wizard (`sras_viewer/align_wizard.py`) + +A `QWizard` rather than another dialog because the three steps are genuinely +sequential and the last one is destructive: correlate, choose a crop, write a +file. It replaces both former Fusion actions, so it also absorbs the old +`ManualAlignmentDialog`'s by-eye nudge editor — otherwise a scan the search +cannot fit would have no fallback at all. + +Shared state lives on the wizard object, not in `registerField`: the pages pass +numpy arrays, `ManualAngleParams` and an `AlignmentResult` between them, none of +which are scalar widget properties. + +`IndependentPages` is deliberately left **off**. With it set Qt never calls +`cleanupPage`, and `cleanupPage` is how the ROI page discards a crop when the +user goes back to re-correlate — a crop is indexed in canvas pixels, and a new +rotation means a different canvas, so stale indices would silently be +reinterpreted against the wrong grid. `geometry_generation` is the belt-and- +braces check for the same hazard. + +The mask-stack preview shares the **final** canvas's origin and uses a pitch +that is an integer multiple of it, unlike the old manual dialog's padded, +unsnapped preview canvas. That is what lets the crop page convert a rectangle +drawn in millimetres into an exact integer window of the real canvas, with no +second coordinate frame to reconcile. + +"Fit to full overlap" uses `largest_rect_at_least`, a largest-rectangle sweep, +not a bounding box of the fully-covered pixels. The full-overlap region of +several rotated scans is roughly a disc, and its bounding box has corners no +angle covers — offering that as the crop would hand the user the padding they +were trying to avoid. + +Every background launch follows the two rules `_run_worker`'s docstring +establishes: disable the trigger *before* the call (so a re-entrant click cannot +start a second thread over the first), and never ignore the returned bool. +Progress is an inline `QProgressBar` on the page rather than a `QProgressDialog` +— a window-modal popup over a wizard both looks wrong and reintroduces the +event-loop pumping hazard that ordering exists to avoid. `reject()` refuses to +close while a job is in flight, since the running worker's signals are connected +to bound methods of the pages Qt would be deleting. + ## Manual-alignment sidecar (`sras_compute.py`) `.sras.align.json` lives next to the scan file. The code lives in diff --git a/pyproject.toml b/pyproject.toml index edad96f..9f7942e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,7 @@ py-modules = [ "sras_format", "sras_compute", "sras_workers", + "sras_align_export", "sras_average", "sras_edit_scans", ] diff --git a/scan_format.md b/scan_format.md index 2194f93..b94c0e1 100644 --- a/scan_format.md +++ b/scan_format.md @@ -173,6 +173,41 @@ sum over angles a of: n_rows[a] × 3 × n_frames[a] × samples_per_frame × byte > Table and check `file_size` against the running total before reshaping — > a fixed `(n_angles, n_rows, ...)` reshape (as in pre-v6 readers) will not > work since row/frame counts are no longer uniform across angles. +> +> A consequence worth stating explicitly: because a short file opens +> *successfully* as a scan with fewer angles, a truncated file is not +> detectably broken. Anything that writes a .sras must therefore stage to a +> temporary name and rename on success — the viewer's aligned export writes +> `.part` and `os.replace`s it — or an interrupted write leaves behind +> something that loads without complaint and silently has the wrong angle count. + +--- + +## Files written by the viewer's Alignment Wizard + +The acquisition app is not the only producer of this format. The viewer's +`Fusion → Alignment Wizard…` writes a **v6** file holding the aligned, cropped +stack, with these properties: + +* Every angle shares one grid — the cropped alignment canvas — so the + Per-Angle Geometry Table is `n_angles` identical records and the ragged Row + Table is `n_angles` identical spans. The raggedness v6 exists for is still + *expressible*, just unused, so any v6 reader works unchanged. +* `x_delta` is the reference angle's own pitch, which is exactly + `velocity_mm_s / laser_freq_hz`, so the derived X axis stays consistent with + the header. +* The `*_nominal` header fields describe the crop. Uniquely for these files they + coincide with the actual per-angle geometry, since after alignment every angle + really does scan the same box. +* The **Angle Table is unchanged**. Alignment removes the sample's spatial + rotation, not the acoustic propagation direction each angle measured — that + direction is the point of a multi-angle scan, so it is preserved. +* Output pixels with no corresponding source pixel (the canvas corners a rotated + scan cannot reach) hold the per-channel ADC code nearest **0 mV**, not zero. + Zero ADC decodes to roughly +100 mV on real calibration and would read as + signal. +* No Cache Tail is written: any cached DC/FFT is indexed by the source's grid + and would be meaningless on the new one. --- diff --git a/sras_align_export.py b/sras_align_export.py index ebba446..18a83a4 100644 --- a/sras_align_export.py +++ b/sras_align_export.py @@ -292,6 +292,14 @@ class _SourceReader: Small angles skip the machinery — if the whole block fits the budget it is materialized once and every band is a view of it. + + One honest caveat: a chunk whose diagonal spans more source rows than the + budget allows still gets the band it asked for, so the budget can be + overshot. The overshoot is bounded by the span of _ROW_CHUNK output rows, + and in the worst case (an extreme rotation on a huge scan) that is the whole + angle — i.e. no worse than the in-RAM path above. Accepted deliberately: + correctness of the gather is not negotiable, and the alternative is the + memmap thrashing this class exists to avoid. """ def __init__(self, sras: SrasFile, angle_idx: int, budget: int): diff --git a/sras_compute.py b/sras_compute.py index 26a8600..37e1696 100644 --- a/sras_compute.py +++ b/sras_compute.py @@ -1466,7 +1466,7 @@ def canvas_for_params(sras: SrasFile, ref_angle_idx: int, 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 + the wizard'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). @@ -1694,7 +1694,7 @@ def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int, """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 + Meant for a background thread — deliberately recomputes CH4 DC images from scratch rather than reading the GUI-thread _dc_cache dict, since background-thread workers must not touch GUI-thread-owned caches. @@ -1757,7 +1757,7 @@ def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int, # 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 +# only the wizard's own downsampled mask stack calls that per keystroke # — see that class's docstring for how it limits each nudge to reprojecting only # the actively-edited angle. # --------------------------------------------------------------------------- @@ -1770,9 +1770,8 @@ def reproject_mask(sras: SrasFile, angle_idx: int, ref_angle_idx: int, canvas_shape: tuple[int, int], *, src_downsample: tuple[int, int] = (1, 1)) -> np.ndarray: """Resample one angle's binary/float mask onto an arbitrary canvas via an - explicit rotation+shift — the single building block ManualAlignmentDialog's - live preview repeatedly calls (once per keystroke, for only the - actively-nudged angle). *src_downsample* must match the (rows, cols) + explicit rotation+shift — the single building block the alignment wizard's + live mask stack repeatedly calls (once per angle per edit). *src_downsample* must match the (rows, cols) block-mean factor already applied to *mask*, or the reprojection lands at the wrong scale. order=0 (nearest) matches apply_alignment's own reasoning: a binary mask must never be blended with zero-padding.""" @@ -1885,7 +1884,7 @@ def delete_manual_alignment(sras: SrasFile) -> bool: """Delete the sidecar if present. Returns whether a file actually existed to delete, so Clear Alignment's status message can say so. Genuine I/O errors (permission denied, read-only share) propagate — the caller - (ManualAlignmentDialog._on_clear) surfaces them rather than silently + (the wizard's Clear path) surfaces them rather than silently pretending the destructive action succeeded.""" path = sidecar_path(sras.path) try: diff --git a/sras_edit_scans.py b/sras_edit_scans.py index 040e0d5..a906480 100644 --- a/sras_edit_scans.py +++ b/sras_edit_scans.py @@ -14,6 +14,11 @@ 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. +This tool only ever *drops* angles — every kept angle's waveform bytes and +geometry are carried across verbatim. To write a file whose angles have been +resampled onto one shared aligned grid and cropped, use the viewer's +Fusion -> Alignment Wizard (sras_align_export.py) instead. + Usage: python sras_edit_scans.py input.sras --list python sras_edit_scans.py input.sras output.sras --drop 2,5 diff --git a/sras_viewer/__init__.py b/sras_viewer/__init__.py index 2f87856..4014385 100644 --- a/sras_viewer/__init__.py +++ b/sras_viewer/__init__.py @@ -18,9 +18,10 @@ import faulthandler faulthandler.enable() # print a native stack trace on SIGSEGV/SIGABRT/etc. -from .canvases import ImageCanvas, RoiQuad, WaveformCanvas # noqa: E402,F401 -from .common import CH_LABELS, CMAPS, VELOCITY_MODE_IDX # noqa: E402,F401 -from .dialogs import ( # noqa: E402,F401 - FftOptionsDialog, ManualAlignmentDialog, RowAverageFftOptionsDialog, +from .align_wizard import AlignmentWizard # noqa: E402,F401 +from .canvases import ( # noqa: E402,F401 + AlignOverlayCanvas, ImageCanvas, RoiQuad, WaveformCanvas, ) +from .common import CH_LABELS, CMAPS, VELOCITY_MODE_IDX # noqa: E402,F401 +from .dialogs import FftOptionsDialog, RowAverageFftOptionsDialog # noqa: E402,F401 from .main_window import SrasViewerWindow, main # noqa: E402,F401 diff --git a/sras_viewer/align_wizard.py b/sras_viewer/align_wizard.py new file mode 100644 index 0000000..1128674 --- /dev/null +++ b/sras_viewer/align_wizard.py @@ -0,0 +1,1411 @@ +"""The alignment wizard: correlate, crop, export. + +Replaces what used to be two disconnected menu actions. `Angle Alignment` ran +the registration with every parameter hard-coded and no way to retry, and +`Manual Alignment…` was a separate dialog that deliberately refused to inherit +the automatic result — so a bad fit meant starting over by hand, and neither +path told you whether the alignment was actually any good. + +The three steps are the three decisions: + + 1. Correlate. Every parameter that affects the fit is on the page, the run is + repeatable, and the verdict is a picture: each angle's DC mask reprojected + onto the shared canvas and summed, so the colour of a pixel is *how many + angles cover it*. A good alignment is one saturated plateau at N; a bad one + is a fan of low-count halos. Angles start pre-rotated by the stage angles + stored in the file, so the page is informative before any correlation runs. + Manual nudging lives here too, since this page is now the only place to + correct an angle the search cannot fit. + 2. Crop. Axis-aligned only, because that is the only crop .sras geometry can + express — a free quadrilateral would have to be squared off behind the + user's back. + 3. Export. Writes the aligned, cropped data to a new .sras. + +State lives on the wizard rather than in QWizardPage.registerField: the pages +share numpy arrays, ManualAngleParams and an AlignmentResult, none of which are +scalar widget properties. +""" + +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING + +import matplotlib as mpl +import numpy as np +from matplotlib.backends.backend_qtagg import NavigationToolbar2QT +from PyQt6.QtCore import QSignalBlocker, Qt, pyqtSignal +from PyQt6.QtWidgets import ( + QCheckBox, QComboBox, QFileDialog, QHBoxLayout, QHeaderView, QLabel, + QLineEdit, QMessageBox, QProgressBar, QPushButton, QSpinBox, + QTableWidget, QTableWidgetItem, QVBoxLayout, QWidget, QWizard, QWizardPage, +) + +import sras_align_export as export +import sras_compute as compute +from sras_compute import ManualAngleParams, build_manual_alignment +from sras_format import SrasFile +from sras_workers import AlignedExportWorker, Ch4MaskWorker, CrossCorrelateWorker + +from .canvases import AlignOverlayCanvas, ImageCanvas, RoiQuad, count_colormap +from .common import ( + _CSS_HINT, _CSS_INFO, _CSS_MUTED, _CSS_WARN, Jobs, _axes_extent, _form, + _group, _make_dspin, _scroll_panel, _wrap_label, +) + +if TYPE_CHECKING: + from .main_window import SrasViewerWindow + +# Longest side of the coarsened preview the mask stack is built on. Same +# reasoning as the old manual dialog: a full-resolution reprojection per angle +# per edit is far more detail than an alignment can be judged by eye at. +_MAX_PREVIEW_DIM = 1024 + +# Alpha for the per-angle colour view (the active angle sits on top, brighter). +_BASE_ALPHA = 0.42 +_ACTIVE_ALPHA = 0.75 + +_PANEL_W = 372 + +# (label, sources for register_angle_to_reference) — "both" costs roughly double +# but removes the failure mode where the one chosen source happens to be +# uninformative for a single angle. +_CORRELATE_SOURCES = ( + ("Both, keep best", ("signal", "mask")), + ("Raw signal", ("signal",)), + ("Thresholded mask", ("mask",)), +) + +# (label, seed_deg override or None for "the file's stage angle", signs to +# sweep, sign used for the pre-rotation *preview*). The stage's rotational sense +# relative to this module's math-positive convention is not knowable from the +# file, which is why "both" exists and is the default; the explicit signs are +# for a user who has established which way their own stage turns. +_PREROTATE_MODES = ( + ("Stage angle, both signs", None, (-1, 1), 1), + ("Stage angle, + delta only", None, (1,), 1), + ("Stage angle, − delta only", None, (-1,), -1), + ("No pre-rotation", 0.0, (1,), 0), +) + + +@dataclass +class WizardState: + """Everything the three pages share.""" + dc4_mv: dict[int, np.ndarray] = field(default_factory=dict) + masks_small: dict[int, np.ndarray] = field(default_factory=dict) + downsample: tuple[int, int] = (1, 1) + params: dict[int, ManualAngleParams] = field(default_factory=dict) + fits: dict[int, tuple[float, str]] = field(default_factory=dict) + ref_angle_idx: int = 0 + threshold_mv: float = 0.0 + result: compute.AlignmentResult | None = None # full, uncropped + counts: np.ndarray | None = None # coarse overlap counts + layers: dict[int, np.ndarray] = field(default_factory=dict) + preview_pitch_mm: tuple[float, float] = (1.0, 1.0) + preview_shape: tuple[int, int] = (1, 1) + geometry_generation: int = 0 + crop: tuple[int, int, int, int] | None = None # canvas (row0,col0,nr,nc) + cropped_result: compute.AlignmentResult | None = None + out_path: str = "" + exported_path: str = "" + + +class AlignmentWizard(QWizard): + """Three-page alignment + export flow (Fusion → Alignment Wizard…). + + Non-modal by design, like the dialog it replaces: an export can take minutes + on a real scan and the user should still be able to look at the data. + + Background work goes through parent._run_worker, which inherits the main + window's job registry, thread joining and shutdown handling. Two rules that + module's docstring establishes and every launch here follows: disable the + trigger *before* calling it (a re-entrant click must not be able to start a + second thread over the first), and never ignore its return value — False + means another job holds the key. + """ + + PAGE_CORRELATE, PAGE_ROI, PAGE_SAVE = 0, 1, 2 + + # cropped AlignmentResult, written path + alignment_ready = pyqtSignal(object, str) + + def __init__(self, parent: "SrasViewerWindow", sras: SrasFile, *, + ref_angle_idx: int, dc_threshold_mv: float, + seed_per_angle: dict[int, ManualAngleParams] | None, + cached_dc4_mv: dict[int, np.ndarray]): + super().__init__(parent) + self._parent = parent + self._sras = sras + self._cached_dc4 = dict(cached_dc4_mv) + self._seed = dict(seed_per_angle or {}) + self._closing = False + + self.state = WizardState( + ref_angle_idx=ref_angle_idx, threshold_mv=dc_threshold_mv, + params={a: ManualAngleParams() for a in range(sras.n_angles)}) + + n = sras.n_angles + cmap = mpl.colormaps["tab10"] if n <= 10 else mpl.colormaps["tab20"] + self.angle_colors = {a: cmap(a % cmap.N)[:3] for a in range(n)} + + self.setWindowTitle(f"Alignment Wizard — {sras.path.name}") + self.setWizardStyle(QWizard.WizardStyle.ModernStyle) + self.setOption(QWizard.WizardOption.HaveHelpButton, False) + self.setOption(QWizard.WizardOption.NoBackButtonOnStartPage, True) + # IndependentPages must stay OFF: it suppresses cleanupPage, which is + # how the ROI page discards a crop whose canvas is about to change. + self.setOption(QWizard.WizardOption.IndependentPages, False) + self.resize(1220, 820) + + self.setPage(self.PAGE_CORRELATE, CorrelatePage(self)) + self.setPage(self.PAGE_ROI, RoiPage(self)) + self.setPage(self.PAGE_SAVE, SavePage(self)) + self.setStartId(self.PAGE_CORRELATE) + + # ------------------------------------------------------------------ + # Shared helpers the pages use + # ------------------------------------------------------------------ + + @property + def sras(self) -> SrasFile: + return self._sras + + def seed_params(self) -> dict[int, ManualAngleParams]: + """Per-angle starting parameters: a saved sidecar if one matches, + otherwise identity (the pre-rotation is applied on top by page 1).""" + return {a: ManualAngleParams(self._seed[a].rotation_deg, + self._seed[a].shift_mm) + if a in self._seed else ManualAngleParams() + for a in range(self._sras.n_angles)} + + def rebuild_result(self): + """Recompute the full AlignmentResult from the current parameters. + + Closed-form matrix and bbox maths with no per-pixel work, so it is cheap + enough to call on every edit — which is what keeps the preview honest + about the canvas a rotation change implies. + """ + st = self.state + st.result = build_manual_alignment( + self._sras, st.ref_angle_idx, st.threshold_mv, st.params) + + def rebuild_stack(self): + """Reproject every angle's mask onto a coarsened view of the *final* + canvas and sum them into an overlap-count image. + + The preview deliberately shares the final canvas's origin and uses a + pitch that is an integer multiple of it, rather than the padded, + unsnapped preview canvas the old manual dialog used. That is what lets + the crop page turn a rectangle drawn in millimetres into an exact + integer window of the real canvas, with no second coordinate frame to + reconcile. + """ + st = self.state + if st.result is None or not st.masks_small: + return + fy, fx = st.downsample + n_rows, n_cols = st.result.canvas_shape + st.preview_pitch_mm = (st.result.canvas_dx_mm * fx, + st.result.canvas_dy_mm * fy) + st.preview_shape = (max(1, -(-n_rows // fy)), max(1, -(-n_cols // fx))) + + st.layers = {} + counts = np.zeros(st.preview_shape, dtype=np.int16) + for a in range(self._sras.n_angles): + p = st.params[a] + layer = compute.reproject_mask( + self._sras, a, st.ref_angle_idx, st.masks_small[a], + p.rotation_deg, p.shift_mm, st.preview_pitch_mm, + st.result.canvas_origin_mm, st.preview_shape, + src_downsample=st.downsample) + st.layers[a] = layer + counts += (layer > 0.5).astype(np.int16) + st.counts = counts + + def preview_extent(self) -> list[float]: + st = self.state + x0, y0 = st.result.canvas_origin_mm + dx, dy = st.preview_pitch_mm + n_rows, n_cols = st.preview_shape + return _axes_extent(x0 + np.arange(n_cols) * dx, + y0 + np.arange(n_rows) * dy, dx, dy) + + def mm_to_canvas(self, x_mm: float, y_mm: float) -> tuple[float, float]: + """(col, row) in *final* canvas pixels, fractional.""" + st = self.state + x0, y0 = st.result.canvas_origin_mm + return ((x_mm - x0) / st.result.canvas_dx_mm, + (y_mm - y0) / st.result.canvas_dy_mm) + + def canvas_to_mm(self, col: float, row: float) -> tuple[float, float]: + st = self.state + x0, y0 = st.result.canvas_origin_mm + return (x0 + col * st.result.canvas_dx_mm, + y0 + row * st.result.canvas_dy_mm) + + def job_running(self) -> bool: + return any(self._parent._job_running(k) for k in + (Jobs.ALIGN_MASKS, Jobs.ALIGN_CORRELATE, Jobs.ALIGN_EXPORT)) + + # ------------------------------------------------------------------ + # Lifetime + # ------------------------------------------------------------------ + + def reject(self): + """Refuse to close while a job is in flight. + + The running worker's signals are connected to bound methods of this + wizard and its pages; letting Qt delete them under a live thread is a + crash, not an inconvenience. Ask the worker to stop and let the job's + own completion close us. + """ + if self.job_running(): + for page_id in (self.PAGE_CORRELATE, self.PAGE_SAVE): + page = self.page(page_id) + if page is not None: + page.request_stop() + self._closing = True + return + super().reject() + + def closeEvent(self, event): + """Window-manager close goes through the same deferral as Cancel.""" + if self.job_running(): + self.reject() + event.ignore() + return + super().closeEvent(event) + + def accept(self): + st = self.state + self.alignment_ready.emit(st.cropped_result, st.exported_path) + super().accept() + + def maybe_close_after_job(self): + """Called by a page when its job finishes; completes a deferred close.""" + if self._closing and not self.job_running(): + self._closing = False + super().reject() + + +# --------------------------------------------------------------------------- +# Page 1 — cross-correlation +# --------------------------------------------------------------------------- + +class CorrelatePage(QWizardPage): + """Register every angle against the reference, and show whether it worked. + + The stack image is the point of the page. Numbers alone ("score 0.42") do + not tell you whether a five-angle fusion is usable; an overlap-count image + does, immediately. + """ + + def __init__(self, wizard: AlignmentWizard): + super().__init__(wizard) + self._wiz = wizard + self._sras = wizard.sras + self._busy = False + self._masks_ready = False + self._active_angle = 0 + self._worker = None + self._done_count = 0 + self._total = 0 + + self.setTitle("Step 1 — Cross-correlate the angles") + self.setSubTitle( + "Angles start pre-rotated by the stage angles stored in the scan. " + "Run the correlation, then check the stack: every pixel is coloured " + "by how many angles cover it, so a good alignment is one solid " + "plateau.") + self._build_ui() + + # ---- construction ------------------------------------------------- + + def _build_ui(self): + root = QHBoxLayout(self) + + self.canvas = AlignOverlayCanvas() + left = QWidget() + ll = QVBoxLayout(left) + ll.setContentsMargins(0, 0, 0, 0) + ll.setSpacing(4) + ll.addWidget(NavigationToolbar2QT(self.canvas, left)) + ll.addWidget(self.canvas, stretch=1) + self.lbl_overlap = _wrap_label("", _CSS_INFO) + ll.addWidget(self.lbl_overlap) + self.lbl_warn = _wrap_label("", _CSS_WARN) + ll.addWidget(self.lbl_warn) + root.addWidget(left, stretch=1) + + panel = QWidget() + pl = QVBoxLayout(panel) + pl.setContentsMargins(0, 0, 0, 0) + pl.setSpacing(8) + pl.addWidget(self._build_view_group()) + pl.addWidget(self._build_reg_group()) + pl.addWidget(self._build_run_group()) + pl.addWidget(self._build_nudge_group()) + pl.addWidget(self._build_table_group()) + self.lbl_status = _wrap_label("", _CSS_MUTED) + pl.addWidget(self.lbl_status) + pl.addStretch() + root.addWidget(_scroll_panel(panel, _PANEL_W)) + + self._connect() + + def _build_view_group(self) -> QWidget: + grp, lay = _group("View") + self.combo_view = _combo(["Overlap count", "Per-angle colours"]) + lay.addWidget(self.combo_view) + return grp + + def _build_reg_group(self) -> QWidget: + grp, lay = _group("Registration") + form = _form() + + self.combo_ref = _combo( + f"Angle {a} ({self._sras.angles_deg[a]:.1f}°)" + for a in range(self._sras.n_angles)) + self.combo_ref.setCurrentIndex(self._wiz.state.ref_angle_idx) + form.addRow("Reference:", self.combo_ref) + + self.spin_threshold = _make_dspin(-500.0, 500.0, 3, suffix=" mV", + value=self._wiz.state.threshold_mv) + form.addRow("DC threshold:", self.spin_threshold) + + self.combo_source = _combo( + (label, sources) for label, sources in _CORRELATE_SOURCES) + form.addRow("Correlate on:", self.combo_source) + + self.combo_prerotate = _combo( + (label, (seed, signs, disp)) + for label, seed, signs, disp in _PREROTATE_MODES) + form.addRow("Pre-rotate:", self.combo_prerotate) + + self.chk_lock_rotation = QCheckBox("Lock rotation to the seed") + form.addRow("", self.chk_lock_rotation) + + self.spin_search_deg = _make_dspin(0.0, 180.0, 1, suffix=" °", + value=6.0, step=1.0) + form.addRow("Rotation search (±):", self.spin_search_deg) + + self.spin_coarse_step = _make_dspin(0.1, 30.0, 2, suffix=" °", value=2.0) + form.addRow("Coarse step:", self.spin_coarse_step) + + self.combo_fine_dim = _combo((f"{dim} px", dim) for dim in (320, 640, 1024)) + self.combo_fine_dim.setCurrentIndex(1) + form.addRow("Fine grid:", self.combo_fine_dim) + + lay.addLayout(form) + lay.addWidget(_wrap_label( + "Pre-rotation only seeds the search — the rotation is still found " + "from image content, and both signs of the stage's reported angle " + "are tried unless you narrow it. Locking instead pins rotation to " + "the stage angle and searches translation only.", _CSS_HINT)) + return grp + + def _build_run_group(self) -> QWidget: + grp, lay = _group("Run") + self.btn_correlate = QPushButton("Run / Re-run Correlation") + lay.addWidget(self.btn_correlate) + self.progress = QProgressBar() + self.progress.setRange(0, 100) + self.progress.setValue(0) + lay.addWidget(self.progress) + self.btn_reset = QPushButton("Reset to pre-rotation only") + lay.addWidget(self.btn_reset) + return grp + + def _build_nudge_group(self) -> QWidget: + self.grp_nudge, lay = _group("Manual Correction") + form = _form() + self.combo_active = _combo( + f"Angle {a} ({self._sras.angles_deg[a]:.1f}°)" + for a in range(self._sras.n_angles)) + form.addRow("Active angle:", self.combo_active) + + self.spin_rot = _make_dspin(-3600.0, 3600.0, 3, suffix=" °") + form.addRow("Rotation:", self.spin_rot) + self.spin_shift_x = _make_dspin(-1e5, 1e5, 4, suffix=" mm") + form.addRow("Shift X:", self.spin_shift_x) + self.spin_shift_y = _make_dspin(-1e5, 1e5, 4, suffix=" mm") + form.addRow("Shift Y:", self.spin_shift_y) + + self.spin_step_translate = _make_dspin(0.0001, 1000.0, 4, suffix=" mm", + value=0.01) + form.addRow("Translate step:", self.spin_step_translate) + self.spin_step_rotate = _make_dspin(0.001, 90.0, 3, suffix=" °", value=0.1) + form.addRow("Rotate step:", self.spin_step_rotate) + self.spin_step_mult = _make_dspin(1.0, 1000.0, 1, value=10.0) + form.addRow("Coarse × (Shift):", self.spin_step_mult) + lay.addLayout(form) + lay.addWidget(_wrap_label( + "Arrow keys nudge translation; Q/E nudge rotation (CCW/CW). Hold " + "Shift for the coarse step. Click the image once to give it " + "keyboard focus. Switch to per-angle colours to see which layer " + "you are moving.", _CSS_HINT)) + self.lbl_active_note = _wrap_label("", _CSS_WARN) + lay.addWidget(self.lbl_active_note) + return self.grp_nudge + + def _build_table_group(self) -> QWidget: + grp, lay = _group("Fit per Angle") + self.table = QTableWidget(self._sras.n_angles, 5) + self.table.setHorizontalHeaderLabels( + ["Angle", "Rot °", "Δ stage °", "Score", "On"]) + self.table.verticalHeader().setVisible(False) + self.table.setEditTriggers(QTableWidget.EditTrigger.NoEditTriggers) + self.table.setSelectionMode(QTableWidget.SelectionMode.NoSelection) + self.table.horizontalHeader().setSectionResizeMode( + QHeaderView.ResizeMode.ResizeToContents) + self.table.setMinimumHeight(150) + lay.addWidget(self.table) + return grp + + def _connect(self): + self.combo_view.currentIndexChanged.connect(self._redraw) + self.combo_ref.currentIndexChanged.connect(self._on_ref_changed) + self.spin_threshold.editingFinished.connect(self._on_threshold_changed) + self.combo_prerotate.currentIndexChanged.connect(self._apply_prerotation) + self.chk_lock_rotation.toggled.connect(self._on_lock_toggled) + self.btn_correlate.clicked.connect(self._on_correlate) + self.btn_reset.clicked.connect(self._apply_prerotation) + self.combo_active.currentIndexChanged.connect(self._on_active_changed) + for spin in (self.spin_rot, self.spin_shift_x, self.spin_shift_y): + spin.editingFinished.connect(self._on_manual_edit) + self.canvas.nudge_translate.connect(self._on_nudge_translate) + self.canvas.nudge_rotate.connect(self._on_nudge_rotate) + + # ---- QWizardPage contract ---------------------------------------- + + def initializePage(self): + if self._masks_ready or self._busy: + return + self._set_controls_enabled(False) + self._start_mask_prep() + + def isComplete(self) -> bool: + """Next is gated on masks being ready and nothing running. + + Returning False while busy also disables Finish, which is what stops a + Back/Next/Finish race from reaching the writer with a half-built result. + """ + return self._masks_ready and not self._busy + + def request_stop(self): + if self._worker is not None: + self._worker.stop() + + # ---- mask preparation -------------------------------------------- + + def _start_mask_prep(self): + st = self._wiz.state + st.dc4_mv = dict(self._wiz._cached_dc4) + missing = [a for a in range(self._sras.n_angles) if a not in st.dc4_mv] + if not missing: + self._finish_mask_prep() + return + self._busy = True + self.completeChanged.emit() + self.lbl_status.setText( + f"Preparing masks: 0/{len(missing)} angle(s) needed…") + worker = Ch4MaskWorker(self._sras, missing) + started = self._wiz._parent._run_worker( + Jobs.ALIGN_MASKS, worker, + connect=(("angle_done", self._on_mask_done), + ("error", lambda m: self.lbl_status.setText( + f"Mask preparation failed: {m}"))), + on_done=self._finish_mask_prep) + if not started: + self._busy = False + self.completeChanged.emit() + self.lbl_status.setText( + "Another alignment step is still running — close and reopen.") + return + self._worker = worker + + def _on_mask_done(self, angle_idx: int, dc4_mv: np.ndarray): + st = self._wiz.state + st.dc4_mv[angle_idx] = dc4_mv + self.lbl_status.setText( + f"Preparing masks: {len(st.dc4_mv)}/{self._sras.n_angles} ready…") + + def _finish_mask_prep(self): + st = self._wiz.state + self._worker = None + self._busy = False + if len(st.dc4_mv) < self._sras.n_angles: + self.completeChanged.emit() + self.lbl_status.setText( + "Mask preparation did not finish for every angle.") + self._wiz.maybe_close_after_job() + return + + # Rows and columns get independent factors. A real scan is ~7500 frames + # wide but only ~750 rows tall, so one shared factor sized for the + # frames would throw away 10x more row detail than the preview needs. + max_rows = max(img.shape[0] for img in st.dc4_mv.values()) + max_cols = max(img.shape[1] for img in st.dc4_mv.values()) + st.downsample = (max(1, -(-max_rows // _MAX_PREVIEW_DIM)), + max(1, -(-max_cols // _MAX_PREVIEW_DIM))) + + st.params = self._wiz.seed_params() + self._recompute_masks() + self._masks_ready = True + self._set_controls_enabled(True) + self._apply_prerotation() + self.lbl_status.setText("Ready.") + self.completeChanged.emit() + self._wiz.maybe_close_after_job() + + def _recompute_masks(self): + """Threshold + downsample each angle's in-memory CH4 image. Cheap (a + compare and a block-mean), so a threshold change re-runs it in full + rather than re-fetching anything.""" + st = self._wiz.state + fy, fx = st.downsample + st.masks_small = { + a: compute.block_mean_2d((img >= st.threshold_mv).astype(np.float32), + fy, fx) + for a, img in st.dc4_mv.items() + } + + # ---- parameter changes ------------------------------------------- + + def _on_ref_changed(self, idx: int): + st = self._wiz.state + st.ref_angle_idx = idx + st.fits = {} + self._apply_prerotation() + + def _on_threshold_changed(self): + st = self._wiz.state + value = self.spin_threshold.value() + if value == st.threshold_mv: + return + st.threshold_mv = value + if self._masks_ready: + self._recompute_masks() + self._refresh() + + def _on_lock_toggled(self, locked: bool): + self.spin_search_deg.setEnabled(not locked) + self.spin_coarse_step.setEnabled(not locked) + + def _apply_prerotation(self): + """Seed every non-reference angle's rotation from the stage angles in + the file, and redraw — before any correlation has run. + + This is the cheapest useful thing the page can show: the stage angles + are usually within a degree or two of the truth, so the stack already + looks close to right, and how close is a fair first read on whether the + scan's own metadata can be trusted. + """ + if not self._masks_ready: + return + st = self._wiz.state + _seed, _signs, disp_sign = self.combo_prerotate.currentData() + for a in range(self._sras.n_angles): + rot = 0.0 if a == st.ref_angle_idx else ( + disp_sign * compute.nominal_delta_deg(self._sras, a, + st.ref_angle_idx)) + st.params[a] = ManualAngleParams(rot, (0.0, 0.0)) + st.fits = {} + self._refresh() + + def _refresh(self): + """Rebuild result, stack and every readout from the current params.""" + st = self._wiz.state + self._wiz.rebuild_result() + st.geometry_generation += 1 + self._wiz.rebuild_stack() + self._sync_spins() + self._update_table() + self._redraw() + + # ---- correlation -------------------------------------------------- + + def _reg_kwargs(self) -> dict: + seed, signs, _disp = self.combo_prerotate.currentData() + kwargs = { + "seed_deg": seed, + "seed_signs": signs, + "coarse_step_deg": self.spin_coarse_step.value(), + "fine_dim": self.combo_fine_dim.currentData(), + } + if self.chk_lock_rotation.isChecked(): + # Exactly one candidate, no hill-climb: rotation is the seed and + # only the translation is searched. + kwargs["refine"] = False + if len(signs) > 1: + kwargs["seed_signs"] = (1,) + return kwargs + + def _on_correlate(self): + if not self._masks_ready or self._busy: + return + st = self._wiz.state + angles = [a for a in range(self._sras.n_angles) if a != st.ref_angle_idx] + if not angles: + self.lbl_status.setText("Only one angle — nothing to correlate.") + return + + search = 0.0 if self.chk_lock_rotation.isChecked() \ + else self.spin_search_deg.value() + worker = CrossCorrelateWorker( + self._sras, st.ref_angle_idx, angles, st.dc4_mv, + sources=self.combo_source.currentData(), + dc_threshold_mv=st.threshold_mv, search_deg=search, + reg_kwargs=self._reg_kwargs()) + + # Claim busy and disable the trigger *before* _run_worker, never after: + # anything that pumps the event loop in between could deliver a second + # click that starts a thread the first assignment then drops. + self._busy = True + self._done_count, self._total = 0, len(angles) + st.fits = {} + self.completeChanged.emit() + self._set_controls_enabled(False) + self.progress.setValue(0) + self.lbl_status.setText(f"Cross-correlating: 0/{self._total} angle(s)…") + + started = self._wiz._parent._run_worker( + Jobs.ALIGN_CORRELATE, worker, + connect=(("angle_done", self._on_angle_done), + ("error", lambda m: self.lbl_status.setText( + f"Cross-correlation failed: {m}"))), + on_done=self._finish_correlate) + if not started: + self._busy = False + self.completeChanged.emit() + self._set_controls_enabled(True) + self.lbl_status.setText( + "Another alignment step is still running — try again shortly.") + return + self._worker = worker + + def _on_angle_done(self, angle_idx: int, rot: float, sx: float, sy: float, + score: float, source: str): + st = self._wiz.state + st.params[angle_idx] = ManualAngleParams(rot, (sx, sy)) + st.fits[angle_idx] = (score, source) + self._done_count += 1 + self.progress.setValue(int(self._done_count / max(1, self._total) * 100)) + self.lbl_status.setText( + f"Cross-correlating: {self._done_count}/{self._total} angle(s)…") + + def _finish_correlate(self): + self._worker = None + self._busy = False + self.progress.setValue(100) + self._refresh() + self._set_controls_enabled(True) + self.completeChanged.emit() + self.lbl_status.setText( + f"Correlated {self._done_count} angle(s) against angle " + f"{self._wiz.state.ref_angle_idx}. " + self._fit_summary()) + self._wiz.maybe_close_after_job() + + def _fit_summary(self) -> str: + """Worst fit and any angle whose rotation disagrees with its stage angle. + + Surfaced rather than buried because a single bad acquisition (stage + glitch, laser dropout) registers poorly and would otherwise be fused in + silently — knowing *which* angle is what makes dropping it with + sras_edit_scans.py actionable. + """ + st = self._wiz.state + if not st.fits: + return "" + rows = sorted(st.fits.items(), key=lambda kv: kv[1][0]) + worst_a, (worst_score, worst_src) = rows[0] + parts = [f"Worst fit: angle {worst_a} (score {worst_score:.3f}, " + f"{worst_src or 'n/a'})."] + failed = [str(a) for a, (score, src) in rows if score < 0 or src == "none"] + if failed: + parts.append( + "Angle(s) " + ", ".join(failed) + " did not register at all and " + "are being treated as unrotated — lower the DC threshold, try " + "Raw signal, nudge them by hand, or drop them with " + "sras_edit_scans.py.") + drifted = [] + for a, _ in rows: + nominal = compute.nominal_delta_deg(self._sras, a, st.ref_angle_idx) + got = st.params[a].rotation_deg + dev = min(abs(got - nominal), abs(got + nominal)) + if dev > 1.0: + drifted.append(f"{a} ({dev:.2f}°)") + if drifted: + parts.append("Rotation differs from the stage angle by >1° for " + "angle(s) " + ", ".join(drifted) + ".") + return " ".join(parts) + + # ---- manual nudging ---------------------------------------------- + + def _on_active_changed(self, idx: int): + self._active_angle = idx + is_ref = idx == self._wiz.state.ref_angle_idx + for spin in (self.spin_rot, self.spin_shift_x, self.spin_shift_y): + spin.setEnabled(self._masks_ready and not is_ref) + self.lbl_active_note.setText( + "Reference angle — defines the shared origin, not adjustable." + if is_ref else "") + self._sync_spins() + self._redraw() + + def _sync_spins(self): + p = self._wiz.state.params.get(self._active_angle, ManualAngleParams()) + for spin, val in ((self.spin_rot, p.rotation_deg), + (self.spin_shift_x, p.shift_mm[0]), + (self.spin_shift_y, p.shift_mm[1])): + with QSignalBlocker(spin): + spin.setValue(val) + + def _on_manual_edit(self): + if not self._masks_ready or self._active_angle == self._wiz.state.ref_angle_idx: + return + self._wiz.state.params[self._active_angle] = ManualAngleParams( + self.spin_rot.value(), + (self.spin_shift_x.value(), self.spin_shift_y.value())) + self._refresh() + + def _on_nudge_translate(self, dir_x: int, dir_y: int, coarse: bool): + if not self._masks_ready or self._active_angle == self._wiz.state.ref_angle_idx: + return + step = self.spin_step_translate.value() + if coarse: + step *= self.spin_step_mult.value() + p = self._wiz.state.params[self._active_angle] + self._wiz.state.params[self._active_angle] = ManualAngleParams( + p.rotation_deg, + (p.shift_mm[0] + dir_x * step, p.shift_mm[1] + dir_y * step)) + self._refresh() + + def _on_nudge_rotate(self, direction: int, coarse: bool): + if not self._masks_ready or self._active_angle == self._wiz.state.ref_angle_idx: + return + step = self.spin_step_rotate.value() + if coarse: + step *= self.spin_step_mult.value() + p = self._wiz.state.params[self._active_angle] + self._wiz.state.params[self._active_angle] = ManualAngleParams( + p.rotation_deg + direction * step, p.shift_mm) + self._refresh() + + # ---- drawing ------------------------------------------------------ + + def _set_controls_enabled(self, enabled: bool): + for w in (self.combo_ref, self.spin_threshold, self.combo_source, + self.combo_prerotate, self.chk_lock_rotation, + self.spin_search_deg, self.spin_coarse_step, + self.combo_fine_dim, self.btn_correlate, self.btn_reset, + self.grp_nudge): + w.setEnabled(enabled) + if enabled: + self._on_lock_toggled(self.chk_lock_rotation.isChecked()) + self._on_active_changed(self._active_angle) + + def _redraw(self): + st = self._wiz.state + if st.counts is None or st.result is None: + return + extent = self._wiz.preview_extent() + if self.combo_view.currentIndex() == 0: + self.canvas.show_counts( + st.counts, self._sras.n_angles, extent, + f"Mask stack — {self._sras.n_angles} angles, " + f"ref angle {st.ref_angle_idx}") + else: + self.canvas.show_overlay( + self._rgba(), extent, + f"Angle {self._active_angle} active " + f"({self._sras.angles_deg[self._active_angle]:.1f}°)") + self._update_overlap_text() + + def _rgba(self) -> np.ndarray: + """Alpha-composite each angle's mask in its own colour, active on top.""" + st = self._wiz.state + n_rows, n_cols = st.preview_shape + rgba = np.zeros((n_rows, n_cols, 4), dtype=np.float32) + order = sorted(range(self._sras.n_angles), + key=lambda a: a == self._active_angle) + for a in order: + layer = st.layers.get(a) + if layer is None: + continue + alpha = _ACTIVE_ALPHA if a == self._active_angle else _BASE_ALPHA + color = self._wiz.angle_colors[a] + fg = layer * alpha + for c in range(3): + rgba[..., c] = color[c] * fg + rgba[..., c] * rgba[..., 3] * (1 - fg) + rgba[..., 3] = fg + rgba[..., 3] * (1 - fg) + return rgba + + def _update_overlap_text(self): + st = self._wiz.state + n = self._sras.n_angles + stats = compute.overlap_stats(st.counts, n) + px = st.preview_pitch_mm + area = abs(px[0] * px[1]) + self.lbl_overlap.setText( + f"Covered by any angle: {stats['union_px']:,} px " + f"({stats['union_px'] * area:.2f} mm²) · " + f"by all {n}: {stats['full_px']:,} px " + f"({stats['full_frac'] * 100:.1f}% of that) · " + f"mean overlap {stats['mean_count']:.2f} angles") + if stats["empty"]: + self.lbl_warn.setText( + "No angle covers any pixel — check the DC threshold.") + elif stats["max_count"] < n: + self.lbl_warn.setText( + f"No pixel is covered by all {n} angles (best is " + f"{stats['max_count']}). Either the alignment is wrong or these " + f"scans genuinely do not all overlap. You can still export.") + else: + self.lbl_warn.setText("") + + def _update_table(self): + st = self._wiz.state + for a in range(self._sras.n_angles): + score, source = st.fits.get(a, (float("nan"), "")) + nominal = compute.nominal_delta_deg(self._sras, a, st.ref_angle_idx) + got = st.params[a].rotation_deg + dev = min(abs(got - nominal), abs(got + nominal)) + is_ref = a == st.ref_angle_idx + cells = [ + f"{a}" + (" [ref]" if is_ref else ""), + f"{got:.3f}", + "—" if is_ref else f"{dev:.2f}", + "—" if is_ref or np.isnan(score) else f"{score:.3f}", + "ref" if is_ref else (source or "—"), + ] + bad = (not is_ref) and (score < 0 or source == "none") + weak = (not is_ref) and (not np.isnan(score)) and 0 <= score < 0.3 + for col, text in enumerate(cells): + item = QTableWidgetItem(text) + item.setTextAlignment(Qt.AlignmentFlag.AlignCenter) + if bad: + item.setForeground(Qt.GlobalColor.red) + elif weak or (not is_ref and dev > 1.0): + item.setForeground(Qt.GlobalColor.darkYellow) + self.table.setItem(a, col, item) + + +# --------------------------------------------------------------------------- +# Page 2 — ROI crop +# --------------------------------------------------------------------------- + +class RoiPage(QWizardPage): + """Pick the rectangle of the shared canvas to keep. + + Axis-aligned, because v6 geometry is (x_start, x_delta, n_frames, n_rows) + plus a Y row table — a rectangle on the canvas grid is the only crop the + file can express, and drawing something else would have to be squared off + without saying so. + """ + + def __init__(self, wizard: AlignmentWizard): + super().__init__(wizard) + self._wiz = wizard + self._sras = wizard.sras + self._syncing = False + self._generation = -1 + + self.setTitle("Step 2 — Choose the region to keep") + self.setSubTitle( + "Drag a rectangle on the stack, or type canvas pixels directly. " + "Cropping to the region the angles actually share is usually a big " + "size saving over the full canvas.") + self._build_ui() + + def _build_ui(self): + root = QHBoxLayout(self) + + self.canvas = ImageCanvas(rect_only=True) + left = QWidget() + ll = QVBoxLayout(left) + ll.setContentsMargins(0, 0, 0, 0) + ll.setSpacing(4) + ll.addWidget(NavigationToolbar2QT(self.canvas, left)) + ll.addWidget(self.canvas, stretch=1) + root.addWidget(left, stretch=1) + + panel = QWidget() + pl = QVBoxLayout(panel) + pl.setContentsMargins(0, 0, 0, 0) + pl.setSpacing(8) + + grp_fit, fl = _group("Preset") + self.btn_fit_overlap = QPushButton("Fit to full overlap") + self.btn_fit_union = QPushButton("Fit to any coverage") + self.btn_whole = QPushButton("Whole canvas (no crop)") + for b in (self.btn_fit_overlap, self.btn_fit_union, self.btn_whole): + fl.addWidget(b) + fl.addWidget(_wrap_label( + "“Full overlap” finds the largest rectangle lying entirely inside " + "the region every angle covers — not its bounding box, which would " + "include corners no angle reaches.", _CSS_HINT)) + pl.addWidget(grp_fit) + + grp_rect, rl = _group("Crop (canvas pixels)") + form = _form() + self.spin_col0 = QSpinBox() + self.spin_row0 = QSpinBox() + self.spin_cols = QSpinBox() + self.spin_rows = QSpinBox() + for spin, label in ((self.spin_col0, "First column:"), + (self.spin_row0, "First row:"), + (self.spin_cols, "Columns:"), + (self.spin_rows, "Rows:")): + spin.setRange(0, 1) + spin.setMinimumWidth(96) + form.addRow(label, spin) + rl.addLayout(form) + self.btn_draw = QPushButton("Draw a new rectangle") + rl.addWidget(self.btn_draw) + pl.addWidget(grp_rect) + + grp_info, il = _group("Result") + self.lbl_extent = _wrap_label("", _CSS_INFO) + self.lbl_size = _wrap_label("", _CSS_INFO) + self.lbl_coverage = _wrap_label("", _CSS_MUTED) + self.lbl_warn = _wrap_label("", _CSS_WARN) + for w in (self.lbl_extent, self.lbl_size, self.lbl_coverage, self.lbl_warn): + il.addWidget(w) + pl.addWidget(grp_info) + + pl.addStretch() + root.addWidget(_scroll_panel(panel, _PANEL_W)) + + self.btn_fit_overlap.clicked.connect(self._on_fit_overlap) + self.btn_fit_union.clicked.connect(self._on_fit_union) + self.btn_whole.clicked.connect(self._on_whole) + self.btn_draw.clicked.connect(self.canvas.start_drawing) + self.canvas.roi_changed.connect(self._on_roi_changed) + for spin in (self.spin_col0, self.spin_row0, self.spin_cols, self.spin_rows): + spin.valueChanged.connect(self._on_spin_changed) + + # ---- QWizardPage contract ---------------------------------------- + + def initializePage(self): + st = self._wiz.state + # A crop is indexed in canvas pixels, so it is meaningless against a + # canvas built from different rotations. Drop a stale one rather than + # silently reinterpreting its indices. + if st.crop is not None and self._generation != st.geometry_generation: + st.crop = None + self._generation = st.geometry_generation + + n_rows, n_cols = st.result.canvas_shape + # Blocked: setRange clamps the current value into the new range and + # emits valueChanged, which would run _on_spin_changed and commit a + # crop built from three not-yet-set spin boxes — landing on 1x1 and + # making the crop look already-chosen, so the preset below is skipped. + for spin, lo, hi in ((self.spin_col0, 0, n_cols - 1), + (self.spin_row0, 0, n_rows - 1), + (self.spin_cols, 1, n_cols), + (self.spin_rows, 1, n_rows)): + with QSignalBlocker(spin): + spin.setRange(lo, max(lo, hi)) + self._draw_counts() + self.btn_fit_overlap.setEnabled( + compute.largest_rect_at_least(st.counts, self._sras.n_angles) + is not None) + if st.crop is None: + self._on_fit_union() + else: + self._set_crop(*st.crop) + + def isComplete(self) -> bool: + crop = self._wiz.state.crop + return crop is not None and crop[2] >= 1 and crop[3] >= 1 + + def validatePage(self) -> bool: + st = self._wiz.state + cropped = compute.crop_alignment_result(st.result, *st.crop) + plan = export.plan_export(self._sras, cropped) + + empty = [a for a in range(self._sras.n_angles) if plan.valid_px[a] == 0] + if len(empty) == self._sras.n_angles: + QMessageBox.warning( + self, "Empty crop", + "This rectangle contains no data from any angle. Move or grow " + "it before continuing.") + return False + if empty: + answer = QMessageBox.question( + self, "Some angles are empty", + f"Angle(s) {', '.join(map(str, empty))} have no data inside " + f"this crop and would be written as padding.\n\nContinue anyway?") + if answer != QMessageBox.StandardButton.Yes: + return False + + st.cropped_result = cropped + return True + + def cleanupPage(self): + """Going back means the canvas geometry may change under this crop.""" + self._wiz.state.crop = None + self._wiz.state.cropped_result = None + + # ---- drawing / presets -------------------------------------------- + + def _draw_counts(self): + st = self._wiz.state + n = self._sras.n_angles + # Same colormap the previous page used, so a count keeps its colour + # across the two pages that show this image on different canvas classes. + cmap, norm, ticks = count_colormap(n) + self.canvas.show_image( + st.counts, self._wiz.preview_extent(), cmap, 0.0, 0.0, + "X (mm)", "Y (mm)", + f"Overlap count — drag the crop rectangle ({n} angles)", + colorbar_label="angles overlapping", cb_ticks=ticks, norm=norm) + + def _coarse_rect_to_canvas(self, rect) -> tuple[int, int, int, int]: + """A rectangle in coarse preview indices, as final canvas pixels. + + Inset by one coarse block on each side. The coarse grid samples canvas + pixels 0, fx, 2fx, …, so a coarse pixel reported as fully covered + stands for a block whose far edge may not be; shrinking by a block keeps + a "fit to full overlap" rectangle honestly inside the overlap region. + """ + st = self._wiz.state + fy, fx = st.downsample + row0, col0, nr, nc = rect + n_rows, n_cols = st.result.canvas_shape + r0 = min(n_rows - 1, row0 * fy + fy) + c0 = min(n_cols - 1, col0 * fx + fx) + nr_c = max(1, min(n_rows - r0, nr * fy - 2 * fy)) + nc_c = max(1, min(n_cols - c0, nc * fx - 2 * fx)) + return r0, c0, nr_c, nc_c + + def _on_fit_overlap(self): + st = self._wiz.state + rect = compute.largest_rect_at_least(st.counts, self._sras.n_angles) + if rect is None: + return + self._set_crop(*self._coarse_rect_to_canvas(rect)) + + def _on_fit_union(self): + st = self._wiz.state + rows, cols = np.nonzero(st.counts > 0) + if rows.size == 0: + self._on_whole() + return + fy, fx = st.downsample + n_rows, n_cols = st.result.canvas_shape + r0 = int(rows.min()) * fy + c0 = int(cols.min()) * fx + r1 = min(n_rows, (int(rows.max()) + 1) * fy) + c1 = min(n_cols, (int(cols.max()) + 1) * fx) + self._set_crop(r0, c0, max(1, r1 - r0), max(1, c1 - c0)) + + def _on_whole(self): + n_rows, n_cols = self._wiz.state.result.canvas_shape + self._set_crop(0, 0, n_rows, n_cols) + + # ---- two-way sync ------------------------------------------------- + + def _set_crop(self, row0: int, col0: int, n_rows: int, n_cols: int): + st = self._wiz.state + cr, cc = st.result.canvas_shape + row0 = int(np.clip(row0, 0, cr - 1)) + col0 = int(np.clip(col0, 0, cc - 1)) + n_rows = int(np.clip(n_rows, 1, cr - row0)) + n_cols = int(np.clip(n_cols, 1, cc - col0)) + st.crop = (row0, col0, n_rows, n_cols) + + self._syncing = True + try: + for spin, val in ((self.spin_row0, row0), (self.spin_col0, col0), + (self.spin_rows, n_rows), (self.spin_cols, n_cols)): + with QSignalBlocker(spin): + spin.setValue(val) + x0, y0 = self._wiz.canvas_to_mm(col0 - 0.5, row0 - 0.5) + x1, y1 = self._wiz.canvas_to_mm(col0 + n_cols - 0.5, + row0 + n_rows - 0.5) + self.canvas.set_roi(RoiQuad.from_bbox(min(x0, x1), min(y0, y1), + max(x0, x1), max(y0, y1))) + finally: + self._syncing = False + self._update_readout() + self.completeChanged.emit() + + def _on_spin_changed(self): + if self._syncing: + return + self._set_crop(self.spin_row0.value(), self.spin_col0.value(), + self.spin_rows.value(), self.spin_cols.value()) + + def _on_roi_changed(self): + # set_roi itself emits roi_changed, so without the guard _set_crop + # would re-enter through its own canvas update. + if self._syncing: + return + roi = self.canvas.get_roi() + if roi is None: + return + pts = roi.corners() + (x0, y0), (x1, y1) = pts.min(axis=0), pts.max(axis=0) + c_a, r_a = self._wiz.mm_to_canvas(x0, y0) + c_b, r_b = self._wiz.mm_to_canvas(x1, y1) + col0, col1 = sorted((c_a, c_b)) + row0, row1 = sorted((r_a, r_b)) + # Both edges snap to the nearest pixel boundary with the same rule, so + # the numeric boxes -> rectangle -> numeric boxes round trip is exact. + # _SNAP_TOL absorbs the float noise of the mm round trip: a boundary + # that should land on 11.0 arrives as 11.000000000000002, and a bare + # ceil() turns that into 12 — one spurious column per edit. + r0, r1 = _snap_edge(row0), _snap_edge(row1) + c0, c1 = _snap_edge(col0), _snap_edge(col1) + self._set_crop(r0, c0, max(1, r1 - r0), max(1, c1 - c0)) + + def _update_readout(self): + st = self._wiz.state + row0, col0, n_rows, n_cols = st.crop + x0, y0 = self._wiz.canvas_to_mm(col0, row0) + x1, y1 = self._wiz.canvas_to_mm(col0 + n_cols - 1, row0 + n_rows - 1) + self.lbl_extent.setText( + f"Output: {n_cols:,} × {n_rows:,} px " + f"X {min(x0, x1):.3f} … {max(x0, x1):.3f} mm " + f"Y {min(y0, y1):.3f} … {max(y0, y1):.3f} mm") + + cropped = compute.crop_alignment_result(st.result, *st.crop) + plan = export.plan_export(self._sras, cropped) + self.lbl_size.setText(f"Estimated file size: {_humanize(plan.total_bytes)}" + f" ({self._sras.n_angles} angles)") + self.lbl_coverage.setText("Real data per angle: " + ", ".join( + f"{a}: {plan.coverage_frac(a) * 100:.0f}%" + for a in range(self._sras.n_angles))) + empty = [a for a in range(self._sras.n_angles) if plan.valid_px[a] == 0] + self.lbl_warn.setText( + f"Angle(s) {', '.join(map(str, empty))} have no data here and would " + f"be written as padding." if empty else "") + + +# --------------------------------------------------------------------------- +# Page 3 — save +# --------------------------------------------------------------------------- + +class SavePage(QWizardPage): + """Choose a destination and write the aligned, cropped scan. + + The write runs on a worker and is started by a button rather than from + validatePage, which must not block the GUI thread for what can be minutes of + I/O. Finish only becomes available once a file has actually been written, so + the wizard cannot be completed on a failed export. + """ + + def __init__(self, wizard: AlignmentWizard): + super().__init__(wizard) + self._wiz = wizard + self._sras = wizard.sras + self._busy = False + self._worker = None + + self.setTitle("Step 3 — Save the aligned scan") + self.setSubTitle( + "Writes a new v6 .sras in which every angle shares the cropped " + "grid, so it opens already aligned. The original is not modified.") + self._build_ui() + + def _build_ui(self): + root = QVBoxLayout(self) + + row = QHBoxLayout() + row.addWidget(QLabel("Save to:")) + self.edit_path = QLineEdit() + self.edit_path.setReadOnly(True) + row.addWidget(self.edit_path, stretch=1) + self.btn_browse = QPushButton("Browse…") + row.addWidget(self.btn_browse) + root.addLayout(row) + + grp, gl = _group("What will be written") + self.lbl_summary = _wrap_label("", _CSS_INFO) + gl.addWidget(self.lbl_summary) + self.lbl_notes = _wrap_label("", _CSS_WARN) + gl.addWidget(self.lbl_notes) + root.addWidget(grp) + + run = QHBoxLayout() + self.btn_export = QPushButton("Write .sras") + run.addWidget(self.btn_export) + self.btn_cancel = QPushButton("Cancel write") + self.btn_cancel.setEnabled(False) + run.addWidget(self.btn_cancel) + run.addStretch() + root.addLayout(run) + + self.progress = QProgressBar() + self.progress.setRange(0, 100) + root.addWidget(self.progress) + + self.lbl_status = _wrap_label("", _CSS_MUTED) + root.addWidget(self.lbl_status) + root.addStretch() + + self.btn_browse.clicked.connect(self._on_browse) + self.btn_export.clicked.connect(self._on_export) + self.btn_cancel.clicked.connect(self.request_stop) + + # ---- QWizardPage contract ---------------------------------------- + + def initializePage(self): + st = self._wiz.state + if not st.out_path: + src = Path(self._sras.path) + st.out_path = str(src.with_name(f"{src.stem}_aligned.sras")) + self.edit_path.setText(st.out_path) + st.exported_path = "" + self.progress.setValue(0) + self.lbl_status.setText("") + self._update_summary() + self._wiz.setButtonText(QWizard.WizardButton.FinishButton, "Done") + self.completeChanged.emit() + + def isComplete(self) -> bool: + return bool(self._wiz.state.exported_path) and not self._busy + + def request_stop(self): + if self._worker is not None: + self._worker.stop() + self.lbl_status.setText("Cancelling…") + + # ---- summary ------------------------------------------------------ + + def _update_summary(self): + st = self._wiz.state + plan = export.plan_export(self._sras, st.cropped_result) + x0, y0 = st.cropped_result.canvas_origin_mm + self.lbl_summary.setText( + f"Format: .sras v6, no precomputed cache (the viewer recomputes " + f"DC/FFT on first open).\n" + f"Angles: {plan.n_angles}, all sharing one grid of " + f"{plan.n_frames:,} × {plan.n_rows:,} px.\n" + f"Origin: X {x0:.4f} mm, Y {y0:.4f} mm · " + f"pitch {st.cropped_result.canvas_dx_mm * 1000:.2f} × " + f"{abs(st.cropped_result.canvas_dy_mm) * 1000:.2f} µm.\n" + f"Stage angles, calibration preambles and the background waveform " + f"are carried over unchanged.\n" + f"Size: {_humanize(plan.bytes_per_angle)} per angle, " + f"{_humanize(plan.total_bytes)} total.\n" + f"Real data per angle: " + ", ".join( + f"{a}: {plan.coverage_frac(a) * 100:.0f}%" + for a in range(plan.n_angles))) + self.lbl_notes.setText("\n".join(plan.warnings)) + + # ---- export ------------------------------------------------------- + + def _on_browse(self): + st = self._wiz.state + path, _ = QFileDialog.getSaveFileName( + self, "Save aligned .sras", st.out_path, + "SRAS scans (*.sras);;All files (*)") + if not path: + return + if not path.lower().endswith(".sras"): + path += ".sras" + st.out_path = path + self.edit_path.setText(path) + st.exported_path = "" + self.completeChanged.emit() + + def _on_export(self): + st = self._wiz.state + if self._busy or not st.out_path: + return + + worker = AlignedExportWorker(self._sras, st.cropped_result, st.out_path) + self._busy = True + self.completeChanged.emit() + self.btn_export.setEnabled(False) + self.btn_browse.setEnabled(False) + self.progress.setValue(0) + self.lbl_status.setText(f"Writing {Path(st.out_path).name}…") + + started = self._wiz._parent._run_worker( + Jobs.ALIGN_EXPORT, worker, + connect=(("progress", self.progress.setValue), + ("finished", self._on_export_finished))) + if not started: + self._busy = False + self.completeChanged.emit() + self.btn_export.setEnabled(True) + self.btn_browse.setEnabled(True) + self.lbl_status.setText( + "Another alignment step is still running — try again shortly.") + return + self._worker = worker + self.btn_cancel.setEnabled(True) + + def _on_export_finished(self, written: str, error: str): + st = self._wiz.state + self._worker = None + self._busy = False + self.btn_export.setEnabled(True) + self.btn_browse.setEnabled(True) + self.btn_cancel.setEnabled(False) + + if error: + self.progress.setValue(0) + self.lbl_status.setText(f"Export failed: {error}") + elif not written: + self.progress.setValue(0) + self.lbl_status.setText("Export cancelled; no file was written.") + else: + st.exported_path = written + self.progress.setValue(100) + self.lbl_status.setText( + f"Wrote {Path(written).name}. Choose Done to apply this " + f"alignment to the open scan as well.") + self.completeChanged.emit() + self._wiz.maybe_close_after_job() + + +# Slack when snapping a rectangle edge, in canvas pixels. The mm round trip is +# two multiplications and a subtraction, so an exact boundary can come back a +# few ULPs either side of the integer. +_SNAP_TOL = 1e-6 + + +def _combo(items) -> QComboBox: + """A combo box whose size hint does not depend on its longest entry. + + By default a QComboBox asks for enough width to show its widest item. These + hold descriptive phrases, and the panel lives in a fixed-width scroll area + with the horizontal scrollbar off (`_scroll_panel`) — so an unconstrained + hint pushes the inner widget past the panel and everything on the right, + including the hint text, is silently clipped instead of scrolling. + + *items* is a sequence of (label, data) pairs, or of plain labels. + """ + combo = QComboBox() + combo.setSizeAdjustPolicy( + QComboBox.SizeAdjustPolicy.AdjustToMinimumContentsLengthWithIcon) + combo.setMinimumContentsLength(10) + for item in items: + if isinstance(item, tuple): + combo.addItem(item[0], item[1]) + else: + combo.addItem(item) + return combo + + +def _snap_edge(coord: float) -> int: + """A fractional pixel-index edge, as the nearest pixel boundary index. + + Pixel k spans [k - 0.5, k + 0.5), so a boundary at fractional coordinate + *coord* is boundary index coord + 0.5. + """ + return int(np.floor(coord + 0.5 + _SNAP_TOL)) + + +def _humanize(n_bytes: int) -> str: + value = float(n_bytes) + for unit in ("B", "KB", "MB", "GB", "TB"): + if value < 1024 or unit == "TB": + return f"{value:.1f} {unit}" if unit != "B" else f"{int(value)} B" + value /= 1024 + return f"{value:.1f} TB" diff --git a/sras_viewer/canvases.py b/sras_viewer/canvases.py index fa73704..8c0a1e4 100644 --- a/sras_viewer/canvases.py +++ b/sras_viewer/canvases.py @@ -13,6 +13,27 @@ from PyQt6.QtWidgets import QSizePolicy from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, CH_NAMES, SrasFile, adc_to_mv +def count_colormap(n_angles: int): + """(cmap, norm, ticks) for an integer "how many angles cover this pixel" + image, 0..n_angles. + + Discrete, not continuous: the judgement the wizard's stack view exists for + is "is this a plateau at N, or a fan of partial overlaps", so a region + covered by one angle too few has to read as its own band rather than a + slightly darker shade. Count 0 is fully transparent so uncovered canvas + cannot be mistaken for a low count. + + Shared by both wizard pages that draw this image — they use different canvas + classes, and the same number must not change colour between them. + """ + n = max(1, int(n_angles)) + base = mpl.colormaps["viridis"].resampled(n) + colors = [(0.0, 0.0, 0.0, 0.0)] + [base(i) for i in range(n)] + return (ListedColormap(colors), + BoundaryNorm(np.arange(-0.5, n + 1), len(colors)), + np.arange(0, n + 1)) + + # --------------------------------------------------------------------------- # ROI (free quadrilateral in data coordinates) # --------------------------------------------------------------------------- @@ -145,10 +166,11 @@ class ImageCanvas(FigureCanvasQTAgg): def show_image(self, img: np.ndarray, extent: list[float], cmap, vmin: float, vmax: float, xlabel: str, ylabel: str, title: str, - colorbar_label: str = "", cb_ticks=None): - """*cmap* may be a name or a Colormap instance; *cb_ticks* pins the - colorbar's ticks, which the wizard's integer overlap-count view needs so - each band reads as a whole number of angles rather than a shade.""" + colorbar_label: str = "", cb_ticks=None, norm=None): + """*cmap* may be a name or a Colormap instance. *norm* (which overrides + vmin/vmax) and *cb_ticks* let a caller draw a discrete integer image — + the wizard's overlap-count view — with whole-number colorbar bands + instead of a continuous shade.""" self.figure.clf() self.ax = self.figure.add_subplot(111) # Patches and lines are destroyed by figure.clf(); drop stale refs. @@ -157,10 +179,11 @@ class ImageCanvas(FigureCanvasQTAgg): self._extent = extent self._img_shape = img.shape + kw = ({"norm": norm} if norm is not None + else {"vmin": vmin, "vmax": vmax}) im = self.ax.imshow( img, aspect="auto", origin="upper", - extent=extent, cmap=cmap, vmin=vmin, vmax=vmax, - interpolation="nearest", + extent=extent, cmap=cmap, interpolation="nearest", **kw, ) cb = self.figure.colorbar(im, ax=self.ax, fraction=0.046, pad=0.04, ticks=cb_ticks) @@ -576,16 +599,12 @@ class AlignOverlayCanvas(FigureCanvasQTAgg): """ self.figure.clf() self.ax = self.figure.add_subplot(111) - n = max(1, int(n_angles)) - colors = [(0.0, 0.0, 0.0, 0.0)] # count 0: nothing - base = mpl.colormaps["viridis"].resampled(n) - colors += [base(i) for i in range(n)] + cmap, norm, ticks = count_colormap(n_angles) im = self.ax.imshow( np.asarray(counts), extent=extent, origin="upper", aspect="auto", - interpolation="nearest", cmap=ListedColormap(colors), - norm=BoundaryNorm(np.arange(-0.5, n + 1), len(colors))) + interpolation="nearest", cmap=cmap, norm=norm) cb = self.figure.colorbar(im, ax=self.ax, fraction=0.046, pad=0.04, - ticks=np.arange(0, n + 1)) + ticks=ticks) cb.set_label("angles overlapping") self._finish(title) diff --git a/sras_viewer/dialogs.py b/sras_viewer/dialogs.py index c1c8352..47c067d 100644 --- a/sras_viewer/dialogs.py +++ b/sras_viewer/dialogs.py @@ -1,33 +1,18 @@ -"""FFT Options and Manual Alignment dialogs.""" +"""FFT option dialogs. -from typing import TYPE_CHECKING +Angle alignment used to live here too, as ManualAlignmentDialog; it is now the +alignment wizard's first page (see align_wizard.py), which needed the same +mask-overlay editor plus the crop and export steps. +""" -import matplotlib as mpl -import numpy as np -from matplotlib.backends.backend_qtagg import NavigationToolbar2QT -from PyQt6.QtCore import QSignalBlocker, pyqtSignal from PyQt6.QtWidgets import ( - QButtonGroup, QComboBox, QDialog, QDialogButtonBox, QGroupBox, - QHBoxLayout, QLabel, QMessageBox, QPushButton, QRadioButton, QSpinBox, - QVBoxLayout, QWidget, + QButtonGroup, QDialog, QDialogButtonBox, QGroupBox, QHBoxLayout, QLabel, + QRadioButton, QSpinBox, QVBoxLayout, ) -import sras_compute as compute -from sras_compute import ( - PYFFTW_AVAILABLE, ManualAngleParams, build_manual_alignment, - delete_manual_alignment, save_manual_alignment, -) -from sras_format import SrasFile -from sras_workers import Ch4MaskWorker, CrossCorrelateWorker +from sras_compute import PYFFTW_AVAILABLE -from .canvases import AlignOverlayCanvas -from .common import ( - _CSS_HINT, _CSS_MUTED, _CSS_WARN, Jobs, _axes_extent, _form, _group, _make_dspin, - _scroll_panel, _wrap_label, -) - -if TYPE_CHECKING: - from .main_window import SrasViewerWindow +from .common import _CSS_HINT, _make_dspin # --------------------------------------------------------------------------- @@ -250,588 +235,3 @@ class RowAverageFftOptionsDialog(QDialog): def get_threshold_mv(self) -> float: return self._spin_threshold.value() - - -class ManualAlignmentDialog(QDialog): - """Non-modal manual angle-alignment editor (Fusion -> Manual Alignment...). - - Shows every angle's binarized CH4 (Bias B) mask overlaid in a distinct - color at partial opacity on one shared canvas, so translation/rotation - 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 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 - SrasViewerWindow two ways: it reuses parent._run_worker/_jobs directly - for its background mask-fetch and cross-correlate steps, so the main - window's existing shutdown/lifecycle plumbing covers both for free, and - it emits alignment_saved / alignment_cleared signals for the two moments - that should actually mutate the main window's persistent state — - everything else (nudging, Auto De-rotate, Auto Cross-Correlate, threshold - edits) stays purely local to this dialog until Save. - """ - - alignment_saved = pyqtSignal(object, str) # AlignmentResult, sidecar path (str) - alignment_cleared = pyqtSignal() - - _PREVIEW_MARGIN_FRAC = 0.15 - _BASE_ALPHA = 0.42 - _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, - cached_dc4_mv: dict[int, np.ndarray]): - super().__init__(parent) - self._parent = parent - self._sras = sras - self._ref_angle_idx = ref_angle_idx - self._downsample = (1, 1) # (rows, cols) block-mean factors - self._dc4_mv: dict[int, np.ndarray] = {} - self._masks_small: dict[int, np.ndarray] = {} - self._preview_layers: dict[int, np.ndarray] = {} - self._preview_origin_mm = (0.0, 0.0) - self._preview_shape = (1, 1) - 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) - - self._seed_initial_params(seed_per_angle) - n = sras.n_angles - cmap = mpl.colormaps["tab10"] if n <= 10 else mpl.colormaps["tab20"] - self._angle_colors = {a: cmap(a % cmap.N)[:3] for a in range(n)} - self._active_angle = 1 if ref_angle_idx == 0 and n > 1 else 0 - - self._build_ui(dc_threshold_mv) - self._set_controls_enabled(False) # re-enabled once masks are ready - self._start_mask_prep(cached_dc4_mv) - - def showEvent(self, event): - super().showEvent(event) - self.canvas.setFocus() - - # ------------------------------------------------------------------ - # Construction - # ------------------------------------------------------------------ - - def _seed_initial_params(self, seed_per_angle: dict[int, ManualAngleParams] | None): - seed = seed_per_angle or {} - self._angle_params: dict[int, ManualAngleParams] = { - a: (ManualAngleParams(seed[a].rotation_deg, seed[a].shift_mm) - if a in seed else ManualAngleParams()) - for a in range(self._sras.n_angles) - } - self._angle_params[self._ref_angle_idx] = ManualAngleParams() - - def _build_ui(self, dc_threshold_mv: float): - root = QHBoxLayout(self) - - self.canvas = AlignOverlayCanvas() - left = QWidget() - left_l = QVBoxLayout(left) - left_l.setContentsMargins(0, 0, 0, 0) - left_l.setSpacing(4) - left_l.addWidget(NavigationToolbar2QT(self.canvas, left)) - left_l.addWidget(self.canvas) - root.addWidget(left, stretch=1) - - panel = QWidget() - panel_l = QVBoxLayout(panel) - panel_l.setContentsMargins(0, 0, 0, 0) - panel_l.setSpacing(8) - panel_l.addWidget(self._build_angle_group()) - panel_l.addWidget(self._build_adjust_group()) - panel_l.addWidget(self._build_step_group()) - panel_l.addWidget(self._build_threshold_group(dc_threshold_mv)) - panel_l.addWidget(self._build_correlate_group()) - panel_l.addWidget(self._build_actions_group()) - self.lbl_status = _wrap_label("", _CSS_MUTED) - panel_l.addWidget(self.lbl_status) - panel_l.addStretch() - - root.addWidget(_scroll_panel(panel, 320)) - self._connect_controls() - - def _build_angle_group(self) -> QWidget: - grp_angle, al = _group("Active Angle") - self.combo_active_angle = QComboBox() - for a in range(self._sras.n_angles): - label = f"Angle {a} ({self._sras.angles_deg[a]:.1f}°)" - if a == self._ref_angle_idx: - label += " [reference]" - self.combo_active_angle.addItem(label) - al.addWidget(self.combo_active_angle) - self.lbl_active_note = _wrap_label("", _CSS_WARN) - al.addWidget(self.lbl_active_note) - return grp_angle - - def _build_adjust_group(self) -> QWidget: - self.grp_manual_adjust, mform_box = _group("Manual Adjustment") - mform = _form() - self.spin_active_rotation_deg = _make_dspin(-3600.0, 3600.0, 3, suffix=" °") - mform.addRow("Rotation:", self.spin_active_rotation_deg) - - self.spin_active_shift_x_mm = _make_dspin(-1e5, 1e5, 4, suffix=" mm") - mform.addRow("Shift X:", self.spin_active_shift_x_mm) - - self.spin_active_shift_y_mm = _make_dspin(-1e5, 1e5, 4, suffix=" mm") - mform.addRow("Shift Y:", self.spin_active_shift_y_mm) - mform_box.addLayout(mform) - return self.grp_manual_adjust - - def _build_step_group(self) -> QWidget: - self.grp_step_sizes, sl = _group("Nudge Step Sizes") - sform = _form() - self.spin_step_translate_mm = _make_dspin(0.0001, 1000.0, 4, - suffix=" mm", value=0.01) - sform.addRow("Translate step:", self.spin_step_translate_mm) - - self.spin_step_rotate_deg = _make_dspin(0.001, 90.0, 3, - suffix=" °", value=0.1) - sform.addRow("Rotate step:", self.spin_step_rotate_deg) - - self.spin_step_multiplier = _make_dspin(1.0, 1000.0, 1, value=10.0) - sform.addRow("Coarse × (Shift):", self.spin_step_multiplier) - sl.addLayout(sform) - sl.addWidget(_wrap_label( - "Arrow keys nudge X/Y translation; Q/E nudge rotation (CCW/CW). " - "Hold Shift for the coarse step. Click the image once so it has " - "keyboard focus.", _CSS_HINT)) - return self.grp_step_sizes - - def _build_threshold_group(self, dc_threshold_mv: float) -> QWidget: - self.grp_mask_threshold, tl = _group("Mask Threshold") - tform = _form() - self.spin_mask_threshold_mv = _make_dspin(-500.0, 500.0, 3, - suffix=" mV", value=dc_threshold_mv) - tform.addRow("DC threshold:", self.spin_mask_threshold_mv) - tl.addLayout(tform) - return self.grp_mask_threshold - - def _build_correlate_group(self) -> QWidget: - self.grp_correlate, cl = _group("Cross-Correlate (FFT)") - cform = _form() - self.combo_correlate_source = QComboBox() - 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_search_deg = _make_dspin(0.0, 180.0, 1, suffix=" °", - value=6.0, step=1.0) - 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( - "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)) - return self.grp_correlate - - def _build_actions_group(self) -> QWidget: - grp_actions, acl = _group("Actions") - self.btn_auto_derotate = QPushButton("Auto De-rotate (use known angles)") - self.btn_save = QPushButton("Save Alignment") - self.btn_clear = QPushButton("Clear Alignment…") - self.btn_close = QPushButton("Close") - for btn in (self.btn_auto_derotate, self.btn_save, self.btn_clear, self.btn_close): - acl.addWidget(btn) - return grp_actions - - def _connect_controls(self): - self.combo_active_angle.currentIndexChanged.connect(self._on_active_angle_changed) - self.spin_active_rotation_deg.editingFinished.connect(self._on_rotation_spin_edited) - self.spin_active_shift_x_mm.editingFinished.connect(self._on_shift_spin_edited) - self.spin_active_shift_y_mm.editingFinished.connect(self._on_shift_spin_edited) - self.spin_mask_threshold_mv.editingFinished.connect(self._on_mask_threshold_edited) - self.btn_auto_derotate.clicked.connect(self._on_auto_derotate) - self.btn_auto_correlate.clicked.connect(self._on_auto_correlate) - self.btn_save.clicked.connect(self._on_save) - self.btn_clear.clicked.connect(self._on_clear) - self.btn_close.clicked.connect(self.close) - self.canvas.nudge_translate.connect(self._on_nudge_translate) - self.canvas.nudge_rotate.connect(self._on_nudge_rotate) - - with QSignalBlocker(self.combo_active_angle): - self.combo_active_angle.setCurrentIndex(self._active_angle) - self._on_active_angle_changed(self._active_angle) - - # ------------------------------------------------------------------ - # Mask preparation (initial CH4 fetch + threshold + downsample) - # ------------------------------------------------------------------ - - def _start_mask_prep(self, cached_dc4_mv: dict[int, np.ndarray]): - self._dc4_mv = dict(cached_dc4_mv) - missing = [a for a in range(self._sras.n_angles) if a not in self._dc4_mv] - if not missing: - self._finish_mask_prep() - return - self.lbl_status.setText(f"Preparing masks: 0/{len(missing)} angle(s) needed…") - started = self._parent._run_worker( - Jobs.MANUAL_ALIGN_MASKS, Ch4MaskWorker(self._sras, missing), - connect=( - ("angle_done", self._on_mask_angle_done), - ("error", lambda msg: self.lbl_status.setText(f"Mask prep error: {msg}")), - ), - on_done=self._finish_mask_prep) - if not started: - self.lbl_status.setText( - "Could not start mask preparation (busy) — close and reopen.") - - def _on_mask_angle_done(self, angle_idx: int, dc4_mv: np.ndarray): - self._dc4_mv[angle_idx] = dc4_mv - self.lbl_status.setText( - f"Preparing masks: {len(self._dc4_mv)}/{self._sras.n_angles} ready…") - - def _finish_mask_prep(self): - if len(self._dc4_mv) < self._sras.n_angles: - return # a mask-worker error left some angles unfetched - # 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() - self._rebuild_preview_canvas() - self._set_controls_enabled(True) - self.lbl_status.setText("Ready.") - - def _recompute_masks_small(self): - """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: no alignment geometry depends on this - threshold, only which pixels the overlay paints.""" - threshold = self.spin_mask_threshold_mv.value() - fy, fx = self._downsample - self._masks_small = { - a: compute.block_mean_2d((img >= threshold).astype(np.float32), fy, fx) - for a, img in self._dc4_mv.items() - } - - # ------------------------------------------------------------------ - # Preview canvas: full rebuild vs. incremental single-layer refresh - # ------------------------------------------------------------------ - - def _rebuild_preview_canvas(self): - """Full geometry rebuild: recomputes the shared preview canvas's - origin/shape (rotation can grow the union bbox — translation alone - cannot, per the padding baked in via _PREVIEW_MARGIN_FRAC) and every - angle's reprojected mask layer. Triggered by: dialog open, - mask-threshold change, Auto De-rotate, a rotation nudge/edit of the - 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) - 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_pitch_mm = pitch - self._preview_layers = { - 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.""" - self._preview_layers[self._active_angle] = self._reproject(self._active_angle) - self._redraw_overlay() - - def _redraw_overlay(self): - """Alpha-composite every angle's colored mask layer into one RGBA - image ("all thresholds overlaid with varying opacity"). Each angle - keeps a fixed, distinct color regardless of which is active; the - active angle is drawn last (on top) at a visibly higher alpha so - it's easy to track while nudging.""" - if not self._preview_layers: - return # mask prep hasn't finished yet — nothing to draw - n_rows, n_cols = self._preview_shape - rgba = np.zeros((n_rows, n_cols, 4), dtype=np.float32) - order = sorted(range(self._sras.n_angles), key=lambda a: a == self._active_angle) - for a in order: - layer = self._preview_layers.get(a) - if layer is None: - continue - alpha = self._ACTIVE_ALPHA if a == self._active_angle else self._BASE_ALPHA - color = self._angle_colors[a] - fg_a = layer * alpha - for c in range(3): - rgba[..., c] = color[c] * fg_a + rgba[..., c] * rgba[..., 3] * (1 - fg_a) - rgba[..., 3] = fg_a + rgba[..., 3] * (1 - fg_a) - - x0, y0 = self._preview_origin_mm - dx, dy = self._preview_pitch_mm - x_axis = x0 + np.arange(n_cols) * dx - y_axis = y0 + np.arange(n_rows) * dy - extent = _axes_extent(x_axis, y_axis, dx, dy) - title = (f"Angle {self._active_angle} active " - f"({self._sras.angles_deg[self._active_angle]:.1f}°)") - self.canvas.show_overlay(rgba, extent, title) - - # ------------------------------------------------------------------ - # Angle selection / nudge / edit handlers - # ------------------------------------------------------------------ - - def _on_active_angle_changed(self, angle_idx: int): - self._active_angle = angle_idx - is_ref = angle_idx == self._ref_angle_idx - self.grp_manual_adjust.setEnabled(self._masks_ready and not is_ref) - self.lbl_active_note.setText( - "Reference angle — defines the shared origin, not adjustable." if is_ref else "") - self._sync_active_spinboxes() - self._redraw_overlay() - - def _sync_active_spinboxes(self): - p = self._angle_params[self._active_angle] - for spin, val in ((self.spin_active_rotation_deg, p.rotation_deg), - (self.spin_active_shift_x_mm, p.shift_mm[0]), - (self.spin_active_shift_y_mm, p.shift_mm[1])): - with QSignalBlocker(spin): - spin.setValue(val) - - def _on_nudge_translate(self, dir_x: int, dir_y: int, coarse: bool): - if not self._masks_ready or self._active_angle == self._ref_angle_idx: - return - step = self.spin_step_translate_mm.value() - if coarse: - step *= self.spin_step_multiplier.value() - p = self._angle_params[self._active_angle] - p.shift_mm = (p.shift_mm[0] + dir_x * step, p.shift_mm[1] + dir_y * step) - self._sync_active_spinboxes() - self._refresh_active_preview_layer() - - def _on_nudge_rotate(self, direction: int, coarse: bool): - if not self._masks_ready or self._active_angle == self._ref_angle_idx: - return - step = self.spin_step_rotate_deg.value() - if coarse: - step *= self.spin_step_multiplier.value() - self._angle_params[self._active_angle].rotation_deg += direction * step - self._sync_active_spinboxes() - self._rebuild_preview_canvas() - - def _on_rotation_spin_edited(self): - if self._active_angle == self._ref_angle_idx: - return - self._angle_params[self._active_angle].rotation_deg = self.spin_active_rotation_deg.value() - self._rebuild_preview_canvas() - - def _on_shift_spin_edited(self): - if self._active_angle == self._ref_angle_idx: - return - p = self._angle_params[self._active_angle] - p.shift_mm = (self.spin_active_shift_x_mm.value(), self.spin_active_shift_y_mm.value()) - self._refresh_active_preview_layer() - - def _on_mask_threshold_edited(self): - if not self._masks_ready: - return - self._recompute_masks_small() - self._rebuild_preview_canvas() - - # ------------------------------------------------------------------ - # Actions - # ------------------------------------------------------------------ - - 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 = 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 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: - return - angles = [a for a in range(self._sras.n_angles) if a != self._ref_angle_idx] - if not angles: - return - worker = CrossCorrelateWorker( - 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( - Jobs.MANUAL_ALIGN_CORRELATE, worker, - connect=( - ("angle_done", self._on_correlate_angle_done), - ("error", self._on_correlate_error), - ), - on_done=self._finish_auto_correlate) - if not started: - self._set_controls_enabled(True) - 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, - 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)…") - - def _on_correlate_error(self, msg: str): - self.lbl_status.setText(f"Cross-correlation error: {msg}") - - def _finish_auto_correlate(self): - self._sync_active_spinboxes() - self._rebuild_preview_canvas() - self._set_controls_enabled(True) - self.lbl_status.setText( - f"Cross-correlated {self._correlate_done_count} angle(s) against " - 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) - except OSError as exc: - QMessageBox.warning(self, "Save Alignment Failed", str(exc)) - return - self.lbl_status.setText(f"Saved to {path.name}.") - self.alignment_saved.emit(result, str(path)) - - def _on_clear(self): - reply = QMessageBox.question( - self, "Clear Alignment", - "This resets every angle back to raw/unaligned (0° rotation, no " - "shift) and deletes the saved alignment file for this scan, if " - "any. This cannot be undone. Continue?", - QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, - QMessageBox.StandardButton.No) - if reply != QMessageBox.StandardButton.Yes: - return - try: - existed = delete_manual_alignment(self._sras) - except OSError as exc: - QMessageBox.warning(self, "Clear Alignment Failed", - 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( - "Alignment cleared; saved file removed." if existed - else "Alignment cleared (there was no saved file).") - self.alignment_cleared.emit() - - def _set_controls_enabled(self, enabled: bool): - self._masks_ready = enabled - self.combo_active_angle.setEnabled(enabled) - self.grp_manual_adjust.setEnabled(enabled and self._active_angle != self._ref_angle_idx) - self.grp_step_sizes.setEnabled(enabled) - self.grp_mask_threshold.setEnabled(enabled) - self.grp_correlate.setEnabled(enabled) - self.btn_auto_derotate.setEnabled(enabled) - self.btn_save.setEnabled(enabled) - self.btn_clear.setEnabled(enabled) - - diff --git a/sras_viewer/main_window.py b/sras_viewer/main_window.py index 6401603..6cee70c 100644 --- a/sras_viewer/main_window.py +++ b/sras_viewer/main_window.py @@ -16,15 +16,14 @@ from PyQt6.QtWidgets import ( import sras_compute as compute from sras_compute import ( ManualAngleParams, apply_alignment, build_manual_alignment, - load_manual_alignment, sidecar_path, + load_manual_alignment, save_manual_alignment, sidecar_path, ) from sras_format import ( CH3_IDX, CH4_IDX, CH_NAMES, SrasFile, mv_to_adc, _FALLBACK_YMULT_MV, _FALLBACK_YOFF_ADC, ) from sras_workers import ( - AngleAlignmentWorker, BatchCacheWorker, ComputeWorker, DcPrecomputeWorker, - LoadWorker, + BatchCacheWorker, ComputeWorker, DcPrecomputeWorker, LoadWorker, ) from .canvases import ImageCanvas, WaveformCanvas @@ -33,7 +32,8 @@ from .common import ( _CSS_BUSY, _axes_extent, _CSS_HINT, _CSS_INFO, _CSS_MUTED, _CSS_WARN, _LEFT_PANEL_W, _RIGHT_PANEL_W, Jobs, _form, _group, _make_dspin, _scroll_panel, _wrap_label, ) -from .dialogs import FftOptionsDialog, ManualAlignmentDialog, RowAverageFftOptionsDialog +from .align_wizard import AlignmentWizard +from .dialogs import FftOptionsDialog, RowAverageFftOptionsDialog # --------------------------------------------------------------------------- # Main window @@ -99,7 +99,7 @@ class SrasViewerWindow(QMainWindow): self._alignment_result = None self._alignment_generation: int = 0 self._aligned_cache: dict[tuple, np.ndarray] = {} - self._manual_align_dialog: ManualAlignmentDialog | None = None + self._align_wizard: AlignmentWizard | None = None self._build_ui() @@ -430,23 +430,13 @@ class SrasViewerWindow(QMainWindow): fft_menu.addAction(fft_act) fusion_menu = menubar.addMenu("&Fusion") - self._alignment_act = QAction("Angle &Alignment", self) - self._alignment_act.setStatusTip( - "Compute a rotation+translation alignment across all angles " - "(from CH4 masks) and enable Aligned View. Requires >1 angle.") - self._alignment_act.setEnabled(False) - self._alignment_act.triggered.connect(self._on_angle_alignment) - fusion_menu.addAction(self._alignment_act) - - self._manual_align_act = QAction("&Manual Alignment…", self) - self._manual_align_act.setStatusTip( - "Open an interactive dialog to align angles by eye: overlaid CH4 " - "threshold masks, keyboard nudge (translate + rotate), auto " - "de-rotate to the known scan angles, and save/clear a persistent " - "alignment.") - self._manual_align_act.setEnabled(False) - self._manual_align_act.triggered.connect(self._on_manual_alignment) - fusion_menu.addAction(self._manual_align_act) + self._wizard_act = QAction("Alignment &Wizard…", self) + self._wizard_act.setStatusTip( + "Align the angles, crop to a region of interest, and save the " + "aligned data as a new .sras file. Requires >1 angle.") + self._wizard_act.setEnabled(False) + self._wizard_act.triggered.connect(self._on_alignment_wizard) + fusion_menu.addAction(self._wizard_act) convert_menu = menubar.addMenu("&Convert") self._batch_dc_act = QAction("Batch Compute DC and &Store…", self) @@ -522,9 +512,9 @@ class SrasViewerWindow(QMainWindow): # A manual-alignment dialog bound to the previous file must not # survive a reload — its per-angle state (and the sras it was # constructed against) no longer matches the new file's geometry. - if self._manual_align_dialog is not None: - self._manual_align_dialog.close() - self._manual_align_dialog = None + if self._align_wizard is not None: + self._align_wizard.close() + self._align_wizard = None # Caches (and any in-flight DC precompute) belong to the previous # file's geometry — discard and start fresh. Bumping the generation @@ -661,10 +651,8 @@ class SrasViewerWindow(QMainWindow): self._batch_fft_act.setEnabled(can_batch) self._batch_fft_rowavg_act.setEnabled(can_batch) - self._alignment_act.setEnabled( - has_file and s.n_angles > 1 and not self._job_running(Jobs.ALIGN)) - self._manual_align_act.setEnabled( - has_file and s.n_angles > 1 and not self._job_running(Jobs.ALIGN)) + self._wizard_act.setEnabled( + has_file and s.n_angles > 1 and self._align_wizard is None) self.chk_aligned_view.setEnabled(enabled and self._alignment_result is not None) self._update_roi_ui() @@ -1143,27 +1131,16 @@ class SrasViewerWindow(QMainWindow): # Progress dialogs # ------------------------------------------------------------------ - def _show_progress(self, key: str, message: str, maximum: int = 0, - on_cancel=None): + def _show_progress(self, key: str, message: str, maximum: int = 0): """Show (or relabel) the progress dialog under *key*. maximum=0 gives - an indeterminate busy indicator. - - *on_cancel* adds a Cancel button wired to it. Almost every job here is - short enough that a cancel button would only invite a click that does - nothing, hence the no-button default; the aligned export is the - exception, since it can spend minutes writing gigabytes. - """ + an indeterminate busy indicator.""" dlg = self._progress_dlgs.get(key) if dlg is not None: dlg.setLabelText(message) return dlg = QProgressDialog(message, "", 0, maximum, self) dlg.setWindowTitle("Please wait…") - if on_cancel is None: - dlg.setCancelButton(None) - else: - dlg.setCancelButtonText("Cancel") - dlg.canceled.connect(on_cancel) + dlg.setCancelButton(None) dlg.setWindowModality(Qt.WindowModality.WindowModal) dlg.setMinimumDuration(300) # only appears if it takes > 300 ms dlg.show() @@ -1289,90 +1266,68 @@ class SrasViewerWindow(QMainWindow): # Fusion: angle alignment # ------------------------------------------------------------------ - def _on_angle_alignment(self): + def _on_alignment_wizard(self): if self._sras is None or self._sras.n_angles <= 1: return - ref_idx = 0 - threshold_mv = self.spin_threshold_mv.value() - generation = self._alignment_generation - - started = self._run_worker( - Jobs.ALIGN, AngleAlignmentWorker(self._sras, ref_idx, threshold_mv), - connect=( - ("progress", lambda pct: self._set_progress("main", pct)), - ("finished", lambda result, err, g=generation: - self._on_alignment_done(g, result, err)), - ), - on_done=lambda: self._update_controls_enabled(self._sras is not None), - ) - if not started: - return - - self._alignment_act.setEnabled(False) - self._show_progress( - "main", - f"Computing angle alignment ({self._sras.n_angles} angles, " - f"ref=angle 0, CH4 mask ≥ {threshold_mv:.3f} mV)…", - maximum=100) - - def _on_alignment_done(self, generation: int, result, error_msg: str): - self._close_progress("main") - if generation != self._alignment_generation: - return # a new file was loaded while this was computing — discard - if error_msg: - self.statusBar().showMessage(f"Angle alignment failed: {error_msg}") - return - # No generation bump: this result *is* the current generation's. - self._apply_alignment_result(result, view_checked=True, - bump_generation=False) - nr, nc = result.canvas_shape - self.statusBar().showMessage( - f"Angle alignment computed ({self._sras.n_angles} angles, " - f"canvas {nc}×{nr} px).") - self._refresh_display() - - # ------------------------------------------------------------------ - # Fusion: manual alignment - # ------------------------------------------------------------------ - - def _on_manual_alignment(self): - if self._sras is None or self._sras.n_angles <= 1: - return - if self._manual_align_dialog is not None: - self._manual_align_dialog.raise_() - self._manual_align_dialog.activateWindow() + if self._align_wizard is not None: + self._align_wizard.raise_() + self._align_wizard.activateWindow() return ref_idx = 0 threshold_mv = self.spin_threshold_mv.value() seed: dict[int, ManualAngleParams] = {} - # Seed only from a previously *saved manual* alignment (this dialog's - # own Save also writes this sidecar) -- never from self._alignment_result - # when it holds the automatic Fusion -> Angle Alignment's output. That - # path's translation comes from FFT phase correlation, which is the - # very thing manual mode exists to work around; inheriting it here - # would silently reintroduce the same bad translations under a - # "manual" label, on top of the (correct) analytic rotation, which is - # exactly what makes manual mode look like it "still does the same - # thing" the automatic one does. + # Seed only from a previously *saved* alignment, never from + # self._alignment_result: the wizard's first page starts every angle + # pre-rotated from the stage angles and expects to own the parameters + # from there, so inheriting a half-edited in-memory state would make + # "Reset to pre-rotation only" mean something different each time. sidecar = load_manual_alignment(self._sras) if sidecar is not None and sidecar.ref_angle_idx == ref_idx: seed = dict(sidecar.per_angle) threshold_mv = sidecar.dc_threshold_mv cached_dc4 = {a: img for (a, ch), img in self._dc_cache.items() if ch == CH4_IDX} - dlg = ManualAlignmentDialog( + wiz = AlignmentWizard( self, self._sras, ref_angle_idx=ref_idx, dc_threshold_mv=threshold_mv, seed_per_angle=seed, cached_dc4_mv=cached_dc4) - dlg.alignment_saved.connect(self._on_manual_alignment_saved) - dlg.alignment_cleared.connect(self._on_manual_alignment_cleared) - dlg.finished.connect(self._on_manual_align_dialog_closed) - dlg.setAttribute(Qt.WidgetAttribute.WA_DeleteOnClose) - self._manual_align_dialog = dlg - dlg.show() + wiz.alignment_ready.connect(self._on_wizard_finished) + wiz.finished.connect(self._on_wizard_closed) + wiz.setAttribute(Qt.WidgetAttribute.WA_DeleteOnClose) + self._align_wizard = wiz + self._update_controls_enabled(True) + wiz.show() - def _on_manual_align_dialog_closed(self, _result_code: int): - self._manual_align_dialog = None + def _on_wizard_closed(self, _result_code: int): + self._align_wizard = None + self._update_controls_enabled(self._sras is not None) + + def _on_wizard_finished(self, result, out_path: str): + """The wizard exported a file; make the session match what was written. + + Applies the *cropped* result, so Aligned View shows exactly the extent + that went into the file rather than the wider uncropped canvas, and + saves the sidecar for the input scan so reopening it lands in the same + place. + """ + if result is None: + return + self._apply_alignment_result(result, view_checked=True) + note = "" + try: + save_manual_alignment( + self._sras, result.ref_angle_idx, result.dc_threshold_mv, + {a: ManualAngleParams(t.rotation_deg, t.shift_mm) + for a, t in result.per_angle.items()}) + except OSError as exc: + note = f" (could not save the sidecar: {exc})" + self._update_controls_enabled(self._sras is not None) + nr, nc = result.canvas_shape + self.statusBar().showMessage( + f"Aligned scan written to {Path(out_path).name}; Aligned View now " + f"shows the exported {nc}×{nr} px region.{note}") + if self._current_image is not None: + self._refresh_display() def _apply_alignment_result(self, result, *, view_checked: bool, bump_generation: bool = True): @@ -1387,20 +1342,6 @@ class SrasViewerWindow(QMainWindow): self.chk_aligned_view.setChecked(view_checked) self.chk_aligned_view.setEnabled(result is not None) - def _manual_alignment_changed(self, result, message: str): - self._apply_alignment_result(result, view_checked=result is not None) - self._update_controls_enabled(self._sras is not None) - self.statusBar().showMessage(message) - if self._current_image is not None: - self._refresh_display() - - def _on_manual_alignment_saved(self, result, sidecar_path_str: str): - self._manual_alignment_changed( - result, f"Manual alignment saved to {Path(sidecar_path_str).name}") - - def _on_manual_alignment_cleared(self): - self._manual_alignment_changed(None, "Manual alignment cleared.") - # ------------------------------------------------------------------ # FFT Options # ------------------------------------------------------------------ @@ -1429,8 +1370,8 @@ class SrasViewerWindow(QMainWindow): # ------------------------------------------------------------------ def closeEvent(self, event): - if self._manual_align_dialog is not None: - self._manual_align_dialog.close() + if self._align_wizard is not None: + self._align_wizard.close() # Signal every cancellable worker first, then wait. Waiting without # signalling means sitting out whatever is in flight — on a large diff --git a/sras_workers.py b/sras_workers.py index 237b37e..0931eab 100644 --- a/sras_workers.py +++ b/sras_workers.py @@ -16,9 +16,7 @@ from PyQt6.QtCore import QObject, pyqtSignal import sras_compute as compute from sras_align_export import write_aligned_sras -from sras_compute import ( - cache_file, compute_angle_alignment, compute_rf_image, dc_image_mv, -) +from sras_compute import cache_file, compute_rf_image, dc_image_mv from sras_format import CH3_IDX, CH4_IDX, SrasFile # Concurrency caps. Batch conversion runs one process per file, and each of @@ -302,34 +300,9 @@ class BatchCacheWorker(QObject): self.finished.emit() -class AngleAlignmentWorker(QObject): - """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) - - def __init__(self, sras: SrasFile, ref_angle_idx: int, dc_threshold_mv: float): - super().__init__() - self._sras = sras - self._ref = ref_angle_idx - self._threshold = dc_threshold_mv - - def run(self): - try: - result = compute_angle_alignment( - self._sras, self._ref, self._threshold, - progress_cb=self.progress.emit) - self.finished.emit(result, "") - except Exception as exc: - self.finished.emit(None, str(exc)) - - class Ch4MaskWorker(_PooledWorker): - """Fetches each requested angle's CH4 (Bias B) DC image in mV, for - ManualAlignmentDialog's initial threshold-mask overlay. + """Fetches each requested angle's CH4 (Bias B) DC image in mV, for the + alignment wizard's initial threshold-mask stack. Reuses dc_image_mv, which prefers a stored v5/v7 cache over recomputing from raw waveforms, so this only does real work for a file that hasn't @@ -338,7 +311,7 @@ class Ch4MaskWorker(_PooledWorker): every file load) hasn't reached yet. In the common case — the user opens Fusion -> Manual Alignment after DC precompute has already finished — *angle_indices* is empty and this worker is never even constructed (see - ManualAlignmentDialog._start_mask_prep). + CorrelatePage._start_mask_prep). """ angle_done = pyqtSignal(int, np.ndarray) # angle_idx, dc4_mv @@ -365,8 +338,8 @@ class Ch4MaskWorker(_PooledWorker): class CrossCorrelateWorker(_PooledWorker): """Rigid registration (rotation + translation, never scale) of each of - *angle_indices* against *ref_angle_idx*, for ManualAlignmentDialog's Auto - Cross-Correlate button. + *angle_indices* against *ref_angle_idx*, for the alignment wizard's + Run/Re-run Correlation button. Runs on a background thread — registering a real many-angle, high-resolution scan takes long enough that doing it on the GUI thread diff --git a/tests/test_align_export.py b/tests/test_align_export.py index 0e6326c..0cf1360 100644 --- a/tests/test_align_export.py +++ b/tests/test_align_export.py @@ -406,6 +406,40 @@ def test_largest_rect_at_least(): assert compute.largest_rect_at_least(np.full((3, 4), 2), 2) == (0, 0, 3, 4) +def test_largest_rect_matches_brute_force(): + """Randomized check against an O(n^4) reference. + + The histogram sweep is short and easy to get subtly wrong — an off-by-one in + the stack unwind yields rectangles that are merely large, and "large but not + maximal" is invisible by eye on real data. + """ + def brute(good): + n_rows, n_cols = good.shape + best = 0 + for r0 in range(n_rows): + for r1 in range(r0 + 1, n_rows + 1): + run = 0 + for g in good[r0:r1].all(axis=0): + run = run + 1 if g else 0 + best = max(best, run * (r1 - r0)) + return best + + rng = np.random.default_rng(0) + for _ in range(200): + counts = rng.integers(0, 3, size=(int(rng.integers(1, 9)), + int(rng.integers(1, 9)))) + got = compute.largest_rect_at_least(counts, 2) + expected = brute(counts >= 2) + if got is None: + assert expected == 0 + continue + row0, col0, nr, nc = got + assert (counts[row0:row0 + nr, col0:col0 + nc] >= 2).all(), \ + f"rectangle is not pure:\n{counts}\n{got}" + assert nr * nc == expected, \ + f"not maximal ({nr * nc} < {expected}):\n{counts}\n{got}" + + def test_largest_rect_is_pure_on_the_real_fixture(rig): """On real overlap counts the returned rectangle must contain only full-overlap pixels — the property a bounding box would violate.""" diff --git a/tests/test_alignment.py b/tests/test_alignment.py index 04bb85f..d9ac0e2 100644 --- a/tests/test_alignment.py +++ b/tests/test_alignment.py @@ -155,7 +155,7 @@ def test_all_angles_stack(rig): def test_downsampled_preview_lands_with_full_res(rig): - # ManualAlignmentDialog reprojects block-mean-downsampled masks, so the + # The wizard 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. diff --git a/tests/test_gui.py b/tests/test_gui.py index ada0b73..599936c 100644 --- a/tests/test_gui.py +++ b/tests/test_gui.py @@ -3,8 +3,8 @@ signals and worker threads under the offscreen platform plugin. Covers the interactions a manual smoke test would: load, switch angles and channels, background DC precompute, lazy FFT compute, threshold and bg-sub -changes, angle alignment, manual angle alignment, aligned view, ROI -draw/move, and CSV export. +changes, the alignment wizard end to end (pre-rotation, correlation, manual +nudging, crop, export), aligned view, ROI draw/move, and CSV export. NOTE: this module is one ordered integration sequence over a single shared window — the tests build on each other's state and must run in definition @@ -218,44 +218,16 @@ def test_roi_survives_switches(ctx): "ROI still present after channel switch" -def test_angle_alignment(ctx): - win, s = ctx.win, ctx.s - win.spin_angle.setValue(0) - win._on_view_changed() - wait_until(lambda: not win._job_running("compute")) - assert win._alignment_act.isEnabled(), "alignment action enabled" - win._on_angle_alignment() - assert wait_until( - lambda: win._alignment_result is not None and not win._job_running("align"), - timeout_ms=60000), "alignment completed" - - r = win._alignment_result - assert len(r.per_angle) == s.n_angles, "transform for every angle" - assert all(r.canvas_shape[0] >= int(s.n_rows[a]) - and r.canvas_shape[1] >= int(s.n_frames[a]) - for a in range(s.n_angles)), \ - f"canvas is at least as large as any single angle: {r.canvas_shape}" - assert r.per_angle[r.ref_angle_idx].shift_mm == (0.0, 0.0), \ - "reference angle has zero shift" - assert win.chk_aligned_view.isEnabled() and win.chk_aligned_view.isChecked(), \ - "Aligned View auto-enabled and checked" - pump(200) - assert win.image_canvas._img_shape == r.canvas_shape, \ - f"{win.image_canvas._img_shape} vs {r.canvas_shape}" - - win.chk_aligned_view.setChecked(False) - pump(200) - assert win.image_canvas._img_shape == s.image_shape(0), \ - f"unchecking returns to the raw per-angle grid: {win.image_canvas._img_shape}" - - -def test_manual_alignment_geometry(ctx): +def test_alignment_geometry_is_stage_independent(ctx): """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 tests/test_alignment.py, which has a synthetic sample to register.)""" win, s = ctx.win, ctx.s - assert win._manual_align_act.isEnabled(), "manual alignment action enabled" + win.spin_angle.setValue(0) + win._on_view_changed() + wait_until(lambda: not win._job_running("compute")) + assert win._wizard_act.isEnabled(), "alignment wizard action enabled" n_rows, n_frames = s.image_shape(0) assert np.allclose(compute._center_idx(s, 0), @@ -276,7 +248,7 @@ def test_manual_alignment_geometry(ctx): ("moving every non-reference angle's scan window must leave the canvas " f"unchanged: {origin_a} {shape_a} vs {origin_b} {shape_b}") - # Both signs of the stage's reported angle are searched. + # Both signs of the stage's reported angle are searched by default. cands = compute._rotation_candidates(30.0, 6.0, 2.0) assert min(cands) < -29.0 and max(cands) > 29.0, f"{min(cands)}..{max(cands)}" @@ -289,128 +261,285 @@ def test_manual_alignment_geometry(ctx): "_shift_into moves content by exactly the requested offset" -def test_manual_dialog_opens_at_identity(ctx): - """Open must NOT seed from the still-live automatic AlignmentResult. - 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.""" +def test_wizard_opens_prerotated(ctx): + """The wizard shows a mask stack before any correlation has run, built from + the stage angles in the file — the "pre-rotate" step. Nothing may seed from + the still-live automatic result; only a saved sidecar.""" win, s = ctx.win, ctx.s - win._on_manual_alignment() - assert win._manual_align_dialog is not None, "dialog opened" - ctx.dlg = dlg = win._manual_align_dialog - assert not win._job_running("manual_align_masks"), \ + win._on_alignment_wizard() + assert win._align_wizard is not None, "wizard opened" + ctx.wiz = wiz = win._align_wizard + ctx.p1 = p1 = wiz.page(wiz.PAGE_CORRELATE) + + assert not win._job_running("align_masks"), \ "mask prep needed no background worker (already DC-cached)" - assert all(dlg._angle_params[a] == compute.ManualAngleParams() - for a in range(s.n_angles)), \ - "no manual sidecar yet -> dialog starts at identity, not the automatic result" + assert wait_until(lambda: p1.isComplete()), "masks ready, Next enabled" + assert not win._wizard_act.isEnabled(), \ + "wizard action disabled while a wizard is open" + + st = wiz.state + assert st.result is not None, "an AlignmentResult exists from pre-rotation alone" + assert st.counts is not None and st.counts.shape == st.preview_shape + assert 0 <= st.counts.max() <= s.n_angles + assert not st.fits, "no fits before a correlation has run" + for a in range(s.n_angles): + nominal = compute.nominal_delta_deg(s, a, st.ref_angle_idx) + expected = 0.0 if a == st.ref_angle_idx else nominal + assert abs(st.params[a].rotation_deg - expected) < 1e-9, \ + f"angle {a} not pre-rotated to its stage angle" + assert st.params[a].shift_mm == (0.0, 0.0), \ + "pre-rotation must not invent a translation" -def test_reference_angle_is_locked(ctx): - dlg = ctx.dlg - dlg.combo_active_angle.setCurrentIndex(dlg._ref_angle_idx) +def test_wizard_reference_angle_is_locked(ctx): + p1, wiz = ctx.p1, ctx.wiz + p1.combo_active.setCurrentIndex(wiz.state.ref_angle_idx) pump(30) - before_ref = dlg._angle_params[dlg._ref_angle_idx] - dlg._on_nudge_translate(1, 0, False) - dlg._on_nudge_rotate(1, False) - assert not dlg.grp_manual_adjust.isEnabled(), "reference angle group disabled" - assert dlg._angle_params[dlg._ref_angle_idx] == before_ref, \ + before = wiz.state.params[wiz.state.ref_angle_idx] + p1._on_nudge_translate(1, 0, False) + p1._on_nudge_rotate(1, False) + assert wiz.state.params[wiz.state.ref_angle_idx] == before, \ "reference angle untouched by nudge attempts" -def test_nudges(ctx): - """Nudging a real angle (fine + coarse, translate + rotate).""" - dlg, s = ctx.dlg, ctx.s +def test_wizard_nudges(ctx): + """Manual correction, which the wizard absorbed from the old dialog.""" + p1, wiz, s = ctx.p1, ctx.wiz, ctx.s ctx.active = active = 1 if s.n_angles > 1 else 0 - dlg.combo_active_angle.setCurrentIndex(active) + p1.combo_active.setCurrentIndex(active) pump(30) - before = dlg._angle_params[active].shift_mm - dlg._on_nudge_translate(1, 0, False) # fine +X - fine_step = dlg.spin_step_translate_mm.value() - assert abs(dlg._angle_params[active].shift_mm[0] - (before[0] + fine_step)) < 1e-9, \ + + before = wiz.state.params[active].shift_mm + p1._on_nudge_translate(1, 0, False) + fine = p1.spin_step_translate.value() + assert abs(wiz.state.params[active].shift_mm[0] - (before[0] + fine)) < 1e-9, \ "fine translate nudge moved shift_x by exactly one fine step" - before = dlg._angle_params[active].shift_mm - dlg._on_nudge_translate(0, -1, True) # coarse -Y - coarse_step = fine_step * dlg.spin_step_multiplier.value() - assert abs(dlg._angle_params[active].shift_mm[1] - (before[1] - coarse_step)) < 1e-9, \ + before = wiz.state.params[active].shift_mm + p1._on_nudge_translate(0, -1, True) + coarse = fine * p1.spin_step_mult.value() + assert abs(wiz.state.params[active].shift_mm[1] - (before[1] - coarse)) < 1e-9, \ "coarse translate nudge uses the multiplier" - before_rot = dlg._angle_params[active].rotation_deg - dlg._on_nudge_rotate(1, False) - assert dlg._angle_params[active].rotation_deg != before_rot, \ - "rotate nudge changed rotation_deg" - assert len(dlg._preview_layers) == s.n_angles, \ - "preview canvas rebuilt for every angle after a rotation nudge" + before_rot = wiz.state.params[active].rotation_deg + p1._on_nudge_rotate(1, False) + assert wiz.state.params[active].rotation_deg != before_rot + assert len(wiz.state.layers) == s.n_angles, \ + "stack rebuilt for every angle after a rotation nudge" - # Real key-event wiring (proves keyPressEvent -> signal -> slot). - before = dlg._angle_params[active].shift_mm - QTest.keyClick(dlg.canvas, Qt.Key.Key_Right) - assert dlg._angle_params[active].shift_mm[0] > before[0], \ + # Real key-event wiring (keyPressEvent -> signal -> slot). + before = wiz.state.params[active].shift_mm + QTest.keyClick(p1.canvas, Qt.Key.Key_Right) + assert wiz.state.params[active].shift_mm[0] > before[0], \ "a real Right-arrow key event nudged shift_x" - -def test_auto_derotate(ctx): - """Auto De-rotate: seeds rotation from the stage angle, no translation.""" - dlg, s, active = ctx.dlg, ctx.s, ctx.active - shift_before_derotate = dlg._angle_params[active].shift_mm - dlg._on_auto_derotate() - nominal = compute.nominal_delta_deg(s, active, dlg._ref_angle_idx) - assert abs(dlg._angle_params[active].rotation_deg - nominal) < 1e-6, \ - "auto de-rotate seeded rotation from the stage's reported angle" - assert dlg._angle_params[active].shift_mm == shift_before_derotate, \ - "auto de-rotate left translation untouched" - assert dlg._angle_params[dlg._ref_angle_idx].rotation_deg == 0.0, \ - "reference angle stays identity after auto de-rotate" - # Clicking again offers the other sign, since which one lines the scans up - # is not knowable from the file. - dlg._on_auto_derotate() - assert abs(dlg._angle_params[active].rotation_deg + nominal) < 1e-6, \ - "auto de-rotate offers the opposite sign on a second click" + # Both views render from the same reprojected layers. + p1.combo_view.setCurrentIndex(1) + pump(50) + p1.combo_view.setCurrentIndex(0) + pump(50) -def test_auto_cross_correlate(ctx): - """Auto Cross-Correlate: searches rotation *and* translation.""" - win, dlg, s = ctx.win, ctx.dlg, ctx.s - assert dlg.btn_auto_correlate.isEnabled(), \ - "cross-correlate action enabled once masks are ready" - for label_idx, (label, _sources) in enumerate(dlg._CORRELATE_SOURCES): - dlg.combo_correlate_source.setCurrentIndex(label_idx) - dlg._on_auto_correlate() - assert wait_until( - lambda: not win._job_running("manual_align_correlate"), - timeout_ms=60000), f"auto cross-correlate completed ({label})" - assert all(a in dlg._fit_notes for a in range(s.n_angles) - if a != dlg._ref_angle_idx), \ +def test_wizard_correlate(ctx): + """Cross-correlation, for every source option, retryable.""" + win, p1, wiz, s = ctx.win, ctx.p1, ctx.wiz, ctx.s + from sras_viewer.align_wizard import _CORRELATE_SOURCES + + for idx, (label, _sources) in enumerate(_CORRELATE_SOURCES): + p1.combo_source.setCurrentIndex(idx) + p1.btn_correlate.click() + assert not p1.isComplete(), \ + f"Next must be disabled while correlating ({label})" + assert wait_until(lambda: not win._job_running("align_correlate"), + timeout_ms=60000), f"correlation finished ({label})" + assert p1.isComplete(), f"Next re-enabled ({label})" + assert all(a in wiz.state.fits for a in range(s.n_angles) + if a != wiz.state.ref_angle_idx), \ f"every non-reference angle got a fit ({label})" - assert dlg._angle_params[dlg._ref_angle_idx] == compute.ManualAngleParams(), \ - "auto cross-correlate reference angle stays identity" - assert dlg.grp_correlate.isEnabled() and dlg.btn_save.isEnabled(), \ - "auto cross-correlate re-enabled controls when done" - assert len(dlg._preview_layers) == s.n_angles, \ - "preview canvas rebuilt after cross-correlate" - assert dlg._fit_report(), "fit quality is reported per angle" + + assert wiz.state.params[wiz.state.ref_angle_idx] == compute.ManualAngleParams(), \ + "reference angle stays identity after correlation" + assert p1.btn_correlate.isEnabled(), "controls re-enabled when done" + assert p1.table.rowCount() == s.n_angles and p1.table.item(0, 0) is not None, \ + "per-angle fit table populated" + assert p1.lbl_overlap.text(), "overlap summary reported" + + r = wiz.state.result + assert len(r.per_angle) == s.n_angles, "transform for every angle" + assert r.per_angle[r.ref_angle_idx].shift_mm == (0.0, 0.0), \ + "reference angle has zero shift" -def test_save_sidecar(ctx): - win, dlg, s, active = ctx.win, ctx.dlg, ctx.s, ctx.active - dlg._on_save() - sidecar = compute.sidecar_path(s.path) - assert sidecar.exists(), "sidecar file written" +def test_wizard_retry_changes_geometry(ctx): + """Editing a parameter and re-running is the retry path, and it must + invalidate anything indexed against the old canvas.""" + p1, wiz = ctx.p1, ctx.wiz + gen_before = wiz.state.geometry_generation + p1.spin_threshold.setValue(p1.spin_threshold.value() + 5.0) + p1.spin_threshold.editingFinished.emit() + pump(60) + assert wiz.state.geometry_generation > gen_before, \ + "a threshold change rebuilt the geometry" + + # Reset drops the fits and returns to pre-rotation only. + p1.btn_reset.click() + pump(60) + assert not wiz.state.fits, "reset cleared the fits" + nominal = compute.nominal_delta_deg(ctx.s, ctx.active, wiz.state.ref_angle_idx) + assert abs(wiz.state.params[ctx.active].rotation_deg - nominal) < 1e-9 + assert wiz.state.params[ctx.active].shift_mm == (0.0, 0.0), \ + "reset also drops nudged translation" + + # Put a real correlation back for the pages that follow. + p1.btn_correlate.click() + assert wait_until(lambda: not ctx.win._job_running("align_correlate"), + timeout_ms=60000) + + +def test_wizard_roi_page(ctx): + """The crop page: presets, and two-way sync between the drawn rectangle and + the numeric canvas-pixel boxes.""" + wiz = ctx.wiz + wiz.next() + pump(150) + assert wiz.currentId() == wiz.PAGE_ROI, "advanced to the ROI page" + ctx.p2 = p2 = wiz.page(wiz.PAGE_ROI) + st = wiz.state + + assert st.crop is not None and p2.isComplete(), \ + "a default crop is offered on entry" + n_rows, n_cols = st.result.canvas_shape + assert st.crop[2] > 1 or n_rows == 1, \ + f"default crop must not collapse to a single row: {st.crop}" + + p2.btn_whole.click() + pump(50) + assert st.crop == (0, 0, n_rows, n_cols), "whole-canvas preset" + + p2.btn_fit_union.click() + pump(50) + assert st.counts[st.crop[0]:st.crop[0] + st.crop[2], + st.crop[1]:st.crop[1] + st.crop[3]].sum() == st.counts.sum(), \ + "fit-to-union must keep every covered pixel" + + if p2.btn_fit_overlap.isEnabled(): + p2.btn_fit_overlap.click() + pump(50) + row0, col0, nr, nc = st.crop + assert (st.counts[row0:row0 + nr, col0:col0 + nc] >= 1).all(), \ + "full-overlap crop must not include uncovered pixels" + + # Numeric -> drawn rectangle. + p2.btn_whole.click() + pump(50) + target = (0, 0, max(1, n_rows // 2), max(1, n_cols // 2)) + p2.spin_rows.setValue(target[2]) + p2.spin_cols.setValue(target[3]) + pump(50) + assert st.crop == target, f"spin boxes drive the crop: {st.crop} vs {target}" + + # Drawn rectangle -> numeric, round-tripping exactly. + x0, y0 = wiz.canvas_to_mm(target[1] - 0.5, target[0] - 0.5) + x1, y1 = wiz.canvas_to_mm(target[1] + target[3] - 0.5, + target[0] + target[2] - 0.5) + p2.canvas.set_roi(RoiQuad.from_bbox(min(x0, x1), min(y0, y1), + max(x0, x1), max(y0, y1))) + pump(80) + assert st.crop == target, \ + f"drawn rectangle round-trips to the same crop: {st.crop} vs {target}" + + # A degenerate crop blocks Next. + st.crop = None + p2.completeChanged.emit() + assert not p2.isComplete(), "an absent crop blocks Next" + p2._set_crop(*target) + assert p2.isComplete() + ctx.crop = target + + +def test_wizard_crop_dropped_when_going_back(ctx): + """A crop is canvas-pixel indexed, so it cannot survive a re-correlation.""" + wiz, p2 = ctx.wiz, ctx.p2 + wiz.back() + pump(120) + assert wiz.currentId() == wiz.PAGE_CORRELATE + assert wiz.state.crop is None, "cleanupPage discarded the stale crop" + assert wiz.state.cropped_result is None + wiz.next() + pump(150) + assert wiz.state.crop is not None, "a fresh default crop is offered again" + p2._set_crop(*ctx.crop) + + +def test_wizard_export(ctx): + """Writing the file: Finish stays unavailable until a write succeeds.""" + win, wiz = ctx.win, ctx.wiz + out = ctx.tmpdir / "wizard_aligned.sras" + with patch("sras_viewer.align_wizard.QMessageBox.question", + return_value=QMessageBox.StandardButton.Yes): + wiz.next() + pump(150) + assert wiz.currentId() == wiz.PAGE_SAVE, "advanced to the save page" + ctx.p3 = p3 = wiz.page(wiz.PAGE_SAVE) + assert wiz.state.cropped_result is not None, "crop applied on leaving page 2" + assert wiz.state.cropped_result.canvas_shape == ctx.crop[2:], \ + "cropped result carries the chosen shape" + assert not p3.isComplete(), "Finish unavailable before anything is written" + assert p3.lbl_summary.text(), "a summary of what will be written is shown" + + with patch("sras_viewer.align_wizard.QFileDialog.getSaveFileName", + return_value=(str(out), "")): + p3.btn_browse.click() + assert wiz.state.out_path == str(out) + + p3.btn_export.click() + assert wait_until(lambda: not win._job_running("align_export"), + timeout_ms=60000), "export finished" + assert wiz.state.exported_path == str(out), p3.lbl_status.text() + assert p3.isComplete(), "Finish available once the file exists" + assert out.exists() + + written = SrasFile(str(out)) + ctx.written = written + assert written.version == 6, "export is a v6 file" + assert written.n_angles == ctx.s.n_angles + assert all(written.image_shape(a) == ctx.crop[2:] + for a in range(written.n_angles)), \ + "every angle shares the cropped grid" + assert not out.with_name(out.name + ".part").exists(), \ + "no staging file left behind" + + +def test_wizard_finish_applies_and_persists(ctx): + """Finish makes the session match the file: Aligned View shows the exported + extent, and the sidecar records it for the input scan.""" + win, wiz = ctx.win, ctx.wiz + wiz.accept() + pump(250) + assert win._align_wizard is None, "wizard reference released" + assert win._wizard_act.isEnabled(), "wizard action available again" + assert win._alignment_result is not None + assert win._alignment_result.canvas_shape == ctx.crop[2:], \ + "the *cropped* result is what the view now uses" + assert win.chk_aligned_view.isEnabled() and win.chk_aligned_view.isChecked() + pump(200) + assert win.image_canvas._img_shape == ctx.crop[2:], \ + f"canvas shows the cropped extent: {win.image_canvas._img_shape}" + + sidecar = compute.sidecar_path(ctx.s.path) + assert sidecar.exists(), "sidecar written for the input scan" ctx.sidecar = sidecar ctx.sidecar_raw = raw = json.loads(sidecar.read_text()) - assert raw.get("schema_version") == compute._SIDECAR_SCHEMA_VERSION, \ - "sidecar schema_version is current" - assert all(raw.get("per_angle", {}).get(str(a), {}).get("rotation_deg") - == dlg._angle_params[a].rotation_deg for a in range(s.n_angles)), \ - "sidecar per_angle round-trips the dialog's resolved params" - assert (win._alignment_result is not None - and win._alignment_result.per_angle[active].rotation_deg - == dlg._angle_params[active].rotation_deg), \ - "main window's alignment_result replaced by the manual build" - assert win.chk_aligned_view.isEnabled() and win.chk_aligned_view.isChecked(), \ - "Aligned View auto-enabled after Save" + assert raw.get("schema_version") == compute._SIDECAR_SCHEMA_VERSION + assert all(raw["per_angle"][str(a)]["rotation_deg"] + == win._alignment_result.per_angle[a].rotation_deg + for a in range(ctx.s.n_angles)), \ + "sidecar round-trips the applied rotations" + + win.chk_aligned_view.setChecked(False) + pump(200) + assert win.image_canvas._img_shape == ctx.s.image_shape(0), \ + f"unchecking returns to the raw per-angle grid: {win.image_canvas._img_shape}" def test_stale_schema_sidecar_ignored(ctx): @@ -424,56 +553,39 @@ def test_stale_schema_sidecar_ignored(ctx): sidecar.write_text(json.dumps(raw)) # restore for the rest of the sequence -def test_clear_with_confirmation(ctx): - win, dlg, s = ctx.win, ctx.dlg, ctx.s - with patch("sras_viewer.dialogs.QMessageBox.question", - return_value=QMessageBox.StandardButton.Yes): - dlg._on_clear() - assert not ctx.sidecar.exists(), "sidecar file deleted" - assert all(dlg._angle_params[a] == compute.ManualAngleParams() - for a in range(s.n_angles)), "dialog params reset to identity" - assert win._alignment_result is None, "main window alignment_result cleared" - assert (not win.chk_aligned_view.isEnabled() - and not win.chk_aligned_view.isChecked()), \ - "Aligned View disabled after Clear" - - dlg.close() - pump(150) - assert win._manual_align_dialog is None, "dialog reference released on close" - - def test_sidecar_restored_on_reload(ctx): win, active = ctx.win, ctx.active - win._on_manual_alignment() - dlg = win._manual_align_dialog - dlg.combo_active_angle.setCurrentIndex(active) - pump(30) - dlg._on_auto_derotate() - dlg._on_nudge_translate(1, 1, True) - saved_rotation = dlg._angle_params[active].rotation_deg - saved_shift = dlg._angle_params[active].shift_mm - dlg._on_save() - dlg.close() - pump(150) - + saved = json.loads(ctx.sidecar.read_text())["per_angle"][str(active)] old_sras_id = id(win._sras) win._load_file(str(ctx.path)) # reload the same file fresh assert wait_until( lambda: win._sras is not None and id(win._sras) != old_sras_id), \ "file reloaded" ctx.s = win._sras - assert win._manual_align_dialog is None, \ - "manual dialog force-closed by a reload" + assert win._align_wizard is None, "no wizard left open across a reload" assert win._alignment_result is not None, \ - "reload restores the saved manual alignment automatically" + "reload restores the saved alignment automatically" assert abs(win._alignment_result.per_angle[active].rotation_deg - - saved_rotation) < 1e-9, "restored rotation matches what was saved" - assert win._alignment_result.per_angle[active].shift_mm == saved_shift, \ - "restored shift matches what was saved" + - saved["rotation_deg"]) < 1e-9, \ + "restored rotation matches what was saved" assert win.chk_aligned_view.isChecked(), \ "Aligned View auto-checked after restoring a saved alignment" +def test_wizard_closes_with_a_reload(ctx): + """An open wizard belongs to the file it was opened on.""" + win = ctx.win + win._on_alignment_wizard() + assert win._align_wizard is not None + assert wait_until( + lambda: win._align_wizard.page(win._align_wizard.PAGE_CORRELATE).isComplete()) + win._load_file(str(ctx.path)) + assert wait_until(lambda: not win._job_running("load")) + pump(200) + assert win._align_wizard is None, "wizard force-closed by a reload" + ctx.s = win._sras + + def test_pixel_inspector(ctx): win = ctx.win win.chk_aligned_view.setChecked(False)