#!/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