From ed0eba4bae4d57b913a355890c3eac55b6b1f9df Mon Sep 17 00:00:00 2001 From: Thomas Ales Date: Thu, 30 Jul 2026 22:27:56 -0500 Subject: [PATCH 01/10] Add v7 format with on-disk DC/FFT cache, replace v5 pre-process with Convert menu Scans come off the scope as v6; a new trailing CACH section (independent SDCB/SFFT blocks, sized per-angle from the existing geometry table) lets computed DC and FFT images be cached in place, bumping the file to v7 the first time either is stored. The Convert menu's "Batch Compute DC and Store"/"Batch Compute FFT and Store" actions run this across multiple files in the background. compute_rf_image and the DC compute workers now reuse cached v7 data instead of recomputing it. Removes "Pre-process and Save as v5", which never supported v6 sources and is superseded by in-place v7 caching. Legacy v2-v5 reading, including v5's PREC section, is unaffected. Co-Authored-By: Claude Sonnet 5 --- scan_format.md | 119 +++++++++- sras_viewer.py | 634 ++++++++++++++++++++++++++++++++----------------- 2 files changed, 530 insertions(+), 223 deletions(-) diff --git a/scan_format.md b/scan_format.md index c864ff4..3c90e16 100644 --- a/scan_format.md +++ b/scan_format.md @@ -1,4 +1,4 @@ -# SRAS Scan Binary Format — Version 6 +# SRAS Scan Binary Format — Version 6 / 7 Each `.sras` file contains **one complete scan**: all GR rotation angles and all Y rows. Files are named `{prefix}.sras`. @@ -10,6 +10,14 @@ rotated by that specific angle** — not the worst case across all angles — so fewer rows than a 45° scan of the same ROI, and the file format reflects that instead of forcing every angle to the largest bounding box. +v7 is byte-identical to v6 (same header, angle table, geometry table, row +table, preamble blocks, background block, waveform data) plus an optional +trailing **Cache Tail** holding precomputed per-angle DC and/or FFT images +(see [Cache Tail (v7)](#cache-tail-v7) below) — only the header's `version` +field and the presence of that trailing section differ. Scans come off the +scope as v6; `sras_viewer.py`'s "Convert" menu batch actions convert a file +to v7 **in place** the first time either cache block is computed and stored. + --- ## File Layout @@ -22,6 +30,7 @@ instead of forcing every angle to the largest bounding box. [Preamble Blocks — n_channels × (uint16 length + UTF-8 WFMOutpre string)] [Background Block — uint32 n_bg_samples + n_bg_samples × int8 bytes] [Waveform Data (ragged) — per angle: n_rows[a] × n_channels × n_frames[a] × samples_per_frame × bps bytes] +[Cache Tail (optional) — "CACH" + DC block (optional) + FFT block (optional); v7 only] ``` All multi-byte integers and floats use **big-endian** byte order @@ -182,6 +191,111 @@ using that angle's `x_start` from the Per-Angle Geometry Table (not --- +## Cache Tail (v7) + +Present iff `version == 7` and `file_size > cache_offset`, where: + +``` +cache_offset = data_offset + Σ over angles a of: + n_rows[a] × n_channels × n_frames[a] × samples_per_frame × bytes_per_sample +``` + +i.e. exactly `data_offset + waveform_bytes` — the same "Total data size" +formula as Waveform Data above. This offset is derivable from the header and +Per-Angle Geometry Table alone and does **not** depend on which cache +block(s) are present, so a writer can always seek straight there without +reading or touching any waveform byte before it. + +Unlike the (removed) v5 `PREC` section, which stored one omnibus per-angle +entry (FFT + both DC channels together) in a dense, uniform-geometry array, +the v7 Cache Tail splits DC and FFT into two **independent** sub-blocks — +each sized per-angle from the Per-Angle Geometry Table, each independently +present, and each independently updatable in any order, any number of +times, without disturbing the other. This matches `sras_viewer.py`'s +"Convert" menu, which exposes DC and FFT store as two separate batch +actions. + +### CACH outer header (6 bytes, `">4sBB"`) + +| Offset | Size | Type | Field | Description | +|--------|------|------|-------|-------------| +| 0 | 4 | `char[4]` | `cach_magic` | `CACH` (ASCII). Missing/wrong magic → treat file as having no cache. | +| 4 | 1 | `u8` | `cach_version` | Cache format version. Currently `1`. Readers must treat the file as uncached if this is not a version they understand (unlike v5's `PREC` section, which read but never validated its version byte). | +| 5 | 1 | `u8` | `block_flags` | Bit 0 = DC block (`SDCB`) follows. Bit 1 = FFT block (`SFFT`) follows, immediately after the DC block if both are present. Bits 2–7 reserved, must be zero on write. | + +### DC block `SDCB` (present iff `block_flags & 0x01`) + +7-byte block header, format `">4sBH"`: + +| Offset (rel) | Size | Type | Field | Description | +|--------------|------|------|-------|-------------| +| 0 | 4 | `char[4]` | `magic` | `SDCB` | +| 4 | 1 | `u8` | `reserved` | `0`, reserved for future use | +| 5 | 2 | `u16` | `n_stored` | Number of angle entries that follow, `0 ≤ n_stored ≤ n_angles` | + +followed by `n_stored` entries, each: + +``` +u16 angle_idx — index into the angle table (0-based) +f32[n_rows[angle_idx] × n_frames[angle_idx]] dc3_mv — CH3 waveform mean, mV, row-major +f32[n_rows[angle_idx] × n_frames[angle_idx]] dc4_mv — CH4 waveform mean, mV, row-major +``` + +Entries may appear in any order and need not be contiguous from angle 0 — +this supports storing (or re-storing) a subset of angles, or an +interrupted batch run leaving only some angles cached. Readers bounds-check +`angle_idx < n_angles` on each entry and stop parsing on an out-of-range +value, same as v5's `PREC` section. + +### FFT block `SFFT` (present iff `block_flags & 0x02`) + +7-byte block header, format `">4sBH"`: + +| Offset (rel) | Size | Type | Field | Description | +|--------------|------|------|-------|-------------| +| 0 | 4 | `char[4]` | `magic` | `SFFT` | +| 4 | 1 | `u8` | `flags` | Bit 0 = `bg_sub_applied` — background waveform was subtracted from CH1 before the FFT when these images were computed. Bits 1–7 reserved. | +| 5 | 2 | `u16` | `n_stored` | Number of angle entries that follow | + +followed by `n_stored` entries, each: + +``` +u16 angle_idx — index into the angle table (0-based) +f32[n_rows[angle_idx] × n_frames[angle_idx]] peak_freq_mhz — CH1 FFT peak frequency, MHz, row-major +``` + +**`peak_freq_mhz`** is computed without any DC-threshold masking (i.e. the +FFT is run on every pixel unconditionally, same as v5's `PREC` convention). +Readers apply the DC4 threshold at display time: + +``` +pixel is valid ⟺ dc4_mv[r][f] ≥ threshold_mv +display_value = peak_freq_mhz[r][f] if valid, else 0 +``` + +using the DC4 image from the DC block if that angle is also cached there, +else computed on demand. + +Readers must fall back to real-time FFT computation (ignoring stored +`peak_freq_mhz`) under the same conditions as v5's PREC fast path: time-domain +gating is active, zero-padding (`n_fft ≠ samples_per_frame`) is requested, or +the reader's background-subtraction setting doesn't match `flags.bg_sub_applied`. + +### In-place write ordering + +A writer updating a file's Cache Tail must write the payload (CACH header + +whichever block(s) are present) **before** flipping the header's `version` +byte to `7`, and `truncate()` the file to the new payload's end immediately +after writing it. If the process is interrupted between the payload write +and the version-byte flip, the file is still valid v6 — v6 parsing only +bounds-checks each angle's `offset + nbytes ≤ file_size`, it never asserts +exactly how many bytes follow the last angle's waveform block — so the +interrupted write leaves harmless trailing bytes rather than a corrupt file, +and the next successful write overwrites them via the same deterministic +`cache_offset`. + +--- + ## Acquisition Settings (fixed by sc3_aui_app.py) | Parameter | Value | @@ -206,5 +320,6 @@ using that angle's `x_start` from the Per-Angle Geometry Table (not | 3 | Added preamble blocks (WFMOutpre strings) after the row table, one length-prefixed UTF-8 block per channel. | | 4 | Added background waveform block (CH1, Helios ON / Genesis OFF) after the preamble blocks; stored as `uint32` sample count followed by raw `int8` ADC bytes. | | 5 | (skipped) | -| 6 | Each angle now scans only the bounding box of the nominal ROI rotated by that angle instead of the AABB-expanded worst case across all angles. Header no longer carries a single global `x_start`/`x_delta`/`n_rows` — replaced with `*_nominal` reference fields plus a new Per-Angle Geometry Table (`x_start`, `x_delta`, `n_frames`, `n_rows` per angle) and a ragged Row Table / Waveform Data block sized per angle. **Not compatible with pre-v6 readers** that assume uniform geometry. `sras_viewer.py` reads v6 natively (per-angle `n_rows`/`n_frames`/`x_start`); the "Pre-process and Save as v5" fast-path export is not offered for v6 files since the flat v5 layout cannot represent per-angle geometry. | +| 6 | Each angle now scans only the bounding box of the nominal ROI rotated by that angle instead of the AABB-expanded worst case across all angles. Header no longer carries a single global `x_start`/`x_delta`/`n_rows` — replaced with `*_nominal` reference fields plus a new Per-Angle Geometry Table (`x_start`, `x_delta`, `n_frames`, `n_rows` per angle) and a ragged Row Table / Waveform Data block sized per angle. **Not compatible with pre-v6 readers** that assume uniform geometry. `sras_viewer.py` reads v6 natively (per-angle `n_rows`/`n_frames`/`x_start`). | +| 7 | Adds an optional trailing **Cache Tail** (`CACH` section, see [Cache Tail (v7)](#cache-tail-v7)) after the ragged waveform data, holding independently-present, independently-updatable per-angle DC (`SDCB`: dc3_mv + dc4_mv) and FFT (`SFFT`: peak_freq_mhz) blocks, so previously-computed images redisplay instantly instead of being recomputed. Header / angle table / geometry table / row table / preamble blocks / background block / waveform data are byte-identical to v6 — only the version byte and the optional Cache Tail differ. Scans still come off the scope as v6; `sras_viewer.py`'s "Convert" menu ("Batch Compute DC and Store" / "Batch Compute FFT and Store") converts a file to v7 **in place** on first use, or updates an existing v7 file's cache blocks, without rewriting any waveform bytes. Supersedes the removed "Pre-process and Save as v5" workflow, which was never available for v6 sources since the flat v5 `PREC` layout can't represent per-angle geometry. | diff --git a/sras_viewer.py b/sras_viewer.py index 9b7e3ca..13c657d 100644 --- a/sras_viewer.py +++ b/sras_viewer.py @@ -15,11 +15,16 @@ header before acquisition, but the scope may acquire fewer frames. The actual count is computed from the file size and used for the reshape so channels are correctly aligned. -Scan geometry: v6 files scan a different bounding box per angle (x_start, +Scan geometry: v6/v7 files scan a different bounding box per angle (x_start, x_delta, n_frames, n_rows all vary by angle), so geometry is exposed per-angle via SrasFile.n_rows / n_frames / x_start_mm arrays and the x_axis_mm() / y_positions_mm() methods. v2–v5 files have uniform geometry across angles, so those arrays simply repeat the same value n_angles times. + +v7 files are v6 files with an optional trailing cache section holding +precomputed per-angle DC and/or FFT images (see scan_format.md), so display +never has to recompute them after the "Convert" menu's batch actions have +stored them once. """ import re @@ -90,6 +95,29 @@ HDR_SIZE_V6 = struct.calcsize(HDR_FMT_V6) # 49 bytes GEO_FMT_V6 = ">ffIH" GEO_SIZE_V6 = struct.calcsize(GEO_FMT_V6) # 14 bytes +# v7: identical to v6 (same header/geometry/waveform layout, version byte +# is the only header difference) plus an optional trailing cache section +# ("CACH") holding precomputed per-angle DC and/or FFT images so a v7 file +# never needs to recompute them on open. See scan_format.md for the full +# spec. Sub-block sizes are per-angle (n_rows[a] * n_frames[a]), taken from +# the Per-Angle Geometry Table already parsed for v6 — no new geometry +# fields are needed. +CACH_MAGIC = b"CACH" +CACH_HDR_FMT = ">4sBB" # magic, cach_version, block_flags +CACH_HDR_SIZE = struct.calcsize(CACH_HDR_FMT) # 6 bytes +CACH_VERSION = 1 +CACH_FLAG_DC = 0x01 +CACH_FLAG_FFT = 0x02 + +SDCB_MAGIC = b"SDCB" +SDCB_HDR_FMT = ">4sBH" # magic, reserved, n_stored +SDCB_HDR_SIZE = struct.calcsize(SDCB_HDR_FMT) # 7 bytes + +SFFT_MAGIC = b"SFFT" +SFFT_HDR_FMT = ">4sBH" # magic, flags, n_stored +SFFT_HDR_SIZE = struct.calcsize(SFFT_HDR_FMT) # 7 bytes +SFFT_FLAG_BG_SUB = 0x01 + # Fixed-order channels in the file: index 0=CH1, 1=CH3, 2=CH4 # Fixed channel indices into the .sras data array (CH1=RF, CH3/CH4=Bias DC) CH1_IDX, CH3_IDX, CH4_IDX = 0, 1, 2 @@ -149,15 +177,21 @@ def adc_to_mv(adc: float, ymult_mv: float = _FALLBACK_YMULT_MV, # --------------------------------------------------------------------------- class SrasFile: - """Parsed in-memory representation of a v2–v6 .sras file. + """Parsed in-memory representation of a v2–v7 .sras file. Scan geometry (rows, frames, x_start) is exposed per-angle via the ``n_rows`` / ``n_frames`` / ``x_start_mm`` arrays and the ``x_axis_mm()`` - / ``y_positions_mm()`` methods, since v6 files scan a different bounding - box per angle. v2–v5 files have uniform geometry, so these arrays just - repeat the same value ``n_angles`` times. Waveform data is likewise - exposed as ``data[angle_idx]``, an array of shape + / ``y_positions_mm()`` methods, since v6/v7 files scan a different + bounding box per angle. v2–v5 files have uniform geometry, so these + arrays just repeat the same value ``n_angles`` times. Waveform data is + likewise exposed as ``data[angle_idx]``, an array of shape ``(n_rows[a], n_channels, n_frames[a], samples_per_frame)``. + + v7 files (v6 plus an optional trailing cache section) expose any stored + precomputed images as ``precomputed_dc3_mv`` / ``precomputed_dc4_mv`` / + ``precomputed_freq_mhz`` — for v6/v7 these are ragged per-angle lists + (``list[np.ndarray | None]``, one entry per angle); for v5 they are dense + ``(n_angles, n_rows, n_frames)`` arrays, since v5 geometry is uniform. """ def __init__(self, path: str): @@ -174,7 +208,7 @@ class SrasFile: self.version = version if version in (2, 3, 4, 5): self._parse_legacy() - elif version == 6: + elif version in (6, 7): self._parse_v6() else: raise ValueError(f"Unsupported version: {version}") @@ -368,11 +402,6 @@ class SrasFile: self.frame_count_mismatch = False self.n_frames_remainder = 0 - self.precomputed_freq_mhz: np.ndarray | None = None - self.precomputed_dc4_mv: np.ndarray | None = None - self.precomputed_dc3_mv: np.ndarray | None = None - self.precomputed_bg_sub: bool = False - angles = np.frombuffer(f.read(n_angles * 4), dtype=">f4").astype(np.float32) x_start = np.empty(n_angles, dtype=np.float64) @@ -409,6 +438,8 @@ class SrasFile: data_offset = f.tell() + self._data_offset = data_offset + # ---- Memory-map each angle's ragged waveform block ------------- # v6 gives each angle its own row/frame count, so waveform data is # no longer one uniform (n_angles, n_rows, ...) block — each angle's @@ -450,108 +481,157 @@ class SrasFile: self.n_rows = n_rows[:n_complete] self._y_pos_per_angle = y_pos_per_angle[:n_complete] + # Precomputed-image cache (v7 only). Ragged per-angle lists — unlike + # v5's dense (n_angles, n_rows, n_frames) arrays, v6/v7 geometry + # varies per angle, so each entry is its own (n_rows[a], n_frames[a]) + # array or None if that angle isn't cached yet. + self.precomputed_freq_mhz: list[np.ndarray | None] = [None] * n_complete + self.precomputed_dc4_mv: list[np.ndarray | None] = [None] * n_complete + self.precomputed_dc3_mv: list[np.ndarray | None] = [None] * n_complete + self.precomputed_bg_sub: bool = False + + if self.version == 7 and offset < file_size: + self._parse_cach_section(offset) + # ------------------------------------------------------------------ - # v5 writer + # v7 cache tail (CACH section: precomputed DC / FFT images) # ------------------------------------------------------------------ - def write_v5(self, dest_path: str, - freq_images: np.ndarray, - dc4_images: np.ndarray, - dc3_images: np.ndarray, - bg_sub_applied: bool, - progress_cb=None): - """Write a v5 .sras file to *dest_path*. + def _cache_tail_offset(self) -> int: + """Deterministic file offset where the CACH tail starts (or would + start), derived purely from the header + Per-Angle Geometry Table — + independent of whether a cache tail is actually present. Used by + both the parser and the in-place writer.""" + n_ch = self.n_channels + spf = self.samples_per_frame + bps = self.bytes_per_sample + waveform_bytes = int(sum( + int(self.n_rows[a]) * n_ch * int(self.n_frames[a]) * spf * bps + for a in range(self.n_angles) + )) + return self._data_offset + waveform_bytes - Copies the raw waveform bytes verbatim from the current file, - bumps the version byte to 5, patches n_frames_hdr to the actual - frame count, then appends the PREC section. + def _parse_cach_section(self, offset: int): + """Parse the v7 CACH tail that holds precomputed DC/FFT images.""" + with open(self.path, "rb") as f: + f.seek(offset) + header_raw = f.read(CACH_HDR_SIZE) + if len(header_raw) < CACH_HDR_SIZE: + return + magic, cach_version, block_flags = struct.unpack(CACH_HDR_FMT, header_raw) + if magic != CACH_MAGIC or cach_version != CACH_VERSION: + return - *freq_images* / *dc4_images* / *dc3_images*: - shape (n_angles, n_rows, n_frames) float32. + if block_flags & CACH_FLAG_DC: + sdcb_raw = f.read(SDCB_HDR_SIZE) + if len(sdcb_raw) < SDCB_HDR_SIZE: + return + sdcb_magic, _reserved, n_stored = struct.unpack(SDCB_HDR_FMT, sdcb_raw) + if sdcb_magic != SDCB_MAGIC: + return + for _ in range(n_stored): + (angle_idx,) = struct.unpack(">H", f.read(2)) + if angle_idx >= self.n_angles: + break + px = int(self.n_rows[angle_idx]) * int(self.n_frames[angle_idx]) + shape = (int(self.n_rows[angle_idx]), int(self.n_frames[angle_idx])) + self.precomputed_dc3_mv[angle_idx] = np.frombuffer( + f.read(px * 4), dtype=">f4").reshape(shape) + self.precomputed_dc4_mv[angle_idx] = np.frombuffer( + f.read(px * 4), dtype=">f4").reshape(shape) - *progress_cb*: optional callable(fraction: float) for UI updates. + if block_flags & CACH_FLAG_FFT: + sfft_raw = f.read(SFFT_HDR_SIZE) + if len(sfft_raw) < SFFT_HDR_SIZE: + return + sfft_magic, flags, n_stored = struct.unpack(SFFT_HDR_FMT, sfft_raw) + if sfft_magic != SFFT_MAGIC: + return + self.precomputed_bg_sub = bool(flags & SFFT_FLAG_BG_SUB) + for _ in range(n_stored): + (angle_idx,) = struct.unpack(">H", f.read(2)) + if angle_idx >= self.n_angles: + break + px = int(self.n_rows[angle_idx]) * int(self.n_frames[angle_idx]) + shape = (int(self.n_rows[angle_idx]), int(self.n_frames[angle_idx])) + self.precomputed_freq_mhz[angle_idx] = np.frombuffer( + f.read(px * 4), dtype=">f4").reshape(shape) - Only valid for v2–v5 source files, which have uniform per-angle - geometry. v6 files scan a different bounding box per angle and - cannot be losslessly represented in the flat v5 layout. + def write_v7_cache(self, *, + new_dc3_mv: list[np.ndarray | None] | None = None, + new_dc4_mv: list[np.ndarray | None] | None = None, + new_freq_mhz: list[np.ndarray | None] | None = None, + new_bg_sub: bool | None = None): + """Store computed DC and/or FFT images into this file's CACH tail, + in place, converting a v6 source to v7 (or updating an existing v7 + file). Only the block(s) passed in are recomputed; whichever block + isn't passed is carried forward unchanged from whatever this + ``SrasFile`` already has in memory (from parsing, or a prior write + in this same session) — its bytes are never re-read from disk. + + The waveform data itself is never touched: the cache tail always + starts at ``_cache_tail_offset()``, a fixed offset derived from the + header and geometry table alone. """ - if self.version == 6: - raise NotImplementedError( - "Pre-process to v5 is not supported for v6 source files " - "(per-angle geometry does not fit the flat v5 layout).") + if self.version not in (6, 7): + raise ValueError( + f"write_v7_cache only supports v6/v7 source files, got v{self.version}") - import shutil + final_dc3 = new_dc3_mv if new_dc3_mv is not None else self.precomputed_dc3_mv + final_dc4 = new_dc4_mv if new_dc4_mv is not None else self.precomputed_dc4_mv + final_freq = new_freq_mhz if new_freq_mhz is not None else self.precomputed_freq_mhz + final_bg_sub = new_bg_sub if new_bg_sub is not None else self.precomputed_bg_sub - dest = Path(dest_path) - src = self.path - n_rows0 = int(self.n_rows[0]) - n_frames0 = int(self.n_frames[0]) + dc_entries = [a for a in range(self.n_angles) if final_dc3[a] is not None] + fft_entries = [a for a in range(self.n_angles) if final_freq[a] is not None] - # --- Copy the source file verbatim, then patch the header ------- - shutil.copy2(str(src), str(dest)) + block_flags = 0 + if dc_entries: + block_flags |= CACH_FLAG_DC + if fft_entries: + block_flags |= CACH_FLAG_FFT - waveform_bytes = (self.n_angles * n_rows0 * self.n_channels - * n_frames0 * self.samples_per_frame - * self.bytes_per_sample) + payload = bytearray() + payload += struct.pack(CACH_HDR_FMT, CACH_MAGIC, CACH_VERSION, block_flags) - with open(str(dest), "r+b") as f: - # Patch version byte (offset 4 in the header struct) - f.seek(4) - f.write(struct.pack("B", 5)) + if dc_entries: + payload += struct.pack(SDCB_HDR_FMT, SDCB_MAGIC, 0, len(dc_entries)) + for a in dc_entries: + payload += struct.pack(">H", a) + payload += final_dc3[a].astype(">f4").tobytes() + payload += final_dc4[a].astype(">f4").tobytes() - # Patch n_frames_hdr (uint32, big-endian) with the actual count. - # Locate its offset: magic(4) + ver(1) + n_angles(2) + n_rows(2) = 9 - # then x_start(4)+x_delta(4)+vel(4)+freq(4) = 16, total = 25 - # then n_frames_hdr is at offset 25 as ">I" (4 bytes) - f.seek(25) - f.write(struct.pack(">I", n_frames0)) + if fft_entries: + fft_flags = SFFT_FLAG_BG_SUB if final_bg_sub else 0 + payload += struct.pack(SFFT_HDR_FMT, SFFT_MAGIC, fft_flags, len(fft_entries)) + for a in fft_entries: + payload += struct.pack(">H", a) + payload += final_freq[a].astype(">f4").tobytes() - # Truncate anything after the waveform data (e.g. old PREC tail) - # and seek to the append position. - waveform_end = self._data_offset_for_write() - f.seek(waveform_end + waveform_bytes) + cache_offset = self._cache_tail_offset() + with open(self.path, "r+b") as f: + f.seek(cache_offset) + f.write(payload) f.truncate() + f.flush() + os.fsync(f.fileno()) + # Version-byte flip last: if the process dies before this point, + # the file is still readable as plain v6 (v6 parsing only + # bounds-checks per-angle offset+nbytes <= file_size, it never + # asserts exactly how many bytes follow the last angle) — so an + # interrupted write can never corrupt the file, only leave + # harmless trailing bytes that the next successful write + # overwrites via this same deterministic cache_offset. + f.seek(4) + f.write(struct.pack("B", 7)) + f.flush() + os.fsync(f.fileno()) - # --- Write PREC section ------------------------------------- - n_stored = self.n_angles - flags = 0x01 if bg_sub_applied else 0x00 - f.write(b"PREC") - f.write(struct.pack("BB", 1, flags)) - f.write(struct.pack(">H", n_stored)) - - for aidx in range(n_stored): - if progress_cb is not None: - progress_cb(aidx / n_stored) - f.write(struct.pack(">H", aidx)) - f.write(freq_images[aidx].astype(">f4").tobytes()) - f.write(dc4_images[aidx].astype(">f4").tobytes()) - f.write(dc3_images[aidx].astype(">f4").tobytes()) - - if progress_cb is not None: - progress_cb(1.0) - - def _data_offset_for_write(self) -> int: - """Return the file offset where waveform data starts (used by write_v5).""" - # Re-derive the offset by walking the header fields, since we do not - # persist data_offset as an attribute from _parse. - with open(self.path, "rb") as f: - fields = struct.unpack(HDR_FMT, f.read(HDR_SIZE)) - ver = fields[1] - n_ch = fields[12] - n_rows0 = int(self.n_rows[0]) - - with open(self.path, "rb") as f: - f.seek(HDR_SIZE) - f.read(self.n_angles * 4) # angles - f.read(n_rows0 * 4) # y_pos - if ver >= 3: - for _ in range(n_ch): - (length,) = struct.unpack(">H", f.read(2)) - f.read(length) - if ver >= 4: - (n_bg,) = struct.unpack(">I", f.read(4)) - f.read(n_bg) - return f.tell() + self.version = 7 + self.precomputed_dc3_mv = final_dc3 + self.precomputed_dc4_mv = final_dc4 + self.precomputed_freq_mhz = final_freq + self.precomputed_bg_sub = final_bg_sub # ------------------------------------------------------------------ # Axes helpers @@ -621,6 +701,21 @@ def compute_dc_image(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray: return img +def _cached_dc_mv(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray | None: + """Return a v7-cached DC image (mV, already converted) for + (angle_idx, ch_idx) if the file's CACH section has it, else None. + + Only applies to the ragged v6/v7 list representation of + ``precomputed_dc3_mv``/``precomputed_dc4_mv`` — v5's dense PREC arrays + don't track per-angle validity (unstored angles are left at 0.0 rather + than a sentinel), so they're intentionally not consulted here. + """ + store = sras.precomputed_dc3_mv if ch_idx == CH3_IDX else sras.precomputed_dc4_mv + if isinstance(store, list) and angle_idx < len(store): + return store[angle_idx] + return None + + def compute_rf_image(sras: SrasFile, angle_idx: int, dc_threshold_mv: float, apply_bg_sub: bool = True, @@ -640,9 +735,10 @@ def compute_rf_image(sras: SrasFile, angle_idx: int, computed chunk-by-chunk internally (which does need every pixel's CH4 data, since that's what determines validity in the first place). - Fast path: if the file contains v5 precomputed peak-frequency images, - and zero-padding is not active, and the bg-sub flag matches, the stored - images are used directly — no FFT is run. + Fast path: if the file contains precomputed peak-frequency images for + this angle (v5's dense per-file PREC section, or v7's ragged per-angle + CACH section), and zero-padding is not active, and the bg-sub flag + matches, the stored image is used directly — no FFT is run. Otherwise, data is processed in row chunks sized to a fixed memory budget (see ``_chunk_rows_for``) to bound peak RAM regardless of scan @@ -652,15 +748,39 @@ def compute_rf_image(sras: SrasFile, angle_idx: int, n_frames = int(sras.n_frames[angle_idx]) data = sras.data[angle_idx] - # ---- Fast path: v5 precomputed images ---------------------------------- + # ---- Fast path: precomputed images (v5 PREC or v7 CACH) ---------------- + # v5 stores a dense (n_angles, n_rows, n_frames) ndarray (uniform + # geometry); v6/v7 store a ragged list, one entry per angle (None where + # that angle hasn't been cached), since geometry varies per angle. + freq_store = sras.precomputed_freq_mhz + if isinstance(freq_store, list): + angle_cached = angle_idx < len(freq_store) and freq_store[angle_idx] is not None + else: + angle_cached = freq_store is not None + can_use_precomputed = ( - sras.precomputed_freq_mhz is not None + angle_cached and n_fft is None # no custom zero-padding and sras.precomputed_bg_sub == (apply_bg_sub and sras.background is not None) ) if can_use_precomputed: - freq_img = sras.precomputed_freq_mhz[angle_idx].copy() - dc4_img = sras.precomputed_dc4_mv[angle_idx] + freq_img = freq_store[angle_idx].copy() + + # DC4 mask, in priority order: already-cached DC block, caller- + # supplied image, or a fresh (cheap — no FFT) recompute. + dc4_store = sras.precomputed_dc4_mv + if isinstance(dc4_store, list): + dc4_img = _cached_dc_mv(sras, angle_idx, CH4_IDX) + else: + dc4_img = dc4_store[angle_idx] if dc4_store is not None else None + if dc4_img is None: + if dc4_mv is not None: + dc4_img = dc4_mv + else: + dc4_img = adc_to_mv(compute_dc_image(sras, angle_idx, CH4_IDX), + sras.ch_ymult_mv[CH4_IDX], sras.ch_yoff_adc[CH4_IDX], + sras.ch_yzero_mv[CH4_IDX]) + freq_img[dc4_img < dc_threshold_mv] = 0.0 return freq_img @@ -897,7 +1017,8 @@ def _compute_angle_alignment(sras: SrasFile, ref_angle_idx: int, """Top-level alignment driver. Runs on a background thread (see AngleAlignmentWorker) — deliberately recomputes CH4 DC images from scratch via compute_dc_image rather than reading the GUI-thread - _dc_cache dict, mirroring PreprocessWorker's existing precedent.""" + _dc_cache dict, since background-thread workers must not touch + GUI-thread-owned caches (BatchCacheWorker follows the same rule).""" n = sras.n_angles # already the *complete*-angle count for aborted v6 scans dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx) @@ -1037,7 +1158,12 @@ class ComputeWorker(QObject): dc4_mv=self._dc4_mv) self.finished.emit(img) else: - # DC channels: convert ADC counts → mV + # DC channels: use the v7 cache if this angle is already + # stored (already mV), else compute fresh and convert. + cached = _cached_dc_mv(self._sras, self._angle, self._ch) + if cached is not None: + self.finished.emit(cached) + return adc_img = compute_dc_image(self._sras, self._angle, self._ch) img = adc_to_mv(adc_img, self._sras.ch_ymult_mv[self._ch], @@ -1076,6 +1202,11 @@ class DcPrecomputeWorker(QObject): for a in range(self._sras.n_angles): if self._stop: break + cached_dc3 = _cached_dc_mv(self._sras, a, CH3_IDX) + cached_dc4 = _cached_dc_mv(self._sras, a, CH4_IDX) + if cached_dc3 is not None and cached_dc4 is not None: + self.angle_done.emit(a, cached_dc3, cached_dc4) + continue dc3_mv = adc_to_mv( compute_dc_image(self._sras, a, CH3_IDX), self._sras.ch_ymult_mv[CH3_IDX], @@ -1092,69 +1223,93 @@ class DcPrecomputeWorker(QObject): self.error.emit(str(exc)) -class PreprocessWorker(QObject): - """Compute all-angle FFT and DC images and write a v5 file. +class BatchCacheWorker(QObject): + """Batch-computes and stores DC or FFT images into each of *paths*'s + v7 CACH tail, in place — converting v6 sources to v7 on first use, or + updating an existing v7 file's cache blocks without disturbing whatever + the other block already holds. - Emits ``progress(int)`` (0–100) as each angle completes and - ``finished(str)`` with an empty string on success or an error message - on failure. Only used for v2–v5 source files (uniform geometry). + *mode* is ``"dc"`` (CH3/CH4 mean images) or ``"fft"`` (CH1 peak-frequency + images, unmasked — masking is applied at display time, same as v5's + PREC convention). + + Emits ``progress(int)`` (0–100, weighted by total angle count across the + whole batch), ``file_done(str, str)`` (path, error message or "" on + success) after each file so one file's failure doesn't abort the batch, + and ``finished()`` once every file has been attempted. """ - progress = pyqtSignal(int) # 0–100 - finished = pyqtSignal(str) # empty = success, else error message + progress = pyqtSignal(int) + file_done = pyqtSignal(str, str) + finished = pyqtSignal() - def __init__(self, sras: SrasFile, dest_path: str, apply_bg_sub: bool): + def __init__(self, paths: list[str], mode: str, apply_bg_sub: bool): super().__init__() - self._sras = sras - self._dest_path = dest_path + self._paths = paths + self._mode = mode self._apply_bg_sub = apply_bg_sub def run(self): - try: - sras = self._sras - n = sras.n_angles - n_rows0 = int(sras.n_rows[0]) - n_frames0 = int(sras.n_frames[0]) - freq_images = np.empty((n, n_rows0, n_frames0), dtype=np.float32) - dc4_images = np.empty_like(freq_images) - dc3_images = np.empty_like(freq_images) + # Pass 1: quick open of each file just to weight progress by total + # angle count. Don't hold every file's memmap open at once — reopen + # fresh per file in pass 2 below. + total_angles = 0 + for path in self._paths: + try: + total_angles += SrasFile(path).n_angles + except Exception: + pass # unreadable files are reported properly in pass 2 + total_angles = max(total_angles, 1) - for aidx in range(n): - # Compute raw peak-frequency (no DC-threshold masking yet) - freq_images[aidx] = compute_rf_image( - sras, aidx, dc_threshold_mv=-1e9, # mask nothing - apply_bg_sub=self._apply_bg_sub) + done_angles = 0 + for path in self._paths: + try: + sras = SrasFile(path) + if sras.version not in (6, 7): + self.file_done.emit( + path, f"unsupported version {sras.version} — only " + "v6/v7 files can be batch-cached") + continue - # DC images (ADC counts → mV) - dc4_images[aidx] = adc_to_mv( - compute_dc_image(sras, aidx, CH4_IDX), - sras.ch_ymult_mv[CH4_IDX], - sras.ch_yoff_adc[CH4_IDX], - sras.ch_yzero_mv[CH4_IDX]) - dc3_images[aidx] = adc_to_mv( - compute_dc_image(sras, aidx, CH3_IDX), - sras.ch_ymult_mv[CH3_IDX], - sras.ch_yoff_adc[CH3_IDX], - sras.ch_yzero_mv[CH3_IDX]) + n = sras.n_angles + if self._mode == "dc": + new_dc3 = [None] * n + new_dc4 = [None] * n + for a in range(n): + new_dc3[a] = adc_to_mv( + compute_dc_image(sras, a, CH3_IDX), + sras.ch_ymult_mv[CH3_IDX], sras.ch_yoff_adc[CH3_IDX], + sras.ch_yzero_mv[CH3_IDX]) + new_dc4[a] = adc_to_mv( + compute_dc_image(sras, a, CH4_IDX), + sras.ch_ymult_mv[CH4_IDX], sras.ch_yoff_adc[CH4_IDX], + sras.ch_yzero_mv[CH4_IDX]) + done_angles += 1 + self.progress.emit(int(done_angles / total_angles * 100)) + sras.write_v7_cache(new_dc3_mv=new_dc3, new_dc4_mv=new_dc4) + else: + effective_bg = self._apply_bg_sub and sras.background is not None + new_freq = [None] * n + for a in range(n): + new_freq[a] = compute_rf_image( + sras, a, dc_threshold_mv=-1e9, # mask nothing + apply_bg_sub=effective_bg) + done_angles += 1 + self.progress.emit(int(done_angles / total_angles * 100)) + sras.write_v7_cache(new_freq_mhz=new_freq, new_bg_sub=effective_bg) - self.progress.emit(int((aidx + 1) / n * 90)) + self.file_done.emit(path, "") + except Exception as exc: + self.file_done.emit(path, str(exc)) - bg_sub_flag = self._apply_bg_sub and sras.background is not None - sras.write_v5( - self._dest_path, - freq_images, dc4_images, dc3_images, - bg_sub_applied=bg_sub_flag, - progress_cb=lambda frac: self.progress.emit(90 + int(frac * 10)), - ) - self.finished.emit("") - except Exception as exc: - self.finished.emit(str(exc)) + self.finished.emit() class AngleAlignmentWorker(QObject): """Computes rotation+translation alignment for every angle in *sras*, referenced to *ref_angle_idx*, from each angle's binarized CH4 mask. Rotation is analytic (from sras.angles_deg); only translation is found - by phase correlation. Mirrors PreprocessWorker's progress/finished shape. + by phase correlation. Uses the same progress(int)/finished(...) worker + shape as the other background-thread workers in this file. """ progress = pyqtSignal(int) # 0–100 finished = pyqtSignal(object, str) # AlignmentResult|None, error msg ("" = success) @@ -1794,7 +1949,11 @@ class SrasViewerWindow(QMainWindow): self._fft_pad_factor: int = 1 # 1 = no padding self._pending_fft_pad_factor: int = 1 - self._preprocess_thread: QThread | None = None + # Convert menu: batch DC/FFT compute-and-store (v6 -> v7) + self._batch_thread: QThread | None = None + self._batch_worker: BatchCacheWorker | None = None + self._batch_errors: list[str] = [] + self._batch_progress_dlg: QProgressDialog | None = None # Display-only settings (colormap, grating) never trigger a # recompute — they're applied to cached data on redraw. DC images @@ -2110,15 +2269,6 @@ class SrasViewerWindow(QMainWindow): fft_act.triggered.connect(self._on_fft_options) fft_menu.addAction(fft_act) - fft_menu.addSeparator() - self._preprocess_act = QAction("&Pre-process and Save as v5…", self) - self._preprocess_act.setStatusTip( - "Compute FFT and DC images for all angles and save to a v5 file " - "for instant re-opening (no FFT on load). Not available for v6 files.") - self._preprocess_act.setEnabled(False) - self._preprocess_act.triggered.connect(self._on_preprocess) - fft_menu.addAction(self._preprocess_act) - fusion_menu = menubar.addMenu("&Fusion") self._alignment_act = QAction("Angle &Alignment", self) self._alignment_act.setStatusTip( @@ -2128,6 +2278,21 @@ class SrasViewerWindow(QMainWindow): self._alignment_act.triggered.connect(self._on_angle_alignment) fusion_menu.addAction(self._alignment_act) + convert_menu = menubar.addMenu("&Convert") + self._batch_dc_act = QAction("Batch Compute DC and &Store…", self) + self._batch_dc_act.setStatusTip( + "Select .sras files and compute+store DC images (CH3/CH4 mean) " + "for every angle, converting v6 files to v7 in place.") + self._batch_dc_act.triggered.connect(lambda: self._on_batch_compute("dc")) + convert_menu.addAction(self._batch_dc_act) + + self._batch_fft_act = QAction("Batch Compute FFT and Sto&re…", self) + self._batch_fft_act.setStatusTip( + "Select .sras files and compute+store FFT peak-frequency images " + "for every angle, converting v6 files to v7 in place.") + self._batch_fft_act.triggered.connect(lambda: self._on_batch_compute("fft")) + convert_menu.addAction(self._batch_fft_act) + # ------------------------------------------------------------------ # Drag-and-drop # ------------------------------------------------------------------ @@ -2265,11 +2430,22 @@ class SrasViewerWindow(QMainWindow): ) if s.background is not None: notes.append(f"Background waveform: {len(s.background)} samples") - if s.precomputed_freq_mhz is not None: + if isinstance(s.precomputed_freq_mhz, np.ndarray): bg_note = " (bg-sub)" if s.precomputed_bg_sub else " (no bg-sub)" notes.append(f"v5: precomputed images present{bg_note} — display is instant") - if s.version == 6: - notes.append("v6 format: rows / frames / x_start are per-angle") + if s.version in (6, 7): + notes.append("v6/v7 format: rows / frames / x_start are per-angle") + if s.version == 7: + n_dc = sum(1 for x in s.precomputed_dc4_mv if x is not None) + n_fft = sum(1 for x in s.precomputed_freq_mhz if x is not None) + if n_dc or n_fft: + bg_note = " (bg-sub)" if s.precomputed_bg_sub else " (no bg-sub)" + notes.append( + f"v7: cached DC ({n_dc}/{s.n_angles} angles), " + f"FFT ({n_fft}/{s.n_angles} angles{bg_note if n_fft else ''}) " + "— display is instant for cached angles") + else: + notes.append("v7 format: no cache blocks stored yet") self.lbl_frame_warn.setText("\n".join(notes)) # ------------------------------------------------------------------ @@ -2299,10 +2475,11 @@ class SrasViewerWindow(QMainWindow): self.btn_export_csv.setEnabled(is_ch1 and self._current_image is not None) # ROI: always usable once a file is loaded (independent of channel) self.btn_draw_roi.setEnabled(enabled and s is not None) - # Pre-process: available for v2–v5 files when loaded and not already running - can_preprocess = (enabled and s is not None and s.version != 6 - and self._preprocess_thread is None) - self._preprocess_act.setEnabled(can_preprocess) + # Batch Convert actions: pick their own files, independent of + # whatever's currently open — only gated on no batch already running + can_batch = self._batch_thread is None + self._batch_dc_act.setEnabled(can_batch) + self._batch_fft_act.setEnabled(can_batch) # Angle alignment: needs >1 angle and no alignment already running can_align = (enabled and s is not None and s.n_angles > 1 and self._alignment_thread is None) @@ -2897,75 +3074,90 @@ class SrasViewerWindow(QMainWindow): self._progress_dlg = None # ------------------------------------------------------------------ - # Pre-process → save v5 + # Convert menu: batch DC/FFT compute-and-store (v6 -> v7) # ------------------------------------------------------------------ - def _on_preprocess(self): - if self._sras is None: + def _on_batch_compute(self, mode: str): + if self._batch_thread is not None: return - if self._preprocess_thread is not None: - return - s = self._sras - if s.version == 6: - self.statusBar().showMessage( - "Pre-process to v5 is not supported for v6 files (per-angle geometry).") - return - default_name = s.path.stem + "_v5" + s.path.suffix - dest, _ = QFileDialog.getSaveFileName( - self, "Save Pre-processed v5 File", - str(s.path.parent / default_name), + label = "DC" if mode == "dc" else "FFT" + paths, _ = QFileDialog.getOpenFileNames( + self, f"Select .sras files to batch-compute {label}", "", "SRAS files (*.sras);;All files (*)", ) - if not dest: + if not paths: return - if self._preprocess_thread is not None: - return # a second trigger snuck in while the save dialog was open + if self._batch_thread is not None: + return # a second trigger snuck in while the file dialog was open - apply_bg = self.chk_bg_sub.isChecked() and s.background is not None + apply_bg = self.chk_bg_sub.isChecked() + self._batch_errors = [] - n_total = s.n_angles - n_px = int(s.n_rows[0]) * int(s.n_frames[0]) - approx_mb = n_total * n_px * 3 * 4 / 1e6 + # Claim self._batch_thread before anything below that can pump the + # Qt event loop — see the comment in _start_compute for why. + self._batch_worker = BatchCacheWorker(paths, mode, apply_bg) + self._batch_thread = QThread() + self._batch_worker.moveToThread(self._batch_thread) + self._batch_thread.started.connect(self._batch_worker.run) + self._batch_worker.progress.connect(self._on_batch_progress) + self._batch_worker.file_done.connect(self._on_batch_file_done) + self._batch_worker.finished.connect(lambda paths=paths: self._on_batch_finished(paths)) + self._batch_worker.finished.connect(self._batch_thread.quit) + self._batch_thread.finished.connect(self._on_batch_thread_finished) - # Claim self._preprocess_thread before any call below that can pump - # the Qt event loop — see the comment in _start_compute for why. - self._preprocess_worker = PreprocessWorker(s, dest, apply_bg) - self._preprocess_thread = QThread() - self._preprocess_worker.moveToThread(self._preprocess_thread) - self._preprocess_thread.started.connect(self._preprocess_worker.run) - self._preprocess_worker.progress.connect(self._on_preprocess_progress) - self._preprocess_worker.finished.connect(self._on_preprocess_done) - self._preprocess_worker.finished.connect(self._preprocess_thread.quit) - self._preprocess_thread.finished.connect(self._on_preprocess_thread_finished) + self._batch_dc_act.setEnabled(False) + self._batch_fft_act.setEnabled(False) - self._preprocess_act.setEnabled(False) - self._show_progress( - f"Pre-processing {n_total} angle(s) " - f"({int(s.n_rows[0])}×{int(s.n_frames[0])} px each, ~{approx_mb:.0f} MB output)…" - ) + self._batch_progress_dlg = QProgressDialog( + f"Batch computing {label} for {len(paths)} file(s)…", "", 0, 100, self) + self._batch_progress_dlg.setWindowTitle("Please wait…") + self._batch_progress_dlg.setCancelButton(None) + self._batch_progress_dlg.setWindowModality(Qt.WindowModality.WindowModal) + self._batch_progress_dlg.setMinimumDuration(300) + self._batch_progress_dlg.show() - self._preprocess_thread.start() + self._batch_thread.start() - def _on_preprocess_progress(self, pct: int): - if self._progress_dlg is not None: - self._progress_dlg.setValue(pct) + def _on_batch_progress(self, pct: int): + if self._batch_progress_dlg is not None: + self._batch_progress_dlg.setValue(pct) - def _on_preprocess_done(self, error_msg: str): - self._close_progress() - if error_msg: - self.statusBar().showMessage(f"Pre-process failed: {error_msg}") + def _on_batch_file_done(self, path: str, err: str): + if err: + self._batch_errors.append(f"{Path(path).name} — {err}") + if self._batch_progress_dlg is not None: + self._batch_progress_dlg.setLabelText(f"Processed {Path(path).name}…") + + def _on_batch_finished(self, paths: list[str]): + if self._batch_progress_dlg is not None: + self._batch_progress_dlg.close() + self._batch_progress_dlg = None + + n_total = len(paths) + n_failed = len(self._batch_errors) + n_ok = n_total - n_failed + if n_failed: + summary = (f"Batch store: {n_ok}/{n_total} file(s) updated, " + f"{n_failed} failed: {'; '.join(self._batch_errors)}") else: - self.statusBar().showMessage("v5 file written — re-open it for instant display.") + summary = f"Batch store: {n_ok}/{n_total} file(s) updated." + self.statusBar().showMessage(summary) + self._batch_errors = [] - def _on_preprocess_thread_finished(self): + # If the currently-open file was in this batch, reload it so the + # GUI picks up the newly-written v7 cache instead of stale state. + if self._sras is not None and str(self._sras.path) in paths: + self._load_file(str(self._sras.path)) + + def _on_batch_thread_finished(self): # See the comment in _on_compute_thread_finished: wait() before # releasing our reference to avoid destroying a QThread whose OS # thread hasn't fully joined yet. - if self._preprocess_thread is not None: - self._preprocess_thread.wait() - self._preprocess_thread = None - self._preprocess_act.setEnabled( - self._sras is not None and self._sras.version != 6) + if self._batch_thread is not None: + self._batch_thread.wait() + self._batch_thread = None + self._batch_dc_act.setEnabled(True) + self._batch_fft_act.setEnabled(True) # ------------------------------------------------------------------ # Fusion: angle alignment @@ -3066,7 +3258,7 @@ class SrasViewerWindow(QMainWindow): def closeEvent(self, event): if self._dc_precompute_worker is not None: self._dc_precompute_worker.stop() - for attr in ("_load_thread", "_compute_thread", "_preprocess_thread", + for attr in ("_load_thread", "_compute_thread", "_batch_thread", "_dc_precompute_thread", "_alignment_thread"): t = getattr(self, attr, None) if t is not None: From 95c273dce542021eab8107b476547c00f44beaf2 Mon Sep 17 00:00:00 2001 From: Thomas Ales Date: Thu, 30 Jul 2026 23:25:45 -0500 Subject: [PATCH 02/10] Add test tooling: synthetic .sras generator and equivalence harness tools/make_test_sras.py writes small v6 files with per-angle-varying geometry, distinct per-channel calibration, and deterministic waveform content (predictable FFT peak per pixel, predictable DC mean per pixel). tools/check_equivalence.py hashes DC/FFT/alignment outputs across channels, angles, bg-sub, pad factors, and thresholds. It imports from either the monolith or the split modules, so the same script captures both sides of a refactor. Hashes canonicalise to native float64 so a dtype or byte-order change that preserves values is not a false failure. Co-Authored-By: Claude Opus 5 --- tools/__init__.py | 0 tools/check_equivalence.py | 197 +++++++++++++++++++++++++++++++++++++ tools/make_test_sras.py | 142 ++++++++++++++++++++++++++ 3 files changed, 339 insertions(+) create mode 100644 tools/__init__.py create mode 100644 tools/check_equivalence.py create mode 100644 tools/make_test_sras.py diff --git a/tools/__init__.py b/tools/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tools/check_equivalence.py b/tools/check_equivalence.py new file mode 100644 index 0000000..fa11968 --- /dev/null +++ b/tools/check_equivalence.py @@ -0,0 +1,197 @@ +#!/usr/bin/env python3 +"""Golden-output equivalence harness for the sras-viewer refactor. + +Computes a battery of DC / FFT / alignment outputs and prints a stable hash +for each. Run it on the pre-refactor commit to capture a baseline, then again +after the refactor and diff the two reports — every line must match. + +Imports work against both the pre-refactor monolith (`sras_viewer`) and the +post-refactor split (`sras_format` + `sras_compute`), so the *same* script +produces both sides of the comparison. + +Hashes canonicalise to native little-endian float64 before hashing, so a +deliberate dtype/byte-order change that preserves values does not show up as +a false mismatch. + +Usage: + python tools/check_equivalence.py --out baseline.txt + python tools/check_equivalence.py --out after.txt --real /path/to/big.sras + diff baseline.txt after.txt +""" + +import argparse +import copy +import hashlib +import sys +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +# --- Import shim: split modules if present, else the monolith -------------- +try: + from sras_format import SrasFile, CH1_IDX, CH3_IDX, CH4_IDX, adc_to_mv + import sras_compute as C + _LAYOUT = "split" +except ImportError: + import sras_viewer as _V + from sras_viewer import SrasFile, CH1_IDX, CH3_IDX, CH4_IDX, adc_to_mv + C = _V + _LAYOUT = "monolith" + +compute_dc_image = C.compute_dc_image +compute_rf_image = C.compute_rf_image +compute_alignment = C._compute_angle_alignment +apply_alignment = C.apply_alignment + +import tools.make_test_sras as gen # noqa: E402 + + +def h(arr) -> str: + """Stable hash of an array's *values*, independent of dtype/byte order.""" + a = np.ascontiguousarray(np.asarray(arr, dtype=np.float64)) + return hashlib.sha256(a.tobytes()).hexdigest()[:16] + + +def row_slice(sras: SrasFile, angle_idx: int, n_rows: int) -> SrasFile: + """A shallow view of *sras* restricted to the first *n_rows* rows of + *angle_idx*, so the huge real file can be exercised in seconds.""" + view = copy.copy(sras) + n = min(int(n_rows), int(sras.n_rows[angle_idx])) + view.n_rows = np.array(sras.n_rows, copy=True) + view.n_rows[angle_idx] = n + view.data = list(sras.data) + view.data[angle_idx] = sras.data[angle_idx][:n] + view._y_pos_per_angle = list(sras._y_pos_per_angle) + view._y_pos_per_angle[angle_idx] = sras._y_pos_per_angle[angle_idx][:n] + return view + + +def report(lines: list[str], label: str, value: str): + lines.append(f"{label:<58} {value}") + + +def check_file(path: Path, lines: list[str], tag: str, + angles: list[int], n_rows: int | None): + sras = SrasFile(str(path)) + report(lines, f"[{tag}] version/n_angles", + f"v{sras.version} n={sras.n_angles}") + report(lines, f"[{tag}] geometry", + f"rows={list(map(int, sras.n_rows))} frames={list(map(int, sras.n_frames))}") + report(lines, f"[{tag}] calibration", + " ".join(f"{m:.6g}/{o:.6g}/{z:.6g}" for m, o, z in + zip(sras.ch_ymult_mv, sras.ch_yoff_adc, sras.ch_yzero_mv))) + report(lines, f"[{tag}] angles_deg", h(sras.angles_deg)) + if sras.background is not None: + report(lines, f"[{tag}] background", h(sras.background)) + + for a in angles: + if a >= sras.n_angles: + continue + s = row_slice(sras, a, n_rows) if n_rows else sras + + report(lines, f"[{tag}] x_axis_mm(a={a})", h(s.x_axis_mm(a))) + report(lines, f"[{tag}] y_positions_mm(a={a})", h(s.y_positions_mm(a))) + + for ch, name in ((CH3_IDX, "CH3"), (CH4_IDX, "CH4")): + raw = compute_dc_image(s, a, ch) + report(lines, f"[{tag}] dc_adc(a={a},{name})", h(raw)) + mv = adc_to_mv(raw, s.ch_ymult_mv[ch], s.ch_yoff_adc[ch], + s.ch_yzero_mv[ch]) + report(lines, f"[{tag}] dc_mv(a={a},{name})", h(mv)) + + dc4 = adc_to_mv(compute_dc_image(s, a, CH4_IDX), + s.ch_ymult_mv[CH4_IDX], s.ch_yoff_adc[CH4_IDX], + s.ch_yzero_mv[CH4_IDX]) + thresholds = [-1e9, float(np.median(dc4))] + + for bg in (False, True): + if bg and s.background is None: + continue + for pad in (1, 2): + n_fft = s.samples_per_frame * pad if pad > 1 else None + for ti, thr in enumerate(thresholds): + img = compute_rf_image(s, a, dc_threshold_mv=thr, + apply_bg_sub=bg, n_fft=n_fft) + report(lines, + f"[{tag}] rf(a={a},bg={int(bg)},pad={pad},thr{ti})", + h(img)) + # dc4_mv passthrough must give the identical result + img2 = compute_rf_image(s, a, dc_threshold_mv=thr, + apply_bg_sub=bg, n_fft=n_fft, + dc4_mv=dc4) + same = "SAME" if h(img) == h(img2) else "DIFFER" + report(lines, + f"[{tag}] rf-dc4arg(a={a},bg={int(bg)},pad={pad},thr{ti})", + same) + + +def check_alignment(path: Path, lines: list[str], tag: str): + sras = SrasFile(str(path)) + if sras.n_angles < 2: + return + dc4 = adc_to_mv(compute_dc_image(sras, 0, CH4_IDX), + sras.ch_ymult_mv[CH4_IDX], sras.ch_yoff_adc[CH4_IDX], + sras.ch_yzero_mv[CH4_IDX]) + thr = float(np.median(dc4)) + res = compute_alignment(sras, 0, thr) + report(lines, f"[{tag}] align canvas_shape", str(res.canvas_shape)) + report(lines, f"[{tag}] align canvas_origin", + f"{res.canvas_origin_mm[0]:.9g},{res.canvas_origin_mm[1]:.9g}") + report(lines, f"[{tag}] align canvas_pitch", + f"{res.canvas_dx_mm:.9g},{res.canvas_dy_mm:.9g}") + for a in sorted(res.per_angle): + t = res.per_angle[a] + report(lines, f"[{tag}] align shift_mm(a={a})", + f"{t.shift_mm[0]:.9g},{t.shift_mm[1]:.9g}") + report(lines, f"[{tag}] align rot(a={a})", f"{t.rotation_deg:.9g}") + report(lines, f"[{tag}] align matrix(a={a})", h(t.matrix)) + report(lines, f"[{tag}] align offset(a={a})", h(t.offset)) + img = compute_dc_image(sras, a, CH4_IDX) + report(lines, f"[{tag}] align resampled(a={a})", + h(apply_alignment(res, a, img))) + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--out", required=True, help="report file to write") + p.add_argument("--real", help="optional path to a real .sras file") + p.add_argument("--real-rows", type=int, default=2, + help="rows per angle to sample from the real file") + p.add_argument("--real-angles", type=int, default=2, + help="how many angles to sample from the real file") + p.add_argument("--scratch", default=".", + help="directory for generated synthetic files") + args = p.parse_args() + + lines = [f"# layout: {_LAYOUT}", f"# numpy: {np.__version__}"] + + scratch = Path(args.scratch) + synth = scratch / "equiv_synth.sras" + gen.write(synth, n_angles=4, seed=0, samples_per_frame=64) + check_file(synth, lines, "synth", angles=[0, 1, 2, 3], n_rows=None) + check_alignment(synth, lines, "synth") + + # A second synthetic with an odd sample count, to catch off-by-one in + # rfft bin handling and chunk-boundary arithmetic. + synth_odd = scratch / "equiv_synth_odd.sras" + gen.write(synth_odd, n_angles=2, seed=7, samples_per_frame=37) + check_file(synth_odd, lines, "odd", angles=[0, 1], n_rows=None) + + if args.real: + real = Path(args.real) + if real.exists(): + check_file(real, lines, "real", + angles=list(range(args.real_angles)), + n_rows=args.real_rows) + else: + lines.append(f"# real file not found: {real}") + + Path(args.out).write_text("\n".join(lines) + "\n") + print("\n".join(lines)) + print(f"\nWrote {args.out} ({len(lines)} lines)") + + +if __name__ == "__main__": + main() diff --git a/tools/make_test_sras.py b/tools/make_test_sras.py new file mode 100644 index 0000000..5ea0e00 --- /dev/null +++ b/tools/make_test_sras.py @@ -0,0 +1,142 @@ +#!/usr/bin/env python3 +"""Generate small synthetic .sras files for testing. + +Writes v6 files (per-angle geometry, ragged waveform blocks) matching +scan_format.md, with deterministic pseudo-random waveform content so a test +can compute expected DC/FFT images independently of the reader under test. + +Usage: + python tools/make_test_sras.py out.sras [--angles 3] [--seed 0] +""" + +import argparse +import struct +from pathlib import Path + +import numpy as np + +HDR_FMT_V6 = ">4sBHfffffffIdBB" +GEO_FMT_V6 = ">ffIH" + +# Per-angle (n_rows, n_frames) — deliberately different per angle so ragged +# geometry handling is actually exercised. +_GEOMETRY = [(5, 7), (4, 11), (6, 9), (3, 13), (7, 6)] + +_SAMPLE_RATE_HZ = 6.25e9 +_VELOCITY_MM_S = 20.0 +_LASER_FREQ_HZ = 1000.0 +_ROW_SPACING_MM = 0.05 + + +def _preamble(ymult_v: float, yoff_adc: float, yzero_v: float) -> bytes: + """A Tektronix WFMOutpre string in verbose (keyword) form — the reader + pulls YMULT/YOFF/YZERO out of it by name, so the keywords must be + present literally. YMULT/YZERO are in volts, as the scope reports them.""" + return ( + ":WFMOUTPRE:BYT_NR 1;BIT_NR 8;ENCDG BIN;BN_FMT RI;BYT_OR MSB;" + 'WFID "Ch1, DC coupling";NR_PT 2500;PT_FMT Y;' + "XINCR 1.6000E-10;XZERO 0.0E0;XUNIT \"s\";" + f"YMULT {ymult_v:.6E};YOFF {yoff_adc:.6E};YZERO {yzero_v:.6E};" + 'YUNIT "V"' + ).encode("utf-8") + + +def build(n_angles: int, seed: int, samples_per_frame: int) -> tuple[bytes, dict]: + rng = np.random.default_rng(seed) + geom = [_GEOMETRY[a % len(_GEOMETRY)] for a in range(n_angles)] + n_ch = 3 + bps = 1 + + angles_deg = np.linspace(0.0, 60.0, n_angles, dtype=np.float32) + # Distinct calibration per channel so a swapped-channel bug is visible. + cal = [ + (1.5625e-3, -87.04, 0.0), + (2.0000e-3, -60.00, 1.0e-3), + (2.5000e-3, -40.00, -2.0e-3), + ] + + out = bytearray() + out += struct.pack( + HDR_FMT_V6, b"SRAS", 6, n_angles, + 0.0, 0.0, 1.0, 1.0, _ROW_SPACING_MM, + _VELOCITY_MM_S, _LASER_FREQ_HZ, + samples_per_frame, _SAMPLE_RATE_HZ, bps, n_ch, + ) + out += angles_deg.astype(">f4").tobytes() + + x_starts = [] + for a, (n_rows, n_frames) in enumerate(geom): + x_start = -0.5 + 0.1 * a + x_starts.append(x_start) + out += struct.pack(GEO_FMT_V6, x_start, 1.0, n_frames, n_rows) + + y_positions = [] + for a, (n_rows, _) in enumerate(geom): + y = (0.2 * a + np.arange(n_rows) * _ROW_SPACING_MM).astype(np.float32) + y_positions.append(y) + out += y.astype(">f4").tobytes() + + for ymult_v, yoff, yzero_v in cal: + p = _preamble(ymult_v, yoff, yzero_v) + out += struct.pack(">H", len(p)) + p + + background = rng.integers(-8, 9, size=samples_per_frame, dtype=np.int8) + out += struct.pack(">I", samples_per_frame) + background.tobytes() + + # Waveform data. CH1 gets a sinusoid at a per-pixel frequency so the FFT + # peak is predictable; CH3/CH4 get per-pixel DC levels so the mean is too. + t = np.arange(samples_per_frame) + waveforms = [] + for a, (n_rows, n_frames) in enumerate(geom): + block = np.empty((n_rows, n_ch, n_frames, samples_per_frame), dtype=np.int8) + for r in range(n_rows): + for f in range(n_frames): + bin_idx = 3 + ((a + r + f) % 17) + phase = 2 * np.pi * bin_idx * t / samples_per_frame + block[r, 0, f] = np.clip( + np.round(60 * np.sin(phase)), -128, 127).astype(np.int8) + block[r, 1, f] = np.int8((a * 7 + r * 3 + f) % 100 - 50) + block[r, 2, f] = np.int8((a * 5 + r * 11 + f * 2) % 120 - 60) + waveforms.append(block) + out += block.tobytes() + + meta = { + "n_angles": n_angles, + "geometry": geom, + "angles_deg": angles_deg, + "x_starts": x_starts, + "y_positions": y_positions, + "cal": cal, + "background": background, + "waveforms": waveforms, + "samples_per_frame": samples_per_frame, + "sample_rate_hz": _SAMPLE_RATE_HZ, + } + return bytes(out), meta + + +def write(path: Path, n_angles: int = 3, seed: int = 0, + samples_per_frame: int = 64) -> dict: + payload, meta = build(n_angles, seed, samples_per_frame) + path.write_bytes(payload) + return meta + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("output") + p.add_argument("--angles", type=int, default=3) + p.add_argument("--seed", type=int, default=0) + p.add_argument("--spf", type=int, default=64, help="samples per frame") + args = p.parse_args() + + out = Path(args.output) + meta = write(out, args.angles, args.seed, args.spf) + print(f"Wrote {out} ({out.stat().st_size:,} bytes)") + print(f" angles : {meta['n_angles']}") + print(f" geometry : {meta['geometry']}") + print(f" spf : {meta['samples_per_frame']}") + + +if __name__ == "__main__": + main() From 55a4c5e42f2574a80a0c6fb6c1b3f055091c07ab Mon Sep 17 00:00:00 2001 From: Thomas Ales Date: Thu, 30 Jul 2026 23:38:08 -0500 Subject: [PATCH 03/10] Split into four modules, deduplicate, and parallelize the compute paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Structure sras_format.py parsing/writing, calibration, axes (numpy + struct) sras_compute.py DC/FFT images, alignment, batch cache (+ scipy) sras_workers.py Qt background workers sras_viewer.py ROI, canvases, dialog, main window format+compute import in 0.59s with no Qt or matplotlib (vs 2.96s for the full app), which is what makes a spawn-based process pool worth using. Deduplication - SrasFile.cal() and compute.dc_image_mv() replace the ADC->mV calibration incantation that appeared at ten call sites. - _read_preambles/_read_background/_set_calibration are shared by the legacy and v6 parsers instead of duplicated. - v5 PREC images are normalised into the same ragged per-angle list v6/v7 uses, removing four isinstance() branches and shortening compute_rf_image's fast path. - _run_worker replaces five copies of the QThread setup and five near-identical teardown methods; the two lifetime hazards they guarded against are now documented once, authoritatively. - _on_channel_changed defers to _update_controls_enabled rather than re-deriving the same six enable rules. - _corners_in_ref_frame, _CHANNEL_DISPLAY dict, unified progress-dialog helper, dead _on_roi_angle_edited stub removed. - sras_average.py builds on SrasFile instead of carrying a second copy of the header format, and streams per angle instead of loading the whole file (was a ~3x file-size peak). Executable lines: 2390 -> 2254, despite adding all of the below. Parallelism - compute_rf_image/compute_dc_image map row chunks over a thread pool. Worker count is derived from the memory budget rather than the core count: on a large scan chunk_rows is already floored at 1 row (~75 MB of float32 at 7507x2500), so only the worker count can bound peak RAM. - DcPrecomputeWorker computes angles on a pool, emitting each result from its own QThread as it lands. - BatchCacheWorker runs one process per file via cache_file(); only paths and scalars cross the boundary. Per-process thread counts are divided so the two levels don't oversubscribe. Drops the old pass 1, which fully parsed every file just to weight a progress bar. - dc_threshold_mv=None means "no mask" and skips the CH4 read entirely — the batch FFT job previously read all of CH4 to compare against a threshold of -1e9. - Alignment mask and correlation stages map over angles; fft2/ifft2 use workers=-1. Fixes found on the way - A partially-cached v5 file showed an all-zero image for uncached angles: the fast-path check was file-wide, not per-angle. - v7 cached images were read-only big-endian views; now native float32. Verified: tools/check_equivalence.py produces byte-identical hashes for 167 outputs (DC/FFT across channels, angles, bg-sub, pad factors and thresholds, plus the full alignment result) against ed0eba4, on both synthetic files and a real 496 GB 17-angle v6 scan. Co-Authored-By: Claude Opus 5 --- sras_average.py | 231 ++-- sras_compute.py | 577 ++++++++++ sras_format.py | 593 ++++++++++ sras_viewer.py | 2868 ++++++++++++----------------------------------- sras_workers.py | 224 ++++ 5 files changed, 2218 insertions(+), 2275 deletions(-) create mode 100644 sras_compute.py create mode 100644 sras_format.py create mode 100644 sras_workers.py diff --git a/sras_average.py b/sras_average.py index 0ddc876..d997a15 100644 --- a/sras_average.py +++ b/sras_average.py @@ -15,25 +15,24 @@ Options: Default: include a partial average for the last group. """ -import sys -import struct import argparse -import numpy as np +import shutil +import struct +import sys from pathlib import Path -# --------------------------------------------------------------------------- -# Header format — must match sras_viewer.py exactly -# --------------------------------------------------------------------------- +import numpy as np -HDR_FMT = ">4sBHHffffIIdBB" -HDR_SIZE = struct.calcsize(HDR_FMT) # 43 bytes +from sras_format import HDR_FMT, HDR_SIZE, SrasFile + +_SUPPORTED = (2, 3, 4) def parse_args(): p = argparse.ArgumentParser( description="Average every N waveforms in a .sras file and write a new file." ) - p.add_argument("input", help="Input .sras file") + p.add_argument("input", help="Input .sras file") p.add_argument("output", help="Output .sras file") p.add_argument("--n", type=int, required=True, metavar="N", help="Number of consecutive frames to average into one") @@ -42,143 +41,68 @@ def parse_args(): return p.parse_args() -def read_sras(path: Path): - """Read all sections of a .sras file and return them as a dict.""" - with open(path, "rb") as f: - header_bytes = f.read(HDR_SIZE) - fields = struct.unpack(HDR_FMT, header_bytes) - - (magic, ver, n_angles, n_rows, x_start, x_delta, vel, freq, - n_frames_hdr, spf, sr, bps, n_ch) = fields - - if magic != b"SRAS": - raise ValueError(f"Not a .sras file (bad magic: {magic!r})") - if ver not in (2, 3, 4): - raise ValueError(f"Unsupported .sras version: {ver}") - - with open(path, "rb") as f: +def read_header_sections(sras: SrasFile) -> bytes: + """The raw bytes between the header and the waveform data (angle table, + row table, preambles, background), copied through verbatim so nothing is + lost in a re-encode.""" + with open(sras.path, "rb") as f: f.seek(HDR_SIZE) - - angles = f.read(n_angles * 4) # big-endian float32 array, raw bytes - y_pos = f.read(n_rows * 4) # big-endian float32 array, raw bytes - - preambles = [] # list of raw bytes (length-prefixed strings) - if ver >= 3: - for _ in range(n_ch): - (length,) = struct.unpack(">H", f.read(2)) - preambles.append(f.read(length)) - - background = b"" # raw bytes for the v4 background block - if ver >= 4: - (n_bg,) = struct.unpack(">I", f.read(4)) - background = f.read(n_bg) - - raw = f.read() # all waveform data - - # ----------------------------------------------------------------------- - # Determine actual frame count from file size (the header value can be - # wrong — the viewer does the same correction) - # ----------------------------------------------------------------------- - total_samples = len(raw) // bps - samples_per_pixel = n_ch * spf # samples in one (angle, row, frame) cell - samples_per_full = n_angles * n_rows * samples_per_pixel - - actual_n_frames = total_samples // (n_angles * n_rows * samples_per_pixel) - good_bytes = actual_n_frames * samples_per_full * bps - - # Decode waveform data - dtype = np.int8 if bps == 1 else ">i2" - data = np.frombuffer(raw[:good_bytes], dtype=dtype) - data = data.reshape(n_angles, n_rows, n_ch, actual_n_frames, spf) - # Work in int16 (safe intermediate for both int8 and int16 inputs) - data = data.astype(np.int16) - - return { - "ver": ver, "n_angles": n_angles, "n_rows": n_rows, - "x_start": x_start, "x_delta": x_delta, "vel": vel, "freq": freq, - "n_frames_hdr": n_frames_hdr, "spf": spf, "sr": sr, - "bps": bps, "n_ch": n_ch, - "angles_raw": angles, "y_pos_raw": y_pos, - "preambles": preambles, "background": background, - "data": data, # shape: (n_angles, n_rows, n_ch, n_frames, spf), int16 - } + return f.read(sras._data_offset - HDR_SIZE) -def average_frames(data: np.ndarray, n: int, discard_remainder: bool) -> np.ndarray: - """Average every N frames along axis 3. +def average_rows(block: np.ndarray, n: int, discard_remainder: bool) -> np.ndarray: + """Average every N frames of one angle's (n_rows, n_ch, n_frames, spf) + block. Returns int16 of shape (n_rows, n_ch, n_out, spf). - data shape: (n_angles, n_rows, n_ch, n_frames, spf) - Returns array of shape (n_angles, n_rows, n_ch, n_out, spf). + Averaging is done in float32 and rounded on cast, matching numpy's mean + followed by an int16 cast in the original implementation. """ - n_frames = data.shape[3] - - n_full = n_frames // n + n_frames = block.shape[2] + n_full = n_frames // n remainder = n_frames % n - # Average full groups using reshape-trick (no Python loop) - if n_full > 0: - full = data[:, :, :, :n_full * n, :] # trim to full groups - full = full.reshape(data.shape[0], data.shape[1], data.shape[2], - n_full, n, data.shape[4]) # (..., n_out, n, spf) - averaged = full.mean(axis=4).astype(np.int16) # (..., n_out, spf) - else: - averaged = np.empty((*data.shape[:3], 0, data.shape[4]), dtype=np.int16) + parts = [] + if n_full: + full = block[:, :, :n_full * n, :].astype(np.float32) + full = full.reshape(block.shape[0], block.shape[1], n_full, n, block.shape[3]) + parts.append(full.mean(axis=3).astype(np.int16)) + if remainder and not discard_remainder: + tail = block[:, :, n_full * n:, :].astype(np.float32) + parts.append(tail.mean(axis=2, keepdims=True).astype(np.int16)) - if remainder > 0 and not discard_remainder: - tail = data[:, :, :, n_full * n:, :] # shape (..., remainder, spf) - tail_avg = tail.mean(axis=3, keepdims=True).astype(np.int16) - averaged = np.concatenate([averaged, tail_avg], axis=3) - - return averaged + if not parts: + return np.empty((*block.shape[:2], 0, block.shape[3]), dtype=np.int16) + return parts[0] if len(parts) == 1 else np.concatenate(parts, axis=2) -def write_sras(path: Path, src: dict, data_out: np.ndarray): - """Write a new .sras file with the averaged waveform data.""" - ver = src["ver"] - bps = src["bps"] - n_out = data_out.shape[3] +def write_averaged(out_path: Path, sras: SrasFile, mid_sections: bytes, + n: int, discard_remainder: bool) -> int: + """Stream each angle through the averager, writing as we go so peak RAM + stays at one angle's block rather than the whole file.""" + bps = sras.bytes_per_sample + n_frames_in = int(sras.n_frames[0]) + n_out = n_frames_in // n + if n_frames_in % n and not discard_remainder: + n_out += 1 - # Pack header — update only n_frames_hdr; everything else stays the same header = struct.pack( - HDR_FMT, - b"SRAS", - ver, - src["n_angles"], - src["n_rows"], - src["x_start"], - src["x_delta"], - src["vel"], - src["freq"], - n_out, # updated frame count - src["spf"], - src["sr"], - bps, - src["n_ch"], + HDR_FMT, b"SRAS", sras.version, sras.n_angles, int(sras.n_rows[0]), + float(sras.x_start_mm[0]), float(sras.x_delta_mm), + sras.velocity_mm_s, sras.laser_freq_hz, + n_out, # updated frame count + sras.samples_per_frame, sras.sample_rate_hz, bps, sras.n_channels, ) - # Encode waveform data back to original dtype - if bps == 1: - raw_out = np.clip(data_out, -128, 127).astype(np.int8).tobytes() - else: - # big-endian int16 - raw_out = data_out.astype(">i2").tobytes() - - with open(path, "wb") as f: + with open(out_path, "wb") as f: f.write(header) - f.write(src["angles_raw"]) - f.write(src["y_pos_raw"]) - - if ver >= 3: - for preamble_bytes in src["preambles"]: - f.write(struct.pack(">H", len(preamble_bytes))) - f.write(preamble_bytes) - - if ver >= 4: - bg = src["background"] - f.write(struct.pack(">I", len(bg))) - f.write(bg) - - f.write(raw_out) + f.write(mid_sections) + for a in range(sras.n_angles): + averaged = average_rows(sras.data[a], n, discard_remainder) + if bps == 1: + f.write(np.clip(averaged, -128, 127).astype(np.int8).tobytes()) + else: + f.write(averaged.astype(">i2").tobytes()) + return n_out def main(): @@ -188,32 +112,35 @@ def main(): print("Error: --n must be at least 1.", file=sys.stderr) sys.exit(1) - in_path = Path(args.input) + in_path = Path(args.input) out_path = Path(args.output) if not in_path.exists(): print(f"Error: input file not found: {in_path}", file=sys.stderr) sys.exit(1) - if out_path.resolve() == in_path.resolve(): print("Error: output path must differ from input path.", file=sys.stderr) sys.exit(1) print(f"Reading {in_path} ...", flush=True) - src = read_sras(in_path) + sras = SrasFile(str(in_path)) + if sras.version not in _SUPPORTED: + print(f"Error: unsupported .sras version: {sras.version} " + f"(this tool handles v{'/v'.join(map(str, _SUPPORTED))})", + file=sys.stderr) + sys.exit(1) - n_frames_in = src["data"].shape[3] - print(f" Version : v{src['ver']}") - print(f" Angles : {src['n_angles']}") - print(f" Rows : {src['n_rows']}") + n_frames_in = int(sras.n_frames[0]) + print(f" Version : v{sras.version}") + print(f" Angles : {sras.n_angles}") + print(f" Rows : {int(sras.n_rows[0])}") print(f" Frames (actual): {n_frames_in}") - print(f" Channels : {src['n_ch']}") - print(f" Samples/frame : {src['spf']}") - print(f" Bytes/sample : {src['bps']}") + print(f" Channels : {sras.n_channels}") + print(f" Samples/frame : {sras.samples_per_frame}") + print(f" Bytes/sample : {sras.bytes_per_sample}") if args.n == 1: print("--n 1: no averaging needed; copying file as-is.") - import shutil shutil.copy2(in_path, out_path) print(f"Wrote {out_path}") return @@ -223,27 +150,25 @@ def main(): "The entire dataset will be averaged into a single frame.") print(f"\nAveraging every {args.n} frames ...", flush=True) - data_out = average_frames(src["data"], args.n, args.discard_remainder) - n_frames_out = data_out.shape[3] + print(f"\nWriting {out_path} ...", flush=True) + mid_sections = read_header_sections(sras) + n_frames_out = write_averaged(out_path, sras, mid_sections, + args.n, args.discard_remainder) - n_full = n_frames_in // args.n - remainder = n_frames_in % args.n + n_full = n_frames_in // args.n + remainder = n_frames_in % args.n if remainder and not args.discard_remainder: status = f"({n_full} full groups + 1 partial group of {remainder})" - elif remainder and args.discard_remainder: + elif remainder: status = f"({n_full} full groups, {remainder} trailing frames discarded)" else: status = f"({n_full} full groups)" - print(f" {n_frames_in} frames -> {n_frames_out} frames {status}") - print(f"\nWriting {out_path} ...", flush=True) - write_sras(out_path, src, data_out) - - in_mb = in_path.stat().st_size / 1024**2 + in_mb = in_path.stat().st_size / 1024**2 out_mb = out_path.stat().st_size / 1024**2 print(f" Input size : {in_mb:.1f} MB") - print(f" Output size: {out_mb:.1f} MB ({out_mb/in_mb*100:.1f}% of input)") + print(f" Output size: {out_mb:.1f} MB ({out_mb / in_mb * 100:.1f}% of input)") print("Done.") diff --git a/sras_compute.py b/sras_compute.py new file mode 100644 index 0000000..280cc87 --- /dev/null +++ b/sras_compute.py @@ -0,0 +1,577 @@ +#!/usr/bin/env python3 +"""Image computation and angle alignment for .sras scans. + +Depends only on numpy/scipy (+ optional pyfftw) and sras_format, so a +multiprocessing child can import it without loading Qt or matplotlib — +which matters because Python 3.14 on macOS spawns rather than forks. +""" + +import os +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass + +import numpy as np +import scipy.fft as scipy_fft +import scipy.ndimage as scipy_ndimage + +from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, SrasFile, adc_to_mv + +# --------------------------------------------------------------------------- +# FFT backend +# --------------------------------------------------------------------------- + +try: + import pyfftw + pyfftw.interfaces.cache.enable() + PYFFTW_AVAILABLE = True +except ImportError: + PYFFTW_AVAILABLE = False + +_fft_backend = "numpy" # "numpy" or "pyfftw"; set via set_fft_backend() + + +def set_fft_backend(name: str): + """Select the rfft implementation. Module-level state, so it must be set + explicitly inside each multiprocessing child — it does not survive a + spawn.""" + global _fft_backend + _fft_backend = name if (name != "pyfftw" or PYFFTW_AVAILABLE) else "numpy" + + +def get_fft_backend() -> str: + return _fft_backend + + +def _do_rfft(x: np.ndarray, n: int | None = None, axis: int = -1, + workers: int = 1) -> np.ndarray: + """Dispatch rfft to the selected backend with optional multithreading.""" + if _fft_backend == "pyfftw" and PYFFTW_AVAILABLE: + return pyfftw.interfaces.numpy_fft.rfft(x, n=n, axis=axis, threads=workers) + return scipy_fft.rfft(x, n=n, axis=axis, workers=workers) + + +# --------------------------------------------------------------------------- +# Chunking / parallel budget +# --------------------------------------------------------------------------- +# +# Rows are batched so the float32 working buffers for one chunk stay under a +# memory budget. A fixed row count (the original design) works fine for small +# legacy scans but is catastrophic for a v6 scan with a large per-angle +# frame/sample count — e.g. a 7500-frame x 2500-sample angle needs ~2.4 GB for +# a single 32-row chunk. +# +# With chunks running concurrently the budget has to cover *all* live chunks at +# once. Note that on a large scan chunk_rows is already clamped to its floor of +# 1 row (one row alone is ~75 MB of float32 at 7507x2500), so shrinking the +# per-chunk size cannot buy more concurrency — the worker count must be derived +# from the budget instead. See _plan_chunks. + +_TOTAL_BYTES_BUDGET = int(os.environ.get("SRAS_MEM_BUDGET_MB", 1536)) * 1024 * 1024 +_CHUNK_ROWS_MAX = 32 # cap for small scans (original behavior) +_MAX_WORKERS = int(os.environ.get("SRAS_MAX_WORKERS", 0)) or (os.cpu_count() or 4) + +# An rfft chunk holds, live at once: the float32 input, the complex64 +# transform, and the float32 power spectrum — roughly 3x the input buffer. +_FFT_LIVE_MULTIPLIER = 3 + + +def _chunk_rows_for(n_frames: int, samples_per_frame: int, + budget: int = _TOTAL_BYTES_BUDGET) -> int: + bytes_per_row = max(1, n_frames * samples_per_frame * 4) # float32 + return int(max(1, min(_CHUNK_ROWS_MAX, budget // bytes_per_row))) + + +def _plan_chunks(n_rows: int, n_frames: int, samples_per_frame: int, + live_multiplier: int = 1, + max_workers: int | None = None) -> tuple[int, int]: + """(chunk_rows, n_workers) sized so live_multiplier x chunk_rows x + n_frames x samples x 4 bytes x n_workers stays within the budget.""" + per_chunk_budget = max(1, _TOTAL_BYTES_BUDGET // live_multiplier) + chunk_rows = _chunk_rows_for(n_frames, samples_per_frame, per_chunk_budget) + + chunk_bytes = max(1, live_multiplier * chunk_rows * n_frames * samples_per_frame * 4) + cap = max_workers if max_workers is not None else _MAX_WORKERS + n_workers = int(max(1, min(cap, + _TOTAL_BYTES_BUDGET // chunk_bytes, + -(-n_rows // chunk_rows)))) + return chunk_rows, n_workers + + +def _map_row_chunks(n_rows: int, chunk_rows: int, n_workers: int, fn): + """Apply fn(r0, r1) over row chunks, in parallel when it pays. + + Chunks write to disjoint output slices, so no locking is needed. numpy + ufuncs, scipy's pocketfft, and memmap page faults all release the GIL, so + threads give real parallelism here — and on a very large file they also + keep many more page-fault requests in flight, which is what the I/O path + wants. + """ + bounds = [(r0, min(r0 + chunk_rows, n_rows)) + for r0 in range(0, n_rows, chunk_rows)] + if n_workers <= 1 or len(bounds) == 1: + for r0, r1 in bounds: + fn(r0, r1) + return + with ThreadPoolExecutor(max_workers=min(n_workers, len(bounds))) as pool: + list(pool.map(lambda b: fn(*b), bounds)) + + +# --------------------------------------------------------------------------- +# Image computation +# --------------------------------------------------------------------------- + +def compute_dc_image(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray: + """Mean of each waveform → (n_rows, n_frames) float32, in ADC counts.""" + n_rows, n_frames = sras.image_shape(angle_idx) + data = sras.data[angle_idx] + chunk_rows, n_workers = _plan_chunks(n_rows, n_frames, sras.samples_per_frame) + img = np.empty((n_rows, n_frames), dtype=np.float32) + + def chunk(r0: int, r1: int): + img[r0:r1] = data[r0:r1, ch_idx, :, :].astype(np.float32).mean(axis=-1) + + _map_row_chunks(n_rows, chunk_rows, n_workers, chunk) + return img + + +def dc_image_mv(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray: + """DC image for (angle, channel) in mV, preferring a stored v5/v7 cache.""" + cached = sras.cached_dc_mv(angle_idx, ch_idx) + if cached is not None: + return cached + return adc_to_mv(compute_dc_image(sras, angle_idx, ch_idx), *sras.cal(ch_idx)) + + +def compute_rf_image(sras: SrasFile, angle_idx: int, + dc_threshold_mv: float | None, + apply_bg_sub: bool = True, + n_fft: int | None = None, + dc4_mv: np.ndarray | None = None) -> np.ndarray: + """FFT of each CH1 waveform; pixel = peak frequency in MHz. + + Pixels where CH4_dc < dc_threshold_mv are set to 0 — and the FFT is + never run for them, since that's the expensive part. The masked-out + CH1 samples are also never *read*: the boolean mask is applied to the + raw memmap slice before any dtype conversion, so numpy only pages in + the bytes for pixels that pass the threshold (an untouched memmap page + is never read from disk). + + *dc_threshold_mv* of None means "mask nothing", and skips reading and + averaging CH4 entirely — worth a third of the I/O when caching a whole + file, where the mask is applied later at display time. + + If the DC4 image for this angle is already known (e.g. from the + DC-channel precompute cache), pass it as *dc4_mv* (mV, shape + (n_rows, n_frames)) to reuse it instead of re-reading CH4 here. + + Fast path: if the file has a precomputed peak-frequency image for this + angle (v5 PREC or v7 CACH), zero-padding is off, and the bg-sub flag + matches, the stored image is used directly — no FFT is run. + """ + n_rows, n_frames = sras.image_shape(angle_idx) + data = sras.data[angle_idx] + + # ---- Fast path: precomputed image (v5 PREC or v7 CACH) ---------------- + cached_freq = sras.precomputed_freq_mhz[angle_idx] + if (cached_freq is not None + and n_fft is None # no custom zero-padding + and sras.precomputed_bg_sub == (apply_bg_sub and sras.background is not None)): + freq_img = cached_freq.copy() + if dc_threshold_mv is not None: + # DC4 mask, in priority order: already-cached DC block, caller- + # supplied image, or a fresh (cheap — no FFT) recompute. + dc4_img = sras.cached_dc_mv(angle_idx, CH4_IDX) + if dc4_img is None: + dc4_img = dc4_mv if dc4_mv is not None else adc_to_mv( + compute_dc_image(sras, angle_idx, CH4_IDX), *sras.cal(CH4_IDX)) + freq_img[dc4_img < dc_threshold_mv] = 0.0 + return freq_img + + # ---- Chunked FFT path -------------------------------------------------- + freq_axis = sras.freq_axis_mhz(n_fft) + img = np.zeros((n_rows, n_frames), dtype=np.float32) + n_fft_bins = n_fft if n_fft is not None else sras.samples_per_frame + chunk_rows, n_workers = _plan_chunks( + n_rows, n_frames, max(sras.samples_per_frame, n_fft_bins), + live_multiplier=_FFT_LIVE_MULTIPLIER) + background = sras.background if (apply_bg_sub and sras.background is not None) else None + cal4 = sras.cal(CH4_IDX) + + def chunk(r0: int, r1: int): + if dc_threshold_mv is None: + valid = None + else: + if dc4_mv is not None: + dc4_chunk = dc4_mv[r0:r1] + else: + dc4_chunk = adc_to_mv( + data[r0:r1, CH4_IDX, :, :].astype(np.float32).mean(axis=-1), *cal4) + valid = dc4_chunk >= dc_threshold_mv # True = run the FFT + if not valid.any(): + return + + # Index the raw memmap slice with the boolean mask *before* + # converting dtype — this is a lazy view until touched, so only the + # selected elements are actually read from disk; masked-out pixels' + # pages are never paged in at all. + raw = data[r0:r1, CH1_IDX, :, :] + waves = (raw[valid] if valid is not None else raw).astype(np.float32) + + if background is not None: + waves -= background # background is 1-D (spf,) + + # Parallelism comes from the outer chunk loop, so keep the inner + # transform single-threaded to avoid oversubscribing the machine. + spectrum = _do_rfft(waves, n=n_fft, axis=-1, workers=1) + del waves + # |z|^2 without np.abs()'s extra full-size temporary. + power = spectrum.real ** 2 + power += spectrum.imag ** 2 + del spectrum + # [..., 0] not [:, 0]: the unmasked path keeps the (rows, frames, + # bins) shape, where [:, 0] would blank a whole frame. + power[..., 0] = 0.0 # suppress DC bin + peak_bins = np.argmax(power, axis=-1) + del power + + if valid is not None: + img[r0:r1][valid] = freq_axis[peak_bins] + else: + img[r0:r1] = freq_axis[peak_bins] + + _map_row_chunks(n_rows, chunk_rows, n_workers, chunk) + return img + + +# --------------------------------------------------------------------------- +# Batch cache (Convert menu) — process-pool entry point +# --------------------------------------------------------------------------- + +def cache_file(path: str, mode: str, apply_bg_sub: bool, + fft_backend: str = "numpy", max_workers: int = 0) -> str: + """Compute and store DC or FFT images for every angle of one file, + converting v6 → v7 in place. Returns "" on success or an error message. + + Module-level and picklable so it can run in a ProcessPoolExecutor. The + FFT backend and worker cap are passed explicitly because module globals + do not survive a spawn. + """ + global _MAX_WORKERS + try: + set_fft_backend(fft_backend) + if max_workers: + _MAX_WORKERS = max_workers + + sras = SrasFile(path) + if sras.version not in (6, 7): + return (f"unsupported version {sras.version} — only v6/v7 files " + "can be batch-cached") + + n = sras.n_angles + if mode == "dc": + dc3 = [adc_to_mv(compute_dc_image(sras, a, CH3_IDX), *sras.cal(CH3_IDX)) + for a in range(n)] + dc4 = [adc_to_mv(compute_dc_image(sras, a, CH4_IDX), *sras.cal(CH4_IDX)) + for a in range(n)] + sras.write_v7_cache(new_dc3_mv=dc3, new_dc4_mv=dc4) + else: + effective_bg = apply_bg_sub and sras.background is not None + # dc_threshold_mv=None: store unmasked images and mask at display + # time (same convention as v5's PREC block). Skipping the mask + # also skips reading CH4 entirely. + freq = [compute_rf_image(sras, a, dc_threshold_mv=None, + apply_bg_sub=effective_bg) + for a in range(n)] + sras.write_v7_cache(new_freq_mhz=freq, new_bg_sub=effective_bg) + return "" + except Exception as exc: + return str(exc) + + +# --------------------------------------------------------------------------- +# Angle alignment (Fusion menu) +# +# Puts every angle's images onto one shared, zero-padded pixel grid using a +# rigid transform only (rotation + translation, never scale). Rotation for +# angle `a` is the *known* scan-angle delta relative to a reference angle — +# never searched. Only the residual translation is found, via FFT phase +# correlation of each angle's binarized CH4 ("dc-mask") image. +# +# Rotation is done in physical mm space rather than on raw pixel indices: +# the x-pixel pitch (SrasFile.pixel_x_mm) is file-wide constant but the +# y-pixel pitch (row spacing) can differ from it, and for v6 files can even +# vary per angle. Rotating the raw index grid directly would implicitly +# assume square pixels and shear a non-square-pixel image — an unwanted +# effective anisotropic scale. Instead each angle gets one affine that maps +# shared-canvas pixel index -> mm -> undo rotation/shift -> that angle's own +# local mm -> that angle's own raw pixel index, matching the output->input +# convention scipy.ndimage.affine_transform expects. +# --------------------------------------------------------------------------- + +@dataclass +class AngleTransform: + rotation_deg: float + shift_mm: tuple[float, float] # (dx_mm, dy_mm) found by phase correlation + matrix: np.ndarray # (2,2): canvas (row,col) -> this angle's raw (row,col) + offset: np.ndarray # (2,) + + +@dataclass +class AlignmentResult: + ref_angle_idx: int + dc_threshold_mv: float + canvas_shape: tuple[int, int] # (n_rows, n_cols) + canvas_dx_mm: float + canvas_dy_mm: float + canvas_origin_mm: tuple[float, float] # mm at canvas pixel index (0, 0) + per_angle: dict[int, AngleTransform] + + +def _pixel_pitch_mm(sras: SrasFile, angle_idx: int) -> tuple[float, float]: + """(dx, dy) mm/pixel for one angle: dx is the file-wide constant + pixel_x_mm; dy is this angle's own row spacing (assumed uniform, the same + assumption _redraw_image makes when it builds the display extent).""" + y = sras.y_positions_mm(angle_idx) + return sras.pixel_x_mm, (float(y[1] - y[0]) if len(y) > 1 else 1.0) + + +def _bbox_center_mm(sras: SrasFile, angle_idx: int) -> tuple[float, float]: + x = sras.x_axis_mm(angle_idx) + y = sras.y_positions_mm(angle_idx) + return float((x[0] + x[-1]) / 2.0), float((y[0] + y[-1]) / 2.0) + + +def _bbox_corners_mm(sras: SrasFile, angle_idx: int) -> np.ndarray: + """4 corners (x, y) of this angle's raw mm bounding box, shape (4, 2).""" + x = sras.x_axis_mm(angle_idx) + y = sras.y_positions_mm(angle_idx) + return np.array([[xx, yy] for xx in (x[0], x[-1]) for yy in (y[0], y[-1])]) + + +def _rotation_matrix(theta_deg: float) -> np.ndarray: + t = np.radians(theta_deg) + c, s = np.cos(t), np.sin(t) + return np.array([[c, -s], [s, c]]) # CCW rotation acting on (x, y) + + +def _theta_deg(sras: SrasFile, angle_idx: int, ref_idx: int) -> float: + return float(sras.angles_deg[angle_idx] - sras.angles_deg[ref_idx]) + + +def _corners_in_ref_frame(sras: SrasFile, angle_idx: int, ref_idx: int, + shift_mm=(0.0, 0.0)) -> np.ndarray: + """Angle *angle_idx*'s bbox corners, rotated about its own centroid into + the reference frame and translated by *shift_mm*. Shape (4, 2).""" + R = _rotation_matrix(_theta_deg(sras, angle_idx, ref_idx)) + c_a = np.array(_bbox_center_mm(sras, angle_idx)) + c_ref = np.array(_bbox_center_mm(sras, ref_idx)) + shift = np.asarray(shift_mm, dtype=np.float64) + return np.array([R @ (corner - c_a) + c_ref + shift + for corner in _bbox_corners_mm(sras, angle_idx)]) + + +def _build_affine_canvas_to_raw(sras: SrasFile, angle_idx: int, ref_idx: int, + shift_mm: tuple[float, float], + canvas_dx: float, canvas_dy: float, + canvas_origin_mm: tuple[float, float] + ) -> tuple[np.ndarray, np.ndarray]: + """matrix, offset s.t. raw_index = matrix @ [row_out, col_out] + offset, + matching scipy.ndimage.affine_transform's output->input convention. + + Pipeline (all mm unless noted): + [X;Y] = A_out @ [row_out;col_out] + b_out # canvas idx -> ref-frame mm + [lx;ly] = R(theta)^T @ ([X;Y]-c_ref-shift) + c_a # undo rotation+shift -> angle a's local mm + [row;col] = D @ ([lx;ly] - [x_start_a; y0_a]) # local mm -> angle a's raw idx + + where theta = angles_deg[angle_idx] - angles_deg[ref_idx], c_ref/c_a are + each angle's own raw-bbox mm centroid (the rotation pivot — this keeps + rotated content centered, minimizing required canvas padding), and + A_out/D are the index<->mm scaling matrices for the canvas pitch and + this angle's own native pitch respectively. + """ + Rinv = _rotation_matrix(_theta_deg(sras, angle_idx, ref_idx)).T + cx_a, cy_a = _bbox_center_mm(sras, angle_idx) + cx_ref, cy_ref = _bbox_center_mm(sras, ref_idx) + dx_a, dy_a = _pixel_pitch_mm(sras, angle_idx) + x0_a = float(sras.x_start_mm[angle_idx]) + y0_a = float(sras.y_positions_mm(angle_idx)[0]) + + A_out = np.array([[0.0, canvas_dx], [canvas_dy, 0.0]]) # [row,col] -> [X,Y] + b_out = np.array(canvas_origin_mm, dtype=np.float64) + D = np.array([[0.0, 1.0 / dy_a], [1.0 / dx_a, 0.0]]) # [x,y] -> [row,col] + shift = np.array(shift_mm, dtype=np.float64) + c_ref_v = np.array([cx_ref, cy_ref]) + c_a_v = np.array([cx_a, cy_a]) + origin_a = np.array([x0_a, y0_a]) + + matrix = D @ Rinv @ A_out + offset = D @ Rinv @ (b_out - c_ref_v - shift) + D @ (c_a_v - origin_a) + return matrix, offset + + +def apply_alignment(result: AlignmentResult, angle_idx: int, img: np.ndarray, + order: int = 0) -> np.ndarray: + """Resample any already-computed 2D image for `angle_idx` (same shape as + that angle's raw (n_rows, n_frames)) onto the shared alignment canvas. + order=0 (nearest) avoids blending real data with zero-padding or with + masked-out (0-valued) CH1/velocity pixels at mask edges. Channel- + agnostic: the same per-angle transform (found from the CH4 mask) works + for any channel's image of that angle.""" + t = result.per_angle[angle_idx] + return scipy_ndimage.affine_transform( + img.astype(np.float32, copy=False), t.matrix, offset=t.offset, + output_shape=result.canvas_shape, order=order, + mode="constant", cval=0.0) + + +def _block_mean_downsample(img: np.ndarray, factor: int) -> np.ndarray: + if factor <= 1: + return img + h, w = img.shape + h2, w2 = (h // factor) * factor, (w // factor) * factor + trimmed = img[:h2, :w2] + return trimmed.reshape(h2 // factor, factor, w2 // factor, factor).mean(axis=(1, 3)) + + +def _phase_correlate_shift(ref_img: np.ndarray, mov_img: np.ndarray) -> tuple[int, int]: + """FFT normalized cross-power-spectrum phase correlation. Returns the + integer (dr, dc) pixel shift of mov_img relative to ref_img; both must + be the same shape. Risk: if the true shift is near +/- half the array + size, wraparound can bias the peak — mitigated by the generous + margin_frac padding in _working_canvas_for_pair, which keeps the true + residual shift small relative to the correlation canvas.""" + F1 = scipy_fft.fft2(ref_img.astype(np.float64), workers=-1) + F2 = scipy_fft.fft2(mov_img.astype(np.float64), workers=-1) + R = F1 * np.conj(F2) + R /= np.maximum(np.abs(R), 1e-12) + corr = scipy_fft.ifft2(R, workers=-1).real + dr, dc = np.unravel_index(np.argmax(corr), corr.shape) + h, w = corr.shape + if dr > h // 2: + dr -= h + if dc > w // 2: + dc -= w + return int(dr), int(dc) + + +def _working_canvas_for_pair(sras: SrasFile, ref_idx: int, a_idx: int, + dx: float, dy: float, margin_frac: float = 0.3 + ) -> tuple[tuple[float, float], tuple[int, int]]: + """Union of the reference's own raw bbox and angle a's raw bbox rotated + (about its own center) into the ref frame with zero shift, padded by + margin_frac on each side — sized generously so the true phase- + correlation shift lands well inside the canvas (see + _phase_correlate_shift's wraparound note).""" + pts = np.vstack([_bbox_corners_mm(sras, ref_idx), + _corners_in_ref_frame(sras, a_idx, ref_idx)]) + x_min, y_min = pts.min(axis=0) + x_max, y_max = pts.max(axis=0) + pad_x, pad_y = (x_max - x_min) * margin_frac, (y_max - y_min) * margin_frac + x_min, x_max = x_min - pad_x, x_max + pad_x + y_min, y_max = y_min - pad_y, y_max + pad_y + n_cols = int(np.ceil((x_max - x_min) / dx)) + 1 + n_rows = int(np.ceil((y_max - y_min) / abs(dy))) + 1 + origin = (x_min, y_min if dy > 0 else y_max) + return origin, (n_rows, n_cols) + + +def _parallel_map(fn, items, n_workers: int) -> list: + """fn over items, in order, threaded when it pays.""" + items = list(items) + if n_workers <= 1 or len(items) <= 1: + return [fn(x) for x in items] + with ThreadPoolExecutor(max_workers=min(n_workers, len(items))) as pool: + return list(pool.map(fn, items)) + + +def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int, + dc_threshold_mv: float, + progress_cb=None) -> AlignmentResult: + """Top-level alignment driver. Runs on a background thread (see + AngleAlignmentWorker) — deliberately recomputes CH4 DC images from + scratch rather than reading the GUI-thread _dc_cache dict, since + background-thread workers must not touch GUI-thread-owned caches.""" + n = sras.n_angles # already the *complete*-angle count for aborted v6 scans + dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx) + n_workers = max(1, min(_MAX_WORKERS, n)) + + # progress_cb fires from pool threads, so the counter behind it must be + # atomic. list.append is, and len() of a list is a consistent read. + _ticks: list[int] = [] + + def tick(base: int, span: int): + if progress_cb: + _ticks.append(1) + progress_cb(base + int(min(len(_ticks), n) / n * span)) + + # ---- Step 1: binarized CH4 mask per angle, native per-angle grid ----- + def mask_for(a: int) -> np.ndarray: + dc4 = adc_to_mv(compute_dc_image(sras, a, CH4_IDX), *sras.cal(CH4_IDX)) + m = (dc4 >= dc_threshold_mv).astype(np.float32) + tick(0, 25) + return m + + masks = dict(enumerate(_parallel_map(mask_for, range(n), n_workers))) + + # ---- Step 2: coarse correlation stage (downsample first, then rotate) + # Downsampling before affine_transform (not after) is what keeps this + # tractable for a v6 scan with thousands of rows/frames per angle. + max_dim = max(max(m.shape) for m in masks.values()) + factor = max(1, int(np.ceil(max_dim / 1024))) + dx_c, dy_c = dx_ref * factor, dy_ref * factor + masks_small = {a: _block_mean_downsample(m, factor) for a, m in masks.items()} + + _ticks.clear() + + def shift_for(a: int) -> tuple[float, float]: + if a == ref_angle_idx: + tick(25, 50) + return (0.0, 0.0) + work_origin, work_shape = _working_canvas_for_pair( + sras, ref_angle_idx, a, dx_c, dy_c, margin_frac=0.3) + + # Coarse canvas->raw affine at native pitch, then rescale by /factor + # so it maps coarse-canvas idx -> coarse (downsampled) raw idx — + # exact for block-mean downsampling (up to the trimmed remainder). + m_a, o_a = _build_affine_canvas_to_raw( + sras, a, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin) + m_ref, o_ref = _build_affine_canvas_to_raw( + sras, ref_angle_idx, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin) + rotated_a = scipy_ndimage.affine_transform( + masks_small[a], m_a / factor, offset=o_a / factor, + output_shape=work_shape, order=0, mode="constant", cval=0.0) + embedded_ref = scipy_ndimage.affine_transform( + masks_small[ref_angle_idx], m_ref / factor, offset=o_ref / factor, + output_shape=work_shape, order=0, mode="constant", cval=0.0) + + dr, dc = _phase_correlate_shift(embedded_ref, rotated_a) + tick(25, 50) + return (dc * dx_c, dr * dy_c) + + shifts_mm = dict(enumerate(_parallel_map(shift_for, range(n), n_workers))) + + # ---- Step 3: union bounding box over all angles (rotation+shift applied) + corners = np.vstack([_corners_in_ref_frame(sras, a, ref_angle_idx, shifts_mm[a]) + for a in range(n)]) + x_min, y_min = corners.min(axis=0) + x_max, y_max = corners.max(axis=0) + n_cols = int(np.ceil((x_max - x_min) / dx_ref)) + 1 + n_rows = int(np.ceil((y_max - y_min) / abs(dy_ref))) + 1 + canvas_origin_mm = (float(x_min), float(y_min if dy_ref > 0 else y_max)) + + # ---- Step 4: final per-angle full-resolution affine (canvas -> raw idx) + per_angle: dict[int, AngleTransform] = {} + for a in range(n): + matrix, offset = _build_affine_canvas_to_raw( + sras, a, ref_angle_idx, shifts_mm[a], dx_ref, dy_ref, canvas_origin_mm) + per_angle[a] = AngleTransform( + _theta_deg(sras, a, ref_angle_idx), shifts_mm[a], matrix, offset) + if progress_cb: + progress_cb(75 + int((a + 1) / n * 25)) + + return AlignmentResult(ref_angle_idx, dc_threshold_mv, (n_rows, n_cols), + dx_ref, dy_ref, canvas_origin_mm, per_angle) + + +# Back-compat alias for the pre-split private name (used by tooling). +_compute_angle_alignment = compute_angle_alignment diff --git a/sras_format.py b/sras_format.py new file mode 100644 index 0000000..0b2715d --- /dev/null +++ b/sras_format.py @@ -0,0 +1,593 @@ +#!/usr/bin/env python3 +"""SRAS binary scan file format — parsing and writing. + +Reads v2–v7 .sras files. Depends only on numpy + struct, so compute workers +(including multiprocessing children) can import it without pulling in Qt or +matplotlib. See scan_format.md for the full v6/v7 spec. + +Channel semantics (fixed by sc3_aui_app.py acquisition settings): + CH1 — RF Acoustic Packet (AC-coupled, 100 mV/div): FFT → peak frequency + CH3 — Bias A (DC-coupled, 50 mV/div): waveform mean + CH4 — Bias B (DC-coupled, 50 mV/div): waveform mean + +Frame-count correction: the scanner writes the *configured* frame count in the +header before acquisition, but the scope may acquire fewer frames. The actual +count is computed from the file size and used for the reshape so channels are +correctly aligned. + +Scan geometry: v6/v7 files scan a different bounding box per angle (x_start, +x_delta, n_frames, n_rows all vary by angle), so geometry is exposed per-angle +via SrasFile.n_rows / n_frames / x_start_mm arrays and the x_axis_mm() / +y_positions_mm() methods. v2–v5 files have uniform geometry across angles, so +those arrays simply repeat the same value n_angles times. + +v7 files are v6 files with an optional trailing cache section holding +precomputed per-angle DC and/or FFT images, so display never has to recompute +them after the "Convert" menu's batch actions have stored them once. +""" + +import os +import re +import struct +from pathlib import Path + +import numpy as np + +# --------------------------------------------------------------------------- +# Format constants +# --------------------------------------------------------------------------- + +# v2–v5: fixed header, uniform geometry across angles (43 bytes) +HDR_FMT = ">4sBHHffffIIdBB" +HDR_SIZE = struct.calcsize(HDR_FMT) + +# v6/v7: fixed header, per-angle geometry in a separate table (49 bytes) +HDR_FMT_V6 = ">4sBHfffffffIdBB" +HDR_SIZE_V6 = struct.calcsize(HDR_FMT_V6) + +# v6/v7: per-angle geometry record (x_start, x_delta, n_frames, n_rows) +GEO_FMT_V6 = ">ffIH" +GEO_SIZE_V6 = struct.calcsize(GEO_FMT_V6) + +# v5 precomputed-image tail +PREC_MAGIC = b"PREC" +PREC_FLAG_BG_SUB = 0x01 + +# v7 cache tail. v7 is byte-identical to v6 (the version byte is the only +# header difference) plus this optional trailing section. Sub-block sizes are +# per-angle (n_rows[a] * n_frames[a]), taken from the Per-Angle Geometry Table +# already parsed for v6 — no new geometry fields are needed. +CACH_MAGIC = b"CACH" +CACH_HDR_FMT = ">4sBB" # magic, cach_version, block_flags +CACH_HDR_SIZE = struct.calcsize(CACH_HDR_FMT) +CACH_VERSION = 1 +CACH_FLAG_DC = 0x01 +CACH_FLAG_FFT = 0x02 + +SDCB_MAGIC = b"SDCB" +SDCB_HDR_FMT = ">4sBH" # magic, reserved, n_stored +SDCB_HDR_SIZE = struct.calcsize(SDCB_HDR_FMT) + +SFFT_MAGIC = b"SFFT" +SFFT_HDR_FMT = ">4sBH" # magic, flags, n_stored +SFFT_HDR_SIZE = struct.calcsize(SFFT_HDR_FMT) +SFFT_FLAG_BG_SUB = 0x01 + +# Fixed channel indices into the .sras data array (CH1=RF, CH3/CH4=Bias DC) +CH1_IDX, CH3_IDX, CH4_IDX = 0, 1, 2 +CH_NAMES = ["CH1", "CH3", "CH4", "VEL"] + +# Fallback scope calibration used only when reading v2 files without embedded +# preambles. v3+ files carry the WFMOutpre string so these are not used. +# 50 mV/div, 8 div full-scale, int8 ADC, position = -2.72 div +# ymult = 50 mV × 8 / 256 = 1.5625 mV/count +# yoff = position × (256/8) = -2.72 × 32 = -87.04 (ADC count for 0 V) +_FALLBACK_YMULT_MV = 1.5625 # mV per ADC count +_FALLBACK_YOFF_ADC = -87.04 # ADC count that represents 0 V + + +# --------------------------------------------------------------------------- +# Calibration +# --------------------------------------------------------------------------- + +def _parse_preamble(preamble: str) -> dict[str, float]: + """Extract YMULT, YOFF, YZERO from a Tektronix WFMOutpre string. + + Returns a dict with float values for whichever keys are present. + YMULT is left in V/count as the scope reports it. + """ + result = {} + for key in ("YMULT", "YOFF", "YZERO"): + m = re.search(rf'\b{key}\s+([-+]?\d*\.?\d+(?:[Ee][+-]?\d+)?)', preamble) + if m: + result[key] = float(m.group(1)) + return result + + +def mv_to_adc(mv: float, ymult_mv: float = _FALLBACK_YMULT_MV, + yoff_adc: float = _FALLBACK_YOFF_ADC, + yzero_mv: float = 0.0) -> float: + return (mv - yzero_mv) / ymult_mv + yoff_adc + + +def adc_to_mv(adc, ymult_mv: float = _FALLBACK_YMULT_MV, + yoff_adc: float = _FALLBACK_YOFF_ADC, + yzero_mv: float = 0.0): + return (adc - yoff_adc) * ymult_mv + yzero_mv + + +# --------------------------------------------------------------------------- +# Binary read helpers +# --------------------------------------------------------------------------- + +def _read_struct(f, fmt: str) -> tuple: + return struct.unpack(fmt, f.read(struct.calcsize(fmt))) + + +def _read_preambles(f, n_ch: int) -> list[str]: + """n_ch length-prefixed UTF-8 WFMOutpre strings.""" + out = [] + for _ in range(n_ch): + (length,) = _read_struct(f, ">H") + out.append(f.read(length).decode("utf-8")) + return out + + +def _read_background(f) -> np.ndarray: + """uint32 sample count followed by that many int8 samples.""" + (n_bg,) = _read_struct(f, ">I") + return np.frombuffer(f.read(n_bg), dtype=np.int8).astype(np.float32) + + +def _read_f32_image(f, shape: tuple[int, int]) -> np.ndarray: + """One big-endian float32 image, converted to native float32. + + The conversion matters: np.frombuffer hands back a read-only big-endian + view, and these arrays flow straight into the display caches and every + downstream arithmetic op. + """ + n_bytes = shape[0] * shape[1] * 4 + return np.frombuffer(f.read(n_bytes), dtype=">f4").reshape(shape).astype(np.float32) + + +# --------------------------------------------------------------------------- +# File parser +# --------------------------------------------------------------------------- + +class SrasFile: + """Parsed in-memory representation of a v2–v7 .sras file. + + Scan geometry (rows, frames, x_start) is exposed per-angle via the + ``n_rows`` / ``n_frames`` / ``x_start_mm`` arrays and the ``x_axis_mm()`` + / ``y_positions_mm()`` methods, since v6/v7 files scan a different + bounding box per angle. v2–v5 files have uniform geometry, so these + arrays just repeat the same value ``n_angles`` times. Waveform data is + likewise exposed as ``data[angle_idx]``, an array of shape + ``(n_rows[a], n_channels, n_frames[a], samples_per_frame)``. + + Precomputed images (v5's PREC tail or v7's CACH tail) are exposed as + ``precomputed_dc3_mv`` / ``precomputed_dc4_mv`` / ``precomputed_freq_mhz``, + always as ragged per-angle lists (``list[np.ndarray | None]``, one entry + per angle, ``None`` where that angle was never stored) regardless of + source version. + """ + + def __init__(self, path: str): + self.path = Path(path) + self._parse() + + def _parse(self): + with open(self.path, "rb") as f: + magic = f.read(4) + if magic != b"SRAS": + raise ValueError(f"Bad magic bytes: {magic!r}") + (version,) = struct.unpack(">B", f.read(1)) + + self.version = version + if version in (2, 3, 4, 5): + self._parse_legacy() + elif version in (6, 7): + self._parse_v6() + else: + raise ValueError(f"Unsupported version: {version}") + + # ------------------------------------------------------------------ + # Calibration + # ------------------------------------------------------------------ + + def _set_calibration(self, preambles: list[str] | None, n_ch: int): + """Populate the per-channel ymult/yoff/yzero lists from preamble + strings, falling back to the hardcoded scope constants for v2 files + that carry no preambles.""" + if preambles is None: + self.preambles = None + self.ch_ymult_mv = [_FALLBACK_YMULT_MV] * n_ch + self.ch_yoff_adc = [_FALLBACK_YOFF_ADC] * n_ch + self.ch_yzero_mv = [0.0] * n_ch + return + + self.preambles = preambles + self.ch_ymult_mv, self.ch_yoff_adc, self.ch_yzero_mv = [], [], [] + for p in preambles: + cal = _parse_preamble(p) + # The scope reports YMULT and YZERO in volts; store both as mV. + self.ch_ymult_mv.append(cal.get("YMULT", _FALLBACK_YMULT_MV / 1000) * 1000) + self.ch_yoff_adc.append(cal.get("YOFF", _FALLBACK_YOFF_ADC)) + self.ch_yzero_mv.append(cal.get("YZERO", 0.0) * 1000) + + def cal(self, ch_idx: int) -> tuple[float, float, float]: + """(ymult_mv, yoff_adc, yzero_mv) for one channel — splat straight + into adc_to_mv / mv_to_adc.""" + return (self.ch_ymult_mv[ch_idx], self.ch_yoff_adc[ch_idx], + self.ch_yzero_mv[ch_idx]) + + def _init_precomputed(self, n_angles: int): + self.precomputed_freq_mhz: list[np.ndarray | None] = [None] * n_angles + self.precomputed_dc4_mv: list[np.ndarray | None] = [None] * n_angles + self.precomputed_dc3_mv: list[np.ndarray | None] = [None] * n_angles + self.precomputed_bg_sub: bool = False + + def cached_dc_mv(self, angle_idx: int, ch_idx: int) -> np.ndarray | None: + """A stored DC image (already in mV) for (angle, channel), or None.""" + store = self.precomputed_dc3_mv if ch_idx == CH3_IDX else self.precomputed_dc4_mv + return store[angle_idx] if angle_idx < len(store) else None + + def image_shape(self, angle_idx: int) -> tuple[int, int]: + return int(self.n_rows[angle_idx]), int(self.n_frames[angle_idx]) + + # ------------------------------------------------------------------ + # v2–v5 parsing (uniform geometry, flat waveform block) + # ------------------------------------------------------------------ + + def _parse_legacy(self): + with open(self.path, "rb") as f: + (magic, ver, n_angles, n_rows, x_start, x_delta, vel, freq, + n_frames_hdr, spf, sr, bps, n_ch) = _read_struct(f, HDR_FMT) + + self.n_angles = n_angles + self.velocity_mm_s = float(vel) + self.laser_freq_hz = float(freq) + self.n_frames_header = n_frames_hdr # configured count (may be wrong) + self.samples_per_frame = spf + self.sample_rate_hz = float(sr) + self.bytes_per_sample = bps + self.n_channels = n_ch + + self._init_precomputed(n_angles) + self.scan_aborted = False + self.n_angles_declared = n_angles + + angles = np.frombuffer(f.read(n_angles * 4), dtype=">f4").astype(np.float32) + y_pos = np.frombuffer(f.read(n_rows * 4), dtype=">f4").astype(np.float32) + + self._set_calibration(_read_preambles(f, n_ch) if ver >= 3 else None, n_ch) + self.background = _read_background(f) if ver >= 4 else None + + # Record where raw waveform data begins; np.memmap maps from here. + data_offset = f.tell() + + # ---- Determine actual frame count from file size --------------- + # For v4 and earlier the header n_frames may be the *configured* + # count before acquisition; the actual count is derived from the + # bytes on disk. For v5 files a PREC tail follows the waveform + # data, so we must not include those extra bytes in the frame count. + file_size = self.path.stat().st_size + samples_per_row_per_ch = n_ch * spf + available_bytes = file_size - data_offset + + if ver == 5: + actual_n_frames = n_frames_hdr + remainder = 0 + else: + total_samples = available_bytes // bps + per_frame = n_angles * n_rows * samples_per_row_per_ch + actual_n_frames = total_samples // per_frame + remainder = total_samples % per_frame + + self.frame_count_mismatch = (actual_n_frames != n_frames_hdr) + self.n_frames_remainder = remainder + + # ---- Memory-map the waveform data (zero RAM cost) -------------- + # Instead of f.read() → astype() (which peaks at 2× file size), + # memmap lets the OS page only the bytes that are actually touched. + data5d = np.memmap( + str(self.path), + dtype=np.int8 if bps == 1 else ">i2", + mode="r", + offset=data_offset, + shape=(n_angles, n_rows, n_ch, actual_n_frames, spf), + ) + # Expose as a list of per-angle views so downstream code shares one + # indexing convention with v6: sras.data[a][row, ch, frame, sample] + self.data = [data5d[a] for a in range(n_angles)] + + # Uniform per-angle geometry, repeated so callers don't need to + # special-case legacy vs. v6 files. + self.n_rows = np.full(n_angles, n_rows, dtype=np.int64) + self.n_frames = np.full(n_angles, actual_n_frames, dtype=np.int64) + self.x_start_mm = np.full(n_angles, float(x_start), dtype=np.float64) + self.x_delta_mm = float(x_delta) # reference only; kept for re-encode + self._y_pos_per_angle = [y_pos] * n_angles + self.angles_deg = angles + self._data_offset = data_offset + + if ver >= 5: + waveform_bytes = actual_n_frames * n_angles * n_rows * n_ch * spf * bps + prec_offset = data_offset + waveform_bytes + if file_size > prec_offset: + self._parse_prec_section(prec_offset) + + def _parse_prec_section(self, offset: int): + """Parse the v5 PREC tail that holds precomputed images. + + Stored per-angle as (freq, dc4, dc3), each a full-image float32 + block prefixed by its uint16 angle index. + """ + with open(self.path, "rb") as f: + f.seek(offset) + header_raw = f.read(6) # magic(4) + fmt_ver(1) + flags(1) + if len(header_raw) < 6 or header_raw[:4] != PREC_MAGIC: + return + self.precomputed_bg_sub = bool(header_raw[5] & PREC_FLAG_BG_SUB) + + (n_stored,) = _read_struct(f, ">H") + for _ in range(n_stored): + (aidx,) = _read_struct(f, ">H") + if aidx >= self.n_angles: + break + shape = self.image_shape(aidx) + self.precomputed_freq_mhz[aidx] = _read_f32_image(f, shape) + self.precomputed_dc4_mv[aidx] = _read_f32_image(f, shape) + self.precomputed_dc3_mv[aidx] = _read_f32_image(f, shape) + + # ------------------------------------------------------------------ + # v6/v7 parsing (per-angle geometry, ragged waveform blocks) + # ------------------------------------------------------------------ + + def _parse_v6(self): + with open(self.path, "rb") as f: + (magic, ver, n_angles, x_start_nom, y_start_nom, x_delta_nom, + y_delta_nom, row_spacing, vel, freq, spf, sr, bps, + n_ch) = _read_struct(f, HDR_FMT_V6) + + n_angles_declared = n_angles + + self.velocity_mm_s = float(vel) + self.laser_freq_hz = float(freq) + self.samples_per_frame = spf + self.sample_rate_hz = float(sr) + self.bytes_per_sample = bps + self.n_channels = n_ch + + # Reference-only fields: the ROI as entered before per-angle + # bounding-box expansion. Actual per-angle geometry used for + # rendering comes from the Per-Angle Geometry Table below. + self.x_start_nominal_mm = float(x_start_nom) + self.y_start_nominal_mm = float(y_start_nom) + self.x_delta_nominal_mm = float(x_delta_nom) + self.y_delta_nominal_mm = float(y_delta_nom) + self.row_spacing_mm = float(row_spacing) + + self.n_frames_header = None + self.frame_count_mismatch = False + self.n_frames_remainder = 0 + + angles = np.frombuffer(f.read(n_angles * 4), dtype=">f4").astype(np.float32) + + x_start = np.empty(n_angles, dtype=np.float64) + n_frames = np.empty(n_angles, dtype=np.int64) + n_rows = np.empty(n_angles, dtype=np.int64) + for a in range(n_angles): + xs, xd, nf, nr = _read_struct(f, GEO_FMT_V6) + x_start[a], n_frames[a], n_rows[a] = xs, nf, nr + + y_pos_per_angle = [ + np.frombuffer(f.read(int(n_rows[a]) * 4), dtype=">f4").astype(np.float32) + for a in range(n_angles) + ] + + self._set_calibration(_read_preambles(f, n_ch), n_ch) + self.background = _read_background(f) + + data_offset = f.tell() + + self._data_offset = data_offset + + # ---- Memory-map each angle's ragged waveform block ------------- + # v6 gives each angle its own row/frame count, so waveform data is + # no longer one uniform (n_angles, n_rows, ...) block — each angle's + # block sits at a different offset with its own shape. An aborted + # scan truncates the file mid-angle; per the format spec we keep + # whatever complete angles are present rather than refusing to open + # the file. + file_size = self.path.stat().st_size + waveform_dtype = np.int8 if bps == 1 else ">i2" + + data = [] + offset = data_offset + for a in range(n_angles): + nr, nf = int(n_rows[a]), int(n_frames[a]) + nbytes = nr * n_ch * nf * spf * bps + if offset + nbytes > file_size: + break + data.append(np.memmap( + str(self.path), dtype=waveform_dtype, mode="r", + offset=offset, shape=(nr, n_ch, nf, spf), + )) + offset += nbytes + + n_complete = len(data) + if n_complete == 0: + raise ValueError( + "v6 file has no complete angle blocks — scan was aborted " + "before the first angle finished.") + + self.data = data + self.n_angles = n_complete + self.n_angles_declared = n_angles_declared + self.scan_aborted = n_complete < n_angles_declared + self.angles_deg = angles[:n_complete] + self.x_start_mm = x_start[:n_complete] + self.n_frames = n_frames[:n_complete] + self.n_rows = n_rows[:n_complete] + self._y_pos_per_angle = y_pos_per_angle[:n_complete] + + self._init_precomputed(n_complete) + if self.version == 7 and offset < file_size: + self._parse_cach_section(offset) + + # ------------------------------------------------------------------ + # v7 cache tail (CACH section: precomputed DC / FFT images) + # ------------------------------------------------------------------ + + def _cache_tail_offset(self) -> int: + """Deterministic file offset where the CACH tail starts (or would + start), derived purely from the header + Per-Angle Geometry Table — + independent of whether a cache tail is actually present. Used by + both the parser and the in-place writer.""" + waveform_bytes = sum( + int(self.n_rows[a]) * self.n_channels * int(self.n_frames[a]) + * self.samples_per_frame * self.bytes_per_sample + for a in range(self.n_angles) + ) + return self._data_offset + int(waveform_bytes) + + def _read_cache_block(self, f, hdr_fmt: str, magic: bytes, + stores: list[list]) -> int | None: + """Read one CACH sub-block header, then its per-angle image entries + into *stores* (one list per image the block stores per angle). + + Returns the header's flags byte, or None if the block is malformed. + """ + raw = f.read(struct.calcsize(hdr_fmt)) + if len(raw) < struct.calcsize(hdr_fmt): + return None + block_magic, flags, n_stored = struct.unpack(hdr_fmt, raw) + if block_magic != magic: + return None + for _ in range(n_stored): + (angle_idx,) = _read_struct(f, ">H") + if angle_idx >= self.n_angles: + break + shape = self.image_shape(angle_idx) + for store in stores: + store[angle_idx] = _read_f32_image(f, shape) + return flags + + def _parse_cach_section(self, offset: int): + """Parse the v7 CACH tail that holds precomputed DC/FFT images.""" + with open(self.path, "rb") as f: + f.seek(offset) + header_raw = f.read(CACH_HDR_SIZE) + if len(header_raw) < CACH_HDR_SIZE: + return + magic, cach_version, block_flags = struct.unpack(CACH_HDR_FMT, header_raw) + if magic != CACH_MAGIC or cach_version != CACH_VERSION: + return + + if block_flags & CACH_FLAG_DC: + if self._read_cache_block( + f, SDCB_HDR_FMT, SDCB_MAGIC, + [self.precomputed_dc3_mv, self.precomputed_dc4_mv]) is None: + return + + if block_flags & CACH_FLAG_FFT: + flags = self._read_cache_block( + f, SFFT_HDR_FMT, SFFT_MAGIC, [self.precomputed_freq_mhz]) + if flags is None: + return + self.precomputed_bg_sub = bool(flags & SFFT_FLAG_BG_SUB) + + def write_v7_cache(self, *, + new_dc3_mv: list[np.ndarray | None] | None = None, + new_dc4_mv: list[np.ndarray | None] | None = None, + new_freq_mhz: list[np.ndarray | None] | None = None, + new_bg_sub: bool | None = None): + """Store computed DC and/or FFT images into this file's CACH tail, + in place, converting a v6 source to v7 (or updating an existing v7 + file). Only the block(s) passed in are recomputed; whichever block + isn't passed is carried forward unchanged from whatever this + ``SrasFile`` already has in memory (from parsing, or a prior write + in this same session) — its bytes are never re-read from disk. + + The waveform data itself is never touched: the cache tail always + starts at ``_cache_tail_offset()``, a fixed offset derived from the + header and geometry table alone. + """ + if self.version not in (6, 7): + raise ValueError( + f"write_v7_cache only supports v6/v7 source files, got v{self.version}") + + final_dc3 = new_dc3_mv if new_dc3_mv is not None else self.precomputed_dc3_mv + final_dc4 = new_dc4_mv if new_dc4_mv is not None else self.precomputed_dc4_mv + final_freq = new_freq_mhz if new_freq_mhz is not None else self.precomputed_freq_mhz + final_bg_sub = new_bg_sub if new_bg_sub is not None else self.precomputed_bg_sub + + dc_entries = [a for a in range(self.n_angles) if final_dc3[a] is not None] + fft_entries = [a for a in range(self.n_angles) if final_freq[a] is not None] + + block_flags = ((CACH_FLAG_DC if dc_entries else 0) + | (CACH_FLAG_FFT if fft_entries else 0)) + + payload = bytearray() + payload += struct.pack(CACH_HDR_FMT, CACH_MAGIC, CACH_VERSION, block_flags) + + if dc_entries: + payload += struct.pack(SDCB_HDR_FMT, SDCB_MAGIC, 0, len(dc_entries)) + for a in dc_entries: + payload += struct.pack(">H", a) + payload += final_dc3[a].astype(">f4").tobytes() + payload += final_dc4[a].astype(">f4").tobytes() + + if fft_entries: + fft_flags = SFFT_FLAG_BG_SUB if final_bg_sub else 0 + payload += struct.pack(SFFT_HDR_FMT, SFFT_MAGIC, fft_flags, len(fft_entries)) + for a in fft_entries: + payload += struct.pack(">H", a) + payload += final_freq[a].astype(">f4").tobytes() + + with open(self.path, "r+b") as f: + f.seek(self._cache_tail_offset()) + f.write(payload) + f.truncate() + f.flush() + os.fsync(f.fileno()) + # Version-byte flip last: if the process dies before this point, + # the file is still readable as plain v6 (v6 parsing only + # bounds-checks per-angle offset+nbytes <= file_size, it never + # asserts exactly how many bytes follow the last angle) — so an + # interrupted write can never corrupt the file, only leave + # harmless trailing bytes that the next successful write + # overwrites via this same deterministic cache offset. + f.seek(4) + f.write(struct.pack("B", 7)) + f.flush() + os.fsync(f.fileno()) + + self.version = 7 + self.precomputed_dc3_mv = final_dc3 + self.precomputed_dc4_mv = final_dc4 + self.precomputed_freq_mhz = final_freq + self.precomputed_bg_sub = final_bg_sub + + # ------------------------------------------------------------------ + # Axes helpers + # ------------------------------------------------------------------ + + @property + def pixel_x_mm(self) -> float: + return self.velocity_mm_s / self.laser_freq_hz + + def x_axis_mm(self, angle_idx: int) -> np.ndarray: + n = int(self.n_frames[angle_idx]) + return self.x_start_mm[angle_idx] + np.arange(n) * self.pixel_x_mm + + def y_positions_mm(self, angle_idx: int) -> np.ndarray: + return self._y_pos_per_angle[angle_idx] + + def time_axis_ns(self) -> np.ndarray: + return np.arange(self.samples_per_frame) / self.sample_rate_hz * 1e9 + + def freq_axis_mhz(self, n_fft: int | None = None) -> np.ndarray: + n = n_fft if n_fft is not None else self.samples_per_frame + return np.fft.rfftfreq(n, d=1.0 / self.sample_rate_hz) / 1e6 diff --git a/sras_viewer.py b/sras_viewer.py index 13c657d..6336e09 100644 --- a/sras_viewer.py +++ b/sras_viewer.py @@ -10,1324 +10,73 @@ Channel semantics (fixed by sc3_aui_app.py acquisition settings): RF images are masked: pixels where CH4_dc < dc_threshold show 0. -Frame-count correction: the scanner writes the *configured* frame count in the -header before acquisition, but the scope may acquire fewer frames. The actual -count is computed from the file size and used for the reshape so channels are -correctly aligned. - -Scan geometry: v6/v7 files scan a different bounding box per angle (x_start, -x_delta, n_frames, n_rows all vary by angle), so geometry is exposed per-angle -via SrasFile.n_rows / n_frames / x_start_mm arrays and the x_axis_mm() / -y_positions_mm() methods. v2–v5 files have uniform geometry across angles, so -those arrays simply repeat the same value n_angles times. - -v7 files are v6 files with an optional trailing cache section holding -precomputed per-angle DC and/or FFT images (see scan_format.md), so display -never has to recompute them after the "Convert" menu's batch actions have -stored them once. +File parsing lives in sras_format, image/alignment math in sras_compute, and +background workers in sras_workers — the first two import neither Qt nor +matplotlib so multiprocessing children can load them cheaply. """ -import re -import sys -import struct import faulthandler -import numpy as np +import sys from pathlib import Path -from dataclasses import dataclass -import os -faulthandler.enable() # print a native stack trace on SIGSEGV/SIGABRT/etc. - -from PyQt6.QtWidgets import ( - QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, - QGroupBox, QLabel, QPushButton, QComboBox, QSpinBox, QDoubleSpinBox, - QFileDialog, QSizePolicy, QSplitter, QCheckBox, QFrame, QProgressDialog, - QDialog, QDialogButtonBox, QRadioButton, QButtonGroup, -) -from PyQt6.QtGui import QAction -from PyQt6.QtCore import Qt, QThread, pyqtSignal, QObject +import numpy as np from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg, NavigationToolbar2QT from matplotlib.figure import Figure from matplotlib.patches import Polygon -from matplotlib.lines import Line2D from matplotlib.path import Path as MplPath +from PyQt6.QtCore import QObject, Qt, QThread, pyqtSignal +from PyQt6.QtGui import QAction +from PyQt6.QtWidgets import ( + QApplication, QButtonGroup, QCheckBox, QComboBox, QDialog, QDialogButtonBox, + QDoubleSpinBox, QFileDialog, QFrame, QGroupBox, QHBoxLayout, QLabel, + QMainWindow, QProgressDialog, QPushButton, QRadioButton, QSizePolicy, + QSpinBox, QSplitter, QVBoxLayout, QWidget, +) + +import sras_compute as compute +from sras_compute import PYFFTW_AVAILABLE, apply_alignment +from sras_format import ( + CH1_IDX, CH3_IDX, CH4_IDX, CH_NAMES, SrasFile, adc_to_mv, mv_to_adc, + _FALLBACK_YMULT_MV, _FALLBACK_YOFF_ADC, +) +from sras_workers import ( + AngleAlignmentWorker, BatchCacheWorker, ComputeWorker, DcPrecomputeWorker, + LoadWorker, +) + +faulthandler.enable() # print a native stack trace on SIGSEGV/SIGABRT/etc. # --------------------------------------------------------------------------- -# FFT backend +# Display constants # --------------------------------------------------------------------------- -_pyfftw_available = False -try: - import pyfftw - pyfftw.interfaces.cache.enable() - _pyfftw_available = True -except ImportError: - pass - -import scipy.fft as scipy_fft -import scipy.ndimage as scipy_ndimage - -# Runtime-mutable settings changed via FftOptionsDialog -_fft_backend = "numpy" # "numpy" or "pyfftw" - - -def _do_rfft(x: np.ndarray, n: int | None = None, axis: int = -1, - workers: int = 1) -> np.ndarray: - """Dispatch rfft to the selected backend with optional multithreading.""" - if _fft_backend == "pyfftw" and _pyfftw_available: - return pyfftw.interfaces.numpy_fft.rfft(x, n=n, axis=axis, threads=workers) - return scipy_fft.rfft(x, n=n, axis=axis, workers=workers) - - -# --------------------------------------------------------------------------- -# SRAS format -# --------------------------------------------------------------------------- - -# v2–v5: fixed header, uniform geometry across angles (43 bytes) -HDR_FMT = ">4sBHHffffIIdBB" -HDR_SIZE = struct.calcsize(HDR_FMT) # 43 bytes - -# v6: fixed header, per-angle geometry in a separate table (49 bytes) -HDR_FMT_V6 = ">4sBHfffffffIdBB" -HDR_SIZE_V6 = struct.calcsize(HDR_FMT_V6) # 49 bytes - -# v6: per-angle geometry table record (x_start, x_delta, n_frames, n_rows) -GEO_FMT_V6 = ">ffIH" -GEO_SIZE_V6 = struct.calcsize(GEO_FMT_V6) # 14 bytes - -# v7: identical to v6 (same header/geometry/waveform layout, version byte -# is the only header difference) plus an optional trailing cache section -# ("CACH") holding precomputed per-angle DC and/or FFT images so a v7 file -# never needs to recompute them on open. See scan_format.md for the full -# spec. Sub-block sizes are per-angle (n_rows[a] * n_frames[a]), taken from -# the Per-Angle Geometry Table already parsed for v6 — no new geometry -# fields are needed. -CACH_MAGIC = b"CACH" -CACH_HDR_FMT = ">4sBB" # magic, cach_version, block_flags -CACH_HDR_SIZE = struct.calcsize(CACH_HDR_FMT) # 6 bytes -CACH_VERSION = 1 -CACH_FLAG_DC = 0x01 -CACH_FLAG_FFT = 0x02 - -SDCB_MAGIC = b"SDCB" -SDCB_HDR_FMT = ">4sBH" # magic, reserved, n_stored -SDCB_HDR_SIZE = struct.calcsize(SDCB_HDR_FMT) # 7 bytes - -SFFT_MAGIC = b"SFFT" -SFFT_HDR_FMT = ">4sBH" # magic, flags, n_stored -SFFT_HDR_SIZE = struct.calcsize(SFFT_HDR_FMT) # 7 bytes -SFFT_FLAG_BG_SUB = 0x01 - -# Fixed-order channels in the file: index 0=CH1, 1=CH3, 2=CH4 -# Fixed channel indices into the .sras data array (CH1=RF, CH3/CH4=Bias DC) -CH1_IDX, CH3_IDX, CH4_IDX = 0, 1, 2 - CH_LABELS = [ "CH1 — RF (FFT peak freq)", "CH3 — Bias A (DC mean)", "CH4 — Bias B (DC mean)", "CH1 — Velocity (SRAS)", ] -CH_NAMES = ["CH1", "CH3", "CH4", "VEL"] # Combo index for the derived velocity mode (uses CH1_IDX data) -VELOCITY_MODE_IDX = 3 +VELOCITY_MODE_IDX = 3 # All modes that operate on CH1 waveforms -CH1_DERIVED_MODES = (CH1_IDX, VELOCITY_MODE_IDX) - -# Fallback scope calibration used only when reading v2 files without embedded -# preambles. v3+ files carry the WFMOutpre string so these are not used. -# 50 mV/div, 8 div full-scale, int8 ADC, position = -2.72 div -# ymult = 50 mV × 8 / 256 = 1.5625 mV/count -# yoff = position × (256/8) = -2.72 × 32 = -87.04 (ADC count for 0 V) -_FALLBACK_YMULT_MV = 1.5625 # mV per ADC count -_FALLBACK_YOFF_ADC = -87.04 # ADC count that represents 0 V +CH1_DERIVED_MODES = (CH1_IDX, VELOCITY_MODE_IDX) CMAPS = ["gray", "viridis", "plasma", "inferno", "hot", "jet", "RdBu_r", "seismic"] - -def _parse_preamble(preamble: str) -> dict[str, float]: - """Extract YMULT, YOFF, YZERO from a Tektronix WFMOutpre string. - - Returns a dict with float values for whichever keys are present. - YMULT is left in V/count as the scope reports it. - """ - result = {} - for key in ("YMULT", "YOFF", "YZERO"): - m = re.search(rf'\b{key}\s+([-+]?\d*\.?\d+(?:[Ee][+-]?\d+)?)', preamble) - if m: - result[key] = float(m.group(1)) - return result - - -def mv_to_adc(mv: float, ymult_mv: float = _FALLBACK_YMULT_MV, - yoff_adc: float = _FALLBACK_YOFF_ADC, - yzero_mv: float = 0.0) -> float: - return (mv - yzero_mv) / ymult_mv + yoff_adc - - -def adc_to_mv(adc: float, ymult_mv: float = _FALLBACK_YMULT_MV, - yoff_adc: float = _FALLBACK_YOFF_ADC, - yzero_mv: float = 0.0) -> float: - return (adc - yoff_adc) * ymult_mv + yzero_mv - - -# --------------------------------------------------------------------------- -# File parser -# --------------------------------------------------------------------------- - -class SrasFile: - """Parsed in-memory representation of a v2–v7 .sras file. - - Scan geometry (rows, frames, x_start) is exposed per-angle via the - ``n_rows`` / ``n_frames`` / ``x_start_mm`` arrays and the ``x_axis_mm()`` - / ``y_positions_mm()`` methods, since v6/v7 files scan a different - bounding box per angle. v2–v5 files have uniform geometry, so these - arrays just repeat the same value ``n_angles`` times. Waveform data is - likewise exposed as ``data[angle_idx]``, an array of shape - ``(n_rows[a], n_channels, n_frames[a], samples_per_frame)``. - - v7 files (v6 plus an optional trailing cache section) expose any stored - precomputed images as ``precomputed_dc3_mv`` / ``precomputed_dc4_mv`` / - ``precomputed_freq_mhz`` — for v6/v7 these are ragged per-angle lists - (``list[np.ndarray | None]``, one entry per angle); for v5 they are dense - ``(n_angles, n_rows, n_frames)`` arrays, since v5 geometry is uniform. - """ - - def __init__(self, path: str): - self.path = Path(path) - self._parse() - - def _parse(self): - with open(self.path, "rb") as f: - magic = f.read(4) - if magic != b"SRAS": - raise ValueError(f"Bad magic bytes: {magic!r}") - (version,) = struct.unpack(">B", f.read(1)) - - self.version = version - if version in (2, 3, 4, 5): - self._parse_legacy() - elif version in (6, 7): - self._parse_v6() - else: - raise ValueError(f"Unsupported version: {version}") - - # ------------------------------------------------------------------ - # v2–v5 parsing (uniform geometry, flat waveform block) - # ------------------------------------------------------------------ - - def _parse_legacy(self): - with open(self.path, "rb") as f: - fields = struct.unpack(HDR_FMT, f.read(HDR_SIZE)) - (magic, ver, n_angles, n_rows, x_start, x_delta, vel, freq, - n_frames_hdr, spf, sr, bps, n_ch) = fields - - self.n_angles = n_angles - self.velocity_mm_s = float(vel) - self.laser_freq_hz = float(freq) - self.n_frames_header = n_frames_hdr # configured count (may be wrong) - self.samples_per_frame = spf - self.sample_rate_hz = float(sr) - self.bytes_per_sample = bps - self.n_channels = n_ch - - # Precomputed-image cache (populated when reading a v5 file). - # These are (n_angles, n_rows, n_frames) float32 arrays or None. - self.precomputed_freq_mhz: np.ndarray | None = None - self.precomputed_dc4_mv: np.ndarray | None = None - self.precomputed_dc3_mv: np.ndarray | None = None - self.precomputed_bg_sub: bool = False - self.scan_aborted = False - self.n_angles_declared = n_angles - - with open(self.path, "rb") as f: - f.seek(HDR_SIZE) - angles = np.frombuffer(f.read(n_angles * 4), dtype=">f4").astype(np.float32) - y_pos = np.frombuffer(f.read(n_rows * 4), dtype=">f4").astype(np.float32) - - if ver >= 3: - preambles = [] - for _ in range(n_ch): - (length,) = struct.unpack(">H", f.read(2)) - preambles.append(f.read(length).decode("utf-8")) - self.preambles = preambles - self.ch_ymult_mv = [] - self.ch_yoff_adc = [] - self.ch_yzero_mv = [] - for p in preambles: - cal = _parse_preamble(p) - # YMULT from scope is V/count; store as mV/count - self.ch_ymult_mv.append(cal.get("YMULT", _FALLBACK_YMULT_MV / 1000) * 1000) - self.ch_yoff_adc.append(cal.get("YOFF", _FALLBACK_YOFF_ADC)) - # YZERO from scope is in V; store as mV - self.ch_yzero_mv.append(cal.get("YZERO", 0.0) * 1000) - else: - self.preambles = None - self.ch_ymult_mv = [_FALLBACK_YMULT_MV] * n_ch - self.ch_yoff_adc = [_FALLBACK_YOFF_ADC] * n_ch - self.ch_yzero_mv = [0.0] * n_ch - - if ver >= 4: - (n_bg,) = struct.unpack(">I", f.read(4)) - self.background = np.frombuffer(f.read(n_bg), dtype=np.int8).astype(np.float32) - else: - self.background = None - - # Record the byte offset where raw waveform data begins. - # np.memmap will use this to map only the waveform section. - data_offset = f.tell() - - # ---- Determine actual frame count from file size --------------- - # For v4 and earlier the header n_frames may be the *configured* - # count before acquisition; the actual count is derived from the - # bytes on disk. For v5 files a PREC tail follows the waveform - # data, so we must not include those extra bytes in the frame count. - file_size = self.path.stat().st_size - samples_per_row_per_ch = n_ch * spf - - # Upper bound: bytes from data_offset to end of file - available_bytes = file_size - data_offset - - if ver == 5: - actual_n_frames = n_frames_hdr - remainder = 0 - else: - total_samples = available_bytes // bps - actual_n_frames = total_samples // (n_angles * n_rows * samples_per_row_per_ch) - remainder = total_samples % (n_angles * n_rows * samples_per_row_per_ch) - - self.frame_count_mismatch = (actual_n_frames != n_frames_hdr) - self.n_frames_remainder = remainder - - # ---- Memory-map the waveform data (zero RAM cost) -------------- - # Instead of f.read() → astype() (which peaks at 2× file size), - # memmap lets the OS page only the bytes that are actually touched. - waveform_dtype = np.int8 if bps == 1 else ">i2" - waveform_shape = (n_angles, n_rows, n_ch, actual_n_frames, spf) - data5d = np.memmap( - str(self.path), - dtype=waveform_dtype, - mode="r", - offset=data_offset, - shape=waveform_shape, - ) - # Expose as a list of per-angle views so downstream code shares one - # indexing convention with v6: sras.data[a][row, ch, frame, sample] - self.data = [data5d[a] for a in range(n_angles)] - - # Uniform per-angle geometry, repeated so callers don't need to - # special-case legacy vs. v6 files. - self.n_rows = np.full(n_angles, n_rows, dtype=np.int64) - self.n_frames = np.full(n_angles, actual_n_frames, dtype=np.int64) - self.x_start_mm = np.full(n_angles, float(x_start), dtype=np.float64) - self._y_pos_per_angle = [y_pos] * n_angles - - self.angles_deg = angles - - # ---- Read v5 precomputed section if present -------------------- - if ver >= 5: - waveform_bytes = actual_n_frames * n_angles * n_rows * n_ch * spf * bps - prec_offset = data_offset + waveform_bytes - if file_size > prec_offset: - self._parse_prec_section(prec_offset, n_angles, n_rows, actual_n_frames) - - def _parse_prec_section(self, offset: int, - n_angles: int, n_rows: int, n_frames: int): - """Parse the v5 PREC tail that holds precomputed images.""" - _PREC_MAGIC = b"PREC" - px = n_rows * n_frames # pixels per angle image - img_bytes = px * 4 # float32 - - with open(self.path, "rb") as f: - f.seek(offset) - header_raw = f.read(6) # magic(4) + fmt_ver(1) + flags(1) - if len(header_raw) < 6 or header_raw[:4] != _PREC_MAGIC: - return - flags = header_raw[5] - self.precomputed_bg_sub = bool(flags & 0x01) - - (n_stored,) = struct.unpack(">H", f.read(2)) - if n_stored == 0: - return - - freq_buf = np.zeros((n_angles, n_rows, n_frames), dtype=np.float32) - dc4_buf = np.zeros((n_angles, n_rows, n_frames), dtype=np.float32) - dc3_buf = np.zeros((n_angles, n_rows, n_frames), dtype=np.float32) - - for _ in range(n_stored): - (aidx,) = struct.unpack(">H", f.read(2)) - if aidx >= n_angles: - break - freq_buf[aidx] = np.frombuffer( - f.read(img_bytes), dtype=">f4").reshape(n_rows, n_frames) - dc4_buf[aidx] = np.frombuffer( - f.read(img_bytes), dtype=">f4").reshape(n_rows, n_frames) - dc3_buf[aidx] = np.frombuffer( - f.read(img_bytes), dtype=">f4").reshape(n_rows, n_frames) - - self.precomputed_freq_mhz = freq_buf - self.precomputed_dc4_mv = dc4_buf - self.precomputed_dc3_mv = dc3_buf - - # ------------------------------------------------------------------ - # v6 parsing (per-angle geometry, ragged waveform blocks) - # ------------------------------------------------------------------ - - def _parse_v6(self): - with open(self.path, "rb") as f: - fields = struct.unpack(HDR_FMT_V6, f.read(HDR_SIZE_V6)) - (magic, ver, n_angles, x_start_nom, y_start_nom, x_delta_nom, - y_delta_nom, row_spacing, vel, freq, spf, sr, bps, n_ch) = fields - - n_angles_declared = n_angles - - self.velocity_mm_s = float(vel) - self.laser_freq_hz = float(freq) - self.samples_per_frame = spf - self.sample_rate_hz = float(sr) - self.bytes_per_sample = bps - self.n_channels = n_ch - - # Reference-only fields: the ROI as entered before per-angle - # bounding-box expansion. Actual per-angle geometry used for - # rendering comes from the Per-Angle Geometry Table below. - self.x_start_nominal_mm = float(x_start_nom) - self.y_start_nominal_mm = float(y_start_nom) - self.x_delta_nominal_mm = float(x_delta_nom) - self.y_delta_nominal_mm = float(y_delta_nom) - self.row_spacing_mm = float(row_spacing) - - self.n_frames_header = None - self.frame_count_mismatch = False - self.n_frames_remainder = 0 - - angles = np.frombuffer(f.read(n_angles * 4), dtype=">f4").astype(np.float32) - - x_start = np.empty(n_angles, dtype=np.float64) - n_frames = np.empty(n_angles, dtype=np.int64) - n_rows = np.empty(n_angles, dtype=np.int64) - for a in range(n_angles): - xs, xd, nf, nr = struct.unpack(GEO_FMT_V6, f.read(GEO_SIZE_V6)) - x_start[a] = xs - n_frames[a] = nf - n_rows[a] = nr - - y_pos_per_angle = [] - for a in range(n_angles): - nr = int(n_rows[a]) - y_pos_per_angle.append( - np.frombuffer(f.read(nr * 4), dtype=">f4").astype(np.float32)) - - preambles = [] - for _ in range(n_ch): - (length,) = struct.unpack(">H", f.read(2)) - preambles.append(f.read(length).decode("utf-8")) - self.preambles = preambles - self.ch_ymult_mv = [] - self.ch_yoff_adc = [] - self.ch_yzero_mv = [] - for p in preambles: - cal = _parse_preamble(p) - self.ch_ymult_mv.append(cal.get("YMULT", _FALLBACK_YMULT_MV / 1000) * 1000) - self.ch_yoff_adc.append(cal.get("YOFF", _FALLBACK_YOFF_ADC)) - self.ch_yzero_mv.append(cal.get("YZERO", 0.0) * 1000) - - (n_bg,) = struct.unpack(">I", f.read(4)) - self.background = np.frombuffer(f.read(n_bg), dtype=np.int8).astype(np.float32) - - data_offset = f.tell() - - self._data_offset = data_offset - - # ---- Memory-map each angle's ragged waveform block ------------- - # v6 gives each angle its own row/frame count, so waveform data is - # no longer one uniform (n_angles, n_rows, ...) block — each angle's - # block sits at a different offset with its own shape. An aborted - # scan truncates the file mid-angle; per the format spec we keep - # whatever complete angles are present rather than refusing to open - # the file. - file_size = self.path.stat().st_size - waveform_dtype = np.int8 if bps == 1 else ">i2" - - data = [] - offset = data_offset - n_complete = 0 - for a in range(n_angles): - nr = int(n_rows[a]) - nf = int(n_frames[a]) - nbytes = nr * n_ch * nf * spf * bps - if offset + nbytes > file_size: - break - data.append(np.memmap( - str(self.path), dtype=waveform_dtype, mode="r", - offset=offset, shape=(nr, n_ch, nf, spf), - )) - offset += nbytes - n_complete += 1 - - if n_complete == 0: - raise ValueError( - "v6 file has no complete angle blocks — scan was aborted " - "before the first angle finished.") - - self.data = data - self.n_angles = n_complete - self.n_angles_declared = n_angles_declared - self.scan_aborted = n_complete < n_angles_declared - self.angles_deg = angles[:n_complete] - self.x_start_mm = x_start[:n_complete] - self.n_frames = n_frames[:n_complete] - self.n_rows = n_rows[:n_complete] - self._y_pos_per_angle = y_pos_per_angle[:n_complete] - - # Precomputed-image cache (v7 only). Ragged per-angle lists — unlike - # v5's dense (n_angles, n_rows, n_frames) arrays, v6/v7 geometry - # varies per angle, so each entry is its own (n_rows[a], n_frames[a]) - # array or None if that angle isn't cached yet. - self.precomputed_freq_mhz: list[np.ndarray | None] = [None] * n_complete - self.precomputed_dc4_mv: list[np.ndarray | None] = [None] * n_complete - self.precomputed_dc3_mv: list[np.ndarray | None] = [None] * n_complete - self.precomputed_bg_sub: bool = False - - if self.version == 7 and offset < file_size: - self._parse_cach_section(offset) - - # ------------------------------------------------------------------ - # v7 cache tail (CACH section: precomputed DC / FFT images) - # ------------------------------------------------------------------ - - def _cache_tail_offset(self) -> int: - """Deterministic file offset where the CACH tail starts (or would - start), derived purely from the header + Per-Angle Geometry Table — - independent of whether a cache tail is actually present. Used by - both the parser and the in-place writer.""" - n_ch = self.n_channels - spf = self.samples_per_frame - bps = self.bytes_per_sample - waveform_bytes = int(sum( - int(self.n_rows[a]) * n_ch * int(self.n_frames[a]) * spf * bps - for a in range(self.n_angles) - )) - return self._data_offset + waveform_bytes - - def _parse_cach_section(self, offset: int): - """Parse the v7 CACH tail that holds precomputed DC/FFT images.""" - with open(self.path, "rb") as f: - f.seek(offset) - header_raw = f.read(CACH_HDR_SIZE) - if len(header_raw) < CACH_HDR_SIZE: - return - magic, cach_version, block_flags = struct.unpack(CACH_HDR_FMT, header_raw) - if magic != CACH_MAGIC or cach_version != CACH_VERSION: - return - - if block_flags & CACH_FLAG_DC: - sdcb_raw = f.read(SDCB_HDR_SIZE) - if len(sdcb_raw) < SDCB_HDR_SIZE: - return - sdcb_magic, _reserved, n_stored = struct.unpack(SDCB_HDR_FMT, sdcb_raw) - if sdcb_magic != SDCB_MAGIC: - return - for _ in range(n_stored): - (angle_idx,) = struct.unpack(">H", f.read(2)) - if angle_idx >= self.n_angles: - break - px = int(self.n_rows[angle_idx]) * int(self.n_frames[angle_idx]) - shape = (int(self.n_rows[angle_idx]), int(self.n_frames[angle_idx])) - self.precomputed_dc3_mv[angle_idx] = np.frombuffer( - f.read(px * 4), dtype=">f4").reshape(shape) - self.precomputed_dc4_mv[angle_idx] = np.frombuffer( - f.read(px * 4), dtype=">f4").reshape(shape) - - if block_flags & CACH_FLAG_FFT: - sfft_raw = f.read(SFFT_HDR_SIZE) - if len(sfft_raw) < SFFT_HDR_SIZE: - return - sfft_magic, flags, n_stored = struct.unpack(SFFT_HDR_FMT, sfft_raw) - if sfft_magic != SFFT_MAGIC: - return - self.precomputed_bg_sub = bool(flags & SFFT_FLAG_BG_SUB) - for _ in range(n_stored): - (angle_idx,) = struct.unpack(">H", f.read(2)) - if angle_idx >= self.n_angles: - break - px = int(self.n_rows[angle_idx]) * int(self.n_frames[angle_idx]) - shape = (int(self.n_rows[angle_idx]), int(self.n_frames[angle_idx])) - self.precomputed_freq_mhz[angle_idx] = np.frombuffer( - f.read(px * 4), dtype=">f4").reshape(shape) - - def write_v7_cache(self, *, - new_dc3_mv: list[np.ndarray | None] | None = None, - new_dc4_mv: list[np.ndarray | None] | None = None, - new_freq_mhz: list[np.ndarray | None] | None = None, - new_bg_sub: bool | None = None): - """Store computed DC and/or FFT images into this file's CACH tail, - in place, converting a v6 source to v7 (or updating an existing v7 - file). Only the block(s) passed in are recomputed; whichever block - isn't passed is carried forward unchanged from whatever this - ``SrasFile`` already has in memory (from parsing, or a prior write - in this same session) — its bytes are never re-read from disk. - - The waveform data itself is never touched: the cache tail always - starts at ``_cache_tail_offset()``, a fixed offset derived from the - header and geometry table alone. - """ - if self.version not in (6, 7): - raise ValueError( - f"write_v7_cache only supports v6/v7 source files, got v{self.version}") - - final_dc3 = new_dc3_mv if new_dc3_mv is not None else self.precomputed_dc3_mv - final_dc4 = new_dc4_mv if new_dc4_mv is not None else self.precomputed_dc4_mv - final_freq = new_freq_mhz if new_freq_mhz is not None else self.precomputed_freq_mhz - final_bg_sub = new_bg_sub if new_bg_sub is not None else self.precomputed_bg_sub - - dc_entries = [a for a in range(self.n_angles) if final_dc3[a] is not None] - fft_entries = [a for a in range(self.n_angles) if final_freq[a] is not None] - - block_flags = 0 - if dc_entries: - block_flags |= CACH_FLAG_DC - if fft_entries: - block_flags |= CACH_FLAG_FFT - - payload = bytearray() - payload += struct.pack(CACH_HDR_FMT, CACH_MAGIC, CACH_VERSION, block_flags) - - if dc_entries: - payload += struct.pack(SDCB_HDR_FMT, SDCB_MAGIC, 0, len(dc_entries)) - for a in dc_entries: - payload += struct.pack(">H", a) - payload += final_dc3[a].astype(">f4").tobytes() - payload += final_dc4[a].astype(">f4").tobytes() - - if fft_entries: - fft_flags = SFFT_FLAG_BG_SUB if final_bg_sub else 0 - payload += struct.pack(SFFT_HDR_FMT, SFFT_MAGIC, fft_flags, len(fft_entries)) - for a in fft_entries: - payload += struct.pack(">H", a) - payload += final_freq[a].astype(">f4").tobytes() - - cache_offset = self._cache_tail_offset() - with open(self.path, "r+b") as f: - f.seek(cache_offset) - f.write(payload) - f.truncate() - f.flush() - os.fsync(f.fileno()) - # Version-byte flip last: if the process dies before this point, - # the file is still readable as plain v6 (v6 parsing only - # bounds-checks per-angle offset+nbytes <= file_size, it never - # asserts exactly how many bytes follow the last angle) — so an - # interrupted write can never corrupt the file, only leave - # harmless trailing bytes that the next successful write - # overwrites via this same deterministic cache_offset. - f.seek(4) - f.write(struct.pack("B", 7)) - f.flush() - os.fsync(f.fileno()) - - self.version = 7 - self.precomputed_dc3_mv = final_dc3 - self.precomputed_dc4_mv = final_dc4 - self.precomputed_freq_mhz = final_freq - self.precomputed_bg_sub = final_bg_sub - - # ------------------------------------------------------------------ - # Axes helpers - # ------------------------------------------------------------------ - - @property - def pixel_x_mm(self) -> float: - return self.velocity_mm_s / self.laser_freq_hz - - def x_axis_mm(self, angle_idx: int) -> np.ndarray: - n = int(self.n_frames[angle_idx]) - return self.x_start_mm[angle_idx] + np.arange(n) * self.pixel_x_mm - - def y_positions_mm(self, angle_idx: int) -> np.ndarray: - return self._y_pos_per_angle[angle_idx] - - def time_axis_ns(self) -> np.ndarray: - return np.arange(self.samples_per_frame) / self.sample_rate_hz * 1e9 - - def freq_axis_mhz(self, n_fft: int | None = None) -> np.ndarray: - n = n_fft if n_fft is not None else self.samples_per_frame - return np.fft.rfftfreq(n, d=1.0 / self.sample_rate_hz) / 1e6 - - -# --------------------------------------------------------------------------- -# Image computation (vectorised) -# --------------------------------------------------------------------------- - - -# Rows are batched so the float32 working buffer for one channel's chunk -# (chunk_rows × n_frames × spf × 4 bytes) stays under this budget. A fixed -# row count (the original design) works fine for small legacy scans but is -# catastrophic for a v6 scan with a large per-angle frame/sample count — -# e.g. a 7500-frame × 2500-sample angle needs ~2.4 GB for a single 32-row -# chunk, times several such buffers alive at once for the FFT step, which -# can exceed physical RAM entirely on its own. Sizing the chunk to the -# actual dimensions keeps peak RAM bounded regardless of scan size. -_CHUNK_BYTES_BUDGET = 128 * 1024 * 1024 # ~128 MB per channel-buffer chunk -_CHUNK_ROWS_MAX = 32 # cap for small scans (old behavior) - - -def _chunk_rows_for(n_frames: int, samples_per_frame: int) -> int: - bytes_per_row = max(1, n_frames * samples_per_frame * 4) # float32 - rows = _CHUNK_BYTES_BUDGET // bytes_per_row - return int(max(1, min(_CHUNK_ROWS_MAX, rows))) - - -def compute_dc_image(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray: - """Mean of each waveform → (n_rows, n_frames) float32. - - Processes in row chunks sized to a fixed memory budget (see - ``_chunk_rows_for``) so the float32 working buffer stays bounded - regardless of scan size. - """ - n_rows = int(sras.n_rows[angle_idx]) - n_frames = int(sras.n_frames[angle_idx]) - data = sras.data[angle_idx] - chunk_rows = _chunk_rows_for(n_frames, sras.samples_per_frame) - img = np.empty((n_rows, n_frames), dtype=np.float32) - for r0 in range(0, n_rows, chunk_rows): - r1 = min(r0 + chunk_rows, n_rows) - img[r0:r1] = ( - data[r0:r1, ch_idx, :, :] - .astype(np.float32) - .mean(axis=-1) - ) - return img - - -def _cached_dc_mv(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray | None: - """Return a v7-cached DC image (mV, already converted) for - (angle_idx, ch_idx) if the file's CACH section has it, else None. - - Only applies to the ragged v6/v7 list representation of - ``precomputed_dc3_mv``/``precomputed_dc4_mv`` — v5's dense PREC arrays - don't track per-angle validity (unstored angles are left at 0.0 rather - than a sentinel), so they're intentionally not consulted here. - """ - store = sras.precomputed_dc3_mv if ch_idx == CH3_IDX else sras.precomputed_dc4_mv - if isinstance(store, list) and angle_idx < len(store): - return store[angle_idx] - return None - - -def compute_rf_image(sras: SrasFile, angle_idx: int, - dc_threshold_mv: float, - apply_bg_sub: bool = True, - n_fft: int | None = None, - dc4_mv: np.ndarray | None = None) -> np.ndarray: - """FFT of each CH1 waveform; pixel = peak frequency in MHz. - - Pixels where CH4_dc < dc_threshold_mv are set to 0 — and the FFT is - never run for them, since that's the expensive part. The masked-out - CH1 samples are also never *read*: the boolean mask is applied to the - raw memmap slice before any dtype conversion, so numpy only pages in - the bytes for pixels that pass the threshold (an untouched memmap - page is never read from disk). If the DC4 image for this angle is - already known (e.g. from the DC-channel precompute cache), pass it as - *dc4_mv* (mV, shape (n_rows, n_frames)) to reuse it directly instead - of re-reading/re-averaging the CH4 channel here; otherwise it's - computed chunk-by-chunk internally (which does need every pixel's - CH4 data, since that's what determines validity in the first place). - - Fast path: if the file contains precomputed peak-frequency images for - this angle (v5's dense per-file PREC section, or v7's ragged per-angle - CACH section), and zero-padding is not active, and the bg-sub flag - matches, the stored image is used directly — no FFT is run. - - Otherwise, data is processed in row chunks sized to a fixed memory - budget (see ``_chunk_rows_for``) to bound peak RAM regardless of scan - size. - """ - n_rows = int(sras.n_rows[angle_idx]) - n_frames = int(sras.n_frames[angle_idx]) - data = sras.data[angle_idx] - - # ---- Fast path: precomputed images (v5 PREC or v7 CACH) ---------------- - # v5 stores a dense (n_angles, n_rows, n_frames) ndarray (uniform - # geometry); v6/v7 store a ragged list, one entry per angle (None where - # that angle hasn't been cached), since geometry varies per angle. - freq_store = sras.precomputed_freq_mhz - if isinstance(freq_store, list): - angle_cached = angle_idx < len(freq_store) and freq_store[angle_idx] is not None - else: - angle_cached = freq_store is not None - - can_use_precomputed = ( - angle_cached - and n_fft is None # no custom zero-padding - and sras.precomputed_bg_sub == (apply_bg_sub and sras.background is not None) - ) - if can_use_precomputed: - freq_img = freq_store[angle_idx].copy() - - # DC4 mask, in priority order: already-cached DC block, caller- - # supplied image, or a fresh (cheap — no FFT) recompute. - dc4_store = sras.precomputed_dc4_mv - if isinstance(dc4_store, list): - dc4_img = _cached_dc_mv(sras, angle_idx, CH4_IDX) - else: - dc4_img = dc4_store[angle_idx] if dc4_store is not None else None - if dc4_img is None: - if dc4_mv is not None: - dc4_img = dc4_mv - else: - dc4_img = adc_to_mv(compute_dc_image(sras, angle_idx, CH4_IDX), - sras.ch_ymult_mv[CH4_IDX], sras.ch_yoff_adc[CH4_IDX], - sras.ch_yzero_mv[CH4_IDX]) - - freq_img[dc4_img < dc_threshold_mv] = 0.0 - return freq_img - - # ---- Chunked FFT path -------------------------------------------------- - freq_axis = sras.freq_axis_mhz(n_fft) - img = np.zeros((n_rows, n_frames), dtype=np.float32) - _n_workers = os.cpu_count() or 4 - n_fft_bins = n_fft if n_fft is not None else sras.samples_per_frame - chunk_rows = _chunk_rows_for(n_frames, max(sras.samples_per_frame, n_fft_bins)) - - for r0 in range(0, n_rows, chunk_rows): - r1 = min(r0 + chunk_rows, n_rows) - - # DC mask for this chunk (float32 expansion is only chunk-sized) - if dc4_mv is not None: - dc4_chunk = dc4_mv[r0:r1] - else: - dc4_raw = data[r0:r1, CH4_IDX, :, :].astype(np.float32) - dc4_chunk = adc_to_mv(dc4_raw.mean(axis=-1), - sras.ch_ymult_mv[CH4_IDX], - sras.ch_yoff_adc[CH4_IDX], - sras.ch_yzero_mv[CH4_IDX]) - del dc4_raw - valid = dc4_chunk >= dc_threshold_mv # True = above threshold = run FFT - - if not valid.any(): - continue - - # Index the raw memmap slice with the boolean mask *before* - # converting dtype — this is a lazy view until touched, so only - # the (n_valid, spf) selected elements are actually read from - # disk; masked-out pixels' pages are never paged in at all. - valid_waves = data[r0:r1, CH1_IDX, :, :][valid].astype(np.float32) - - if apply_bg_sub and sras.background is not None: - valid_waves -= sras.background # background is 1-D (spf,) - - fft_pow = np.abs(_do_rfft(valid_waves, n=n_fft, axis=-1, workers=_n_workers)) ** 2 - del valid_waves - fft_pow[:, 0] = 0.0 # suppress DC bin - peak_bins = np.argmax(fft_pow, axis=-1) - del fft_pow - - img[r0:r1][valid] = freq_axis[peak_bins] - - return img - - -# --------------------------------------------------------------------------- -# Angle alignment (Fusion menu) -# -# Puts every angle's images onto one shared, zero-padded pixel grid using a -# rigid transform only (rotation + translation, never scale). Rotation for -# angle `a` is the *known* scan-angle delta relative to a reference angle — -# never searched. Only the residual translation is found, via FFT phase -# correlation of each angle's binarized CH4 ("dc-mask") image. -# -# Rotation is done in physical mm space rather than on raw pixel indices: -# the x-pixel pitch (SrasFile.pixel_x_mm) is file-wide constant but the -# y-pixel pitch (row spacing) can differ from it, and for v6 files can even -# vary per angle. Rotating the raw index grid directly would implicitly -# assume square pixels and shear a non-square-pixel image — an unwanted -# effective anisotropic scale. Instead each angle gets one affine that maps -# shared-canvas pixel index -> mm -> undo rotation/shift -> that angle's own -# local mm -> that angle's own raw pixel index, matching the output->input -# convention scipy.ndimage.affine_transform expects. -# --------------------------------------------------------------------------- - -@dataclass -class AngleTransform: - rotation_deg: float - shift_mm: tuple[float, float] # (dx_mm, dy_mm) found by phase correlation - matrix: np.ndarray # (2,2): canvas (row,col) -> this angle's raw (row,col) - offset: np.ndarray # (2,) - - -@dataclass -class AlignmentResult: - ref_angle_idx: int - dc_threshold_mv: float - canvas_shape: tuple[int, int] # (n_rows, n_cols) - canvas_dx_mm: float - canvas_dy_mm: float - canvas_origin_mm: tuple[float, float] # mm at canvas pixel index (0, 0) - per_angle: dict[int, AngleTransform] - - -def _pixel_pitch_mm(sras: SrasFile, angle_idx: int) -> tuple[float, float]: - """(dx, dy) mm/pixel for one angle: dx is the file-wide constant - pixel_x_mm; dy is this angle's own row spacing (assumed uniform, same - assumption _redraw_image already makes when it builds the display - extent).""" - dx = sras.pixel_x_mm - y = sras.y_positions_mm(angle_idx) - dy = float(y[1] - y[0]) if len(y) > 1 else 1.0 - return dx, dy - - -def _bbox_center_mm(sras: SrasFile, angle_idx: int) -> tuple[float, float]: - x = sras.x_axis_mm(angle_idx) - y = sras.y_positions_mm(angle_idx) - return float((x[0] + x[-1]) / 2.0), float((y[0] + y[-1]) / 2.0) - - -def _bbox_corners_mm(sras: SrasFile, angle_idx: int) -> np.ndarray: - """4 corners (x, y) of this angle's raw mm bounding box, shape (4, 2).""" - x = sras.x_axis_mm(angle_idx) - y = sras.y_positions_mm(angle_idx) - return np.array([[xx, yy] for xx in (x[0], x[-1]) for yy in (y[0], y[-1])]) - - -def _rotation_matrix(theta_deg: float) -> np.ndarray: - t = np.radians(theta_deg) - c, s = np.cos(t), np.sin(t) - return np.array([[c, -s], [s, c]]) # CCW rotation acting on (x, y) - - -def _build_affine_canvas_to_raw(sras: SrasFile, angle_idx: int, ref_idx: int, - shift_mm: tuple[float, float], - canvas_dx: float, canvas_dy: float, - canvas_origin_mm: tuple[float, float] - ) -> tuple[np.ndarray, np.ndarray]: - """matrix, offset s.t. raw_index = matrix @ [row_out, col_out] + offset, - matching scipy.ndimage.affine_transform's output->input convention. - - Pipeline (all mm unless noted): - [X;Y] = A_out @ [row_out;col_out] + b_out # canvas idx -> ref-frame mm - [lx;ly] = R(theta)^T @ ([X;Y]-c_ref-shift) + c_a # undo rotation+shift -> angle a's local mm - [row;col] = D @ ([lx;ly] - [x_start_a; y0_a]) # local mm -> angle a's raw idx - - where theta = angles_deg[angle_idx] - angles_deg[ref_idx], c_ref/c_a are - each angle's own raw-bbox mm centroid (the rotation pivot — this keeps - rotated content centered, minimizing required canvas padding), and - A_out/D are the index<->mm scaling matrices for the canvas pitch and - this angle's own native pitch respectively. - """ - theta = float(sras.angles_deg[angle_idx] - sras.angles_deg[ref_idx]) - Rinv = _rotation_matrix(theta).T - cx_a, cy_a = _bbox_center_mm(sras, angle_idx) - cx_ref, cy_ref = _bbox_center_mm(sras, ref_idx) - dx_a, dy_a = _pixel_pitch_mm(sras, angle_idx) - x0_a = float(sras.x_start_mm[angle_idx]) - y0_a = float(sras.y_positions_mm(angle_idx)[0]) - - A_out = np.array([[0.0, canvas_dx], [canvas_dy, 0.0]]) # [row,col] -> [X,Y] - b_out = np.array(canvas_origin_mm, dtype=np.float64) - D = np.array([[0.0, 1.0 / dy_a], [1.0 / dx_a, 0.0]]) # [x,y] -> [row,col] - shift = np.array(shift_mm, dtype=np.float64) - c_ref_v = np.array([cx_ref, cy_ref]) - c_a_v = np.array([cx_a, cy_a]) - origin_a = np.array([x0_a, y0_a]) - - matrix = D @ Rinv @ A_out - offset = D @ Rinv @ (b_out - c_ref_v - shift) + D @ (c_a_v - origin_a) - return matrix, offset - - -def apply_alignment(result: AlignmentResult, angle_idx: int, img: np.ndarray, - order: int = 0) -> np.ndarray: - """Resample any already-computed 2D image for `angle_idx` (same shape as - that angle's raw (n_rows, n_frames) — e.g. compute_dc_image / - compute_rf_image output) onto the shared alignment canvas. order=0 - (nearest) avoids blending real data with zero-padding or with - masked-out (0-valued) CH1/velocity pixels at mask edges. Channel- - agnostic: the same per-angle transform (found from the CH4 mask) works - for any channel's image of that angle.""" - t = result.per_angle[angle_idx] - return scipy_ndimage.affine_transform( - img.astype(np.float32, copy=False), t.matrix, offset=t.offset, - output_shape=result.canvas_shape, order=order, - mode="constant", cval=0.0) - - -def _block_mean_downsample(img: np.ndarray, factor: int) -> np.ndarray: - if factor <= 1: - return img - h, w = img.shape - h2, w2 = (h // factor) * factor, (w // factor) * factor - trimmed = img[:h2, :w2] - return trimmed.reshape(h2 // factor, factor, w2 // factor, factor).mean(axis=(1, 3)) - - -def _phase_correlate_shift(ref_img: np.ndarray, mov_img: np.ndarray) -> tuple[int, int]: - """FFT normalized cross-power-spectrum phase correlation. Returns the - integer (dr, dc) pixel shift of mov_img relative to ref_img; both must - be the same shape. Risk: if the true shift is near +/- half the array - size, wraparound can bias the peak — mitigated by the generous - margin_frac padding in _working_canvas_for_pair, which keeps the true - residual shift small relative to the correlation canvas.""" - F1 = scipy_fft.fft2(ref_img.astype(np.float64)) - F2 = scipy_fft.fft2(mov_img.astype(np.float64)) - R = F1 * np.conj(F2) - R /= np.maximum(np.abs(R), 1e-12) - corr = scipy_fft.ifft2(R).real - dr, dc = np.unravel_index(np.argmax(corr), corr.shape) - h, w = corr.shape - if dr > h // 2: - dr -= h - if dc > w // 2: - dc -= w - return int(dr), int(dc) - - -def _working_canvas_for_pair(sras: SrasFile, ref_idx: int, a_idx: int, - dx: float, dy: float, margin_frac: float = 0.3 - ) -> tuple[tuple[float, float], tuple[int, int]]: - """Union of the reference's own raw bbox and angle a's raw bbox rotated - (about its own center) into the ref frame with zero shift, padded by - margin_frac on each side — sized generously so the true phase- - correlation shift lands well inside the canvas (see - _phase_correlate_shift's wraparound note).""" - theta = float(sras.angles_deg[a_idx] - sras.angles_deg[ref_idx]) - R = _rotation_matrix(theta) - c_a = np.array(_bbox_center_mm(sras, a_idx)) - c_ref = np.array(_bbox_center_mm(sras, ref_idx)) - pts = list(_bbox_corners_mm(sras, ref_idx)) - for corner in _bbox_corners_mm(sras, a_idx): - pts.append(R @ (corner - c_a) + c_ref) - pts = np.array(pts) - x_min, y_min = pts.min(axis=0) - x_max, y_max = pts.max(axis=0) - pad_x, pad_y = (x_max - x_min) * margin_frac, (y_max - y_min) * margin_frac - x_min, x_max = x_min - pad_x, x_max + pad_x - y_min, y_max = y_min - pad_y, y_max + pad_y - n_cols = int(np.ceil((x_max - x_min) / dx)) + 1 - n_rows = int(np.ceil((y_max - y_min) / abs(dy))) + 1 - origin = (x_min, y_min if dy > 0 else y_max) - return origin, (n_rows, n_cols) - - -def _compute_angle_alignment(sras: SrasFile, ref_angle_idx: int, - dc_threshold_mv: float, - progress_cb=None) -> AlignmentResult: - """Top-level alignment driver. Runs on a background thread (see - AngleAlignmentWorker) — deliberately recomputes CH4 DC images from - scratch via compute_dc_image rather than reading the GUI-thread - _dc_cache dict, since background-thread workers must not touch - GUI-thread-owned caches (BatchCacheWorker follows the same rule).""" - n = sras.n_angles # already the *complete*-angle count for aborted v6 scans - dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx) - - # ---- Step 1: binarized CH4 mask per angle, native per-angle grid ----- - masks: dict[int, np.ndarray] = {} - for a in range(n): - dc4 = adc_to_mv(compute_dc_image(sras, a, CH4_IDX), - sras.ch_ymult_mv[CH4_IDX], sras.ch_yoff_adc[CH4_IDX], - sras.ch_yzero_mv[CH4_IDX]) - masks[a] = (dc4 >= dc_threshold_mv).astype(np.float32) - if progress_cb: - progress_cb(int((a + 1) / n * 25)) - - # ---- Step 2: coarse correlation stage (downsample first, then rotate) - # Downsampling before affine_transform (not after) is what keeps this - # tractable for a v6 scan with thousands of rows/frames per angle. - max_dim = max(max(m.shape) for m in masks.values()) - factor = max(1, int(np.ceil(max_dim / 1024))) - dx_c, dy_c = dx_ref * factor, dy_ref * factor - masks_small = {a: _block_mean_downsample(m, factor) for a, m in masks.items()} - - shifts_mm: dict[int, tuple[float, float]] = {ref_angle_idx: (0.0, 0.0)} - for a in range(n): - if a == ref_angle_idx: - continue - work_origin, work_shape = _working_canvas_for_pair( - sras, ref_angle_idx, a, dx_c, dy_c, margin_frac=0.3) - - # Coarse canvas->raw affine at native pitch, then rescale by /factor - # so it maps coarse-canvas idx -> coarse (downsampled) raw idx — - # exact for block-mean downsampling (up to the trimmed remainder). - m_a, o_a = _build_affine_canvas_to_raw( - sras, a, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin) - m_ref, o_ref = _build_affine_canvas_to_raw( - sras, ref_angle_idx, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin) - rotated_a = scipy_ndimage.affine_transform( - masks_small[a], m_a / factor, offset=o_a / factor, - output_shape=work_shape, order=0, mode="constant", cval=0.0) - embedded_ref = scipy_ndimage.affine_transform( - masks_small[ref_angle_idx], m_ref / factor, offset=o_ref / factor, - output_shape=work_shape, order=0, mode="constant", cval=0.0) - - dr, dc = _phase_correlate_shift(embedded_ref, rotated_a) - shifts_mm[a] = (dc * dx_c, dr * dy_c) - if progress_cb: - progress_cb(25 + int((a + 1) / n * 50)) - - # ---- Step 3: union bounding box over all angles (rotation+shift applied) - corners_ref_frame = [] - for a in range(n): - theta = float(sras.angles_deg[a] - sras.angles_deg[ref_angle_idx]) - R = _rotation_matrix(theta) - c_a = np.array(_bbox_center_mm(sras, a)) - c_ref = np.array(_bbox_center_mm(sras, ref_angle_idx)) - shift = np.array(shifts_mm[a]) - for corner in _bbox_corners_mm(sras, a): - corners_ref_frame.append(R @ (corner - c_a) + c_ref + shift) - corners_ref_frame = np.array(corners_ref_frame) - x_min, y_min = corners_ref_frame.min(axis=0) - x_max, y_max = corners_ref_frame.max(axis=0) - n_cols = int(np.ceil((x_max - x_min) / dx_ref)) + 1 - n_rows = int(np.ceil((y_max - y_min) / abs(dy_ref))) + 1 - canvas_origin_mm = (float(x_min), float(y_min if dy_ref > 0 else y_max)) - canvas_shape = (n_rows, n_cols) - - # ---- Step 4: final per-angle full-resolution affine (canvas -> raw idx) - per_angle: dict[int, AngleTransform] = {} - for a in range(n): - matrix, offset = _build_affine_canvas_to_raw( - sras, a, ref_angle_idx, shifts_mm[a], dx_ref, dy_ref, canvas_origin_mm) - theta = float(sras.angles_deg[a] - sras.angles_deg[ref_angle_idx]) - per_angle[a] = AngleTransform(theta, shifts_mm[a], matrix, offset) - if progress_cb: - progress_cb(75 + int((a + 1) / n * 25)) - - return AlignmentResult(ref_angle_idx, dc_threshold_mv, canvas_shape, - dx_ref, dy_ref, canvas_origin_mm, per_angle) - - -# --------------------------------------------------------------------------- -# Background workers -# --------------------------------------------------------------------------- - -class LoadWorker(QObject): - finished = pyqtSignal(object) # SrasFile | None - error = pyqtSignal(str) - - def __init__(self, path: str): - super().__init__() - self._path = path - - def run(self): - try: - self.finished.emit(SrasFile(self._path)) - except Exception as exc: - self.error.emit(str(exc)) - self.finished.emit(None) - - -class ComputeWorker(QObject): - """Computes one displayable image for (angle, channel). - - For CH1/Velocity (FFT-derived) channels, the FFT is only run for - pixels whose DC4 (Bias B) mean is at or above dc_threshold_mv — masked - pixels are left at 0 MHz without ever being FFT'd, since that's the - expensive part of a scan. If the DC4 image for this angle is already - known (e.g. from the DC-channel precompute cache), pass it in as - *dc4_mv* to skip re-reading the CH4 channel from disk entirely. - - Emits a plain ``np.ndarray`` (already in display units, masked for - CH1/Velocity) for both DC and FFT-derived channels. - """ - finished = pyqtSignal(object) - error = pyqtSignal(str) - - def __init__(self, sras: SrasFile, angle_idx: int, - ch_idx: int, - apply_bg_sub: bool = True, - n_fft: int | None = None, - dc_threshold_mv: float = 0.0, - dc4_mv: np.ndarray | None = None): - super().__init__() - self._sras = sras - self._angle = angle_idx - self._ch = ch_idx - self._apply_bg_sub = apply_bg_sub - self._n_fft = n_fft - self._dc_threshold = dc_threshold_mv - self._dc4_mv = dc4_mv - - def run(self): - try: - if self._ch in (CH1_IDX, VELOCITY_MODE_IDX): - img = compute_rf_image( - self._sras, self._angle, dc_threshold_mv=self._dc_threshold, - apply_bg_sub=self._apply_bg_sub, n_fft=self._n_fft, - dc4_mv=self._dc4_mv) - self.finished.emit(img) - else: - # DC channels: use the v7 cache if this angle is already - # stored (already mV), else compute fresh and convert. - cached = _cached_dc_mv(self._sras, self._angle, self._ch) - if cached is not None: - self.finished.emit(cached) - return - adc_img = compute_dc_image(self._sras, self._angle, self._ch) - img = adc_to_mv(adc_img, - self._sras.ch_ymult_mv[self._ch], - self._sras.ch_yoff_adc[self._ch], - self._sras.ch_yzero_mv[self._ch]) - self.finished.emit(img) - except Exception as exc: - self.error.emit(str(exc)) - - -class DcPrecomputeWorker(QObject): - """Computes CH3/CH4 DC images for every angle in the background. - - DC images are cheap (a per-waveform mean, no FFT) compared to the - CH1/Velocity FFT, so precomputing them for the whole file right after - load makes switching angles instant while on a DC channel, and also - means the FFT masking step (which needs a DC4 image) rarely has to - wait on anything. Emits one ``angle_done`` signal per angle as it - completes rather than waiting for the whole file, so the cache fills - in progressively. - """ - angle_done = pyqtSignal(int, np.ndarray, np.ndarray) # angle_idx, dc3_mv, dc4_mv - finished = pyqtSignal() - error = pyqtSignal(str) - - def __init__(self, sras: SrasFile): - super().__init__() - self._sras = sras - self._stop = False - - def stop(self): - self._stop = True - - def run(self): - try: - for a in range(self._sras.n_angles): - if self._stop: - break - cached_dc3 = _cached_dc_mv(self._sras, a, CH3_IDX) - cached_dc4 = _cached_dc_mv(self._sras, a, CH4_IDX) - if cached_dc3 is not None and cached_dc4 is not None: - self.angle_done.emit(a, cached_dc3, cached_dc4) - continue - dc3_mv = adc_to_mv( - compute_dc_image(self._sras, a, CH3_IDX), - self._sras.ch_ymult_mv[CH3_IDX], - self._sras.ch_yoff_adc[CH3_IDX], - self._sras.ch_yzero_mv[CH3_IDX]) - dc4_mv = adc_to_mv( - compute_dc_image(self._sras, a, CH4_IDX), - self._sras.ch_ymult_mv[CH4_IDX], - self._sras.ch_yoff_adc[CH4_IDX], - self._sras.ch_yzero_mv[CH4_IDX]) - self.angle_done.emit(a, dc3_mv, dc4_mv) - self.finished.emit() - except Exception as exc: - self.error.emit(str(exc)) - - -class BatchCacheWorker(QObject): - """Batch-computes and stores DC or FFT images into each of *paths*'s - v7 CACH tail, in place — converting v6 sources to v7 on first use, or - updating an existing v7 file's cache blocks without disturbing whatever - the other block already holds. - - *mode* is ``"dc"`` (CH3/CH4 mean images) or ``"fft"`` (CH1 peak-frequency - images, unmasked — masking is applied at display time, same as v5's - PREC convention). - - Emits ``progress(int)`` (0–100, weighted by total angle count across the - whole batch), ``file_done(str, str)`` (path, error message or "" on - success) after each file so one file's failure doesn't abort the batch, - and ``finished()`` once every file has been attempted. - """ - progress = pyqtSignal(int) - file_done = pyqtSignal(str, str) - finished = pyqtSignal() - - def __init__(self, paths: list[str], mode: str, apply_bg_sub: bool): - super().__init__() - self._paths = paths - self._mode = mode - self._apply_bg_sub = apply_bg_sub - - def run(self): - # Pass 1: quick open of each file just to weight progress by total - # angle count. Don't hold every file's memmap open at once — reopen - # fresh per file in pass 2 below. - total_angles = 0 - for path in self._paths: - try: - total_angles += SrasFile(path).n_angles - except Exception: - pass # unreadable files are reported properly in pass 2 - total_angles = max(total_angles, 1) - - done_angles = 0 - for path in self._paths: - try: - sras = SrasFile(path) - if sras.version not in (6, 7): - self.file_done.emit( - path, f"unsupported version {sras.version} — only " - "v6/v7 files can be batch-cached") - continue - - n = sras.n_angles - if self._mode == "dc": - new_dc3 = [None] * n - new_dc4 = [None] * n - for a in range(n): - new_dc3[a] = adc_to_mv( - compute_dc_image(sras, a, CH3_IDX), - sras.ch_ymult_mv[CH3_IDX], sras.ch_yoff_adc[CH3_IDX], - sras.ch_yzero_mv[CH3_IDX]) - new_dc4[a] = adc_to_mv( - compute_dc_image(sras, a, CH4_IDX), - sras.ch_ymult_mv[CH4_IDX], sras.ch_yoff_adc[CH4_IDX], - sras.ch_yzero_mv[CH4_IDX]) - done_angles += 1 - self.progress.emit(int(done_angles / total_angles * 100)) - sras.write_v7_cache(new_dc3_mv=new_dc3, new_dc4_mv=new_dc4) - else: - effective_bg = self._apply_bg_sub and sras.background is not None - new_freq = [None] * n - for a in range(n): - new_freq[a] = compute_rf_image( - sras, a, dc_threshold_mv=-1e9, # mask nothing - apply_bg_sub=effective_bg) - done_angles += 1 - self.progress.emit(int(done_angles / total_angles * 100)) - sras.write_v7_cache(new_freq_mhz=new_freq, new_bg_sub=effective_bg) - - self.file_done.emit(path, "") - except Exception as exc: - self.file_done.emit(path, str(exc)) - - self.finished.emit() - - -class AngleAlignmentWorker(QObject): - """Computes rotation+translation alignment for every angle in *sras*, - referenced to *ref_angle_idx*, from each angle's binarized CH4 mask. - Rotation is analytic (from sras.angles_deg); only translation is found - by phase correlation. Uses the same progress(int)/finished(...) worker - shape as the other background-thread workers in this file. - """ - progress = pyqtSignal(int) # 0–100 - finished = pyqtSignal(object, str) # AlignmentResult|None, error msg ("" = 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=lambda pct: self.progress.emit(pct)) - self.finished.emit(result, "") - except Exception as exc: - self.finished.emit(None, str(exc)) +# (mode_str, status-bar unit, colorbar label) per channel index +_CHANNEL_DISPLAY = { + CH1_IDX: ("RF", "Peak frequency (MHz)", "MHz"), + CH3_IDX: ("DC", "DC mean (mV)", "mV"), + CH4_IDX: ("DC", "DC mean (mV)", "mV"), + VELOCITY_MODE_IDX: ("Velocity", "Velocity (m/s)", "m/s"), +} + +_CSS_HINT = "font-size: 11px; color: #aaa;" +_CSS_INFO = "font-size: 11px;" +_CSS_MUTED = "color: #888; font-size: 11px;" +_CSS_WARN = "color: #e07000; font-size: 11px;" +_CSS_BUSY = "color: #4a90d9; font-size: 11px;" # --------------------------------------------------------------------------- @@ -1348,13 +97,11 @@ class RoiQuad: self._pts = np.asarray(pts, dtype=np.float64).reshape(4, 2).copy() @classmethod - def from_bbox(cls, x0: float, y0: float, - x1: float, y1: float) -> "RoiQuad": + def from_bbox(cls, x0: float, y0: float, x1: float, y1: float) -> "RoiQuad": """Create an axis-aligned rectangle from two opposite corners.""" lx, rx = min(x0, x1), max(x0, x1) by, ty = min(y0, y1), max(y0, y1) - pts = np.array([[lx, by], [rx, by], [rx, ty], [lx, ty]]) - return cls(pts) + return cls(np.array([[lx, by], [rx, by], [rx, ty], [lx, ty]])) def copy(self) -> "RoiQuad": return RoiQuad(self._pts.copy()) @@ -1364,7 +111,6 @@ class RoiQuad: return self._pts.copy() def centroid(self) -> np.ndarray: - """Mean of the four corners.""" return self._pts.mean(axis=0) def bbox_size(self) -> np.ndarray: @@ -1378,12 +124,29 @@ class RoiQuad: y_axis: np.ndarray) -> np.ndarray: """Boolean mask (n_rows, n_frames) of pixels whose centres lie inside the quadrilateral. + + Only the quad's axis-aligned bounding box is tested — meshgrid and + contains_points over the *whole* grid would be tens of millions of + point-in-polygon tests (and hundreds of MB of float64 temporaries) + on a large scan, on every ROI edit. """ - X, Y = np.meshgrid(np.asarray(x_axis, dtype=np.float64), - np.asarray(y_axis, dtype=np.float64)) - points = np.column_stack([X.ravel(), Y.ravel()]) - inside = MplPath(self._pts).contains_points(points) - return inside.reshape(X.shape) + x = np.asarray(x_axis, dtype=np.float64) + y = np.asarray(y_axis, dtype=np.float64) + mask = np.zeros((y.size, x.size), dtype=bool) + + (x0, y0), (x1, y1) = self._pts.min(axis=0), self._pts.max(axis=0) + cols = np.nonzero((x >= x0) & (x <= x1))[0] + rows = np.nonzero((y >= y0) & (y <= y1))[0] + if cols.size == 0 or rows.size == 0: + return mask + + c0, c1 = int(cols[0]), int(cols[-1]) + 1 + r0, r1 = int(rows[0]), int(rows[-1]) + 1 + X, Y = np.meshgrid(x[c0:c1], y[r0:r1]) + inside = MplPath(self._pts).contains_points( + np.column_stack([X.ravel(), Y.ravel()])) + mask[r0:r1, c0:c1] = inside.reshape(X.shape) + return mask # --------------------------------------------------------------------------- @@ -1391,18 +154,18 @@ class RoiQuad: # --------------------------------------------------------------------------- class ImageCanvas(FigureCanvasQTAgg): - pixel_clicked = pyqtSignal(int, int) # row_idx, frame_idx - roi_changed = pyqtSignal() # emitted when ROI is created / edited / cleared - draw_mode_changed = pyqtSignal(bool) # emitted when "draw new ROI" arm toggles + pixel_clicked = pyqtSignal(int, int) # row_idx, frame_idx + roi_changed = pyqtSignal() # ROI created / edited / cleared + draw_mode_changed = pyqtSignal(bool) # "draw new ROI" arm toggled # Interaction state values - _IDLE = "idle" - _DRAW_NEW = "draw_new" - _MOVE = "move" + _IDLE = "idle" + _DRAW_NEW = "draw_new" + _MOVE = "move" _DRAG_CORNER = "drag_corner" # Hit tolerance (display pixels) for handles. - _HANDLE_PX = 12 + _HANDLE_PX = 12 _CLICK_THRESH_PX = 4 # releases within this of press count as a click def __init__(self, parent=None): @@ -1411,7 +174,7 @@ class ImageCanvas(FigureCanvasQTAgg): super().__init__(fig) self.setParent(parent) self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding) - self._extent = None + self._extent = None self._img_shape = None # ROI state @@ -1421,16 +184,16 @@ class ImageCanvas(FigureCanvasQTAgg): self._draw_mode = False # Per-interaction snapshots / anchors - self._press_xy : tuple[float, float] | None = None - self._press_pixel : tuple[float, float] | None = None - self._press_button = None - self._snapshot : RoiQuad | None = None - self._drag_corner_idx: int = -1 - self._move_anchor = None # press-point in world coords - self._draw_previous : RoiQuad | None = None + self._press_xy: tuple[float, float] | None = None + self._press_pixel: tuple[float, float] | None = None + self._press_button = None + self._snapshot: RoiQuad | None = None + self._drag_corner_idx: int = -1 + self._move_anchor = None # press-point in world coords + self._draw_previous: RoiQuad | None = None - self.mpl_connect("button_press_event", self._on_press) - self.mpl_connect("motion_notify_event", self._on_motion) + self.mpl_connect("button_press_event", self._on_press) + self.mpl_connect("motion_notify_event", self._on_motion) self.mpl_connect("button_release_event", self._on_release) # ------------------------------------------------------------------ @@ -1445,7 +208,7 @@ class ImageCanvas(FigureCanvasQTAgg): # Patches and lines are destroyed by figure.clf(); drop stale refs. self._roi_artists = [] - self._extent = extent + self._extent = extent self._img_shape = img.shape im = self.ax.imshow( @@ -1512,27 +275,23 @@ class ImageCanvas(FigureCanvasQTAgg): return corners = self._roi.corners() - # Filled quad outline - poly = Polygon(corners, closed=True, fill=True, - facecolor="#ffd93a", edgecolor="#e53935", - alpha=0.22, linewidth=2.0, zorder=10) - self.ax.add_patch(poly) - self._roi_artists.append(poly) + # Filled quad, then a sharp unfilled edge for visibility over bright + # images, then draggable corner handles. + for kwargs in ( + dict(fill=True, facecolor="#ffd93a", edgecolor="#e53935", + alpha=0.22, linewidth=2.0, zorder=10), + dict(fill=False, edgecolor="#e53935", linewidth=1.8, zorder=11), + ): + patch = Polygon(corners, closed=True, **kwargs) + self.ax.add_patch(patch) + self._roi_artists.append(patch) - # Sharp edge (no fill) for better visibility over bright images - edge = Polygon(corners, closed=True, fill=False, - edgecolor="#e53935", linewidth=1.8, zorder=11) - self.ax.add_patch(edge) - self._roi_artists.append(edge) - - # Corner handles (white fill, red edge) — drag each independently - handles = self.ax.scatter(corners[:, 0], corners[:, 1], - s=60, c="white", edgecolors="#e53935", - linewidths=1.6, zorder=13) - self._roi_artists.append(handles) + self._roi_artists.append(self.ax.scatter( + corners[:, 0], corners[:, 1], s=60, c="white", + edgecolors="#e53935", linewidths=1.6, zorder=13)) # ------------------------------------------------------------------ - # Hit testing (uses display pixels for handles, data coords for "inside") + # Hit testing (display pixels for handles, data coords for "inside") # ------------------------------------------------------------------ def _hit_test(self, event) -> tuple[str, int | None] | None: @@ -1540,9 +299,8 @@ class ImageCanvas(FigureCanvasQTAgg): return None if event.x is None or event.y is None: return None - corners = self._roi.corners() - corners_disp = self.ax.transData.transform(corners) - click = np.array([event.x, event.y]) + corners_disp = self.ax.transData.transform(self._roi.corners()) + click = np.array([event.x, event.y]) for i in range(4): if np.hypot(*(corners_disp[i] - click)) <= self._HANDLE_PX: @@ -1568,8 +326,8 @@ class ImageCanvas(FigureCanvasQTAgg): if tb is not None and getattr(tb, "mode", ""): return - self._press_xy = (event.xdata, event.ydata) - self._press_pixel = (event.x, event.y) + self._press_xy = (event.xdata, event.ydata) + self._press_pixel = (event.x, event.y) self._press_button = event.button if self._draw_mode: @@ -1588,11 +346,10 @@ class ImageCanvas(FigureCanvasQTAgg): kind, idx = hit self._snapshot = self._roi.copy() - if kind == "corner": self._state = self._DRAG_CORNER self._drag_corner_idx = idx - elif kind == "inside": + else: self._state = self._MOVE self._move_anchor = (event.xdata, event.ydata) @@ -1606,14 +363,11 @@ class ImageCanvas(FigureCanvasQTAgg): if self._state == self._DRAW_NEW: x0, y0 = self._press_xy - x1, y1 = event.xdata, event.ydata - self._roi = RoiQuad.from_bbox(x0, y0, x1, y1) - + self._roi = RoiQuad.from_bbox(x0, y0, event.xdata, event.ydata) elif self._state == self._MOVE: - dx = event.xdata - self._move_anchor[0] - dy = event.ydata - self._move_anchor[1] - self._roi._pts = self._snapshot.corners() + np.array([dx, dy]) - + delta = np.array([event.xdata - self._move_anchor[0], + event.ydata - self._move_anchor[1]]) + self._roi._pts = self._snapshot.corners() + delta elif self._state == self._DRAG_CORNER: self._roi._pts[self._drag_corner_idx] = [event.xdata, event.ydata] @@ -1625,59 +379,59 @@ class ImageCanvas(FigureCanvasQTAgg): return prev_state = self._state self._state = self._IDLE - - if prev_state == self._DRAW_NEW: - # Reject zero-area or vanishingly-small quads - if self._extent is not None: - x0, x1, y_bot, y_top = self._extent - # Minimum: 1% of each axis range - min_w = abs(x1 - x0) * 0.01 - min_h = abs(y_bot - y_top) * 0.01 + try: + if prev_state == self._DRAW_NEW: + self._finish_draw() + elif prev_state in (self._MOVE, self._DRAG_CORNER): + self._draw_roi() + self.draw_idle() + self.roi_changed.emit() else: - min_w = min_h = 1e-6 - if self._roi is not None: - bbox = self._roi.bbox_size() - too_small = bbox[0] < min_w or bbox[1] < min_h - else: - too_small = True - if too_small: - self._roi = self._draw_previous - self._draw_previous = None - self.cancel_drawing() - self._draw_roi() - self.draw_idle() - self.roi_changed.emit() + self._maybe_emit_pixel_click(event) + finally: self._press_xy = self._press_pixel = None self._press_button = None + + def _finish_draw(self): + """Commit (or reject) a freshly-dragged quad.""" + if self._extent is not None: + x0, x1, y_bot, y_top = self._extent + min_w = abs(x1 - x0) * 0.01 # minimum: 1% of each axis range + min_h = abs(y_bot - y_top) * 0.01 + else: + min_w = min_h = 1e-6 + + if self._roi is None: + too_small = True + else: + bbox = self._roi.bbox_size() + too_small = bbox[0] < min_w or bbox[1] < min_h + if too_small: + self._roi = self._draw_previous + + self._draw_previous = None + self.cancel_drawing() + self._draw_roi() + self.draw_idle() + self.roi_changed.emit() + + def _maybe_emit_pixel_click(self, event): + """A release close enough to its press counts as a pixel click.""" + if (self._press_pixel is None or event.x is None or event.y is None + or self._extent is None or event.inaxes is not self.ax + or event.xdata is None): + return + dx_px = event.x - self._press_pixel[0] + dy_px = event.y - self._press_pixel[1] + if dx_px * dx_px + dy_px * dy_px > self._CLICK_THRESH_PX ** 2: return - if prev_state in (self._MOVE, self._DRAG_CORNER): - self._draw_roi() - self.draw_idle() - self.roi_changed.emit() - self._press_xy = self._press_pixel = None - self._press_button = None - return - - # IDLE → treat as pixel click if release is close to press - if (self._press_pixel is not None and event.x is not None and - event.y is not None and self._extent is not None): - dx_px = event.x - self._press_pixel[0] - dy_px = event.y - self._press_pixel[1] - if (dx_px * dx_px + dy_px * dy_px - <= self._CLICK_THRESH_PX * self._CLICK_THRESH_PX - and event.inaxes is self.ax - and event.xdata is not None): - x0, x1, y_bot, y_top = self._extent - n_rows, n_frames = self._img_shape - col = int((event.xdata - x0) / (x1 - x0) * n_frames) - row = int((event.ydata - y_top) / (y_bot - y_top) * n_rows) - col = max(0, min(col, n_frames - 1)) - row = max(0, min(row, n_rows - 1)) - self.pixel_clicked.emit(row, col) - - self._press_xy = self._press_pixel = None - self._press_button = None + x0, x1, y_bot, y_top = self._extent + n_rows, n_frames = self._img_shape + col = int((event.xdata - x0) / (x1 - x0) * n_frames) + row = int((event.ydata - y_top) / (y_bot - y_top) * n_rows) + self.pixel_clicked.emit(max(0, min(row, n_rows - 1)), + max(0, min(col, n_frames - 1))) class WaveformCanvas(FigureCanvasQTAgg): @@ -1699,12 +453,12 @@ class WaveformCanvas(FigureCanvasQTAgg): on the subtracted signal. The unsubtracted FFT is also shown faintly for comparison. """ - data = sras.data[angle_idx] + data = sras.data[angle_idx] waveform = data[row_idx, CH1_IDX, frame_idx, :].astype(np.float32) - t_ns = sras.time_axis_ns() - f_mhz = sras.freq_axis_mhz() - dc3_val = data[row_idx, CH3_IDX, frame_idx, :].astype(np.float32).mean() - dc4_val = data[row_idx, CH4_IDX, frame_idx, :].astype(np.float32).mean() + t_ns = sras.time_axis_ns() + f_mhz = sras.freq_axis_mhz() + dc3_val = data[row_idx, CH3_IDX, frame_idx, :].astype(np.float32).mean() + dc4_val = data[row_idx, CH4_IDX, frame_idx, :].astype(np.float32).mean() bg = sras.background if (apply_bg_sub and sras.background is not None) else None waveform_plot = waveform - bg if bg is not None else waveform @@ -1726,19 +480,19 @@ class WaveformCanvas(FigureCanvasQTAgg): self.ax_wave.set_xlabel("Time (ns)") self.ax_wave.set_ylabel("ADC counts") bg_tag = " [bg sub]" if bg is not None else "" + dc3_mv = adc_to_mv(dc3_val, *sras.cal(CH3_IDX)) + dc4_mv = adc_to_mv(dc4_val, *sras.cal(CH4_IDX)) self.ax_wave.set_title( f"CH1 RF row={row_idx} frame={frame_idx}{bg_tag}\n" f"CH3={dc3_val:.1f} CH4={dc4_val:.1f} " - f"({adc_to_mv(dc3_val, sras.ch_ymult_mv[CH3_IDX], sras.ch_yoff_adc[CH3_IDX], sras.ch_yzero_mv[CH3_IDX]):.2f} / " - f"{adc_to_mv(dc4_val, sras.ch_ymult_mv[CH4_IDX], sras.ch_yoff_adc[CH4_IDX], sras.ch_yzero_mv[CH4_IDX]):.2f} mV)", + f"({dc3_mv:.2f} / {dc4_mv:.2f} mV)", fontsize=8, ) # FFT of the (possibly subtracted) waveform power_sub = np.abs(np.fft.rfft(waveform_plot)) ** 2 power_sub[0] = 0.0 - peak_idx = int(np.argmax(power_sub)) - peak_mhz = f_mhz[peak_idx] + peak_mhz = f_mhz[int(np.argmax(power_sub))] if bg is not None: # Also show the unsubtracted FFT for reference @@ -1763,29 +517,24 @@ class WaveformCanvas(FigureCanvasQTAgg): row_idx: int, frame_idx: int): """CH3 or CH4 DC: time-domain + mean annotation.""" waveform = sras.data[angle_idx][row_idx, ch_idx, frame_idx, :].astype(np.float32) - t_ns = sras.time_axis_ns() mean_val = float(waveform.mean()) - mean_mv = adc_to_mv(mean_val, sras.ch_ymult_mv[ch_idx], sras.ch_yoff_adc[ch_idx], - sras.ch_yzero_mv[ch_idx]) + mean_mv = adc_to_mv(mean_val, *sras.cal(ch_idx)) self.ax_wave.cla() self.ax_right.cla() - self.ax_wave.plot(t_ns, waveform, linewidth=0.7, color="#4488cc") + self.ax_wave.plot(sras.time_axis_ns(), waveform, linewidth=0.7, color="#4488cc") self.ax_wave.axhline(mean_val, color="tomato", linestyle="--", linewidth=1.2, label=f"mean = {mean_val:.2f} ADC") self.ax_wave.set_xlabel("Time (ns)") self.ax_wave.set_ylabel("ADC counts") self.ax_wave.set_title( - f"{CH_NAMES[ch_idx]} DC row={row_idx} frame={frame_idx}" - ) + f"{CH_NAMES[ch_idx]} DC row={row_idx} frame={frame_idx}") self.ax_wave.legend(fontsize=8) self.ax_right.text( 0.5, 0.5, - f"DC mode\n\n" - f"mean = {mean_val:.3f} ADC\n" - f" = {mean_mv:.3f} mV", + f"DC mode\n\nmean = {mean_val:.3f} ADC\n = {mean_mv:.3f} mV", ha="center", va="center", transform=self.ax_right.transAxes, fontsize=11, ) @@ -1819,8 +568,8 @@ class FftOptionsDialog(QDialog): self.setMinimumWidth(380) self._samples_per_frame = samples_per_frame - self._sample_rate_hz = sample_rate_hz - self._grating_um = grating_um + self._sample_rate_hz = sample_rate_hz + self._grating_um = grating_um layout = QVBoxLayout(self) @@ -1828,18 +577,17 @@ class FftOptionsDialog(QDialog): grp_backend = QGroupBox("FFT Backend") bl = QVBoxLayout(grp_backend) - self._btn_numpy = QRadioButton( - "NumPy FFT (always available)") + self._btn_numpy = QRadioButton("NumPy FFT (always available)") self._btn_pyfftw = QRadioButton( - "pyFFTW (faster for large arrays)" if _pyfftw_available + "pyFFTW (faster for large arrays)" if PYFFTW_AVAILABLE else "pyFFTW (not installed — run: pip install pyfftw)") - self._btn_pyfftw.setEnabled(_pyfftw_available) + self._btn_pyfftw.setEnabled(PYFFTW_AVAILABLE) self._backend_group = QButtonGroup(self) - self._backend_group.addButton(self._btn_numpy, id=0) + self._backend_group.addButton(self._btn_numpy, id=0) self._backend_group.addButton(self._btn_pyfftw, id=1) - if current_backend == "pyfftw" and _pyfftw_available: + if current_backend == "pyfftw" and PYFFTW_AVAILABLE: self._btn_pyfftw.setChecked(True) else: self._btn_numpy.setChecked(True) @@ -1869,30 +617,26 @@ class FftOptionsDialog(QDialog): self._lbl_nfft = QLabel() self._lbl_freq_res = QLabel() - self._lbl_vel_res = QLabel() + self._lbl_vel_res = QLabel() for lbl in (self._lbl_nfft, self._lbl_freq_res, self._lbl_vel_res): - lbl.setStyleSheet("font-size: 11px; color: #aaa;") + lbl.setStyleSheet(_CSS_HINT) zl.addWidget(lbl) layout.addWidget(grp_zp) # ---- Buttons --------------------------------------------------- buttons = QDialogButtonBox() - self._apply_btn = buttons.addButton( - "Apply", QDialogButtonBox.ButtonRole.AcceptRole) - self._cancel_btn = buttons.addButton( - "Cancel", QDialogButtonBox.ButtonRole.RejectRole) - self._apply_btn.clicked.connect(self.accept) - self._cancel_btn.clicked.connect(self.reject) + buttons.addButton("Apply", QDialogButtonBox.ButtonRole.AcceptRole + ).clicked.connect(self.accept) + buttons.addButton("Cancel", QDialogButtonBox.ButtonRole.RejectRole + ).clicked.connect(self.reject) layout.addWidget(buttons) self._update_info() - # ------------------------------------------------------------------ - def _update_info(self): spf = self._samples_per_frame - sr = self._sample_rate_hz + sr = self._sample_rate_hz pad = self._spin_pad.value() if spf is None or sr is None: @@ -1902,13 +646,12 @@ class FftOptionsDialog(QDialog): return n_fft = spf * pad - freq_res_hz = sr / n_fft + freq_res_hz = sr / n_fft freq_res_mhz = freq_res_hz / 1e6 # v (m/s) = freq (MHz) × grating (µm) - vel_res_ms = freq_res_mhz * self._grating_um + vel_res_ms = freq_res_mhz * self._grating_um - self._lbl_nfft.setText( - f"FFT points: {spf} × {pad} = {n_fft:,}") + self._lbl_nfft.setText(f"FFT points: {spf} × {pad} = {n_fft:,}") self._lbl_freq_res.setText( f"Frequency bin: {freq_res_mhz:.4f} MHz ({freq_res_hz / 1e3:.2f} kHz)") self._lbl_vel_res.setText( @@ -1916,7 +659,7 @@ class FftOptionsDialog(QDialog): f"(at grating = {self._grating_um:.2f} µm)") def get_backend(self) -> str: - return "pyfftw" if self._btn_pyfftw.isChecked() and _pyfftw_available else "numpy" + return "pyfftw" if self._btn_pyfftw.isChecked() and PYFFTW_AVAILABLE else "numpy" def get_pad_factor(self) -> int: return max(1, self._spin_pad.value()) @@ -1937,23 +680,21 @@ class SrasViewerWindow(QMainWindow): self._current_image: np.ndarray | None = None self._current_angle: int = 0 self._current_ch: int = 0 - self._load_thread: QThread | None = None - self._compute_thread: QThread | None = None self._pending_angle: int = 0 self._pending_ch: int = 0 self._pending_bg_sub: bool = True self._pending_threshold: float = 50.0 # mV - self._progress_dlg: QProgressDialog | None = None + self._pending_fft_pad_factor: int = 1 + + # Live background jobs, keyed by role — see _run_worker. + self._jobs: dict[str, tuple] = {} + self._progress_dlgs: dict[str, QProgressDialog] = {} # FFT settings (configured via FFT Options dialog) self._fft_pad_factor: int = 1 # 1 = no padding - self._pending_fft_pad_factor: int = 1 # Convert menu: batch DC/FFT compute-and-store (v6 -> v7) - self._batch_thread: QThread | None = None - self._batch_worker: BatchCacheWorker | None = None self._batch_errors: list[str] = [] - self._batch_progress_dlg: QProgressDialog | None = None # Display-only settings (colormap, grating) never trigger a # recompute — they're applied to cached data on redraw. DC images @@ -1966,14 +707,11 @@ class SrasViewerWindow(QMainWindow): # same combination is free. self._dc_cache: dict[tuple[int, int], np.ndarray] = {} self._fft_cache: dict[tuple[int, bool, int | None, float], np.ndarray] = {} - self._dc_precompute_thread: QThread | None = None self._dc_precompute_worker: DcPrecomputeWorker | None = None self._dc_generation: int = 0 # Angle alignment ("Fusion" menu) - self._alignment_result: AlignmentResult | None = None - self._alignment_thread: QThread | None = None - self._alignment_worker: AngleAlignmentWorker | None = None + self._alignment_result = None self._alignment_generation: int = 0 self._aligned_cache: dict[tuple, np.ndarray] = {} @@ -1982,6 +720,57 @@ class SrasViewerWindow(QMainWindow): if initial_path: self._load_file(initial_path) + # ------------------------------------------------------------------ + # Background job plumbing + # ------------------------------------------------------------------ + + def _run_worker(self, key: str, worker: QObject, *, + connect: tuple = (), quit_on: tuple = ("finished",), + on_done=None) -> bool: + """Move *worker* onto its own QThread and start it. Returns False if + a job under *key* is already running. + + Centralises two lifetime hazards that each cost a process abort: + + 1. The job is claimed in self._jobs *before* start() and before + anything below that can pump the Qt event loop (a + QProgressDialog.show() does on first display). If it weren't, a + re-entrant editingFinished could slip past the busy check, start a + second thread, and then have the first call's own assignment + clobber — and destroy while still running — that second QThread. + + 2. thread.finished fires as the thread winds down but does not + guarantee the OS thread has joined. Dropping the last reference to + a QThread whose thread is still running logs "QThread: Destroyed + while thread is still running" and aborts, so wait() first. + """ + if key in self._jobs: + return False + + thread = QThread() + self._jobs[key] = (thread, worker, on_done) # claim before anything pumps + worker.moveToThread(thread) + thread.started.connect(worker.run) + for signal_name, slot in connect: + getattr(worker, signal_name).connect(slot) + for signal_name in quit_on: + getattr(worker, signal_name).connect(thread.quit) + thread.finished.connect(lambda k=key: self._on_job_finished(k)) + thread.start() + return True + + def _on_job_finished(self, key: str): + job = self._jobs.pop(key, None) + if job is None: + return + thread, _worker, on_done = job + thread.wait() # join before releasing our last reference + if on_done is not None: + on_done() + + def _job_running(self, key: str) -> bool: + return key in self._jobs + # ------------------------------------------------------------------ # UI construction # ------------------------------------------------------------------ @@ -1993,27 +782,33 @@ class SrasViewerWindow(QMainWindow): root.setContentsMargins(8, 8, 8, 8) root.setSpacing(8) - # ---- Left control panel ---------------------------------------- + root.addWidget(self._build_left_panel()) + root.addWidget(self._build_canvases(), stretch=1) + root.addWidget(self._build_right_panel()) + + self.statusBar().showMessage("Open an .sras file to begin.") + self._build_menus() + + def _build_left_panel(self) -> QWidget: panel = QWidget() panel.setFixedWidth(260) panel_layout = QVBoxLayout(panel) panel_layout.setContentsMargins(0, 0, 0, 0) panel_layout.setSpacing(6) - root.addWidget(panel) - # File + # ---- File ------------------------------------------------------- grp_file = QGroupBox("File") fl = QVBoxLayout(grp_file) self.btn_open = QPushButton("Open .sras…") self.btn_open.clicked.connect(self._on_open) self.lbl_filename = QLabel("No file loaded") self.lbl_filename.setWordWrap(True) - self.lbl_filename.setStyleSheet("color: #888; font-size: 11px;") + self.lbl_filename.setStyleSheet(_CSS_MUTED) fl.addWidget(self.btn_open) fl.addWidget(self.lbl_filename) panel_layout.addWidget(grp_file) - # Scan info + # ---- Scan info -------------------------------------------------- grp_info = QGroupBox("Scan Info") il = QVBoxLayout(grp_info) self._info = {} @@ -2021,26 +816,25 @@ class SrasViewerWindow(QMainWindow): "Sample rate", "X start", "Pixel Δx", "Laser freq"): lbl = QLabel(f"{key}: —") lbl.setWordWrap(True) - lbl.setStyleSheet("font-size: 11px;") + lbl.setStyleSheet(_CSS_INFO) il.addWidget(lbl) self._info[key] = lbl - # Frame count warning (hidden until needed) - self.lbl_frame_warn = QLabel("") + + self.lbl_frame_warn = QLabel("") # frame-count / format notes self.lbl_frame_warn.setWordWrap(True) - self.lbl_frame_warn.setStyleSheet("color: #e07000; font-size: 11px;") + self.lbl_frame_warn.setStyleSheet(_CSS_WARN) il.addWidget(self.lbl_frame_warn) - # Background DC-precompute progress (hidden until a file is loaded) - self.lbl_dc_precompute = QLabel("") + + self.lbl_dc_precompute = QLabel("") # background DC-precompute progress self.lbl_dc_precompute.setWordWrap(True) - self.lbl_dc_precompute.setStyleSheet("color: #4a90d9; font-size: 11px;") + self.lbl_dc_precompute.setStyleSheet(_CSS_BUSY) il.addWidget(self.lbl_dc_precompute) panel_layout.addWidget(grp_info) - # View settings + # ---- View settings ---------------------------------------------- grp_view = QGroupBox("View Settings") vl = QVBoxLayout(grp_view) - # Angle ar = QHBoxLayout() ar.addWidget(QLabel("Angle:")) self.spin_angle = QSpinBox() @@ -2052,7 +846,6 @@ class SrasViewerWindow(QMainWindow): ar.addWidget(self.lbl_angle_deg) vl.addLayout(ar) - # Channel cr = QHBoxLayout() cr.addWidget(QLabel("Channel:")) self.combo_channel = QComboBox() @@ -2062,12 +855,12 @@ class SrasViewerWindow(QMainWindow): cr.addWidget(self.combo_channel) vl.addLayout(cr) - # DC threshold (for RF / CH1 masking) sep = QFrame() sep.setFrameShape(QFrame.Shape.HLine) sep.setStyleSheet("color: #555;") vl.addWidget(sep) + # DC threshold (for RF / CH1 masking) self.grp_threshold = QGroupBox("RF Mask Threshold (CH1 only)") tl = QVBoxLayout(self.grp_threshold) thr_row = QHBoxLayout() @@ -2082,8 +875,8 @@ class SrasViewerWindow(QMainWindow): self.spin_threshold_mv.editingFinished.connect(self._on_threshold_changed) thr_row.addWidget(self.spin_threshold_mv) tl.addLayout(thr_row) - self.lbl_threshold_adc = QLabel(f"≈ {mv_to_adc(50.0):.1f} ADC counts") # updated on file load - self.lbl_threshold_adc.setStyleSheet("font-size: 11px; color: #888;") + self.lbl_threshold_adc = QLabel(f"≈ {mv_to_adc(50.0):.1f} ADC counts") + self.lbl_threshold_adc.setStyleSheet(_CSS_MUTED) tl.addWidget(self.lbl_threshold_adc) vl.addWidget(self.grp_threshold) @@ -2110,39 +903,16 @@ class SrasViewerWindow(QMainWindow): self.chk_aligned_view.toggled.connect(self._on_aligned_view_toggled) vl.addWidget(self.chk_aligned_view) - # Velocity settings (visible only in velocity mode) - self.grp_velocity = QGroupBox("Velocity Settings (CH1 only)") - vel_l = QVBoxLayout(self.grp_velocity) - grat_row = QHBoxLayout() - grat_row.addWidget(QLabel("Grating size:")) - self.spin_grating_um = QDoubleSpinBox() - self.spin_grating_um.setRange(0.1, 1000.0) - self.spin_grating_um.setDecimals(2) - self.spin_grating_um.setSingleStep(0.5) - self.spin_grating_um.setSuffix(" µm") - self.spin_grating_um.setValue(25) - self.spin_grating_um.setEnabled(False) - self.spin_grating_um.editingFinished.connect(self._on_grating_changed) - grat_row.addWidget(self.spin_grating_um) - vel_l.addLayout(grat_row) - self.lbl_velocity_formula = QLabel("v (m/s) = freq (MHz) × grating (µm)") - self.lbl_velocity_formula.setStyleSheet("font-size: 10px; color: #888;") - vel_l.addWidget(self.lbl_velocity_formula) - self.grp_velocity.setVisible(False) - # (grp_velocity will be added to the right panel below) - - # Export self.btn_export_csv = QPushButton("Export Image as CSV…") self.btn_export_csv.setEnabled(False) self.btn_export_csv.setToolTip( - "Save the current CH1 image (one scan row per CSV line)." - ) + "Save the current CH1 image (one scan row per CSV line).") self.btn_export_csv.clicked.connect(self._on_export_csv) vl.addWidget(self.btn_export_csv) panel_layout.addWidget(grp_view) - # ---- ROI (Region of Interest) --------------------------------- + # ---- ROI --------------------------------------------------------- grp_roi = QGroupBox("ROI (Region of Interest)") rl = QVBoxLayout(grp_roi) @@ -2152,8 +922,7 @@ class SrasViewerWindow(QMainWindow): self.btn_draw_roi.setToolTip( "Arm next click+drag on the image to draw a new ROI\n" "(replaces any existing one). Click again to cancel.\n" - "After drawing, drag inside to move, grab corners to resize,\n" - "or use the handle above the top edge to rotate.\n" + "After drawing, drag inside to move, or grab corners to reshape.\n" "The ROI is persistent across channels / modes / angles." ) self.btn_draw_roi.toggled.connect(self._on_draw_roi_toggled) @@ -2175,16 +944,66 @@ class SrasViewerWindow(QMainWindow): rl.addWidget(self.btn_export_roi) self.lbl_roi_center = QLabel("centroid: —") - self.lbl_roi_size = QLabel("bbox: —") - self.lbl_roi_npix = QLabel("pixels inside: —") + self.lbl_roi_size = QLabel("bbox: —") + self.lbl_roi_npix = QLabel("pixels inside: —") for lbl in (self.lbl_roi_center, self.lbl_roi_size, self.lbl_roi_npix): - lbl.setStyleSheet("font-size: 11px; color: #aaa;") + lbl.setStyleSheet(_CSS_HINT) rl.addWidget(lbl) panel_layout.addWidget(grp_roi) panel_layout.addStretch() + return panel + + def _build_canvases(self) -> QWidget: + splitter = QSplitter(Qt.Orientation.Vertical) + + img_widget = QWidget() + img_vl = QVBoxLayout(img_widget) + img_vl.setContentsMargins(0, 0, 0, 0) + self.image_canvas = ImageCanvas() + self.image_canvas.pixel_clicked.connect(self._on_pixel_clicked) + self.image_canvas.roi_changed.connect(self._update_roi_ui) + self.image_canvas.draw_mode_changed.connect(self._on_draw_mode_changed) + img_vl.addWidget(NavigationToolbar2QT(self.image_canvas, img_widget)) + img_vl.addWidget(self.image_canvas) + splitter.addWidget(img_widget) + + wave_widget = QWidget() + wave_vl = QVBoxLayout(wave_widget) + wave_vl.setContentsMargins(0, 0, 0, 0) + self.lbl_wave_hint = QLabel( + "Click a pixel in the image above to inspect its waveform.") + self.lbl_wave_hint.setAlignment(Qt.AlignmentFlag.AlignCenter) + self.lbl_wave_hint.setStyleSheet(_CSS_MUTED) + self.wave_canvas = WaveformCanvas() + wave_vl.addWidget(self.lbl_wave_hint) + wave_vl.addWidget(self.wave_canvas) + splitter.addWidget(wave_widget) + + splitter.setSizes([580, 250]) + return splitter + + def _build_right_panel(self) -> QWidget: + # Velocity settings (visible only in velocity mode) + self.grp_velocity = QGroupBox("Velocity Settings (CH1 only)") + vel_l = QVBoxLayout(self.grp_velocity) + grat_row = QHBoxLayout() + grat_row.addWidget(QLabel("Grating size:")) + self.spin_grating_um = QDoubleSpinBox() + self.spin_grating_um.setRange(0.1, 1000.0) + self.spin_grating_um.setDecimals(2) + self.spin_grating_um.setSingleStep(0.5) + self.spin_grating_um.setSuffix(" µm") + self.spin_grating_um.setValue(25) + self.spin_grating_um.setEnabled(False) + self.spin_grating_um.editingFinished.connect(self._on_grating_changed) + grat_row.addWidget(self.spin_grating_um) + vel_l.addLayout(grat_row) + lbl_formula = QLabel("v (m/s) = freq (MHz) × grating (µm)") + lbl_formula.setStyleSheet("font-size: 10px; color: #888;") + vel_l.addWidget(lbl_formula) + self.grp_velocity.setVisible(False) - # ---- Display Options group (added to right panel below) ------------ grp_display = QGroupBox("Display Options") dl = QVBoxLayout(grp_display) @@ -2215,54 +1034,19 @@ class SrasViewerWindow(QMainWindow): row.addWidget(spin) dl.addLayout(row) - # ---- Right: image + waveform splitter -------------------------- - splitter = QSplitter(Qt.Orientation.Vertical) - root.addWidget(splitter, stretch=1) - - # Image canvas - img_widget = QWidget() - img_vl = QVBoxLayout(img_widget) - img_vl.setContentsMargins(0, 0, 0, 0) - self.image_canvas = ImageCanvas() - self.image_canvas.pixel_clicked.connect(self._on_pixel_clicked) - self.image_canvas.roi_changed.connect(self._on_roi_changed) - self.image_canvas.draw_mode_changed.connect(self._on_draw_mode_changed) - toolbar = NavigationToolbar2QT(self.image_canvas, img_widget) - img_vl.addWidget(toolbar) - img_vl.addWidget(self.image_canvas) - splitter.addWidget(img_widget) - - # Waveform inspector - wave_widget = QWidget() - wave_vl = QVBoxLayout(wave_widget) - wave_vl.setContentsMargins(0, 0, 0, 0) - self.lbl_wave_hint = QLabel( - "Click a pixel in the image above to inspect its waveform." - ) - self.lbl_wave_hint.setAlignment(Qt.AlignmentFlag.AlignCenter) - self.lbl_wave_hint.setStyleSheet("color: #888; font-size: 11px;") - self.wave_canvas = WaveformCanvas() - wave_vl.addWidget(self.lbl_wave_hint) - wave_vl.addWidget(self.wave_canvas) - splitter.addWidget(wave_widget) - - splitter.setSizes([580, 250]) - - # ---- Right control panel ------------------------------------------- right_panel = QWidget() right_panel.setFixedWidth(270) - right_panel_layout = QVBoxLayout(right_panel) - right_panel_layout.setContentsMargins(0, 0, 0, 0) - right_panel_layout.setSpacing(6) - right_panel_layout.addWidget(self.grp_velocity) - right_panel_layout.addWidget(grp_display) - right_panel_layout.addStretch() - root.addWidget(right_panel) + layout = QVBoxLayout(right_panel) + layout.setContentsMargins(0, 0, 0, 0) + layout.setSpacing(6) + layout.addWidget(self.grp_velocity) + layout.addWidget(grp_display) + layout.addStretch() + return right_panel - self.statusBar().showMessage("Open an .sras file to begin.") - - # ---- Menu bar ---------------------------------------------------------- + def _build_menus(self): menubar = self.menuBar() + fft_menu = menubar.addMenu("&FFT") fft_act = QAction("FFT &Options…", self) fft_act.setStatusTip("Configure FFT backend and zero-padding") @@ -2311,43 +1095,26 @@ class SrasViewerWindow(QMainWindow): def _on_open(self): path, _ = QFileDialog.getOpenFileName( - self, "Open SRAS File", "", "SRAS Files (*.sras);;All Files (*)" - ) + self, "Open SRAS File", "", "SRAS Files (*.sras);;All Files (*)") if path: self._load_file(path) def _load_file(self, path: str): - if self._load_thread is not None: - return - # Claim self._load_thread before any call below that can pump the - # Qt event loop — see the comment in _start_compute for why. - self._load_worker = LoadWorker(path) - self._load_thread = QThread() - self._load_worker.moveToThread(self._load_thread) - self._load_thread.started.connect(self._load_worker.run) - self._load_worker.finished.connect(self._on_load_done) - self._load_worker.error.connect( - lambda msg: self.statusBar().showMessage(f"Error: {msg}") + started = self._run_worker( + "load", LoadWorker(path), + connect=( + ("finished", self._on_load_done), + ("error", lambda msg: self.statusBar().showMessage(f"Error: {msg}")), + ), ) - self._load_worker.finished.connect(self._load_thread.quit) - self._load_thread.finished.connect(self._on_load_thread_finished) - + if not started: + return self.btn_open.setEnabled(False) self.statusBar().showMessage(f"Loading {Path(path).name}…") - self._show_progress(f"Loading {Path(path).name}…") - - self._load_thread.start() - - def _on_load_thread_finished(self): - # See the comment in _on_compute_thread_finished: wait() before - # releasing our reference to avoid destroying a QThread whose OS - # thread hasn't fully joined yet. - if self._load_thread is not None: - self._load_thread.wait() - self._load_thread = None + self._show_progress("main", f"Loading {Path(path).name}…") def _on_load_done(self, sras): - self._close_progress() + self._close_progress("main") self.btn_open.setEnabled(True) if sras is None: return @@ -2355,15 +1122,14 @@ class SrasViewerWindow(QMainWindow): self._current_image = None # Caches (and any in-flight DC precompute) belong to the previous - # file's geometry — discard and start fresh. + # file's geometry — discard and start fresh. Bumping the generation + # counters makes any still-running worker's result get dropped when + # it lands. self._dc_cache = {} self._fft_cache = {} + self._dc_generation += 1 self.lbl_dc_precompute.setText("") - # Same for any alignment result — belongs to the previous file's - # geometry. Bump the generation counter so a still-running - # alignment worker's result is discarded when it lands (see - # _on_alignment_done). self._alignment_result = None self._aligned_cache = {} self._alignment_generation += 1 @@ -2391,13 +1157,8 @@ class SrasViewerWindow(QMainWindow): self.combo_channel.blockSignals(False) self._update_controls_enabled(True) - # Refresh the ADC-count label now that we have file calibration - self._on_threshold_changed() + self._on_threshold_changed() # refresh ADC label with file calibration self._on_view_changed() - - # Precompute DC images for every angle in the background so - # switching angles while viewing a DC channel is instant, and the - # FFT masking step rarely has to wait on a DC4 image either. self._start_dc_precompute() # ------------------------------------------------------------------ @@ -2408,44 +1169,41 @@ class SrasViewerWindow(QMainWindow): s = self._sras if s is None: return - angle_idx = self.spin_angle.value() - self._info["Angles"].setText(f"Angles: {s.n_angles}") - self._info["Rows"].setText(f"Rows: {s.n_rows[angle_idx]}") - self._info["Frames / row"].setText(f"Frames / row: {s.n_frames[angle_idx]}") - self._info["Samples / frame"].setText(f"Samples / frame: {s.samples_per_frame}") - self._info["Sample rate"].setText(f"Sample rate: {s.sample_rate_hz/1e9:.4g} GS/s") - self._info["X start"].setText(f"X start: {s.x_start_mm[angle_idx]:.4g} mm") - self._info["Pixel Δx"].setText(f"Pixel Δx: {s.pixel_x_mm*1e3:.3g} µm") - self._info["Laser freq"].setText(f"Laser freq: {s.laser_freq_hz/1e3:.4g} kHz") + a = self.spin_angle.value() + for key, text in ( + ("Angles", f"{s.n_angles}"), + ("Rows", f"{s.n_rows[a]}"), + ("Frames / row", f"{s.n_frames[a]}"), + ("Samples / frame", f"{s.samples_per_frame}"), + ("Sample rate", f"{s.sample_rate_hz / 1e9:.4g} GS/s"), + ("X start", f"{s.x_start_mm[a]:.4g} mm"), + ("Pixel Δx", f"{s.pixel_x_mm * 1e3:.3g} µm"), + ("Laser freq", f"{s.laser_freq_hz / 1e3:.4g} kHz"), + ): + self._info[key].setText(f"{key}: {text}") notes = [] if s.frame_count_mismatch: - notes.append( - f"! Header n_frames={s.n_frames_header}, " - f"actual={s.n_frames[angle_idx]} (scanner bug — corrected)" - ) + notes.append(f"! Header n_frames={s.n_frames_header}, " + f"actual={s.n_frames[a]} (scanner bug — corrected)") if s.scan_aborted: - notes.append( - f"! Scan aborted: {s.n_angles}/{s.n_angles_declared} angles complete" - ) + notes.append(f"! Scan aborted: {s.n_angles}/{s.n_angles_declared} " + "angles complete") if s.background is not None: notes.append(f"Background waveform: {len(s.background)} samples") - if isinstance(s.precomputed_freq_mhz, np.ndarray): - bg_note = " (bg-sub)" if s.precomputed_bg_sub else " (no bg-sub)" - notes.append(f"v5: precomputed images present{bg_note} — display is instant") if s.version in (6, 7): notes.append("v6/v7 format: rows / frames / x_start are per-angle") - if s.version == 7: - n_dc = sum(1 for x in s.precomputed_dc4_mv if x is not None) - n_fft = sum(1 for x in s.precomputed_freq_mhz if x is not None) - if n_dc or n_fft: - bg_note = " (bg-sub)" if s.precomputed_bg_sub else " (no bg-sub)" - notes.append( - f"v7: cached DC ({n_dc}/{s.n_angles} angles), " - f"FFT ({n_fft}/{s.n_angles} angles{bg_note if n_fft else ''}) " - "— display is instant for cached angles") - else: - notes.append("v7 format: no cache blocks stored yet") + + n_dc = sum(1 for x in s.precomputed_dc4_mv if x is not None) + n_fft = sum(1 for x in s.precomputed_freq_mhz if x is not None) + if n_dc or n_fft: + bg_note = " (bg-sub)" if s.precomputed_bg_sub else " (no bg-sub)" + notes.append( + f"Cached images: DC {n_dc}/{s.n_angles} angles, " + f"FFT {n_fft}/{s.n_angles} angles{bg_note if n_fft else ''} " + "— display is instant for cached angles") + elif s.version == 7: + notes.append("v7 format: no cache blocks stored yet") self.lbl_frame_warn.setText("\n".join(notes)) # ------------------------------------------------------------------ @@ -2454,61 +1212,50 @@ class SrasViewerWindow(QMainWindow): def _update_controls_enabled(self, enabled: bool): s = self._sras - self.spin_angle.setEnabled(enabled and s is not None and s.n_angles > 1) + has_file = enabled and s is not None + ch_idx = self.combo_channel.currentIndex() + is_ch1 = enabled and ch_idx in CH1_DERIVED_MODES + is_vel = enabled and ch_idx == VELOCITY_MODE_IDX + + self.spin_angle.setEnabled(has_file and s.n_angles > 1) self.combo_channel.setEnabled(enabled) self.combo_cmap.setEnabled(enabled) self.chk_auto.setEnabled(enabled) manual = enabled and not self.chk_auto.isChecked() self.spin_vmin.setEnabled(manual) self.spin_vmax.setEnabled(manual) - ch_idx = self.combo_channel.currentIndex() - is_ch1 = enabled and ch_idx in CH1_DERIVED_MODES + # Threshold and bg-sub apply to all CH1 modes self.spin_threshold_mv.setEnabled(is_ch1) - has_bg = enabled and s is not None and s.background is not None - self.chk_bg_sub.setEnabled(has_bg and is_ch1) - # Velocity grating spinbox - is_vel = enabled and ch_idx == VELOCITY_MODE_IDX + self.chk_bg_sub.setEnabled(has_file and s.background is not None and is_ch1) self.spin_grating_um.setEnabled(is_vel) self.grp_velocity.setVisible(is_vel) - # CSV export: enabled when a CH1-derived image is displayed + self.btn_export_csv.setEnabled(is_ch1 and self._current_image is not None) # ROI: always usable once a file is loaded (independent of channel) - self.btn_draw_roi.setEnabled(enabled and s is not None) - # Batch Convert actions: pick their own files, independent of - # whatever's currently open — only gated on no batch already running - can_batch = self._batch_thread is None + self.btn_draw_roi.setEnabled(has_file) + + # Batch Convert actions pick their own files, independent of + # whatever's currently open — only gated on no batch already running. + can_batch = not self._job_running("batch") self._batch_dc_act.setEnabled(can_batch) self._batch_fft_act.setEnabled(can_batch) - # Angle alignment: needs >1 angle and no alignment already running - can_align = (enabled and s is not None and s.n_angles > 1 - and self._alignment_thread is None) - self._alignment_act.setEnabled(can_align) - self.chk_aligned_view.setEnabled( - enabled and self._alignment_result is not None) + + self._alignment_act.setEnabled( + has_file and s.n_angles > 1 and not self._job_running("align")) + self.chk_aligned_view.setEnabled(enabled and self._alignment_result is not None) self._update_roi_ui() def _on_channel_changed(self): - ch_idx = self.combo_channel.currentIndex() - has_file = self._sras is not None - is_ch1 = ch_idx in CH1_DERIVED_MODES - self.spin_threshold_mv.setEnabled(is_ch1 and has_file) - has_bg = has_file and self._sras.background is not None - self.chk_bg_sub.setEnabled(has_bg and is_ch1) - is_vel = ch_idx == VELOCITY_MODE_IDX - self.spin_grating_um.setEnabled(is_vel and has_file) - self.grp_velocity.setVisible(is_vel) - self.btn_export_csv.setEnabled(is_ch1 and has_file and self._current_image is not None) + self._update_controls_enabled(self._sras is not None) self._on_view_changed() def _on_bg_sub_toggled(self): # Background subtraction changes the FFT input, so it genuinely # invalidates the cached raw FFT (the cache key includes it) — - # _refresh_display() will recompute only if there's no cache hit - # for the new bg-sub state. - if self._sras is not None: - if self.combo_channel.currentIndex() in CH1_DERIVED_MODES: - self._refresh_display() + # _refresh_display() recomputes only on a miss for the new state. + if self._sras is not None and self.combo_channel.currentIndex() in CH1_DERIVED_MODES: + self._refresh_display() def _on_grating_changed(self): # Grating is a pure post-multiply on the cached frequency image — @@ -2516,27 +1263,119 @@ class SrasViewerWindow(QMainWindow): if self._sras is not None and self.combo_channel.currentIndex() == VELOCITY_MODE_IDX: self._refresh_display() + def _on_threshold_changed(self): + mv = self.spin_threshold_mv.value() + cal = (self._sras.cal(CH4_IDX) if self._sras is not None + else (_FALLBACK_YMULT_MV, _FALLBACK_YOFF_ADC, 0.0)) + self.lbl_threshold_adc.setText(f"≈ {mv_to_adc(mv, *cal):.1f} ADC counts") + # Threshold decides which pixels get an FFT at all, so changing it is + # a genuine cache-key change — but the recompute reuses the cached DC4 + # image to skip masked-out pixels. + if self._sras is not None and self.combo_channel.currentIndex() in CH1_DERIVED_MODES: + self._refresh_display() + + def _on_autoscale_toggled(self, checked: bool): + manual = not checked + self.spin_vmin.setEnabled(manual and self._sras is not None) + self.spin_vmax.setEnabled(manual and self._sras is not None) + if self._sras is not None and self._current_image is not None: + self._redraw_image(self._current_image) + + def _on_manual_range_changed(self): + if not self.chk_auto.isChecked() and self._current_image is not None: + self._redraw_image(self._current_image) + + def _on_cmap_changed(self): + # Colormap is purely how the existing image is rendered. + if self._current_image is not None: + self._redraw_image(self._current_image) + + def _on_view_changed(self): + if self._sras is None: + return + idx = self.spin_angle.value() + self.lbl_angle_deg.setText(f"({self._sras.angles_deg[idx]:.1f}°)") + self._update_scan_info_labels() + self._refresh_display() + + def _on_aligned_view_toggled(self, checked: bool): + if self._current_image is not None: + self._redraw_image(self._current_image) + + # ------------------------------------------------------------------ + # CSV export + # ------------------------------------------------------------------ + def _on_export_csv(self): if self._current_image is None or self._sras is None: return - ch_idx = self._current_ch - angle = self._current_angle - ch_name = CH_NAMES[ch_idx] - default_name = ( - f"{self._sras.path.stem}_angle{angle}_{ch_name}.csv" - ) + default_name = (f"{self._sras.path.stem}_angle{self._current_angle}" + f"_{CH_NAMES[self._current_ch]}.csv") path, _ = QFileDialog.getSaveFileName( self, "Export Image as CSV", str(self._sras.path.parent / default_name), - "CSV files (*.csv);;All files (*)", - ) + "CSV files (*.csv);;All files (*)") if not path: return np.savetxt(path, self._current_image, delimiter=",", fmt="%.6g") self.statusBar().showMessage(f"Exported {Path(path).name}") + def _on_export_roi_csv(self): + if self._current_image is None or self._sras is None: + return + roi = self.image_canvas.get_roi() + if roi is None: + self.statusBar().showMessage("No ROI — draw one first") + return + s = self._sras + x_axis = s.x_axis_mm(self._current_angle) + y_axis = s.y_positions_mm(self._current_angle) + mask = roi.mask_for_grid(x_axis, y_axis) + if not mask.any(): + self.statusBar().showMessage("ROI does not overlap any pixel") + return + img = self._current_image + if img.shape != mask.shape: + self.statusBar().showMessage( + f"ROI shape {mask.shape} does not match image {img.shape}") + return + + X, Y = np.meshgrid(np.asarray(x_axis, dtype=np.float64), + np.asarray(y_axis, dtype=np.float64)) + rows_idx, frames_idx = np.where(mask) + n_pix = int(mask.sum()) + + ch_name = CH_NAMES[self._current_ch] + angle = self._current_angle + default_name = f"{s.path.stem}_angle{angle}_{ch_name}_ROI.csv" + path, _ = QFileDialog.getSaveFileName( + self, "Export ROI as CSV", + str(s.path.parent / default_name), + "CSV files (*.csv);;All files (*)") + if not path: + return + + corners_str = " ".join(f"({p[0]:.6g},{p[1]:.6g})" for p in roi.corners()) + header = ( + f"# ROI quad corners (BL BR TR TL) mm: {corners_str}\n" + f"# source: {s.path.name}, channel={ch_name}, " + f"angle_idx={angle}, angle_deg={s.angles_deg[angle]:.4g}\n" + f"# n_pixels={n_pix}\n" + "row,frame,x_mm,y_mm,value" + ) + data = np.column_stack([ + rows_idx.astype(np.int64), frames_idx.astype(np.int64), + X[mask], Y[mask], img[mask].astype(np.float64), + ]) + # integer columns first, floats after — use a per-column format list + np.savetxt(path, data, delimiter=",", + fmt=["%d", "%d", "%.6g", "%.6g", "%.6g"], + header=header, comments="") + self.statusBar().showMessage( + f"Exported ROI ({n_pix} pixels) to {Path(path).name}") + # ------------------------------------------------------------------ - # ROI (rectangle on the image) + # ROI # ------------------------------------------------------------------ def _on_draw_roi_toggled(self, checked: bool): @@ -2553,8 +1392,9 @@ class SrasViewerWindow(QMainWindow): self.btn_draw_roi.setChecked(active) self.btn_draw_roi.blockSignals(False) - def _on_roi_changed(self): - self._update_roi_ui() + def _on_clear_roi(self): + self.image_canvas.clear_roi() + self.statusBar().showMessage("ROI cleared") def _update_roi_ui(self): roi = self.image_canvas.get_roi() @@ -2565,12 +1405,12 @@ class SrasViewerWindow(QMainWindow): self.btn_clear_roi.setEnabled(False) self.btn_export_roi.setEnabled(False) return - cen = roi.centroid() + + cen = roi.centroid() bbox = roi.bbox_size() - self.lbl_roi_center.setText( - f"centroid: ({cen[0]:.3f}, {cen[1]:.3f}) mm") - self.lbl_roi_size.setText( - f"bbox: {bbox[0]:.3f} × {bbox[1]:.3f} mm") + self.lbl_roi_center.setText(f"centroid: ({cen[0]:.3f}, {cen[1]:.3f}) mm") + self.lbl_roi_size.setText(f"bbox: {bbox[0]:.3f} × {bbox[1]:.3f} mm") + npix = 0 if self._sras is not None: try: @@ -2578,133 +1418,24 @@ class SrasViewerWindow(QMainWindow): # Aligned View is on: _on_export_roi_csv also exports on the # raw grid (never synthetically-resampled pixels), so this # readout must match what Export ROI actually writes. - mask = roi.mask_for_grid(self._sras.x_axis_mm(self._current_angle), - self._sras.y_positions_mm(self._current_angle)) + mask = roi.mask_for_grid( + self._sras.x_axis_mm(self._current_angle), + self._sras.y_positions_mm(self._current_angle)) npix = int(mask.sum()) except Exception: npix = 0 self.lbl_roi_npix.setText(f"pixels inside: {npix}") self.btn_clear_roi.setEnabled(True) - self.btn_export_roi.setEnabled( - self._current_image is not None and npix > 0) - - def _on_clear_roi(self): - self.image_canvas.clear_roi() - self.statusBar().showMessage("ROI cleared") - - def _on_roi_angle_edited(self): - pass # rotation control removed — corners are dragged individually - - def _on_export_roi_csv(self): - if self._current_image is None or self._sras is None: - return - roi = self.image_canvas.get_roi() - if roi is None: - self.statusBar().showMessage("No ROI — draw one first") - return - s = self._sras - x_axis = s.x_axis_mm(self._current_angle) - y_axis = s.y_positions_mm(self._current_angle) - X, Y = np.meshgrid(np.asarray(x_axis, dtype=np.float64), - np.asarray(y_axis, dtype=np.float64)) - mask = roi.mask_for_grid(x_axis, y_axis) - if not mask.any(): - self.statusBar().showMessage("ROI does not overlap any pixel") - return - img = self._current_image - if img.shape != mask.shape: - self.statusBar().showMessage( - f"ROI shape {mask.shape} does not match image {img.shape}") - return - - rows_idx, frames_idx = np.where(mask) - xs = X[mask] - ys = Y[mask] - vals = img[mask] - - ch_idx = self._current_ch - ch_name = CH_NAMES[ch_idx] - angle = self._current_angle - default_name = (f"{s.path.stem}_angle{angle}_{ch_name}_ROI.csv") - path, _ = QFileDialog.getSaveFileName( - self, "Export ROI as CSV", - str(s.path.parent / default_name), - "CSV files (*.csv);;All files (*)", - ) - if not path: - return - - pts = roi.corners() - corners_str = " ".join(f"({p[0]:.6g},{p[1]:.6g})" for p in pts) - header = ( - f"# ROI quad corners (BL BR TR TL) mm: {corners_str}\n" - f"# source: {s.path.name}, channel={ch_name}, " - f"angle_idx={angle}, angle_deg={s.angles_deg[angle]:.4g}\n" - f"# n_pixels={int(mask.sum())}\n" - "row,frame,x_mm,y_mm,value" - ) - data = np.column_stack([ - rows_idx.astype(np.int64), - frames_idx.astype(np.int64), - xs, ys, vals.astype(np.float64), - ]) - # integer columns first, floats after — use a per-column format list - np.savetxt(path, data, delimiter=",", - fmt=["%d", "%d", "%.6g", "%.6g", "%.6g"], - header=header, comments="") - self.statusBar().showMessage( - f"Exported ROI ({int(mask.sum())} pixels) to {Path(path).name}") - - def _on_threshold_changed(self): - mv = self.spin_threshold_mv.value() - if self._sras is not None: - ymult = self._sras.ch_ymult_mv[CH4_IDX] - yoff = self._sras.ch_yoff_adc[CH4_IDX] - yzero = self._sras.ch_yzero_mv[CH4_IDX] - else: - ymult, yoff, yzero = _FALLBACK_YMULT_MV, _FALLBACK_YOFF_ADC, 0.0 - self.lbl_threshold_adc.setText(f"≈ {mv_to_adc(mv, ymult, yoff, yzero):.1f} ADC counts") - # Threshold decides which pixels get an FFT at all (see - # ComputeWorker), so changing it is a genuine cache key change — - # _refresh_display() recomputes only on a miss, and that recompute - # reuses the cached DC4 image to skip masked-out pixels. - if self._sras is not None and self.combo_channel.currentIndex() in CH1_DERIVED_MODES: - self._refresh_display() - - def _on_autoscale_toggled(self, checked: bool): - manual = not checked - self.spin_vmin.setEnabled(manual and self._sras is not None) - self.spin_vmax.setEnabled(manual and self._sras is not None) - if self._sras is not None and self._current_image is not None: - self._redraw_image(self._current_image) - - def _on_manual_range_changed(self): - if not self.chk_auto.isChecked() and self._current_image is not None: - self._redraw_image(self._current_image) - - def _on_cmap_changed(self): - # Colormap is purely how the existing image is rendered — never - # needs a recompute. - if self._current_image is not None: - self._redraw_image(self._current_image) - - def _on_view_changed(self): - if self._sras is None: - return - idx = self.spin_angle.value() - self.lbl_angle_deg.setText(f"({self._sras.angles_deg[idx]:.1f}°)") - self._update_scan_info_labels() - self._refresh_display() + self.btn_export_roi.setEnabled(self._current_image is not None and npix > 0) # ------------------------------------------------------------------ - # Computation + # Display # ------------------------------------------------------------------ def _current_n_fft(self) -> int | None: - pad_factor = self._fft_pad_factor - if pad_factor <= 1 or self._sras is None: + if self._fft_pad_factor <= 1 or self._sras is None: return None - return self._sras.samples_per_frame * pad_factor + return self._sras.samples_per_frame * self._fft_pad_factor def _scale_for_display(self, freq_mhz: np.ndarray, ch_idx: int) -> np.ndarray: """Velocity is a pure post-multiply of the (already DC-masked) @@ -2713,28 +1444,26 @@ class SrasViewerWindow(QMainWindow): return freq_mhz * self.spin_grating_um.value() return freq_mhz - # ------------------------------------------------------------------ - # Aligned View (Fusion) display helpers - # ------------------------------------------------------------------ + def _fft_cache_key(self, angle_idx: int) -> tuple: + return (angle_idx, self.chk_bg_sub.isChecked(), self._current_n_fft(), + self.spin_threshold_mv.value()) def _aligned_cache_key(self, angle_idx: int, ch_idx: int) -> tuple: """Mirrors _fft_cache's key granularity so a stale aligned image is never shown after bg_sub/threshold/pad/grating changes.""" if ch_idx in CH1_DERIVED_MODES: - return (angle_idx, ch_idx, self.chk_bg_sub.isChecked(), - self._current_n_fft(), self.spin_threshold_mv.value(), + return (*self._fft_cache_key(angle_idx), ch_idx, self.spin_grating_um.value() if ch_idx == VELOCITY_MODE_IDX else None) return (angle_idx, ch_idx) def _aligned_canvas_axes(self) -> tuple[np.ndarray, np.ndarray]: r = self._alignment_result n_rows, n_cols = r.canvas_shape - x_axis = r.canvas_origin_mm[0] + np.arange(n_cols) * r.canvas_dx_mm - y_axis = r.canvas_origin_mm[1] + np.arange(n_rows) * r.canvas_dy_mm - return x_axis, y_axis + return (r.canvas_origin_mm[0] + np.arange(n_cols) * r.canvas_dx_mm, + r.canvas_origin_mm[1] + np.arange(n_rows) * r.canvas_dy_mm) def _get_aligned_display_image(self, raw_img: np.ndarray, angle_idx: int, - ch_idx: int) -> np.ndarray: + ch_idx: int) -> np.ndarray: key = self._aligned_cache_key(angle_idx, ch_idx) cached = self._aligned_cache.get(key) if cached is None: @@ -2744,138 +1473,164 @@ class SrasViewerWindow(QMainWindow): def _refresh_display(self): """Show the image for the current angle/channel/threshold, using - cached data whenever possible and only falling back to a - background compute (with progress popup) when genuinely nothing - is cached yet for these settings.""" + cached data whenever possible and only falling back to a background + compute (with progress popup) when genuinely nothing is cached yet.""" if self._sras is None: return angle_idx = self.spin_angle.value() - ch_idx = self.combo_channel.currentIndex() + ch_idx = self.combo_channel.currentIndex() - if ch_idx in (CH3_IDX, CH4_IDX): + if ch_idx in CH1_DERIVED_MODES: + raw = self._fft_cache.get(self._fft_cache_key(angle_idx)) + if raw is not None: + self._show_image_now(self._scale_for_display(raw, ch_idx), + angle_idx, ch_idx) + return + else: cached = self._dc_cache.get((angle_idx, ch_idx)) if cached is not None: self._show_image_now(cached, angle_idx, ch_idx) return - elif ch_idx in (CH1_IDX, VELOCITY_MODE_IDX): - apply_bg_sub = self.chk_bg_sub.isChecked() - threshold_mv = self.spin_threshold_mv.value() - key = (angle_idx, apply_bg_sub, self._current_n_fft(), threshold_mv) - raw = self._fft_cache.get(key) - if raw is not None: - img = self._scale_for_display(raw, ch_idx) - self._show_image_now(img, angle_idx, ch_idx) - return # Nothing cached for these settings — need a real compute. Changing # the DC threshold changes *which* pixels get an FFT at all, so it # can't be satisfied from the cache — but with the DC map already - # known, the recompute skips the FFT for masked-out pixels and is - # much cheaper than a full-image compute would be. + # known, the recompute skips the FFT for masked-out pixels. self._start_compute() def _show_image_now(self, img: np.ndarray, angle_idx: int, ch_idx: int): """Display an already-available image with no compute involved.""" self._current_image = img self._current_angle = angle_idx - self._current_ch = ch_idx + self._current_ch = ch_idx self.btn_export_csv.setEnabled(ch_idx in CH1_DERIVED_MODES) self._redraw_image(img) self._update_roi_ui() + def _redraw_image(self, img: np.ndarray): + s = self._sras + angle_idx = self._current_angle + ch_idx = self._current_ch + + aligned = (self.chk_aligned_view.isChecked() + and self._alignment_result is not None + and angle_idx in self._alignment_result.per_angle) + if aligned: + display_img = self._get_aligned_display_image(img, angle_idx, ch_idx) + x_axis, y_axis = self._aligned_canvas_axes() + else: + display_img = img + x_axis = s.x_axis_mm(angle_idx) + y_axis = s.y_positions_mm(angle_idx) + + dx = x_axis[1] - x_axis[0] if len(x_axis) > 1 else s.pixel_x_mm + dy = float(y_axis[1] - y_axis[0]) if len(y_axis) > 1 else 1.0 + extent = [x_axis[0] - dx / 2, x_axis[-1] + dx / 2, + y_axis[-1] + dy / 2, y_axis[0] - dy / 2] + + if self.chk_auto.isChecked(): + vmin, vmax = float(display_img.min()), float(display_img.max()) + for spin, val in ((self.spin_vmin, vmin), (self.spin_vmax, vmax)): + spin.blockSignals(True) + spin.setValue(val) + spin.blockSignals(False) + else: + vmin, vmax = self.spin_vmin.value(), self.spin_vmax.value() + + angle_deg = s.angles_deg[angle_idx] + mode_str, unit, colorbar_label = _CHANNEL_DISPLAY[ch_idx] + if ch_idx == VELOCITY_MODE_IDX: + ch_label = f"Velocity [grating={self.spin_grating_um.value():.2f} µm]" + else: + ch_label = CH_LABELS[ch_idx] + + title = f"{CH_NAMES[ch_idx]} | {mode_str} | {angle_deg:.1f}°" + if aligned: + title += " [Aligned]" + + self.image_canvas.show_image( + display_img, extent, + cmap=self.combo_cmap.currentText(), + vmin=vmin, vmax=vmax, + xlabel="X (mm)", ylabel="Y (mm)", + title=title, colorbar_label=colorbar_label, + ) + self.statusBar().showMessage( + f"{s.path.name} | {ch_label} @ {angle_deg:.1f}° " + f"| {display_img.shape[1]} × {display_img.shape[0]} px | {unit}" + f"{' | Aligned' if aligned else ''}" + ) + # ------------------------------------------------------------------ # Background compute (only reached on a genuine cache miss) # ------------------------------------------------------------------ def _start_compute(self): - if self._sras is None: + if self._sras is None or self._job_running("compute"): + return # re-checked when the running compute finishes + + angle_idx = self.spin_angle.value() + ch_idx = self.combo_channel.currentIndex() + is_fft = ch_idx in CH1_DERIVED_MODES + + self._pending_angle = angle_idx + self._pending_ch = ch_idx + self._pending_bg_sub = self.chk_bg_sub.isChecked() + self._pending_threshold = self.spin_threshold_mv.value() + self._pending_fft_pad_factor = self._fft_pad_factor + + worker = ComputeWorker( + self._sras, angle_idx, ch_idx, + apply_bg_sub=self._pending_bg_sub, + n_fft=self._current_n_fft(), + dc_threshold_mv=self._pending_threshold, + # Reuse the cached DC4 image (if the precompute has reached this + # angle) so the FFT skips masked-out pixels entirely and doesn't + # need to re-read the CH4 channel from disk. + dc4_mv=self._dc_cache.get((angle_idx, CH4_IDX)), + is_fft_mode=is_fft, + ) + if not self._run_worker( + "compute", worker, + connect=( + ("finished", self._on_compute_done), + ("error", lambda msg: self.statusBar().showMessage( + f"Compute error: {msg}")), + ), + on_done=self._after_compute): return - if self._compute_thread is not None: - return # re-check in _on_compute_thread_finished - angle_idx = self.spin_angle.value() - ch_idx = self.combo_channel.currentIndex() - apply_bg_sub = self.chk_bg_sub.isChecked() - threshold_mv = self.spin_threshold_mv.value() - pad_factor = self._fft_pad_factor - n_fft = self._current_n_fft() - # Reuse the cached DC4 image (if the precompute has reached this - # angle) so the FFT compute skips masked-out pixels entirely and - # doesn't need to re-read the CH4 channel from disk. - dc4_mv = self._dc_cache.get((angle_idx, CH4_IDX)) - - self._pending_angle = angle_idx - self._pending_ch = ch_idx - self._pending_bg_sub = apply_bg_sub - self._pending_threshold = threshold_mv - self._pending_fft_pad_factor = pad_factor - - # Claim self._compute_thread *before* anything below that can pump - # the Qt event loop (e.g. QProgressDialog.show() on first display). - # If that happened first, a re-entrant editingFinished/signal could - # slip past the guard above, start a second thread, and then have - # this call's own assignment clobber (and destroy while still - # running) that second thread's QThread object — which aborts the - # process. Assigning immediately closes that window. - self._compute_worker = ComputeWorker( - self._sras, angle_idx, ch_idx, apply_bg_sub, n_fft=n_fft, - dc_threshold_mv=threshold_mv, dc4_mv=dc4_mv, - ) - self._compute_thread = QThread() - self._compute_worker.moveToThread(self._compute_thread) - self._compute_thread.started.connect(self._compute_worker.run) - self._compute_worker.finished.connect(self._on_compute_done) - self._compute_worker.error.connect( - lambda msg: self.statusBar().showMessage(f"Compute error: {msg}") - ) - self._compute_worker.finished.connect(self._compute_thread.quit) - self._compute_thread.finished.connect(self._on_compute_thread_finished) - - if ch_idx in (CH1_IDX, VELOCITY_MODE_IDX): + if is_fft: self.statusBar().showMessage("Computing FFT…") self._show_progress( + "main", f"Computing FFT for angle {angle_idx}…\n" "This can take a while on a large scan — result is cached " - "so revisiting this angle/mode/threshold will be instant." - ) + "so revisiting this angle/mode/threshold will be instant.") else: self.statusBar().showMessage("Computing DC image…") - self._show_progress(f"Computing DC image for angle {angle_idx}…") + self._show_progress("main", f"Computing DC image for angle {angle_idx}…") - self._compute_thread.start() - - def _on_compute_thread_finished(self): - # Block until the OS thread has actually joined before dropping our - # last reference — deallocating a QThread whose thread hasn't fully - # terminated yet logs "QThread: Destroyed while thread is still - # running" and aborts the process. The finished() signal fires as - # the thread is winding down but does not guarantee it has joined. - if self._compute_thread is not None: - self._compute_thread.wait() - self._compute_thread = None - angle_idx = self.spin_angle.value() - ch_idx = self.combo_channel.currentIndex() - apply_bg_sub = self.chk_bg_sub.isChecked() - threshold_mv = self.spin_threshold_mv.value() - if (angle_idx, ch_idx, apply_bg_sub, threshold_mv, self._fft_pad_factor) != ( + def _after_compute(self): + """If settings changed while the compute was running, re-dispatch + through the cache-aware path — the now-current combination may + already be cached.""" + if (self.spin_angle.value(), self.combo_channel.currentIndex(), + self.chk_bg_sub.isChecked(), self.spin_threshold_mv.value(), + self._fft_pad_factor) != ( self._pending_angle, self._pending_ch, self._pending_bg_sub, self._pending_threshold, self._pending_fft_pad_factor): - # Settings changed while this compute was running — re-dispatch - # through the cache-aware path in case that now-current - # combination happens to already be cached. self._refresh_display() def _on_compute_done(self, result): - self._close_progress() + self._close_progress("main") angle_idx = self._pending_angle - ch_idx = self._pending_ch + ch_idx = self._pending_ch - if ch_idx in (CH1_IDX, VELOCITY_MODE_IDX): - masked_freq_mhz = result - key = (angle_idx, self._pending_bg_sub, self._current_n_fft(), - self._pending_threshold) - self._fft_cache[key] = masked_freq_mhz - img = self._scale_for_display(masked_freq_mhz, ch_idx) + if ch_idx in CH1_DERIVED_MODES: + self._fft_cache[(angle_idx, self._pending_bg_sub, + self._current_n_fft(), self._pending_threshold)] = result + img = self._scale_for_display(result, ch_idx) else: img = result self._dc_cache[(angle_idx, ch_idx)] = img @@ -2889,135 +1644,48 @@ class SrasViewerWindow(QMainWindow): def _start_dc_precompute(self): if self._sras is None: return - self._dc_generation += 1 generation = self._dc_generation - n_angles = self._sras.n_angles + n_angles = self._sras.n_angles worker = DcPrecomputeWorker(self._sras) - thread = QThread() - worker.moveToThread(thread) - thread.started.connect(worker.run) - worker.angle_done.connect( - lambda a, dc3, dc4, g=generation: self._on_dc_precompute_angle_done( - g, a, dc3, dc4, n_angles) - ) - worker.error.connect( - lambda msg: self.statusBar().showMessage(f"DC precompute error: {msg}", 5000) - ) - worker.finished.connect(thread.quit) - worker.error.connect(thread.quit) - thread.finished.connect(lambda: self._on_dc_precompute_thread_finished(thread)) - - # See _start_compute for why the thread/worker are claimed on self - # before .start() rather than after. self._dc_precompute_worker = worker - self._dc_precompute_thread = thread - thread.start() + started = self._run_worker( + "dc_precompute", worker, + connect=( + ("angle_done", lambda a, dc3, dc4, g=generation: + self._on_dc_precompute_angle_done(g, a, dc3, dc4, n_angles)), + ("error", lambda msg: self.statusBar().showMessage( + f"DC precompute error: {msg}", 5000)), + ), + quit_on=("finished", "error"), + on_done=lambda: setattr(self, "_dc_precompute_worker", None), + ) + if not started: + self._dc_precompute_worker = None def _on_dc_precompute_angle_done(self, generation: int, angle_idx: int, - dc3_mv: np.ndarray, dc4_mv: np.ndarray, - n_angles: int): + dc3_mv: np.ndarray, dc4_mv: np.ndarray, + n_angles: int): if generation != self._dc_generation: return # stale result from a previously-loaded file — discard self._dc_cache[(angle_idx, CH3_IDX)] = dc3_mv self._dc_cache[(angle_idx, CH4_IDX)] = dc4_mv + done = sum(1 for a in range(n_angles) if (a, CH4_IDX) in self._dc_cache) - if done < n_angles: - self.lbl_dc_precompute.setText( - f"Precomputing DC images: {done}/{n_angles} angles ready…") - else: - self.lbl_dc_precompute.setText("DC images ready for all angles.") - # If we just finished the angle/channel the user is currently - # looking at and it wasn't shown yet (e.g. they switched here - # before precompute caught up and are still waiting), show it now. - if (self._sras is not None and angle_idx == self.spin_angle.value() - and self._compute_thread is None - and self.combo_channel.currentIndex() in (CH3_IDX, CH4_IDX) - and (self._current_angle != angle_idx - or self._current_ch != self.combo_channel.currentIndex())): + self.lbl_dc_precompute.setText( + f"Precomputing DC images: {done}/{n_angles} angles ready…" + if done < n_angles else "DC images ready for all angles.") + + # If we just finished the angle/channel the user is currently looking + # at and it wasn't shown yet (they switched here before the precompute + # caught up and are still waiting), show it now. + current_ch = self.combo_channel.currentIndex() + if (angle_idx == self.spin_angle.value() + and not self._job_running("compute") + and current_ch in (CH3_IDX, CH4_IDX) + and (self._current_angle != angle_idx or self._current_ch != current_ch)): self._refresh_display() - def _on_dc_precompute_thread_finished(self, thread: QThread): - # See the comment in _on_compute_thread_finished: wait() before - # releasing the reference to avoid destroying a QThread whose OS - # thread hasn't fully joined yet. - thread.wait() - if self._dc_precompute_thread is thread: - self._dc_precompute_thread = None - self._dc_precompute_worker = None - - def _redraw_image(self, img: np.ndarray): - s = self._sras - angle_idx = self._current_angle - ch_idx = self._current_ch - - aligned = (self.chk_aligned_view.isChecked() - and self._alignment_result is not None - and angle_idx in self._alignment_result.per_angle) - if aligned: - display_img = self._get_aligned_display_image(img, angle_idx, ch_idx) - x_axis, y_axis = self._aligned_canvas_axes() - else: - display_img = img - x_axis = s.x_axis_mm(angle_idx) - y_axis = s.y_positions_mm(angle_idx) - - dx = x_axis[1] - x_axis[0] if len(x_axis) > 1 else s.pixel_x_mm - dy = float(y_axis[1] - y_axis[0]) if len(y_axis) > 1 else 1.0 - - extent = [ - x_axis[0] - dx / 2, - x_axis[-1] + dx / 2, - y_axis[-1] + dy / 2, - y_axis[0] - dy / 2, - ] - - if self.chk_auto.isChecked(): - vmin, vmax = float(display_img.min()), float(display_img.max()) - for spin, val in ((self.spin_vmin, vmin), (self.spin_vmax, vmax)): - spin.blockSignals(True) - spin.setValue(val) - spin.blockSignals(False) - else: - vmin = self.spin_vmin.value() - vmax = self.spin_vmax.value() - - angle_deg = s.angles_deg[self._current_angle] - ch_label = CH_LABELS[ch_idx] - - if ch_idx == CH1_IDX: - mode_str = "RF" - unit = "Peak frequency (MHz)" - colorbar_label = "MHz" - elif ch_idx == VELOCITY_MODE_IDX: - grating = self.spin_grating_um.value() - mode_str = "Velocity" - unit = "Velocity (m/s)" - colorbar_label = "m/s" - ch_label = f"Velocity [grating={grating:.2f} µm]" - else: - mode_str = "DC" - unit = "DC mean (mV)" - colorbar_label = "mV" - - title = f"{CH_NAMES[ch_idx]} | {mode_str} | {angle_deg:.1f}°" - if aligned: - title += " [Aligned]" - - self.image_canvas.show_image( - display_img, extent, - cmap=self.combo_cmap.currentText(), - vmin=vmin, vmax=vmax, - xlabel="X (mm)", ylabel="Y (mm)", - title=title, - colorbar_label=colorbar_label, - ) - self.statusBar().showMessage( - f"{s.path.name} | {ch_label} @ {angle_deg:.1f}° " - f"| {display_img.shape[1]} × {display_img.shape[0]} px | {unit}" - f"{' | Aligned' if aligned else ''}" - ) - # ------------------------------------------------------------------ # Pixel inspector # ------------------------------------------------------------------ @@ -3028,110 +1696,97 @@ class SrasViewerWindow(QMainWindow): angle_idx = self._current_angle if (self.chk_aligned_view.isChecked() and self._alignment_result is not None and angle_idx in self._alignment_result.per_angle): - # The click landed on the shared aligned canvas — invert the - # same canvas->raw affine used to display it back to a raw - # (row, frame) index before looking up the waveform. + # The click landed on the shared aligned canvas — invert the same + # canvas->raw affine used to display it back to a raw (row, frame) + # index before looking up the waveform. t = self._alignment_result.per_angle[angle_idx] raw = t.matrix @ np.array([row_idx, frame_idx], dtype=np.float64) + t.offset row_idx, frame_idx = int(round(raw[0])), int(round(raw[1])) - n_rows_a = int(self._sras.n_rows[angle_idx]) - n_frames_a = int(self._sras.n_frames[angle_idx]) + n_rows_a, n_frames_a = self._sras.image_shape(angle_idx) if not (0 <= row_idx < n_rows_a and 0 <= frame_idx < n_frames_a): self.statusBar().showMessage( "No source waveform here (padding region of the aligned canvas).") return + self.lbl_wave_hint.hide() - ch_idx = self._current_ch - if ch_idx in CH1_DERIVED_MODES: + if self._current_ch in CH1_DERIVED_MODES: self.wave_canvas.show_rf_waveform( self._sras, angle_idx, row_idx, frame_idx, - apply_bg_sub=self.chk_bg_sub.isChecked(), - ) + apply_bg_sub=self.chk_bg_sub.isChecked()) else: self.wave_canvas.show_dc_waveform( - self._sras, angle_idx, ch_idx, row_idx, frame_idx - ) + self._sras, angle_idx, self._current_ch, row_idx, frame_idx) # ------------------------------------------------------------------ - # Progress dialog helpers + # Progress dialogs # ------------------------------------------------------------------ - def _show_progress(self, message: str): - if self._progress_dlg is not None: - self._progress_dlg.setLabelText(message) + 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.""" + dlg = self._progress_dlgs.get(key) + if dlg is not None: + dlg.setLabelText(message) return - dlg = QProgressDialog(message, "", 0, 0, self) + dlg = QProgressDialog(message, "", 0, maximum, self) dlg.setWindowTitle("Please wait…") dlg.setCancelButton(None) dlg.setWindowModality(Qt.WindowModality.WindowModal) - dlg.setMinimumDuration(300) # only appears if operation takes > 300 ms + dlg.setMinimumDuration(300) # only appears if it takes > 300 ms dlg.show() - self._progress_dlg = dlg + self._progress_dlgs[key] = dlg - def _close_progress(self): - if self._progress_dlg is not None: - self._progress_dlg.close() - self._progress_dlg = None + def _set_progress(self, key: str, pct: int): + dlg = self._progress_dlgs.get(key) + if dlg is not None: + dlg.setValue(pct) + + def _close_progress(self, key: str): + dlg = self._progress_dlgs.pop(key, None) + if dlg is not None: + dlg.close() # ------------------------------------------------------------------ # Convert menu: batch DC/FFT compute-and-store (v6 -> v7) # ------------------------------------------------------------------ def _on_batch_compute(self, mode: str): - if self._batch_thread is not None: + if self._job_running("batch"): return label = "DC" if mode == "dc" else "FFT" paths, _ = QFileDialog.getOpenFileNames( self, f"Select .sras files to batch-compute {label}", "", - "SRAS files (*.sras);;All files (*)", - ) + "SRAS files (*.sras);;All files (*)") if not paths: return - if self._batch_thread is not None: - return # a second trigger snuck in while the file dialog was open - apply_bg = self.chk_bg_sub.isChecked() self._batch_errors = [] - - # Claim self._batch_thread before anything below that can pump the - # Qt event loop — see the comment in _start_compute for why. - self._batch_worker = BatchCacheWorker(paths, mode, apply_bg) - self._batch_thread = QThread() - self._batch_worker.moveToThread(self._batch_thread) - self._batch_thread.started.connect(self._batch_worker.run) - self._batch_worker.progress.connect(self._on_batch_progress) - self._batch_worker.file_done.connect(self._on_batch_file_done) - self._batch_worker.finished.connect(lambda paths=paths: self._on_batch_finished(paths)) - self._batch_worker.finished.connect(self._batch_thread.quit) - self._batch_thread.finished.connect(self._on_batch_thread_finished) + worker = BatchCacheWorker(paths, mode, self.chk_bg_sub.isChecked()) + started = self._run_worker( + "batch", worker, + connect=( + ("progress", lambda pct: self._set_progress("batch", pct)), + ("file_done", self._on_batch_file_done), + ("finished", lambda p=paths: self._on_batch_finished(p)), + ), + on_done=self._after_batch, + ) + if not started: + return # a second trigger snuck in while the file dialog was open self._batch_dc_act.setEnabled(False) self._batch_fft_act.setEnabled(False) - - self._batch_progress_dlg = QProgressDialog( - f"Batch computing {label} for {len(paths)} file(s)…", "", 0, 100, self) - self._batch_progress_dlg.setWindowTitle("Please wait…") - self._batch_progress_dlg.setCancelButton(None) - self._batch_progress_dlg.setWindowModality(Qt.WindowModality.WindowModal) - self._batch_progress_dlg.setMinimumDuration(300) - self._batch_progress_dlg.show() - - self._batch_thread.start() - - def _on_batch_progress(self, pct: int): - if self._batch_progress_dlg is not None: - self._batch_progress_dlg.setValue(pct) + self._show_progress( + "batch", f"Batch computing {label} for {len(paths)} file(s)…", + maximum=100) def _on_batch_file_done(self, path: str, err: str): if err: self._batch_errors.append(f"{Path(path).name} — {err}") - if self._batch_progress_dlg is not None: - self._batch_progress_dlg.setLabelText(f"Processed {Path(path).name}…") + self._show_progress("batch", f"Processed {Path(path).name}…") def _on_batch_finished(self, paths: list[str]): - if self._batch_progress_dlg is not None: - self._batch_progress_dlg.close() - self._batch_progress_dlg = None + self._close_progress("batch") n_total = len(paths) n_failed = len(self._batch_errors) @@ -3144,18 +1799,12 @@ class SrasViewerWindow(QMainWindow): self.statusBar().showMessage(summary) self._batch_errors = [] - # If the currently-open file was in this batch, reload it so the - # GUI picks up the newly-written v7 cache instead of stale state. + # If the currently-open file was in this batch, reload it so the GUI + # picks up the newly-written v7 cache instead of stale state. if self._sras is not None and str(self._sras.path) in paths: self._load_file(str(self._sras.path)) - def _on_batch_thread_finished(self): - # See the comment in _on_compute_thread_finished: wait() before - # releasing our reference to avoid destroying a QThread whose OS - # thread hasn't fully joined yet. - if self._batch_thread is not None: - self._batch_thread.wait() - self._batch_thread = None + def _after_batch(self): self._batch_dc_act.setEnabled(True) self._batch_fft_act.setEnabled(True) @@ -3166,46 +1815,31 @@ class SrasViewerWindow(QMainWindow): def _on_angle_alignment(self): if self._sras is None or self._sras.n_angles <= 1: return - if self._alignment_thread is not None: - return - ref_idx = 0 + ref_idx = 0 threshold_mv = self.spin_threshold_mv.value() - generation = self._alignment_generation + generation = self._alignment_generation - # Claim self._alignment_thread before anything below that can pump - # the Qt event loop — see the comment in _start_compute for why. - self._alignment_worker = AngleAlignmentWorker(self._sras, ref_idx, threshold_mv) - self._alignment_thread = QThread() - self._alignment_worker.moveToThread(self._alignment_thread) - self._alignment_thread.started.connect(self._alignment_worker.run) - self._alignment_worker.progress.connect(self._on_alignment_progress) - self._alignment_worker.finished.connect( - lambda result, err, g=generation: self._on_alignment_done(g, result, err)) - self._alignment_worker.finished.connect(self._alignment_thread.quit) - self._alignment_thread.finished.connect(self._on_alignment_thread_finished) + started = self._run_worker( + "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)…" - ) - self._alignment_thread.start() - - def _on_alignment_progress(self, pct: int): - if self._progress_dlg is not None: - self._progress_dlg.setValue(pct) - - def _on_alignment_thread_finished(self): - # See the comment in _on_compute_thread_finished: wait() before - # releasing our reference to avoid destroying a QThread whose OS - # thread hasn't fully joined yet. - if self._alignment_thread is not None: - self._alignment_thread.wait() - self._alignment_thread = None - self._update_controls_enabled(self._sras is not None) + 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() + self._close_progress("main") if generation != self._alignment_generation: return # a new file was loaded while this was computing — discard if error_msg: @@ -3223,47 +1857,37 @@ class SrasViewerWindow(QMainWindow): f"canvas {nc}×{nr} px).") self._refresh_display() - def _on_aligned_view_toggled(self, checked: bool): - if self._current_image is not None: - self._redraw_image(self._current_image) - # ------------------------------------------------------------------ # FFT Options # ------------------------------------------------------------------ def _on_fft_options(self): - global _fft_backend - spf = self._sras.samples_per_frame if self._sras is not None else None - sr = self._sras.sample_rate_hz if self._sras is not None else None dlg = FftOptionsDialog( self, - current_backend=_fft_backend, + current_backend=compute.get_fft_backend(), current_pad_factor=self._fft_pad_factor, - samples_per_frame=spf, - sample_rate_hz=sr, + samples_per_frame=self._sras.samples_per_frame if self._sras else None, + sample_rate_hz=self._sras.sample_rate_hz if self._sras else None, grating_um=self.spin_grating_um.value(), ) - if dlg.exec() == QDialog.DialogCode.Accepted: - _fft_backend = dlg.get_backend() - self._fft_pad_factor = dlg.get_pad_factor() - # Pad factor changes the FFT bin count, so it genuinely - # invalidates the cached raw FFT (part of the cache key) — - # _refresh_display() will recompute only on a cache miss. - if self._sras is not None and self.combo_channel.currentIndex() in ( - CH1_IDX, VELOCITY_MODE_IDX): - self._refresh_display() + if dlg.exec() != QDialog.DialogCode.Accepted: + return + compute.set_fft_backend(dlg.get_backend()) + self._fft_pad_factor = dlg.get_pad_factor() + # Pad factor changes the FFT bin count, so it genuinely invalidates + # the cached raw FFT (part of the cache key) — _refresh_display() + # recomputes only on a cache miss. + if self._sras is not None and self.combo_channel.currentIndex() in CH1_DERIVED_MODES: + self._refresh_display() # ------------------------------------------------------------------ def closeEvent(self, event): if self._dc_precompute_worker is not None: self._dc_precompute_worker.stop() - for attr in ("_load_thread", "_compute_thread", "_batch_thread", - "_dc_precompute_thread", "_alignment_thread"): - t = getattr(self, attr, None) - if t is not None: - t.quit() - t.wait(2000) + for thread, _worker, _on_done in list(self._jobs.values()): + thread.quit() + thread.wait(2000) super().closeEvent(event) @@ -3271,8 +1895,8 @@ class SrasViewerWindow(QMainWindow): def main(): app = QApplication(sys.argv) - initial = sys.argv[1] if len(sys.argv) > 1 else None - window = SrasViewerWindow(initial_path=initial) + window = SrasViewerWindow( + initial_path=sys.argv[1] if len(sys.argv) > 1 else None) window.show() sys.exit(app.exec()) diff --git a/sras_workers.py b/sras_workers.py new file mode 100644 index 0000000..1445f9e --- /dev/null +++ b/sras_workers.py @@ -0,0 +1,224 @@ +#!/usr/bin/env python3 +"""Background workers for the SRAS viewer. + +Every worker is a plain QObject moved onto its own QThread by +SrasViewerWindow._run_worker, exposing signals only. Workers must never touch +GUI-thread-owned state (the display caches in particular) — they take +everything they need through their constructor and hand results back by signal. +""" + +import os +from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed + +import numpy as np +from PyQt6.QtCore import QObject, pyqtSignal + +import sras_compute as compute +from sras_compute import ( + cache_file, compute_angle_alignment, 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 +# those processes threads internally, so the two must be divided rather than +# both set to the core count. Files also commonly sit on one external drive, +# where a dozen concurrent readers is slower than a few — hence the low +# default, overridable from the environment. +_BATCH_MAX_PROCS = int(os.environ.get("SRAS_BATCH_PROCS", 0)) or min( + 4, os.cpu_count() or 2) + + +class LoadWorker(QObject): + finished = pyqtSignal(object) # SrasFile | None + error = pyqtSignal(str) + + def __init__(self, path: str): + super().__init__() + self._path = path + + def run(self): + try: + self.finished.emit(SrasFile(self._path)) + except Exception as exc: + self.error.emit(str(exc)) + self.finished.emit(None) + + +class ComputeWorker(QObject): + """Computes one displayable image for (angle, channel). + + For CH1/Velocity (FFT-derived) channels, the FFT is only run for pixels + whose DC4 (Bias B) mean is at or above dc_threshold_mv — masked pixels are + left at 0 MHz without ever being FFT'd, since that's the expensive part of + a scan. If the DC4 image for this angle is already known, pass it in as + *dc4_mv* to skip re-reading the CH4 channel from disk entirely. + + Emits a plain ``np.ndarray`` already in display units. + """ + finished = pyqtSignal(object) + error = pyqtSignal(str) + + def __init__(self, sras: SrasFile, angle_idx: int, ch_idx: int, + apply_bg_sub: bool = True, n_fft: int | None = None, + dc_threshold_mv: float = 0.0, + dc4_mv: np.ndarray | None = None, + is_fft_mode: bool = False): + super().__init__() + self._sras = sras + self._angle = angle_idx + self._ch = ch_idx + self._apply_bg_sub = apply_bg_sub + self._n_fft = n_fft + self._dc_threshold = dc_threshold_mv + self._dc4_mv = dc4_mv + self._is_fft_mode = is_fft_mode + + def run(self): + try: + if self._is_fft_mode: + img = compute_rf_image( + self._sras, self._angle, dc_threshold_mv=self._dc_threshold, + apply_bg_sub=self._apply_bg_sub, n_fft=self._n_fft, + dc4_mv=self._dc4_mv) + else: + img = dc_image_mv(self._sras, self._angle, self._ch) + self.finished.emit(img) + except Exception as exc: + self.error.emit(str(exc)) + + +class DcPrecomputeWorker(QObject): + """Computes CH3/CH4 DC images for every angle in the background. + + DC images are cheap (a per-waveform mean, no FFT) compared to the + CH1/Velocity FFT, so precomputing them for the whole file right after load + makes switching angles instant while on a DC channel, and also means the + FFT masking step (which needs a DC4 image) rarely has to wait on anything. + + Angles are computed on a thread pool — the work is a pure mean over the + waveform block, so it is I/O- and bandwidth-bound and embarrassingly + parallel. Results are emitted one at a time as they land (out of angle + order), and always from this worker's own thread: nothing emits a Qt + signal from a pool thread. + """ + angle_done = pyqtSignal(int, np.ndarray, np.ndarray) # angle_idx, dc3_mv, dc4_mv + finished = pyqtSignal() + error = pyqtSignal(str) + + def __init__(self, sras: SrasFile): + super().__init__() + self._sras = sras + self._stop = False + + def stop(self): + self._stop = True + + def _one_angle(self, a: int) -> tuple[int, np.ndarray, np.ndarray]: + return a, dc_image_mv(self._sras, a, CH3_IDX), dc_image_mv(self._sras, a, CH4_IDX) + + def run(self): + try: + n = self._sras.n_angles + n_workers = max(1, min(compute._MAX_WORKERS, n)) + with ThreadPoolExecutor(max_workers=n_workers) as pool: + futures = {pool.submit(self._one_angle, a): a for a in range(n)} + try: + for fut in as_completed(futures): + if self._stop: + break + a, dc3, dc4 = fut.result() + self.angle_done.emit(a, dc3, dc4) + finally: + if self._stop: + for fut in futures: + fut.cancel() + self.finished.emit() + except Exception as exc: + self.error.emit(str(exc)) + + +class BatchCacheWorker(QObject): + """Batch-computes and stores DC or FFT images into each of *paths*'s v7 + CACH tail, in place — converting v6 sources to v7 on first use, or updating + an existing v7 file's cache blocks without disturbing whatever the other + block already holds. + + *mode* is ``"dc"`` (CH3/CH4 mean images) or ``"fft"`` (CH1 peak-frequency + images, unmasked — masking is applied at display time, same as v5's PREC + convention). + + Files are processed one per subprocess: they are fully independent, each + opens its own memmap and writes only its own bytes, and only path strings + and scalars cross the process boundary. Emits ``progress(int)`` (0–100 by + files completed), ``file_done(str, str)`` (path, error message or "") so + one file's failure doesn't abort the batch, and ``finished()``. + """ + progress = pyqtSignal(int) + file_done = pyqtSignal(str, str) + finished = pyqtSignal() + + def __init__(self, paths: list[str], mode: str, apply_bg_sub: bool): + super().__init__() + self._paths = paths + self._mode = mode + self._apply_bg_sub = apply_bg_sub + + def _pool(self, n_files: int): + """A process pool, falling back to threads if the platform refuses to + spawn (still a win: the compute releases the GIL).""" + n = max(1, min(_BATCH_MAX_PROCS, n_files)) + try: + return ProcessPoolExecutor(max_workers=n), n + except (OSError, ValueError): + return ThreadPoolExecutor(max_workers=n), n + + def run(self): + paths = self._paths + executor, n_procs = self._pool(len(paths)) + # Each child threads internally; divide the machine rather than + # letting every process claim every core. + per_proc_workers = max(1, (os.cpu_count() or 4) // n_procs) + + done = 0 + with executor: + futures = { + executor.submit(cache_file, p, self._mode, self._apply_bg_sub, + compute.get_fft_backend(), per_proc_workers): p + for p in paths + } + for fut in as_completed(futures): + path = futures[fut] + try: + err = fut.result() + except Exception as exc: + err = str(exc) + done += 1 + self.file_done.emit(path, err) + self.progress.emit(int(done / max(1, len(paths)) * 100)) + + self.finished.emit() + + +class AngleAlignmentWorker(QObject): + """Computes rotation+translation alignment for every angle in *sras*, + referenced to *ref_angle_idx*, from each angle's binarized CH4 mask. + Rotation is analytic (from sras.angles_deg); only translation is found by + phase correlation. + """ + 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)) From 9a7cc396e7126d02c8d7a9c4f97e9226f0712328 Mon Sep 17 00:00:00 2001 From: Thomas Ales Date: Fri, 31 Jul 2026 00:09:46 -0500 Subject: [PATCH 04/10] Bound nested parallelism, make compute cancellable, harden batch pool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three defects found by testing the parallel paths end to end. Nested parallelism was unbounded. DcPrecomputeWorker and the alignment driver both parallelise over angles, and each angle's compute_dc_image then parallelised over rows again: 14 x 14 threads, and — worse — every angle sized its chunk against the *whole* memory budget, so worst-case live buffers were ~13.9 GB against a 1 GB budget. Capping the inner worker count alone does not fix it; the chunk is still sized against the full budget. plan_angle_level now hands each angle both max_workers=1 and its share of the budget, which bounds the real file at 1002 MB. Chunk planning could not produce concurrency at all. _plan_chunks sized chunk_rows first, and _chunk_rows_for spends whatever budget it is given — so a single chunk consumed the lot and budget//chunk_bytes came back as one worker. It now picks the worker count first and sizes the chunk to it, giving 7-16 workers on the real file instead of 1. Cancellation was too coarse to shut down. stop() was only checked between angles, and one angle of the real scan is ~40 s, so closing the window sat through it and returned with threads still running. should_stop is now polled per row chunk and closeEvent signals every worker before waiting: close went from 4.0 s with two live threads to 1.04 s with both finished. A cancelled ComputeWorker emits None so a half-filled image is never cached. Also: BatchCacheWorker no longer reports every file as failed when the process pool itself dies (unguarded __main__, frozen build, sandbox) — it retries those files in-process. Batches below 512 MB skip the pool entirely, since interpreter startup would otherwise make small jobs slower. cache_file rejects an unknown mode instead of silently treating it as "fft". Measured on the real 496 GB scan, angle 3, 32 rows: DC 2.94s -> 0.75s RF 6.10s -> 1.99s peak RSS 1.40 -> 1.93 GB Batch convert, 24 files / 1 GB: 4.03s -> 2.36s. Memory budget defaults to 1024 MB (SRAS_MEM_BUDGET_MB), the measured knee. Testing: tools/test_refactor.py (70 checks: cache round-trip and block carry-forward, partial-cache handling, parallel-vs-serial identity, no-mask path, ROI mask vs full-grid, v2/v3/v4 parsing, sras_average round-trip incl. remainder handling) and tools/test_gui.py (46 checks driving the real widgets and workers headlessly). check_equivalence.py still reports all 167 outputs identical to ed0eba4. Co-Authored-By: Claude Opus 5 --- sras_compute.py | 123 +++++++++++--- sras_viewer.py | 16 +- sras_workers.py | 154 +++++++++++++---- tools/make_test_sras.py | 57 ++++++- tools/test_gui.py | 270 ++++++++++++++++++++++++++++++ tools/test_refactor.py | 358 ++++++++++++++++++++++++++++++++++++++++ 6 files changed, 909 insertions(+), 69 deletions(-) create mode 100644 tools/test_gui.py create mode 100644 tools/test_refactor.py diff --git a/sras_compute.py b/sras_compute.py index 280cc87..fc01b06 100644 --- a/sras_compute.py +++ b/sras_compute.py @@ -66,7 +66,11 @@ def _do_rfft(x: np.ndarray, n: int | None = None, axis: int = -1, # per-chunk size cannot buy more concurrency — the worker count must be derived # from the budget instead. See _plan_chunks. -_TOTAL_BYTES_BUDGET = int(os.environ.get("SRAS_MEM_BUDGET_MB", 1536)) * 1024 * 1024 +# 1024 MB is the measured knee on a 16-core machine against a 7507-frame x +# 2500-sample angle: 512 MB leaves ~20% of the FFT speedup on the table, and +# 1536+ MB costs ~0.4 GB more resident for no further gain. Override with +# SRAS_MEM_BUDGET_MB on a smaller machine. +_TOTAL_BYTES_BUDGET = int(os.environ.get("SRAS_MEM_BUDGET_MB", 1024)) * 1024 * 1024 _CHUNK_ROWS_MAX = 32 # cap for small scans (original behavior) _MAX_WORKERS = int(os.environ.get("SRAS_MAX_WORKERS", 0)) or (os.cpu_count() or 4) @@ -83,21 +87,38 @@ def _chunk_rows_for(n_frames: int, samples_per_frame: int, def _plan_chunks(n_rows: int, n_frames: int, samples_per_frame: int, live_multiplier: int = 1, - max_workers: int | None = None) -> tuple[int, int]: - """(chunk_rows, n_workers) sized so live_multiplier x chunk_rows x - n_frames x samples x 4 bytes x n_workers stays within the budget.""" - per_chunk_budget = max(1, _TOTAL_BYTES_BUDGET // live_multiplier) - chunk_rows = _chunk_rows_for(n_frames, samples_per_frame, per_chunk_budget) + max_workers: int | None = None, + budget: int | None = None) -> tuple[int, int]: + """(chunk_rows, n_workers) such that all concurrent chunks together fit + the budget: n_workers x live_multiplier x chunk_rows x n_frames x + samples x 4 bytes <= budget. - chunk_bytes = max(1, live_multiplier * chunk_rows * n_frames * samples_per_frame * 4) + The worker count is chosen *first* and the chunk sized to it. Sizing the + chunk first is the trap: _chunk_rows_for spends whatever budget it is + given, so a single chunk would always consume the lot and leave room for + exactly one worker — no concurrency, on precisely the large scans that + need it most. + + A caller that is itself running several of these concurrently must pass + *both* max_workers=1 and its share of the budget. Capping the workers + alone is not enough: the chunk would still be sized against the whole + budget, and N concurrent callers would each allocate all of it. + """ + total = _TOTAL_BYTES_BUDGET if budget is None else max(1, budget) cap = max_workers if max_workers is not None else _MAX_WORKERS - n_workers = int(max(1, min(cap, - _TOTAL_BYTES_BUDGET // chunk_bytes, - -(-n_rows // chunk_rows)))) + bytes_per_row = max(1, live_multiplier * n_frames * samples_per_frame * 4) + + # A chunk is at least one row, so that alone caps how many can be live. + n_workers = int(max(1, min(cap, n_rows, total // bytes_per_row))) + chunk_rows = _chunk_rows_for(n_frames, samples_per_frame, + max(1, total // (n_workers * live_multiplier))) + # With chunk_rows known, more workers than chunks buys nothing. + n_workers = int(max(1, min(n_workers, -(-n_rows // chunk_rows)))) return chunk_rows, n_workers -def _map_row_chunks(n_rows: int, chunk_rows: int, n_workers: int, fn): +def _map_row_chunks(n_rows: int, chunk_rows: int, n_workers: int, fn, + should_stop=None): """Apply fn(r0, r1) over row chunks, in parallel when it pays. Chunks write to disjoint output slices, so no locking is needed. numpy @@ -105,48 +126,94 @@ def _map_row_chunks(n_rows: int, chunk_rows: int, n_workers: int, fn): threads give real parallelism here — and on a very large file they also keep many more page-fault requests in flight, which is what the I/O path wants. + + *should_stop* is polled per chunk so a cancelled job abandons work within + one chunk rather than one whole angle — on a large scan an angle is + ~40 s, which is far too long to block application shutdown. """ bounds = [(r0, min(r0 + chunk_rows, n_rows)) for r0 in range(0, n_rows, chunk_rows)] + + def run(b): + if should_stop is not None and should_stop(): + return + fn(*b) + if n_workers <= 1 or len(bounds) == 1: - for r0, r1 in bounds: - fn(r0, r1) + for b in bounds: + run(b) return with ThreadPoolExecutor(max_workers=min(n_workers, len(bounds))) as pool: - list(pool.map(lambda b: fn(*b), bounds)) + list(pool.map(run, bounds)) # --------------------------------------------------------------------------- # Image computation # --------------------------------------------------------------------------- -def compute_dc_image(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray: - """Mean of each waveform → (n_rows, n_frames) float32, in ADC counts.""" +def plan_angle_level(sras: SrasFile, + live_multiplier: int = 1) -> tuple[int, int]: + """(n_angle_workers, per_angle_budget) for a caller that parallelises + over angles instead of over rows. + + Callers that parallelise over angles must not also let each angle + parallelise over rows: the two levels multiply, in threads and in memory. + Each angle is therefore given max_workers=1 and the returned budget share, + which together keep total live buffers within _TOTAL_BYTES_BUDGET. + """ + biggest = max(range(sras.n_angles), key=lambda a: int(sras.n_frames[a])) + _, n_workers = _plan_chunks(int(sras.n_rows[biggest]), + int(sras.n_frames[biggest]), + sras.samples_per_frame, + live_multiplier=live_multiplier) + n_workers = max(1, min(n_workers, sras.n_angles)) + return n_workers, max(1, _TOTAL_BYTES_BUDGET // n_workers) + + +def compute_dc_image(sras: SrasFile, angle_idx: int, ch_idx: int, + max_workers: int | None = None, + budget: int | None = None, + should_stop=None) -> np.ndarray: + """Mean of each waveform → (n_rows, n_frames) float32, in ADC counts. + + Pass max_workers=1 when the caller is already parallelising over angles. + If *should_stop* ever returns True the result is incomplete — the caller + is expected to be abandoning it. + """ n_rows, n_frames = sras.image_shape(angle_idx) data = sras.data[angle_idx] - chunk_rows, n_workers = _plan_chunks(n_rows, n_frames, sras.samples_per_frame) + chunk_rows, n_workers = _plan_chunks(n_rows, n_frames, sras.samples_per_frame, + max_workers=max_workers, budget=budget) img = np.empty((n_rows, n_frames), dtype=np.float32) def chunk(r0: int, r1: int): img[r0:r1] = data[r0:r1, ch_idx, :, :].astype(np.float32).mean(axis=-1) - _map_row_chunks(n_rows, chunk_rows, n_workers, chunk) + _map_row_chunks(n_rows, chunk_rows, n_workers, chunk, should_stop=should_stop) return img -def dc_image_mv(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray: +def dc_image_mv(sras: SrasFile, angle_idx: int, ch_idx: int, + max_workers: int | None = None, budget: int | None = None, + should_stop=None) -> np.ndarray: """DC image for (angle, channel) in mV, preferring a stored v5/v7 cache.""" cached = sras.cached_dc_mv(angle_idx, ch_idx) if cached is not None: return cached - return adc_to_mv(compute_dc_image(sras, angle_idx, ch_idx), *sras.cal(ch_idx)) + return adc_to_mv( + compute_dc_image(sras, angle_idx, ch_idx, max_workers=max_workers, + budget=budget, should_stop=should_stop), + *sras.cal(ch_idx)) def compute_rf_image(sras: SrasFile, angle_idx: int, dc_threshold_mv: float | None, apply_bg_sub: bool = True, n_fft: int | None = None, - dc4_mv: np.ndarray | None = None) -> np.ndarray: + dc4_mv: np.ndarray | None = None, + max_workers: int | None = None, + budget: int | None = None, + should_stop=None) -> np.ndarray: """FFT of each CH1 waveform; pixel = peak frequency in MHz. Pixels where CH4_dc < dc_threshold_mv are set to 0 — and the FFT is @@ -193,7 +260,8 @@ def compute_rf_image(sras: SrasFile, angle_idx: int, n_fft_bins = n_fft if n_fft is not None else sras.samples_per_frame chunk_rows, n_workers = _plan_chunks( n_rows, n_frames, max(sras.samples_per_frame, n_fft_bins), - live_multiplier=_FFT_LIVE_MULTIPLIER) + live_multiplier=_FFT_LIVE_MULTIPLIER, max_workers=max_workers, + budget=budget) background = sras.background if (apply_bg_sub and sras.background is not None) else None cal4 = sras.cal(CH4_IDX) @@ -239,7 +307,7 @@ def compute_rf_image(sras: SrasFile, angle_idx: int, else: img[r0:r1] = freq_axis[peak_bins] - _map_row_chunks(n_rows, chunk_rows, n_workers, chunk) + _map_row_chunks(n_rows, chunk_rows, n_workers, chunk, should_stop=should_stop) return img @@ -258,6 +326,8 @@ def cache_file(path: str, mode: str, apply_bg_sub: bool, """ global _MAX_WORKERS try: + if mode not in ("dc", "fft"): + return f"unknown cache mode {mode!r} (expected 'dc' or 'fft')" set_fft_backend(fft_backend) if max_workers: _MAX_WORKERS = max_workers @@ -493,7 +563,8 @@ def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int, background-thread workers must not touch GUI-thread-owned caches.""" n = sras.n_angles # already the *complete*-angle count for aborted v6 scans dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx) - n_workers = max(1, min(_MAX_WORKERS, n)) + # Parallel over angles, serial within each — see plan_angle_level. + n_workers, angle_budget = plan_angle_level(sras) # progress_cb fires from pool threads, so the counter behind it must be # atomic. list.append is, and len() of a list is a consistent read. @@ -506,7 +577,9 @@ def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int, # ---- Step 1: binarized CH4 mask per angle, native per-angle grid ----- def mask_for(a: int) -> np.ndarray: - dc4 = adc_to_mv(compute_dc_image(sras, a, CH4_IDX), *sras.cal(CH4_IDX)) + dc4 = adc_to_mv( + compute_dc_image(sras, a, CH4_IDX, max_workers=1, budget=angle_budget), + *sras.cal(CH4_IDX)) m = (dc4 >= dc_threshold_mv).astype(np.float32) tick(0, 25) return m diff --git a/sras_viewer.py b/sras_viewer.py index 6336e09..77a9a0c 100644 --- a/sras_viewer.py +++ b/sras_viewer.py @@ -1624,6 +1624,8 @@ class SrasViewerWindow(QMainWindow): def _on_compute_done(self, result): self._close_progress("main") + if result is None: + return # cancelled mid-compute; the partial image must not cache angle_idx = self._pending_angle ch_idx = self._pending_ch @@ -1883,11 +1885,17 @@ class SrasViewerWindow(QMainWindow): # ------------------------------------------------------------------ def closeEvent(self, event): - if self._dc_precompute_worker is not None: - self._dc_precompute_worker.stop() - for thread, _worker, _on_done in list(self._jobs.values()): + # Signal every cancellable worker first, then wait. Waiting without + # signalling means sitting out whatever is in flight — on a large + # scan a single angle is ~40 s. + jobs = list(self._jobs.values()) + for _thread, worker, _on_done in jobs: + stop = getattr(worker, "stop", None) + if callable(stop): + stop() + for thread, _worker, _on_done in jobs: thread.quit() - thread.wait(2000) + thread.wait(5000) super().closeEvent(event) diff --git a/sras_workers.py b/sras_workers.py index 1445f9e..1a21e50 100644 --- a/sras_workers.py +++ b/sras_workers.py @@ -9,6 +9,7 @@ everything they need through their constructor and hand results back by signal. import os from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed +from concurrent.futures.process import BrokenProcessPool import numpy as np from PyQt6.QtCore import QObject, pyqtSignal @@ -27,6 +28,31 @@ from sras_format import CH3_IDX, CH4_IDX, SrasFile _BATCH_MAX_PROCS = int(os.environ.get("SRAS_BATCH_PROCS", 0)) or min( 4, os.cpu_count() or 2) +# Spawning a pool costs roughly a second of interpreter startup (each child +# re-imports the entry module). That is noise against a multi-GB scan but +# dominates a batch of small files, where it would make the job *slower* — +# so below this total size the batch just runs in the worker thread. +_BATCH_POOL_MIN_BYTES = int(os.environ.get("SRAS_BATCH_POOL_MIN_MB", 512)) * 1024 * 1024 + + +class CancellableWorker(QObject): + """A worker whose compute polls stop() between row chunks. + + Without this a shutdown has to wait out whatever is in flight, and on a + large scan a single angle is ~40 s — far too long to block closing the + window. Chunk-level polling bounds the wait to one chunk instead. + """ + + def __init__(self): + super().__init__() + self._stop = False + + def stop(self): + self._stop = True + + def _stopped(self) -> bool: + return self._stop + class LoadWorker(QObject): finished = pyqtSignal(object) # SrasFile | None @@ -44,7 +70,7 @@ class LoadWorker(QObject): self.finished.emit(None) -class ComputeWorker(QObject): +class ComputeWorker(CancellableWorker): """Computes one displayable image for (angle, channel). For CH1/Velocity (FFT-derived) channels, the FFT is only run for pixels @@ -79,15 +105,19 @@ class ComputeWorker(QObject): img = compute_rf_image( self._sras, self._angle, dc_threshold_mv=self._dc_threshold, apply_bg_sub=self._apply_bg_sub, n_fft=self._n_fft, - dc4_mv=self._dc4_mv) + dc4_mv=self._dc4_mv, should_stop=self._stopped) else: - img = dc_image_mv(self._sras, self._angle, self._ch) - self.finished.emit(img) + img = dc_image_mv(self._sras, self._angle, self._ch, + should_stop=self._stopped) + # On cancellation the image is only partly filled, so hand back + # None rather than something that would be cached as real. The + # signal still fires either way — it is what quits the thread. + self.finished.emit(None if self._stop else img) except Exception as exc: self.error.emit(str(exc)) -class DcPrecomputeWorker(QObject): +class DcPrecomputeWorker(CancellableWorker): """Computes CH3/CH4 DC images for every angle in the background. DC images are cheap (a per-waveform mean, no FFT) compared to the @@ -108,30 +138,34 @@ class DcPrecomputeWorker(QObject): def __init__(self, sras: SrasFile): super().__init__() self._sras = sras - self._stop = False - - def stop(self): - self._stop = True def _one_angle(self, a: int) -> tuple[int, np.ndarray, np.ndarray]: - return a, dc_image_mv(self._sras, a, CH3_IDX), dc_image_mv(self._sras, a, CH4_IDX) + # max_workers=1 *and* a budget share: this call is one of several + # concurrent angles, and both the thread count and the buffer size + # have to be divided (see compute.plan_angle_level). + kw = dict(max_workers=1, budget=self._angle_budget, + should_stop=self._stopped) + return (a, + dc_image_mv(self._sras, a, CH3_IDX, **kw), + dc_image_mv(self._sras, a, CH4_IDX, **kw)) def run(self): try: n = self._sras.n_angles - n_workers = max(1, min(compute._MAX_WORKERS, n)) - with ThreadPoolExecutor(max_workers=n_workers) as pool: + n_workers, self._angle_budget = compute.plan_angle_level(self._sras) + pool = ThreadPoolExecutor(max_workers=n_workers) + try: futures = {pool.submit(self._one_angle, a): a for a in range(n)} - try: - for fut in as_completed(futures): - if self._stop: - break - a, dc3, dc4 = fut.result() - self.angle_done.emit(a, dc3, dc4) - finally: + for fut in as_completed(futures): if self._stop: - for fut in futures: - fut.cancel() + break + a, dc3, dc4 = fut.result() + self.angle_done.emit(a, dc3, dc4) + finally: + # cancel_futures drops the queued angles; should_stop lets the + # in-flight ones bail within a chunk. Not waiting here is what + # keeps closing the window responsive on a large scan. + pool.shutdown(wait=not self._stop, cancel_futures=True) self.finished.emit() except Exception as exc: self.error.emit(str(exc)) @@ -163,24 +197,27 @@ class BatchCacheWorker(QObject): self._mode = mode self._apply_bg_sub = apply_bg_sub - def _pool(self, n_files: int): - """A process pool, falling back to threads if the platform refuses to - spawn (still a win: the compute releases the GIL).""" - n = max(1, min(_BATCH_MAX_PROCS, n_files)) - try: - return ProcessPoolExecutor(max_workers=n), n - except (OSError, ValueError): - return ThreadPoolExecutor(max_workers=n), n + def _report(self, path: str, err: str, done: int, total: int): + self.file_done.emit(path, err) + self.progress.emit(int(done / max(1, total) * 100)) - def run(self): - paths = self._paths - executor, n_procs = self._pool(len(paths)) + def _run_pooled(self, paths: list[str], n_procs: int) -> list[str]: + """Process the batch across *n_procs* subprocesses. Returns the paths + that never got a real answer because the pool itself died, so the + caller can retry them in-process. + + Under spawn each child re-imports the entry module, so the batch must + survive that going wrong (an unguarded __main__, a frozen build, a + sandbox that forbids subprocesses) rather than reporting every file as + failed — hence the retry list instead of a per-file error. + """ # Each child threads internally; divide the machine rather than # letting every process claim every core. per_proc_workers = max(1, (os.cpu_count() or 4) // n_procs) - + unresolved: list[str] = [] done = 0 - with executor: + + with ProcessPoolExecutor(max_workers=n_procs) as executor: futures = { executor.submit(cache_file, p, self._mode, self._apply_bg_sub, compute.get_fft_backend(), per_proc_workers): p @@ -190,11 +227,56 @@ class BatchCacheWorker(QObject): path = futures[fut] try: err = fut.result() + except BrokenProcessPool: + unresolved.append(path) + continue except Exception as exc: err = str(exc) done += 1 - self.file_done.emit(path, err) - self.progress.emit(int(done / max(1, len(paths)) * 100)) + self._report(path, err, done, len(paths)) + + return unresolved + + def _run_inline(self, paths: list[str], done: int, total: int): + """Fallback / single-file path: compute in this thread. Still uses the + full core count internally, since nothing else is competing.""" + for path in paths: + try: + err = cache_file(path, self._mode, self._apply_bg_sub, + compute.get_fft_backend(), compute._MAX_WORKERS) + except Exception as exc: + err = str(exc) + done += 1 + self._report(path, err, done, total) + + def _worth_pooling(self, paths: list[str]) -> bool: + if len(paths) < 2: + return False + total = 0 + for p in paths: + try: + total += os.path.getsize(p) + except OSError: + pass # unreadable files are reported by cache_file + return total >= _BATCH_POOL_MIN_BYTES + + def run(self): + paths = self._paths + n_procs = max(1, min(_BATCH_MAX_PROCS, len(paths))) + + if not self._worth_pooling(paths): + self._run_inline(paths, 0, len(paths)) + self.finished.emit() + return + + try: + unresolved = self._run_pooled(paths, n_procs) + except Exception: + # The pool could not be created or collapsed wholesale. + unresolved = list(paths) + + if unresolved: + self._run_inline(unresolved, len(paths) - len(unresolved), len(paths)) self.finished.emit() diff --git a/tools/make_test_sras.py b/tools/make_test_sras.py index 5ea0e00..07386c3 100644 --- a/tools/make_test_sras.py +++ b/tools/make_test_sras.py @@ -41,9 +41,11 @@ def _preamble(ymult_v: float, yoff_adc: float, yzero_v: float) -> bytes: ).encode("utf-8") -def build(n_angles: int, seed: int, samples_per_frame: int) -> tuple[bytes, dict]: +def build(n_angles: int, seed: int, samples_per_frame: int, + geometry: list[tuple[int, int]] | None = None) -> tuple[bytes, dict]: rng = np.random.default_rng(seed) - geom = [_GEOMETRY[a % len(_GEOMETRY)] for a in range(n_angles)] + src_geom = geometry or _GEOMETRY + geom = [src_geom[a % len(src_geom)] for a in range(n_angles)] n_ch = 3 bps = 1 @@ -116,12 +118,59 @@ def build(n_angles: int, seed: int, samples_per_frame: int) -> tuple[bytes, dict def write(path: Path, n_angles: int = 3, seed: int = 0, - samples_per_frame: int = 64) -> dict: - payload, meta = build(n_angles, seed, samples_per_frame) + samples_per_frame: int = 64, + geometry: list[tuple[int, int]] | None = None) -> dict: + payload, meta = build(n_angles, seed, samples_per_frame, geometry) path.write_bytes(payload) return meta +HDR_FMT_LEGACY = ">4sBHHffffIIdBB" + + +def write_legacy(path: Path, version: int = 4, n_angles: int = 2, + n_rows: int = 4, n_frames: int = 10, + samples_per_frame: int = 32, seed: int = 0) -> dict: + """Write a v2/v3/v4 file: uniform geometry, one flat waveform block. + + Used to exercise sras_average.py, which only handles the legacy formats. + """ + rng = np.random.default_rng(seed) + n_ch, bps = 3, 1 + + out = bytearray() + out += struct.pack( + HDR_FMT_LEGACY, b"SRAS", version, n_angles, n_rows, + -0.5, 1.0, _VELOCITY_MM_S, _LASER_FREQ_HZ, + n_frames, samples_per_frame, _SAMPLE_RATE_HZ, bps, n_ch, + ) + angles = np.linspace(0.0, 45.0, n_angles, dtype=np.float32) + out += angles.astype(">f4").tobytes() + y = (np.arange(n_rows) * _ROW_SPACING_MM).astype(np.float32) + out += y.astype(">f4").tobytes() + + if version >= 3: + for ymult_v, yoff, yzero_v in ((1.5625e-3, -87.04, 0.0), + (2.0e-3, -60.0, 1.0e-3), + (2.5e-3, -40.0, -2.0e-3))[:n_ch]: + p = _preamble(ymult_v, yoff, yzero_v) + out += struct.pack(">H", len(p)) + p + + background = rng.integers(-8, 9, size=samples_per_frame, dtype=np.int8) + if version >= 4: + out += struct.pack(">I", samples_per_frame) + background.tobytes() + + data = rng.integers(-100, 101, + size=(n_angles, n_rows, n_ch, n_frames, samples_per_frame), + dtype=np.int8) + out += data.tobytes() + path.write_bytes(bytes(out)) + return {"version": version, "n_angles": n_angles, "n_rows": n_rows, + "n_frames": n_frames, "samples_per_frame": samples_per_frame, + "n_channels": n_ch, "data": data, "angles_deg": angles, + "y_positions": y, "background": background} + + def main(): p = argparse.ArgumentParser(description=__doc__) p.add_argument("output") diff --git a/tools/test_gui.py b/tools/test_gui.py new file mode 100644 index 0000000..c44e6da --- /dev/null +++ b/tools/test_gui.py @@ -0,0 +1,270 @@ +#!/usr/bin/env python3 +"""Headless GUI test: drives SrasViewerWindow through the real Qt widgets, +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, aligned view, ROI draw/move, and CSV export. + +Usage: QT_QPA_PLATFORM=offscreen python tools/test_gui.py [file.sras] +""" + +import os +import sys +import tempfile +from pathlib import Path + +os.environ.setdefault("QT_QPA_PLATFORM", "offscreen") + +import numpy as np # noqa: E402 +from PyQt6.QtCore import QEventLoop, QTimer # noqa: E402 +from PyQt6.QtWidgets import QApplication # noqa: E402 + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from sras_format import CH1_IDX, CH3_IDX, CH4_IDX # noqa: E402 +from sras_viewer import RoiQuad, SrasViewerWindow, VELOCITY_MODE_IDX # noqa: E402 +import tools.make_test_sras as gen # noqa: E402 + +_failures: list[str] = [] + + +def check(name: str, ok: bool, detail: str = ""): + print(f" {'PASS' if ok else 'FAIL'} {name}" + (f" — {detail}" if detail else "")) + if not ok: + _failures.append(name) + + +def pump(ms: int = 250): + """Run the event loop for a while so queued signals and worker threads + make progress.""" + loop = QEventLoop() + QTimer.singleShot(ms, loop.quit) + loop.exec() + + +def wait_until(pred, timeout_ms: int = 20000, step: int = 100) -> bool: + waited = 0 + while waited < timeout_ms: + if pred(): + return True + pump(step) + waited += step + return pred() + + +def main(): + app = QApplication(sys.argv) + errors: list[str] = [] + + tmpdir = Path(tempfile.mkdtemp(prefix="sras_gui_")) + path = Path(sys.argv[1]) if len(sys.argv) > 1 else tmpdir / "gui.sras" + if len(sys.argv) <= 1: + gen.write(path, n_angles=4, seed=11, samples_per_frame=256) + + print(f"\nloading {path.name}") + win = SrasViewerWindow() + win.show() + # Capture anything the app reports as an error via the status bar. + win.statusBar().messageChanged.connect( + lambda m: errors.append(m) if m and "error" in m.lower() else None) + + win._load_file(str(path)) + check("file loaded", wait_until(lambda: win._sras is not None)) + s = win._sras + check("parsed as v6", s.version == 6, f"v{s.version}") + check("defaults to CH4", win.combo_channel.currentIndex() == CH4_IDX) + check("image displayed", win._current_image is not None) + check("angle spinbox ranges over all angles", + win.spin_angle.maximum() == s.n_angles - 1) + check("scan info populated", + win._info["Angles"].text() == f"Angles: {s.n_angles}", + win._info["Angles"].text()) + + print("\nbackground DC precompute (all angles)") + ok = wait_until(lambda: all((a, CH4_IDX) in win._dc_cache + and (a, CH3_IDX) in win._dc_cache + for a in range(s.n_angles))) + check("every angle cached for CH3 and CH4", ok, + f"{len(win._dc_cache)} entries") + check("status label reports completion", + "ready for all angles" in win.lbl_dc_precompute.text(), + win.lbl_dc_precompute.text()) + + print("\nangle switching (DC, should be served from cache)") + for a in range(s.n_angles): + win.spin_angle.setValue(a) + win._on_view_changed() + pump(60) + expected = win._sras.image_shape(a) + check(f"angle {a} shows its own geometry {expected}", + win._current_image.shape == expected, + str(win._current_image.shape)) + check("no compute job needed for cached DC angles", + not win._job_running("compute")) + + print("\nchannel switching") + win.spin_angle.setValue(0) + win._on_view_changed() + pump(60) + win.combo_channel.setCurrentIndex(CH3_IDX) + check("CH3 displayed", wait_until(lambda: win._current_ch == CH3_IDX)) + + win.combo_channel.setCurrentIndex(CH1_IDX) + check("CH1 (FFT) computed", wait_until( + lambda: win._current_ch == CH1_IDX and not win._job_running("compute"))) + check("FFT result cached", len(win._fft_cache) > 0, f"{len(win._fft_cache)} keys") + rf_img = win._current_image + check("FFT image is non-degenerate", len(np.unique(rf_img)) > 1, + f"{len(np.unique(rf_img))} unique values") + + print("\nvelocity mode (pure post-multiply, no recompute)") + n_fft_before = len(win._fft_cache) + win.combo_channel.setCurrentIndex(VELOCITY_MODE_IDX) + check("velocity displayed", wait_until( + lambda: win._current_ch == VELOCITY_MODE_IDX and not win._job_running("compute"))) + grating = win.spin_grating_um.value() + check("velocity == freq x grating", + np.allclose(win._current_image, rf_img * grating, atol=1e-3)) + check("velocity reused the cached FFT", len(win._fft_cache) == n_fft_before, + f"{n_fft_before} -> {len(win._fft_cache)}") + check("grating spinbox visible in velocity mode", win.grp_velocity.isVisible()) + + print("\nthreshold change (genuine cache-key change)") + win.combo_channel.setCurrentIndex(CH1_IDX) + wait_until(lambda: not win._job_running("compute")) + dc4 = win._dc_cache[(0, CH4_IDX)] + win.spin_threshold_mv.setValue(float(np.median(dc4))) + win._on_threshold_changed() + check("recomputed at new threshold", wait_until( + lambda: not win._job_running("compute") and len(win._fft_cache) > n_fft_before)) + check("masking zeroed some pixels", + int((win._current_image == 0).sum()) > 0, + f"{int((win._current_image == 0).sum())} of {win._current_image.size}") + + print("\nbackground subtraction toggle") + n_before = len(win._fft_cache) + win.chk_bg_sub.setChecked(False) + check("recomputed without bg-sub", wait_until( + lambda: not win._job_running("compute") and len(win._fft_cache) > n_before)) + win.chk_bg_sub.setChecked(True) + pump(200) + check("returning to bg-sub was a cache hit (no recompute)", + not win._job_running("compute")) + + print("\nROI") + x = s.x_axis_mm(0) + y = s.y_positions_mm(0) + roi = RoiQuad.from_bbox(float(x[1]), float(y[1]), + float(x[-2]), float(y[-2])) + win.image_canvas.set_roi(roi) + pump(120) + check("ROI registered", win.image_canvas.get_roi() is not None) + check("pixel count reported", + "pixels inside" in win.lbl_roi_npix.text() + and win.lbl_roi_npix.text() != "pixels inside: —", + win.lbl_roi_npix.text()) + npix = int(win.lbl_roi_npix.text().split(":")[1]) + check("ROI pixel count is plausible", + 0 < npix <= win._current_image.size, f"{npix}") + check("Export ROI enabled", win.btn_export_roi.isEnabled()) + + csv_path = tmpdir / "roi.csv" + from unittest.mock import patch + with patch("sras_viewer.QFileDialog.getSaveFileName", + return_value=(str(csv_path), "")): + win._on_export_roi_csv() + check("ROI CSV written", csv_path.exists()) + if csv_path.exists(): + body = [l for l in csv_path.read_text().splitlines() if not l.startswith("#")] + check("ROI CSV has header + one line per pixel", + len(body) == npix + 1, f"{len(body)} lines for {npix} pixels") + + img_csv = tmpdir / "img.csv" + with patch("sras_viewer.QFileDialog.getSaveFileName", + return_value=(str(img_csv), "")): + win._on_export_csv() + check("image CSV written", img_csv.exists()) + if img_csv.exists(): + arr = np.loadtxt(img_csv, delimiter=",") + check("image CSV round-trips the displayed image", + arr.shape == win._current_image.shape + and np.allclose(arr, win._current_image, rtol=1e-5, atol=1e-4)) + + print("\nROI survives angle and channel switches") + win.spin_angle.setValue(1) + win._on_view_changed() + wait_until(lambda: not win._job_running("compute")) + check("ROI still present after angle switch", + win.image_canvas.get_roi() is not None) + win.combo_channel.setCurrentIndex(CH4_IDX) + wait_until(lambda: win._current_ch == CH4_IDX) + check("ROI still present after channel switch", + win.image_canvas.get_roi() is not None) + + print("\nangle alignment (Fusion)") + win.spin_angle.setValue(0) + win._on_view_changed() + wait_until(lambda: not win._job_running("compute")) + check("alignment action enabled", win._alignment_act.isEnabled()) + win._on_angle_alignment() + check("alignment completed", wait_until( + lambda: win._alignment_result is not None and not win._job_running("align"), + timeout_ms=60000)) + if win._alignment_result is not None: + r = win._alignment_result + check("transform for every angle", len(r.per_angle) == s.n_angles) + check("canvas is at least as large as any single angle", + 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)), str(r.canvas_shape)) + check("reference angle has zero shift", + r.per_angle[r.ref_angle_idx].shift_mm == (0.0, 0.0)) + check("Aligned View auto-enabled and checked", + win.chk_aligned_view.isEnabled() and win.chk_aligned_view.isChecked()) + pump(200) + check("displayed image is on the alignment canvas", + 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) + check("unchecking returns to the raw per-angle grid", + win.image_canvas._img_shape == s.image_shape(0), + str(win.image_canvas._img_shape)) + + print("\npixel inspector") + win.chk_aligned_view.setChecked(False) + pump(100) + win._on_pixel_clicked(0, 0) + pump(150) + check("waveform hint hidden after a click", win.lbl_wave_hint.isHidden()) + win.combo_channel.setCurrentIndex(CH1_IDX) + wait_until(lambda: not win._job_running("compute")) + win._on_pixel_clicked(1, 1) + pump(150) + check("RF waveform panel rendered", + len(win.wave_canvas.ax_wave.lines) > 0, + f"{len(win.wave_canvas.ax_wave.lines)} lines") + + print("\nshutdown") + win.close() + pump(400) + check("all background jobs released", len(win._jobs) == 0, + f"{list(win._jobs)}") + + print() + unexpected = [e for e in errors if e] + if unexpected: + print(f"status-bar errors seen: {unexpected}") + _failures.append("status-bar errors") + + if _failures: + print(f"{len(_failures)} FAILURE(S): " + ", ".join(_failures)) + return 1 + print("All GUI checks passed.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/test_refactor.py b/tools/test_refactor.py new file mode 100644 index 0000000..0cc50f6 --- /dev/null +++ b/tools/test_refactor.py @@ -0,0 +1,358 @@ +#!/usr/bin/env python3 +"""Behavioural tests for the sras-viewer refactor. + +Covers what the golden-hash harness can't: the v6->v7 cache round-trip +(including block carry-forward), parallel-vs-serial identity, the no-mask +fast path, and the ROI bounding-box mask optimisation. + +Usage: python tools/test_refactor.py [--scratch DIR] +""" + +import argparse +import shutil +import sys +import tempfile +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +import sras_compute as compute # noqa: E402 +from sras_compute import ( # noqa: E402 + cache_file, compute_dc_image, compute_rf_image, dc_image_mv, +) +from sras_format import CH3_IDX, CH4_IDX, SrasFile, adc_to_mv # noqa: E402 +import tools.make_test_sras as gen # noqa: E402 + +_failures: list[str] = [] + + +def check(name: str, ok: bool, detail: str = ""): + print(f" {'PASS' if ok else 'FAIL'} {name}" + (f" — {detail}" if detail else "")) + if not ok: + _failures.append(name) + + +def test_cache_roundtrip(scratch: Path): + """v6 -> v7 for DC, then FFT, asserting the first block survives the + second write (the carry-forward path in write_v7_cache).""" + print("\ncache round-trip (v6 -> v7, both blocks)") + path = scratch / "roundtrip.sras" + gen.write(path, n_angles=3, seed=1, samples_per_frame=64) + + src = SrasFile(str(path)) + check("source is v6", src.version == 6, f"got v{src.version}") + expect_dc3 = [dc_image_mv(src, a, CH3_IDX) for a in range(src.n_angles)] + expect_dc4 = [dc_image_mv(src, a, CH4_IDX) for a in range(src.n_angles)] + expect_fft = [compute_rf_image(src, a, dc_threshold_mv=None, apply_bg_sub=True) + for a in range(src.n_angles)] + + err = cache_file(str(path), "dc", True) + check("dc cache_file succeeded", err == "", err) + + after_dc = SrasFile(str(path)) + check("version flipped to 7", after_dc.version == 7, f"got v{after_dc.version}") + check("dc3 stored for every angle", + all(x is not None for x in after_dc.precomputed_dc3_mv)) + check("dc3 values round-trip", + all(np.allclose(after_dc.precomputed_dc3_mv[a], expect_dc3[a], atol=1e-4) + for a in range(after_dc.n_angles))) + check("dc4 values round-trip", + all(np.allclose(after_dc.precomputed_dc4_mv[a], expect_dc4[a], atol=1e-4) + for a in range(after_dc.n_angles))) + check("no fft block yet", + all(x is None for x in after_dc.precomputed_freq_mhz)) + check("cached images are native float32", + after_dc.precomputed_dc3_mv[0].dtype == np.float32 + and after_dc.precomputed_dc3_mv[0].dtype.byteorder in ("=", "|"), + str(after_dc.precomputed_dc3_mv[0].dtype.byteorder)) + check("cached images are writable", + after_dc.precomputed_dc3_mv[0].flags.writeable) + + err = cache_file(str(path), "fft", True) + check("fft cache_file succeeded", err == "", err) + + both = SrasFile(str(path)) + check("fft stored for every angle", + all(x is not None for x in both.precomputed_freq_mhz)) + check("fft values round-trip", + all(np.allclose(both.precomputed_freq_mhz[a], expect_fft[a], atol=1e-3) + for a in range(both.n_angles))) + check("DC block carried forward through the FFT write", + all(np.allclose(both.precomputed_dc3_mv[a], expect_dc3[a], atol=1e-4) + for a in range(both.n_angles))) + check("bg_sub flag persisted", both.precomputed_bg_sub is True) + + # The fast path must reproduce a fresh compute, and masking must still + # apply on top of a cached (unmasked) image. + fresh = SrasFile(str(path)) + fresh.precomputed_freq_mhz = [None] * fresh.n_angles + dc4 = dc_image_mv(both, 0, CH4_IDX) + thr = float(np.median(dc4)) + check("cached fast path == fresh compute (unmasked)", + np.allclose(compute_rf_image(both, 0, dc_threshold_mv=None, apply_bg_sub=True), + compute_rf_image(fresh, 0, dc_threshold_mv=None, apply_bg_sub=True), + atol=1e-3)) + check("cached fast path == fresh compute (masked)", + np.allclose(compute_rf_image(both, 0, dc_threshold_mv=thr, apply_bg_sub=True), + compute_rf_image(fresh, 0, dc_threshold_mv=thr, apply_bg_sub=True), + atol=1e-3)) + + # Waveform data must be byte-identical to the pre-cache file. + orig = scratch / "roundtrip_orig.sras" + gen.write(orig, n_angles=3, seed=1, samples_per_frame=64) + o, n = SrasFile(str(orig)), SrasFile(str(path)) + check("waveform data untouched by the cache write", + all(np.array_equal(np.asarray(o.data[a]), np.asarray(n.data[a])) + for a in range(o.n_angles))) + + +def test_partial_v7_cache(scratch: Path): + """Only some angles cached: uncached angles must compute, not read zeros. + This is the v5 bug the ragged normalisation fixed, checked via v7.""" + print("\npartial cache (only some angles stored)") + path = scratch / "partial.sras" + gen.write(path, n_angles=3, seed=2, samples_per_frame=64) + + src = SrasFile(str(path)) + expected = [compute_rf_image(src, a, dc_threshold_mv=None, apply_bg_sub=True) + for a in range(src.n_angles)] + partial = [expected[0], None, expected[2]] # angle 1 deliberately absent + src.write_v7_cache(new_freq_mhz=partial, new_bg_sub=True) + + reread = SrasFile(str(path)) + check("angle 1 is not cached", reread.precomputed_freq_mhz[1] is None) + check("angles 0 and 2 are cached", + reread.precomputed_freq_mhz[0] is not None + and reread.precomputed_freq_mhz[2] is not None) + img1 = compute_rf_image(reread, 1, dc_threshold_mv=None, apply_bg_sub=True) + check("uncached angle computes rather than returning zeros", + np.any(img1 != 0) and np.allclose(img1, expected[1], atol=1e-3)) + + +def test_parallel_identity(scratch: Path): + """Forcing 1 worker vs many must give identical output — catches + chunk-boundary and race bugs.""" + print("\nparallel vs serial identity") + path = scratch / "parallel.sras" + # Many rows, so the row loop actually splits into several chunks. + n_rows, n_frames, spf = 48, 9, 256 + gen.write(path, n_angles=1, seed=3, samples_per_frame=spf, + geometry=[(n_rows, n_frames)]) + sras = SrasFile(str(path)) + + saved_budget, saved_workers = compute._TOTAL_BYTES_BUDGET, compute._MAX_WORKERS + try: + # Shrink the budget so chunk_rows collapses to 1 and every row is + # its own chunk — the worst case for boundary bugs. + compute._TOTAL_BYTES_BUDGET = 8 * n_frames * spf * 4 + + compute._MAX_WORKERS = 1 + dc_serial = compute_dc_image(sras, 0, CH4_IDX) + rf_serial = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True) + dc4 = adc_to_mv(dc_serial, *sras.cal(CH4_IDX)) + thr = float(np.median(dc4)) + rf_masked_serial = compute_rf_image(sras, 0, dc_threshold_mv=thr, + apply_bg_sub=True) + + chunk_rows, n_workers = compute._plan_chunks( + n_rows, n_frames, spf, live_multiplier=compute._FFT_LIVE_MULTIPLIER) + check("serial plan uses 1 worker", n_workers == 1, f"chunk_rows={chunk_rows}") + check("work actually splits into multiple chunks", chunk_rows < n_rows, + f"chunk_rows={chunk_rows} of {n_rows} rows") + + compute._MAX_WORKERS = 8 + chunk_rows, n_workers = compute._plan_chunks( + n_rows, n_frames, spf, live_multiplier=compute._FFT_LIVE_MULTIPLIER) + check("parallel plan uses >1 worker", n_workers > 1, + f"chunk_rows={chunk_rows} workers={n_workers}") + + dc_par = compute_dc_image(sras, 0, CH4_IDX) + rf_par = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True) + rf_masked_par = compute_rf_image(sras, 0, dc_threshold_mv=thr, apply_bg_sub=True) + + check("dc image identical", np.array_equal(dc_serial, dc_par)) + check("rf image identical (unmasked)", np.array_equal(rf_serial, rf_par)) + check("rf image identical (masked)", + np.array_equal(rf_masked_serial, rf_masked_par)) + finally: + compute._TOTAL_BYTES_BUDGET, compute._MAX_WORKERS = saved_budget, saved_workers + + +def test_nomask_equals_low_threshold(scratch: Path): + """dc_threshold_mv=None must equal a threshold below every pixel, while + skipping the CH4 read.""" + print("\nno-mask path") + path = scratch / "nomask.sras" + gen.write(path, n_angles=2, seed=4, samples_per_frame=128) + sras = SrasFile(str(path)) + for a in range(sras.n_angles): + none_img = compute_rf_image(sras, a, dc_threshold_mv=None, apply_bg_sub=True) + low_img = compute_rf_image(sras, a, dc_threshold_mv=-1e9, apply_bg_sub=True) + check(f"angle {a}: None == -1e9 threshold", + np.array_equal(none_img, low_img)) + check(f"angle {a}: image is non-degenerate", + len(np.unique(none_img)) > 1, f"{len(np.unique(none_img))} unique") + + +def test_roi_mask(): + """The bbox-restricted mask must equal a full-grid point-in-polygon test.""" + print("\nROI mask (bbox fast path vs full grid)") + from matplotlib.path import Path as MplPath + from sras_viewer import RoiQuad + + rng = np.random.default_rng(0) + x = np.linspace(-2.0, 3.0, 137) + y = np.linspace(1.0, 4.0, 91) + + cases = { + "axis-aligned rect": np.array([[0.0, 1.5], [1.0, 1.5], [1.0, 3.0], [0.0, 3.0]]), + "skewed quad": np.array([[-0.5, 1.2], [1.7, 1.9], [1.2, 3.4], [-1.0, 2.6]]), + "entirely outside": np.array([[8.0, 8.0], [9.0, 8.0], [9.0, 9.0], [8.0, 9.0]]), + "covers whole grid": np.array([[-9.0, -9.0], [9.0, -9.0], [9.0, 9.0], [-9.0, 9.0]]), + "straddles left edge": np.array([[-4.0, 2.0], [0.5, 2.0], [0.5, 3.0], [-4.0, 3.0]]), + } + for _ in range(5): + cases[f"random {_}"] = rng.uniform([-2.5, 0.5], [3.5, 4.5], size=(4, 2)) + + for name, pts in cases.items(): + roi = RoiQuad(pts) + fast = roi.mask_for_grid(x, y) + X, Y = np.meshgrid(x.astype(np.float64), y.astype(np.float64)) + slow = MplPath(pts).contains_points( + np.column_stack([X.ravel(), Y.ravel()])).reshape(X.shape) + check(f"{name} ({int(slow.sum())} px inside)", np.array_equal(fast, slow)) + + # Descending y axis (images are stored top-down in some scans). + roi = RoiQuad(cases["skewed quad"]) + y_desc = y[::-1] + fast = roi.mask_for_grid(x, y_desc) + X, Y = np.meshgrid(x.astype(np.float64), y_desc.astype(np.float64)) + slow = MplPath(cases["skewed quad"]).contains_points( + np.column_stack([X.ravel(), Y.ravel()])).reshape(X.shape) + check("descending y axis", np.array_equal(fast, slow)) + + +def test_legacy_parse_and_average(scratch: Path): + """v2-v4 parsing plus the sras_average.py rewrite (which now streams via + SrasFile rather than slurping the whole file).""" + import subprocess + print("\nlegacy formats (v2-v4) and sras_average") + repo = Path(__file__).resolve().parent.parent + + for version in (2, 3, 4): + path = scratch / f"legacy_v{version}.sras" + meta = gen.write_legacy(path, version=version, n_angles=2, n_rows=4, + n_frames=12, samples_per_frame=32, seed=version) + s = SrasFile(str(path)) + check(f"v{version} parses", s.version == version, f"got v{s.version}") + check(f"v{version} geometry uniform across angles", + list(s.n_rows) == [4, 4] and list(s.n_frames) == [12, 12], + f"rows={list(s.n_rows)} frames={list(s.n_frames)}") + check(f"v{version} waveform data matches what was written", + all(np.array_equal(np.asarray(s.data[a]), meta["data"][a]) + for a in range(s.n_angles))) + check(f"v{version} background {'present' if version >= 4 else 'absent'}", + (s.background is not None) == (version >= 4)) + check(f"v{version} precomputed stores are ragged lists", + isinstance(s.precomputed_freq_mhz, list) + and len(s.precomputed_freq_mhz) == s.n_angles) + # DC image must equal a direct mean of the known input. + expect = meta["data"][0][:, CH3_IDX, :, :].astype(np.float64).mean(axis=-1) + check(f"v{version} DC image equals a direct mean", + np.allclose(compute_dc_image(s, 0, CH3_IDX), expect, atol=1e-3)) + + src = scratch / "legacy_v4.sras" + meta = gen.write_legacy(src, version=4, n_angles=2, n_rows=4, n_frames=12, + samples_per_frame=32, seed=4) + dst = scratch / "legacy_v4_avg.sras" + if dst.exists(): + dst.unlink() + proc = subprocess.run( + [sys.executable, str(repo / "sras_average.py"), str(src), str(dst), "--n", "4"], + capture_output=True, text=True, cwd=repo) + check("sras_average ran", proc.returncode == 0, + (proc.stderr or proc.stdout).strip()[-200:]) + + if dst.exists(): + avg = SrasFile(str(dst)) + check("averaged file parses", avg.version == 4) + check("frame count divided by 4", + list(avg.n_frames) == [3, 3], f"{list(avg.n_frames)}") + check("angles/rows/channels unchanged", + avg.n_angles == 2 and list(avg.n_rows) == [4, 4] + and avg.n_channels == meta["n_channels"]) + check("calibration preserved", + np.allclose(avg.ch_ymult_mv, SrasFile(str(src)).ch_ymult_mv)) + check("background preserved", + np.array_equal(avg.background, SrasFile(str(src)).background)) + src_data = meta["data"] + expect0 = src_data[0][:, :, 0:4, :].astype(np.float32).mean(axis=2).astype(np.int16) + check("first averaged group equals the mean of its 4 source frames", + np.array_equal(np.asarray(avg.data[0])[:, :, 0, :], expect0)) + + # Remainder handling: 12 frames / 5 -> 2 full groups + 1 partial. + dst2 = scratch / "legacy_v4_avg5.sras" + subprocess.run([sys.executable, str(repo / "sras_average.py"), + str(src), str(dst2), "--n", "5"], + capture_output=True, text=True, cwd=repo) + if dst2.exists(): + check("partial trailing group kept by default", + list(SrasFile(str(dst2)).n_frames) == [3, 3], + f"{list(SrasFile(str(dst2)).n_frames)}") + dst3 = scratch / "legacy_v4_avg5d.sras" + subprocess.run([sys.executable, str(repo / "sras_average.py"), + str(src), str(dst3), "--n", "5", "--discard-remainder"], + capture_output=True, text=True, cwd=repo) + if dst3.exists(): + check("--discard-remainder drops the partial group", + list(SrasFile(str(dst3)).n_frames) == [2, 2], + f"{list(SrasFile(str(dst3)).n_frames)}") + + +def test_unsupported_version_reported(scratch: Path): + """cache_file must report, not raise, for a file it can't handle.""" + print("\nerror reporting") + bogus = scratch / "bogus.sras" + bogus.write_bytes(b"SRAS" + bytes([99]) + b"\x00" * 200) + err = cache_file(str(bogus), "dc", True) + check("bad version returns an error string", bool(err), err) + missing = cache_file(str(scratch / "does_not_exist.sras"), "dc", True) + check("missing file returns an error string", bool(missing), missing) + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--scratch") + args = p.parse_args() + + tmp = None + if args.scratch: + scratch = Path(args.scratch) + scratch.mkdir(parents=True, exist_ok=True) + else: + tmp = tempfile.mkdtemp(prefix="sras_test_") + scratch = Path(tmp) + + try: + test_cache_roundtrip(scratch) + test_partial_v7_cache(scratch) + test_parallel_identity(scratch) + test_nomask_equals_low_threshold(scratch) + test_roi_mask() + test_legacy_parse_and_average(scratch) + test_unsupported_version_reported(scratch) + finally: + if tmp: + shutil.rmtree(tmp, ignore_errors=True) + + print() + if _failures: + print(f"{len(_failures)} FAILURE(S): " + ", ".join(_failures)) + sys.exit(1) + print("All checks passed.") + + +if __name__ == "__main__": + main() From 013c7739a14addd2be41265e003704b8d71416e9 Mon Sep 17 00:00:00 2001 From: Thomas Ales Date: Fri, 31 Jul 2026 08:18:57 -0500 Subject: [PATCH 05/10] Fix clipped and overlapping controls in the side panels The left panel was a plain fixed-width widget sized by the window, so once the Scan Info notes grew past the window height Qt compressed every control below its minimum: the DC threshold spin box was sliced in half, the ADC hint collided with it, and the cache note was cut off mid-line. Word-wrapped labels made it worse by advertising a single-line minimum height, which clipped the filename and the format/cache notes. - Put both side panels in fixed-width QScrollAreas so a short window scrolls instead of crushing the controls. - Add _wrap_label(), which enables height-for-width so wrapped labels report their real height; use it for the filename, Scan Info rows, format/cache notes, DC-precompute status, ADC hint and ROI readouts. - Replace the ad-hoc QHBoxLayout label/field rows with a shared QFormLayout (Angle, Channel, DC threshold, Grating size, Colormap, min, max) so labels align and fields keep a usable minimum width. - Add _group() for consistent group-box margins; widen the left column 260 -> 288 px so the channel combo and checkbox text stop eliding. - Keep the canvases usable: non-collapsible splitter, minimum canvas heights and a 960x560 window minimum. Co-Authored-By: Claude Opus 5 --- sras_viewer.py | 209 +++++++++++++++++++++++++++++++++---------------- 1 file changed, 142 insertions(+), 67 deletions(-) diff --git a/sras_viewer.py b/sras_viewer.py index 77a9a0c..c09d578 100644 --- a/sras_viewer.py +++ b/sras_viewer.py @@ -28,9 +28,9 @@ from PyQt6.QtCore import QObject, Qt, QThread, pyqtSignal from PyQt6.QtGui import QAction from PyQt6.QtWidgets import ( QApplication, QButtonGroup, QCheckBox, QComboBox, QDialog, QDialogButtonBox, - QDoubleSpinBox, QFileDialog, QFrame, QGroupBox, QHBoxLayout, QLabel, - QMainWindow, QProgressDialog, QPushButton, QRadioButton, QSizePolicy, - QSpinBox, QSplitter, QVBoxLayout, QWidget, + QDoubleSpinBox, QFileDialog, QFormLayout, QFrame, QGroupBox, QHBoxLayout, + QLabel, QMainWindow, QProgressDialog, QPushButton, QRadioButton, QScrollArea, + QSizePolicy, QSpinBox, QSplitter, QVBoxLayout, QWidget, ) import sras_compute as compute @@ -78,6 +78,79 @@ _CSS_MUTED = "color: #888; font-size: 11px;" _CSS_WARN = "color: #e07000; font-size: 11px;" _CSS_BUSY = "color: #4a90d9; font-size: 11px;" +# Side-panel column widths (the scroll areas that hold the controls). +_LEFT_PANEL_W = 288 +_RIGHT_PANEL_W = 272 + +# Minimum width for a spin box so its value + suffix are never clipped. +_SPIN_MIN_W = 96 + + +# --------------------------------------------------------------------------- +# Small layout helpers +# --------------------------------------------------------------------------- + +def _wrap_label(text: str = "", css: str | None = None) -> QLabel: + """A word-wrapped QLabel that reports its *wrapped* height to the layout. + + A plain word-wrapped QLabel advertises a single-line minimum height, so in a + fixed-width column the layout happily shrinks it and the extra lines get + clipped. Enabling height-for-width makes the box layout ask for the real + height at the column's width instead. + """ + lbl = QLabel(text) + lbl.setWordWrap(True) + sp = lbl.sizePolicy() + sp.setVerticalPolicy(QSizePolicy.Policy.Minimum) + sp.setHeightForWidth(True) + lbl.setSizePolicy(sp) + if css: + lbl.setStyleSheet(css) + return lbl + + +def _group(title: str) -> tuple[QGroupBox, QVBoxLayout]: + """A group box with consistent, non-cramped internal margins.""" + grp = QGroupBox(title) + lay = QVBoxLayout(grp) + lay.setContentsMargins(10, 8, 10, 10) + lay.setSpacing(6) + return grp, lay + + +def _form() -> QFormLayout: + """A label/field form layout for a narrow side panel.""" + form = QFormLayout() + form.setContentsMargins(0, 0, 0, 0) + form.setHorizontalSpacing(8) + form.setVerticalSpacing(6) + form.setLabelAlignment(Qt.AlignmentFlag.AlignRight + | Qt.AlignmentFlag.AlignVCenter) + form.setFormAlignment(Qt.AlignmentFlag.AlignLeft | Qt.AlignmentFlag.AlignTop) + form.setFieldGrowthPolicy( + QFormLayout.FieldGrowthPolicy.AllNonFixedFieldsGrow) + form.setRowWrapPolicy(QFormLayout.RowWrapPolicy.DontWrapRows) + return form + + +def _scroll_panel(inner: QWidget, width: int) -> QScrollArea: + """Put a side panel in a fixed-width scroll area. + + Without this the panels are sized by the window: a short window squeezes the + controls past their minimum heights, which is what makes text overlap the + widget below it. Scrolling keeps every control at its natural size. + """ + area = QScrollArea() + area.setWidget(inner) + area.setWidgetResizable(True) + area.setFrameShape(QFrame.Shape.NoFrame) + area.setHorizontalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAlwaysOff) + area.setVerticalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAsNeeded) + area.setFixedWidth(width) + area.viewport().setAutoFillBackground(False) + inner.setAutoFillBackground(False) + return area + # --------------------------------------------------------------------------- # ROI (free quadrilateral in data coordinates) @@ -674,6 +747,7 @@ class SrasViewerWindow(QMainWindow): super().__init__() self.setWindowTitle("SRAS Scan Viewer") self.resize(1560, 840) + self.setMinimumSize(960, 560) self.setAcceptDrops(True) self._sras: SrasFile | None = None @@ -791,69 +865,69 @@ class SrasViewerWindow(QMainWindow): def _build_left_panel(self) -> QWidget: panel = QWidget() - panel.setFixedWidth(260) panel_layout = QVBoxLayout(panel) panel_layout.setContentsMargins(0, 0, 0, 0) - panel_layout.setSpacing(6) + panel_layout.setSpacing(8) # ---- File ------------------------------------------------------- - grp_file = QGroupBox("File") - fl = QVBoxLayout(grp_file) + grp_file, fl = _group("File") self.btn_open = QPushButton("Open .sras…") self.btn_open.clicked.connect(self._on_open) - self.lbl_filename = QLabel("No file loaded") - self.lbl_filename.setWordWrap(True) - self.lbl_filename.setStyleSheet(_CSS_MUTED) + self.lbl_filename = _wrap_label("No file loaded", _CSS_MUTED) fl.addWidget(self.btn_open) fl.addWidget(self.lbl_filename) panel_layout.addWidget(grp_file) # ---- Scan info -------------------------------------------------- - grp_info = QGroupBox("Scan Info") - il = QVBoxLayout(grp_info) + grp_info, il = _group("Scan Info") + il.setSpacing(3) self._info = {} for key in ("Angles", "Rows", "Frames / row", "Samples / frame", "Sample rate", "X start", "Pixel Δx", "Laser freq"): - lbl = QLabel(f"{key}: —") - lbl.setWordWrap(True) - lbl.setStyleSheet(_CSS_INFO) + lbl = _wrap_label(f"{key}: —", _CSS_INFO) il.addWidget(lbl) self._info[key] = lbl - self.lbl_frame_warn = QLabel("") # frame-count / format notes - self.lbl_frame_warn.setWordWrap(True) - self.lbl_frame_warn.setStyleSheet(_CSS_WARN) + # frame-count / format notes + self.lbl_frame_warn = _wrap_label("", _CSS_WARN) il.addWidget(self.lbl_frame_warn) - self.lbl_dc_precompute = QLabel("") # background DC-precompute progress - self.lbl_dc_precompute.setWordWrap(True) - self.lbl_dc_precompute.setStyleSheet(_CSS_BUSY) + # background DC-precompute progress + self.lbl_dc_precompute = _wrap_label("", _CSS_BUSY) il.addWidget(self.lbl_dc_precompute) panel_layout.addWidget(grp_info) # ---- View settings ---------------------------------------------- - grp_view = QGroupBox("View Settings") - vl = QVBoxLayout(grp_view) + grp_view, vl = _group("View Settings") + + view_form = _form() - ar = QHBoxLayout() - ar.addWidget(QLabel("Angle:")) self.spin_angle = QSpinBox() self.spin_angle.setRange(0, 0) self.spin_angle.setEnabled(False) + self.spin_angle.setMinimumWidth(64) self.spin_angle.editingFinished.connect(self._on_view_changed) self.lbl_angle_deg = QLabel("—") + angle_field = QWidget() + ar = QHBoxLayout(angle_field) + ar.setContentsMargins(0, 0, 0, 0) + ar.setSpacing(6) ar.addWidget(self.spin_angle) ar.addWidget(self.lbl_angle_deg) - vl.addLayout(ar) + ar.addStretch() + view_form.addRow("Angle:", angle_field) - cr = QHBoxLayout() - cr.addWidget(QLabel("Channel:")) self.combo_channel = QComboBox() self.combo_channel.addItems(CH_LABELS) self.combo_channel.setEnabled(False) + self.combo_channel.setSizePolicy(QSizePolicy.Policy.Expanding, + QSizePolicy.Policy.Fixed) + self.combo_channel.setSizeAdjustPolicy( + QComboBox.SizeAdjustPolicy.AdjustToMinimumContentsLengthWithIcon) + self.combo_channel.setMinimumContentsLength(12) self.combo_channel.currentIndexChanged.connect(self._on_channel_changed) - cr.addWidget(self.combo_channel) - vl.addLayout(cr) + view_form.addRow("Channel:", self.combo_channel) + vl.addLayout(view_form) sep = QFrame() sep.setFrameShape(QFrame.Shape.HLine) @@ -861,10 +935,8 @@ class SrasViewerWindow(QMainWindow): vl.addWidget(sep) # DC threshold (for RF / CH1 masking) - self.grp_threshold = QGroupBox("RF Mask Threshold (CH1 only)") - tl = QVBoxLayout(self.grp_threshold) - thr_row = QHBoxLayout() - thr_row.addWidget(QLabel("DC threshold:")) + self.grp_threshold, tl = _group("RF Mask Threshold (CH1 only)") + thr_form = _form() self.spin_threshold_mv = QDoubleSpinBox() self.spin_threshold_mv.setRange(-500.0, 500.0) self.spin_threshold_mv.setDecimals(3) @@ -872,11 +944,12 @@ class SrasViewerWindow(QMainWindow): self.spin_threshold_mv.setSuffix(" mV") self.spin_threshold_mv.setValue(50.0) self.spin_threshold_mv.setEnabled(False) + self.spin_threshold_mv.setMinimumWidth(_SPIN_MIN_W) self.spin_threshold_mv.editingFinished.connect(self._on_threshold_changed) - thr_row.addWidget(self.spin_threshold_mv) - tl.addLayout(thr_row) - self.lbl_threshold_adc = QLabel(f"≈ {mv_to_adc(50.0):.1f} ADC counts") - self.lbl_threshold_adc.setStyleSheet(_CSS_MUTED) + thr_form.addRow("DC threshold:", self.spin_threshold_mv) + tl.addLayout(thr_form) + self.lbl_threshold_adc = _wrap_label( + f"≈ {mv_to_adc(50.0):.1f} ADC counts", _CSS_MUTED) tl.addWidget(self.lbl_threshold_adc) vl.addWidget(self.grp_threshold) @@ -913,8 +986,7 @@ class SrasViewerWindow(QMainWindow): panel_layout.addWidget(grp_view) # ---- ROI --------------------------------------------------------- - grp_roi = QGroupBox("ROI (Region of Interest)") - rl = QVBoxLayout(grp_roi) + grp_roi, rl = _group("ROI (Region of Interest)") self.btn_draw_roi = QPushButton("Draw ROI") self.btn_draw_roi.setCheckable(True) @@ -943,24 +1015,26 @@ class SrasViewerWindow(QMainWindow): self.btn_export_roi.clicked.connect(self._on_export_roi_csv) rl.addWidget(self.btn_export_roi) - self.lbl_roi_center = QLabel("centroid: —") - self.lbl_roi_size = QLabel("bbox: —") - self.lbl_roi_npix = QLabel("pixels inside: —") + self.lbl_roi_center = _wrap_label("centroid: —", _CSS_HINT) + self.lbl_roi_size = _wrap_label("bbox: —", _CSS_HINT) + self.lbl_roi_npix = _wrap_label("pixels inside: —", _CSS_HINT) for lbl in (self.lbl_roi_center, self.lbl_roi_size, self.lbl_roi_npix): - lbl.setStyleSheet(_CSS_HINT) rl.addWidget(lbl) panel_layout.addWidget(grp_roi) panel_layout.addStretch() - return panel + return _scroll_panel(panel, _LEFT_PANEL_W) def _build_canvases(self) -> QWidget: splitter = QSplitter(Qt.Orientation.Vertical) + splitter.setChildrenCollapsible(False) img_widget = QWidget() img_vl = QVBoxLayout(img_widget) img_vl.setContentsMargins(0, 0, 0, 0) + img_vl.setSpacing(4) self.image_canvas = ImageCanvas() + self.image_canvas.setMinimumHeight(220) self.image_canvas.pixel_clicked.connect(self._on_pixel_clicked) self.image_canvas.roi_changed.connect(self._update_roi_ui) self.image_canvas.draw_mode_changed.connect(self._on_draw_mode_changed) @@ -971,24 +1045,26 @@ class SrasViewerWindow(QMainWindow): wave_widget = QWidget() wave_vl = QVBoxLayout(wave_widget) wave_vl.setContentsMargins(0, 0, 0, 0) + wave_vl.setSpacing(4) self.lbl_wave_hint = QLabel( "Click a pixel in the image above to inspect its waveform.") self.lbl_wave_hint.setAlignment(Qt.AlignmentFlag.AlignCenter) self.lbl_wave_hint.setStyleSheet(_CSS_MUTED) self.wave_canvas = WaveformCanvas() + self.wave_canvas.setMinimumHeight(150) wave_vl.addWidget(self.lbl_wave_hint) wave_vl.addWidget(self.wave_canvas) splitter.addWidget(wave_widget) + splitter.setStretchFactor(0, 3) + splitter.setStretchFactor(1, 1) splitter.setSizes([580, 250]) return splitter def _build_right_panel(self) -> QWidget: # Velocity settings (visible only in velocity mode) - self.grp_velocity = QGroupBox("Velocity Settings (CH1 only)") - vel_l = QVBoxLayout(self.grp_velocity) - grat_row = QHBoxLayout() - grat_row.addWidget(QLabel("Grating size:")) + self.grp_velocity, vel_l = _group("Velocity Settings (CH1 only)") + vel_form = _form() self.spin_grating_um = QDoubleSpinBox() self.spin_grating_um.setRange(0.1, 1000.0) self.spin_grating_um.setDecimals(2) @@ -996,53 +1072,52 @@ class SrasViewerWindow(QMainWindow): self.spin_grating_um.setSuffix(" µm") self.spin_grating_um.setValue(25) self.spin_grating_um.setEnabled(False) + self.spin_grating_um.setMinimumWidth(_SPIN_MIN_W) self.spin_grating_um.editingFinished.connect(self._on_grating_changed) - grat_row.addWidget(self.spin_grating_um) - vel_l.addLayout(grat_row) - lbl_formula = QLabel("v (m/s) = freq (MHz) × grating (µm)") - lbl_formula.setStyleSheet("font-size: 10px; color: #888;") - vel_l.addWidget(lbl_formula) + vel_form.addRow("Grating size:", self.spin_grating_um) + vel_l.addLayout(vel_form) + vel_l.addWidget(_wrap_label("v (m/s) = freq (MHz) × grating (µm)", + "font-size: 10px; color: #888;")) self.grp_velocity.setVisible(False) - grp_display = QGroupBox("Display Options") - dl = QVBoxLayout(grp_display) + grp_display, dl = _group("Display Options") - cmr = QHBoxLayout() - cmr.addWidget(QLabel("Colormap:")) + cmap_form = _form() self.combo_cmap = QComboBox() self.combo_cmap.addItems(CMAPS) self.combo_cmap.setCurrentText("gray") self.combo_cmap.setEnabled(False) + self.combo_cmap.setSizePolicy(QSizePolicy.Policy.Expanding, + QSizePolicy.Policy.Fixed) self.combo_cmap.currentIndexChanged.connect(self._on_cmap_changed) - cmr.addWidget(self.combo_cmap) - dl.addLayout(cmr) + cmap_form.addRow("Colormap:", self.combo_cmap) + dl.addLayout(cmap_form) self.chk_auto = QCheckBox("Auto-scale colormap") self.chk_auto.setChecked(True) self.chk_auto.toggled.connect(self._on_autoscale_toggled) dl.addWidget(self.chk_auto) + range_form = _form() for label, attr in (("min:", "spin_vmin"), ("max:", "spin_vmax")): - row = QHBoxLayout() - row.addWidget(QLabel(label)) spin = QDoubleSpinBox() spin.setRange(-1e9, 1e9) spin.setDecimals(4) spin.setEnabled(False) + spin.setMinimumWidth(_SPIN_MIN_W) spin.editingFinished.connect(self._on_manual_range_changed) setattr(self, attr, spin) - row.addWidget(spin) - dl.addLayout(row) + range_form.addRow(label, spin) + dl.addLayout(range_form) right_panel = QWidget() - right_panel.setFixedWidth(270) layout = QVBoxLayout(right_panel) layout.setContentsMargins(0, 0, 0, 0) - layout.setSpacing(6) + layout.setSpacing(8) layout.addWidget(self.grp_velocity) layout.addWidget(grp_display) layout.addStretch() - return right_panel + return _scroll_panel(right_panel, _RIGHT_PANEL_W) def _build_menus(self): menubar = self.menuBar() From 3c835f45927d28865f6625f788607dd9119c9898 Mon Sep 17 00:00:00 2001 From: Thomas Ales Date: Fri, 31 Jul 2026 09:19:17 -0500 Subject: [PATCH 06/10] Add manual angle alignment mode with overlay, keyboard nudge, and save/clear Automatic alignment's phase-correlation translation search is unreliable, so add a Fusion -> Manual Alignment dialog: every angle's CH4 threshold mask overlaid at once with distinct colors/opacity, a selector for the active angle, arrow keys to nudge translation and Q/E to nudge rotation, an Auto De-rotate button that snaps to the known scan angle, and Save/Clear controls backed by a JSON sidecar that's restored automatically on reload. Co-Authored-By: Claude Sonnet 5 --- sras_compute.py | 265 +++++++++++++++++- sras_viewer.py | 663 +++++++++++++++++++++++++++++++++++++++++++++- sras_workers.py | 42 +++ tools/test_gui.py | 142 +++++++++- 4 files changed, 1092 insertions(+), 20 deletions(-) diff --git a/sras_compute.py b/sras_compute.py index fc01b06..37f7b16 100644 --- a/sras_compute.py +++ b/sras_compute.py @@ -6,9 +6,11 @@ multiprocessing child can import it without loading Qt or matplotlib — which matters because Python 3.14 on macOS spawns rather than forks. """ +import json import os from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass +from pathlib import Path import numpy as np import scipy.fft as scipy_fft @@ -397,6 +399,21 @@ class AlignmentResult: per_angle: dict[int, AngleTransform] +@dataclass +class ManualAngleParams: + """One angle's manual-alignment state, independent of any canvas. + + rotation_deg/shift_mm are exactly AngleTransform's non-derived fields — + the pair a canvas-bound AngleTransform's matrix/offset get built from + once a canvas is decided (build_manual_alignment). Defaults to identity + (no rotation, no shift): a fresh angle with no prior alignment is shown + raw, exactly as scanned — the same "fully unaligned" state Clear + Alignment resets back to. + """ + rotation_deg: float = 0.0 + shift_mm: tuple[float, float] = (0.0, 0.0) + + def _pixel_pitch_mm(sras: SrasFile, angle_idx: int) -> tuple[float, float]: """(dx, dy) mm/pixel for one angle: dx is the file-wide constant pixel_x_mm; dy is this angle's own row spacing (assumed uniform, the same @@ -429,10 +446,18 @@ def _theta_deg(sras: SrasFile, angle_idx: int, ref_idx: int) -> float: def _corners_in_ref_frame(sras: SrasFile, angle_idx: int, ref_idx: int, - shift_mm=(0.0, 0.0)) -> np.ndarray: + shift_mm=(0.0, 0.0), + theta_deg: float | None = None) -> np.ndarray: """Angle *angle_idx*'s bbox corners, rotated about its own centroid into - the reference frame and translated by *shift_mm*. Shape (4, 2).""" - R = _rotation_matrix(_theta_deg(sras, angle_idx, ref_idx)) + the reference frame and translated by *shift_mm*. Shape (4, 2). + + theta_deg overrides the analytic angles_deg-derived rotation used by + default — the manual-alignment path (union_canvas_mm) passes a + user-chosen rotation here (which may differ from the known scan-angle + delta) without needing a parallel code path. + """ + theta = _theta_deg(sras, angle_idx, ref_idx) if theta_deg is None else theta_deg + R = _rotation_matrix(theta) c_a = np.array(_bbox_center_mm(sras, angle_idx)) c_ref = np.array(_bbox_center_mm(sras, ref_idx)) shift = np.asarray(shift_mm, dtype=np.float64) @@ -443,7 +468,8 @@ def _corners_in_ref_frame(sras: SrasFile, angle_idx: int, ref_idx: int, def _build_affine_canvas_to_raw(sras: SrasFile, angle_idx: int, ref_idx: int, shift_mm: tuple[float, float], canvas_dx: float, canvas_dy: float, - canvas_origin_mm: tuple[float, float] + canvas_origin_mm: tuple[float, float], + theta_deg: float | None = None ) -> tuple[np.ndarray, np.ndarray]: """matrix, offset s.t. raw_index = matrix @ [row_out, col_out] + offset, matching scipy.ndimage.affine_transform's output->input convention. @@ -453,13 +479,16 @@ def _build_affine_canvas_to_raw(sras: SrasFile, angle_idx: int, ref_idx: int, [lx;ly] = R(theta)^T @ ([X;Y]-c_ref-shift) + c_a # undo rotation+shift -> angle a's local mm [row;col] = D @ ([lx;ly] - [x_start_a; y0_a]) # local mm -> angle a's raw idx - where theta = angles_deg[angle_idx] - angles_deg[ref_idx], c_ref/c_a are - each angle's own raw-bbox mm centroid (the rotation pivot — this keeps - rotated content centered, minimizing required canvas padding), and - A_out/D are the index<->mm scaling matrices for the canvas pitch and - this angle's own native pitch respectively. + where theta = angles_deg[angle_idx] - angles_deg[ref_idx] (unless + *theta_deg* overrides it — see _corners_in_ref_frame's note, used by the + manual-alignment path), c_ref/c_a are each angle's own raw-bbox mm + centroid (the rotation pivot — this keeps rotated content centered, + minimizing required canvas padding), and A_out/D are the index<->mm + scaling matrices for the canvas pitch and this angle's own native pitch + respectively. """ - Rinv = _rotation_matrix(_theta_deg(sras, angle_idx, ref_idx)).T + theta = _theta_deg(sras, angle_idx, ref_idx) if theta_deg is None else theta_deg + Rinv = _rotation_matrix(theta).T cx_a, cy_a = _bbox_center_mm(sras, angle_idx) cx_ref, cy_ref = _bbox_center_mm(sras, ref_idx) dx_a, dy_a = _pixel_pitch_mm(sras, angle_idx) @@ -648,3 +677,219 @@ def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int, # Back-compat alias for the pre-split private name (used by tooling). _compute_angle_alignment = compute_angle_alignment + + +# --------------------------------------------------------------------------- +# Manual alignment (Fusion menu -> Manual Alignment... dialog) +# +# Skips compute_angle_alignment's mask + phase-correlation search entirely: +# every angle's rotation_deg/shift_mm is supplied directly by the caller +# (nudged by eye against a live multi-angle mask overlay, or pre-seeded from +# a previous compute_angle_alignment run or a saved sidecar). Building the +# final AlignmentResult from already-known per-angle parameters is pure +# closed-form matrix math (union_canvas_mm + _build_affine_canvas_to_raw) +# with no per-pixel image work at all, so build_manual_alignment is cheap +# enough to call synchronously on the GUI thread on every edit. The only +# genuinely expensive per-pixel operation anywhere in this flow is +# reproject_mask, and only ManualAlignmentDialog's own downsampled preview +# calls that per keystroke — see that class's docstring for how it limits +# each nudge to reprojecting only the actively-edited angle. +# --------------------------------------------------------------------------- + +def union_canvas_mm(sras: SrasFile, ref_angle_idx: int, dx: float, dy: float, + per_angle_params: dict[int, ManualAngleParams], + margin_frac: float = 0.0 + ) -> tuple[tuple[float, float], tuple[int, int]]: + """Shared-canvas origin (mm) and (n_rows, n_cols) at pitch (dx, dy) that + contains every angle's footprint after applying its own rotation+shift — + the generalisation of compute_angle_alignment's Step 3 to arbitrary (not + just _theta_deg-analytic) per-angle rotation. Angles missing from + per_angle_params default to identity (e.g. a sidecar saved before a + rescan added more angles). + + margin_frac pads the box on every side: 0 for a final canvas (this then + reproduces compute_angle_alignment's own Step-3 math exactly, when every + angle's rotation_deg equals the analytic delta and shift_mm matches); + nonzero for ManualAlignmentDialog's downsampled preview canvas, which + needs headroom so an ordinary translation nudge of the active angle never + has to trigger a full canvas resize (see that class's docstring — an + extreme nudge can still, in principle, push content past this padding; + accepted as a known edge case, same as _phase_correlate_shift's + wraparound risk note). + """ + n = sras.n_angles + corners = np.vstack([ + _corners_in_ref_frame( + sras, a, ref_angle_idx, + per_angle_params.get(a, ManualAngleParams()).shift_mm, + theta_deg=per_angle_params.get(a, ManualAngleParams()).rotation_deg) + for a in range(n)]) + x_min, y_min = corners.min(axis=0) + x_max, y_max = corners.max(axis=0) + if margin_frac: + pad_x, pad_y = (x_max - x_min) * margin_frac, (y_max - y_min) * margin_frac + x_min, x_max = x_min - pad_x, x_max + pad_x + y_min, y_max = y_min - pad_y, y_max + pad_y + n_cols = int(np.ceil((x_max - x_min) / dx)) + 1 + n_rows = int(np.ceil((y_max - y_min) / abs(dy))) + 1 + origin = (float(x_min), float(y_min if dy > 0 else y_max)) + return origin, (n_rows, n_cols) + + +def reproject_mask(sras: SrasFile, angle_idx: int, ref_angle_idx: int, + mask: np.ndarray, rotation_deg: float, + shift_mm: tuple[float, float], + canvas_dx: float, canvas_dy: float, + canvas_origin_mm: tuple[float, float], + canvas_shape: tuple[int, int]) -> np.ndarray: + """Resample one angle's binary/float mask onto an arbitrary canvas via an + explicit rotation+shift — the single building block + ManualAlignmentDialog's live preview repeatedly calls (once per + keystroke, for only the actively-nudged angle), since it bypasses + compute_angle_alignment's phase-correlation search entirely and just + takes rotation_deg/shift_mm as given. order=0 (nearest) matches + apply_alignment's own reasoning: a binary mask must never be blended with + zero-padding. + """ + matrix, offset = _build_affine_canvas_to_raw( + sras, angle_idx, ref_angle_idx, shift_mm, canvas_dx, canvas_dy, + canvas_origin_mm, theta_deg=rotation_deg) + return scipy_ndimage.affine_transform( + mask.astype(np.float32, copy=False), matrix, offset=offset, + output_shape=canvas_shape, order=0, mode="constant", cval=0.0) + + +def build_manual_alignment(sras: SrasFile, ref_angle_idx: int, + dc_threshold_mv: float, + per_angle_params: dict[int, ManualAngleParams] + ) -> AlignmentResult: + """Build a full, full-resolution AlignmentResult from user-supplied + per-angle rotation+shift — the Manual Alignment counterpart to + compute_angle_alignment, skipping its mask/phase-correlation search + entirely (every angle's transform here is exactly what the caller + supplied). Pure matrix/bbox math, no image data touched anywhere in this + function, so it is cheap enough to call synchronously on the GUI thread. + The reference angle's params are always forced to identity, regardless + of what per_angle_params holds for it — it defines the shared origin and + must never be transformed. + """ + n = sras.n_angles + dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx) + params = {a: per_angle_params.get(a, ManualAngleParams()) for a in range(n)} + params[ref_angle_idx] = ManualAngleParams() + + canvas_origin_mm, canvas_shape = union_canvas_mm( + sras, ref_angle_idx, dx_ref, dy_ref, params, margin_frac=0.0) + + per_angle: dict[int, AngleTransform] = {} + for a in range(n): + p = params[a] + matrix, offset = _build_affine_canvas_to_raw( + sras, a, ref_angle_idx, p.shift_mm, dx_ref, dy_ref, + canvas_origin_mm, theta_deg=p.rotation_deg) + per_angle[a] = AngleTransform(p.rotation_deg, p.shift_mm, matrix, offset) + + return AlignmentResult(ref_angle_idx, dc_threshold_mv, canvas_shape, + dx_ref, dy_ref, canvas_origin_mm, per_angle) + + +# ---- Sidecar persistence (.sras.align.json) ------------------------- +# +# Lives here, not sras_format.py: sras_format.py is scoped to the versioned +# binary .sras spec itself (see scan_format.md); a manual alignment is a +# viewer-computed *derived* artifact, analogous in kind to AlignmentResult — +# so it belongs with the alignment math it serialises, which already lives +# in this module. json + pathlib are both stdlib, so this doesn't add a new +# dependency to a module whose only load-bearing constraint is staying free +# of Qt/matplotlib for cheap multiprocessing-child imports. + +@dataclass +class ManualAlignmentSidecar: + ref_angle_idx: int + dc_threshold_mv: float + per_angle: dict[int, ManualAngleParams] + + +def sidecar_path(sras_path) -> Path: + """.sras.align.json next to the scan file. A thin, independently + testable function since save/load/delete and the GUI's status messages + all need the identical path.""" + p = Path(sras_path) + return p.with_name(p.name + ".align.json") + + +def save_manual_alignment(sras: SrasFile, ref_angle_idx: int, + dc_threshold_mv: float, + per_angle: dict[int, ManualAngleParams]) -> Path: + """Write the sidecar JSON for sras.path (overwriting any existing one) + and return the path written. + + Schema (schema_version 1): + { + "schema_version": 1, + "ref_angle_idx": , + "dc_threshold_mv": , + "per_angle": { + "": {"rotation_deg": , "shift_mm": [, ]}, + ... + } + } + Angle indices are JSON object keys, so they round-trip as strings — + load_manual_alignment converts them back to int. + """ + path = sidecar_path(sras.path) + payload = { + "schema_version": 1, + "ref_angle_idx": ref_angle_idx, + "dc_threshold_mv": dc_threshold_mv, + "per_angle": { + str(a): {"rotation_deg": p.rotation_deg, "shift_mm": list(p.shift_mm)} + for a, p in per_angle.items() + }, + } + path.write_text(json.dumps(payload, indent=2)) + return path + + +def load_manual_alignment(sras: SrasFile) -> ManualAlignmentSidecar | None: + """Read 's sidecar JSON if present, else None — never raises: + a missing, corrupt, foreign, or version-mismatched JSON file must not + block opening the .sras file itself (see SrasViewerWindow._on_load_done). + Per-angle entries for an angle index no longer present in *sras* (e.g. + re-scanned with fewer angles) are silently dropped. + """ + path = sidecar_path(sras.path) + if not path.exists(): + return None + try: + raw = json.loads(path.read_text()) + if raw.get("schema_version") != 1: + return None + per_angle = { + int(a): ManualAngleParams( + float(v["rotation_deg"]), + (float(v["shift_mm"][0]), float(v["shift_mm"][1]))) + for a, v in raw.get("per_angle", {}).items() + if int(a) < sras.n_angles + } + return ManualAlignmentSidecar( + ref_angle_idx=int(raw.get("ref_angle_idx", 0)), + dc_threshold_mv=float(raw.get("dc_threshold_mv", 0.0)), + per_angle=per_angle) + except (OSError, ValueError, KeyError, TypeError, IndexError, + json.JSONDecodeError): + return None + + +def delete_manual_alignment(sras: SrasFile) -> bool: + """Delete the sidecar if present. Returns whether a file actually existed + to delete, so Clear Alignment's status message can say so. Genuine I/O + errors (permission denied, read-only share) propagate — the caller + (ManualAlignmentDialog._on_clear) surfaces them rather than silently + pretending the destructive action succeeded.""" + path = sidecar_path(sras.path) + try: + path.unlink() + return True + except FileNotFoundError: + return False diff --git a/sras_viewer.py b/sras_viewer.py index c09d578..e6f2730 100644 --- a/sras_viewer.py +++ b/sras_viewer.py @@ -19,29 +19,34 @@ import faulthandler import sys from pathlib import Path +import matplotlib as mpl import numpy as np from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg, NavigationToolbar2QT from matplotlib.figure import Figure from matplotlib.patches import Polygon from matplotlib.path import Path as MplPath from PyQt6.QtCore import QObject, Qt, QThread, pyqtSignal -from PyQt6.QtGui import QAction +from PyQt6.QtGui import QAction, QKeyEvent from PyQt6.QtWidgets import ( QApplication, QButtonGroup, QCheckBox, QComboBox, QDialog, QDialogButtonBox, QDoubleSpinBox, QFileDialog, QFormLayout, QFrame, QGroupBox, QHBoxLayout, - QLabel, QMainWindow, QProgressDialog, QPushButton, QRadioButton, QScrollArea, - QSizePolicy, QSpinBox, QSplitter, QVBoxLayout, QWidget, + QLabel, QMainWindow, QMessageBox, QProgressDialog, QPushButton, QRadioButton, + QScrollArea, QSizePolicy, QSpinBox, QSplitter, QVBoxLayout, QWidget, ) import sras_compute as compute -from sras_compute import PYFFTW_AVAILABLE, apply_alignment +from sras_compute import ( + PYFFTW_AVAILABLE, ManualAngleParams, apply_alignment, build_manual_alignment, + delete_manual_alignment, load_manual_alignment, save_manual_alignment, + sidecar_path, +) from sras_format import ( CH1_IDX, CH3_IDX, CH4_IDX, CH_NAMES, SrasFile, adc_to_mv, mv_to_adc, _FALLBACK_YMULT_MV, _FALLBACK_YOFF_ADC, ) from sras_workers import ( - AngleAlignmentWorker, BatchCacheWorker, ComputeWorker, DcPrecomputeWorker, - LoadWorker, + AngleAlignmentWorker, BatchCacheWorker, Ch4MaskWorker, ComputeWorker, + DcPrecomputeWorker, LoadWorker, ) faulthandler.enable() # print a native stack trace on SIGSEGV/SIGABRT/etc. @@ -738,6 +743,542 @@ class FftOptionsDialog(QDialog): return max(1, self._spin_pad.value()) +# --------------------------------------------------------------------------- +# Manual alignment dialog (Fusion -> Manual Alignment...) +# --------------------------------------------------------------------------- + +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. + + 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 + full preview-canvas rebuild is needed. + + FigureCanvasQTAgg is a real QWidget, so keyPressEvent works like on any + other widget, but Qt only ever delivers key events to whichever widget + currently has focus — StrongFocus, plus grabbing focus on click and once + right after the dialog is shown, are both required or arrow keys + silently do nothing. + + Rotate keys are letters (Q/E), not punctuation (comma/period or + brackets): Shift+letter still reports the same Qt.Key on every platform, + whereas Shift+comma/bracket can report a different virtual key + (Key_Less / Key_BraceLeft) depending on platform and keyboard layout — + which would silently break the "Shift = coarse step" modifier for + rotation specifically. Arrow keys have no such hazard. + """ + nudge_translate = pyqtSignal(int, int, bool) # dir_x, dir_y in {-1,0,1}; coarse + nudge_rotate = pyqtSignal(int, bool) # dir in {-1,1} (CCW/CW); coarse + + _TRANSLATE_KEYS = { + Qt.Key.Key_Left: (-1, 0), + Qt.Key.Key_Right: (1, 0), + Qt.Key.Key_Up: (0, -1), + Qt.Key.Key_Down: (0, 1), + } + _ROTATE_KEYS = {Qt.Key.Key_Q: 1, Qt.Key.Key_E: -1} # CCW, CW + + def __init__(self, parent=None): + fig = Figure(figsize=(6, 6), tight_layout=True) + self.ax = fig.add_subplot(111) + super().__init__(fig) + self.setParent(parent) + self.setFocusPolicy(Qt.FocusPolicy.StrongFocus) + self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding) + self.mpl_connect("button_press_event", lambda _e: self.setFocus()) + + def show_overlay(self, rgba: np.ndarray, extent: list[float], title: str): + self.figure.clf() + self.ax = self.figure.add_subplot(111) + self.ax.imshow(rgba, extent=extent, origin="upper", aspect="auto") + self.ax.set_xlabel("X (mm)") + self.ax.set_ylabel("Y (mm)") + self.ax.set_title(title) + self.draw_idle() # coalesces rapid redraws — matters for key-repeat. + + def keyPressEvent(self, event: QKeyEvent): + key = event.key() + coarse = bool(event.modifiers() & Qt.KeyboardModifier.ShiftModifier) + if key in self._TRANSLATE_KEYS: + dx, dy = self._TRANSLATE_KEYS[key] + self.nudge_translate.emit(dx, dy, coarse) + event.accept() + elif key in self._ROTATE_KEYS: + self.nudge_rotate.emit(self._ROTATE_KEYS[key], coarse) + event.accept() + else: + super().keyPressEvent(event) + + +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 — the thing the automatic phase- + correlation step (compute_angle_alignment) sometimes gets wrong. The user + picks an "active" angle and nudges its rotation+translation with the + keyboard; Auto De-rotate sets every non-reference angle's rotation to the + known, analytic scan-angle delta without touching any translation the + user has already dialed in. 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 own (rare) background mask-fetch step, so the main window's + existing shutdown/lifecycle plumbing covers it 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, 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 + + 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_factor = 1 + 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_dx_mm = self._preview_dy_mm = 1.0 + self._masks_ready = 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 = ManualAlignOverlayCanvas() + 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) + + # ---- Active Angle ------------------------------------------------- + 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) + panel_l.addWidget(grp_angle) + + # ---- Manual Adjustment --------------------------------------------- + self.grp_manual_adjust, mform_box = _group("Manual Adjustment") + mform = _form() + self.spin_active_rotation_deg = QDoubleSpinBox() + self.spin_active_rotation_deg.setRange(-3600.0, 3600.0) + self.spin_active_rotation_deg.setDecimals(3) + self.spin_active_rotation_deg.setSuffix(" °") + self.spin_active_rotation_deg.setMinimumWidth(_SPIN_MIN_W) + mform.addRow("Rotation:", self.spin_active_rotation_deg) + + self.spin_active_shift_x_mm = QDoubleSpinBox() + self.spin_active_shift_x_mm.setRange(-1e5, 1e5) + self.spin_active_shift_x_mm.setDecimals(4) + self.spin_active_shift_x_mm.setSuffix(" mm") + self.spin_active_shift_x_mm.setMinimumWidth(_SPIN_MIN_W) + mform.addRow("Shift X:", self.spin_active_shift_x_mm) + + self.spin_active_shift_y_mm = QDoubleSpinBox() + self.spin_active_shift_y_mm.setRange(-1e5, 1e5) + self.spin_active_shift_y_mm.setDecimals(4) + self.spin_active_shift_y_mm.setSuffix(" mm") + self.spin_active_shift_y_mm.setMinimumWidth(_SPIN_MIN_W) + mform.addRow("Shift Y:", self.spin_active_shift_y_mm) + mform_box.addLayout(mform) + panel_l.addWidget(self.grp_manual_adjust) + + # ---- Nudge Step Sizes ------------------------------------------------ + self.grp_step_sizes, sl = _group("Nudge Step Sizes") + sform = _form() + self.spin_step_translate_mm = QDoubleSpinBox() + self.spin_step_translate_mm.setRange(0.0001, 1000.0) + self.spin_step_translate_mm.setDecimals(4) + self.spin_step_translate_mm.setSuffix(" mm") + self.spin_step_translate_mm.setValue(0.01) + self.spin_step_translate_mm.setMinimumWidth(_SPIN_MIN_W) + sform.addRow("Translate step:", self.spin_step_translate_mm) + + self.spin_step_rotate_deg = QDoubleSpinBox() + self.spin_step_rotate_deg.setRange(0.001, 90.0) + self.spin_step_rotate_deg.setDecimals(3) + self.spin_step_rotate_deg.setSuffix(" °") + self.spin_step_rotate_deg.setValue(0.1) + self.spin_step_rotate_deg.setMinimumWidth(_SPIN_MIN_W) + sform.addRow("Rotate step:", self.spin_step_rotate_deg) + + self.spin_step_multiplier = QDoubleSpinBox() + self.spin_step_multiplier.setRange(1.0, 1000.0) + self.spin_step_multiplier.setDecimals(1) + self.spin_step_multiplier.setValue(10.0) + self.spin_step_multiplier.setMinimumWidth(_SPIN_MIN_W) + 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)) + panel_l.addWidget(self.grp_step_sizes) + + # ---- Mask Threshold --------------------------------------------------- + self.grp_mask_threshold, tl = _group("Mask Threshold") + tform = _form() + self.spin_mask_threshold_mv = QDoubleSpinBox() + self.spin_mask_threshold_mv.setRange(-500.0, 500.0) + self.spin_mask_threshold_mv.setDecimals(3) + self.spin_mask_threshold_mv.setSuffix(" mV") + self.spin_mask_threshold_mv.setValue(dc_threshold_mv) + self.spin_mask_threshold_mv.setMinimumWidth(_SPIN_MIN_W) + tform.addRow("DC threshold:", self.spin_mask_threshold_mv) + tl.addLayout(tform) + panel_l.addWidget(self.grp_mask_threshold) + + # ---- Actions ------------------------------------------------------ + 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) + panel_l.addWidget(grp_actions) + + self.lbl_status = _wrap_label("", _CSS_MUTED) + panel_l.addWidget(self.lbl_status) + panel_l.addStretch() + + root.addWidget(_scroll_panel(panel, 300)) + + 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_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) + + self.combo_active_angle.blockSignals(True) + self.combo_active_angle.setCurrentIndex(self._active_angle) + self.combo_active_angle.blockSignals(False) + 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( + "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 + max_dim = max(max(img.shape) for img in self._dc4_mv.values()) + self._downsample_factor = max(1, int(np.ceil(max_dim / self._MAX_PREVIEW_DIM))) + 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.""" + threshold = self.spin_mask_threshold_mv.value() + factor = self._downsample_factor + self._masks_small = { + a: compute._block_mean_downsample( + (img >= threshold).astype(np.float32), factor) + 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) + factor = self._downsample_factor + dx_c, dy_c = dx_ref * factor, dy_ref * factor + origin, shape = compute.union_canvas_mm( + self._sras, self._ref_angle_idx, dx_c, dy_c, self._angle_params, + margin_frac=self._PREVIEW_MARGIN_FRAC) + self._preview_origin_mm, self._preview_shape = origin, shape + self._preview_dx_mm, self._preview_dy_mm = dx_c, dy_c + self._preview_layers = { + a: compute.reproject_mask( + self._sras, a, self._ref_angle_idx, self._masks_small[a], + self._angle_params[a].rotation_deg, self._angle_params[a].shift_mm, + dx_c, dy_c, origin, shape) + for a in range(self._sras.n_angles) + } + self._redraw_overlay() + + def _refresh_active_preview_layer(self): + """Cheap path for a translation-only nudge/edit of the active angle: + reproject just that one angle's downsampled mask onto the *existing* + preview canvas — every other angle's cached layer is untouched.""" + a = self._active_angle + self._preview_layers[a] = compute.reproject_mask( + self._sras, a, self._ref_angle_idx, self._masks_small[a], + self._angle_params[a].rotation_deg, self._angle_params[a].shift_mm, + self._preview_dx_mm, self._preview_dy_mm, + self._preview_origin_mm, self._preview_shape) + self._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_dx_mm, self._preview_dy_mm + x_axis = x0 + np.arange(n_cols) * dx + y_axis = y0 + np.arange(n_rows) * dy + extent = [x_axis[0] - dx / 2, x_axis[-1] + dx / 2, + y_axis[-1] + dy / 2, y_axis[0] - dy / 2] + 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])): + spin.blockSignals(True) + spin.setValue(val) + spin.blockSignals(False) + + 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): + n_changed = 0 + for a in range(self._sras.n_angles): + if a == self._ref_angle_idx: + continue + self._angle_params[a].rotation_deg = compute._theta_deg( + self._sras, a, self._ref_angle_idx) + n_changed += 1 + self._sync_active_spinboxes() + self._rebuild_preview_canvas() + self.lbl_status.setText( + f"Rotation set to the known scan angle for {n_changed} angle(s) " + "(translation left untouched).") + + 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._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.btn_auto_derotate.setEnabled(enabled) + self.btn_save.setEnabled(enabled) + self.btn_clear.setEnabled(enabled) + + # --------------------------------------------------------------------------- # Main window # --------------------------------------------------------------------------- @@ -788,6 +1329,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._build_ui() @@ -1137,6 +1679,16 @@ class SrasViewerWindow(QMainWindow): 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) + convert_menu = menubar.addMenu("&Convert") self._batch_dc_act = QAction("Batch Compute DC and &Store…", self) self._batch_dc_act.setStatusTip( @@ -1196,6 +1748,13 @@ class SrasViewerWindow(QMainWindow): self._sras = sras self._current_image = None + # 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 + # Caches (and any in-flight DC precompute) belong to the previous # file's geometry — discard and start fresh. Bumping the generation # counters makes any still-running worker's result get dropped when @@ -1213,6 +1772,26 @@ class SrasViewerWindow(QMainWindow): self.chk_aligned_view.setEnabled(False) self.chk_aligned_view.blockSignals(False) + # Silently restore a previously-saved manual alignment, if any, so + # the work survives closing and reopening the file. + sidecar = load_manual_alignment(sras) + if sidecar is not None: + try: + self._alignment_result = build_manual_alignment( + sras, sidecar.ref_angle_idx, sidecar.dc_threshold_mv, + sidecar.per_angle) + self.chk_aligned_view.blockSignals(True) + self.chk_aligned_view.setChecked(True) + self.chk_aligned_view.blockSignals(False) + self.statusBar().showMessage( + f"Restored saved manual alignment from " + f"{sidecar_path(sras.path).name}") + except Exception as exc: + # A corrupt/foreign sidecar or a rescan that shrank n_angles + # below ref_angle_idx must not block opening the .sras file. + self.statusBar().showMessage( + f"Could not restore saved alignment: {exc}") + # A ROI from the previous file no longer matches the new scan's # geometry, so discard it on every load. self.image_canvas.clear_roi() @@ -1318,6 +1897,8 @@ class SrasViewerWindow(QMainWindow): self._alignment_act.setEnabled( has_file and s.n_angles > 1 and not self._job_running("align")) + self._manual_align_act.setEnabled( + has_file and s.n_angles > 1 and not self._job_running("align")) self.chk_aligned_view.setEnabled(enabled and self._alignment_result is not None) self._update_roi_ui() @@ -1934,6 +2515,73 @@ class SrasViewerWindow(QMainWindow): 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() + return + + ref_idx = 0 + threshold_mv = self.spin_threshold_mv.value() + seed: dict[int, ManualAngleParams] = {} + if (self._alignment_result is not None + and self._alignment_result.ref_angle_idx == ref_idx): + seed = {a: ManualAngleParams(t.rotation_deg, t.shift_mm) + for a, t in self._alignment_result.per_angle.items()} + threshold_mv = self._alignment_result.dc_threshold_mv + else: + 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( + 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() + + def _on_manual_align_dialog_closed(self, _result_code: int): + self._manual_align_dialog = None + + def _on_manual_alignment_saved(self, result, sidecar_path_str: str): + self._alignment_result = result + self._aligned_cache = {} + self._alignment_generation += 1 + self.chk_aligned_view.setEnabled(True) + self.chk_aligned_view.blockSignals(True) + self.chk_aligned_view.setChecked(True) + self.chk_aligned_view.blockSignals(False) + self._update_controls_enabled(self._sras is not None) + self.statusBar().showMessage( + f"Manual alignment saved to {Path(sidecar_path_str).name}") + if self._current_image is not None: + self._refresh_display() + + def _on_manual_alignment_cleared(self): + self._alignment_result = None + self._aligned_cache = {} + self._alignment_generation += 1 + self.chk_aligned_view.blockSignals(True) + self.chk_aligned_view.setChecked(False) + self.chk_aligned_view.setEnabled(False) + self.chk_aligned_view.blockSignals(False) + self._update_controls_enabled(self._sras is not None) + self.statusBar().showMessage("Manual alignment cleared.") + if self._current_image is not None: + self._refresh_display() + # ------------------------------------------------------------------ # FFT Options # ------------------------------------------------------------------ @@ -1960,6 +2608,9 @@ class SrasViewerWindow(QMainWindow): # ------------------------------------------------------------------ def closeEvent(self, event): + if self._manual_align_dialog is not None: + self._manual_align_dialog.close() + # Signal every cancellable worker first, then wait. Waiting without # signalling means sitting out whatever is in flight — on a large # scan a single angle is ~40 s. diff --git a/sras_workers.py b/sras_workers.py index 1a21e50..0068136 100644 --- a/sras_workers.py +++ b/sras_workers.py @@ -304,3 +304,45 @@ class AngleAlignmentWorker(QObject): self.finished.emit(result, "") except Exception as exc: self.finished.emit(None, str(exc)) + + +class Ch4MaskWorker(QObject): + """Fetches each requested angle's CH4 (Bias B) DC image in mV, for + ManualAlignmentDialog's initial threshold-mask overlay. + + 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 + gone through the v7 "Convert" batch step and for angles the main + window's own DcPrecomputeWorker (which runs automatically right after + 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). + """ + angle_done = pyqtSignal(int, np.ndarray) # angle_idx, dc4_mv + finished = pyqtSignal() + error = pyqtSignal(str) + + def __init__(self, sras: SrasFile, angle_indices: list[int]): + super().__init__() + self._sras = sras + self._angles = angle_indices + + def run(self): + try: + n_workers, budget = compute.plan_angle_level(self._sras) + pool = ThreadPoolExecutor(max_workers=n_workers) + try: + futures = { + pool.submit(dc_image_mv, self._sras, a, CH4_IDX, + max_workers=1, budget=budget): a + for a in self._angles + } + for fut in as_completed(futures): + a = futures[fut] + self.angle_done.emit(a, fut.result()) + finally: + pool.shutdown(wait=True) + self.finished.emit() + except Exception as exc: + self.error.emit(str(exc)) diff --git a/tools/test_gui.py b/tools/test_gui.py index c44e6da..915dbd9 100644 --- a/tools/test_gui.py +++ b/tools/test_gui.py @@ -4,24 +4,29 @@ 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, aligned view, ROI draw/move, and CSV export. +changes, angle alignment, manual angle alignment, aligned view, ROI +draw/move, and CSV export. Usage: QT_QPA_PLATFORM=offscreen python tools/test_gui.py [file.sras] """ +import json import os import sys import tempfile from pathlib import Path +from unittest.mock import patch os.environ.setdefault("QT_QPA_PLATFORM", "offscreen") import numpy as np # noqa: E402 -from PyQt6.QtCore import QEventLoop, QTimer # noqa: E402 -from PyQt6.QtWidgets import QApplication # noqa: E402 +from PyQt6.QtCore import QEventLoop, Qt, QTimer # noqa: E402 +from PyQt6.QtTest import QTest # noqa: E402 +from PyQt6.QtWidgets import QApplication, QMessageBox # noqa: E402 sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +import sras_compute as compute # noqa: E402 from sras_format import CH1_IDX, CH3_IDX, CH4_IDX # noqa: E402 from sras_viewer import RoiQuad, SrasViewerWindow, VELOCITY_MODE_IDX # noqa: E402 import tools.make_test_sras as gen # noqa: E402 @@ -170,7 +175,6 @@ def main(): check("Export ROI enabled", win.btn_export_roi.isEnabled()) csv_path = tmpdir / "roi.csv" - from unittest.mock import patch with patch("sras_viewer.QFileDialog.getSaveFileName", return_value=(str(csv_path), "")): win._on_export_roi_csv() @@ -233,6 +237,136 @@ def main(): win.image_canvas._img_shape == s.image_shape(0), str(win.image_canvas._img_shape)) + print("\nmanual alignment (Fusion)") + check("manual alignment action enabled", win._manual_align_act.isEnabled()) + + # --- Open, seeding from the still-live automatic AlignmentResult -------- + win._on_manual_alignment() + check("dialog opened", win._manual_align_dialog is not None) + dlg = win._manual_align_dialog + check("mask prep needed no background worker (already DC-cached)", + not win._job_running("manual_align_masks")) + r = win._alignment_result # still the automatic result from the block above + if r is not None: + check("seeded rotation matches the automatic result", + all(dlg._angle_params[a].rotation_deg == r.per_angle[a].rotation_deg + for a in range(s.n_angles))) + check("seeded shift matches the automatic result", + all(dlg._angle_params[a].shift_mm == r.per_angle[a].shift_mm + for a in range(s.n_angles))) + + # --- Reference angle is locked ------------------------------------------- + dlg.combo_active_angle.setCurrentIndex(dlg._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) + check("reference angle group disabled", not dlg.grp_manual_adjust.isEnabled()) + check("reference angle untouched by nudge attempts", + dlg._angle_params[dlg._ref_angle_idx] == before_ref) + + # --- Nudging a real angle (fine + coarse, translate + rotate) ----------- + active = 1 if s.n_angles > 1 else 0 + dlg.combo_active_angle.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() + check("fine translate nudge moved shift_x by exactly one fine step", + abs(dlg._angle_params[active].shift_mm[0] - (before[0] + fine_step)) < 1e-9) + + 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() + check("coarse translate nudge uses the multiplier", + abs(dlg._angle_params[active].shift_mm[1] - (before[1] - coarse_step)) < 1e-9) + + before_rot = dlg._angle_params[active].rotation_deg + dlg._on_nudge_rotate(1, False) + check("rotate nudge changed rotation_deg", + dlg._angle_params[active].rotation_deg != before_rot) + check("preview canvas rebuilt for every angle after a rotation nudge", + len(dlg._preview_layers) == s.n_angles) + + # --- Real key-event wiring (proves keyPressEvent -> signal -> slot) ----- + before = dlg._angle_params[active].shift_mm + QTest.keyClick(dlg.canvas, Qt.Key.Key_Right) + check("a real Right-arrow key event nudged shift_x", + dlg._angle_params[active].shift_mm[0] > before[0]) + + # --- Auto De-rotate: rotation only, translation untouched --------------- + shift_before_derotate = dlg._angle_params[active].shift_mm + dlg._on_auto_derotate() + expected_theta = compute._theta_deg(s, active, dlg._ref_angle_idx) + check("auto de-rotate set the known analytic angle", + abs(dlg._angle_params[active].rotation_deg - expected_theta) < 1e-6) + check("auto de-rotate left translation untouched", + dlg._angle_params[active].shift_mm == shift_before_derotate) + check("reference angle stays identity after auto de-rotate", + dlg._angle_params[dlg._ref_angle_idx].rotation_deg == 0.0) + + # --- Save ----------------------------------------------------------------- + dlg._on_save() + sidecar = compute.sidecar_path(s.path) + check("sidecar file written", sidecar.exists()) + raw = json.loads(sidecar.read_text()) if sidecar.exists() else {} + check("sidecar schema_version is 1", raw.get("schema_version") == 1) + check("sidecar per_angle round-trips the dialog's resolved params", + all(raw.get("per_angle", {}).get(str(a), {}).get("rotation_deg") + == dlg._angle_params[a].rotation_deg for a in range(s.n_angles))) + check("main window's alignment_result replaced by the manual build", + win._alignment_result is not None + and win._alignment_result.per_angle[active].rotation_deg + == dlg._angle_params[active].rotation_deg) + check("Aligned View auto-enabled after Save", + win.chk_aligned_view.isEnabled() and win.chk_aligned_view.isChecked()) + + # --- Clear (with confirmation) -------------------------------------------- + with patch("sras_viewer.QMessageBox.question", + return_value=QMessageBox.StandardButton.Yes): + dlg._on_clear() + check("sidecar file deleted", not sidecar.exists()) + check("dialog params reset to identity", + all(dlg._angle_params[a] == compute.ManualAngleParams() + for a in range(s.n_angles))) + check("main window alignment_result cleared", win._alignment_result is None) + check("Aligned View disabled after Clear", + not win.chk_aligned_view.isEnabled() and not win.chk_aligned_view.isChecked()) + + dlg.close() + pump(150) + check("dialog reference released on close", win._manual_align_dialog is None) + + # --- Sidecar auto-restore on next load ------------------------------------ + 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) + + old_sras_id = id(win._sras) + win._load_file(str(path)) # reload the same file fresh + check("file reloaded", wait_until( + lambda: win._sras is not None and id(win._sras) != old_sras_id)) + s = win._sras + check("manual dialog force-closed by a reload", win._manual_align_dialog is None) + check("reload restores the saved manual alignment automatically", + win._alignment_result is not None) + if win._alignment_result is not None: + check("restored rotation matches what was saved", + abs(win._alignment_result.per_angle[active].rotation_deg + - saved_rotation) < 1e-9) + check("restored shift matches what was saved", + win._alignment_result.per_angle[active].shift_mm == saved_shift) + check("Aligned View auto-checked after restoring a saved alignment", + win.chk_aligned_view.isChecked()) + print("\npixel inspector") win.chk_aligned_view.setChecked(False) pump(100) From ca0c736c28013b249f537ccf705b8e3007fdf3b7 Mon Sep 17 00:00:00 2001 From: Thomas Ales Date: Fri, 31 Jul 2026 09:36:12 -0500 Subject: [PATCH 07/10] Fix alignment rotation pivot and sign convention Rotation was pivoting on each angle's raw scan-window bbox center, which ties alignment to wherever the window happened to sit in microscope/global XY space rather than to the sample itself, leaving a residual orbital drift between angles. Pivot on each angle's own CH4-threshold mask centroid instead (its own sample footprint), for both the automatic phase- correlation path and the manual dialog's Auto De-rotate/live preview. Also negate the rotation used everywhere (_theta_deg): the GR stage's reported angle increases in the opposite rotational sense from this module's CCW math convention, so de-rotating by the raw angles_deg delta was making misalignment worse rather than better. Co-Authored-By: Claude Sonnet 5 --- sras_compute.py | 158 +++++++++++++++++++++++++++++++++++----------- sras_viewer.py | 25 ++++++-- tools/test_gui.py | 29 +++++++++ 3 files changed, 172 insertions(+), 40 deletions(-) diff --git a/sras_compute.py b/sras_compute.py index 37f7b16..803d2c1 100644 --- a/sras_compute.py +++ b/sras_compute.py @@ -442,14 +442,75 @@ def _rotation_matrix(theta_deg: float) -> np.ndarray: def _theta_deg(sras: SrasFile, angle_idx: int, ref_idx: int) -> float: - return float(sras.angles_deg[angle_idx] - sras.angles_deg[ref_idx]) + """CCW rotation, in degrees and in _rotation_matrix's convention, that + maps angle_idx's own local mm frame onto ref_idx's. + + This is the *negative* of the raw angles_deg delta: the GR rotation + stage's reported angle increases in the opposite rotational sense from + this module's math-positive (CCW, x toward y) convention in scan mm + space. Rotating by +(angles_deg[a] - angles_deg[ref]) therefore turns + misalignment the wrong way — confirmed empirically (Auto De-rotate made + real scans worse, not better, before this negation). + """ + return -float(sras.angles_deg[angle_idx] - sras.angles_deg[ref_idx]) + + +def _mask_centroid_mm(sras: SrasFile, angle_idx: int, + mask: np.ndarray) -> tuple[float, float]: + """Centroid (mean x, mean y) of *mask*'s True pixels, in angle_idx's own + local mm frame — the alignment pivot used in place of the raw scan- + window bbox center (see compute_pivot_points_mm). Falls back to the bbox + center if the mask has no pixels above threshold, since a centroid of + nothing is undefined.""" + if not mask.any(): + return _bbox_center_mm(sras, angle_idx) + x = sras.x_axis_mm(angle_idx) + y = sras.y_positions_mm(angle_idx) + rows, cols = np.nonzero(mask) + return float(x[cols].mean()), float(y[rows].mean()) + + +def compute_pivot_points_mm(sras: SrasFile, dc_threshold_mv: float, + masks: dict[int, np.ndarray] | None = None + ) -> dict[int, tuple[float, float]]: + """Per-angle alignment pivot, in each angle's own local mm frame: the + centroid of its own CH4-threshold mask (its own sample footprint), + rather than the raw scan-window bbox center. + + Pivoting on each angle's own content — instead of on wherever its + scanned window happened to sit in microscope/global XY space — is what + makes alignment purely relative *between scans* rather than to a global + coordinate system: the rotation stage's true mechanical axis need not + coincide with the scan window's geometric center, and the sample need + not be perfectly centered on that axis either, so a bbox-center pivot + leaves a residual orbital motion between angles that a content-centroid + pivot does not. + + *masks* lets a caller that has already computed CH4 masks at this + threshold (compute_angle_alignment's Step 1) reuse them instead of + recomputing the DC image; angles missing from it are computed fresh via + dc_image_mv, which prefers a stored v5/v7 cache over recomputing from + raw waveforms. + """ + masks = masks or {} + pivots: dict[int, tuple[float, float]] = {} + for a in range(sras.n_angles): + mask = masks.get(a) + if mask is None: + mask = dc_image_mv(sras, a, CH4_IDX) >= dc_threshold_mv + else: + mask = mask > 0 # compute_angle_alignment's masks are float32 0.0/1.0 + pivots[a] = _mask_centroid_mm(sras, a, mask) + return pivots def _corners_in_ref_frame(sras: SrasFile, angle_idx: int, ref_idx: int, + pivot_mm: dict[int, tuple[float, float]], shift_mm=(0.0, 0.0), theta_deg: float | None = None) -> np.ndarray: - """Angle *angle_idx*'s bbox corners, rotated about its own centroid into - the reference frame and translated by *shift_mm*. Shape (4, 2). + """Angle *angle_idx*'s bbox corners, rotated about its own alignment + pivot (pivot_mm[angle_idx] — see compute_pivot_points_mm) into the + reference frame and translated by *shift_mm*. Shape (4, 2). theta_deg overrides the analytic angles_deg-derived rotation used by default — the manual-alignment path (union_canvas_mm) passes a @@ -458,8 +519,8 @@ def _corners_in_ref_frame(sras: SrasFile, angle_idx: int, ref_idx: int, """ theta = _theta_deg(sras, angle_idx, ref_idx) if theta_deg is None else theta_deg R = _rotation_matrix(theta) - c_a = np.array(_bbox_center_mm(sras, angle_idx)) - c_ref = np.array(_bbox_center_mm(sras, ref_idx)) + c_a = np.array(pivot_mm[angle_idx]) + c_ref = np.array(pivot_mm[ref_idx]) shift = np.asarray(shift_mm, dtype=np.float64) return np.array([R @ (corner - c_a) + c_ref + shift for corner in _bbox_corners_mm(sras, angle_idx)]) @@ -469,6 +530,7 @@ def _build_affine_canvas_to_raw(sras: SrasFile, angle_idx: int, ref_idx: int, shift_mm: tuple[float, float], canvas_dx: float, canvas_dy: float, canvas_origin_mm: tuple[float, float], + pivot_mm: dict[int, tuple[float, float]], theta_deg: float | None = None ) -> tuple[np.ndarray, np.ndarray]: """matrix, offset s.t. raw_index = matrix @ [row_out, col_out] + offset, @@ -479,18 +541,21 @@ def _build_affine_canvas_to_raw(sras: SrasFile, angle_idx: int, ref_idx: int, [lx;ly] = R(theta)^T @ ([X;Y]-c_ref-shift) + c_a # undo rotation+shift -> angle a's local mm [row;col] = D @ ([lx;ly] - [x_start_a; y0_a]) # local mm -> angle a's raw idx - where theta = angles_deg[angle_idx] - angles_deg[ref_idx] (unless - *theta_deg* overrides it — see _corners_in_ref_frame's note, used by the - manual-alignment path), c_ref/c_a are each angle's own raw-bbox mm - centroid (the rotation pivot — this keeps rotated content centered, - minimizing required canvas padding), and A_out/D are the index<->mm - scaling matrices for the canvas pitch and this angle's own native pitch - respectively. + where theta is _theta_deg(angle_idx, ref_idx) unless *theta_deg* + overrides it (see _corners_in_ref_frame's note, used by the manual- + alignment path), c_ref/c_a are each angle's own alignment pivot + (pivot_mm — see compute_pivot_points_mm; the centroid of its own + CH4-threshold mask, not the raw scan-window bbox center — this keeps + the sample itself centered post-rotation, minimizing required canvas + padding and, more importantly, keeping alignment relative to the sample + rather than to wherever the scan window sat in microscope/global XY + space), and A_out/D are the index<->mm scaling matrices for the canvas + pitch and this angle's own native pitch respectively. """ theta = _theta_deg(sras, angle_idx, ref_idx) if theta_deg is None else theta_deg Rinv = _rotation_matrix(theta).T - cx_a, cy_a = _bbox_center_mm(sras, angle_idx) - cx_ref, cy_ref = _bbox_center_mm(sras, ref_idx) + cx_a, cy_a = pivot_mm[angle_idx] + cx_ref, cy_ref = pivot_mm[ref_idx] dx_a, dy_a = _pixel_pitch_mm(sras, angle_idx) x0_a = float(sras.x_start_mm[angle_idx]) y0_a = float(sras.y_positions_mm(angle_idx)[0]) @@ -554,15 +619,17 @@ def _phase_correlate_shift(ref_img: np.ndarray, mov_img: np.ndarray) -> tuple[in def _working_canvas_for_pair(sras: SrasFile, ref_idx: int, a_idx: int, - dx: float, dy: float, margin_frac: float = 0.3 + dx: float, dy: float, + pivot_mm: dict[int, tuple[float, float]], + margin_frac: float = 0.3 ) -> tuple[tuple[float, float], tuple[int, int]]: """Union of the reference's own raw bbox and angle a's raw bbox rotated - (about its own center) into the ref frame with zero shift, padded by - margin_frac on each side — sized generously so the true phase- - correlation shift lands well inside the canvas (see + (about its own alignment pivot) into the ref frame with zero shift, + padded by margin_frac on each side — sized generously so the true + phase-correlation shift lands well inside the canvas (see _phase_correlate_shift's wraparound note).""" pts = np.vstack([_bbox_corners_mm(sras, ref_idx), - _corners_in_ref_frame(sras, a_idx, ref_idx)]) + _corners_in_ref_frame(sras, a_idx, ref_idx, pivot_mm)]) x_min, y_min = pts.min(axis=0) x_max, y_max = pts.max(axis=0) pad_x, pad_y = (x_max - x_min) * margin_frac, (y_max - y_min) * margin_frac @@ -615,6 +682,11 @@ def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int, masks = dict(enumerate(_parallel_map(mask_for, range(n), n_workers))) + # Each angle's own alignment pivot — the centroid of its own mask, not + # the raw scan-window bbox center (see compute_pivot_points_mm) — reuses + # the masks just computed above, so this is free (no extra DC compute). + pivot_mm = compute_pivot_points_mm(sras, dc_threshold_mv, masks=masks) + # ---- Step 2: coarse correlation stage (downsample first, then rotate) # Downsampling before affine_transform (not after) is what keeps this # tractable for a v6 scan with thousands of rows/frames per angle. @@ -630,15 +702,16 @@ def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int, tick(25, 50) return (0.0, 0.0) work_origin, work_shape = _working_canvas_for_pair( - sras, ref_angle_idx, a, dx_c, dy_c, margin_frac=0.3) + sras, ref_angle_idx, a, dx_c, dy_c, pivot_mm, margin_frac=0.3) # Coarse canvas->raw affine at native pitch, then rescale by /factor # so it maps coarse-canvas idx -> coarse (downsampled) raw idx — # exact for block-mean downsampling (up to the trimmed remainder). m_a, o_a = _build_affine_canvas_to_raw( - sras, a, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin) + sras, a, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin, pivot_mm) m_ref, o_ref = _build_affine_canvas_to_raw( - sras, ref_angle_idx, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin) + sras, ref_angle_idx, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin, + pivot_mm) rotated_a = scipy_ndimage.affine_transform( masks_small[a], m_a / factor, offset=o_a / factor, output_shape=work_shape, order=0, mode="constant", cval=0.0) @@ -653,8 +726,9 @@ def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int, shifts_mm = dict(enumerate(_parallel_map(shift_for, range(n), n_workers))) # ---- Step 3: union bounding box over all angles (rotation+shift applied) - corners = np.vstack([_corners_in_ref_frame(sras, a, ref_angle_idx, shifts_mm[a]) - for a in range(n)]) + corners = np.vstack([ + _corners_in_ref_frame(sras, a, ref_angle_idx, pivot_mm, shifts_mm[a]) + for a in range(n)]) x_min, y_min = corners.min(axis=0) x_max, y_max = corners.max(axis=0) n_cols = int(np.ceil((x_max - x_min) / dx_ref)) + 1 @@ -665,7 +739,8 @@ def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int, per_angle: dict[int, AngleTransform] = {} for a in range(n): matrix, offset = _build_affine_canvas_to_raw( - sras, a, ref_angle_idx, shifts_mm[a], dx_ref, dy_ref, canvas_origin_mm) + sras, a, ref_angle_idx, shifts_mm[a], dx_ref, dy_ref, canvas_origin_mm, + pivot_mm) per_angle[a] = AngleTransform( _theta_deg(sras, a, ref_angle_idx), shifts_mm[a], matrix, offset) if progress_cb: @@ -698,6 +773,7 @@ _compute_angle_alignment = compute_angle_alignment def union_canvas_mm(sras: SrasFile, ref_angle_idx: int, dx: float, dy: float, per_angle_params: dict[int, ManualAngleParams], + pivot_mm: dict[int, tuple[float, float]], margin_frac: float = 0.0 ) -> tuple[tuple[float, float], tuple[int, int]]: """Shared-canvas origin (mm) and (n_rows, n_cols) at pitch (dx, dy) that @@ -720,7 +796,7 @@ def union_canvas_mm(sras: SrasFile, ref_angle_idx: int, dx: float, dy: float, n = sras.n_angles corners = np.vstack([ _corners_in_ref_frame( - sras, a, ref_angle_idx, + sras, a, ref_angle_idx, pivot_mm, per_angle_params.get(a, ManualAngleParams()).shift_mm, theta_deg=per_angle_params.get(a, ManualAngleParams()).rotation_deg) for a in range(n)]) @@ -741,7 +817,8 @@ def reproject_mask(sras: SrasFile, angle_idx: int, ref_angle_idx: int, shift_mm: tuple[float, float], canvas_dx: float, canvas_dy: float, canvas_origin_mm: tuple[float, float], - canvas_shape: tuple[int, int]) -> np.ndarray: + canvas_shape: tuple[int, int], + pivot_mm: dict[int, tuple[float, float]]) -> 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 @@ -753,7 +830,7 @@ def reproject_mask(sras: SrasFile, angle_idx: int, ref_angle_idx: int, """ matrix, offset = _build_affine_canvas_to_raw( sras, angle_idx, ref_angle_idx, shift_mm, canvas_dx, canvas_dy, - canvas_origin_mm, theta_deg=rotation_deg) + canvas_origin_mm, pivot_mm, theta_deg=rotation_deg) return scipy_ndimage.affine_transform( mask.astype(np.float32, copy=False), matrix, offset=offset, output_shape=canvas_shape, order=0, mode="constant", cval=0.0) @@ -761,32 +838,41 @@ def reproject_mask(sras: SrasFile, angle_idx: int, ref_angle_idx: int, def build_manual_alignment(sras: SrasFile, ref_angle_idx: int, dc_threshold_mv: float, - per_angle_params: dict[int, ManualAngleParams] + per_angle_params: dict[int, ManualAngleParams], + pivot_mm: dict[int, tuple[float, float]] | None = None ) -> AlignmentResult: """Build a full, full-resolution AlignmentResult from user-supplied per-angle rotation+shift — the Manual Alignment counterpart to compute_angle_alignment, skipping its mask/phase-correlation search entirely (every angle's transform here is exactly what the caller - supplied). Pure matrix/bbox math, no image data touched anywhere in this - function, so it is cheap enough to call synchronously on the GUI thread. - The reference angle's params are always forced to identity, regardless - of what per_angle_params holds for it — it defines the shared origin and - must never be transformed. + supplied). Pure matrix/bbox math once pivot_mm is known, so it is cheap + enough to call synchronously on the GUI thread. The reference angle's + params are always forced to identity, regardless of what + per_angle_params holds for it — it defines the shared origin and must + never be transformed. + + Pass an already-computed *pivot_mm* (e.g. ManualAlignmentDialog's own + cache, built once from its live CH4 masks) to skip recomputing every + angle's DC image here; otherwise it's computed fresh via + compute_pivot_points_mm, which is still fine for a one-off Save or + sidecar restore (just not free on a very large, not-yet-cached scan). """ n = sras.n_angles dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx) params = {a: per_angle_params.get(a, ManualAngleParams()) for a in range(n)} params[ref_angle_idx] = ManualAngleParams() + if pivot_mm is None: + pivot_mm = compute_pivot_points_mm(sras, dc_threshold_mv) canvas_origin_mm, canvas_shape = union_canvas_mm( - sras, ref_angle_idx, dx_ref, dy_ref, params, margin_frac=0.0) + sras, ref_angle_idx, dx_ref, dy_ref, params, pivot_mm, margin_frac=0.0) per_angle: dict[int, AngleTransform] = {} for a in range(n): p = params[a] matrix, offset = _build_affine_canvas_to_raw( sras, a, ref_angle_idx, p.shift_mm, dx_ref, dy_ref, - canvas_origin_mm, theta_deg=p.rotation_deg) + canvas_origin_mm, pivot_mm, theta_deg=p.rotation_deg) per_angle[a] = AngleTransform(p.rotation_deg, p.shift_mm, matrix, offset) return AlignmentResult(ref_angle_idx, dc_threshold_mv, canvas_shape, diff --git a/sras_viewer.py b/sras_viewer.py index e6f2730..2c26b71 100644 --- a/sras_viewer.py +++ b/sras_viewer.py @@ -860,6 +860,7 @@ class ManualAlignmentDialog(QDialog): self._downsample_factor = 1 self._dc4_mv: dict[int, np.ndarray] = {} self._masks_small: dict[int, np.ndarray] = {} + self._pivot_mm: dict[int, tuple[float, float]] = {} self._preview_layers: dict[int, np.ndarray] = {} self._preview_origin_mm = (0.0, 0.0) self._preview_shape = (1, 1) @@ -1063,6 +1064,7 @@ class ManualAlignmentDialog(QDialog): max_dim = max(max(img.shape) for img in self._dc4_mv.values()) self._downsample_factor = max(1, int(np.ceil(max_dim / self._MAX_PREVIEW_DIM))) self._recompute_masks_small() + self._recompute_pivot_mm() self._rebuild_preview_canvas() self._set_controls_enabled(True) self.lbl_status.setText("Ready.") @@ -1079,6 +1081,19 @@ class ManualAlignmentDialog(QDialog): for a, img in self._dc4_mv.items() } + def _recompute_pivot_mm(self): + """Each angle's alignment pivot: the centroid of its own full- + resolution CH4-threshold mask (its own sample footprint), not the + raw scan-window bbox center — see compute.compute_pivot_points_mm. + Recomputed alongside _recompute_masks_small whenever the mask + threshold changes, from the full-res masks (not the downsampled + preview ones) since this feeds the canonical geometry, not just the + preview.""" + threshold = self.spin_mask_threshold_mv.value() + masks = {a: img >= threshold for a, img in self._dc4_mv.items()} + self._pivot_mm = compute.compute_pivot_points_mm( + self._sras, threshold, masks=masks) + # ------------------------------------------------------------------ # Preview canvas: full rebuild vs. incremental single-layer refresh # ------------------------------------------------------------------ @@ -1096,14 +1111,14 @@ class ManualAlignmentDialog(QDialog): dx_c, dy_c = dx_ref * factor, dy_ref * factor origin, shape = compute.union_canvas_mm( self._sras, self._ref_angle_idx, dx_c, dy_c, self._angle_params, - margin_frac=self._PREVIEW_MARGIN_FRAC) + self._pivot_mm, margin_frac=self._PREVIEW_MARGIN_FRAC) self._preview_origin_mm, self._preview_shape = origin, shape self._preview_dx_mm, self._preview_dy_mm = dx_c, dy_c self._preview_layers = { a: compute.reproject_mask( self._sras, a, self._ref_angle_idx, self._masks_small[a], self._angle_params[a].rotation_deg, self._angle_params[a].shift_mm, - dx_c, dy_c, origin, shape) + dx_c, dy_c, origin, shape, self._pivot_mm) for a in range(self._sras.n_angles) } self._redraw_overlay() @@ -1117,7 +1132,7 @@ class ManualAlignmentDialog(QDialog): self._sras, a, self._ref_angle_idx, self._masks_small[a], self._angle_params[a].rotation_deg, self._angle_params[a].shift_mm, self._preview_dx_mm, self._preview_dy_mm, - self._preview_origin_mm, self._preview_shape) + self._preview_origin_mm, self._preview_shape, self._pivot_mm) self._redraw_overlay() def _redraw_overlay(self): @@ -1212,6 +1227,7 @@ class ManualAlignmentDialog(QDialog): if not self._masks_ready: return self._recompute_masks_small() + self._recompute_pivot_mm() self._rebuild_preview_canvas() # ------------------------------------------------------------------ @@ -1237,7 +1253,8 @@ class ManualAlignmentDialog(QDialog): 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) + result = build_manual_alignment(self._sras, self._ref_angle_idx, threshold, + resolved, self._pivot_mm) except OSError as exc: QMessageBox.warning(self, "Save Alignment Failed", str(exc)) return diff --git a/tools/test_gui.py b/tools/test_gui.py index 915dbd9..8db05f6 100644 --- a/tools/test_gui.py +++ b/tools/test_gui.py @@ -240,6 +240,35 @@ def main(): print("\nmanual alignment (Fusion)") check("manual alignment action enabled", win._manual_align_act.isEnabled()) + # --- Alignment pivot is the mask centroid, not the raw bbox center ----- + corner_mask = np.zeros(s.image_shape(0), dtype=bool) + corner_mask[0, 0] = True + expected_corner = (float(s.x_axis_mm(0)[0]), float(s.y_positions_mm(0)[0])) + centroid = compute._mask_centroid_mm(s, 0, corner_mask) + check("mask centroid of a single corner pixel is that corner exactly", + np.allclose(centroid, expected_corner), f"{centroid} vs {expected_corner}") + bbox_center = compute._bbox_center_mm(s, 0) + check("mask centroid differs from the raw scan-window bbox center", + not np.allclose(centroid, bbox_center), + f"centroid {centroid} vs bbox center {bbox_center}") + + # A threshold that would yield an *empty* mask if recomputed from scratch + # (real DC values never reach 999 mV) — so this only matches expected_corner + # if compute_pivot_points_mm actually reused the pre-computed masks dict + # instead of silently recomputing (and falling back to the bbox center). + reused_pivot = compute.compute_pivot_points_mm( + s, 999.0, masks={0: corner_mask.astype(np.float32)})[0] + check("compute_pivot_points_mm reuses a pre-computed masks dict", + np.allclose(reused_pivot, expected_corner)) + + # --- Rotation sign convention: negative of the raw angles_deg delta ---- + check("_theta_deg negates the raw angles_deg delta (GR stage's positive " + "angle is the opposite rotational sense from this module's CCW " + "math convention)", + all(np.isclose(compute._theta_deg(s, a, 0), + -(float(s.angles_deg[a]) - float(s.angles_deg[0]))) + for a in range(s.n_angles))) + # --- Open, seeding from the still-live automatic AlignmentResult -------- win._on_manual_alignment() check("dialog opened", win._manual_align_dialog is not None) From 8312668c027ed358920ac0c94ba94562f26b38f0 Mon Sep 17 00:00:00 2001 From: Thomas Ales Date: Fri, 31 Jul 2026 09:58:15 -0500 Subject: [PATCH 08/10] Make alignment pivot independent of the DC mask threshold MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The pivot was still derived from the binary CH4 >= dc_threshold_mv mask, so a threshold that happens to leave a real angle's mask empty (signal levels vary scan to scan across a many-angle acquisition) silently fell back to the raw scan-window bbox center — reproducing the exact "aligned to the scan window, not the sample" scatter the centroid pivot exists to fix, without any visible error. Replace the binary-mask centroid with an intensity-weighted centroid of the continuous CH4 signal, which is always well-defined regardless of the RF-mask threshold in effect. The threshold still controls what the manual dialog's overlay and CH1 masking show; it no longer has any bearing on where alignment pivots. Also bump the sidecar schema version: a file saved before today's pivot and rotation-sign fixes stores numbers for a since-corrected transform, so loading it unchanged would reintroduce the same scatter. Old sidecars are now treated as absent rather than silently reused. Co-Authored-By: Claude Sonnet 5 --- sras_compute.py | 131 +++++++++++++++++++++++++++++----------------- sras_viewer.py | 25 ++++----- tools/test_gui.py | 46 +++++++++++----- 3 files changed, 124 insertions(+), 78 deletions(-) diff --git a/sras_compute.py b/sras_compute.py index 803d2c1..6b49a6f 100644 --- a/sras_compute.py +++ b/sras_compute.py @@ -455,27 +455,41 @@ def _theta_deg(sras: SrasFile, angle_idx: int, ref_idx: int) -> float: return -float(sras.angles_deg[angle_idx] - sras.angles_deg[ref_idx]) -def _mask_centroid_mm(sras: SrasFile, angle_idx: int, - mask: np.ndarray) -> tuple[float, float]: - """Centroid (mean x, mean y) of *mask*'s True pixels, in angle_idx's own - local mm frame — the alignment pivot used in place of the raw scan- - window bbox center (see compute_pivot_points_mm). Falls back to the bbox - center if the mask has no pixels above threshold, since a centroid of - nothing is undefined.""" - if not mask.any(): +def _signal_centroid_mm(sras: SrasFile, angle_idx: int, + dc4_mv: np.ndarray) -> tuple[float, float]: + """Intensity-weighted centroid (mean x, mean y, weighted by CH4 signal + after subtracting this angle's own minimum) in angle_idx's own local mm + frame — the alignment pivot used in place of the raw scan-window bbox + center (see compute_pivot_points_mm). + + Weighting by the continuous DC signal, rather than a binary >= + dc_threshold_mv mask, means the pivot never depends on how well one + shared threshold happens to suit this particular angle: real signal + levels vary scan to scan, so a threshold tuned for one angle can leave + another angle's binary mask empty — and a centroid of an empty mask has + nothing to fall back to *except* the raw bbox center, silently + reproducing the exact "aligned to the scan window, not the sample" + problem this pivot exists to avoid. Falls back to the bbox center only + in the fully-degenerate case of a perfectly flat signal (nothing to + weight by at all). + """ + weights = dc4_mv - dc4_mv.min() + total = float(weights.sum()) + if total <= 0.0: return _bbox_center_mm(sras, angle_idx) x = sras.x_axis_mm(angle_idx) y = sras.y_positions_mm(angle_idx) - rows, cols = np.nonzero(mask) - return float(x[cols].mean()), float(y[rows].mean()) + cx = float((x * weights.sum(axis=0)).sum() / total) + cy = float((y * weights.sum(axis=1)).sum() / total) + return cx, cy -def compute_pivot_points_mm(sras: SrasFile, dc_threshold_mv: float, - masks: dict[int, np.ndarray] | None = None +def compute_pivot_points_mm(sras: SrasFile, + dc4_mv: dict[int, np.ndarray] | None = None ) -> dict[int, tuple[float, float]]: """Per-angle alignment pivot, in each angle's own local mm frame: the - centroid of its own CH4-threshold mask (its own sample footprint), - rather than the raw scan-window bbox center. + CH4-signal-weighted centroid of its own footprint (see + _signal_centroid_mm), rather than the raw scan-window bbox center. Pivoting on each angle's own content — instead of on wherever its scanned window happened to sit in microscope/global XY space — is what @@ -486,21 +500,23 @@ def compute_pivot_points_mm(sras: SrasFile, dc_threshold_mv: float, leaves a residual orbital motion between angles that a content-centroid pivot does not. - *masks* lets a caller that has already computed CH4 masks at this - threshold (compute_angle_alignment's Step 1) reuse them instead of - recomputing the DC image; angles missing from it are computed fresh via - dc_image_mv, which prefers a stored v5/v7 cache over recomputing from - raw waveforms. + Deliberately independent of dc_threshold_mv (the RF-mask / overlay + threshold): that value is a display/masking choice and must never + silently change where alignment pivots. + + *dc4_mv* lets a caller that has already computed each angle's CH4 mV + image (compute_angle_alignment's Step 1, or ManualAlignmentDialog's own + cache) reuse it instead of recomputing; angles missing from it are + computed fresh via dc_image_mv, which prefers a stored v5/v7 cache over + recomputing from raw waveforms. """ - masks = masks or {} + dc4_mv = dc4_mv or {} pivots: dict[int, tuple[float, float]] = {} for a in range(sras.n_angles): - mask = masks.get(a) - if mask is None: - mask = dc_image_mv(sras, a, CH4_IDX) >= dc_threshold_mv - else: - mask = mask > 0 # compute_angle_alignment's masks are float32 0.0/1.0 - pivots[a] = _mask_centroid_mm(sras, a, mask) + img = dc4_mv.get(a) + if img is None: + img = dc_image_mv(sras, a, CH4_IDX) + pivots[a] = _signal_centroid_mm(sras, a, img) return pivots @@ -544,13 +560,14 @@ def _build_affine_canvas_to_raw(sras: SrasFile, angle_idx: int, ref_idx: int, where theta is _theta_deg(angle_idx, ref_idx) unless *theta_deg* overrides it (see _corners_in_ref_frame's note, used by the manual- alignment path), c_ref/c_a are each angle's own alignment pivot - (pivot_mm — see compute_pivot_points_mm; the centroid of its own - CH4-threshold mask, not the raw scan-window bbox center — this keeps - the sample itself centered post-rotation, minimizing required canvas - padding and, more importantly, keeping alignment relative to the sample - rather than to wherever the scan window sat in microscope/global XY - space), and A_out/D are the index<->mm scaling matrices for the canvas - pitch and this angle's own native pitch respectively. + (pivot_mm — see compute_pivot_points_mm; the CH4-signal-weighted + centroid of its own footprint, not the raw scan-window bbox center — + this keeps the sample itself centered post-rotation, minimizing + required canvas padding and, more importantly, keeping alignment + relative to the sample rather than to wherever the scan window sat in + microscope/global XY space), and A_out/D are the index<->mm scaling + matrices for the canvas pitch and this angle's own native pitch + respectively. """ theta = _theta_deg(sras, angle_idx, ref_idx) if theta_deg is None else theta_deg Rinv = _rotation_matrix(theta).T @@ -671,21 +688,25 @@ def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int, _ticks.append(1) progress_cb(base + int(min(len(_ticks), n) / n * span)) - # ---- Step 1: binarized CH4 mask per angle, native per-angle grid ----- - def mask_for(a: int) -> np.ndarray: + # ---- Step 1: CH4 DC image + binarized mask per angle, native grid ----- + def dc4_for(a: int) -> np.ndarray: dc4 = adc_to_mv( compute_dc_image(sras, a, CH4_IDX, max_workers=1, budget=angle_budget), *sras.cal(CH4_IDX)) - m = (dc4 >= dc_threshold_mv).astype(np.float32) tick(0, 25) - return m + return dc4 - masks = dict(enumerate(_parallel_map(mask_for, range(n), n_workers))) + dc4_mv = dict(enumerate(_parallel_map(dc4_for, range(n), n_workers))) + masks = {a: (img >= dc_threshold_mv).astype(np.float32) + for a, img in dc4_mv.items()} - # Each angle's own alignment pivot — the centroid of its own mask, not - # the raw scan-window bbox center (see compute_pivot_points_mm) — reuses - # the masks just computed above, so this is free (no extra DC compute). - pivot_mm = compute_pivot_points_mm(sras, dc_threshold_mv, masks=masks) + # Each angle's own alignment pivot — the CH4-signal-weighted centroid, + # not the raw scan-window bbox center (see compute_pivot_points_mm). + # Reuses the dc4_mv images just computed above, so this is free, and is + # deliberately independent of dc_threshold_mv (see that function's + # docstring) so a threshold that happens to leave some angle's binary + # mask empty can't silently degrade the pivot back to the bbox center. + pivot_mm = compute_pivot_points_mm(sras, dc4_mv=dc4_mv) # ---- Step 2: coarse correlation stage (downsample first, then rotate) # Downsampling before affine_transform (not after) is what keeps this @@ -852,17 +873,20 @@ def build_manual_alignment(sras: SrasFile, ref_angle_idx: int, never be transformed. Pass an already-computed *pivot_mm* (e.g. ManualAlignmentDialog's own - cache, built once from its live CH4 masks) to skip recomputing every + cache, built once from its live CH4 images) to skip recomputing every angle's DC image here; otherwise it's computed fresh via compute_pivot_points_mm, which is still fine for a one-off Save or sidecar restore (just not free on a very large, not-yet-cached scan). + dc_threshold_mv itself plays no part in the pivot (see that function's + docstring) — it's stored on the returned AlignmentResult purely as a + record of the RF-mask threshold in effect at the time. """ n = sras.n_angles dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx) params = {a: per_angle_params.get(a, ManualAngleParams()) for a in range(n)} params[ref_angle_idx] = ManualAngleParams() if pivot_mm is None: - pivot_mm = compute_pivot_points_mm(sras, dc_threshold_mv) + pivot_mm = compute_pivot_points_mm(sras) canvas_origin_mm, canvas_shape = union_canvas_mm( sras, ref_angle_idx, dx_ref, dy_ref, params, pivot_mm, margin_frac=0.0) @@ -904,15 +928,26 @@ def sidecar_path(sras_path) -> Path: return p.with_name(p.name + ".align.json") +# Bumped from 1 -> 2 when the rotation pivot changed from the raw scan- +# window bbox center to a content-derived centroid, *and* _theta_deg's sign +# convention was corrected — either change alone makes a version-1 file's +# stored rotation_deg/shift_mm numbers describe a different (and, for the +# pivot bug, actively wrong) transform than they would today. Loading one +# unchanged would silently reproduce exactly the "scans show up everywhere" +# symptom these fixes address, so version-1 sidecars are treated as absent +# rather than migrated. +_SIDECAR_SCHEMA_VERSION = 2 + + def save_manual_alignment(sras: SrasFile, ref_angle_idx: int, dc_threshold_mv: float, per_angle: dict[int, ManualAngleParams]) -> Path: """Write the sidecar JSON for sras.path (overwriting any existing one) and return the path written. - Schema (schema_version 1): + Schema (schema_version 2): { - "schema_version": 1, + "schema_version": 2, "ref_angle_idx": , "dc_threshold_mv": , "per_angle": { @@ -925,7 +960,7 @@ def save_manual_alignment(sras: SrasFile, ref_angle_idx: int, """ path = sidecar_path(sras.path) payload = { - "schema_version": 1, + "schema_version": _SIDECAR_SCHEMA_VERSION, "ref_angle_idx": ref_angle_idx, "dc_threshold_mv": dc_threshold_mv, "per_angle": { @@ -949,7 +984,7 @@ def load_manual_alignment(sras: SrasFile) -> ManualAlignmentSidecar | None: return None try: raw = json.loads(path.read_text()) - if raw.get("schema_version") != 1: + if raw.get("schema_version") != _SIDECAR_SCHEMA_VERSION: return None per_angle = { int(a): ManualAngleParams( diff --git a/sras_viewer.py b/sras_viewer.py index 2c26b71..ea9a663 100644 --- a/sras_viewer.py +++ b/sras_viewer.py @@ -1064,7 +1064,12 @@ class ManualAlignmentDialog(QDialog): max_dim = max(max(img.shape) for img in self._dc4_mv.values()) self._downsample_factor = max(1, int(np.ceil(max_dim / self._MAX_PREVIEW_DIM))) self._recompute_masks_small() - self._recompute_pivot_mm() + # Alignment pivot: the CH4-signal-weighted centroid of each angle's + # own footprint (see compute.compute_pivot_points_mm) — computed once + # from the full-res CH4 images and deliberately independent of the + # mask threshold, so it never needs recomputing when that changes + # (unlike _masks_small, which is purely for the overlay's visuals). + self._pivot_mm = compute.compute_pivot_points_mm(self._sras, self._dc4_mv) self._rebuild_preview_canvas() self._set_controls_enabled(True) self.lbl_status.setText("Ready.") @@ -1072,7 +1077,9 @@ class ManualAlignmentDialog(QDialog): 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.""" + full whenever the mask-threshold spin box changes — no re-fetch. + Purely for the overlay's visuals — the alignment pivot does not + depend on this threshold (see _pivot_mm / compute_pivot_points_mm).""" threshold = self.spin_mask_threshold_mv.value() factor = self._downsample_factor self._masks_small = { @@ -1081,19 +1088,6 @@ class ManualAlignmentDialog(QDialog): for a, img in self._dc4_mv.items() } - def _recompute_pivot_mm(self): - """Each angle's alignment pivot: the centroid of its own full- - resolution CH4-threshold mask (its own sample footprint), not the - raw scan-window bbox center — see compute.compute_pivot_points_mm. - Recomputed alongside _recompute_masks_small whenever the mask - threshold changes, from the full-res masks (not the downsampled - preview ones) since this feeds the canonical geometry, not just the - preview.""" - threshold = self.spin_mask_threshold_mv.value() - masks = {a: img >= threshold for a, img in self._dc4_mv.items()} - self._pivot_mm = compute.compute_pivot_points_mm( - self._sras, threshold, masks=masks) - # ------------------------------------------------------------------ # Preview canvas: full rebuild vs. incremental single-layer refresh # ------------------------------------------------------------------ @@ -1227,7 +1221,6 @@ class ManualAlignmentDialog(QDialog): if not self._masks_ready: return self._recompute_masks_small() - self._recompute_pivot_mm() self._rebuild_preview_canvas() # ------------------------------------------------------------------ diff --git a/tools/test_gui.py b/tools/test_gui.py index 8db05f6..83a0366 100644 --- a/tools/test_gui.py +++ b/tools/test_gui.py @@ -240,27 +240,36 @@ def main(): print("\nmanual alignment (Fusion)") check("manual alignment action enabled", win._manual_align_act.isEnabled()) - # --- Alignment pivot is the mask centroid, not the raw bbox center ----- - corner_mask = np.zeros(s.image_shape(0), dtype=bool) - corner_mask[0, 0] = True + # --- Alignment pivot is a signal-weighted centroid, not the raw bbox -- + # center, and is independent of any DC threshold (so a threshold that + # happens to leave a real angle's binary mask empty can't silently + # degrade the pivot back to the bbox center). + corner_signal = np.zeros(s.image_shape(0), dtype=np.float32) + corner_signal[0, 0] = 1.0 # single spike -> weighted centroid is exact expected_corner = (float(s.x_axis_mm(0)[0]), float(s.y_positions_mm(0)[0])) - centroid = compute._mask_centroid_mm(s, 0, corner_mask) - check("mask centroid of a single corner pixel is that corner exactly", + centroid = compute._signal_centroid_mm(s, 0, corner_signal) + check("signal-weighted centroid of a single spike pixel is that pixel exactly", np.allclose(centroid, expected_corner), f"{centroid} vs {expected_corner}") bbox_center = compute._bbox_center_mm(s, 0) - check("mask centroid differs from the raw scan-window bbox center", + check("signal centroid differs from the raw scan-window bbox center", not np.allclose(centroid, bbox_center), f"centroid {centroid} vs bbox center {bbox_center}") - # A threshold that would yield an *empty* mask if recomputed from scratch - # (real DC values never reach 999 mV) — so this only matches expected_corner - # if compute_pivot_points_mm actually reused the pre-computed masks dict - # instead of silently recomputing (and falling back to the bbox center). - reused_pivot = compute.compute_pivot_points_mm( - s, 999.0, masks={0: corner_mask.astype(np.float32)})[0] - check("compute_pivot_points_mm reuses a pre-computed masks dict", + # compute_pivot_points_mm should reuse a pre-computed dc4_mv dict rather + # than recomputing from the real DC4 image (which has no such spike and + # would give a different answer if silently recomputed). + reused_pivot = compute.compute_pivot_points_mm(s, dc4_mv={0: corner_signal})[0] + check("compute_pivot_points_mm reuses a pre-computed dc4_mv dict", np.allclose(reused_pivot, expected_corner)) + # A perfectly flat signal carries no information to weight by, so it + # falls back to the bbox center rather than producing a NaN/degenerate + # centroid. + flat_signal = np.full(s.image_shape(0), 5.0, dtype=np.float32) + flat_centroid = compute._signal_centroid_mm(s, 0, flat_signal) + check("a perfectly flat signal falls back to the bbox center", + np.allclose(flat_centroid, bbox_center)) + # --- Rotation sign convention: negative of the raw angles_deg delta ---- check("_theta_deg negates the raw angles_deg delta (GR stage's positive " "angle is the opposite rotational sense from this module's CCW " @@ -339,7 +348,8 @@ def main(): sidecar = compute.sidecar_path(s.path) check("sidecar file written", sidecar.exists()) raw = json.loads(sidecar.read_text()) if sidecar.exists() else {} - check("sidecar schema_version is 1", raw.get("schema_version") == 1) + check("sidecar schema_version is current", + raw.get("schema_version") == compute._SIDECAR_SCHEMA_VERSION) check("sidecar per_angle round-trips the dialog's resolved params", all(raw.get("per_angle", {}).get(str(a), {}).get("rotation_deg") == dlg._angle_params[a].rotation_deg for a in range(s.n_angles))) @@ -350,6 +360,14 @@ def main(): check("Aligned View auto-enabled after Save", win.chk_aligned_view.isEnabled() and win.chk_aligned_view.isChecked()) + # --- An old-schema sidecar (pre-pivot/sign fix) is treated as absent ------ + stale = dict(raw) + stale["schema_version"] = compute._SIDECAR_SCHEMA_VERSION - 1 + sidecar.write_text(json.dumps(stale)) + check("a sidecar with an old schema_version is not loaded", + compute.load_manual_alignment(s) is None) + sidecar.write_text(json.dumps(raw)) # restore for the rest of this section + # --- Clear (with confirmation) -------------------------------------------- with patch("sras_viewer.QMessageBox.question", return_value=QMessageBox.StandardButton.Yes): From 52e91d46fdfb4b4a0b0e3946012750c6695dae77 Mon Sep 17 00:00:00 2001 From: Thomas Ales Date: Fri, 31 Jul 2026 11:07:55 -0500 Subject: [PATCH 09/10] Stop seeding Manual Alignment translation from the automatic result The dialog was seeding rotation_deg/shift_mm from self._alignment_result whenever it was non-None, with no way to tell whether that result came from a user's own manual Save or from the automatic Fusion -> Angle Alignment run. The automatic path's translation comes from FFT phase correlation - the exact thing manual mode exists to work around - so opening Manual Alignment after running the automatic pass silently inherited its unreliable shifts on top of the (correct) analytic rotation, reproducing the same scattered-scan symptom under a "manual" label. Manual mode should always start from identity (each angle's centroid already coincides with every other's, per the prior pivot fix) unless the user previously used this same dialog to save their own deliberate translation - which is exactly what the sidecar file already records. Seed exclusively from load_manual_alignment() and drop the self._alignment_result branch entirely; the automatic result is never an appropriate seed for a manual editing session. Co-Authored-By: Claude Sonnet 5 --- sras_viewer.py | 23 +++++++++++++---------- tools/test_gui.py | 19 ++++++++++--------- 2 files changed, 23 insertions(+), 19 deletions(-) diff --git a/sras_viewer.py b/sras_viewer.py index ea9a663..df1bb0b 100644 --- a/sras_viewer.py +++ b/sras_viewer.py @@ -2540,16 +2540,19 @@ class SrasViewerWindow(QMainWindow): ref_idx = 0 threshold_mv = self.spin_threshold_mv.value() seed: dict[int, ManualAngleParams] = {} - if (self._alignment_result is not None - and self._alignment_result.ref_angle_idx == ref_idx): - seed = {a: ManualAngleParams(t.rotation_deg, t.shift_mm) - for a, t in self._alignment_result.per_angle.items()} - threshold_mv = self._alignment_result.dc_threshold_mv - else: - 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 + # 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. + 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( diff --git a/tools/test_gui.py b/tools/test_gui.py index 83a0366..4fd705f 100644 --- a/tools/test_gui.py +++ b/tools/test_gui.py @@ -278,20 +278,21 @@ def main(): -(float(s.angles_deg[a]) - float(s.angles_deg[0]))) for a in range(s.n_angles))) - # --- Open, seeding from the still-live automatic AlignmentResult -------- + # --- Open: must NOT seed from the still-live automatic AlignmentResult -- + # The automatic result's translation comes from FFT phase correlation -- + # the very thing manual mode exists to work around -- so manual mode + # must start from identity (centroids coincide, zero shift) regardless + # of whatever the automatic run last computed. Only a previously *saved + # manual* alignment (sidecar) should ever seed this dialog. win._on_manual_alignment() check("dialog opened", win._manual_align_dialog is not None) dlg = win._manual_align_dialog check("mask prep needed no background worker (already DC-cached)", not win._job_running("manual_align_masks")) - r = win._alignment_result # still the automatic result from the block above - if r is not None: - check("seeded rotation matches the automatic result", - all(dlg._angle_params[a].rotation_deg == r.per_angle[a].rotation_deg - for a in range(s.n_angles))) - check("seeded shift matches the automatic result", - all(dlg._angle_params[a].shift_mm == r.per_angle[a].shift_mm - for a in range(s.n_angles))) + check("no manual sidecar yet -> dialog starts at identity, not the " + "automatic result", + all(dlg._angle_params[a] == compute.ManualAngleParams() + for a in range(s.n_angles))) # --- Reference angle is locked ------------------------------------------- dlg.combo_active_angle.setCurrentIndex(dlg._ref_angle_idx) From bbb075ff340ad09bc630b58d4555309f01cfc316 Mon Sep 17 00:00:00 2001 From: Thomas Ales Date: Fri, 31 Jul 2026 11:31:14 -0500 Subject: [PATCH 10/10] Add FFT cross-correlation to Manual Alignment, with tunable options MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Manual mode previously offered only Auto De-rotate (rotation from known scan angles) plus keyboard nudging, so every angle's translation had to be found entirely by eye — fine for a small correction, unworkable for the tens-of-mm scatter a real many-angle scan can have between angles before any translation search runs at all. Add an "Auto Cross-Correlate (vs Reference)" action that, for every non-reference angle, sets rotation to the known analytic angle (same as Auto De-rotate) and translation to the FFT-phase-correlation best fit against angle #0 (always the reference/ground truth). This is meant to land every angle's outline roughly stacked on the others so keyboard nudging only has to make small corrections afterward, per the intended workflow: cross-correlate first, then nudge. Exposes two tunable options, since real data may correlate better one way than the other: "Correlate on" (raw CH4 signal, minus its own minimum -- the default, robust to a shared threshold not suiting every angle's real signal level -- or the thresholded binary mask), and "Search margin" (how far a shift the phase correlation searches before wraparound bias becomes a risk). Runs on a background thread (new CrossCorrelateWorker in sras_workers.py) since correlating many high-resolution angles can take long enough to visibly freeze the dialog otherwise. The underlying per-angle correlation math is factored out of compute_angle_alignment's Step 2 into a new, reusable correlate_translation_mm, so the existing automatic Fusion -> Angle Alignment action and the new manual button now share one implementation instead of two copies of the same phase-correlation logic. Co-Authored-By: Claude Sonnet 5 --- sras_compute.py | 109 +++++++++++++++++++++++++++++-------------- sras_viewer.py | 116 ++++++++++++++++++++++++++++++++++++++-------- sras_workers.py | 54 +++++++++++++++++++++ tools/test_gui.py | 33 +++++++++++++ 4 files changed, 258 insertions(+), 54 deletions(-) diff --git a/sras_compute.py b/sras_compute.py index 6b49a6f..ee34166 100644 --- a/sras_compute.py +++ b/sras_compute.py @@ -658,6 +658,71 @@ def _working_canvas_for_pair(sras: SrasFile, ref_idx: int, a_idx: int, return origin, (n_rows, n_cols) +def correlate_translation_mm(sras: SrasFile, angle_idx: int, ref_angle_idx: int, + signal_mv: dict[int, np.ndarray], + pivot_mm: dict[int, tuple[float, float]], *, + use_mask: bool = False, + dc_threshold_mv: float = 0.0, + margin_frac: float = 0.3, + max_corr_dim: int = 1024) -> tuple[float, float]: + """FFT phase-correlation translation (shift_x_mm, shift_y_mm) that best + lines up angle_idx's CH4 content onto ref_angle_idx's, given each angle's + own alignment pivot (pivot_mm — see compute_pivot_points_mm) already + rotated about that pivot by the analytic scan-angle delta with zero + shift. Returns (0.0, 0.0) unconditionally for the reference angle. + + By default (use_mask=False) correlates on each angle's own raw CH4 + signal minus its own minimum — the same weighting _signal_centroid_mm + uses — rather than a dc_threshold_mv binary mask: a shared threshold + that doesn't suit every angle's real signal level can empty or distort + one angle's mask (the same failure mode the alignment pivot was fixed + to avoid), and correlating on the raw signal also lets phase + correlation lock onto real internal sample structure rather than just + the scan window's silhouette. Subtracting each angle's own minimum + (rather than using the signal as-is) keeps the zero-padding surrounding + the rotated image from reading as a spurious high-contrast edge against + a nonzero DC baseline. use_mask=True switches to the binary + >= dc_threshold_mv mask instead (compute_angle_alignment's original + behavior), for cases where raw-signal correlation locks onto noise. + """ + if angle_idx == ref_angle_idx: + return (0.0, 0.0) + + def corr_img(a: int) -> np.ndarray: + img = signal_mv[a] + if use_mask: + return (img >= dc_threshold_mv).astype(np.float32) + return (img - img.min()).astype(np.float32) + + img_a = corr_img(angle_idx) + img_ref = corr_img(ref_angle_idx) + + dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx) + max_dim = max(img_a.shape + img_ref.shape) + factor = max(1, int(np.ceil(max_dim / max_corr_dim))) + dx_c, dy_c = dx_ref * factor, dy_ref * factor + small_a = _block_mean_downsample(img_a, factor) + small_ref = _block_mean_downsample(img_ref, factor) + + work_origin, work_shape = _working_canvas_for_pair( + sras, ref_angle_idx, angle_idx, dx_c, dy_c, pivot_mm, margin_frac=margin_frac) + + m_a, o_a = _build_affine_canvas_to_raw( + sras, angle_idx, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin, pivot_mm) + m_ref, o_ref = _build_affine_canvas_to_raw( + sras, ref_angle_idx, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin, pivot_mm) + + rotated_a = scipy_ndimage.affine_transform( + small_a, m_a / factor, offset=o_a / factor, + output_shape=work_shape, order=0, mode="constant", cval=0.0) + embedded_ref = scipy_ndimage.affine_transform( + small_ref, m_ref / factor, offset=o_ref / factor, + output_shape=work_shape, order=0, mode="constant", cval=0.0) + + dr, dc = _phase_correlate_shift(embedded_ref, rotated_a) + return (dc * dx_c, dr * dy_c) + + def _parallel_map(fn, items, n_workers: int) -> list: """fn over items, in order, threaded when it pays.""" items = list(items) @@ -697,8 +762,6 @@ def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int, return dc4 dc4_mv = dict(enumerate(_parallel_map(dc4_for, range(n), n_workers))) - masks = {a: (img >= dc_threshold_mv).astype(np.float32) - for a, img in dc4_mv.items()} # Each angle's own alignment pivot — the CH4-signal-weighted centroid, # not the raw scan-window bbox center (see compute_pivot_points_mm). @@ -708,41 +771,17 @@ def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int, # mask empty can't silently degrade the pivot back to the bbox center. pivot_mm = compute_pivot_points_mm(sras, dc4_mv=dc4_mv) - # ---- Step 2: coarse correlation stage (downsample first, then rotate) - # Downsampling before affine_transform (not after) is what keeps this - # tractable for a v6 scan with thousands of rows/frames per angle. - max_dim = max(max(m.shape) for m in masks.values()) - factor = max(1, int(np.ceil(max_dim / 1024))) - dx_c, dy_c = dx_ref * factor, dy_ref * factor - masks_small = {a: _block_mean_downsample(m, factor) for a, m in masks.items()} - - _ticks.clear() - + # ---- Step 2: coarse translation via FFT phase correlation, on each + # angle's binarized CH4 mask against the reference's (see + # correlate_translation_mm — also the engine behind ManualAlignmentDialog's + # Auto Cross-Correlate button, there defaulting to the raw signal instead + # of a mask). def shift_for(a: int) -> tuple[float, float]: - if a == ref_angle_idx: - tick(25, 50) - return (0.0, 0.0) - work_origin, work_shape = _working_canvas_for_pair( - sras, ref_angle_idx, a, dx_c, dy_c, pivot_mm, margin_frac=0.3) - - # Coarse canvas->raw affine at native pitch, then rescale by /factor - # so it maps coarse-canvas idx -> coarse (downsampled) raw idx — - # exact for block-mean downsampling (up to the trimmed remainder). - m_a, o_a = _build_affine_canvas_to_raw( - sras, a, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin, pivot_mm) - m_ref, o_ref = _build_affine_canvas_to_raw( - sras, ref_angle_idx, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin, - pivot_mm) - rotated_a = scipy_ndimage.affine_transform( - masks_small[a], m_a / factor, offset=o_a / factor, - output_shape=work_shape, order=0, mode="constant", cval=0.0) - embedded_ref = scipy_ndimage.affine_transform( - masks_small[ref_angle_idx], m_ref / factor, offset=o_ref / factor, - output_shape=work_shape, order=0, mode="constant", cval=0.0) - - dr, dc = _phase_correlate_shift(embedded_ref, rotated_a) + shift = correlate_translation_mm( + sras, a, ref_angle_idx, dc4_mv, pivot_mm, + use_mask=True, dc_threshold_mv=dc_threshold_mv) tick(25, 50) - return (dc * dx_c, dr * dy_c) + return shift shifts_mm = dict(enumerate(_parallel_map(shift_for, range(n), n_workers))) diff --git a/sras_viewer.py b/sras_viewer.py index df1bb0b..fb7e8b5 100644 --- a/sras_viewer.py +++ b/sras_viewer.py @@ -46,7 +46,7 @@ from sras_format import ( ) from sras_workers import ( AngleAlignmentWorker, BatchCacheWorker, Ch4MaskWorker, ComputeWorker, - DcPrecomputeWorker, LoadWorker, + CrossCorrelateWorker, DcPrecomputeWorker, LoadWorker, ) faulthandler.enable() # print a native stack trace on SIGSEGV/SIGABRT/etc. @@ -818,27 +818,31 @@ class ManualAlignmentDialog(QDialog): 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 — the thing the automatic phase- - correlation step (compute_angle_alignment) sometimes gets wrong. The user - picks an "active" angle and nudges its rotation+translation with the - keyboard; Auto De-rotate sets every non-reference angle's rotation to the - known, analytic scan-angle delta without touching any translation the - user has already dialed in. 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. + misalignment is visible by eye. Reference angle (always index 0) is + ground truth and never moves; every other angle is aligned to it. The + user picks an "active" angle and nudges its rotation+translation with + the keyboard; Auto De-rotate sets every non-reference angle's rotation to + the known, analytic scan-angle delta without touching any translation; + Auto Cross-Correlate does the same rotation and additionally sets + translation to the FFT-phase-correlation best fit against the reference + (see compute.correlate_translation_mm) — meant to get every angle roughly + stacked on top of each other so keyboard nudging only has to make small + corrections, not find a coarse alignment from scratch. Save writes a + JSON sidecar next to the .sras file and hands a freshly-built, full- + resolution AlignmentResult back to the main window — the exact same + object shape compute_angle_alignment produces, so every existing + Aligned-View code path (apply_alignment, _aligned_canvas_axes, the + pixel-inspector inverse-transform) works completely unmodified. 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 own (rare) background mask-fetch step, so the main window's - existing shutdown/lifecycle plumbing covers it 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, threshold edits) stays purely local to - this dialog until Save. + 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) @@ -998,6 +1002,31 @@ class ManualAlignmentDialog(QDialog): tl.addLayout(tform) panel_l.addWidget(self.grp_mask_threshold) + # ---- Cross-Correlate (FFT) ----------------------------------------- + self.grp_correlate, cl = _group("Cross-Correlate (FFT)") + cform = _form() + self.combo_correlate_source = QComboBox() + self.combo_correlate_source.addItems( + ["Raw signal (recommended)", "Thresholded mask"]) + cform.addRow("Correlate on:", self.combo_correlate_source) + + self.spin_correlate_margin = QDoubleSpinBox() + self.spin_correlate_margin.setRange(0.05, 2.0) + self.spin_correlate_margin.setSingleStep(0.05) + self.spin_correlate_margin.setDecimals(2) + self.spin_correlate_margin.setValue(0.30) + self.spin_correlate_margin.setMinimumWidth(_SPIN_MIN_W) + cform.addRow("Search margin (× extent):", self.spin_correlate_margin) + cl.addLayout(cform) + self.btn_auto_correlate = QPushButton("Auto Cross-Correlate (vs Reference)") + cl.addWidget(self.btn_auto_correlate) + cl.addWidget(_wrap_label( + "Sets rotation to the known scan angle and translation to the " + "FFT-correlated best fit for every non-reference angle. Run this " + "first, then use manual nudging only for small corrections.", + _CSS_HINT)) + panel_l.addWidget(self.grp_correlate) + # ---- Actions ------------------------------------------------------ grp_actions, acl = _group("Actions") self.btn_auto_derotate = QPushButton("Auto De-rotate (use known angles)") @@ -1012,7 +1041,7 @@ class ManualAlignmentDialog(QDialog): panel_l.addWidget(self.lbl_status) panel_l.addStretch() - root.addWidget(_scroll_panel(panel, 300)) + root.addWidget(_scroll_panel(panel, 320)) self.combo_active_angle.currentIndexChanged.connect(self._on_active_angle_changed) self.spin_active_rotation_deg.editingFinished.connect(self._on_rotation_spin_edited) @@ -1020,6 +1049,7 @@ class ManualAlignmentDialog(QDialog): 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) @@ -1241,6 +1271,53 @@ class ManualAlignmentDialog(QDialog): f"Rotation set to the known scan angle for {n_changed} angle(s) " "(translation left untouched).") + 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 + use_mask = self.combo_correlate_source.currentIndex() == 1 + worker = CrossCorrelateWorker( + self._sras, self._ref_angle_idx, angles, self._dc4_mv, self._pivot_mm, + use_mask=use_mask, dc_threshold_mv=self.spin_mask_threshold_mv.value(), + margin_frac=self.spin_correlate_margin.value()) + self._correlate_done_count = 0 + self._correlate_total = len(angles) + self._set_controls_enabled(False) + self.lbl_status.setText(f"Cross-correlating: 0/{self._correlate_total} angle(s)…") + started = self._parent._run_worker( + "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): + self._angle_params[angle_idx] = ManualAngleParams(rotation_deg, (shift_x_mm, shift_y_mm)) + 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) + source = "thresholded mask" if self.combo_correlate_source.currentIndex() == 1 \ + else "raw signal" + self.lbl_status.setText( + f"Cross-correlated {self._correlate_done_count} angle(s) against " + f"Angle {self._ref_angle_idx} using the {source}. Nudge from here " + "for any remaining fine correction.") + def _on_save(self): threshold = self.spin_mask_threshold_mv.value() resolved = dict(self._angle_params) # already concrete floats @@ -1284,6 +1361,7 @@ class ManualAlignmentDialog(QDialog): 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_workers.py b/sras_workers.py index 0068136..5696b0e 100644 --- a/sras_workers.py +++ b/sras_workers.py @@ -346,3 +346,57 @@ class Ch4MaskWorker(QObject): self.finished.emit() except Exception as exc: self.error.emit(str(exc)) + + +class CrossCorrelateWorker(QObject): + """FFT phase-correlation translation for each of *angle_indices* against + *ref_angle_idx*, for ManualAlignmentDialog's Auto Cross-Correlate button. + + Runs on a background thread — a real many-angle, high-resolution scan's + correlation (even at its downsampled working resolution) can take long + enough that doing all of them on the GUI thread would visibly freeze the + dialog. Rotation is set to the same analytic scan-angle delta Auto + De-rotate uses alongside the correlated shift, since a translation + search is only meaningful once both angles' content is already oriented + the same way. dc4_mv/pivot_mm are the dialog's own already-in-memory + per-angle images/pivots — this worker does no fetching of its own. + """ + angle_done = pyqtSignal(int, float, float, float) # angle_idx, rotation_deg, shift_x_mm, shift_y_mm + finished = pyqtSignal() + error = pyqtSignal(str) + + def __init__(self, sras: SrasFile, ref_angle_idx: int, angle_indices: list[int], + dc4_mv: dict[int, np.ndarray], pivot_mm: dict[int, tuple[float, float]], + *, use_mask: bool, dc_threshold_mv: float, margin_frac: float): + super().__init__() + self._sras = sras + self._ref = ref_angle_idx + self._angles = angle_indices + self._dc4_mv = dc4_mv + self._pivot_mm = pivot_mm + self._use_mask = use_mask + self._threshold = dc_threshold_mv + self._margin = margin_frac + + def _one(self, a: int) -> tuple[int, float, float, float]: + theta = compute._theta_deg(self._sras, a, self._ref) + dx, dy = compute.correlate_translation_mm( + self._sras, a, self._ref, self._dc4_mv, self._pivot_mm, + use_mask=self._use_mask, dc_threshold_mv=self._threshold, + margin_frac=self._margin) + return a, theta, dx, dy + + def run(self): + try: + n_workers, _budget = compute.plan_angle_level(self._sras) + pool = ThreadPoolExecutor(max_workers=max(1, n_workers)) + try: + futures = [pool.submit(self._one, a) for a in self._angles] + for fut in as_completed(futures): + a, theta, dx, dy = fut.result() + self.angle_done.emit(a, theta, dx, dy) + finally: + pool.shutdown(wait=True) + self.finished.emit() + except Exception as exc: + self.error.emit(str(exc)) diff --git a/tools/test_gui.py b/tools/test_gui.py index 4fd705f..482e11a 100644 --- a/tools/test_gui.py +++ b/tools/test_gui.py @@ -278,6 +278,16 @@ def main(): -(float(s.angles_deg[a]) - float(s.angles_deg[0]))) for a in range(s.n_angles))) + # --- FFT phase correlation recovers a known synthetic pixel shift ------ + rng = np.random.default_rng(0) + corr_ref = np.zeros((40, 50), dtype=np.float32) + corr_ref[10:25, 15:35] = 1.0 + corr_ref += 0.05 * rng.standard_normal(corr_ref.shape).astype(np.float32) + corr_mov = np.roll(corr_ref, shift=(4, -7), axis=(0, 1)) + dr, dc = compute._phase_correlate_shift(corr_ref, corr_mov) + check("phase correlation recovers the shift that aligns mov onto ref", + (dr, dc) == (-4, 7), f"got (dr, dc)={(dr, dc)}") + # --- Open: must NOT seed from the still-live automatic AlignmentResult -- # The automatic result's translation comes from FFT phase correlation -- # the very thing manual mode exists to work around -- so manual mode @@ -344,6 +354,29 @@ def main(): check("reference angle stays identity after auto de-rotate", dlg._angle_params[dlg._ref_angle_idx].rotation_deg == 0.0) + # --- Auto Cross-Correlate: rotation + FFT-correlated shift, backgrounded - + check("cross-correlate action enabled once masks are ready", + dlg.btn_auto_correlate.isEnabled()) + dlg._on_auto_correlate() + check("auto cross-correlate completed", wait_until( + lambda: not win._job_running("manual_align_correlate"), timeout_ms=30000)) + check("auto cross-correlate set the known analytic angle for every angle", + all(abs(dlg._angle_params[a].rotation_deg + - compute._theta_deg(s, a, dlg._ref_angle_idx)) < 1e-6 + for a in range(s.n_angles) if a != dlg._ref_angle_idx)) + check("auto cross-correlate reference angle stays identity", + dlg._angle_params[dlg._ref_angle_idx] == compute.ManualAngleParams()) + check("auto cross-correlate re-enabled controls when done", + dlg.grp_correlate.isEnabled() and dlg.btn_save.isEnabled()) + check("preview canvas rebuilt after cross-correlate", + len(dlg._preview_layers) == s.n_angles) + + # The thresholded-mask option should also work end to end. + dlg.combo_correlate_source.setCurrentIndex(1) # thresholded mask + dlg._on_auto_correlate() + check("auto cross-correlate (thresholded-mask option) completed", wait_until( + lambda: not win._job_running("manual_align_correlate"), timeout_ms=30000)) + # --- Save ----------------------------------------------------------------- dlg._on_save() sidecar = compute.sidecar_path(s.path)