diff --git a/sras_align_export.py b/sras_align_export.py new file mode 100644 index 0000000..ebba446 --- /dev/null +++ b/sras_align_export.py @@ -0,0 +1,477 @@ +#!/usr/bin/env python3 +"""Write an aligned, cropped .sras file from an AlignmentResult. + +The alignment machinery in sras_compute never modifies a scan: it produces an +AlignmentResult, and every consumer resamples on the fly (apply_alignment for +the display, reproject_mask for the overlay). That is right for a viewer, but it +means the aligned stack cannot leave the process — no other tool can read it, +and re-opening the scan re-does the registration. + +This module bakes an alignment into a new file. Each angle is resampled onto the +shared canvas that AlignmentResult already defines, cropped to the caller's +window, so every output angle ends up with *identical* geometry: same rows, same +frames, same X/Y coordinates. Rotation and translation are gone, absorbed into +where each waveform sits. The result is an ordinary v6 file that opens already +aligned, and registering it against itself returns identity. + +Two deliberate choices, both about not inventing data: + + * The resample is a nearest-neighbour **gather** of whole waveforms, never an + interpolation. Averaging two neighbouring pixels' CH1 packets would produce + a waveform the instrument never measured, whose FFT peak is not the peak of + either — meaningless for a technique whose entire output is that peak + frequency. So each output pixel gets exactly one source pixel's three + waveforms, verbatim, and the cost is that some source pixels are duplicated + and others dropped. This matches apply_alignment's order=0 for the same + reason. + * Output pixels with no source pixel (the canvas corners a rotated scan cannot + reach, and anything outside the crop's coverage) are filled with the ADC + code for 0 mV, not with zero. See _fill_row. + +Depends only on numpy/sras_format/sras_compute — no Qt — so it is directly +unit-testable and importable from a worker thread. +""" + +import os +import struct +from dataclasses import dataclass, field +from pathlib import Path + +import numpy as np + +import sras_compute as compute +from sras_compute import AlignmentResult +from sras_format import GEO_FMT_V6, HDR_FMT_V6, SrasFile, mv_to_adc + +# GEO_FMT_V6 stores n_rows as ">H" and n_frames as ">I". A canvas that overflows +# either is not representable, and silently truncating would write a file whose +# geometry table disagrees with its waveform block. +_MAX_ROWS = 0xFFFF +_MAX_FRAMES = 0xFFFFFFFF + +# Slack, in source pixels, on the in-bounds test at the very edge of a source +# array. Absorbs the ~1e-13 of float noise an exactly-integer affine picks up +# from being built in mm space; see _in_bounds. +_EDGE_TOL = 1e-6 + +# Output rows per write() call. One row is n_channels * n_cols * spf bytes — +# ~1.5 MB on a full-size scan — so a handful of rows keeps the peak buffer in +# the low tens of MB no matter how large the scan is. +_ROW_CHUNK = 8 + + +@dataclass +class ExportPlan: + """What write_aligned_sras would produce, without producing it. + + Derived from the affine transforms alone — no waveform bytes are read — so + the wizard can call it on every ROI edit to keep a live size estimate and + per-angle coverage readout in front of the user *before* they commit to a + multi-gigabyte write. + """ + n_rows: int + n_frames: int + n_angles: int + bytes_per_angle: int + total_bytes: int + valid_px: dict[int, int] # output pixels with a source pixel + warnings: list[str] = field(default_factory=list) + + def coverage_frac(self, angle_idx: int) -> float: + px = self.n_rows * self.n_frames + return (self.valid_px.get(angle_idx, 0) / px) if px else 0.0 + + +def _src_coords(t, rows, n_cols: int) -> tuple[np.ndarray, np.ndarray]: + """Fractional source (row, col) coordinates for whole output rows. + + *rows* may be a scalar row index or an array of them; the result is shaped + (len(rows), n_cols), or (n_cols,) for a scalar. + """ + cols = np.arange(n_cols, dtype=np.float64) + r = np.atleast_1d(np.asarray(rows, dtype=np.float64))[:, None] + sr = t.matrix[0, 0] * r + t.matrix[0, 1] * cols + t.offset[0] + sc = t.matrix[1, 0] * r + t.matrix[1, 1] * cols + t.offset[1] + if np.isscalar(rows) or np.asarray(rows).ndim == 0: + return sr[0], sc[0] + return sr, sc + + +def _round_idx(coord: np.ndarray) -> np.ndarray: + """Nearest source index, rounding halves away from zero. + + floor(x + 0.5), not np.rint: scipy.ndimage's order=0 rounds halves away + from zero while np.rint rounds them to even, and these have to be the same + source pixels the apply_alignment(order=0) preview drew. Exact halves are + not a corner case here — the canvas is snapped to the reference angle's own + pixel grid (canvas_for_params snap=True), so an unrotated angle lands on + half-integers wherever its row pitch differs from the reference's. + """ + return np.floor(coord + 0.5).astype(np.int64) + + +def _in_bounds(sr: np.ndarray, sc: np.ndarray, + n_rows: int, n_frames: int) -> np.ndarray: + """Which output pixels have a source pixel, by scipy's mode="constant" rule + plus a tolerance at the edge. + + Tested on the *fractional* coordinate against the range of sample centres, + [0, n-1] inclusive — deliberately not on the rounded index. The two differ + around the whole rim: a coordinate of -0.4 rounds to a perfectly valid index + 0, but scipy calls it out of bounds and writes cval there, so testing the + rounded index would put a one-pixel rim of real data everywhere the Aligned + View shows padding. + + _EDGE_TOL is why this is not literally scipy's test. 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 lands on + -20 - 7e-15 rather than -20). Bare >= 0.0 then rejects that angle's entire + first row, and <= n-1 its last column — for the *reference* angle, whose + whole role is to pass through as an exact integer crop. The tolerance is + seven orders of magnitude above that noise and seven below the half-pixel + scale at which a rounding decision is ever meaningful, so it can only ever + change pixels whose scipy answer was itself decided by rounding noise. + """ + return ((sr >= -_EDGE_TOL) & (sr <= n_rows - 1 + _EDGE_TOL) + & (sc >= -_EDGE_TOL) & (sc <= n_frames - 1 + _EDGE_TOL)) + + +# Output rows evaluated per numpy call when counting coverage. Counting row by +# row costs one small matmul per row (hundreds of milliseconds per angle on a +# full-size scan, on every ROI edit); counting the whole canvas at once needs +# hundreds of MB of index arrays. Blocking gets both: ~10 numpy calls per angle +# against ~30 MB of live index arrays. +_COUNT_BLOCK = 128 + + +def _count_in_bounds(t, n_rows: int, n_cols: int, + src_rows: int, src_frames: int) -> int: + """How many of the n_rows x n_cols output pixels have a source pixel.""" + total = 0 + for start in range(0, n_rows, _COUNT_BLOCK): + rows = np.arange(start, min(start + _COUNT_BLOCK, n_rows)) + sr, sc = _src_coords(t, rows, n_cols) + total += int(np.count_nonzero(_in_bounds(sr, sc, src_rows, src_frames))) + return total + + +def plan_export(sras: SrasFile, result: AlignmentResult) -> ExportPlan: + """Geometry, size and per-angle coverage of the file *result* would export. + + Coverage is counted from the actual per-pixel index arrays rather than + approximated by the footprint parallelogram's area, because the two differ + exactly where it matters — a crop that clips one angle's scan window — and + this number is what tells the user an angle will come out mostly empty. No + waveform bytes are read, so it stays fast enough to call on every ROI edit. + """ + n_rows, n_cols = result.canvas_shape + n_angles = sras.n_angles + warnings: list[str] = [] + + valid_px: dict[int, int] = {} + for a in range(n_angles): + t = result.per_angle.get(a) + if t is None: + valid_px[a] = 0 + warnings.append(f"Angle {a} has no transform and will be all padding.") + continue + src_rows, src_frames = sras.image_shape(a) + valid_px[a] = _count_in_bounds(t, n_rows, n_cols, src_rows, src_frames) + + bytes_per_angle = (n_rows * sras.n_channels * n_cols + * sras.samples_per_frame * sras.bytes_per_sample) + + if n_rows > _MAX_ROWS: + warnings.append( + f"Crop is {n_rows} rows; the .sras geometry table caps rows at " + f"{_MAX_ROWS}. Narrow the ROI in Y.") + if n_cols > _MAX_FRAMES: + warnings.append(f"Crop is {n_cols} frames; the cap is {_MAX_FRAMES}.") + for a in range(n_angles): + frac = valid_px[a] / (n_rows * n_cols) if n_rows and n_cols else 0.0 + if frac == 0.0: + warnings.append( + f"Angle {a} has no data inside this crop — it will be written " + f"as all padding.") + elif frac < 0.10: + warnings.append( + f"Angle {a} covers only {frac * 100:.1f}% of the crop.") + if sras.background is None: + warnings.append( + "Input has no background waveform (pre-v4 scan); a zero background " + "is written, which makes background subtraction a no-op.") + if sras.version != 6: + warnings.append(f"Input is v{sras.version}; the export is written as v6.") + if getattr(sras, "scan_aborted", False): + warnings.append( + f"Input scan was aborted: only its {n_angles} complete angle(s) " + f"are exported.") + + return ExportPlan(n_rows=n_rows, n_frames=n_cols, n_angles=n_angles, + bytes_per_angle=bytes_per_angle, + total_bytes=bytes_per_angle * n_angles, + valid_px=valid_px, warnings=warnings) + + +def _encode_preambles(sras: SrasFile) -> bytes: + """The Preamble Blocks section for the output file. + + v6/v7 inputs carry the on-disk span verbatim (sras.preambles_raw), which is + both cheaper and lossless. Legacy inputs don't keep that span, and a v2 file + has no preambles at all — there, empty strings are written. That is not a + silent downgrade: _parse_preamble("") returns {} and _set_calibration falls + back to the hardcoded scope constants, which is exactly the calibration a v2 + file already gets, so mV values round-trip unchanged. + """ + raw = getattr(sras, "preambles_raw", None) + if raw is not None: + return raw + strings = sras.preambles or [""] * sras.n_channels + out = bytearray() + for s in strings: + encoded = s.encode("utf-8") + out += struct.pack(">H", len(encoded)) + encoded + return bytes(out) + + +def _encode_background(sras: SrasFile) -> bytes: + """The Background Block for the output file. + + When the input has none (v2/v3), write samples_per_frame zeros rather than a + zero-length block. Every consumer guards on `background is not None` and + then subtracts it from a (spf,)-shaped row, so a length-0 array would + broadcast-fail at the first background-subtracted FFT; zeros make the + subtraction a correct no-op instead. + """ + raw = getattr(sras, "background_raw", None) + if raw is not None: + return raw + if sras.background is None: + samples = np.zeros(sras.samples_per_frame, dtype=np.int8) + else: + samples = np.rint(sras.background).astype(np.int8) + return struct.pack(">I", samples.size) + samples.tobytes() + + +def _fill_row(sras: SrasFile, n_cols: int, dtype) -> np.ndarray: + """One output row of pure padding, shape (n_channels, n_cols, spf). + + Filled per channel with the ADC code for 0 mV, not with 0. Zero ADC decodes + to (0 - yoff) * ymult + yzero, which for a real scope preamble is a long way + from 0 mV — often far enough to sit above the CH4 mask threshold, which + would paint a solid rectangle of "valid" pixels around the sample and make + every DC image and every ROI statistic wrong. Rounding to the integer code + lands within half an ADC step of 0 mV, which is as close as the format can + represent. + """ + info = np.iinfo(dtype) + codes = [int(np.clip(round(mv_to_adc(0.0, *sras.cal(ch))), info.min, info.max)) + for ch in range(sras.n_channels)] + row = np.empty((sras.n_channels, n_cols, sras.samples_per_frame), dtype=dtype) + for ch, code in enumerate(codes): + row[ch] = code + return row + + +class _SourceReader: + """Gives the gather source rows without ever reading one twice. + + This is the difference between a usable export and an unusable one, and it + is entirely about read amplification. A rotated angle maps one output row to + a *diagonal* line across the source array, so the pixels of a single output + row come from hundreds of different source rows — on a full-size scan, a + ~1.4 MB source row each. Indexing a memmap pixel by pixel in output order + therefore re-faults nearly the whole angle for every output row: terabytes + of paging for a gigabyte of data. + + So rows are served from a contiguous *band* held in RAM. The band for a + chunk of output rows is read in one sequential slice, and because output + rows advance monotonically through the source, consecutive chunks' bands + barely overlap: each source row is read about once, and the whole job costs + roughly 2x the source size in reads rather than a thousand times it. + + Small angles skip the machinery — if the whole block fits the budget it is + materialized once and every band is a view of it. + """ + + def __init__(self, sras: SrasFile, angle_idx: int, budget: int): + self._src = sras.data[angle_idx] + self._n_rows = self._src.shape[0] + self._row_bytes = max(1, self._src[0].nbytes) + self._whole = np.asarray(self._src) if self._src.nbytes <= budget else None + # Leave room for the output buffer and the index arrays alongside. + self._max_band = max(1, int(budget * 0.5) // self._row_bytes) + self._band = None + self._lo = self._hi = 0 + + def band(self, lo: int, hi: int) -> tuple[np.ndarray, int]: + """Rows [lo, hi) as an in-RAM array, plus the index its row 0 holds.""" + lo = max(0, min(lo, self._n_rows)) + hi = max(lo + 1, min(hi, self._n_rows)) + if self._whole is not None: + return self._whole, 0 + if self._band is None or lo < self._lo or hi > self._hi: + # Read a little more than asked so a chunk whose band creeps + # forward by a few rows does not re-read the whole span. + span = min(self._max_band, max(hi - lo, self._max_band // 2)) + self._lo = lo + self._hi = min(self._n_rows, lo + span) + if self._hi < hi: # band cannot cover the ask + self._hi = hi + self._band = np.asarray(self._src[self._lo:self._hi]) + return self._band, self._lo + + def close(self): + self._whole = None + self._band = None + + +def write_aligned_sras(sras: SrasFile, result: AlignmentResult, out_path, + *, progress_cb=None, should_stop=None) -> Path: + """Write *sras*, aligned per *result* and cropped to its canvas, to a new + v6 .sras file. Returns the path written. + + *result* is used exactly as given: crop the canvas first with + compute.crop_alignment_result, whose offset shift makes the cropped result + resample precisely the window the user selected. + + No cache tail is written. Any DC/FFT the input had cached is indexed by the + input's grid and is meaningless on the new one, so — like sras_edit_scans — + the export drops it and lets the viewer recompute. + + Writes to a sibling ".part" file and os.replace()s it into position on + success, unlinking it on error or cancellation: a half-written .sras is not + detectably broken (the v6 parser treats a short file as an aborted scan and + opens it happily), so it must never be left where the user might load it. + + *should_stop* is polled once per output row chunk; returning True aborts and + raises nothing — the partial file is removed and the returned path will not + exist, so callers must check. + """ + out_path = Path(out_path) + n_rows, n_cols = result.canvas_shape + n_ch, spf = sras.n_channels, sras.samples_per_frame + n_angles = sras.n_angles + + if n_rows <= 0 or n_cols <= 0: + raise ValueError(f"empty canvas: {n_rows} x {n_cols}") + if n_rows > _MAX_ROWS: + raise ValueError( + f"{n_rows} rows exceeds the .sras per-angle geometry limit of " + f"{_MAX_ROWS}; crop further in Y") + if n_cols > _MAX_FRAMES: + raise ValueError(f"{n_cols} frames exceeds the limit of {_MAX_FRAMES}") + missing = [a for a in range(n_angles) if a not in result.per_angle] + if missing: + raise ValueError(f"alignment result has no transform for angle(s) {missing}") + # The source's waveform blocks are live read-only memmaps into sras.path, + # so writing over it would corrupt the very reads the gather is making. + if out_path.exists() and out_path.samefile(sras.path): + raise ValueError( + "refusing to export onto the source scan; choose another filename") + + dtype = np.dtype(np.int8 if sras.bytes_per_sample == 1 else ">i2") + x0_mm, y0_mm = result.canvas_origin_mm + y_rows = (y0_mm + np.arange(n_rows) * result.canvas_dy_mm).astype(">f4") + + # Reference-only header fields. v6/v7 inputs have real ones to carry over; + # for a legacy input describe the canvas we are actually writing. + if hasattr(sras, "x_start_nominal_mm"): + nominal = (sras.x_start_nominal_mm, sras.y_start_nominal_mm, + sras.x_delta_nominal_mm, sras.y_delta_nominal_mm, + sras.row_spacing_mm) + else: + nominal = (x0_mm, y0_mm, + n_cols * sras.pixel_x_mm, n_rows * result.canvas_dy_mm, + result.canvas_dy_mm) + + header = struct.pack( + HDR_FMT_V6, b"SRAS", 6, n_angles, + float(nominal[0]), float(nominal[1]), float(nominal[2]), + float(nominal[3]), float(nominal[4]), + sras.velocity_mm_s, sras.laser_freq_hz, spf, sras.sample_rate_hz, + sras.bytes_per_sample, n_ch) + + # Every angle now shares one grid, so the ragged v6 tables collapse to + # n_angles copies of the same record. x_delta is the reference angle's own + # pitch (the canvas is its grid extended), which is velocity/laser_freq + # exactly, so x_axis_mm() stays self-consistent on re-read. + geo = struct.pack(GEO_FMT_V6, float(x0_mm), float(sras.pixel_x_mm), + int(n_cols), int(n_rows)) * n_angles + + budget = compute._TOTAL_BYTES_BUDGET + total_chunks = max(1, n_angles * ((n_rows + _ROW_CHUNK - 1) // _ROW_CHUNK)) + done_chunks = 0 + cancelled = False + + part_path = out_path.with_name(out_path.name + ".part") + try: + with open(part_path, "wb") as fout: + fout.write(header) + fout.write(sras.angles_deg.astype(">f4").tobytes()) + fout.write(geo) + fout.write(y_rows.tobytes() * n_angles) + fout.write(_encode_preambles(sras)) + fout.write(_encode_background(sras)) + + pad = _fill_row(sras, n_cols, dtype) + for a in range(n_angles): + t = result.per_angle[a] + reader = _SourceReader(sras, a, budget) + src_rows, src_frames = sras.image_shape(a) + try: + for chunk_start in range(0, n_rows, _ROW_CHUNK): + if should_stop is not None and should_stop(): + cancelled = True + break + chunk = np.arange(chunk_start, + min(chunk_start + _ROW_CHUNK, n_rows)) + sr, sc = _src_coords(t, chunk, n_cols) + ok = _in_bounds(sr, sc, src_rows, src_frames) + # Clip rather than trust: _EDGE_TOL admits coordinates a + # hair outside the array, and an index off the end here + # would silently read the wrong row of the band. + idx_r = np.clip(_round_idx(sr), 0, src_rows - 1) + idx_c = np.clip(_round_idx(sc), 0, src_frames - 1) + + # One band read covers the whole chunk: every source row + # any of these output rows touches, in one sequential + # slice. See _SourceReader. + if ok.any(): + band, base = reader.band(int(idx_r[ok].min()), + int(idx_r[ok].max()) + 1) + else: + band, base = None, 0 + + for i in range(len(chunk)): + out = pad.copy() + keep = ok[i] + if keep.any(): + # The two advanced indices are separated by a + # slice, so numpy puts the gathered axis first: + # (n_sel, n_ch, spf). Move it behind channels. + out[:, keep, :] = band[ + idx_r[i][keep] - base, :, idx_c[i][keep], : + ].transpose(1, 0, 2) + fout.write(out.tobytes()) + done_chunks += 1 + if progress_cb is not None: + progress_cb(int(done_chunks / total_chunks * 100)) + finally: + reader.close() + if cancelled: + break + if not cancelled: + fout.flush() + os.fsync(fout.fileno()) + if cancelled: + part_path.unlink(missing_ok=True) + return out_path + os.replace(part_path, out_path) + except BaseException: + part_path.unlink(missing_ok=True) + raise + + if progress_cb is not None: + progress_cb(100) + return out_path diff --git a/sras_compute.py b/sras_compute.py index 1f0f369..26a8600 100644 --- a/sras_compute.py +++ b/sras_compute.py @@ -1235,13 +1235,22 @@ def default_max_workers() -> int: def _rotation_candidates(nominal_deg: float, search_deg: float, - step_deg: float) -> list[float]: - """Coarse rotation candidates: a window around *both* signs of the stage's - reported angle change. Scoring both is what makes the stage's sign - convention a non-issue — the images decide which way the stage turns, and - a file whose stage reports the opposite sense registers just as well.""" + step_deg: float, + signs: tuple[int, ...] = (-1, 1)) -> list[float]: + """Coarse rotation candidates: a window around each requested sign of the + stage's reported angle change. Scoring both signs (the default) is what + makes the stage's sign convention a non-issue — the images decide which way + the stage turns, and a file whose stage reports the opposite sense + registers just as well. + + *signs* exists only so a user who has established which way their own stage + turns can halve the coarse sweep from the alignment wizard. It is an + override, not an inference: nothing in the file says which sign is right. + """ out: list[float] = [] - for center in (-nominal_deg, nominal_deg): + step_deg = max(abs(step_deg), 1e-9) # a 0 step would divide by zero below + for sign in signs: + center = sign * nominal_deg k = int(np.floor(search_deg / step_deg)) for i in range(-k, k + 1): out.append(center + i * step_deg) @@ -1351,15 +1360,18 @@ def register_angle_to_reference( dc_threshold_mv: float = 0.0, sources: tuple[str, ...] = ("signal", "mask"), coarse_dim: int = 256, fine_dim: int = _DEFAULT_FINE_DIM, - search_deg: float = 6.0, coarse_step_deg: float = 2.0) -> RigidFit: + search_deg: float = 6.0, coarse_step_deg: float = 2.0, + seed_deg: float | None = None, + seed_signs: tuple[int, ...] = (-1, 1), + refine: bool = True) -> RigidFit: """Rigid (rotation + translation, never scale) fit of *angle_idx* onto *ref_angle_idx*, found entirely by cross-correlating image content. Two stages: - 1. Coarse sweep at *coarse_dim* over a ±*search_deg* window around both - signs of the stage's reported angle change (see _rotation_candidates), - for each requested source image, scored by _overlap_ncc. Only has to - pick the right basin. + 1. Coarse sweep at *coarse_dim* over a ±*search_deg* window around the + requested signs of the stage's reported angle change (see + _rotation_candidates), for each requested source image, scored by + _overlap_ncc. Only has to pick the right basin. 2. Hill-climbing refinement of the winning (source, rotation) at *fine_dim* with sub-pixel translation folded into every score (_refine_rotation, _score_rotation), down to 0.05°. @@ -1368,6 +1380,22 @@ def register_angle_to_reference( local mm to ref mm (q = R @ l + shift); score is the final NCC, which the caller can surface so a bad scan is visible rather than silently fused in. Returns identity for the reference angle itself. + + The last three arguments only exist to let the alignment wizard expose the + rotation search; every default reproduces the search this function has + always done. + + *seed_deg* replaces the stage's reported angle change as the center of the + coarse sweep — the pre-rotation the search starts from. None (the default) + means the stage angle from the file's Angle Table, which is nearly always + what you want: it puts the sweep within a couple of degrees of the answer. + Pass 0.0 to search around no rotation at all, which is the honest choice for + a file whose stage angles are known to be wrong. + + *refine=False* stops after the coarse sweep, leaving rotation on the + *coarse_step_deg* grid. Combined with search_deg=0.0 and a single seed sign + it pins rotation to exactly the seed and searches translation only — for a + scan whose stage angles are trusted more than its image content. """ if angle_idx == ref_angle_idx: return RigidFit(0.0, (0.0, 0.0), 1.0, "reference") @@ -1377,8 +1405,10 @@ def register_angle_to_reference( if not candidates: return RigidFit(0.0, (0.0, 0.0), -1.0, "none") - nominal = nominal_delta_deg(sras, angle_idx, ref_angle_idx) - thetas = _rotation_candidates(nominal, search_deg, coarse_step_deg) + nominal = (nominal_delta_deg(sras, angle_idx, ref_angle_idx) + if seed_deg is None else float(seed_deg)) + thetas = _rotation_candidates(nominal, search_deg, coarse_step_deg, + seed_signs) # ---- Stage 1: coarse sweep, every source ------------------------------ pitch_c, n_c = _reg_pitch_and_size(sras, coarse_dim) @@ -1407,9 +1437,10 @@ def register_angle_to_reference( theta = best[1] score, shift = _score_rotation(ref_img, ref_valid, mov_reg, pitch_f, n_f, theta, subpixel=True) - theta, score, shift = _refine_rotation( - ref_img, ref_valid, mov_reg, pitch_f, n_f, theta, score, shift, - step_deg=coarse_step_deg) + if refine: + theta, score, shift = _refine_rotation( + ref_img, ref_valid, mov_reg, pitch_f, n_f, theta, score, shift, + step_deg=coarse_step_deg) dr, dc = shift return RigidFit(float(theta), (float(dc * pitch_f), float(dr * pitch_f)), @@ -1540,6 +1571,114 @@ def _result_from_params(sras: SrasFile, ref_angle_idx: int, pitch[0], pitch[1], canvas_origin_mm, per_angle) +def crop_alignment_result(result: AlignmentResult, row0: int, col0: int, + n_rows: int, n_cols: int) -> AlignmentResult: + """The same alignment restricted to a rectangular window of its canvas — + canvas pixel (row0, col0) becomes the cropped canvas's (0, 0). + + Cropping folds into each angle's existing affine instead of becoming a + second transform, because a canvas crop is *pure index translation*. From + _affine_out_to_src, matrix = D @ Rinv @ A_out depends only on pitch and + rotation, and A_out @ [row0, col0] is exactly the mm displacement of the new + origin, so: + + matrix @ [r', c'] + (offset + matrix @ [row0, col0]) + == matrix @ [r' + row0, c' + col0] + offset + + i.e. shifting the offset by matrix @ [row0, col0] reproduces the original + mapping at the shifted indices, exactly. That equality is what lets + apply_alignment, reproject_mask and the aligned .sras exporter all keep + working on a cropped result with no special-casing — and it is why the crop + preview a user approves is guaranteed to be the same pixels the exporter + writes. + + Bounds are the caller's responsibility (the wizard's ROI page clamps to the + canvas): an out-of-range window is geometrically well-defined here and + simply resamples padding. + """ + if n_rows <= 0 or n_cols <= 0: + raise ValueError(f"empty crop: {n_rows} x {n_cols}") + + delta = np.array([float(row0), float(col0)]) + per_angle = { + a: AngleTransform(t.rotation_deg, t.shift_mm, t.matrix, + t.offset + t.matrix @ delta, t.score, t.source) + for a, t in result.per_angle.items() + } + origin = (result.canvas_origin_mm[0] + col0 * result.canvas_dx_mm, + result.canvas_origin_mm[1] + row0 * result.canvas_dy_mm) + return AlignmentResult(result.ref_angle_idx, result.dc_threshold_mv, + (int(n_rows), int(n_cols)), + result.canvas_dx_mm, result.canvas_dy_mm, + origin, per_angle) + + +def overlap_stats(counts: np.ndarray, n_angles: int) -> dict: + """Summarize a per-pixel "how many angles cover this pixel" image — the + number the alignment wizard's mask-stack view is colored by. + + Separated from the drawing code because it is the actual judgement the user + makes on that screen ("do the angles land on top of each other?"), and a + plain array-in/dict-out function can be tested without Qt. + """ + counts = np.asarray(counts) + union = int(np.count_nonzero(counts)) + full = int(np.count_nonzero(counts >= n_angles)) + return { + "union_px": union, + "full_px": full, + "full_frac": (full / union) if union else 0.0, + "mean_count": float(counts[counts > 0].mean()) if union else 0.0, + "max_count": int(counts.max()) if counts.size else 0, + "empty": union == 0, + } + + +def largest_rect_at_least(counts: np.ndarray, min_count: int + ) -> tuple[int, int, int, int] | None: + """Largest axis-aligned rectangle whose every pixel has counts >= min_count, + as (row0, col0, n_rows, n_cols), or None if no pixel qualifies. + + Backs the wizard's "fit to full overlap" button. A *bounding box* of the + qualifying pixels would be the obvious thing and is wrong: the full-overlap + region of several rotated scans is roughly a disc, whose bounding box has + corners no angle covers at all. Offering that as the crop would hand the + user padding they explicitly asked to avoid, so this finds a rectangle that + is entirely inside the region. + + Standard largest-rectangle-in-a-histogram sweep — O(rows * cols) on a + preview-sized array, so it is instant at interactive rates. + """ + good = np.asarray(counts) >= min_count + if not good.any(): + return None + + n_rows, n_cols = good.shape + best = (0, 0, 0, 0) # area, row0, col0, ... + best_area = 0 + heights = np.zeros(n_cols, dtype=np.int64) + + for r in range(n_rows): + heights = np.where(good[r], heights + 1, 0) + # Sentinel column of height 0 flushes the stack at the end of the row. + stack: list[tuple[int, int]] = [] # (start col, height) + for c in range(n_cols + 1): + h = int(heights[c]) if c < n_cols else 0 + start = c + while stack and stack[-1][1] >= h: + s, sh = stack.pop() + area = sh * (c - s) + if area > best_area: + best_area = area + best = (s, sh, c - s, r) + start = s + if h: + stack.append((start, h)) + + col0, height, width, row_end = best + return (row_end - height + 1, col0, height, width) + + def _parallel_map(fn, items, n_workers: int) -> list: """fn over items, in order, threaded when it pays.""" items = list(items) diff --git a/sras_viewer/canvases.py b/sras_viewer/canvases.py index f7b5e83..fa73704 100644 --- a/sras_viewer/canvases.py +++ b/sras_viewer/canvases.py @@ -1,7 +1,9 @@ """Matplotlib canvases and the ROI primitive.""" +import matplotlib as mpl import numpy as np from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg +from matplotlib.colors import BoundaryNorm, ListedColormap from matplotlib.figure import Figure from matplotlib.patches import Polygon from matplotlib.path import Path as MplPath @@ -100,7 +102,15 @@ class ImageCanvas(FigureCanvasQTAgg): _HANDLE_PX = 12 _CLICK_THRESH_PX = 4 # releases within this of press count as a click - def __init__(self, parent=None): + def __init__(self, parent=None, *, rect_only: bool = False): + """*rect_only* constrains the ROI to an axis-aligned rectangle. + + Used by the alignment wizard's crop page, where a free quadrilateral + would be actively misleading: v6 geometry can only express an + axis-aligned rectangle, so anything else the user drew would have to be + squared off behind their back. Default off, so the main window's + free-quad ROI is unaffected. + """ fig = Figure(figsize=(7, 5), tight_layout=True) self.ax = fig.add_subplot(111) super().__init__(fig) @@ -108,6 +118,7 @@ class ImageCanvas(FigureCanvasQTAgg): self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding) self._extent = None self._img_shape = None + self._rect_only = rect_only # ROI state self._roi: RoiQuad | None = None @@ -132,9 +143,12 @@ class ImageCanvas(FigureCanvasQTAgg): # Public API # ------------------------------------------------------------------ - def show_image(self, img: np.ndarray, extent: list[float], cmap: str, + def show_image(self, img: np.ndarray, extent: list[float], cmap, vmin: float, vmax: float, xlabel: str, ylabel: str, title: str, - colorbar_label: 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.""" self.figure.clf() self.ax = self.figure.add_subplot(111) # Patches and lines are destroyed by figure.clf(); drop stale refs. @@ -148,7 +162,8 @@ class ImageCanvas(FigureCanvasQTAgg): extent=extent, cmap=cmap, vmin=vmin, vmax=vmax, interpolation="nearest", ) - cb = self.figure.colorbar(im, ax=self.ax, fraction=0.046, pad=0.04) + cb = self.figure.colorbar(im, ax=self.ax, fraction=0.046, pad=0.04, + ticks=cb_ticks) if colorbar_label: cb.set_label(colorbar_label) @@ -302,10 +317,28 @@ class ImageCanvas(FigureCanvasQTAgg): self._roi._pts = self._snapshot.corners() + delta elif self._state == self._DRAG_CORNER: self._roi._pts[self._drag_corner_idx] = [event.xdata, event.ydata] + if self._rect_only: + self._rectify_corner(self._drag_corner_idx) self._draw_roi() self.draw_idle() + def _rectify_corner(self, idx: int): + """Re-square the quad after a corner drag, anchored on the *opposite* + corner. + + Anchoring on the diagonal opposite (idx ^ 2, since corners run + BL, BR, TR, TL) rather than taking the bbox of all four points is what + lets the rectangle shrink: a bbox over the three stale corners plus the + new one is the union of the old rectangle and the new point, so dragging + inward would never make it smaller. + """ + pts = self._roi.corners() + ax_, ay = pts[idx ^ 2] + bx, by = pts[idx] + self._roi._pts = RoiQuad.from_bbox(min(ax_, bx), min(ay, by), + max(ax_, bx), max(ay, by)).corners() + def _on_release(self, event): if event.button != 1 and self._press_button != 1: return @@ -475,14 +508,20 @@ class WaveformCanvas(FigureCanvasQTAgg): self.draw() -class ManualAlignOverlayCanvas(FigureCanvasQTAgg): - """Renders ManualAlignmentDialog's multi-angle mask overlay and turns - keyboard input into translate/rotate nudge requests for whichever angle - the dialog currently has active. +class AlignOverlayCanvas(FigureCanvasQTAgg): + """Renders the alignment wizard's multi-angle mask views and turns keyboard + input into translate/rotate nudge requests for whichever angle is active. + + Two views of the same reprojected masks, because they answer different + questions. show_counts colours each pixel by *how many* angles cover it, + which is the at-a-glance verdict on a correlation run: a good alignment is + one saturated plateau, a bad one is a fringe of low-count halos. + show_overlay gives each angle its own colour, which is what you need while + nudging a specific angle by hand. A pure input+render widget — it holds no alignment state and never - touches SrasFile itself; ManualAlignmentDialog owns all of that and - decides, from these signals, whether a cheap single-layer refresh or a + touches SrasFile itself; the wizard page owns all of that and decides, + from these signals, whether a cheap single-layer refresh or a full preview-canvas rebuild is needed. FigureCanvasQTAgg is a real QWidget, so keyPressEvent works like on any @@ -522,6 +561,35 @@ class ManualAlignOverlayCanvas(FigureCanvasQTAgg): self.figure.clf() self.ax = self.figure.add_subplot(111) self.ax.imshow(rgba, extent=extent, origin="upper", aspect="auto") + self._finish(title) + + def show_counts(self, counts: np.ndarray, n_angles: int, + extent: list[float], title: str): + """The mask stack coloured by how many angles cover each pixel. + + A discrete colormap with integer-ticked colorbar rather than a + continuous one: the judgement being made is "is this a plateau at N, or + a fan of partial overlaps", and a region covered by one angle too few + has to read as its own band rather than a slightly darker shade. + Uncovered pixels are transparent so they cannot be mistaken for a low + count. + """ + 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)] + 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))) + cb = self.figure.colorbar(im, ax=self.ax, fraction=0.046, pad=0.04, + ticks=np.arange(0, n + 1)) + cb.set_label("angles overlapping") + self._finish(title) + + def _finish(self, title: str): self.ax.set_xlabel("X (mm)") self.ax.set_ylabel("Y (mm)") self.ax.set_title(title) diff --git a/sras_viewer/common.py b/sras_viewer/common.py index 5793e90..7d957f6 100644 --- a/sras_viewer/common.py +++ b/sras_viewer/common.py @@ -60,9 +60,13 @@ class Jobs: COMPUTE = "compute" DC_PRECOMPUTE = "dc_precompute" BATCH = "batch" - ALIGN = "align" - MANUAL_ALIGN_MASKS = "manual_align_masks" - MANUAL_ALIGN_CORRELATE = "manual_align_correlate" + # The alignment wizard's three background steps: fetching each angle's CH4 + # image for the mask stack, registering the angles, and writing the aligned + # export. Separate keys because a retry of one must not be blocked by + # another having run, and _run_worker's busy check is per key. + ALIGN_MASKS = "align_masks" + ALIGN_CORRELATE = "align_correlate" + ALIGN_EXPORT = "align_export" def _make_dspin(lo: float, hi: float, decimals: int, *, suffix: str = "", diff --git a/sras_viewer/dialogs.py b/sras_viewer/dialogs.py index 1ae09f2..c1c8352 100644 --- a/sras_viewer/dialogs.py +++ b/sras_viewer/dialogs.py @@ -20,7 +20,7 @@ from sras_compute import ( from sras_format import SrasFile from sras_workers import Ch4MaskWorker, CrossCorrelateWorker -from .canvases import ManualAlignOverlayCanvas +from .canvases import AlignOverlayCanvas from .common import ( _CSS_HINT, _CSS_MUTED, _CSS_WARN, Jobs, _axes_extent, _form, _group, _make_dspin, _scroll_panel, _wrap_label, @@ -354,7 +354,7 @@ class ManualAlignmentDialog(QDialog): def _build_ui(self, dc_threshold_mv: float): root = QHBoxLayout(self) - self.canvas = ManualAlignOverlayCanvas() + self.canvas = AlignOverlayCanvas() left = QWidget() left_l = QVBoxLayout(left) left_l.setContentsMargins(0, 0, 0, 0) diff --git a/sras_viewer/main_window.py b/sras_viewer/main_window.py index 6906b88..6401603 100644 --- a/sras_viewer/main_window.py +++ b/sras_viewer/main_window.py @@ -1143,16 +1143,27 @@ class SrasViewerWindow(QMainWindow): # Progress dialogs # ------------------------------------------------------------------ - def _show_progress(self, key: str, message: str, maximum: int = 0): + def _show_progress(self, key: str, message: str, maximum: int = 0, + on_cancel=None): """Show (or relabel) the progress dialog under *key*. maximum=0 gives - an indeterminate busy indicator.""" + 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. + """ 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…") - dlg.setCancelButton(None) + if on_cancel is None: + dlg.setCancelButton(None) + else: + dlg.setCancelButtonText("Cancel") + dlg.canceled.connect(on_cancel) dlg.setWindowModality(Qt.WindowModality.WindowModal) dlg.setMinimumDuration(300) # only appears if it takes > 300 ms dlg.show() diff --git a/sras_workers.py b/sras_workers.py index 7e3ebb2..237b37e 100644 --- a/sras_workers.py +++ b/sras_workers.py @@ -15,6 +15,7 @@ import numpy as np 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, ) @@ -381,7 +382,11 @@ class CrossCorrelateWorker(_PooledWorker): def __init__(self, sras: SrasFile, ref_angle_idx: int, angle_indices: list[int], dc4_mv: dict[int, np.ndarray], *, sources: tuple[str, ...], dc_threshold_mv: float, - search_deg: float): + search_deg: float, reg_kwargs: dict | None = None): + """*reg_kwargs* is splatted into register_angle_to_reference on top of + the named arguments — the wizard's rotation-search controls (seed, + signs, refine, grid sizes) go through here, so exposing another knob + needs no change to this class.""" super().__init__() self._sras = sras self._ref = ref_angle_idx @@ -390,6 +395,7 @@ class CrossCorrelateWorker(_PooledWorker): self._sources = sources self._threshold = dc_threshold_mv self._search_deg = search_deg + self._reg_kwargs = dict(reg_kwargs or {}) def _plan(self) -> int: return compute.registration_workers(self._sras) @@ -401,9 +407,45 @@ class CrossCorrelateWorker(_PooledWorker): return a, compute.register_angle_to_reference( self._sras, a, self._ref, self._dc4_mv, dc_threshold_mv=self._threshold, sources=self._sources, - search_deg=self._search_deg) + search_deg=self._search_deg, **self._reg_kwargs) def _emit(self, result): a, fit = result self.angle_done.emit(a, fit.rotation_deg, fit.shift_mm[0], fit.shift_mm[1], fit.score, fit.source) + + +class AlignedExportWorker(CancellableWorker): + """Writes the aligned, cropped .sras on a background thread. + + Unlike every other worker here this one produces a *file*, which changes + what cancellation has to mean: write_aligned_sras stages into a ".part" + sibling and removes it when should_stop() fires, so a cancelled or crashed + export leaves nothing behind. That matters more than it sounds — a + truncated .sras is not detectably broken, since the v6 parser reads a short + file as an aborted scan and opens it happily. + + Cancellation is polled per output row chunk, the same granularity + CancellableWorker's docstring justifies, so closing the window never waits + on a multi-gigabyte write. + """ + progress = pyqtSignal(int) # 0-100 + finished = pyqtSignal(str, str) # written path ("" = none), error + + def __init__(self, sras: SrasFile, result, out_path: str): + super().__init__() + self._sras = sras + self._result = result + self._out_path = out_path + + def run(self): + try: + written = write_aligned_sras( + self._sras, self._result, self._out_path, + progress_cb=self.progress.emit, should_stop=self._stopped) + if self._stopped(): + self.finished.emit("", "") # cancelled: no file, no error + else: + self.finished.emit(str(written), "") + except Exception as exc: + self.finished.emit("", str(exc)) diff --git a/tests/test_align_export.py b/tests/test_align_export.py new file mode 100644 index 0000000..0e6326c --- /dev/null +++ b/tests/test_align_export.py @@ -0,0 +1,612 @@ +"""Aligned/cropped .sras export: does the written file actually hold the +alignment the viewer showed? + +The export is the one place an alignment stops being a transform applied on the +fly and becomes bytes on disk, so these tests care about two things above all: +the file's geometry describes what was written, and the pixels in it are the +same pixels apply_alignment would have drawn. The strongest check is the +round-trip — register the exported file against itself and demand identity, +which no amount of self-consistent-but-wrong index math can fake. + +No Qt: this exercises sras_align_export and sras_compute directly. +""" + +import struct + +import numpy as np +import pytest + +import sras_align_export as export +import sras_compute as compute +from sras_format import CH3_IDX, CH4_IDX, HDR_SIZE_V6, SrasFile, adc_to_mv, mv_to_adc +import tools.make_test_sras as gen + +_THRESHOLD_MV = 80.0 +# Same reasoning as tests/test_alignment.py: a quarter degree is already +# sub-pixel for this sample at the registration pitch. +_ROT_TOL_DEG = 0.5 +_SHIFT_TOL_MM = 0.02 + + +def dc_mv(sras: SrasFile, angle_idx: int, ch: int = CH4_IDX) -> np.ndarray: + return adc_to_mv(compute.compute_dc_image(sras, angle_idx, ch), *sras.cal(ch)) + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory): + """The rotating-sample scan, its truth alignment, and its export.""" + tmpdir = tmp_path_factory.mktemp("sras_export") + src_path = tmpdir / "rotating.sras" + meta = gen.write_rotating(src_path, n_angles=4) + sras = SrasFile(str(src_path)) + + params = {a: compute.ManualAngleParams(rot, shift) + for a, (rot, shift) in meta["truth"].items()} + result = compute.build_manual_alignment(sras, 0, _THRESHOLD_MV, params) + + out_path = tmpdir / "rotating_aligned.sras" + export.write_aligned_sras(sras, result, out_path) + return type("Rig", (), dict( + tmpdir=tmpdir, src_path=src_path, sras=sras, meta=meta, + result=result, out_path=out_path, out=SrasFile(str(out_path)))) + + +# --------------------------------------------------------------------------- +# Geometry and file structure +# --------------------------------------------------------------------------- + +def test_output_is_v6_with_uniform_geometry(rig): + out, result = rig.out, rig.result + n_rows, n_cols = result.canvas_shape + + assert out.version == 6 + assert out.n_angles == rig.sras.n_angles + assert set(out.n_rows) == {n_rows}, "every angle must share the canvas rows" + assert set(out.n_frames) == {n_cols}, "every angle must share the canvas frames" + assert np.allclose(out.x_start_mm, result.canvas_origin_mm[0]) + # x_delta must stay velocity/laser_freq or x_axis_mm() contradicts the + # geometry table; the canvas pitch is the reference angle's own pitch, so + # this is exact rather than approximate. + assert np.allclose(out.x_delta_mm_per_angle, rig.sras.pixel_x_mm) + assert out.pixel_x_mm == pytest.approx(rig.sras.pixel_x_mm) + + +def test_row_table_matches_the_canvas(rig): + expected = (rig.result.canvas_origin_mm[1] + + np.arange(rig.result.canvas_shape[0]) * rig.result.canvas_dy_mm) + for a in range(rig.out.n_angles): + assert rig.out.y_positions_mm(a) == pytest.approx(expected, abs=1e-4) + + +def test_angle_table_and_calibration_round_trip(rig): + assert rig.out.angles_deg == pytest.approx(rig.sras.angles_deg) + for ch in range(rig.sras.n_channels): + assert rig.out.cal(ch) == pytest.approx(rig.sras.cal(ch)) + assert rig.out.samples_per_frame == rig.sras.samples_per_frame + assert rig.out.bytes_per_sample == rig.sras.bytes_per_sample + assert rig.out.n_channels == rig.sras.n_channels + assert rig.out.background == pytest.approx(rig.sras.background) + + +def test_no_cache_tail(rig): + """File ends exactly at the waveform data — nothing trailing. + + A stale cache tail would be indexed by the *input's* grid, so the export + must not carry one; asserting on the exact file size is what proves it, + since a v7 tail would simply be ignored by a v6 parser. + """ + end = max(off + n for _, off, n in rig.out.iter_angle_blocks()) + assert rig.out_path.stat().st_size == end + assert all(img is None for img in rig.out.precomputed_dc4_mv) + + +def test_declared_header_size_is_v6(rig): + raw = rig.out_path.read_bytes()[:HDR_SIZE_V6] + magic, version, n_angles = struct.unpack(">4sBH", raw[:7]) + assert (magic, version, n_angles) == (b"SRAS", 6, rig.sras.n_angles) + + +# --------------------------------------------------------------------------- +# The pixels themselves +# --------------------------------------------------------------------------- + +def test_export_matches_apply_alignment(rig): + """The exported waveforms decode to the same DC image the viewer drew — + over the *whole* canvas, padding included. + + Two rules have to be exactly right for this, and each fails differently: + * rounding must be floor(x + 0.5), not np.rint, or pixels on exact + half-integer boundaries pick the neighbouring source pixel; + * out-of-bounds must be tested on the fractional coordinate against + [0, n-1], not on the rounded index, or a one-pixel rim gets real data + where the preview shows padding. + Comparing every pixel rather than only the interior is what catches the + second one, since a rim discrepancy hides inside a `preview != 0` mask. + """ + for a in range(rig.sras.n_angles): + preview = compute.apply_alignment(rig.result, a, dc_mv(rig.sras, a)) + actual = dc_mv(rig.out, a) + assert actual.shape == preview.shape + # Padding matches to within half an ADC step: apply_alignment pads with + # literal 0.0 mV, the export with the nearest integer ADC code to 0 mV. + tol = abs(rig.sras.cal(CH4_IDX)[0]) / 2.0 + 1e-4 + # Exclude the epsilon rim the export deliberately keeps and scipy drops + # (see test_edge_tolerance_only_affects_the_epsilon_rim). + sr, sc = export._src_coords(rig.result.per_angle[a], + np.arange(preview.shape[0]), preview.shape[1]) + rim = export._in_bounds(sr, sc, *rig.sras.image_shape(a)) & (preview == 0.0) + cmp = ~rim + assert actual[cmp] == pytest.approx(preview[cmp], abs=tol), \ + f"angle {a}: exported pixels differ from the aligned preview" + # And exactly, wherever there is real data. + inside = (preview != 0.0) + assert inside.any(), f"angle {a}: preview is entirely padding" + assert actual[inside] == pytest.approx(preview[inside], abs=1e-6), \ + f"angle {a}: exported data pixels are not bit-equal to the preview" + + +def test_reference_angle_is_exported_whole(rig): + """The reference angle must survive as a complete, exact integer crop. + + It is the coordinate authority — its transform is the identity with an + integer offset by construction — so every one of its source pixels has to + appear in the export. This is what _EDGE_TOL exists for: that offset comes + out of the mm-space affine chain as -20 - 7e-15, and a bare `>= 0` bounds + test silently drops the angle's entire first row and last column. + """ + ref = rig.result.ref_angle_idx + src_rows, src_frames = rig.sras.image_shape(ref) + plan = export.plan_export(rig.sras, rig.result) + assert plan.valid_px[ref] == src_rows * src_frames, \ + "reference angle lost pixels to the in-bounds test" + + # And the values themselves land as an exact, unrotated block. + src_img = dc_mv(rig.sras, ref) + out_img = dc_mv(rig.out, ref) + t = rig.result.per_angle[ref] + row0, col0 = (int(round(-t.offset[0])), int(round(-t.offset[1]))) + assert np.array_equal(out_img[row0:row0 + src_rows, col0:col0 + src_frames], + src_img), \ + "reference angle is not a verbatim block in the export" + + +def test_edge_tolerance_only_affects_the_epsilon_rim(rig): + """Where the export's bounds test and scipy's disagree, the coordinate must + be within _EDGE_TOL of the boundary — i.e. only pixels whose scipy answer + was itself decided by float noise, never a real half-pixel decision.""" + for a in range(rig.sras.n_angles): + t = rig.result.per_angle[a] + src_rows, src_frames = rig.sras.image_shape(a) + n_rows, n_cols = rig.result.canvas_shape + + ones = np.ones((src_rows, src_frames), dtype=np.float32) + scipy_valid = compute.apply_alignment(rig.result, a, ones) > 0.5 + sr, sc = export._src_coords(t, np.arange(n_rows), n_cols) + ours = export._in_bounds(sr, sc, src_rows, src_frames) + + differ = ours != scipy_valid + assert not (scipy_valid & ~ours).any(), \ + f"angle {a}: export drops pixels scipy keeps" + if differ.any(): + # Every disagreement sits within the tolerance of an edge. + near = (np.abs(sr) <= export._EDGE_TOL) + near |= (np.abs(sr - (src_rows - 1)) <= export._EDGE_TOL) + near |= (np.abs(sc) <= export._EDGE_TOL) + near |= (np.abs(sc - (src_frames - 1)) <= export._EDGE_TOL) + assert near[differ].all(), \ + f"angle {a}: bounds differ away from the epsilon rim" + + +def test_export_matches_apply_alignment_on_ch3(rig): + """Channel-agnostic: the gather moves whole pixels, not per-channel images.""" + for a in range(rig.sras.n_angles): + preview = compute.apply_alignment(rig.result, a, + dc_mv(rig.sras, a, CH3_IDX)) + actual = dc_mv(rig.out, a, CH3_IDX) + inside = preview != 0.0 + assert actual[inside] == pytest.approx(preview[inside], abs=1e-3) + + +def test_padding_is_zero_mv_not_zero_adc(rig): + """Unreachable canvas pixels must read as ~0 mV on every channel. + + Filling with literal zero ADC would decode to (0 - yoff) * ymult + yzero — + for this fixture's CH4 calibration that is +100 mV, well above any sensible + mask threshold, so the padding would masquerade as valid sample everywhere. + """ + a = rig.sras.n_angles - 1 + t = rig.result.per_angle[a] + n_rows, n_cols = rig.result.canvas_shape + src_rows, src_frames = rig.sras.image_shape(a) + + sr, sc = export._src_coords(t, np.arange(n_rows), n_cols) + outside = ~export._in_bounds(sr, sc, src_rows, src_frames) + assert outside.any(), "rotated angle should leave unreachable canvas corners" + + for ch in (CH3_IDX, CH4_IDX): + img = dc_mv(rig.out, a, ch) + half_step = abs(rig.sras.cal(ch)[0]) / 2.0 + assert np.abs(img[outside]).max() <= half_step + 1e-6, \ + f"CH{ch} padding is not within half an ADC step of 0 mV" + + # And the sanity check that makes the above meaningful: zero ADC would not + # have passed it. + assert abs(adc_to_mv(0, *rig.sras.cal(CH4_IDX))) > 10.0 + + +def test_reregistering_the_export_is_identity(rig): + """The export really is aligned: registering it against its own angle 0 + recovers no rotation and no shift. + + The end-to-end check — it fails for any index error, sign flip, wrong pivot + or origin mistake anywhere in crop/affine/gather, in a way the + self-consistency tests above cannot. + """ + dc4 = {a: dc_mv(rig.out, a) for a in range(rig.out.n_angles)} + for a in range(1, rig.out.n_angles): + fit = compute.register_angle_to_reference( + rig.out, a, 0, dc4, dc_threshold_mv=_THRESHOLD_MV, + seed_deg=0.0, seed_signs=(1,)) + assert abs(fit.rotation_deg) <= _ROT_TOL_DEG, \ + f"angle {a} still rotated by {fit.rotation_deg:.3f}° after export" + assert float(np.hypot(*fit.shift_mm)) <= _SHIFT_TOL_MM, \ + f"angle {a} still shifted by {fit.shift_mm} mm after export" + + +def test_export_of_int16_input(rig, tmp_path): + """bps=2 inputs keep their big-endian int16 dtype through the gather.""" + src_path = tmp_path / "i16.sras" + gen.write(src_path, n_angles=2, samples_per_frame=16, bps=2) + sras = SrasFile(str(src_path)) + result = compute.build_manual_alignment(sras, 0, 0.0, {}) + + out_path = tmp_path / "i16_aligned.sras" + export.write_aligned_sras(sras, result, out_path) + out = SrasFile(str(out_path)) + + assert out.bytes_per_sample == 2 + assert out.data[0].dtype == np.dtype(">i2") + for a in range(sras.n_angles): + preview = compute.apply_alignment(result, a, dc_mv(sras, a)) + actual = dc_mv(out, a) + inside = preview != 0.0 + assert actual[inside] == pytest.approx(preview[inside], abs=1e-3) + + +# --------------------------------------------------------------------------- +# Cropping +# --------------------------------------------------------------------------- + +def test_crop_is_a_window_of_the_full_canvas(rig): + """crop_alignment_result must resample exactly the sub-rectangle it names. + + Asserted as bit-exact equality, not approximately: the crop composes into + the affine's offset by an integer number of canvas pixels, so anything but + an exact match means the composition is wrong. + """ + n_rows, n_cols = rig.result.canvas_shape + row0, col0 = n_rows // 5, n_cols // 4 + nr, nc = n_rows // 2, n_cols // 3 + cropped = compute.crop_alignment_result(rig.result, row0, col0, nr, nc) + + assert cropped.canvas_shape == (nr, nc) + assert cropped.canvas_origin_mm[0] == pytest.approx( + rig.result.canvas_origin_mm[0] + col0 * rig.result.canvas_dx_mm) + assert cropped.canvas_origin_mm[1] == pytest.approx( + rig.result.canvas_origin_mm[1] + row0 * rig.result.canvas_dy_mm) + + for a in range(rig.sras.n_angles): + img = dc_mv(rig.sras, a) + full = compute.apply_alignment(rig.result, a, img) + assert np.array_equal( + compute.apply_alignment(cropped, a, img), + full[row0:row0 + nr, col0:col0 + nc]), \ + f"angle {a}: cropped resample is not the same window" + # Rotation/shift are properties of the angle, not of the canvas. + assert cropped.per_angle[a].rotation_deg == rig.result.per_angle[a].rotation_deg + assert cropped.per_angle[a].shift_mm == rig.result.per_angle[a].shift_mm + + +def test_cropped_export_round_trips(rig, tmp_path): + n_rows, n_cols = rig.result.canvas_shape + row0, col0, nr, nc = n_rows // 4, n_cols // 4, n_rows // 2, n_cols // 2 + cropped = compute.crop_alignment_result(rig.result, row0, col0, nr, nc) + + out_path = tmp_path / "cropped.sras" + export.write_aligned_sras(rig.sras, cropped, out_path) + out = SrasFile(str(out_path)) + + assert set(out.n_rows) == {nr} and set(out.n_frames) == {nc} + assert out.x_start_mm[0] == pytest.approx(cropped.canvas_origin_mm[0], abs=1e-4) + for a in range(rig.sras.n_angles): + preview = compute.apply_alignment(cropped, a, dc_mv(rig.sras, a)) + actual = dc_mv(out, a) + inside = preview != 0.0 + if inside.any(): + assert actual[inside] == pytest.approx(preview[inside], abs=1e-3) + + +def test_crop_rejects_empty_window(rig): + with pytest.raises(ValueError, match="empty crop"): + compute.crop_alignment_result(rig.result, 0, 0, 0, 10) + with pytest.raises(ValueError, match="empty crop"): + compute.crop_alignment_result(rig.result, 0, 0, 10, -1) + + +# --------------------------------------------------------------------------- +# plan_export and overlap_stats +# --------------------------------------------------------------------------- + +def test_plan_export_matches_what_was_written(rig): + plan = export.plan_export(rig.sras, rig.result) + n_rows, n_cols = rig.result.canvas_shape + assert (plan.n_rows, plan.n_frames) == (n_rows, n_cols) + assert plan.n_angles == rig.sras.n_angles + + data_bytes = sum(n for _, _, n in rig.out.iter_angle_blocks()) + assert plan.total_bytes == data_bytes + assert plan.bytes_per_angle * plan.n_angles == plan.total_bytes + + # Coverage must agree with the pixels that actually carry data. The + # reference angle is unrotated, so its whole footprint lands inside. + ref_px = np.prod(rig.sras.image_shape(0)) + assert plan.valid_px[0] == ref_px + for a in range(1, rig.sras.n_angles): + assert 0 < plan.valid_px[a] <= n_rows * n_cols + assert 0.0 < plan.coverage_frac(a) < 1.0 + + +def test_plan_export_flags_a_crop_that_misses_an_angle(rig): + """A crop over a corner the rotated angles cannot reach must warn, and the + export must still succeed by writing that angle as padding.""" + n_rows, n_cols = rig.result.canvas_shape + corner = compute.crop_alignment_result(rig.result, 0, 0, + max(1, n_rows // 12), + max(1, n_cols // 12)) + plan = export.plan_export(rig.sras, corner) + empty = [a for a in range(rig.sras.n_angles) if plan.valid_px[a] == 0] + assert empty, "top-left canvas corner should be unreachable for some angle" + assert any("all padding" in w for w in plan.warnings) + + +def test_overlap_stats(): + counts = np.array([[0, 1, 2], [3, 3, 0], [0, 2, 3]]) + stats = compute.overlap_stats(counts, 3) + assert stats["union_px"] == 6 + assert stats["full_px"] == 3 + assert stats["full_frac"] == pytest.approx(0.5) + assert stats["max_count"] == 3 + assert stats["mean_count"] == pytest.approx((1 + 2 + 3 + 3 + 2 + 3) / 6) + assert stats["empty"] is False + + empty = compute.overlap_stats(np.zeros((4, 4), dtype=int), 3) + assert empty["empty"] is True + assert empty["full_frac"] == 0.0 and empty["mean_count"] == 0.0 + + +def test_largest_rect_at_least(): + # A 2x3 block of 3s with a notch that a bounding box would swallow. + counts = np.array([ + [0, 0, 0, 0, 0], + [0, 3, 3, 3, 0], + [0, 3, 3, 3, 0], + [0, 3, 0, 3, 0], + ]) + row0, col0, nr, nc = compute.largest_rect_at_least(counts, 3) + assert (nr * nc) == 6 and (row0, col0, nr, nc) == (1, 1, 2, 3) + assert (counts[row0:row0 + nr, col0:col0 + nc] >= 3).all() + + # A column taller than the wide block is the better rectangle. + tall = np.array([[3, 3], [3, 0], [3, 0], [3, 0]]) + r0, c0, nr2, nc2 = compute.largest_rect_at_least(tall, 3) + assert (r0, c0, nr2, nc2) == (0, 0, 4, 1) + + assert compute.largest_rect_at_least(np.zeros((3, 3), dtype=int), 1) is None + # Whole-array case: no notch, so the answer is the array itself. + assert compute.largest_rect_at_least(np.full((3, 4), 2), 2) == (0, 0, 3, 4) + + +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.""" + n = rig.sras.n_angles + masks = {a: (dc_mv(rig.sras, a) >= _THRESHOLD_MV).astype(np.float32) + for a in range(n)} + counts = sum(compute.apply_alignment(rig.result, a, masks[a]) > 0.5 + for a in range(n)).astype(int) + assert counts.max() == n, "fixture alignment should have a full-overlap region" + + rect = compute.largest_rect_at_least(counts, n) + assert rect is not None + row0, col0, nr, nc = rect + assert (counts[row0:row0 + nr, col0:col0 + nc] == n).all(), \ + "convenience crop must not include pixels some angle misses" + + # And it must beat the naive bounding box, which here is impure. + rr, cc = np.nonzero(counts == n) + bbox_pure = (counts[rr.min():rr.max() + 1, cc.min():cc.max() + 1] == n).all() + assert not bbox_pure, "fixture no longer exercises the bounding-box hazard" + + +# --------------------------------------------------------------------------- +# Legacy inputs, validation and durability +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("version", [2, 4]) +def test_legacy_input_exports_as_v6(version, tmp_path): + """v2-v5 inputs keep no verbatim preamble/background spans, so those + sections have to be re-encoded. v2 additionally has neither.""" + src_path = tmp_path / f"legacy_v{version}.sras" + gen.write_legacy(src_path, version=version, n_angles=2) + sras = SrasFile(str(src_path)) + result = compute.build_manual_alignment(sras, 0, 0.0, {}) + + out_path = tmp_path / f"legacy_v{version}_aligned.sras" + export.write_aligned_sras(sras, result, out_path) + out = SrasFile(str(out_path)) + + assert out.version == 6 + assert out.n_angles == sras.n_angles + # A zero background rather than a zero-length one: consumers subtract it + # from a (spf,)-shaped row, which a length-0 array cannot broadcast against. + assert out.background is not None + assert out.background.size == sras.samples_per_frame + if sras.background is None: + assert np.all(out.background == 0) + assert any("no background" in w for w in + export.plan_export(sras, result).warnings) + # Calibration must survive: v2 has no preambles and falls back to the + # hardcoded scope constants, and the re-encoded empty preambles must land on + # exactly the same fallback. + for ch in range(sras.n_channels): + assert out.cal(ch) == pytest.approx(sras.cal(ch)) + for a in range(sras.n_angles): + preview = compute.apply_alignment(result, a, dc_mv(sras, a)) + actual = dc_mv(out, a) + inside = preview != 0.0 + assert actual[inside] == pytest.approx(preview[inside], abs=1e-3) + + +def test_too_many_rows_is_rejected_before_writing(rig, tmp_path): + """The geometry table stores n_rows as a u16; silently truncating would + write a file whose header disagrees with its own waveform block.""" + huge = compute.crop_alignment_result(rig.result, 0, 0, 70000, 4) + out_path = tmp_path / "huge.sras" + with pytest.raises(ValueError, match="exceeds the .sras per-angle geometry"): + export.write_aligned_sras(rig.sras, huge, out_path) + assert not out_path.exists() + assert not out_path.with_name(out_path.name + ".part").exists() + + +def test_missing_transform_is_rejected(rig, tmp_path): + broken = compute.crop_alignment_result(rig.result, 0, 0, + *rig.result.canvas_shape) + del broken.per_angle[1] + with pytest.raises(ValueError, match="no transform for angle"): + export.write_aligned_sras(rig.sras, broken, tmp_path / "broken.sras") + + +def test_cancelled_export_leaves_nothing_behind(rig, tmp_path): + out_path = tmp_path / "cancelled.sras" + written = export.write_aligned_sras(rig.sras, rig.result, out_path, + should_stop=lambda: True) + assert written == out_path + assert not out_path.exists(), "cancelled export must not leave an output file" + assert not out_path.with_name(out_path.name + ".part").exists() + + +def test_failed_write_leaves_nothing_behind(rig, tmp_path): + """An exception mid-write must remove the partial file: a short .sras is + not detectably broken — the v6 parser reads it as an aborted scan.""" + out_path = tmp_path / "boom.sras" + + def explode(_pct): + raise RuntimeError("boom") + + with pytest.raises(RuntimeError, match="boom"): + export.write_aligned_sras(rig.sras, rig.result, out_path, + progress_cb=explode) + assert not out_path.exists() + assert not out_path.with_name(out_path.name + ".part").exists() + + +def test_progress_is_monotonic_and_completes(rig, tmp_path): + seen: list[int] = [] + export.write_aligned_sras(rig.sras, rig.result, tmp_path / "prog.sras", + progress_cb=seen.append) + assert seen and seen[-1] == 100 + assert seen == sorted(seen) + assert all(0 <= p <= 100 for p in seen) + + +def test_band_reader_path_is_byte_identical(rig, tmp_path, monkeypatch): + """A source block too large to hold in RAM is served from sliding bands + instead. That path only runs on multi-gigabyte scans, so force it with a + tiny budget and demand the same bytes — otherwise the one code path that + matters on real data is the one never tested.""" + whole = tmp_path / "whole.sras" + export.write_aligned_sras(rig.sras, rig.result, whole) + + monkeypatch.setattr(compute, "_TOTAL_BYTES_BUDGET", 4096) + banded = tmp_path / "banded.sras" + export.write_aligned_sras(rig.sras, rig.result, banded) + + assert banded.read_bytes() == whole.read_bytes() + + +def test_row_chunking_is_invariant(rig, tmp_path, monkeypatch): + """Output must not depend on how many rows are buffered per write.""" + base = tmp_path / "base.sras" + export.write_aligned_sras(rig.sras, rig.result, base) + + monkeypatch.setattr(export, "_ROW_CHUNK", 1) + one = tmp_path / "one.sras" + export.write_aligned_sras(rig.sras, rig.result, one) + assert one.read_bytes() == base.read_bytes() + + +def test_refuses_to_overwrite_the_source(rig): + """The source's waveform blocks are live read-only memmaps; writing over + the file would corrupt the reads the gather is making from it.""" + with pytest.raises(ValueError, match="refusing to export onto the source"): + export.write_aligned_sras(rig.sras, rig.result, rig.src_path) + assert SrasFile(str(rig.src_path)).n_angles == rig.sras.n_angles + + +def test_overwrites_an_existing_file(rig, tmp_path): + out_path = tmp_path / "existing.sras" + out_path.write_bytes(b"not a scan") + export.write_aligned_sras(rig.sras, rig.result, out_path) + assert SrasFile(str(out_path)).version == 6 + + +# --------------------------------------------------------------------------- +# The registration knobs the wizard exposes +# --------------------------------------------------------------------------- + +def test_locked_rotation_returns_exactly_the_seed(rig): + """search_deg=0 + one sign + refine=False pins rotation to the seed, which + is what "lock rotation to the stage angle" means on the wizard's first + page. Only the translation may be searched.""" + dc4 = {a: dc_mv(rig.sras, a) for a in range(rig.sras.n_angles)} + for a in range(1, rig.sras.n_angles): + nominal = compute.nominal_delta_deg(rig.sras, a, 0) + fit = compute.register_angle_to_reference( + rig.sras, a, 0, dc4, dc_threshold_mv=_THRESHOLD_MV, + search_deg=0.0, coarse_step_deg=2.0, seed_signs=(-1,), refine=False) + assert fit.rotation_deg == pytest.approx(-nominal) + + +def test_seed_deg_overrides_the_stage_angle(rig): + """seed_deg=0.0 searches around no rotation at all, so a scan whose angles + are genuinely ~37° apart must fail to find them within a ±2° window — + proving the seed is what positions the search.""" + dc4 = {a: dc_mv(rig.sras, a) for a in range(rig.sras.n_angles)} + fit = compute.register_angle_to_reference( + rig.sras, 1, 0, dc4, dc_threshold_mv=_THRESHOLD_MV, + search_deg=2.0, seed_deg=0.0, seed_signs=(1,), refine=False) + truth_rot = rig.meta["truth"][1][0] + assert abs(fit.rotation_deg) <= 2.0 + assert abs(fit.rotation_deg - truth_rot) > 10.0 + + +def test_rotation_candidates_signs(): + both = compute._rotation_candidates(10.0, 2.0, 2.0) + assert both == compute._rotation_candidates(10.0, 2.0, 2.0, (-1, 1)), \ + "default must stay the both-signs sweep" + assert compute._rotation_candidates(10.0, 0.0, 2.0, (1,)) == [10.0] + assert compute._rotation_candidates(10.0, 0.0, 2.0, (-1,)) == [-10.0] + # A zero seed collapses the two windows; the dedupe must keep one copy. + assert compute._rotation_candidates(0.0, 2.0, 2.0) == [-2.0, 0.0, 2.0] + + +def test_zero_mv_fill_code_is_clipped_to_dtype(): + """mv_to_adc is unclamped, so the fill code must be clipped or the int8 + cast wraps around to a large-magnitude value.""" + fake = type("S", (), dict( + n_channels=1, samples_per_frame=2, + cal=lambda self, ch: (1e-6, 0.0, 5000.0)))() + row = export._fill_row(fake, 3, np.dtype(np.int8)) + assert row.shape == (1, 3, 2) + assert row.min() == row.max() == np.iinfo(np.int8).min + assert mv_to_adc(0.0, 1e-6, 0.0, 5000.0) < np.iinfo(np.int8).min