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_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..ee34166 --- /dev/null +++ b/sras_compute.py @@ -0,0 +1,1055 @@ +#!/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 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 +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. + +# 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) + +# 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, + 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. + + 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 + 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, + 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 + 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. + + *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 b in bounds: + run(b) + return + with ThreadPoolExecutor(max_workers=min(n_workers, len(bounds))) as pool: + list(pool.map(run, bounds)) + + +# --------------------------------------------------------------------------- +# Image computation +# --------------------------------------------------------------------------- + +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, + 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, should_stop=should_stop) + return img + + +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, 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, + 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 + 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, 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) + + 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, should_stop=should_stop) + 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: + 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 + + 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] + + +@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 + 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: + """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 _signal_centroid_mm(sras: SrasFile, angle_idx: int, + dc4_mv: np.ndarray) -> tuple[float, float]: + """Intensity-weighted centroid (mean x, mean y, weighted by CH4 signal + after subtracting this angle's own minimum) in angle_idx's own local mm + frame — the alignment pivot used in place of the raw scan-window bbox + center (see compute_pivot_points_mm). + + Weighting by the continuous DC signal, rather than a binary >= + dc_threshold_mv mask, means the pivot never depends on how well one + shared threshold happens to suit this particular angle: real signal + levels vary scan to scan, so a threshold tuned for one angle can leave + another angle's binary mask empty — and a centroid of an empty mask has + nothing to fall back to *except* the raw bbox center, silently + reproducing the exact "aligned to the scan window, not the sample" + problem this pivot exists to avoid. Falls back to the bbox center only + in the fully-degenerate case of a perfectly flat signal (nothing to + weight by at all). + """ + weights = dc4_mv - dc4_mv.min() + total = float(weights.sum()) + if total <= 0.0: + return _bbox_center_mm(sras, angle_idx) + x = sras.x_axis_mm(angle_idx) + y = sras.y_positions_mm(angle_idx) + cx = float((x * weights.sum(axis=0)).sum() / total) + cy = float((y * weights.sum(axis=1)).sum() / total) + return cx, cy + + +def compute_pivot_points_mm(sras: SrasFile, + dc4_mv: dict[int, np.ndarray] | None = None + ) -> dict[int, tuple[float, float]]: + """Per-angle alignment pivot, in each angle's own local mm frame: the + CH4-signal-weighted centroid of its own footprint (see + _signal_centroid_mm), rather than the raw scan-window bbox center. + + Pivoting on each angle's own content — instead of on wherever its + scanned window happened to sit in microscope/global XY space — is what + makes alignment purely relative *between scans* rather than to a global + coordinate system: the rotation stage's true mechanical axis need not + coincide with the scan window's geometric center, and the sample need + not be perfectly centered on that axis either, so a bbox-center pivot + leaves a residual orbital motion between angles that a content-centroid + pivot does not. + + Deliberately independent of dc_threshold_mv (the RF-mask / overlay + threshold): that value is a display/masking choice and must never + silently change where alignment pivots. + + *dc4_mv* lets a caller that has already computed each angle's CH4 mV + image (compute_angle_alignment's Step 1, or ManualAlignmentDialog's own + cache) reuse it instead of recomputing; angles missing from it are + computed fresh via dc_image_mv, which prefers a stored v5/v7 cache over + recomputing from raw waveforms. + """ + dc4_mv = dc4_mv or {} + pivots: dict[int, tuple[float, float]] = {} + for a in range(sras.n_angles): + img = dc4_mv.get(a) + if img is None: + img = dc_image_mv(sras, a, CH4_IDX) + pivots[a] = _signal_centroid_mm(sras, a, img) + return pivots + + +def _corners_in_ref_frame(sras: SrasFile, angle_idx: int, ref_idx: int, + pivot_mm: dict[int, tuple[float, float]], + shift_mm=(0.0, 0.0), + theta_deg: float | None = None) -> np.ndarray: + """Angle *angle_idx*'s bbox corners, rotated about its own alignment + pivot (pivot_mm[angle_idx] — see compute_pivot_points_mm) into the + reference frame and translated by *shift_mm*. Shape (4, 2). + + theta_deg overrides the analytic angles_deg-derived rotation used by + default — the manual-alignment path (union_canvas_mm) passes a + user-chosen rotation here (which may differ from the known scan-angle + delta) without needing a parallel code path. + """ + theta = _theta_deg(sras, angle_idx, ref_idx) if theta_deg is None else theta_deg + R = _rotation_matrix(theta) + c_a = np.array(pivot_mm[angle_idx]) + c_ref = np.array(pivot_mm[ref_idx]) + shift = np.asarray(shift_mm, dtype=np.float64) + return np.array([R @ (corner - c_a) + c_ref + shift + for corner in _bbox_corners_mm(sras, angle_idx)]) + + +def _build_affine_canvas_to_raw(sras: SrasFile, angle_idx: int, ref_idx: int, + shift_mm: tuple[float, float], + canvas_dx: float, canvas_dy: float, + canvas_origin_mm: tuple[float, float], + pivot_mm: dict[int, tuple[float, float]], + theta_deg: float | None = None + ) -> tuple[np.ndarray, np.ndarray]: + """matrix, offset s.t. raw_index = matrix @ [row_out, col_out] + offset, + 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 is _theta_deg(angle_idx, ref_idx) unless *theta_deg* + overrides it (see _corners_in_ref_frame's note, used by the manual- + alignment path), c_ref/c_a are each angle's own alignment pivot + (pivot_mm — see compute_pivot_points_mm; the CH4-signal-weighted + centroid of its own footprint, not the raw scan-window bbox center — + this keeps the sample itself centered post-rotation, minimizing + required canvas padding and, more importantly, keeping alignment + relative to the sample rather than to wherever the scan window sat in + microscope/global XY space), and A_out/D are the index<->mm scaling + matrices for the canvas pitch and this angle's own native pitch + respectively. + """ + theta = _theta_deg(sras, angle_idx, ref_idx) if theta_deg is None else theta_deg + Rinv = _rotation_matrix(theta).T + cx_a, cy_a = pivot_mm[angle_idx] + cx_ref, cy_ref = pivot_mm[ref_idx] + dx_a, dy_a = _pixel_pitch_mm(sras, angle_idx) + x0_a = float(sras.x_start_mm[angle_idx]) + y0_a = float(sras.y_positions_mm(angle_idx)[0]) + + A_out = np.array([[0.0, canvas_dx], [canvas_dy, 0.0]]) # [row,col] -> [X,Y] + b_out = np.array(canvas_origin_mm, dtype=np.float64) + D = np.array([[0.0, 1.0 / dy_a], [1.0 / dx_a, 0.0]]) # [x,y] -> [row,col] + shift = np.array(shift_mm, dtype=np.float64) + c_ref_v = np.array([cx_ref, cy_ref]) + c_a_v = np.array([cx_a, cy_a]) + origin_a = np.array([x0_a, y0_a]) + + 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, + pivot_mm: dict[int, tuple[float, float]], + margin_frac: float = 0.3 + ) -> tuple[tuple[float, float], tuple[int, int]]: + """Union of the reference's own raw bbox and angle a's raw bbox rotated + (about its own alignment pivot) into the ref frame with zero shift, + padded by margin_frac on each side — sized generously so the true + phase-correlation shift lands well inside the canvas (see + _phase_correlate_shift's wraparound note).""" + pts = np.vstack([_bbox_corners_mm(sras, ref_idx), + _corners_in_ref_frame(sras, a_idx, ref_idx, pivot_mm)]) + x_min, y_min = pts.min(axis=0) + x_max, y_max = pts.max(axis=0) + pad_x, pad_y = (x_max - x_min) * margin_frac, (y_max - y_min) * margin_frac + x_min, x_max = x_min - pad_x, x_max + pad_x + y_min, y_max = y_min - pad_y, y_max + pad_y + n_cols = int(np.ceil((x_max - x_min) / dx)) + 1 + n_rows = int(np.ceil((y_max - y_min) / abs(dy))) + 1 + origin = (x_min, y_min if dy > 0 else y_max) + return origin, (n_rows, n_cols) + + +def 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) + 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) + # 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. + _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: 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)) + tick(0, 25) + return dc4 + + dc4_mv = dict(enumerate(_parallel_map(dc4_for, range(n), n_workers))) + + # Each angle's own alignment pivot — the CH4-signal-weighted centroid, + # not the raw scan-window bbox center (see compute_pivot_points_mm). + # Reuses the dc4_mv images just computed above, so this is free, and is + # deliberately independent of dc_threshold_mv (see that function's + # docstring) so a threshold that happens to leave some angle's binary + # mask empty can't silently degrade the pivot back to the bbox center. + pivot_mm = compute_pivot_points_mm(sras, dc4_mv=dc4_mv) + + # ---- Step 2: coarse translation via FFT phase correlation, on each + # angle's binarized CH4 mask against the reference's (see + # correlate_translation_mm — also the engine behind ManualAlignmentDialog's + # Auto Cross-Correlate button, there defaulting to the raw signal instead + # of a mask). + def shift_for(a: int) -> tuple[float, float]: + shift = correlate_translation_mm( + sras, a, ref_angle_idx, dc4_mv, pivot_mm, + use_mask=True, dc_threshold_mv=dc_threshold_mv) + tick(25, 50) + return shift + + 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, pivot_mm, shifts_mm[a]) + for a in range(n)]) + x_min, y_min = corners.min(axis=0) + x_max, y_max = corners.max(axis=0) + n_cols = int(np.ceil((x_max - x_min) / dx_ref)) + 1 + n_rows = int(np.ceil((y_max - y_min) / abs(dy_ref))) + 1 + canvas_origin_mm = (float(x_min), float(y_min if dy_ref > 0 else y_max)) + + # ---- Step 4: final per-angle full-resolution affine (canvas -> raw idx) + per_angle: dict[int, AngleTransform] = {} + for a in range(n): + matrix, offset = _build_affine_canvas_to_raw( + sras, a, ref_angle_idx, shifts_mm[a], dx_ref, dy_ref, canvas_origin_mm, + pivot_mm) + per_angle[a] = AngleTransform( + _theta_deg(sras, a, ref_angle_idx), shifts_mm[a], matrix, offset) + if progress_cb: + progress_cb(75 + int((a + 1) / n * 25)) + + return AlignmentResult(ref_angle_idx, dc_threshold_mv, (n_rows, n_cols), + dx_ref, dy_ref, canvas_origin_mm, per_angle) + + +# 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], + pivot_mm: dict[int, tuple[float, float]], + margin_frac: float = 0.0 + ) -> tuple[tuple[float, float], tuple[int, int]]: + """Shared-canvas origin (mm) and (n_rows, n_cols) at pitch (dx, dy) that + contains every angle's footprint after applying its own rotation+shift — + the generalisation of compute_angle_alignment's Step 3 to arbitrary (not + just _theta_deg-analytic) per-angle rotation. Angles missing from + per_angle_params default to identity (e.g. a sidecar saved before a + rescan added more angles). + + margin_frac pads the box on every side: 0 for a final canvas (this then + reproduces compute_angle_alignment's own Step-3 math exactly, when every + angle's rotation_deg equals the analytic delta and shift_mm matches); + nonzero for ManualAlignmentDialog's downsampled preview canvas, which + needs headroom so an ordinary translation nudge of the active angle never + has to trigger a full canvas resize (see that class's docstring — an + extreme nudge can still, in principle, push content past this padding; + accepted as a known edge case, same as _phase_correlate_shift's + wraparound risk note). + """ + n = sras.n_angles + corners = np.vstack([ + _corners_in_ref_frame( + sras, a, ref_angle_idx, pivot_mm, + per_angle_params.get(a, ManualAngleParams()).shift_mm, + theta_deg=per_angle_params.get(a, ManualAngleParams()).rotation_deg) + for a in range(n)]) + x_min, y_min = corners.min(axis=0) + x_max, y_max = corners.max(axis=0) + if margin_frac: + pad_x, pad_y = (x_max - x_min) * margin_frac, (y_max - y_min) * margin_frac + x_min, x_max = x_min - pad_x, x_max + pad_x + y_min, y_max = y_min - pad_y, y_max + pad_y + n_cols = int(np.ceil((x_max - x_min) / dx)) + 1 + n_rows = int(np.ceil((y_max - y_min) / abs(dy))) + 1 + origin = (float(x_min), float(y_min if dy > 0 else y_max)) + return origin, (n_rows, n_cols) + + +def reproject_mask(sras: SrasFile, angle_idx: int, ref_angle_idx: int, + mask: np.ndarray, rotation_deg: float, + shift_mm: tuple[float, float], + canvas_dx: float, canvas_dy: float, + canvas_origin_mm: tuple[float, float], + 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 + keystroke, for only the actively-nudged angle), since it bypasses + compute_angle_alignment's phase-correlation search entirely and just + takes rotation_deg/shift_mm as given. order=0 (nearest) matches + apply_alignment's own reasoning: a binary mask must never be blended with + zero-padding. + """ + matrix, offset = _build_affine_canvas_to_raw( + sras, angle_idx, ref_angle_idx, shift_mm, canvas_dx, canvas_dy, + canvas_origin_mm, pivot_mm, theta_deg=rotation_deg) + 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], + 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 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 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) + + canvas_origin_mm, canvas_shape = union_canvas_mm( + sras, ref_angle_idx, dx_ref, dy_ref, params, pivot_mm, margin_frac=0.0) + + per_angle: dict[int, AngleTransform] = {} + for a in range(n): + p = params[a] + matrix, offset = _build_affine_canvas_to_raw( + sras, a, ref_angle_idx, p.shift_mm, dx_ref, dy_ref, + canvas_origin_mm, pivot_mm, theta_deg=p.rotation_deg) + per_angle[a] = AngleTransform(p.rotation_deg, p.shift_mm, matrix, offset) + + return AlignmentResult(ref_angle_idx, dc_threshold_mv, canvas_shape, + dx_ref, dy_ref, canvas_origin_mm, per_angle) + + +# ---- 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") + + +# 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 2): + { + "schema_version": 2, + "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": _SIDECAR_SCHEMA_VERSION, + "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") != _SIDECAR_SCHEMA_VERSION: + 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_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 9b7e3ca..fb7e8b5 100644 --- a/sras_viewer.py +++ b/sras_viewer.py @@ -10,1169 +10,151 @@ 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 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. +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 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.lines import Line2D from matplotlib.path import Path as MplPath +from PyQt6.QtCore import QObject, Qt, QThread, pyqtSignal +from PyQt6.QtGui import QAction, QKeyEvent +from PyQt6.QtWidgets import ( + QApplication, QButtonGroup, QCheckBox, QComboBox, QDialog, QDialogButtonBox, + QDoubleSpinBox, QFileDialog, QFormLayout, QFrame, QGroupBox, QHBoxLayout, + QLabel, QMainWindow, QMessageBox, QProgressDialog, QPushButton, QRadioButton, + QScrollArea, QSizePolicy, QSpinBox, QSplitter, QVBoxLayout, QWidget, +) + +import sras_compute as compute +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, Ch4MaskWorker, ComputeWorker, + CrossCorrelateWorker, 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 - -# 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"] +# (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"), +} -def _parse_preamble(preamble: str) -> dict[str, float]: - """Extract YMULT, YOFF, YZERO from a Tektronix WFMOutpre string. +_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;" - 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 +# Side-panel column widths (the scroll areas that hold the controls). +_LEFT_PANEL_W = 288 +_RIGHT_PANEL_W = 272 - -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 +# Minimum width for a spin box so its value + suffix are never clipped. +_SPIN_MIN_W = 96 # --------------------------------------------------------------------------- -# File parser +# Small layout helpers # --------------------------------------------------------------------------- -class SrasFile: - """Parsed in-memory representation of a v2–v6 .sras file. +def _wrap_label(text: str = "", css: str | None = None) -> QLabel: + """A word-wrapped QLabel that reports its *wrapped* height to the layout. - 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 - ``(n_rows[a], n_channels, n_frames[a], samples_per_frame)``. + 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. """ - - 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 == 6: - 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 - - 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) - 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() - - # ---- 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] - - # ------------------------------------------------------------------ - # v5 writer - # ------------------------------------------------------------------ - - 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*. - - 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. - - *freq_images* / *dc4_images* / *dc3_images*: - shape (n_angles, n_rows, n_frames) float32. - - *progress_cb*: optional callable(fraction: float) for UI updates. - - 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. - """ - 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).") - - import shutil - - dest = Path(dest_path) - src = self.path - n_rows0 = int(self.n_rows[0]) - n_frames0 = int(self.n_frames[0]) - - # --- Copy the source file verbatim, then patch the header ------- - shutil.copy2(str(src), str(dest)) - - waveform_bytes = (self.n_angles * n_rows0 * self.n_channels - * n_frames0 * self.samples_per_frame - * self.bytes_per_sample) - - 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)) - - # 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)) - - # 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) - f.truncate() - - # --- 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() - - # ------------------------------------------------------------------ - # 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. + 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. """ - 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 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 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. - - 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: v5 precomputed images ---------------------------------- - can_use_precomputed = ( - sras.precomputed_freq_mhz 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) - ) - if can_use_precomputed: - freq_img = sras.precomputed_freq_mhz[angle_idx].copy() - dc4_img = sras.precomputed_dc4_mv[angle_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, mirroring PreprocessWorker's existing precedent.""" - 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: convert ADC counts → mV - 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 - 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 PreprocessWorker(QObject): - """Compute all-angle FFT and DC images and write a v5 file. - - 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). - """ - progress = pyqtSignal(int) # 0–100 - finished = pyqtSignal(str) # empty = success, else error message - - def __init__(self, sras: SrasFile, dest_path: str, apply_bg_sub: bool): - super().__init__() - self._sras = sras - self._dest_path = dest_path - 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) - - 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) - - # 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]) - - self.progress.emit(int((aidx + 1) / n * 90)) - - 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)) - - -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. - """ - 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)) + 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 # --------------------------------------------------------------------------- @@ -1193,13 +175,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()) @@ -1209,7 +189,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: @@ -1223,12 +202,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 # --------------------------------------------------------------------------- @@ -1236,18 +232,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): @@ -1256,7 +252,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 @@ -1266,16 +262,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) # ------------------------------------------------------------------ @@ -1290,7 +286,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( @@ -1357,27 +353,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: @@ -1385,9 +377,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: @@ -1413,8 +404,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: @@ -1433,11 +424,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) @@ -1451,14 +441,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] @@ -1470,59 +457,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): @@ -1544,12 +531,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 @@ -1571,19 +558,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 @@ -1608,29 +595,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, ) @@ -1664,8 +646,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) @@ -1673,18 +655,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) @@ -1714,30 +695,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: @@ -1747,13 +724,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( @@ -1761,12 +737,636 @@ 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()) +# --------------------------------------------------------------------------- +# 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. 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 background mask-fetch and cross-correlate steps, so the main + window's existing shutdown/lifecycle plumbing covers both for free, and + it emits alignment_saved / alignment_cleared signals for the two moments + that should actually mutate the main window's persistent state — + everything else (nudging, Auto De-rotate, Auto Cross-Correlate, threshold + edits) stays purely local to this dialog until Save. + """ + + alignment_saved = pyqtSignal(object, str) # AlignmentResult, sidecar path (str) + alignment_cleared = pyqtSignal() + + _PREVIEW_MARGIN_FRAC = 0.15 + _BASE_ALPHA = 0.42 + _ACTIVE_ALPHA = 0.75 + _MAX_PREVIEW_DIM = 1024 + + 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._pivot_mm: dict[int, tuple[float, float]] = {} + self._preview_layers: dict[int, np.ndarray] = {} + self._preview_origin_mm = (0.0, 0.0) + self._preview_shape = (1, 1) + self._preview_dx_mm = self._preview_dy_mm = 1.0 + self._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) + + # ---- 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)") + 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, 320)) + + self.combo_active_angle.currentIndexChanged.connect(self._on_active_angle_changed) + self.spin_active_rotation_deg.editingFinished.connect(self._on_rotation_spin_edited) + self.spin_active_shift_x_mm.editingFinished.connect(self._on_shift_spin_edited) + self.spin_active_shift_y_mm.editingFinished.connect(self._on_shift_spin_edited) + self.spin_mask_threshold_mv.editingFinished.connect(self._on_mask_threshold_edited) + self.btn_auto_derotate.clicked.connect(self._on_auto_derotate) + self.btn_auto_correlate.clicked.connect(self._on_auto_correlate) + self.btn_save.clicked.connect(self._on_save) + self.btn_clear.clicked.connect(self._on_clear) + self.btn_close.clicked.connect(self.close) + self.canvas.nudge_translate.connect(self._on_nudge_translate) + self.canvas.nudge_rotate.connect(self._on_nudge_rotate) + + 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() + # 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.") + + def _recompute_masks_small(self): + """Threshold + downsample every angle's already-in-memory full-res + CH4 mV image. Cheap (a compare + block-mean), so this re-runs in + full whenever the mask-threshold spin box changes — no re-fetch. + Purely for the overlay's visuals — 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 = { + 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, + 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, self._pivot_mm) + 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._pivot_mm) + 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_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 + try: + path = save_manual_alignment(self._sras, self._ref_angle_idx, threshold, resolved) + result = build_manual_alignment(self._sras, self._ref_angle_idx, threshold, + resolved, self._pivot_mm) + 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.grp_correlate.setEnabled(enabled) + self.btn_auto_derotate.setEnabled(enabled) + self.btn_save.setEnabled(enabled) + self.btn_clear.setEnabled(enabled) + + # --------------------------------------------------------------------------- # Main window # --------------------------------------------------------------------------- @@ -1776,25 +1376,28 @@ 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 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 - self._preprocess_thread: QThread | None = None + # Convert menu: batch DC/FFT compute-and-store (v6 -> v7) + self._batch_errors: list[str] = [] # Display-only settings (colormap, grating) never trigger a # recompute — they're applied to cached data on redraw. DC images @@ -1807,22 +1410,71 @@ 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] = {} + self._manual_align_dialog: ManualAlignmentDialog | None = None self._build_ui() 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 # ------------------------------------------------------------------ @@ -1834,85 +1486,87 @@ 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) + panel_layout.setSpacing(8) - # File - grp_file = QGroupBox("File") - fl = QVBoxLayout(grp_file) + # ---- 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("color: #888; font-size: 11px;") + 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) + # ---- Scan 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("font-size: 11px;") + lbl = _wrap_label(f"{key}: —", _CSS_INFO) il.addWidget(lbl) self._info[key] = lbl - # Frame count warning (hidden until needed) - self.lbl_frame_warn = QLabel("") - self.lbl_frame_warn.setWordWrap(True) - self.lbl_frame_warn.setStyleSheet("color: #e07000; font-size: 11px;") + + # frame-count / format notes + self.lbl_frame_warn = _wrap_label("", _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.setWordWrap(True) - self.lbl_dc_precompute.setStyleSheet("color: #4a90d9; font-size: 11px;") + + # 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) + # ---- View settings ---------------------------------------------- + grp_view, vl = _group("View Settings") + + view_form = _form() - # Angle - 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) - # Channel - 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) - # DC threshold (for RF / CH1 masking) sep = QFrame() sep.setFrameShape(QFrame.Shape.HLine) sep.setStyleSheet("color: #555;") vl.addWidget(sep) - self.grp_threshold = QGroupBox("RF Mask Threshold (CH1 only)") - tl = QVBoxLayout(self.grp_threshold) - thr_row = QHBoxLayout() - thr_row.addWidget(QLabel("DC threshold:")) + # DC threshold (for RF / CH1 masking) + 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) @@ -1920,11 +1574,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") # updated on file load - self.lbl_threshold_adc.setStyleSheet("font-size: 11px; color: #888;") + 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) @@ -1951,41 +1606,17 @@ 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) --------------------------------- - grp_roi = QGroupBox("ROI (Region of Interest)") - rl = QVBoxLayout(grp_roi) + # ---- ROI --------------------------------------------------------- + grp_roi, rl = _group("ROI (Region of Interest)") self.btn_draw_roi = QPushButton("Draw ROI") self.btn_draw_roi.setCheckable(True) @@ -1993,8 +1624,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) @@ -2015,110 +1645,119 @@ 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("font-size: 11px; color: #aaa;") rl.addWidget(lbl) panel_layout.addWidget(grp_roi) panel_layout.addStretch() + return _scroll_panel(panel, _LEFT_PANEL_W) - # ---- Display Options group (added to right panel below) ------------ - grp_display = QGroupBox("Display Options") - dl = QVBoxLayout(grp_display) + def _build_canvases(self) -> QWidget: + splitter = QSplitter(Qt.Orientation.Vertical) + splitter.setChildrenCollapsible(False) - cmr = QHBoxLayout() - cmr.addWidget(QLabel("Colormap:")) + 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) + 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) + 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, 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) + 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.setMinimumWidth(_SPIN_MIN_W) + self.spin_grating_um.editingFinished.connect(self._on_grating_changed) + 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, dl = _group("Display Options") + + 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: 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(8) + layout.addWidget(self.grp_velocity) + layout.addWidget(grp_display) + layout.addStretch() + return _scroll_panel(right_panel, _RIGHT_PANEL_W) - 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") 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 +1767,31 @@ 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( + "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 # ------------------------------------------------------------------ @@ -2146,59 +1810,48 @@ 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 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. + # 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 @@ -2207,6 +1860,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() @@ -2226,13 +1899,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() # ------------------------------------------------------------------ @@ -2243,33 +1911,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 s.precomputed_freq_mhz is not None: + if s.version in (6, 7): + notes.append("v6/v7 format: rows / frames / x_start are per-angle") + + 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"v5: precomputed images present{bg_note} — display is instant") - if s.version == 6: - notes.append("v6 format: rows / frames / x_start are per-angle") + 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)) # ------------------------------------------------------------------ @@ -2278,60 +1954,52 @@ 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) - # 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) - # 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.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) + + 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() 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 — @@ -2339,27 +2007,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): @@ -2376,8 +2136,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() @@ -2388,12 +2149,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: @@ -2401,133 +2162,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) @@ -2536,28 +2188,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: @@ -2567,138 +2217,166 @@ 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") + if result is None: + return # cancelled mid-compute; the partial image must not cache 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 @@ -2712,135 +2390,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 # ------------------------------------------------------------------ @@ -2851,121 +2442,117 @@ 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() # ------------------------------------------------------------------ - # 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._job_running("batch"): return - if self._preprocess_thread is not None: + 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 paths: 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), - "SRAS files (*.sras);;All files (*)", + + self._batch_errors = [] + 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 dest: - return - if self._preprocess_thread is not None: - return # a second trigger snuck in while the save dialog was open + if not started: + 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 - - 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._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._preprocess_act.setEnabled(False) + self._batch_dc_act.setEnabled(False) + self._batch_fft_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)…" - ) + "batch", f"Batch computing {label} for {len(paths)} file(s)…", + maximum=100) - self._preprocess_thread.start() + def _on_batch_file_done(self, path: str, err: str): + if err: + self._batch_errors.append(f"{Path(path).name} — {err}") + self._show_progress("batch", f"Processed {Path(path).name}…") - def _on_preprocess_progress(self, pct: int): - if self._progress_dlg is not None: - self._progress_dlg.setValue(pct) + def _on_batch_finished(self, paths: list[str]): + self._close_progress("batch") - def _on_preprocess_done(self, error_msg: str): - self._close_progress() - if error_msg: - self.statusBar().showMessage(f"Pre-process failed: {error_msg}") + 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): - # 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 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 _after_batch(self): + self._batch_dc_act.setEnabled(True) + self._batch_fft_act.setEnabled(True) # ------------------------------------------------------------------ # Fusion: angle alignment @@ -2974,46 +2561,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: @@ -3031,47 +2603,116 @@ class SrasViewerWindow(QMainWindow): f"canvas {nc}×{nr} px).") self._refresh_display() - def _on_aligned_view_toggled(self, checked: bool): + # ------------------------------------------------------------------ + # 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] = {} + # 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( + 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._redraw_image(self._current_image) + 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 # ------------------------------------------------------------------ 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", "_preprocess_thread", - "_dc_precompute_thread", "_alignment_thread"): - t = getattr(self, attr, None) - if t is not None: - t.quit() - t.wait(2000) + 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. + 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(5000) super().closeEvent(event) @@ -3079,8 +2720,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..5696b0e --- /dev/null +++ b/sras_workers.py @@ -0,0 +1,402 @@ +#!/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 +from concurrent.futures.process import BrokenProcessPool + +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) + +# 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 + 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(CancellableWorker): + """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, should_stop=self._stopped) + else: + 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(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 + 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 + + def _one_angle(self, a: int) -> tuple[int, np.ndarray, np.ndarray]: + # 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, 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)} + for fut in as_completed(futures): + if self._stop: + 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)) + + +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 _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_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 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 + for p in paths + } + for fut in as_completed(futures): + path = futures[fut] + try: + err = fut.result() + except BrokenProcessPool: + unresolved.append(path) + continue + except Exception as exc: + err = str(exc) + done += 1 + 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() + + +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)) + + +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)) + + +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/__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..07386c3 --- /dev/null +++ b/tools/make_test_sras.py @@ -0,0 +1,191 @@ +#!/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, + geometry: list[tuple[int, int]] | None = None) -> tuple[bytes, dict]: + rng = np.random.default_rng(seed) + src_geom = geometry or _GEOMETRY + geom = [src_geom[a % len(src_geom)] 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, + 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") + 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() diff --git a/tools/test_gui.py b/tools/test_gui.py new file mode 100644 index 0000000..482e11a --- /dev/null +++ b/tools/test_gui.py @@ -0,0 +1,485 @@ +#!/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, 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, 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 + +_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" + 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("\nmanual alignment (Fusion)") + check("manual alignment action enabled", win._manual_align_act.isEnabled()) + + # --- Alignment pivot is a signal-weighted centroid, not the raw bbox -- + # center, and is independent of any DC threshold (so a threshold that + # happens to leave a real angle's binary mask empty can't silently + # degrade the pivot back to the bbox center). + corner_signal = np.zeros(s.image_shape(0), dtype=np.float32) + corner_signal[0, 0] = 1.0 # single spike -> weighted centroid is exact + expected_corner = (float(s.x_axis_mm(0)[0]), float(s.y_positions_mm(0)[0])) + centroid = compute._signal_centroid_mm(s, 0, corner_signal) + check("signal-weighted centroid of a single spike pixel is that pixel exactly", + np.allclose(centroid, expected_corner), f"{centroid} vs {expected_corner}") + bbox_center = compute._bbox_center_mm(s, 0) + check("signal centroid differs from the raw scan-window bbox center", + not np.allclose(centroid, bbox_center), + f"centroid {centroid} vs bbox center {bbox_center}") + + # 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 " + "math convention)", + all(np.isclose(compute._theta_deg(s, a, 0), + -(float(s.angles_deg[a]) - float(s.angles_deg[0]))) + for a in range(s.n_angles))) + + # --- FFT phase correlation recovers a known synthetic pixel shift ------ + rng = np.random.default_rng(0) + corr_ref = np.zeros((40, 50), dtype=np.float32) + corr_ref[10:25, 15:35] = 1.0 + corr_ref += 0.05 * rng.standard_normal(corr_ref.shape).astype(np.float32) + corr_mov = np.roll(corr_ref, shift=(4, -7), axis=(0, 1)) + dr, dc = compute._phase_correlate_shift(corr_ref, corr_mov) + check("phase correlation recovers the shift that aligns mov onto ref", + (dr, dc) == (-4, 7), f"got (dr, dc)={(dr, dc)}") + + # --- 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")) + 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) + 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) + + # --- 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) + check("sidecar file written", sidecar.exists()) + raw = json.loads(sidecar.read_text()) if sidecar.exists() else {} + 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))) + 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()) + + # --- 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): + 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) + 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()