diff --git a/sras_average.py b/sras_average.py index 0ddc876..d997a15 100644 --- a/sras_average.py +++ b/sras_average.py @@ -15,25 +15,24 @@ Options: Default: include a partial average for the last group. """ -import sys -import struct import argparse -import numpy as np +import shutil +import struct +import sys from pathlib import Path -# --------------------------------------------------------------------------- -# Header format — must match sras_viewer.py exactly -# --------------------------------------------------------------------------- +import numpy as np -HDR_FMT = ">4sBHHffffIIdBB" -HDR_SIZE = struct.calcsize(HDR_FMT) # 43 bytes +from sras_format import HDR_FMT, HDR_SIZE, SrasFile + +_SUPPORTED = (2, 3, 4) def parse_args(): p = argparse.ArgumentParser( description="Average every N waveforms in a .sras file and write a new file." ) - p.add_argument("input", help="Input .sras file") + p.add_argument("input", help="Input .sras file") p.add_argument("output", help="Output .sras file") p.add_argument("--n", type=int, required=True, metavar="N", help="Number of consecutive frames to average into one") @@ -42,143 +41,68 @@ def parse_args(): return p.parse_args() -def read_sras(path: Path): - """Read all sections of a .sras file and return them as a dict.""" - with open(path, "rb") as f: - header_bytes = f.read(HDR_SIZE) - fields = struct.unpack(HDR_FMT, header_bytes) - - (magic, ver, n_angles, n_rows, x_start, x_delta, vel, freq, - n_frames_hdr, spf, sr, bps, n_ch) = fields - - if magic != b"SRAS": - raise ValueError(f"Not a .sras file (bad magic: {magic!r})") - if ver not in (2, 3, 4): - raise ValueError(f"Unsupported .sras version: {ver}") - - with open(path, "rb") as f: +def read_header_sections(sras: SrasFile) -> bytes: + """The raw bytes between the header and the waveform data (angle table, + row table, preambles, background), copied through verbatim so nothing is + lost in a re-encode.""" + with open(sras.path, "rb") as f: f.seek(HDR_SIZE) - - angles = f.read(n_angles * 4) # big-endian float32 array, raw bytes - y_pos = f.read(n_rows * 4) # big-endian float32 array, raw bytes - - preambles = [] # list of raw bytes (length-prefixed strings) - if ver >= 3: - for _ in range(n_ch): - (length,) = struct.unpack(">H", f.read(2)) - preambles.append(f.read(length)) - - background = b"" # raw bytes for the v4 background block - if ver >= 4: - (n_bg,) = struct.unpack(">I", f.read(4)) - background = f.read(n_bg) - - raw = f.read() # all waveform data - - # ----------------------------------------------------------------------- - # Determine actual frame count from file size (the header value can be - # wrong — the viewer does the same correction) - # ----------------------------------------------------------------------- - total_samples = len(raw) // bps - samples_per_pixel = n_ch * spf # samples in one (angle, row, frame) cell - samples_per_full = n_angles * n_rows * samples_per_pixel - - actual_n_frames = total_samples // (n_angles * n_rows * samples_per_pixel) - good_bytes = actual_n_frames * samples_per_full * bps - - # Decode waveform data - dtype = np.int8 if bps == 1 else ">i2" - data = np.frombuffer(raw[:good_bytes], dtype=dtype) - data = data.reshape(n_angles, n_rows, n_ch, actual_n_frames, spf) - # Work in int16 (safe intermediate for both int8 and int16 inputs) - data = data.astype(np.int16) - - return { - "ver": ver, "n_angles": n_angles, "n_rows": n_rows, - "x_start": x_start, "x_delta": x_delta, "vel": vel, "freq": freq, - "n_frames_hdr": n_frames_hdr, "spf": spf, "sr": sr, - "bps": bps, "n_ch": n_ch, - "angles_raw": angles, "y_pos_raw": y_pos, - "preambles": preambles, "background": background, - "data": data, # shape: (n_angles, n_rows, n_ch, n_frames, spf), int16 - } + return f.read(sras._data_offset - HDR_SIZE) -def average_frames(data: np.ndarray, n: int, discard_remainder: bool) -> np.ndarray: - """Average every N frames along axis 3. +def average_rows(block: np.ndarray, n: int, discard_remainder: bool) -> np.ndarray: + """Average every N frames of one angle's (n_rows, n_ch, n_frames, spf) + block. Returns int16 of shape (n_rows, n_ch, n_out, spf). - data shape: (n_angles, n_rows, n_ch, n_frames, spf) - Returns array of shape (n_angles, n_rows, n_ch, n_out, spf). + Averaging is done in float32 and rounded on cast, matching numpy's mean + followed by an int16 cast in the original implementation. """ - n_frames = data.shape[3] - - n_full = n_frames // n + n_frames = block.shape[2] + n_full = n_frames // n remainder = n_frames % n - # Average full groups using reshape-trick (no Python loop) - if n_full > 0: - full = data[:, :, :, :n_full * n, :] # trim to full groups - full = full.reshape(data.shape[0], data.shape[1], data.shape[2], - n_full, n, data.shape[4]) # (..., n_out, n, spf) - averaged = full.mean(axis=4).astype(np.int16) # (..., n_out, spf) - else: - averaged = np.empty((*data.shape[:3], 0, data.shape[4]), dtype=np.int16) + parts = [] + if n_full: + full = block[:, :, :n_full * n, :].astype(np.float32) + full = full.reshape(block.shape[0], block.shape[1], n_full, n, block.shape[3]) + parts.append(full.mean(axis=3).astype(np.int16)) + if remainder and not discard_remainder: + tail = block[:, :, n_full * n:, :].astype(np.float32) + parts.append(tail.mean(axis=2, keepdims=True).astype(np.int16)) - if remainder > 0 and not discard_remainder: - tail = data[:, :, :, n_full * n:, :] # shape (..., remainder, spf) - tail_avg = tail.mean(axis=3, keepdims=True).astype(np.int16) - averaged = np.concatenate([averaged, tail_avg], axis=3) - - return averaged + if not parts: + return np.empty((*block.shape[:2], 0, block.shape[3]), dtype=np.int16) + return parts[0] if len(parts) == 1 else np.concatenate(parts, axis=2) -def write_sras(path: Path, src: dict, data_out: np.ndarray): - """Write a new .sras file with the averaged waveform data.""" - ver = src["ver"] - bps = src["bps"] - n_out = data_out.shape[3] +def write_averaged(out_path: Path, sras: SrasFile, mid_sections: bytes, + n: int, discard_remainder: bool) -> int: + """Stream each angle through the averager, writing as we go so peak RAM + stays at one angle's block rather than the whole file.""" + bps = sras.bytes_per_sample + n_frames_in = int(sras.n_frames[0]) + n_out = n_frames_in // n + if n_frames_in % n and not discard_remainder: + n_out += 1 - # Pack header — update only n_frames_hdr; everything else stays the same header = struct.pack( - HDR_FMT, - b"SRAS", - ver, - src["n_angles"], - src["n_rows"], - src["x_start"], - src["x_delta"], - src["vel"], - src["freq"], - n_out, # updated frame count - src["spf"], - src["sr"], - bps, - src["n_ch"], + HDR_FMT, b"SRAS", sras.version, sras.n_angles, int(sras.n_rows[0]), + float(sras.x_start_mm[0]), float(sras.x_delta_mm), + sras.velocity_mm_s, sras.laser_freq_hz, + n_out, # updated frame count + sras.samples_per_frame, sras.sample_rate_hz, bps, sras.n_channels, ) - # Encode waveform data back to original dtype - if bps == 1: - raw_out = np.clip(data_out, -128, 127).astype(np.int8).tobytes() - else: - # big-endian int16 - raw_out = data_out.astype(">i2").tobytes() - - with open(path, "wb") as f: + with open(out_path, "wb") as f: f.write(header) - f.write(src["angles_raw"]) - f.write(src["y_pos_raw"]) - - if ver >= 3: - for preamble_bytes in src["preambles"]: - f.write(struct.pack(">H", len(preamble_bytes))) - f.write(preamble_bytes) - - if ver >= 4: - bg = src["background"] - f.write(struct.pack(">I", len(bg))) - f.write(bg) - - f.write(raw_out) + f.write(mid_sections) + for a in range(sras.n_angles): + averaged = average_rows(sras.data[a], n, discard_remainder) + if bps == 1: + f.write(np.clip(averaged, -128, 127).astype(np.int8).tobytes()) + else: + f.write(averaged.astype(">i2").tobytes()) + return n_out def main(): @@ -188,32 +112,35 @@ def main(): print("Error: --n must be at least 1.", file=sys.stderr) sys.exit(1) - in_path = Path(args.input) + in_path = Path(args.input) out_path = Path(args.output) if not in_path.exists(): print(f"Error: input file not found: {in_path}", file=sys.stderr) sys.exit(1) - if out_path.resolve() == in_path.resolve(): print("Error: output path must differ from input path.", file=sys.stderr) sys.exit(1) print(f"Reading {in_path} ...", flush=True) - src = read_sras(in_path) + sras = SrasFile(str(in_path)) + if sras.version not in _SUPPORTED: + print(f"Error: unsupported .sras version: {sras.version} " + f"(this tool handles v{'/v'.join(map(str, _SUPPORTED))})", + file=sys.stderr) + sys.exit(1) - n_frames_in = src["data"].shape[3] - print(f" Version : v{src['ver']}") - print(f" Angles : {src['n_angles']}") - print(f" Rows : {src['n_rows']}") + n_frames_in = int(sras.n_frames[0]) + print(f" Version : v{sras.version}") + print(f" Angles : {sras.n_angles}") + print(f" Rows : {int(sras.n_rows[0])}") print(f" Frames (actual): {n_frames_in}") - print(f" Channels : {src['n_ch']}") - print(f" Samples/frame : {src['spf']}") - print(f" Bytes/sample : {src['bps']}") + print(f" Channels : {sras.n_channels}") + print(f" Samples/frame : {sras.samples_per_frame}") + print(f" Bytes/sample : {sras.bytes_per_sample}") if args.n == 1: print("--n 1: no averaging needed; copying file as-is.") - import shutil shutil.copy2(in_path, out_path) print(f"Wrote {out_path}") return @@ -223,27 +150,25 @@ def main(): "The entire dataset will be averaged into a single frame.") print(f"\nAveraging every {args.n} frames ...", flush=True) - data_out = average_frames(src["data"], args.n, args.discard_remainder) - n_frames_out = data_out.shape[3] + print(f"\nWriting {out_path} ...", flush=True) + mid_sections = read_header_sections(sras) + n_frames_out = write_averaged(out_path, sras, mid_sections, + args.n, args.discard_remainder) - n_full = n_frames_in // args.n - remainder = n_frames_in % args.n + n_full = n_frames_in // args.n + remainder = n_frames_in % args.n if remainder and not args.discard_remainder: status = f"({n_full} full groups + 1 partial group of {remainder})" - elif remainder and args.discard_remainder: + elif remainder: status = f"({n_full} full groups, {remainder} trailing frames discarded)" else: status = f"({n_full} full groups)" - print(f" {n_frames_in} frames -> {n_frames_out} frames {status}") - print(f"\nWriting {out_path} ...", flush=True) - write_sras(out_path, src, data_out) - - in_mb = in_path.stat().st_size / 1024**2 + in_mb = in_path.stat().st_size / 1024**2 out_mb = out_path.stat().st_size / 1024**2 print(f" Input size : {in_mb:.1f} MB") - print(f" Output size: {out_mb:.1f} MB ({out_mb/in_mb*100:.1f}% of input)") + print(f" Output size: {out_mb:.1f} MB ({out_mb / in_mb * 100:.1f}% of input)") print("Done.") diff --git a/sras_compute.py b/sras_compute.py new file mode 100644 index 0000000..280cc87 --- /dev/null +++ b/sras_compute.py @@ -0,0 +1,577 @@ +#!/usr/bin/env python3 +"""Image computation and angle alignment for .sras scans. + +Depends only on numpy/scipy (+ optional pyfftw) and sras_format, so a +multiprocessing child can import it without loading Qt or matplotlib — +which matters because Python 3.14 on macOS spawns rather than forks. +""" + +import os +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass + +import numpy as np +import scipy.fft as scipy_fft +import scipy.ndimage as scipy_ndimage + +from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, SrasFile, adc_to_mv + +# --------------------------------------------------------------------------- +# FFT backend +# --------------------------------------------------------------------------- + +try: + import pyfftw + pyfftw.interfaces.cache.enable() + PYFFTW_AVAILABLE = True +except ImportError: + PYFFTW_AVAILABLE = False + +_fft_backend = "numpy" # "numpy" or "pyfftw"; set via set_fft_backend() + + +def set_fft_backend(name: str): + """Select the rfft implementation. Module-level state, so it must be set + explicitly inside each multiprocessing child — it does not survive a + spawn.""" + global _fft_backend + _fft_backend = name if (name != "pyfftw" or PYFFTW_AVAILABLE) else "numpy" + + +def get_fft_backend() -> str: + return _fft_backend + + +def _do_rfft(x: np.ndarray, n: int | None = None, axis: int = -1, + workers: int = 1) -> np.ndarray: + """Dispatch rfft to the selected backend with optional multithreading.""" + if _fft_backend == "pyfftw" and PYFFTW_AVAILABLE: + return pyfftw.interfaces.numpy_fft.rfft(x, n=n, axis=axis, threads=workers) + return scipy_fft.rfft(x, n=n, axis=axis, workers=workers) + + +# --------------------------------------------------------------------------- +# Chunking / parallel budget +# --------------------------------------------------------------------------- +# +# Rows are batched so the float32 working buffers for one chunk stay under a +# memory budget. A fixed row count (the original design) works fine for small +# legacy scans but is catastrophic for a v6 scan with a large per-angle +# frame/sample count — e.g. a 7500-frame x 2500-sample angle needs ~2.4 GB for +# a single 32-row chunk. +# +# With chunks running concurrently the budget has to cover *all* live chunks at +# once. Note that on a large scan chunk_rows is already clamped to its floor of +# 1 row (one row alone is ~75 MB of float32 at 7507x2500), so shrinking the +# per-chunk size cannot buy more concurrency — the worker count must be derived +# from the budget instead. See _plan_chunks. + +_TOTAL_BYTES_BUDGET = int(os.environ.get("SRAS_MEM_BUDGET_MB", 1536)) * 1024 * 1024 +_CHUNK_ROWS_MAX = 32 # cap for small scans (original behavior) +_MAX_WORKERS = int(os.environ.get("SRAS_MAX_WORKERS", 0)) or (os.cpu_count() or 4) + +# An rfft chunk holds, live at once: the float32 input, the complex64 +# transform, and the float32 power spectrum — roughly 3x the input buffer. +_FFT_LIVE_MULTIPLIER = 3 + + +def _chunk_rows_for(n_frames: int, samples_per_frame: int, + budget: int = _TOTAL_BYTES_BUDGET) -> int: + bytes_per_row = max(1, n_frames * samples_per_frame * 4) # float32 + return int(max(1, min(_CHUNK_ROWS_MAX, budget // bytes_per_row))) + + +def _plan_chunks(n_rows: int, n_frames: int, samples_per_frame: int, + live_multiplier: int = 1, + max_workers: int | None = None) -> tuple[int, int]: + """(chunk_rows, n_workers) sized so live_multiplier x chunk_rows x + n_frames x samples x 4 bytes x n_workers stays within the budget.""" + per_chunk_budget = max(1, _TOTAL_BYTES_BUDGET // live_multiplier) + chunk_rows = _chunk_rows_for(n_frames, samples_per_frame, per_chunk_budget) + + chunk_bytes = max(1, live_multiplier * chunk_rows * n_frames * samples_per_frame * 4) + cap = max_workers if max_workers is not None else _MAX_WORKERS + n_workers = int(max(1, min(cap, + _TOTAL_BYTES_BUDGET // chunk_bytes, + -(-n_rows // chunk_rows)))) + return chunk_rows, n_workers + + +def _map_row_chunks(n_rows: int, chunk_rows: int, n_workers: int, fn): + """Apply fn(r0, r1) over row chunks, in parallel when it pays. + + Chunks write to disjoint output slices, so no locking is needed. numpy + ufuncs, scipy's pocketfft, and memmap page faults all release the GIL, so + threads give real parallelism here — and on a very large file they also + keep many more page-fault requests in flight, which is what the I/O path + wants. + """ + bounds = [(r0, min(r0 + chunk_rows, n_rows)) + for r0 in range(0, n_rows, chunk_rows)] + if n_workers <= 1 or len(bounds) == 1: + for r0, r1 in bounds: + fn(r0, r1) + return + with ThreadPoolExecutor(max_workers=min(n_workers, len(bounds))) as pool: + list(pool.map(lambda b: fn(*b), bounds)) + + +# --------------------------------------------------------------------------- +# Image computation +# --------------------------------------------------------------------------- + +def compute_dc_image(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray: + """Mean of each waveform → (n_rows, n_frames) float32, in ADC counts.""" + n_rows, n_frames = sras.image_shape(angle_idx) + data = sras.data[angle_idx] + chunk_rows, n_workers = _plan_chunks(n_rows, n_frames, sras.samples_per_frame) + img = np.empty((n_rows, n_frames), dtype=np.float32) + + def chunk(r0: int, r1: int): + img[r0:r1] = data[r0:r1, ch_idx, :, :].astype(np.float32).mean(axis=-1) + + _map_row_chunks(n_rows, chunk_rows, n_workers, chunk) + return img + + +def dc_image_mv(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray: + """DC image for (angle, channel) in mV, preferring a stored v5/v7 cache.""" + cached = sras.cached_dc_mv(angle_idx, ch_idx) + if cached is not None: + return cached + return adc_to_mv(compute_dc_image(sras, angle_idx, ch_idx), *sras.cal(ch_idx)) + + +def compute_rf_image(sras: SrasFile, angle_idx: int, + dc_threshold_mv: float | None, + apply_bg_sub: bool = True, + n_fft: int | None = None, + dc4_mv: np.ndarray | None = None) -> np.ndarray: + """FFT of each CH1 waveform; pixel = peak frequency in MHz. + + Pixels where CH4_dc < dc_threshold_mv are set to 0 — and the FFT is + never run for them, since that's the expensive part. The masked-out + CH1 samples are also never *read*: the boolean mask is applied to the + raw memmap slice before any dtype conversion, so numpy only pages in + the bytes for pixels that pass the threshold (an untouched memmap page + is never read from disk). + + *dc_threshold_mv* of None means "mask nothing", and skips reading and + averaging CH4 entirely — worth a third of the I/O when caching a whole + file, where the mask is applied later at display time. + + If the DC4 image for this angle is already known (e.g. from the + DC-channel precompute cache), pass it as *dc4_mv* (mV, shape + (n_rows, n_frames)) to reuse it instead of re-reading CH4 here. + + Fast path: if the file has a precomputed peak-frequency image for this + angle (v5 PREC or v7 CACH), zero-padding is off, and the bg-sub flag + matches, the stored image is used directly — no FFT is run. + """ + n_rows, n_frames = sras.image_shape(angle_idx) + data = sras.data[angle_idx] + + # ---- Fast path: precomputed image (v5 PREC or v7 CACH) ---------------- + cached_freq = sras.precomputed_freq_mhz[angle_idx] + if (cached_freq is not None + and n_fft is None # no custom zero-padding + and sras.precomputed_bg_sub == (apply_bg_sub and sras.background is not None)): + freq_img = cached_freq.copy() + if dc_threshold_mv is not None: + # DC4 mask, in priority order: already-cached DC block, caller- + # supplied image, or a fresh (cheap — no FFT) recompute. + dc4_img = sras.cached_dc_mv(angle_idx, CH4_IDX) + if dc4_img is None: + dc4_img = dc4_mv if dc4_mv is not None else adc_to_mv( + compute_dc_image(sras, angle_idx, CH4_IDX), *sras.cal(CH4_IDX)) + freq_img[dc4_img < dc_threshold_mv] = 0.0 + return freq_img + + # ---- Chunked FFT path -------------------------------------------------- + freq_axis = sras.freq_axis_mhz(n_fft) + img = np.zeros((n_rows, n_frames), dtype=np.float32) + n_fft_bins = n_fft if n_fft is not None else sras.samples_per_frame + chunk_rows, n_workers = _plan_chunks( + n_rows, n_frames, max(sras.samples_per_frame, n_fft_bins), + live_multiplier=_FFT_LIVE_MULTIPLIER) + background = sras.background if (apply_bg_sub and sras.background is not None) else None + cal4 = sras.cal(CH4_IDX) + + def chunk(r0: int, r1: int): + if dc_threshold_mv is None: + valid = None + else: + if dc4_mv is not None: + dc4_chunk = dc4_mv[r0:r1] + else: + dc4_chunk = adc_to_mv( + data[r0:r1, CH4_IDX, :, :].astype(np.float32).mean(axis=-1), *cal4) + valid = dc4_chunk >= dc_threshold_mv # True = run the FFT + if not valid.any(): + return + + # Index the raw memmap slice with the boolean mask *before* + # converting dtype — this is a lazy view until touched, so only the + # selected elements are actually read from disk; masked-out pixels' + # pages are never paged in at all. + raw = data[r0:r1, CH1_IDX, :, :] + waves = (raw[valid] if valid is not None else raw).astype(np.float32) + + if background is not None: + waves -= background # background is 1-D (spf,) + + # Parallelism comes from the outer chunk loop, so keep the inner + # transform single-threaded to avoid oversubscribing the machine. + spectrum = _do_rfft(waves, n=n_fft, axis=-1, workers=1) + del waves + # |z|^2 without np.abs()'s extra full-size temporary. + power = spectrum.real ** 2 + power += spectrum.imag ** 2 + del spectrum + # [..., 0] not [:, 0]: the unmasked path keeps the (rows, frames, + # bins) shape, where [:, 0] would blank a whole frame. + power[..., 0] = 0.0 # suppress DC bin + peak_bins = np.argmax(power, axis=-1) + del power + + if valid is not None: + img[r0:r1][valid] = freq_axis[peak_bins] + else: + img[r0:r1] = freq_axis[peak_bins] + + _map_row_chunks(n_rows, chunk_rows, n_workers, chunk) + return img + + +# --------------------------------------------------------------------------- +# Batch cache (Convert menu) — process-pool entry point +# --------------------------------------------------------------------------- + +def cache_file(path: str, mode: str, apply_bg_sub: bool, + fft_backend: str = "numpy", max_workers: int = 0) -> str: + """Compute and store DC or FFT images for every angle of one file, + converting v6 → v7 in place. Returns "" on success or an error message. + + Module-level and picklable so it can run in a ProcessPoolExecutor. The + FFT backend and worker cap are passed explicitly because module globals + do not survive a spawn. + """ + global _MAX_WORKERS + try: + set_fft_backend(fft_backend) + if max_workers: + _MAX_WORKERS = max_workers + + sras = SrasFile(path) + if sras.version not in (6, 7): + return (f"unsupported version {sras.version} — only v6/v7 files " + "can be batch-cached") + + n = sras.n_angles + if mode == "dc": + dc3 = [adc_to_mv(compute_dc_image(sras, a, CH3_IDX), *sras.cal(CH3_IDX)) + for a in range(n)] + dc4 = [adc_to_mv(compute_dc_image(sras, a, CH4_IDX), *sras.cal(CH4_IDX)) + for a in range(n)] + sras.write_v7_cache(new_dc3_mv=dc3, new_dc4_mv=dc4) + else: + effective_bg = apply_bg_sub and sras.background is not None + # dc_threshold_mv=None: store unmasked images and mask at display + # time (same convention as v5's PREC block). Skipping the mask + # also skips reading CH4 entirely. + freq = [compute_rf_image(sras, a, dc_threshold_mv=None, + apply_bg_sub=effective_bg) + for a in range(n)] + sras.write_v7_cache(new_freq_mhz=freq, new_bg_sub=effective_bg) + return "" + except Exception as exc: + return str(exc) + + +# --------------------------------------------------------------------------- +# Angle alignment (Fusion menu) +# +# Puts every angle's images onto one shared, zero-padded pixel grid using a +# rigid transform only (rotation + translation, never scale). Rotation for +# angle `a` is the *known* scan-angle delta relative to a reference angle — +# never searched. Only the residual translation is found, via FFT phase +# correlation of each angle's binarized CH4 ("dc-mask") image. +# +# Rotation is done in physical mm space rather than on raw pixel indices: +# the x-pixel pitch (SrasFile.pixel_x_mm) is file-wide constant but the +# y-pixel pitch (row spacing) can differ from it, and for v6 files can even +# vary per angle. Rotating the raw index grid directly would implicitly +# assume square pixels and shear a non-square-pixel image — an unwanted +# effective anisotropic scale. Instead each angle gets one affine that maps +# shared-canvas pixel index -> mm -> undo rotation/shift -> that angle's own +# local mm -> that angle's own raw pixel index, matching the output->input +# convention scipy.ndimage.affine_transform expects. +# --------------------------------------------------------------------------- + +@dataclass +class AngleTransform: + rotation_deg: float + shift_mm: tuple[float, float] # (dx_mm, dy_mm) found by phase correlation + matrix: np.ndarray # (2,2): canvas (row,col) -> this angle's raw (row,col) + offset: np.ndarray # (2,) + + +@dataclass +class AlignmentResult: + ref_angle_idx: int + dc_threshold_mv: float + canvas_shape: tuple[int, int] # (n_rows, n_cols) + canvas_dx_mm: float + canvas_dy_mm: float + canvas_origin_mm: tuple[float, float] # mm at canvas pixel index (0, 0) + per_angle: dict[int, AngleTransform] + + +def _pixel_pitch_mm(sras: SrasFile, angle_idx: int) -> tuple[float, float]: + """(dx, dy) mm/pixel for one angle: dx is the file-wide constant + pixel_x_mm; dy is this angle's own row spacing (assumed uniform, the same + assumption _redraw_image makes when it builds the display extent).""" + y = sras.y_positions_mm(angle_idx) + return sras.pixel_x_mm, (float(y[1] - y[0]) if len(y) > 1 else 1.0) + + +def _bbox_center_mm(sras: SrasFile, angle_idx: int) -> tuple[float, float]: + x = sras.x_axis_mm(angle_idx) + y = sras.y_positions_mm(angle_idx) + return float((x[0] + x[-1]) / 2.0), float((y[0] + y[-1]) / 2.0) + + +def _bbox_corners_mm(sras: SrasFile, angle_idx: int) -> np.ndarray: + """4 corners (x, y) of this angle's raw mm bounding box, shape (4, 2).""" + x = sras.x_axis_mm(angle_idx) + y = sras.y_positions_mm(angle_idx) + return np.array([[xx, yy] for xx in (x[0], x[-1]) for yy in (y[0], y[-1])]) + + +def _rotation_matrix(theta_deg: float) -> np.ndarray: + t = np.radians(theta_deg) + c, s = np.cos(t), np.sin(t) + return np.array([[c, -s], [s, c]]) # CCW rotation acting on (x, y) + + +def _theta_deg(sras: SrasFile, angle_idx: int, ref_idx: int) -> float: + return float(sras.angles_deg[angle_idx] - sras.angles_deg[ref_idx]) + + +def _corners_in_ref_frame(sras: SrasFile, angle_idx: int, ref_idx: int, + shift_mm=(0.0, 0.0)) -> np.ndarray: + """Angle *angle_idx*'s bbox corners, rotated about its own centroid into + the reference frame and translated by *shift_mm*. Shape (4, 2).""" + R = _rotation_matrix(_theta_deg(sras, angle_idx, ref_idx)) + c_a = np.array(_bbox_center_mm(sras, angle_idx)) + c_ref = np.array(_bbox_center_mm(sras, ref_idx)) + shift = np.asarray(shift_mm, dtype=np.float64) + return np.array([R @ (corner - c_a) + c_ref + shift + for corner in _bbox_corners_mm(sras, angle_idx)]) + + +def _build_affine_canvas_to_raw(sras: SrasFile, angle_idx: int, ref_idx: int, + shift_mm: tuple[float, float], + canvas_dx: float, canvas_dy: float, + canvas_origin_mm: tuple[float, float] + ) -> tuple[np.ndarray, np.ndarray]: + """matrix, offset s.t. raw_index = matrix @ [row_out, col_out] + offset, + matching scipy.ndimage.affine_transform's output->input convention. + + Pipeline (all mm unless noted): + [X;Y] = A_out @ [row_out;col_out] + b_out # canvas idx -> ref-frame mm + [lx;ly] = R(theta)^T @ ([X;Y]-c_ref-shift) + c_a # undo rotation+shift -> angle a's local mm + [row;col] = D @ ([lx;ly] - [x_start_a; y0_a]) # local mm -> angle a's raw idx + + where theta = angles_deg[angle_idx] - angles_deg[ref_idx], c_ref/c_a are + each angle's own raw-bbox mm centroid (the rotation pivot — this keeps + rotated content centered, minimizing required canvas padding), and + A_out/D are the index<->mm scaling matrices for the canvas pitch and + this angle's own native pitch respectively. + """ + Rinv = _rotation_matrix(_theta_deg(sras, angle_idx, ref_idx)).T + cx_a, cy_a = _bbox_center_mm(sras, angle_idx) + cx_ref, cy_ref = _bbox_center_mm(sras, ref_idx) + dx_a, dy_a = _pixel_pitch_mm(sras, angle_idx) + x0_a = float(sras.x_start_mm[angle_idx]) + y0_a = float(sras.y_positions_mm(angle_idx)[0]) + + A_out = np.array([[0.0, canvas_dx], [canvas_dy, 0.0]]) # [row,col] -> [X,Y] + b_out = np.array(canvas_origin_mm, dtype=np.float64) + D = np.array([[0.0, 1.0 / dy_a], [1.0 / dx_a, 0.0]]) # [x,y] -> [row,col] + shift = np.array(shift_mm, dtype=np.float64) + c_ref_v = np.array([cx_ref, cy_ref]) + c_a_v = np.array([cx_a, cy_a]) + origin_a = np.array([x0_a, y0_a]) + + matrix = D @ Rinv @ A_out + offset = D @ Rinv @ (b_out - c_ref_v - shift) + D @ (c_a_v - origin_a) + return matrix, offset + + +def apply_alignment(result: AlignmentResult, angle_idx: int, img: np.ndarray, + order: int = 0) -> np.ndarray: + """Resample any already-computed 2D image for `angle_idx` (same shape as + that angle's raw (n_rows, n_frames)) onto the shared alignment canvas. + order=0 (nearest) avoids blending real data with zero-padding or with + masked-out (0-valued) CH1/velocity pixels at mask edges. Channel- + agnostic: the same per-angle transform (found from the CH4 mask) works + for any channel's image of that angle.""" + t = result.per_angle[angle_idx] + return scipy_ndimage.affine_transform( + img.astype(np.float32, copy=False), t.matrix, offset=t.offset, + output_shape=result.canvas_shape, order=order, + mode="constant", cval=0.0) + + +def _block_mean_downsample(img: np.ndarray, factor: int) -> np.ndarray: + if factor <= 1: + return img + h, w = img.shape + h2, w2 = (h // factor) * factor, (w // factor) * factor + trimmed = img[:h2, :w2] + return trimmed.reshape(h2 // factor, factor, w2 // factor, factor).mean(axis=(1, 3)) + + +def _phase_correlate_shift(ref_img: np.ndarray, mov_img: np.ndarray) -> tuple[int, int]: + """FFT normalized cross-power-spectrum phase correlation. Returns the + integer (dr, dc) pixel shift of mov_img relative to ref_img; both must + be the same shape. Risk: if the true shift is near +/- half the array + size, wraparound can bias the peak — mitigated by the generous + margin_frac padding in _working_canvas_for_pair, which keeps the true + residual shift small relative to the correlation canvas.""" + F1 = scipy_fft.fft2(ref_img.astype(np.float64), workers=-1) + F2 = scipy_fft.fft2(mov_img.astype(np.float64), workers=-1) + R = F1 * np.conj(F2) + R /= np.maximum(np.abs(R), 1e-12) + corr = scipy_fft.ifft2(R, workers=-1).real + dr, dc = np.unravel_index(np.argmax(corr), corr.shape) + h, w = corr.shape + if dr > h // 2: + dr -= h + if dc > w // 2: + dc -= w + return int(dr), int(dc) + + +def _working_canvas_for_pair(sras: SrasFile, ref_idx: int, a_idx: int, + dx: float, dy: float, margin_frac: float = 0.3 + ) -> tuple[tuple[float, float], tuple[int, int]]: + """Union of the reference's own raw bbox and angle a's raw bbox rotated + (about its own center) into the ref frame with zero shift, padded by + margin_frac on each side — sized generously so the true phase- + correlation shift lands well inside the canvas (see + _phase_correlate_shift's wraparound note).""" + pts = np.vstack([_bbox_corners_mm(sras, ref_idx), + _corners_in_ref_frame(sras, a_idx, ref_idx)]) + x_min, y_min = pts.min(axis=0) + x_max, y_max = pts.max(axis=0) + pad_x, pad_y = (x_max - x_min) * margin_frac, (y_max - y_min) * margin_frac + x_min, x_max = x_min - pad_x, x_max + pad_x + y_min, y_max = y_min - pad_y, y_max + pad_y + n_cols = int(np.ceil((x_max - x_min) / dx)) + 1 + n_rows = int(np.ceil((y_max - y_min) / abs(dy))) + 1 + origin = (x_min, y_min if dy > 0 else y_max) + return origin, (n_rows, n_cols) + + +def _parallel_map(fn, items, n_workers: int) -> list: + """fn over items, in order, threaded when it pays.""" + items = list(items) + if n_workers <= 1 or len(items) <= 1: + return [fn(x) for x in items] + with ThreadPoolExecutor(max_workers=min(n_workers, len(items))) as pool: + return list(pool.map(fn, items)) + + +def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int, + dc_threshold_mv: float, + progress_cb=None) -> AlignmentResult: + """Top-level alignment driver. Runs on a background thread (see + AngleAlignmentWorker) — deliberately recomputes CH4 DC images from + scratch rather than reading the GUI-thread _dc_cache dict, since + background-thread workers must not touch GUI-thread-owned caches.""" + n = sras.n_angles # already the *complete*-angle count for aborted v6 scans + dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx) + n_workers = max(1, min(_MAX_WORKERS, n)) + + # progress_cb fires from pool threads, so the counter behind it must be + # atomic. list.append is, and len() of a list is a consistent read. + _ticks: list[int] = [] + + def tick(base: int, span: int): + if progress_cb: + _ticks.append(1) + progress_cb(base + int(min(len(_ticks), n) / n * span)) + + # ---- Step 1: binarized CH4 mask per angle, native per-angle grid ----- + def mask_for(a: int) -> np.ndarray: + dc4 = adc_to_mv(compute_dc_image(sras, a, CH4_IDX), *sras.cal(CH4_IDX)) + m = (dc4 >= dc_threshold_mv).astype(np.float32) + tick(0, 25) + return m + + masks = dict(enumerate(_parallel_map(mask_for, range(n), n_workers))) + + # ---- Step 2: coarse correlation stage (downsample first, then rotate) + # Downsampling before affine_transform (not after) is what keeps this + # tractable for a v6 scan with thousands of rows/frames per angle. + max_dim = max(max(m.shape) for m in masks.values()) + factor = max(1, int(np.ceil(max_dim / 1024))) + dx_c, dy_c = dx_ref * factor, dy_ref * factor + masks_small = {a: _block_mean_downsample(m, factor) for a, m in masks.items()} + + _ticks.clear() + + def shift_for(a: int) -> tuple[float, float]: + if a == ref_angle_idx: + tick(25, 50) + return (0.0, 0.0) + work_origin, work_shape = _working_canvas_for_pair( + sras, ref_angle_idx, a, dx_c, dy_c, margin_frac=0.3) + + # Coarse canvas->raw affine at native pitch, then rescale by /factor + # so it maps coarse-canvas idx -> coarse (downsampled) raw idx — + # exact for block-mean downsampling (up to the trimmed remainder). + m_a, o_a = _build_affine_canvas_to_raw( + sras, a, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin) + m_ref, o_ref = _build_affine_canvas_to_raw( + sras, ref_angle_idx, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin) + rotated_a = scipy_ndimage.affine_transform( + masks_small[a], m_a / factor, offset=o_a / factor, + output_shape=work_shape, order=0, mode="constant", cval=0.0) + embedded_ref = scipy_ndimage.affine_transform( + masks_small[ref_angle_idx], m_ref / factor, offset=o_ref / factor, + output_shape=work_shape, order=0, mode="constant", cval=0.0) + + dr, dc = _phase_correlate_shift(embedded_ref, rotated_a) + tick(25, 50) + return (dc * dx_c, dr * dy_c) + + shifts_mm = dict(enumerate(_parallel_map(shift_for, range(n), n_workers))) + + # ---- Step 3: union bounding box over all angles (rotation+shift applied) + corners = np.vstack([_corners_in_ref_frame(sras, a, ref_angle_idx, shifts_mm[a]) + for a in range(n)]) + x_min, y_min = corners.min(axis=0) + x_max, y_max = corners.max(axis=0) + n_cols = int(np.ceil((x_max - x_min) / dx_ref)) + 1 + n_rows = int(np.ceil((y_max - y_min) / abs(dy_ref))) + 1 + canvas_origin_mm = (float(x_min), float(y_min if dy_ref > 0 else y_max)) + + # ---- Step 4: final per-angle full-resolution affine (canvas -> raw idx) + per_angle: dict[int, AngleTransform] = {} + for a in range(n): + matrix, offset = _build_affine_canvas_to_raw( + sras, a, ref_angle_idx, shifts_mm[a], dx_ref, dy_ref, canvas_origin_mm) + per_angle[a] = AngleTransform( + _theta_deg(sras, a, ref_angle_idx), shifts_mm[a], matrix, offset) + if progress_cb: + progress_cb(75 + int((a + 1) / n * 25)) + + return AlignmentResult(ref_angle_idx, dc_threshold_mv, (n_rows, n_cols), + dx_ref, dy_ref, canvas_origin_mm, per_angle) + + +# Back-compat alias for the pre-split private name (used by tooling). +_compute_angle_alignment = compute_angle_alignment diff --git a/sras_format.py b/sras_format.py new file mode 100644 index 0000000..0b2715d --- /dev/null +++ b/sras_format.py @@ -0,0 +1,593 @@ +#!/usr/bin/env python3 +"""SRAS binary scan file format — parsing and writing. + +Reads v2–v7 .sras files. Depends only on numpy + struct, so compute workers +(including multiprocessing children) can import it without pulling in Qt or +matplotlib. See scan_format.md for the full v6/v7 spec. + +Channel semantics (fixed by sc3_aui_app.py acquisition settings): + CH1 — RF Acoustic Packet (AC-coupled, 100 mV/div): FFT → peak frequency + CH3 — Bias A (DC-coupled, 50 mV/div): waveform mean + CH4 — Bias B (DC-coupled, 50 mV/div): waveform mean + +Frame-count correction: the scanner writes the *configured* frame count in the +header before acquisition, but the scope may acquire fewer frames. The actual +count is computed from the file size and used for the reshape so channels are +correctly aligned. + +Scan geometry: v6/v7 files scan a different bounding box per angle (x_start, +x_delta, n_frames, n_rows all vary by angle), so geometry is exposed per-angle +via SrasFile.n_rows / n_frames / x_start_mm arrays and the x_axis_mm() / +y_positions_mm() methods. v2–v5 files have uniform geometry across angles, so +those arrays simply repeat the same value n_angles times. + +v7 files are v6 files with an optional trailing cache section holding +precomputed per-angle DC and/or FFT images, so display never has to recompute +them after the "Convert" menu's batch actions have stored them once. +""" + +import os +import re +import struct +from pathlib import Path + +import numpy as np + +# --------------------------------------------------------------------------- +# Format constants +# --------------------------------------------------------------------------- + +# v2–v5: fixed header, uniform geometry across angles (43 bytes) +HDR_FMT = ">4sBHHffffIIdBB" +HDR_SIZE = struct.calcsize(HDR_FMT) + +# v6/v7: fixed header, per-angle geometry in a separate table (49 bytes) +HDR_FMT_V6 = ">4sBHfffffffIdBB" +HDR_SIZE_V6 = struct.calcsize(HDR_FMT_V6) + +# v6/v7: per-angle geometry record (x_start, x_delta, n_frames, n_rows) +GEO_FMT_V6 = ">ffIH" +GEO_SIZE_V6 = struct.calcsize(GEO_FMT_V6) + +# v5 precomputed-image tail +PREC_MAGIC = b"PREC" +PREC_FLAG_BG_SUB = 0x01 + +# v7 cache tail. v7 is byte-identical to v6 (the version byte is the only +# header difference) plus this optional trailing section. Sub-block sizes are +# per-angle (n_rows[a] * n_frames[a]), taken from the Per-Angle Geometry Table +# already parsed for v6 — no new geometry fields are needed. +CACH_MAGIC = b"CACH" +CACH_HDR_FMT = ">4sBB" # magic, cach_version, block_flags +CACH_HDR_SIZE = struct.calcsize(CACH_HDR_FMT) +CACH_VERSION = 1 +CACH_FLAG_DC = 0x01 +CACH_FLAG_FFT = 0x02 + +SDCB_MAGIC = b"SDCB" +SDCB_HDR_FMT = ">4sBH" # magic, reserved, n_stored +SDCB_HDR_SIZE = struct.calcsize(SDCB_HDR_FMT) + +SFFT_MAGIC = b"SFFT" +SFFT_HDR_FMT = ">4sBH" # magic, flags, n_stored +SFFT_HDR_SIZE = struct.calcsize(SFFT_HDR_FMT) +SFFT_FLAG_BG_SUB = 0x01 + +# Fixed channel indices into the .sras data array (CH1=RF, CH3/CH4=Bias DC) +CH1_IDX, CH3_IDX, CH4_IDX = 0, 1, 2 +CH_NAMES = ["CH1", "CH3", "CH4", "VEL"] + +# Fallback scope calibration used only when reading v2 files without embedded +# preambles. v3+ files carry the WFMOutpre string so these are not used. +# 50 mV/div, 8 div full-scale, int8 ADC, position = -2.72 div +# ymult = 50 mV × 8 / 256 = 1.5625 mV/count +# yoff = position × (256/8) = -2.72 × 32 = -87.04 (ADC count for 0 V) +_FALLBACK_YMULT_MV = 1.5625 # mV per ADC count +_FALLBACK_YOFF_ADC = -87.04 # ADC count that represents 0 V + + +# --------------------------------------------------------------------------- +# Calibration +# --------------------------------------------------------------------------- + +def _parse_preamble(preamble: str) -> dict[str, float]: + """Extract YMULT, YOFF, YZERO from a Tektronix WFMOutpre string. + + Returns a dict with float values for whichever keys are present. + YMULT is left in V/count as the scope reports it. + """ + result = {} + for key in ("YMULT", "YOFF", "YZERO"): + m = re.search(rf'\b{key}\s+([-+]?\d*\.?\d+(?:[Ee][+-]?\d+)?)', preamble) + if m: + result[key] = float(m.group(1)) + return result + + +def mv_to_adc(mv: float, ymult_mv: float = _FALLBACK_YMULT_MV, + yoff_adc: float = _FALLBACK_YOFF_ADC, + yzero_mv: float = 0.0) -> float: + return (mv - yzero_mv) / ymult_mv + yoff_adc + + +def adc_to_mv(adc, ymult_mv: float = _FALLBACK_YMULT_MV, + yoff_adc: float = _FALLBACK_YOFF_ADC, + yzero_mv: float = 0.0): + return (adc - yoff_adc) * ymult_mv + yzero_mv + + +# --------------------------------------------------------------------------- +# Binary read helpers +# --------------------------------------------------------------------------- + +def _read_struct(f, fmt: str) -> tuple: + return struct.unpack(fmt, f.read(struct.calcsize(fmt))) + + +def _read_preambles(f, n_ch: int) -> list[str]: + """n_ch length-prefixed UTF-8 WFMOutpre strings.""" + out = [] + for _ in range(n_ch): + (length,) = _read_struct(f, ">H") + out.append(f.read(length).decode("utf-8")) + return out + + +def _read_background(f) -> np.ndarray: + """uint32 sample count followed by that many int8 samples.""" + (n_bg,) = _read_struct(f, ">I") + return np.frombuffer(f.read(n_bg), dtype=np.int8).astype(np.float32) + + +def _read_f32_image(f, shape: tuple[int, int]) -> np.ndarray: + """One big-endian float32 image, converted to native float32. + + The conversion matters: np.frombuffer hands back a read-only big-endian + view, and these arrays flow straight into the display caches and every + downstream arithmetic op. + """ + n_bytes = shape[0] * shape[1] * 4 + return np.frombuffer(f.read(n_bytes), dtype=">f4").reshape(shape).astype(np.float32) + + +# --------------------------------------------------------------------------- +# File parser +# --------------------------------------------------------------------------- + +class SrasFile: + """Parsed in-memory representation of a v2–v7 .sras file. + + Scan geometry (rows, frames, x_start) is exposed per-angle via the + ``n_rows`` / ``n_frames`` / ``x_start_mm`` arrays and the ``x_axis_mm()`` + / ``y_positions_mm()`` methods, since v6/v7 files scan a different + bounding box per angle. v2–v5 files have uniform geometry, so these + arrays just repeat the same value ``n_angles`` times. Waveform data is + likewise exposed as ``data[angle_idx]``, an array of shape + ``(n_rows[a], n_channels, n_frames[a], samples_per_frame)``. + + Precomputed images (v5's PREC tail or v7's CACH tail) are exposed as + ``precomputed_dc3_mv`` / ``precomputed_dc4_mv`` / ``precomputed_freq_mhz``, + always as ragged per-angle lists (``list[np.ndarray | None]``, one entry + per angle, ``None`` where that angle was never stored) regardless of + source version. + """ + + def __init__(self, path: str): + self.path = Path(path) + self._parse() + + def _parse(self): + with open(self.path, "rb") as f: + magic = f.read(4) + if magic != b"SRAS": + raise ValueError(f"Bad magic bytes: {magic!r}") + (version,) = struct.unpack(">B", f.read(1)) + + self.version = version + if version in (2, 3, 4, 5): + self._parse_legacy() + elif version in (6, 7): + self._parse_v6() + else: + raise ValueError(f"Unsupported version: {version}") + + # ------------------------------------------------------------------ + # Calibration + # ------------------------------------------------------------------ + + def _set_calibration(self, preambles: list[str] | None, n_ch: int): + """Populate the per-channel ymult/yoff/yzero lists from preamble + strings, falling back to the hardcoded scope constants for v2 files + that carry no preambles.""" + if preambles is None: + self.preambles = None + self.ch_ymult_mv = [_FALLBACK_YMULT_MV] * n_ch + self.ch_yoff_adc = [_FALLBACK_YOFF_ADC] * n_ch + self.ch_yzero_mv = [0.0] * n_ch + return + + self.preambles = preambles + self.ch_ymult_mv, self.ch_yoff_adc, self.ch_yzero_mv = [], [], [] + for p in preambles: + cal = _parse_preamble(p) + # The scope reports YMULT and YZERO in volts; store both as mV. + self.ch_ymult_mv.append(cal.get("YMULT", _FALLBACK_YMULT_MV / 1000) * 1000) + self.ch_yoff_adc.append(cal.get("YOFF", _FALLBACK_YOFF_ADC)) + self.ch_yzero_mv.append(cal.get("YZERO", 0.0) * 1000) + + def cal(self, ch_idx: int) -> tuple[float, float, float]: + """(ymult_mv, yoff_adc, yzero_mv) for one channel — splat straight + into adc_to_mv / mv_to_adc.""" + return (self.ch_ymult_mv[ch_idx], self.ch_yoff_adc[ch_idx], + self.ch_yzero_mv[ch_idx]) + + def _init_precomputed(self, n_angles: int): + self.precomputed_freq_mhz: list[np.ndarray | None] = [None] * n_angles + self.precomputed_dc4_mv: list[np.ndarray | None] = [None] * n_angles + self.precomputed_dc3_mv: list[np.ndarray | None] = [None] * n_angles + self.precomputed_bg_sub: bool = False + + def cached_dc_mv(self, angle_idx: int, ch_idx: int) -> np.ndarray | None: + """A stored DC image (already in mV) for (angle, channel), or None.""" + store = self.precomputed_dc3_mv if ch_idx == CH3_IDX else self.precomputed_dc4_mv + return store[angle_idx] if angle_idx < len(store) else None + + def image_shape(self, angle_idx: int) -> tuple[int, int]: + return int(self.n_rows[angle_idx]), int(self.n_frames[angle_idx]) + + # ------------------------------------------------------------------ + # v2–v5 parsing (uniform geometry, flat waveform block) + # ------------------------------------------------------------------ + + def _parse_legacy(self): + with open(self.path, "rb") as f: + (magic, ver, n_angles, n_rows, x_start, x_delta, vel, freq, + n_frames_hdr, spf, sr, bps, n_ch) = _read_struct(f, HDR_FMT) + + self.n_angles = n_angles + self.velocity_mm_s = float(vel) + self.laser_freq_hz = float(freq) + self.n_frames_header = n_frames_hdr # configured count (may be wrong) + self.samples_per_frame = spf + self.sample_rate_hz = float(sr) + self.bytes_per_sample = bps + self.n_channels = n_ch + + self._init_precomputed(n_angles) + self.scan_aborted = False + self.n_angles_declared = n_angles + + angles = np.frombuffer(f.read(n_angles * 4), dtype=">f4").astype(np.float32) + y_pos = np.frombuffer(f.read(n_rows * 4), dtype=">f4").astype(np.float32) + + self._set_calibration(_read_preambles(f, n_ch) if ver >= 3 else None, n_ch) + self.background = _read_background(f) if ver >= 4 else None + + # Record where raw waveform data begins; np.memmap maps from here. + data_offset = f.tell() + + # ---- Determine actual frame count from file size --------------- + # For v4 and earlier the header n_frames may be the *configured* + # count before acquisition; the actual count is derived from the + # bytes on disk. For v5 files a PREC tail follows the waveform + # data, so we must not include those extra bytes in the frame count. + file_size = self.path.stat().st_size + samples_per_row_per_ch = n_ch * spf + available_bytes = file_size - data_offset + + if ver == 5: + actual_n_frames = n_frames_hdr + remainder = 0 + else: + total_samples = available_bytes // bps + per_frame = n_angles * n_rows * samples_per_row_per_ch + actual_n_frames = total_samples // per_frame + remainder = total_samples % per_frame + + self.frame_count_mismatch = (actual_n_frames != n_frames_hdr) + self.n_frames_remainder = remainder + + # ---- Memory-map the waveform data (zero RAM cost) -------------- + # Instead of f.read() → astype() (which peaks at 2× file size), + # memmap lets the OS page only the bytes that are actually touched. + data5d = np.memmap( + str(self.path), + dtype=np.int8 if bps == 1 else ">i2", + mode="r", + offset=data_offset, + shape=(n_angles, n_rows, n_ch, actual_n_frames, spf), + ) + # Expose as a list of per-angle views so downstream code shares one + # indexing convention with v6: sras.data[a][row, ch, frame, sample] + self.data = [data5d[a] for a in range(n_angles)] + + # Uniform per-angle geometry, repeated so callers don't need to + # special-case legacy vs. v6 files. + self.n_rows = np.full(n_angles, n_rows, dtype=np.int64) + self.n_frames = np.full(n_angles, actual_n_frames, dtype=np.int64) + self.x_start_mm = np.full(n_angles, float(x_start), dtype=np.float64) + self.x_delta_mm = float(x_delta) # reference only; kept for re-encode + self._y_pos_per_angle = [y_pos] * n_angles + self.angles_deg = angles + self._data_offset = data_offset + + if ver >= 5: + waveform_bytes = actual_n_frames * n_angles * n_rows * n_ch * spf * bps + prec_offset = data_offset + waveform_bytes + if file_size > prec_offset: + self._parse_prec_section(prec_offset) + + def _parse_prec_section(self, offset: int): + """Parse the v5 PREC tail that holds precomputed images. + + Stored per-angle as (freq, dc4, dc3), each a full-image float32 + block prefixed by its uint16 angle index. + """ + with open(self.path, "rb") as f: + f.seek(offset) + header_raw = f.read(6) # magic(4) + fmt_ver(1) + flags(1) + if len(header_raw) < 6 or header_raw[:4] != PREC_MAGIC: + return + self.precomputed_bg_sub = bool(header_raw[5] & PREC_FLAG_BG_SUB) + + (n_stored,) = _read_struct(f, ">H") + for _ in range(n_stored): + (aidx,) = _read_struct(f, ">H") + if aidx >= self.n_angles: + break + shape = self.image_shape(aidx) + self.precomputed_freq_mhz[aidx] = _read_f32_image(f, shape) + self.precomputed_dc4_mv[aidx] = _read_f32_image(f, shape) + self.precomputed_dc3_mv[aidx] = _read_f32_image(f, shape) + + # ------------------------------------------------------------------ + # v6/v7 parsing (per-angle geometry, ragged waveform blocks) + # ------------------------------------------------------------------ + + def _parse_v6(self): + with open(self.path, "rb") as f: + (magic, ver, n_angles, x_start_nom, y_start_nom, x_delta_nom, + y_delta_nom, row_spacing, vel, freq, spf, sr, bps, + n_ch) = _read_struct(f, HDR_FMT_V6) + + n_angles_declared = n_angles + + self.velocity_mm_s = float(vel) + self.laser_freq_hz = float(freq) + self.samples_per_frame = spf + self.sample_rate_hz = float(sr) + self.bytes_per_sample = bps + self.n_channels = n_ch + + # Reference-only fields: the ROI as entered before per-angle + # bounding-box expansion. Actual per-angle geometry used for + # rendering comes from the Per-Angle Geometry Table below. + self.x_start_nominal_mm = float(x_start_nom) + self.y_start_nominal_mm = float(y_start_nom) + self.x_delta_nominal_mm = float(x_delta_nom) + self.y_delta_nominal_mm = float(y_delta_nom) + self.row_spacing_mm = float(row_spacing) + + self.n_frames_header = None + self.frame_count_mismatch = False + self.n_frames_remainder = 0 + + angles = np.frombuffer(f.read(n_angles * 4), dtype=">f4").astype(np.float32) + + x_start = np.empty(n_angles, dtype=np.float64) + n_frames = np.empty(n_angles, dtype=np.int64) + n_rows = np.empty(n_angles, dtype=np.int64) + for a in range(n_angles): + xs, xd, nf, nr = _read_struct(f, GEO_FMT_V6) + x_start[a], n_frames[a], n_rows[a] = xs, nf, nr + + y_pos_per_angle = [ + np.frombuffer(f.read(int(n_rows[a]) * 4), dtype=">f4").astype(np.float32) + for a in range(n_angles) + ] + + self._set_calibration(_read_preambles(f, n_ch), n_ch) + self.background = _read_background(f) + + data_offset = f.tell() + + self._data_offset = data_offset + + # ---- Memory-map each angle's ragged waveform block ------------- + # v6 gives each angle its own row/frame count, so waveform data is + # no longer one uniform (n_angles, n_rows, ...) block — each angle's + # block sits at a different offset with its own shape. An aborted + # scan truncates the file mid-angle; per the format spec we keep + # whatever complete angles are present rather than refusing to open + # the file. + file_size = self.path.stat().st_size + waveform_dtype = np.int8 if bps == 1 else ">i2" + + data = [] + offset = data_offset + for a in range(n_angles): + nr, nf = int(n_rows[a]), int(n_frames[a]) + nbytes = nr * n_ch * nf * spf * bps + if offset + nbytes > file_size: + break + data.append(np.memmap( + str(self.path), dtype=waveform_dtype, mode="r", + offset=offset, shape=(nr, n_ch, nf, spf), + )) + offset += nbytes + + n_complete = len(data) + if n_complete == 0: + raise ValueError( + "v6 file has no complete angle blocks — scan was aborted " + "before the first angle finished.") + + self.data = data + self.n_angles = n_complete + self.n_angles_declared = n_angles_declared + self.scan_aborted = n_complete < n_angles_declared + self.angles_deg = angles[:n_complete] + self.x_start_mm = x_start[:n_complete] + self.n_frames = n_frames[:n_complete] + self.n_rows = n_rows[:n_complete] + self._y_pos_per_angle = y_pos_per_angle[:n_complete] + + self._init_precomputed(n_complete) + if self.version == 7 and offset < file_size: + self._parse_cach_section(offset) + + # ------------------------------------------------------------------ + # v7 cache tail (CACH section: precomputed DC / FFT images) + # ------------------------------------------------------------------ + + def _cache_tail_offset(self) -> int: + """Deterministic file offset where the CACH tail starts (or would + start), derived purely from the header + Per-Angle Geometry Table — + independent of whether a cache tail is actually present. Used by + both the parser and the in-place writer.""" + waveform_bytes = sum( + int(self.n_rows[a]) * self.n_channels * int(self.n_frames[a]) + * self.samples_per_frame * self.bytes_per_sample + for a in range(self.n_angles) + ) + return self._data_offset + int(waveform_bytes) + + def _read_cache_block(self, f, hdr_fmt: str, magic: bytes, + stores: list[list]) -> int | None: + """Read one CACH sub-block header, then its per-angle image entries + into *stores* (one list per image the block stores per angle). + + Returns the header's flags byte, or None if the block is malformed. + """ + raw = f.read(struct.calcsize(hdr_fmt)) + if len(raw) < struct.calcsize(hdr_fmt): + return None + block_magic, flags, n_stored = struct.unpack(hdr_fmt, raw) + if block_magic != magic: + return None + for _ in range(n_stored): + (angle_idx,) = _read_struct(f, ">H") + if angle_idx >= self.n_angles: + break + shape = self.image_shape(angle_idx) + for store in stores: + store[angle_idx] = _read_f32_image(f, shape) + return flags + + def _parse_cach_section(self, offset: int): + """Parse the v7 CACH tail that holds precomputed DC/FFT images.""" + with open(self.path, "rb") as f: + f.seek(offset) + header_raw = f.read(CACH_HDR_SIZE) + if len(header_raw) < CACH_HDR_SIZE: + return + magic, cach_version, block_flags = struct.unpack(CACH_HDR_FMT, header_raw) + if magic != CACH_MAGIC or cach_version != CACH_VERSION: + return + + if block_flags & CACH_FLAG_DC: + if self._read_cache_block( + f, SDCB_HDR_FMT, SDCB_MAGIC, + [self.precomputed_dc3_mv, self.precomputed_dc4_mv]) is None: + return + + if block_flags & CACH_FLAG_FFT: + flags = self._read_cache_block( + f, SFFT_HDR_FMT, SFFT_MAGIC, [self.precomputed_freq_mhz]) + if flags is None: + return + self.precomputed_bg_sub = bool(flags & SFFT_FLAG_BG_SUB) + + def write_v7_cache(self, *, + new_dc3_mv: list[np.ndarray | None] | None = None, + new_dc4_mv: list[np.ndarray | None] | None = None, + new_freq_mhz: list[np.ndarray | None] | None = None, + new_bg_sub: bool | None = None): + """Store computed DC and/or FFT images into this file's CACH tail, + in place, converting a v6 source to v7 (or updating an existing v7 + file). Only the block(s) passed in are recomputed; whichever block + isn't passed is carried forward unchanged from whatever this + ``SrasFile`` already has in memory (from parsing, or a prior write + in this same session) — its bytes are never re-read from disk. + + The waveform data itself is never touched: the cache tail always + starts at ``_cache_tail_offset()``, a fixed offset derived from the + header and geometry table alone. + """ + if self.version not in (6, 7): + raise ValueError( + f"write_v7_cache only supports v6/v7 source files, got v{self.version}") + + final_dc3 = new_dc3_mv if new_dc3_mv is not None else self.precomputed_dc3_mv + final_dc4 = new_dc4_mv if new_dc4_mv is not None else self.precomputed_dc4_mv + final_freq = new_freq_mhz if new_freq_mhz is not None else self.precomputed_freq_mhz + final_bg_sub = new_bg_sub if new_bg_sub is not None else self.precomputed_bg_sub + + dc_entries = [a for a in range(self.n_angles) if final_dc3[a] is not None] + fft_entries = [a for a in range(self.n_angles) if final_freq[a] is not None] + + block_flags = ((CACH_FLAG_DC if dc_entries else 0) + | (CACH_FLAG_FFT if fft_entries else 0)) + + payload = bytearray() + payload += struct.pack(CACH_HDR_FMT, CACH_MAGIC, CACH_VERSION, block_flags) + + if dc_entries: + payload += struct.pack(SDCB_HDR_FMT, SDCB_MAGIC, 0, len(dc_entries)) + for a in dc_entries: + payload += struct.pack(">H", a) + payload += final_dc3[a].astype(">f4").tobytes() + payload += final_dc4[a].astype(">f4").tobytes() + + if fft_entries: + fft_flags = SFFT_FLAG_BG_SUB if final_bg_sub else 0 + payload += struct.pack(SFFT_HDR_FMT, SFFT_MAGIC, fft_flags, len(fft_entries)) + for a in fft_entries: + payload += struct.pack(">H", a) + payload += final_freq[a].astype(">f4").tobytes() + + with open(self.path, "r+b") as f: + f.seek(self._cache_tail_offset()) + f.write(payload) + f.truncate() + f.flush() + os.fsync(f.fileno()) + # Version-byte flip last: if the process dies before this point, + # the file is still readable as plain v6 (v6 parsing only + # bounds-checks per-angle offset+nbytes <= file_size, it never + # asserts exactly how many bytes follow the last angle) — so an + # interrupted write can never corrupt the file, only leave + # harmless trailing bytes that the next successful write + # overwrites via this same deterministic cache offset. + f.seek(4) + f.write(struct.pack("B", 7)) + f.flush() + os.fsync(f.fileno()) + + self.version = 7 + self.precomputed_dc3_mv = final_dc3 + self.precomputed_dc4_mv = final_dc4 + self.precomputed_freq_mhz = final_freq + self.precomputed_bg_sub = final_bg_sub + + # ------------------------------------------------------------------ + # Axes helpers + # ------------------------------------------------------------------ + + @property + def pixel_x_mm(self) -> float: + return self.velocity_mm_s / self.laser_freq_hz + + def x_axis_mm(self, angle_idx: int) -> np.ndarray: + n = int(self.n_frames[angle_idx]) + return self.x_start_mm[angle_idx] + np.arange(n) * self.pixel_x_mm + + def y_positions_mm(self, angle_idx: int) -> np.ndarray: + return self._y_pos_per_angle[angle_idx] + + def time_axis_ns(self) -> np.ndarray: + return np.arange(self.samples_per_frame) / self.sample_rate_hz * 1e9 + + def freq_axis_mhz(self, n_fft: int | None = None) -> np.ndarray: + n = n_fft if n_fft is not None else self.samples_per_frame + return np.fft.rfftfreq(n, d=1.0 / self.sample_rate_hz) / 1e6 diff --git a/sras_viewer.py b/sras_viewer.py index 13c657d..6336e09 100644 --- a/sras_viewer.py +++ b/sras_viewer.py @@ -10,1324 +10,73 @@ Channel semantics (fixed by sc3_aui_app.py acquisition settings): RF images are masked: pixels where CH4_dc < dc_threshold show 0. -Frame-count correction: the scanner writes the *configured* frame count in the -header before acquisition, but the scope may acquire fewer frames. The actual -count is computed from the file size and used for the reshape so channels are -correctly aligned. - -Scan geometry: v6/v7 files scan a different bounding box per angle (x_start, -x_delta, n_frames, n_rows all vary by angle), so geometry is exposed per-angle -via SrasFile.n_rows / n_frames / x_start_mm arrays and the x_axis_mm() / -y_positions_mm() methods. v2–v5 files have uniform geometry across angles, so -those arrays simply repeat the same value n_angles times. - -v7 files are v6 files with an optional trailing cache section holding -precomputed per-angle DC and/or FFT images (see scan_format.md), so display -never has to recompute them after the "Convert" menu's batch actions have -stored them once. +File parsing lives in sras_format, image/alignment math in sras_compute, and +background workers in sras_workers — the first two import neither Qt nor +matplotlib so multiprocessing children can load them cheaply. """ -import re -import sys -import struct import faulthandler -import numpy as np +import sys from pathlib import Path -from dataclasses import dataclass -import os -faulthandler.enable() # print a native stack trace on SIGSEGV/SIGABRT/etc. - -from PyQt6.QtWidgets import ( - QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, - QGroupBox, QLabel, QPushButton, QComboBox, QSpinBox, QDoubleSpinBox, - QFileDialog, QSizePolicy, QSplitter, QCheckBox, QFrame, QProgressDialog, - QDialog, QDialogButtonBox, QRadioButton, QButtonGroup, -) -from PyQt6.QtGui import QAction -from PyQt6.QtCore import Qt, QThread, pyqtSignal, QObject +import numpy as np from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg, NavigationToolbar2QT from matplotlib.figure import Figure from matplotlib.patches import Polygon -from matplotlib.lines import Line2D from matplotlib.path import Path as MplPath +from PyQt6.QtCore import QObject, Qt, QThread, pyqtSignal +from PyQt6.QtGui import QAction +from PyQt6.QtWidgets import ( + QApplication, QButtonGroup, QCheckBox, QComboBox, QDialog, QDialogButtonBox, + QDoubleSpinBox, QFileDialog, QFrame, QGroupBox, QHBoxLayout, QLabel, + QMainWindow, QProgressDialog, QPushButton, QRadioButton, QSizePolicy, + QSpinBox, QSplitter, QVBoxLayout, QWidget, +) + +import sras_compute as compute +from sras_compute import PYFFTW_AVAILABLE, apply_alignment +from sras_format import ( + CH1_IDX, CH3_IDX, CH4_IDX, CH_NAMES, SrasFile, adc_to_mv, mv_to_adc, + _FALLBACK_YMULT_MV, _FALLBACK_YOFF_ADC, +) +from sras_workers import ( + AngleAlignmentWorker, BatchCacheWorker, ComputeWorker, DcPrecomputeWorker, + LoadWorker, +) + +faulthandler.enable() # print a native stack trace on SIGSEGV/SIGABRT/etc. # --------------------------------------------------------------------------- -# FFT backend +# Display constants # --------------------------------------------------------------------------- -_pyfftw_available = False -try: - import pyfftw - pyfftw.interfaces.cache.enable() - _pyfftw_available = True -except ImportError: - pass - -import scipy.fft as scipy_fft -import scipy.ndimage as scipy_ndimage - -# Runtime-mutable settings changed via FftOptionsDialog -_fft_backend = "numpy" # "numpy" or "pyfftw" - - -def _do_rfft(x: np.ndarray, n: int | None = None, axis: int = -1, - workers: int = 1) -> np.ndarray: - """Dispatch rfft to the selected backend with optional multithreading.""" - if _fft_backend == "pyfftw" and _pyfftw_available: - return pyfftw.interfaces.numpy_fft.rfft(x, n=n, axis=axis, threads=workers) - return scipy_fft.rfft(x, n=n, axis=axis, workers=workers) - - -# --------------------------------------------------------------------------- -# SRAS format -# --------------------------------------------------------------------------- - -# v2–v5: fixed header, uniform geometry across angles (43 bytes) -HDR_FMT = ">4sBHHffffIIdBB" -HDR_SIZE = struct.calcsize(HDR_FMT) # 43 bytes - -# v6: fixed header, per-angle geometry in a separate table (49 bytes) -HDR_FMT_V6 = ">4sBHfffffffIdBB" -HDR_SIZE_V6 = struct.calcsize(HDR_FMT_V6) # 49 bytes - -# v6: per-angle geometry table record (x_start, x_delta, n_frames, n_rows) -GEO_FMT_V6 = ">ffIH" -GEO_SIZE_V6 = struct.calcsize(GEO_FMT_V6) # 14 bytes - -# v7: identical to v6 (same header/geometry/waveform layout, version byte -# is the only header difference) plus an optional trailing cache section -# ("CACH") holding precomputed per-angle DC and/or FFT images so a v7 file -# never needs to recompute them on open. See scan_format.md for the full -# spec. Sub-block sizes are per-angle (n_rows[a] * n_frames[a]), taken from -# the Per-Angle Geometry Table already parsed for v6 — no new geometry -# fields are needed. -CACH_MAGIC = b"CACH" -CACH_HDR_FMT = ">4sBB" # magic, cach_version, block_flags -CACH_HDR_SIZE = struct.calcsize(CACH_HDR_FMT) # 6 bytes -CACH_VERSION = 1 -CACH_FLAG_DC = 0x01 -CACH_FLAG_FFT = 0x02 - -SDCB_MAGIC = b"SDCB" -SDCB_HDR_FMT = ">4sBH" # magic, reserved, n_stored -SDCB_HDR_SIZE = struct.calcsize(SDCB_HDR_FMT) # 7 bytes - -SFFT_MAGIC = b"SFFT" -SFFT_HDR_FMT = ">4sBH" # magic, flags, n_stored -SFFT_HDR_SIZE = struct.calcsize(SFFT_HDR_FMT) # 7 bytes -SFFT_FLAG_BG_SUB = 0x01 - -# Fixed-order channels in the file: index 0=CH1, 1=CH3, 2=CH4 -# Fixed channel indices into the .sras data array (CH1=RF, CH3/CH4=Bias DC) -CH1_IDX, CH3_IDX, CH4_IDX = 0, 1, 2 - CH_LABELS = [ "CH1 — RF (FFT peak freq)", "CH3 — Bias A (DC mean)", "CH4 — Bias B (DC mean)", "CH1 — Velocity (SRAS)", ] -CH_NAMES = ["CH1", "CH3", "CH4", "VEL"] # Combo index for the derived velocity mode (uses CH1_IDX data) -VELOCITY_MODE_IDX = 3 +VELOCITY_MODE_IDX = 3 # All modes that operate on CH1 waveforms -CH1_DERIVED_MODES = (CH1_IDX, VELOCITY_MODE_IDX) - -# Fallback scope calibration used only when reading v2 files without embedded -# preambles. v3+ files carry the WFMOutpre string so these are not used. -# 50 mV/div, 8 div full-scale, int8 ADC, position = -2.72 div -# ymult = 50 mV × 8 / 256 = 1.5625 mV/count -# yoff = position × (256/8) = -2.72 × 32 = -87.04 (ADC count for 0 V) -_FALLBACK_YMULT_MV = 1.5625 # mV per ADC count -_FALLBACK_YOFF_ADC = -87.04 # ADC count that represents 0 V +CH1_DERIVED_MODES = (CH1_IDX, VELOCITY_MODE_IDX) CMAPS = ["gray", "viridis", "plasma", "inferno", "hot", "jet", "RdBu_r", "seismic"] - -def _parse_preamble(preamble: str) -> dict[str, float]: - """Extract YMULT, YOFF, YZERO from a Tektronix WFMOutpre string. - - Returns a dict with float values for whichever keys are present. - YMULT is left in V/count as the scope reports it. - """ - result = {} - for key in ("YMULT", "YOFF", "YZERO"): - m = re.search(rf'\b{key}\s+([-+]?\d*\.?\d+(?:[Ee][+-]?\d+)?)', preamble) - if m: - result[key] = float(m.group(1)) - return result - - -def mv_to_adc(mv: float, ymult_mv: float = _FALLBACK_YMULT_MV, - yoff_adc: float = _FALLBACK_YOFF_ADC, - yzero_mv: float = 0.0) -> float: - return (mv - yzero_mv) / ymult_mv + yoff_adc - - -def adc_to_mv(adc: float, ymult_mv: float = _FALLBACK_YMULT_MV, - yoff_adc: float = _FALLBACK_YOFF_ADC, - yzero_mv: float = 0.0) -> float: - return (adc - yoff_adc) * ymult_mv + yzero_mv - - -# --------------------------------------------------------------------------- -# File parser -# --------------------------------------------------------------------------- - -class SrasFile: - """Parsed in-memory representation of a v2–v7 .sras file. - - Scan geometry (rows, frames, x_start) is exposed per-angle via the - ``n_rows`` / ``n_frames`` / ``x_start_mm`` arrays and the ``x_axis_mm()`` - / ``y_positions_mm()`` methods, since v6/v7 files scan a different - bounding box per angle. v2–v5 files have uniform geometry, so these - arrays just repeat the same value ``n_angles`` times. Waveform data is - likewise exposed as ``data[angle_idx]``, an array of shape - ``(n_rows[a], n_channels, n_frames[a], samples_per_frame)``. - - v7 files (v6 plus an optional trailing cache section) expose any stored - precomputed images as ``precomputed_dc3_mv`` / ``precomputed_dc4_mv`` / - ``precomputed_freq_mhz`` — for v6/v7 these are ragged per-angle lists - (``list[np.ndarray | None]``, one entry per angle); for v5 they are dense - ``(n_angles, n_rows, n_frames)`` arrays, since v5 geometry is uniform. - """ - - def __init__(self, path: str): - self.path = Path(path) - self._parse() - - def _parse(self): - with open(self.path, "rb") as f: - magic = f.read(4) - if magic != b"SRAS": - raise ValueError(f"Bad magic bytes: {magic!r}") - (version,) = struct.unpack(">B", f.read(1)) - - self.version = version - if version in (2, 3, 4, 5): - self._parse_legacy() - elif version in (6, 7): - self._parse_v6() - else: - raise ValueError(f"Unsupported version: {version}") - - # ------------------------------------------------------------------ - # v2–v5 parsing (uniform geometry, flat waveform block) - # ------------------------------------------------------------------ - - def _parse_legacy(self): - with open(self.path, "rb") as f: - fields = struct.unpack(HDR_FMT, f.read(HDR_SIZE)) - (magic, ver, n_angles, n_rows, x_start, x_delta, vel, freq, - n_frames_hdr, spf, sr, bps, n_ch) = fields - - self.n_angles = n_angles - self.velocity_mm_s = float(vel) - self.laser_freq_hz = float(freq) - self.n_frames_header = n_frames_hdr # configured count (may be wrong) - self.samples_per_frame = spf - self.sample_rate_hz = float(sr) - self.bytes_per_sample = bps - self.n_channels = n_ch - - # Precomputed-image cache (populated when reading a v5 file). - # These are (n_angles, n_rows, n_frames) float32 arrays or None. - self.precomputed_freq_mhz: np.ndarray | None = None - self.precomputed_dc4_mv: np.ndarray | None = None - self.precomputed_dc3_mv: np.ndarray | None = None - self.precomputed_bg_sub: bool = False - self.scan_aborted = False - self.n_angles_declared = n_angles - - with open(self.path, "rb") as f: - f.seek(HDR_SIZE) - angles = np.frombuffer(f.read(n_angles * 4), dtype=">f4").astype(np.float32) - y_pos = np.frombuffer(f.read(n_rows * 4), dtype=">f4").astype(np.float32) - - if ver >= 3: - preambles = [] - for _ in range(n_ch): - (length,) = struct.unpack(">H", f.read(2)) - preambles.append(f.read(length).decode("utf-8")) - self.preambles = preambles - self.ch_ymult_mv = [] - self.ch_yoff_adc = [] - self.ch_yzero_mv = [] - for p in preambles: - cal = _parse_preamble(p) - # YMULT from scope is V/count; store as mV/count - self.ch_ymult_mv.append(cal.get("YMULT", _FALLBACK_YMULT_MV / 1000) * 1000) - self.ch_yoff_adc.append(cal.get("YOFF", _FALLBACK_YOFF_ADC)) - # YZERO from scope is in V; store as mV - self.ch_yzero_mv.append(cal.get("YZERO", 0.0) * 1000) - else: - self.preambles = None - self.ch_ymult_mv = [_FALLBACK_YMULT_MV] * n_ch - self.ch_yoff_adc = [_FALLBACK_YOFF_ADC] * n_ch - self.ch_yzero_mv = [0.0] * n_ch - - if ver >= 4: - (n_bg,) = struct.unpack(">I", f.read(4)) - self.background = np.frombuffer(f.read(n_bg), dtype=np.int8).astype(np.float32) - else: - self.background = None - - # Record the byte offset where raw waveform data begins. - # np.memmap will use this to map only the waveform section. - data_offset = f.tell() - - # ---- Determine actual frame count from file size --------------- - # For v4 and earlier the header n_frames may be the *configured* - # count before acquisition; the actual count is derived from the - # bytes on disk. For v5 files a PREC tail follows the waveform - # data, so we must not include those extra bytes in the frame count. - file_size = self.path.stat().st_size - samples_per_row_per_ch = n_ch * spf - - # Upper bound: bytes from data_offset to end of file - available_bytes = file_size - data_offset - - if ver == 5: - actual_n_frames = n_frames_hdr - remainder = 0 - else: - total_samples = available_bytes // bps - actual_n_frames = total_samples // (n_angles * n_rows * samples_per_row_per_ch) - remainder = total_samples % (n_angles * n_rows * samples_per_row_per_ch) - - self.frame_count_mismatch = (actual_n_frames != n_frames_hdr) - self.n_frames_remainder = remainder - - # ---- Memory-map the waveform data (zero RAM cost) -------------- - # Instead of f.read() → astype() (which peaks at 2× file size), - # memmap lets the OS page only the bytes that are actually touched. - waveform_dtype = np.int8 if bps == 1 else ">i2" - waveform_shape = (n_angles, n_rows, n_ch, actual_n_frames, spf) - data5d = np.memmap( - str(self.path), - dtype=waveform_dtype, - mode="r", - offset=data_offset, - shape=waveform_shape, - ) - # Expose as a list of per-angle views so downstream code shares one - # indexing convention with v6: sras.data[a][row, ch, frame, sample] - self.data = [data5d[a] for a in range(n_angles)] - - # Uniform per-angle geometry, repeated so callers don't need to - # special-case legacy vs. v6 files. - self.n_rows = np.full(n_angles, n_rows, dtype=np.int64) - self.n_frames = np.full(n_angles, actual_n_frames, dtype=np.int64) - self.x_start_mm = np.full(n_angles, float(x_start), dtype=np.float64) - self._y_pos_per_angle = [y_pos] * n_angles - - self.angles_deg = angles - - # ---- Read v5 precomputed section if present -------------------- - if ver >= 5: - waveform_bytes = actual_n_frames * n_angles * n_rows * n_ch * spf * bps - prec_offset = data_offset + waveform_bytes - if file_size > prec_offset: - self._parse_prec_section(prec_offset, n_angles, n_rows, actual_n_frames) - - def _parse_prec_section(self, offset: int, - n_angles: int, n_rows: int, n_frames: int): - """Parse the v5 PREC tail that holds precomputed images.""" - _PREC_MAGIC = b"PREC" - px = n_rows * n_frames # pixels per angle image - img_bytes = px * 4 # float32 - - with open(self.path, "rb") as f: - f.seek(offset) - header_raw = f.read(6) # magic(4) + fmt_ver(1) + flags(1) - if len(header_raw) < 6 or header_raw[:4] != _PREC_MAGIC: - return - flags = header_raw[5] - self.precomputed_bg_sub = bool(flags & 0x01) - - (n_stored,) = struct.unpack(">H", f.read(2)) - if n_stored == 0: - return - - freq_buf = np.zeros((n_angles, n_rows, n_frames), dtype=np.float32) - dc4_buf = np.zeros((n_angles, n_rows, n_frames), dtype=np.float32) - dc3_buf = np.zeros((n_angles, n_rows, n_frames), dtype=np.float32) - - for _ in range(n_stored): - (aidx,) = struct.unpack(">H", f.read(2)) - if aidx >= n_angles: - break - freq_buf[aidx] = np.frombuffer( - f.read(img_bytes), dtype=">f4").reshape(n_rows, n_frames) - dc4_buf[aidx] = np.frombuffer( - f.read(img_bytes), dtype=">f4").reshape(n_rows, n_frames) - dc3_buf[aidx] = np.frombuffer( - f.read(img_bytes), dtype=">f4").reshape(n_rows, n_frames) - - self.precomputed_freq_mhz = freq_buf - self.precomputed_dc4_mv = dc4_buf - self.precomputed_dc3_mv = dc3_buf - - # ------------------------------------------------------------------ - # v6 parsing (per-angle geometry, ragged waveform blocks) - # ------------------------------------------------------------------ - - def _parse_v6(self): - with open(self.path, "rb") as f: - fields = struct.unpack(HDR_FMT_V6, f.read(HDR_SIZE_V6)) - (magic, ver, n_angles, x_start_nom, y_start_nom, x_delta_nom, - y_delta_nom, row_spacing, vel, freq, spf, sr, bps, n_ch) = fields - - n_angles_declared = n_angles - - self.velocity_mm_s = float(vel) - self.laser_freq_hz = float(freq) - self.samples_per_frame = spf - self.sample_rate_hz = float(sr) - self.bytes_per_sample = bps - self.n_channels = n_ch - - # Reference-only fields: the ROI as entered before per-angle - # bounding-box expansion. Actual per-angle geometry used for - # rendering comes from the Per-Angle Geometry Table below. - self.x_start_nominal_mm = float(x_start_nom) - self.y_start_nominal_mm = float(y_start_nom) - self.x_delta_nominal_mm = float(x_delta_nom) - self.y_delta_nominal_mm = float(y_delta_nom) - self.row_spacing_mm = float(row_spacing) - - self.n_frames_header = None - self.frame_count_mismatch = False - self.n_frames_remainder = 0 - - angles = np.frombuffer(f.read(n_angles * 4), dtype=">f4").astype(np.float32) - - x_start = np.empty(n_angles, dtype=np.float64) - n_frames = np.empty(n_angles, dtype=np.int64) - n_rows = np.empty(n_angles, dtype=np.int64) - for a in range(n_angles): - xs, xd, nf, nr = struct.unpack(GEO_FMT_V6, f.read(GEO_SIZE_V6)) - x_start[a] = xs - n_frames[a] = nf - n_rows[a] = nr - - y_pos_per_angle = [] - for a in range(n_angles): - nr = int(n_rows[a]) - y_pos_per_angle.append( - np.frombuffer(f.read(nr * 4), dtype=">f4").astype(np.float32)) - - preambles = [] - for _ in range(n_ch): - (length,) = struct.unpack(">H", f.read(2)) - preambles.append(f.read(length).decode("utf-8")) - self.preambles = preambles - self.ch_ymult_mv = [] - self.ch_yoff_adc = [] - self.ch_yzero_mv = [] - for p in preambles: - cal = _parse_preamble(p) - self.ch_ymult_mv.append(cal.get("YMULT", _FALLBACK_YMULT_MV / 1000) * 1000) - self.ch_yoff_adc.append(cal.get("YOFF", _FALLBACK_YOFF_ADC)) - self.ch_yzero_mv.append(cal.get("YZERO", 0.0) * 1000) - - (n_bg,) = struct.unpack(">I", f.read(4)) - self.background = np.frombuffer(f.read(n_bg), dtype=np.int8).astype(np.float32) - - data_offset = f.tell() - - self._data_offset = data_offset - - # ---- Memory-map each angle's ragged waveform block ------------- - # v6 gives each angle its own row/frame count, so waveform data is - # no longer one uniform (n_angles, n_rows, ...) block — each angle's - # block sits at a different offset with its own shape. An aborted - # scan truncates the file mid-angle; per the format spec we keep - # whatever complete angles are present rather than refusing to open - # the file. - file_size = self.path.stat().st_size - waveform_dtype = np.int8 if bps == 1 else ">i2" - - data = [] - offset = data_offset - n_complete = 0 - for a in range(n_angles): - nr = int(n_rows[a]) - nf = int(n_frames[a]) - nbytes = nr * n_ch * nf * spf * bps - if offset + nbytes > file_size: - break - data.append(np.memmap( - str(self.path), dtype=waveform_dtype, mode="r", - offset=offset, shape=(nr, n_ch, nf, spf), - )) - offset += nbytes - n_complete += 1 - - if n_complete == 0: - raise ValueError( - "v6 file has no complete angle blocks — scan was aborted " - "before the first angle finished.") - - self.data = data - self.n_angles = n_complete - self.n_angles_declared = n_angles_declared - self.scan_aborted = n_complete < n_angles_declared - self.angles_deg = angles[:n_complete] - self.x_start_mm = x_start[:n_complete] - self.n_frames = n_frames[:n_complete] - self.n_rows = n_rows[:n_complete] - self._y_pos_per_angle = y_pos_per_angle[:n_complete] - - # Precomputed-image cache (v7 only). Ragged per-angle lists — unlike - # v5's dense (n_angles, n_rows, n_frames) arrays, v6/v7 geometry - # varies per angle, so each entry is its own (n_rows[a], n_frames[a]) - # array or None if that angle isn't cached yet. - self.precomputed_freq_mhz: list[np.ndarray | None] = [None] * n_complete - self.precomputed_dc4_mv: list[np.ndarray | None] = [None] * n_complete - self.precomputed_dc3_mv: list[np.ndarray | None] = [None] * n_complete - self.precomputed_bg_sub: bool = False - - if self.version == 7 and offset < file_size: - self._parse_cach_section(offset) - - # ------------------------------------------------------------------ - # v7 cache tail (CACH section: precomputed DC / FFT images) - # ------------------------------------------------------------------ - - def _cache_tail_offset(self) -> int: - """Deterministic file offset where the CACH tail starts (or would - start), derived purely from the header + Per-Angle Geometry Table — - independent of whether a cache tail is actually present. Used by - both the parser and the in-place writer.""" - n_ch = self.n_channels - spf = self.samples_per_frame - bps = self.bytes_per_sample - waveform_bytes = int(sum( - int(self.n_rows[a]) * n_ch * int(self.n_frames[a]) * spf * bps - for a in range(self.n_angles) - )) - return self._data_offset + waveform_bytes - - def _parse_cach_section(self, offset: int): - """Parse the v7 CACH tail that holds precomputed DC/FFT images.""" - with open(self.path, "rb") as f: - f.seek(offset) - header_raw = f.read(CACH_HDR_SIZE) - if len(header_raw) < CACH_HDR_SIZE: - return - magic, cach_version, block_flags = struct.unpack(CACH_HDR_FMT, header_raw) - if magic != CACH_MAGIC or cach_version != CACH_VERSION: - return - - if block_flags & CACH_FLAG_DC: - sdcb_raw = f.read(SDCB_HDR_SIZE) - if len(sdcb_raw) < SDCB_HDR_SIZE: - return - sdcb_magic, _reserved, n_stored = struct.unpack(SDCB_HDR_FMT, sdcb_raw) - if sdcb_magic != SDCB_MAGIC: - return - for _ in range(n_stored): - (angle_idx,) = struct.unpack(">H", f.read(2)) - if angle_idx >= self.n_angles: - break - px = int(self.n_rows[angle_idx]) * int(self.n_frames[angle_idx]) - shape = (int(self.n_rows[angle_idx]), int(self.n_frames[angle_idx])) - self.precomputed_dc3_mv[angle_idx] = np.frombuffer( - f.read(px * 4), dtype=">f4").reshape(shape) - self.precomputed_dc4_mv[angle_idx] = np.frombuffer( - f.read(px * 4), dtype=">f4").reshape(shape) - - if block_flags & CACH_FLAG_FFT: - sfft_raw = f.read(SFFT_HDR_SIZE) - if len(sfft_raw) < SFFT_HDR_SIZE: - return - sfft_magic, flags, n_stored = struct.unpack(SFFT_HDR_FMT, sfft_raw) - if sfft_magic != SFFT_MAGIC: - return - self.precomputed_bg_sub = bool(flags & SFFT_FLAG_BG_SUB) - for _ in range(n_stored): - (angle_idx,) = struct.unpack(">H", f.read(2)) - if angle_idx >= self.n_angles: - break - px = int(self.n_rows[angle_idx]) * int(self.n_frames[angle_idx]) - shape = (int(self.n_rows[angle_idx]), int(self.n_frames[angle_idx])) - self.precomputed_freq_mhz[angle_idx] = np.frombuffer( - f.read(px * 4), dtype=">f4").reshape(shape) - - def write_v7_cache(self, *, - new_dc3_mv: list[np.ndarray | None] | None = None, - new_dc4_mv: list[np.ndarray | None] | None = None, - new_freq_mhz: list[np.ndarray | None] | None = None, - new_bg_sub: bool | None = None): - """Store computed DC and/or FFT images into this file's CACH tail, - in place, converting a v6 source to v7 (or updating an existing v7 - file). Only the block(s) passed in are recomputed; whichever block - isn't passed is carried forward unchanged from whatever this - ``SrasFile`` already has in memory (from parsing, or a prior write - in this same session) — its bytes are never re-read from disk. - - The waveform data itself is never touched: the cache tail always - starts at ``_cache_tail_offset()``, a fixed offset derived from the - header and geometry table alone. - """ - if self.version not in (6, 7): - raise ValueError( - f"write_v7_cache only supports v6/v7 source files, got v{self.version}") - - final_dc3 = new_dc3_mv if new_dc3_mv is not None else self.precomputed_dc3_mv - final_dc4 = new_dc4_mv if new_dc4_mv is not None else self.precomputed_dc4_mv - final_freq = new_freq_mhz if new_freq_mhz is not None else self.precomputed_freq_mhz - final_bg_sub = new_bg_sub if new_bg_sub is not None else self.precomputed_bg_sub - - dc_entries = [a for a in range(self.n_angles) if final_dc3[a] is not None] - fft_entries = [a for a in range(self.n_angles) if final_freq[a] is not None] - - block_flags = 0 - if dc_entries: - block_flags |= CACH_FLAG_DC - if fft_entries: - block_flags |= CACH_FLAG_FFT - - payload = bytearray() - payload += struct.pack(CACH_HDR_FMT, CACH_MAGIC, CACH_VERSION, block_flags) - - if dc_entries: - payload += struct.pack(SDCB_HDR_FMT, SDCB_MAGIC, 0, len(dc_entries)) - for a in dc_entries: - payload += struct.pack(">H", a) - payload += final_dc3[a].astype(">f4").tobytes() - payload += final_dc4[a].astype(">f4").tobytes() - - if fft_entries: - fft_flags = SFFT_FLAG_BG_SUB if final_bg_sub else 0 - payload += struct.pack(SFFT_HDR_FMT, SFFT_MAGIC, fft_flags, len(fft_entries)) - for a in fft_entries: - payload += struct.pack(">H", a) - payload += final_freq[a].astype(">f4").tobytes() - - cache_offset = self._cache_tail_offset() - with open(self.path, "r+b") as f: - f.seek(cache_offset) - f.write(payload) - f.truncate() - f.flush() - os.fsync(f.fileno()) - # Version-byte flip last: if the process dies before this point, - # the file is still readable as plain v6 (v6 parsing only - # bounds-checks per-angle offset+nbytes <= file_size, it never - # asserts exactly how many bytes follow the last angle) — so an - # interrupted write can never corrupt the file, only leave - # harmless trailing bytes that the next successful write - # overwrites via this same deterministic cache_offset. - f.seek(4) - f.write(struct.pack("B", 7)) - f.flush() - os.fsync(f.fileno()) - - self.version = 7 - self.precomputed_dc3_mv = final_dc3 - self.precomputed_dc4_mv = final_dc4 - self.precomputed_freq_mhz = final_freq - self.precomputed_bg_sub = final_bg_sub - - # ------------------------------------------------------------------ - # Axes helpers - # ------------------------------------------------------------------ - - @property - def pixel_x_mm(self) -> float: - return self.velocity_mm_s / self.laser_freq_hz - - def x_axis_mm(self, angle_idx: int) -> np.ndarray: - n = int(self.n_frames[angle_idx]) - return self.x_start_mm[angle_idx] + np.arange(n) * self.pixel_x_mm - - def y_positions_mm(self, angle_idx: int) -> np.ndarray: - return self._y_pos_per_angle[angle_idx] - - def time_axis_ns(self) -> np.ndarray: - return np.arange(self.samples_per_frame) / self.sample_rate_hz * 1e9 - - def freq_axis_mhz(self, n_fft: int | None = None) -> np.ndarray: - n = n_fft if n_fft is not None else self.samples_per_frame - return np.fft.rfftfreq(n, d=1.0 / self.sample_rate_hz) / 1e6 - - -# --------------------------------------------------------------------------- -# Image computation (vectorised) -# --------------------------------------------------------------------------- - - -# Rows are batched so the float32 working buffer for one channel's chunk -# (chunk_rows × n_frames × spf × 4 bytes) stays under this budget. A fixed -# row count (the original design) works fine for small legacy scans but is -# catastrophic for a v6 scan with a large per-angle frame/sample count — -# e.g. a 7500-frame × 2500-sample angle needs ~2.4 GB for a single 32-row -# chunk, times several such buffers alive at once for the FFT step, which -# can exceed physical RAM entirely on its own. Sizing the chunk to the -# actual dimensions keeps peak RAM bounded regardless of scan size. -_CHUNK_BYTES_BUDGET = 128 * 1024 * 1024 # ~128 MB per channel-buffer chunk -_CHUNK_ROWS_MAX = 32 # cap for small scans (old behavior) - - -def _chunk_rows_for(n_frames: int, samples_per_frame: int) -> int: - bytes_per_row = max(1, n_frames * samples_per_frame * 4) # float32 - rows = _CHUNK_BYTES_BUDGET // bytes_per_row - return int(max(1, min(_CHUNK_ROWS_MAX, rows))) - - -def compute_dc_image(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray: - """Mean of each waveform → (n_rows, n_frames) float32. - - Processes in row chunks sized to a fixed memory budget (see - ``_chunk_rows_for``) so the float32 working buffer stays bounded - regardless of scan size. - """ - n_rows = int(sras.n_rows[angle_idx]) - n_frames = int(sras.n_frames[angle_idx]) - data = sras.data[angle_idx] - chunk_rows = _chunk_rows_for(n_frames, sras.samples_per_frame) - img = np.empty((n_rows, n_frames), dtype=np.float32) - for r0 in range(0, n_rows, chunk_rows): - r1 = min(r0 + chunk_rows, n_rows) - img[r0:r1] = ( - data[r0:r1, ch_idx, :, :] - .astype(np.float32) - .mean(axis=-1) - ) - return img - - -def _cached_dc_mv(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray | None: - """Return a v7-cached DC image (mV, already converted) for - (angle_idx, ch_idx) if the file's CACH section has it, else None. - - Only applies to the ragged v6/v7 list representation of - ``precomputed_dc3_mv``/``precomputed_dc4_mv`` — v5's dense PREC arrays - don't track per-angle validity (unstored angles are left at 0.0 rather - than a sentinel), so they're intentionally not consulted here. - """ - store = sras.precomputed_dc3_mv if ch_idx == CH3_IDX else sras.precomputed_dc4_mv - if isinstance(store, list) and angle_idx < len(store): - return store[angle_idx] - return None - - -def compute_rf_image(sras: SrasFile, angle_idx: int, - dc_threshold_mv: float, - apply_bg_sub: bool = True, - n_fft: int | None = None, - dc4_mv: np.ndarray | None = None) -> np.ndarray: - """FFT of each CH1 waveform; pixel = peak frequency in MHz. - - Pixels where CH4_dc < dc_threshold_mv are set to 0 — and the FFT is - never run for them, since that's the expensive part. The masked-out - CH1 samples are also never *read*: the boolean mask is applied to the - raw memmap slice before any dtype conversion, so numpy only pages in - the bytes for pixels that pass the threshold (an untouched memmap - page is never read from disk). If the DC4 image for this angle is - already known (e.g. from the DC-channel precompute cache), pass it as - *dc4_mv* (mV, shape (n_rows, n_frames)) to reuse it directly instead - of re-reading/re-averaging the CH4 channel here; otherwise it's - computed chunk-by-chunk internally (which does need every pixel's - CH4 data, since that's what determines validity in the first place). - - Fast path: if the file contains precomputed peak-frequency images for - this angle (v5's dense per-file PREC section, or v7's ragged per-angle - CACH section), and zero-padding is not active, and the bg-sub flag - matches, the stored image is used directly — no FFT is run. - - Otherwise, data is processed in row chunks sized to a fixed memory - budget (see ``_chunk_rows_for``) to bound peak RAM regardless of scan - size. - """ - n_rows = int(sras.n_rows[angle_idx]) - n_frames = int(sras.n_frames[angle_idx]) - data = sras.data[angle_idx] - - # ---- Fast path: precomputed images (v5 PREC or v7 CACH) ---------------- - # v5 stores a dense (n_angles, n_rows, n_frames) ndarray (uniform - # geometry); v6/v7 store a ragged list, one entry per angle (None where - # that angle hasn't been cached), since geometry varies per angle. - freq_store = sras.precomputed_freq_mhz - if isinstance(freq_store, list): - angle_cached = angle_idx < len(freq_store) and freq_store[angle_idx] is not None - else: - angle_cached = freq_store is not None - - can_use_precomputed = ( - angle_cached - and n_fft is None # no custom zero-padding - and sras.precomputed_bg_sub == (apply_bg_sub and sras.background is not None) - ) - if can_use_precomputed: - freq_img = freq_store[angle_idx].copy() - - # DC4 mask, in priority order: already-cached DC block, caller- - # supplied image, or a fresh (cheap — no FFT) recompute. - dc4_store = sras.precomputed_dc4_mv - if isinstance(dc4_store, list): - dc4_img = _cached_dc_mv(sras, angle_idx, CH4_IDX) - else: - dc4_img = dc4_store[angle_idx] if dc4_store is not None else None - if dc4_img is None: - if dc4_mv is not None: - dc4_img = dc4_mv - else: - dc4_img = adc_to_mv(compute_dc_image(sras, angle_idx, CH4_IDX), - sras.ch_ymult_mv[CH4_IDX], sras.ch_yoff_adc[CH4_IDX], - sras.ch_yzero_mv[CH4_IDX]) - - freq_img[dc4_img < dc_threshold_mv] = 0.0 - return freq_img - - # ---- Chunked FFT path -------------------------------------------------- - freq_axis = sras.freq_axis_mhz(n_fft) - img = np.zeros((n_rows, n_frames), dtype=np.float32) - _n_workers = os.cpu_count() or 4 - n_fft_bins = n_fft if n_fft is not None else sras.samples_per_frame - chunk_rows = _chunk_rows_for(n_frames, max(sras.samples_per_frame, n_fft_bins)) - - for r0 in range(0, n_rows, chunk_rows): - r1 = min(r0 + chunk_rows, n_rows) - - # DC mask for this chunk (float32 expansion is only chunk-sized) - if dc4_mv is not None: - dc4_chunk = dc4_mv[r0:r1] - else: - dc4_raw = data[r0:r1, CH4_IDX, :, :].astype(np.float32) - dc4_chunk = adc_to_mv(dc4_raw.mean(axis=-1), - sras.ch_ymult_mv[CH4_IDX], - sras.ch_yoff_adc[CH4_IDX], - sras.ch_yzero_mv[CH4_IDX]) - del dc4_raw - valid = dc4_chunk >= dc_threshold_mv # True = above threshold = run FFT - - if not valid.any(): - continue - - # Index the raw memmap slice with the boolean mask *before* - # converting dtype — this is a lazy view until touched, so only - # the (n_valid, spf) selected elements are actually read from - # disk; masked-out pixels' pages are never paged in at all. - valid_waves = data[r0:r1, CH1_IDX, :, :][valid].astype(np.float32) - - if apply_bg_sub and sras.background is not None: - valid_waves -= sras.background # background is 1-D (spf,) - - fft_pow = np.abs(_do_rfft(valid_waves, n=n_fft, axis=-1, workers=_n_workers)) ** 2 - del valid_waves - fft_pow[:, 0] = 0.0 # suppress DC bin - peak_bins = np.argmax(fft_pow, axis=-1) - del fft_pow - - img[r0:r1][valid] = freq_axis[peak_bins] - - return img - - -# --------------------------------------------------------------------------- -# Angle alignment (Fusion menu) -# -# Puts every angle's images onto one shared, zero-padded pixel grid using a -# rigid transform only (rotation + translation, never scale). Rotation for -# angle `a` is the *known* scan-angle delta relative to a reference angle — -# never searched. Only the residual translation is found, via FFT phase -# correlation of each angle's binarized CH4 ("dc-mask") image. -# -# Rotation is done in physical mm space rather than on raw pixel indices: -# the x-pixel pitch (SrasFile.pixel_x_mm) is file-wide constant but the -# y-pixel pitch (row spacing) can differ from it, and for v6 files can even -# vary per angle. Rotating the raw index grid directly would implicitly -# assume square pixels and shear a non-square-pixel image — an unwanted -# effective anisotropic scale. Instead each angle gets one affine that maps -# shared-canvas pixel index -> mm -> undo rotation/shift -> that angle's own -# local mm -> that angle's own raw pixel index, matching the output->input -# convention scipy.ndimage.affine_transform expects. -# --------------------------------------------------------------------------- - -@dataclass -class AngleTransform: - rotation_deg: float - shift_mm: tuple[float, float] # (dx_mm, dy_mm) found by phase correlation - matrix: np.ndarray # (2,2): canvas (row,col) -> this angle's raw (row,col) - offset: np.ndarray # (2,) - - -@dataclass -class AlignmentResult: - ref_angle_idx: int - dc_threshold_mv: float - canvas_shape: tuple[int, int] # (n_rows, n_cols) - canvas_dx_mm: float - canvas_dy_mm: float - canvas_origin_mm: tuple[float, float] # mm at canvas pixel index (0, 0) - per_angle: dict[int, AngleTransform] - - -def _pixel_pitch_mm(sras: SrasFile, angle_idx: int) -> tuple[float, float]: - """(dx, dy) mm/pixel for one angle: dx is the file-wide constant - pixel_x_mm; dy is this angle's own row spacing (assumed uniform, same - assumption _redraw_image already makes when it builds the display - extent).""" - dx = sras.pixel_x_mm - y = sras.y_positions_mm(angle_idx) - dy = float(y[1] - y[0]) if len(y) > 1 else 1.0 - return dx, dy - - -def _bbox_center_mm(sras: SrasFile, angle_idx: int) -> tuple[float, float]: - x = sras.x_axis_mm(angle_idx) - y = sras.y_positions_mm(angle_idx) - return float((x[0] + x[-1]) / 2.0), float((y[0] + y[-1]) / 2.0) - - -def _bbox_corners_mm(sras: SrasFile, angle_idx: int) -> np.ndarray: - """4 corners (x, y) of this angle's raw mm bounding box, shape (4, 2).""" - x = sras.x_axis_mm(angle_idx) - y = sras.y_positions_mm(angle_idx) - return np.array([[xx, yy] for xx in (x[0], x[-1]) for yy in (y[0], y[-1])]) - - -def _rotation_matrix(theta_deg: float) -> np.ndarray: - t = np.radians(theta_deg) - c, s = np.cos(t), np.sin(t) - return np.array([[c, -s], [s, c]]) # CCW rotation acting on (x, y) - - -def _build_affine_canvas_to_raw(sras: SrasFile, angle_idx: int, ref_idx: int, - shift_mm: tuple[float, float], - canvas_dx: float, canvas_dy: float, - canvas_origin_mm: tuple[float, float] - ) -> tuple[np.ndarray, np.ndarray]: - """matrix, offset s.t. raw_index = matrix @ [row_out, col_out] + offset, - matching scipy.ndimage.affine_transform's output->input convention. - - Pipeline (all mm unless noted): - [X;Y] = A_out @ [row_out;col_out] + b_out # canvas idx -> ref-frame mm - [lx;ly] = R(theta)^T @ ([X;Y]-c_ref-shift) + c_a # undo rotation+shift -> angle a's local mm - [row;col] = D @ ([lx;ly] - [x_start_a; y0_a]) # local mm -> angle a's raw idx - - where theta = angles_deg[angle_idx] - angles_deg[ref_idx], c_ref/c_a are - each angle's own raw-bbox mm centroid (the rotation pivot — this keeps - rotated content centered, minimizing required canvas padding), and - A_out/D are the index<->mm scaling matrices for the canvas pitch and - this angle's own native pitch respectively. - """ - theta = float(sras.angles_deg[angle_idx] - sras.angles_deg[ref_idx]) - Rinv = _rotation_matrix(theta).T - cx_a, cy_a = _bbox_center_mm(sras, angle_idx) - cx_ref, cy_ref = _bbox_center_mm(sras, ref_idx) - dx_a, dy_a = _pixel_pitch_mm(sras, angle_idx) - x0_a = float(sras.x_start_mm[angle_idx]) - y0_a = float(sras.y_positions_mm(angle_idx)[0]) - - A_out = np.array([[0.0, canvas_dx], [canvas_dy, 0.0]]) # [row,col] -> [X,Y] - b_out = np.array(canvas_origin_mm, dtype=np.float64) - D = np.array([[0.0, 1.0 / dy_a], [1.0 / dx_a, 0.0]]) # [x,y] -> [row,col] - shift = np.array(shift_mm, dtype=np.float64) - c_ref_v = np.array([cx_ref, cy_ref]) - c_a_v = np.array([cx_a, cy_a]) - origin_a = np.array([x0_a, y0_a]) - - matrix = D @ Rinv @ A_out - offset = D @ Rinv @ (b_out - c_ref_v - shift) + D @ (c_a_v - origin_a) - return matrix, offset - - -def apply_alignment(result: AlignmentResult, angle_idx: int, img: np.ndarray, - order: int = 0) -> np.ndarray: - """Resample any already-computed 2D image for `angle_idx` (same shape as - that angle's raw (n_rows, n_frames) — e.g. compute_dc_image / - compute_rf_image output) onto the shared alignment canvas. order=0 - (nearest) avoids blending real data with zero-padding or with - masked-out (0-valued) CH1/velocity pixels at mask edges. Channel- - agnostic: the same per-angle transform (found from the CH4 mask) works - for any channel's image of that angle.""" - t = result.per_angle[angle_idx] - return scipy_ndimage.affine_transform( - img.astype(np.float32, copy=False), t.matrix, offset=t.offset, - output_shape=result.canvas_shape, order=order, - mode="constant", cval=0.0) - - -def _block_mean_downsample(img: np.ndarray, factor: int) -> np.ndarray: - if factor <= 1: - return img - h, w = img.shape - h2, w2 = (h // factor) * factor, (w // factor) * factor - trimmed = img[:h2, :w2] - return trimmed.reshape(h2 // factor, factor, w2 // factor, factor).mean(axis=(1, 3)) - - -def _phase_correlate_shift(ref_img: np.ndarray, mov_img: np.ndarray) -> tuple[int, int]: - """FFT normalized cross-power-spectrum phase correlation. Returns the - integer (dr, dc) pixel shift of mov_img relative to ref_img; both must - be the same shape. Risk: if the true shift is near +/- half the array - size, wraparound can bias the peak — mitigated by the generous - margin_frac padding in _working_canvas_for_pair, which keeps the true - residual shift small relative to the correlation canvas.""" - F1 = scipy_fft.fft2(ref_img.astype(np.float64)) - F2 = scipy_fft.fft2(mov_img.astype(np.float64)) - R = F1 * np.conj(F2) - R /= np.maximum(np.abs(R), 1e-12) - corr = scipy_fft.ifft2(R).real - dr, dc = np.unravel_index(np.argmax(corr), corr.shape) - h, w = corr.shape - if dr > h // 2: - dr -= h - if dc > w // 2: - dc -= w - return int(dr), int(dc) - - -def _working_canvas_for_pair(sras: SrasFile, ref_idx: int, a_idx: int, - dx: float, dy: float, margin_frac: float = 0.3 - ) -> tuple[tuple[float, float], tuple[int, int]]: - """Union of the reference's own raw bbox and angle a's raw bbox rotated - (about its own center) into the ref frame with zero shift, padded by - margin_frac on each side — sized generously so the true phase- - correlation shift lands well inside the canvas (see - _phase_correlate_shift's wraparound note).""" - theta = float(sras.angles_deg[a_idx] - sras.angles_deg[ref_idx]) - R = _rotation_matrix(theta) - c_a = np.array(_bbox_center_mm(sras, a_idx)) - c_ref = np.array(_bbox_center_mm(sras, ref_idx)) - pts = list(_bbox_corners_mm(sras, ref_idx)) - for corner in _bbox_corners_mm(sras, a_idx): - pts.append(R @ (corner - c_a) + c_ref) - pts = np.array(pts) - x_min, y_min = pts.min(axis=0) - x_max, y_max = pts.max(axis=0) - pad_x, pad_y = (x_max - x_min) * margin_frac, (y_max - y_min) * margin_frac - x_min, x_max = x_min - pad_x, x_max + pad_x - y_min, y_max = y_min - pad_y, y_max + pad_y - n_cols = int(np.ceil((x_max - x_min) / dx)) + 1 - n_rows = int(np.ceil((y_max - y_min) / abs(dy))) + 1 - origin = (x_min, y_min if dy > 0 else y_max) - return origin, (n_rows, n_cols) - - -def _compute_angle_alignment(sras: SrasFile, ref_angle_idx: int, - dc_threshold_mv: float, - progress_cb=None) -> AlignmentResult: - """Top-level alignment driver. Runs on a background thread (see - AngleAlignmentWorker) — deliberately recomputes CH4 DC images from - scratch via compute_dc_image rather than reading the GUI-thread - _dc_cache dict, since background-thread workers must not touch - GUI-thread-owned caches (BatchCacheWorker follows the same rule).""" - n = sras.n_angles # already the *complete*-angle count for aborted v6 scans - dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx) - - # ---- Step 1: binarized CH4 mask per angle, native per-angle grid ----- - masks: dict[int, np.ndarray] = {} - for a in range(n): - dc4 = adc_to_mv(compute_dc_image(sras, a, CH4_IDX), - sras.ch_ymult_mv[CH4_IDX], sras.ch_yoff_adc[CH4_IDX], - sras.ch_yzero_mv[CH4_IDX]) - masks[a] = (dc4 >= dc_threshold_mv).astype(np.float32) - if progress_cb: - progress_cb(int((a + 1) / n * 25)) - - # ---- Step 2: coarse correlation stage (downsample first, then rotate) - # Downsampling before affine_transform (not after) is what keeps this - # tractable for a v6 scan with thousands of rows/frames per angle. - max_dim = max(max(m.shape) for m in masks.values()) - factor = max(1, int(np.ceil(max_dim / 1024))) - dx_c, dy_c = dx_ref * factor, dy_ref * factor - masks_small = {a: _block_mean_downsample(m, factor) for a, m in masks.items()} - - shifts_mm: dict[int, tuple[float, float]] = {ref_angle_idx: (0.0, 0.0)} - for a in range(n): - if a == ref_angle_idx: - continue - work_origin, work_shape = _working_canvas_for_pair( - sras, ref_angle_idx, a, dx_c, dy_c, margin_frac=0.3) - - # Coarse canvas->raw affine at native pitch, then rescale by /factor - # so it maps coarse-canvas idx -> coarse (downsampled) raw idx — - # exact for block-mean downsampling (up to the trimmed remainder). - m_a, o_a = _build_affine_canvas_to_raw( - sras, a, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin) - m_ref, o_ref = _build_affine_canvas_to_raw( - sras, ref_angle_idx, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin) - rotated_a = scipy_ndimage.affine_transform( - masks_small[a], m_a / factor, offset=o_a / factor, - output_shape=work_shape, order=0, mode="constant", cval=0.0) - embedded_ref = scipy_ndimage.affine_transform( - masks_small[ref_angle_idx], m_ref / factor, offset=o_ref / factor, - output_shape=work_shape, order=0, mode="constant", cval=0.0) - - dr, dc = _phase_correlate_shift(embedded_ref, rotated_a) - shifts_mm[a] = (dc * dx_c, dr * dy_c) - if progress_cb: - progress_cb(25 + int((a + 1) / n * 50)) - - # ---- Step 3: union bounding box over all angles (rotation+shift applied) - corners_ref_frame = [] - for a in range(n): - theta = float(sras.angles_deg[a] - sras.angles_deg[ref_angle_idx]) - R = _rotation_matrix(theta) - c_a = np.array(_bbox_center_mm(sras, a)) - c_ref = np.array(_bbox_center_mm(sras, ref_angle_idx)) - shift = np.array(shifts_mm[a]) - for corner in _bbox_corners_mm(sras, a): - corners_ref_frame.append(R @ (corner - c_a) + c_ref + shift) - corners_ref_frame = np.array(corners_ref_frame) - x_min, y_min = corners_ref_frame.min(axis=0) - x_max, y_max = corners_ref_frame.max(axis=0) - n_cols = int(np.ceil((x_max - x_min) / dx_ref)) + 1 - n_rows = int(np.ceil((y_max - y_min) / abs(dy_ref))) + 1 - canvas_origin_mm = (float(x_min), float(y_min if dy_ref > 0 else y_max)) - canvas_shape = (n_rows, n_cols) - - # ---- Step 4: final per-angle full-resolution affine (canvas -> raw idx) - per_angle: dict[int, AngleTransform] = {} - for a in range(n): - matrix, offset = _build_affine_canvas_to_raw( - sras, a, ref_angle_idx, shifts_mm[a], dx_ref, dy_ref, canvas_origin_mm) - theta = float(sras.angles_deg[a] - sras.angles_deg[ref_angle_idx]) - per_angle[a] = AngleTransform(theta, shifts_mm[a], matrix, offset) - if progress_cb: - progress_cb(75 + int((a + 1) / n * 25)) - - return AlignmentResult(ref_angle_idx, dc_threshold_mv, canvas_shape, - dx_ref, dy_ref, canvas_origin_mm, per_angle) - - -# --------------------------------------------------------------------------- -# Background workers -# --------------------------------------------------------------------------- - -class LoadWorker(QObject): - finished = pyqtSignal(object) # SrasFile | None - error = pyqtSignal(str) - - def __init__(self, path: str): - super().__init__() - self._path = path - - def run(self): - try: - self.finished.emit(SrasFile(self._path)) - except Exception as exc: - self.error.emit(str(exc)) - self.finished.emit(None) - - -class ComputeWorker(QObject): - """Computes one displayable image for (angle, channel). - - For CH1/Velocity (FFT-derived) channels, the FFT is only run for - pixels whose DC4 (Bias B) mean is at or above dc_threshold_mv — masked - pixels are left at 0 MHz without ever being FFT'd, since that's the - expensive part of a scan. If the DC4 image for this angle is already - known (e.g. from the DC-channel precompute cache), pass it in as - *dc4_mv* to skip re-reading the CH4 channel from disk entirely. - - Emits a plain ``np.ndarray`` (already in display units, masked for - CH1/Velocity) for both DC and FFT-derived channels. - """ - finished = pyqtSignal(object) - error = pyqtSignal(str) - - def __init__(self, sras: SrasFile, angle_idx: int, - ch_idx: int, - apply_bg_sub: bool = True, - n_fft: int | None = None, - dc_threshold_mv: float = 0.0, - dc4_mv: np.ndarray | None = None): - super().__init__() - self._sras = sras - self._angle = angle_idx - self._ch = ch_idx - self._apply_bg_sub = apply_bg_sub - self._n_fft = n_fft - self._dc_threshold = dc_threshold_mv - self._dc4_mv = dc4_mv - - def run(self): - try: - if self._ch in (CH1_IDX, VELOCITY_MODE_IDX): - img = compute_rf_image( - self._sras, self._angle, dc_threshold_mv=self._dc_threshold, - apply_bg_sub=self._apply_bg_sub, n_fft=self._n_fft, - dc4_mv=self._dc4_mv) - self.finished.emit(img) - else: - # DC channels: use the v7 cache if this angle is already - # stored (already mV), else compute fresh and convert. - cached = _cached_dc_mv(self._sras, self._angle, self._ch) - if cached is not None: - self.finished.emit(cached) - return - adc_img = compute_dc_image(self._sras, self._angle, self._ch) - img = adc_to_mv(adc_img, - self._sras.ch_ymult_mv[self._ch], - self._sras.ch_yoff_adc[self._ch], - self._sras.ch_yzero_mv[self._ch]) - self.finished.emit(img) - except Exception as exc: - self.error.emit(str(exc)) - - -class DcPrecomputeWorker(QObject): - """Computes CH3/CH4 DC images for every angle in the background. - - DC images are cheap (a per-waveform mean, no FFT) compared to the - CH1/Velocity FFT, so precomputing them for the whole file right after - load makes switching angles instant while on a DC channel, and also - means the FFT masking step (which needs a DC4 image) rarely has to - wait on anything. Emits one ``angle_done`` signal per angle as it - completes rather than waiting for the whole file, so the cache fills - in progressively. - """ - angle_done = pyqtSignal(int, np.ndarray, np.ndarray) # angle_idx, dc3_mv, dc4_mv - finished = pyqtSignal() - error = pyqtSignal(str) - - def __init__(self, sras: SrasFile): - super().__init__() - self._sras = sras - self._stop = False - - def stop(self): - self._stop = True - - def run(self): - try: - for a in range(self._sras.n_angles): - if self._stop: - break - cached_dc3 = _cached_dc_mv(self._sras, a, CH3_IDX) - cached_dc4 = _cached_dc_mv(self._sras, a, CH4_IDX) - if cached_dc3 is not None and cached_dc4 is not None: - self.angle_done.emit(a, cached_dc3, cached_dc4) - continue - dc3_mv = adc_to_mv( - compute_dc_image(self._sras, a, CH3_IDX), - self._sras.ch_ymult_mv[CH3_IDX], - self._sras.ch_yoff_adc[CH3_IDX], - self._sras.ch_yzero_mv[CH3_IDX]) - dc4_mv = adc_to_mv( - compute_dc_image(self._sras, a, CH4_IDX), - self._sras.ch_ymult_mv[CH4_IDX], - self._sras.ch_yoff_adc[CH4_IDX], - self._sras.ch_yzero_mv[CH4_IDX]) - self.angle_done.emit(a, dc3_mv, dc4_mv) - self.finished.emit() - except Exception as exc: - self.error.emit(str(exc)) - - -class BatchCacheWorker(QObject): - """Batch-computes and stores DC or FFT images into each of *paths*'s - v7 CACH tail, in place — converting v6 sources to v7 on first use, or - updating an existing v7 file's cache blocks without disturbing whatever - the other block already holds. - - *mode* is ``"dc"`` (CH3/CH4 mean images) or ``"fft"`` (CH1 peak-frequency - images, unmasked — masking is applied at display time, same as v5's - PREC convention). - - Emits ``progress(int)`` (0–100, weighted by total angle count across the - whole batch), ``file_done(str, str)`` (path, error message or "" on - success) after each file so one file's failure doesn't abort the batch, - and ``finished()`` once every file has been attempted. - """ - progress = pyqtSignal(int) - file_done = pyqtSignal(str, str) - finished = pyqtSignal() - - def __init__(self, paths: list[str], mode: str, apply_bg_sub: bool): - super().__init__() - self._paths = paths - self._mode = mode - self._apply_bg_sub = apply_bg_sub - - def run(self): - # Pass 1: quick open of each file just to weight progress by total - # angle count. Don't hold every file's memmap open at once — reopen - # fresh per file in pass 2 below. - total_angles = 0 - for path in self._paths: - try: - total_angles += SrasFile(path).n_angles - except Exception: - pass # unreadable files are reported properly in pass 2 - total_angles = max(total_angles, 1) - - done_angles = 0 - for path in self._paths: - try: - sras = SrasFile(path) - if sras.version not in (6, 7): - self.file_done.emit( - path, f"unsupported version {sras.version} — only " - "v6/v7 files can be batch-cached") - continue - - n = sras.n_angles - if self._mode == "dc": - new_dc3 = [None] * n - new_dc4 = [None] * n - for a in range(n): - new_dc3[a] = adc_to_mv( - compute_dc_image(sras, a, CH3_IDX), - sras.ch_ymult_mv[CH3_IDX], sras.ch_yoff_adc[CH3_IDX], - sras.ch_yzero_mv[CH3_IDX]) - new_dc4[a] = adc_to_mv( - compute_dc_image(sras, a, CH4_IDX), - sras.ch_ymult_mv[CH4_IDX], sras.ch_yoff_adc[CH4_IDX], - sras.ch_yzero_mv[CH4_IDX]) - done_angles += 1 - self.progress.emit(int(done_angles / total_angles * 100)) - sras.write_v7_cache(new_dc3_mv=new_dc3, new_dc4_mv=new_dc4) - else: - effective_bg = self._apply_bg_sub and sras.background is not None - new_freq = [None] * n - for a in range(n): - new_freq[a] = compute_rf_image( - sras, a, dc_threshold_mv=-1e9, # mask nothing - apply_bg_sub=effective_bg) - done_angles += 1 - self.progress.emit(int(done_angles / total_angles * 100)) - sras.write_v7_cache(new_freq_mhz=new_freq, new_bg_sub=effective_bg) - - self.file_done.emit(path, "") - except Exception as exc: - self.file_done.emit(path, str(exc)) - - self.finished.emit() - - -class AngleAlignmentWorker(QObject): - """Computes rotation+translation alignment for every angle in *sras*, - referenced to *ref_angle_idx*, from each angle's binarized CH4 mask. - Rotation is analytic (from sras.angles_deg); only translation is found - by phase correlation. Uses the same progress(int)/finished(...) worker - shape as the other background-thread workers in this file. - """ - progress = pyqtSignal(int) # 0–100 - finished = pyqtSignal(object, str) # AlignmentResult|None, error msg ("" = success) - - def __init__(self, sras: SrasFile, ref_angle_idx: int, dc_threshold_mv: float): - super().__init__() - self._sras = sras - self._ref = ref_angle_idx - self._threshold = dc_threshold_mv - - def run(self): - try: - result = _compute_angle_alignment( - self._sras, self._ref, self._threshold, - progress_cb=lambda pct: self.progress.emit(pct)) - self.finished.emit(result, "") - except Exception as exc: - self.finished.emit(None, str(exc)) +# (mode_str, status-bar unit, colorbar label) per channel index +_CHANNEL_DISPLAY = { + CH1_IDX: ("RF", "Peak frequency (MHz)", "MHz"), + CH3_IDX: ("DC", "DC mean (mV)", "mV"), + CH4_IDX: ("DC", "DC mean (mV)", "mV"), + VELOCITY_MODE_IDX: ("Velocity", "Velocity (m/s)", "m/s"), +} + +_CSS_HINT = "font-size: 11px; color: #aaa;" +_CSS_INFO = "font-size: 11px;" +_CSS_MUTED = "color: #888; font-size: 11px;" +_CSS_WARN = "color: #e07000; font-size: 11px;" +_CSS_BUSY = "color: #4a90d9; font-size: 11px;" # --------------------------------------------------------------------------- @@ -1348,13 +97,11 @@ class RoiQuad: self._pts = np.asarray(pts, dtype=np.float64).reshape(4, 2).copy() @classmethod - def from_bbox(cls, x0: float, y0: float, - x1: float, y1: float) -> "RoiQuad": + def from_bbox(cls, x0: float, y0: float, x1: float, y1: float) -> "RoiQuad": """Create an axis-aligned rectangle from two opposite corners.""" lx, rx = min(x0, x1), max(x0, x1) by, ty = min(y0, y1), max(y0, y1) - pts = np.array([[lx, by], [rx, by], [rx, ty], [lx, ty]]) - return cls(pts) + return cls(np.array([[lx, by], [rx, by], [rx, ty], [lx, ty]])) def copy(self) -> "RoiQuad": return RoiQuad(self._pts.copy()) @@ -1364,7 +111,6 @@ class RoiQuad: return self._pts.copy() def centroid(self) -> np.ndarray: - """Mean of the four corners.""" return self._pts.mean(axis=0) def bbox_size(self) -> np.ndarray: @@ -1378,12 +124,29 @@ class RoiQuad: y_axis: np.ndarray) -> np.ndarray: """Boolean mask (n_rows, n_frames) of pixels whose centres lie inside the quadrilateral. + + Only the quad's axis-aligned bounding box is tested — meshgrid and + contains_points over the *whole* grid would be tens of millions of + point-in-polygon tests (and hundreds of MB of float64 temporaries) + on a large scan, on every ROI edit. """ - X, Y = np.meshgrid(np.asarray(x_axis, dtype=np.float64), - np.asarray(y_axis, dtype=np.float64)) - points = np.column_stack([X.ravel(), Y.ravel()]) - inside = MplPath(self._pts).contains_points(points) - return inside.reshape(X.shape) + x = np.asarray(x_axis, dtype=np.float64) + y = np.asarray(y_axis, dtype=np.float64) + mask = np.zeros((y.size, x.size), dtype=bool) + + (x0, y0), (x1, y1) = self._pts.min(axis=0), self._pts.max(axis=0) + cols = np.nonzero((x >= x0) & (x <= x1))[0] + rows = np.nonzero((y >= y0) & (y <= y1))[0] + if cols.size == 0 or rows.size == 0: + return mask + + c0, c1 = int(cols[0]), int(cols[-1]) + 1 + r0, r1 = int(rows[0]), int(rows[-1]) + 1 + X, Y = np.meshgrid(x[c0:c1], y[r0:r1]) + inside = MplPath(self._pts).contains_points( + np.column_stack([X.ravel(), Y.ravel()])) + mask[r0:r1, c0:c1] = inside.reshape(X.shape) + return mask # --------------------------------------------------------------------------- @@ -1391,18 +154,18 @@ class RoiQuad: # --------------------------------------------------------------------------- class ImageCanvas(FigureCanvasQTAgg): - pixel_clicked = pyqtSignal(int, int) # row_idx, frame_idx - roi_changed = pyqtSignal() # emitted when ROI is created / edited / cleared - draw_mode_changed = pyqtSignal(bool) # emitted when "draw new ROI" arm toggles + pixel_clicked = pyqtSignal(int, int) # row_idx, frame_idx + roi_changed = pyqtSignal() # ROI created / edited / cleared + draw_mode_changed = pyqtSignal(bool) # "draw new ROI" arm toggled # Interaction state values - _IDLE = "idle" - _DRAW_NEW = "draw_new" - _MOVE = "move" + _IDLE = "idle" + _DRAW_NEW = "draw_new" + _MOVE = "move" _DRAG_CORNER = "drag_corner" # Hit tolerance (display pixels) for handles. - _HANDLE_PX = 12 + _HANDLE_PX = 12 _CLICK_THRESH_PX = 4 # releases within this of press count as a click def __init__(self, parent=None): @@ -1411,7 +174,7 @@ class ImageCanvas(FigureCanvasQTAgg): super().__init__(fig) self.setParent(parent) self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding) - self._extent = None + self._extent = None self._img_shape = None # ROI state @@ -1421,16 +184,16 @@ class ImageCanvas(FigureCanvasQTAgg): self._draw_mode = False # Per-interaction snapshots / anchors - self._press_xy : tuple[float, float] | None = None - self._press_pixel : tuple[float, float] | None = None - self._press_button = None - self._snapshot : RoiQuad | None = None - self._drag_corner_idx: int = -1 - self._move_anchor = None # press-point in world coords - self._draw_previous : RoiQuad | None = None + self._press_xy: tuple[float, float] | None = None + self._press_pixel: tuple[float, float] | None = None + self._press_button = None + self._snapshot: RoiQuad | None = None + self._drag_corner_idx: int = -1 + self._move_anchor = None # press-point in world coords + self._draw_previous: RoiQuad | None = None - self.mpl_connect("button_press_event", self._on_press) - self.mpl_connect("motion_notify_event", self._on_motion) + self.mpl_connect("button_press_event", self._on_press) + self.mpl_connect("motion_notify_event", self._on_motion) self.mpl_connect("button_release_event", self._on_release) # ------------------------------------------------------------------ @@ -1445,7 +208,7 @@ class ImageCanvas(FigureCanvasQTAgg): # Patches and lines are destroyed by figure.clf(); drop stale refs. self._roi_artists = [] - self._extent = extent + self._extent = extent self._img_shape = img.shape im = self.ax.imshow( @@ -1512,27 +275,23 @@ class ImageCanvas(FigureCanvasQTAgg): return corners = self._roi.corners() - # Filled quad outline - poly = Polygon(corners, closed=True, fill=True, - facecolor="#ffd93a", edgecolor="#e53935", - alpha=0.22, linewidth=2.0, zorder=10) - self.ax.add_patch(poly) - self._roi_artists.append(poly) + # Filled quad, then a sharp unfilled edge for visibility over bright + # images, then draggable corner handles. + for kwargs in ( + dict(fill=True, facecolor="#ffd93a", edgecolor="#e53935", + alpha=0.22, linewidth=2.0, zorder=10), + dict(fill=False, edgecolor="#e53935", linewidth=1.8, zorder=11), + ): + patch = Polygon(corners, closed=True, **kwargs) + self.ax.add_patch(patch) + self._roi_artists.append(patch) - # Sharp edge (no fill) for better visibility over bright images - edge = Polygon(corners, closed=True, fill=False, - edgecolor="#e53935", linewidth=1.8, zorder=11) - self.ax.add_patch(edge) - self._roi_artists.append(edge) - - # Corner handles (white fill, red edge) — drag each independently - handles = self.ax.scatter(corners[:, 0], corners[:, 1], - s=60, c="white", edgecolors="#e53935", - linewidths=1.6, zorder=13) - self._roi_artists.append(handles) + self._roi_artists.append(self.ax.scatter( + corners[:, 0], corners[:, 1], s=60, c="white", + edgecolors="#e53935", linewidths=1.6, zorder=13)) # ------------------------------------------------------------------ - # Hit testing (uses display pixels for handles, data coords for "inside") + # Hit testing (display pixels for handles, data coords for "inside") # ------------------------------------------------------------------ def _hit_test(self, event) -> tuple[str, int | None] | None: @@ -1540,9 +299,8 @@ class ImageCanvas(FigureCanvasQTAgg): return None if event.x is None or event.y is None: return None - corners = self._roi.corners() - corners_disp = self.ax.transData.transform(corners) - click = np.array([event.x, event.y]) + corners_disp = self.ax.transData.transform(self._roi.corners()) + click = np.array([event.x, event.y]) for i in range(4): if np.hypot(*(corners_disp[i] - click)) <= self._HANDLE_PX: @@ -1568,8 +326,8 @@ class ImageCanvas(FigureCanvasQTAgg): if tb is not None and getattr(tb, "mode", ""): return - self._press_xy = (event.xdata, event.ydata) - self._press_pixel = (event.x, event.y) + self._press_xy = (event.xdata, event.ydata) + self._press_pixel = (event.x, event.y) self._press_button = event.button if self._draw_mode: @@ -1588,11 +346,10 @@ class ImageCanvas(FigureCanvasQTAgg): kind, idx = hit self._snapshot = self._roi.copy() - if kind == "corner": self._state = self._DRAG_CORNER self._drag_corner_idx = idx - elif kind == "inside": + else: self._state = self._MOVE self._move_anchor = (event.xdata, event.ydata) @@ -1606,14 +363,11 @@ class ImageCanvas(FigureCanvasQTAgg): if self._state == self._DRAW_NEW: x0, y0 = self._press_xy - x1, y1 = event.xdata, event.ydata - self._roi = RoiQuad.from_bbox(x0, y0, x1, y1) - + self._roi = RoiQuad.from_bbox(x0, y0, event.xdata, event.ydata) elif self._state == self._MOVE: - dx = event.xdata - self._move_anchor[0] - dy = event.ydata - self._move_anchor[1] - self._roi._pts = self._snapshot.corners() + np.array([dx, dy]) - + delta = np.array([event.xdata - self._move_anchor[0], + event.ydata - self._move_anchor[1]]) + self._roi._pts = self._snapshot.corners() + delta elif self._state == self._DRAG_CORNER: self._roi._pts[self._drag_corner_idx] = [event.xdata, event.ydata] @@ -1625,59 +379,59 @@ class ImageCanvas(FigureCanvasQTAgg): return prev_state = self._state self._state = self._IDLE - - if prev_state == self._DRAW_NEW: - # Reject zero-area or vanishingly-small quads - if self._extent is not None: - x0, x1, y_bot, y_top = self._extent - # Minimum: 1% of each axis range - min_w = abs(x1 - x0) * 0.01 - min_h = abs(y_bot - y_top) * 0.01 + try: + if prev_state == self._DRAW_NEW: + self._finish_draw() + elif prev_state in (self._MOVE, self._DRAG_CORNER): + self._draw_roi() + self.draw_idle() + self.roi_changed.emit() else: - min_w = min_h = 1e-6 - if self._roi is not None: - bbox = self._roi.bbox_size() - too_small = bbox[0] < min_w or bbox[1] < min_h - else: - too_small = True - if too_small: - self._roi = self._draw_previous - self._draw_previous = None - self.cancel_drawing() - self._draw_roi() - self.draw_idle() - self.roi_changed.emit() + self._maybe_emit_pixel_click(event) + finally: self._press_xy = self._press_pixel = None self._press_button = None + + def _finish_draw(self): + """Commit (or reject) a freshly-dragged quad.""" + if self._extent is not None: + x0, x1, y_bot, y_top = self._extent + min_w = abs(x1 - x0) * 0.01 # minimum: 1% of each axis range + min_h = abs(y_bot - y_top) * 0.01 + else: + min_w = min_h = 1e-6 + + if self._roi is None: + too_small = True + else: + bbox = self._roi.bbox_size() + too_small = bbox[0] < min_w or bbox[1] < min_h + if too_small: + self._roi = self._draw_previous + + self._draw_previous = None + self.cancel_drawing() + self._draw_roi() + self.draw_idle() + self.roi_changed.emit() + + def _maybe_emit_pixel_click(self, event): + """A release close enough to its press counts as a pixel click.""" + if (self._press_pixel is None or event.x is None or event.y is None + or self._extent is None or event.inaxes is not self.ax + or event.xdata is None): + return + dx_px = event.x - self._press_pixel[0] + dy_px = event.y - self._press_pixel[1] + if dx_px * dx_px + dy_px * dy_px > self._CLICK_THRESH_PX ** 2: return - if prev_state in (self._MOVE, self._DRAG_CORNER): - self._draw_roi() - self.draw_idle() - self.roi_changed.emit() - self._press_xy = self._press_pixel = None - self._press_button = None - return - - # IDLE → treat as pixel click if release is close to press - if (self._press_pixel is not None and event.x is not None and - event.y is not None and self._extent is not None): - dx_px = event.x - self._press_pixel[0] - dy_px = event.y - self._press_pixel[1] - if (dx_px * dx_px + dy_px * dy_px - <= self._CLICK_THRESH_PX * self._CLICK_THRESH_PX - and event.inaxes is self.ax - and event.xdata is not None): - x0, x1, y_bot, y_top = self._extent - n_rows, n_frames = self._img_shape - col = int((event.xdata - x0) / (x1 - x0) * n_frames) - row = int((event.ydata - y_top) / (y_bot - y_top) * n_rows) - col = max(0, min(col, n_frames - 1)) - row = max(0, min(row, n_rows - 1)) - self.pixel_clicked.emit(row, col) - - self._press_xy = self._press_pixel = None - self._press_button = None + x0, x1, y_bot, y_top = self._extent + n_rows, n_frames = self._img_shape + col = int((event.xdata - x0) / (x1 - x0) * n_frames) + row = int((event.ydata - y_top) / (y_bot - y_top) * n_rows) + self.pixel_clicked.emit(max(0, min(row, n_rows - 1)), + max(0, min(col, n_frames - 1))) class WaveformCanvas(FigureCanvasQTAgg): @@ -1699,12 +453,12 @@ class WaveformCanvas(FigureCanvasQTAgg): on the subtracted signal. The unsubtracted FFT is also shown faintly for comparison. """ - data = sras.data[angle_idx] + data = sras.data[angle_idx] waveform = data[row_idx, CH1_IDX, frame_idx, :].astype(np.float32) - t_ns = sras.time_axis_ns() - f_mhz = sras.freq_axis_mhz() - dc3_val = data[row_idx, CH3_IDX, frame_idx, :].astype(np.float32).mean() - dc4_val = data[row_idx, CH4_IDX, frame_idx, :].astype(np.float32).mean() + t_ns = sras.time_axis_ns() + f_mhz = sras.freq_axis_mhz() + dc3_val = data[row_idx, CH3_IDX, frame_idx, :].astype(np.float32).mean() + dc4_val = data[row_idx, CH4_IDX, frame_idx, :].astype(np.float32).mean() bg = sras.background if (apply_bg_sub and sras.background is not None) else None waveform_plot = waveform - bg if bg is not None else waveform @@ -1726,19 +480,19 @@ class WaveformCanvas(FigureCanvasQTAgg): self.ax_wave.set_xlabel("Time (ns)") self.ax_wave.set_ylabel("ADC counts") bg_tag = " [bg sub]" if bg is not None else "" + dc3_mv = adc_to_mv(dc3_val, *sras.cal(CH3_IDX)) + dc4_mv = adc_to_mv(dc4_val, *sras.cal(CH4_IDX)) self.ax_wave.set_title( f"CH1 RF row={row_idx} frame={frame_idx}{bg_tag}\n" f"CH3={dc3_val:.1f} CH4={dc4_val:.1f} " - f"({adc_to_mv(dc3_val, sras.ch_ymult_mv[CH3_IDX], sras.ch_yoff_adc[CH3_IDX], sras.ch_yzero_mv[CH3_IDX]):.2f} / " - f"{adc_to_mv(dc4_val, sras.ch_ymult_mv[CH4_IDX], sras.ch_yoff_adc[CH4_IDX], sras.ch_yzero_mv[CH4_IDX]):.2f} mV)", + f"({dc3_mv:.2f} / {dc4_mv:.2f} mV)", fontsize=8, ) # FFT of the (possibly subtracted) waveform power_sub = np.abs(np.fft.rfft(waveform_plot)) ** 2 power_sub[0] = 0.0 - peak_idx = int(np.argmax(power_sub)) - peak_mhz = f_mhz[peak_idx] + peak_mhz = f_mhz[int(np.argmax(power_sub))] if bg is not None: # Also show the unsubtracted FFT for reference @@ -1763,29 +517,24 @@ class WaveformCanvas(FigureCanvasQTAgg): row_idx: int, frame_idx: int): """CH3 or CH4 DC: time-domain + mean annotation.""" waveform = sras.data[angle_idx][row_idx, ch_idx, frame_idx, :].astype(np.float32) - t_ns = sras.time_axis_ns() mean_val = float(waveform.mean()) - mean_mv = adc_to_mv(mean_val, sras.ch_ymult_mv[ch_idx], sras.ch_yoff_adc[ch_idx], - sras.ch_yzero_mv[ch_idx]) + mean_mv = adc_to_mv(mean_val, *sras.cal(ch_idx)) self.ax_wave.cla() self.ax_right.cla() - self.ax_wave.plot(t_ns, waveform, linewidth=0.7, color="#4488cc") + self.ax_wave.plot(sras.time_axis_ns(), waveform, linewidth=0.7, color="#4488cc") self.ax_wave.axhline(mean_val, color="tomato", linestyle="--", linewidth=1.2, label=f"mean = {mean_val:.2f} ADC") self.ax_wave.set_xlabel("Time (ns)") self.ax_wave.set_ylabel("ADC counts") self.ax_wave.set_title( - f"{CH_NAMES[ch_idx]} DC row={row_idx} frame={frame_idx}" - ) + f"{CH_NAMES[ch_idx]} DC row={row_idx} frame={frame_idx}") self.ax_wave.legend(fontsize=8) self.ax_right.text( 0.5, 0.5, - f"DC mode\n\n" - f"mean = {mean_val:.3f} ADC\n" - f" = {mean_mv:.3f} mV", + f"DC mode\n\nmean = {mean_val:.3f} ADC\n = {mean_mv:.3f} mV", ha="center", va="center", transform=self.ax_right.transAxes, fontsize=11, ) @@ -1819,8 +568,8 @@ class FftOptionsDialog(QDialog): self.setMinimumWidth(380) self._samples_per_frame = samples_per_frame - self._sample_rate_hz = sample_rate_hz - self._grating_um = grating_um + self._sample_rate_hz = sample_rate_hz + self._grating_um = grating_um layout = QVBoxLayout(self) @@ -1828,18 +577,17 @@ class FftOptionsDialog(QDialog): grp_backend = QGroupBox("FFT Backend") bl = QVBoxLayout(grp_backend) - self._btn_numpy = QRadioButton( - "NumPy FFT (always available)") + self._btn_numpy = QRadioButton("NumPy FFT (always available)") self._btn_pyfftw = QRadioButton( - "pyFFTW (faster for large arrays)" if _pyfftw_available + "pyFFTW (faster for large arrays)" if PYFFTW_AVAILABLE else "pyFFTW (not installed — run: pip install pyfftw)") - self._btn_pyfftw.setEnabled(_pyfftw_available) + self._btn_pyfftw.setEnabled(PYFFTW_AVAILABLE) self._backend_group = QButtonGroup(self) - self._backend_group.addButton(self._btn_numpy, id=0) + self._backend_group.addButton(self._btn_numpy, id=0) self._backend_group.addButton(self._btn_pyfftw, id=1) - if current_backend == "pyfftw" and _pyfftw_available: + if current_backend == "pyfftw" and PYFFTW_AVAILABLE: self._btn_pyfftw.setChecked(True) else: self._btn_numpy.setChecked(True) @@ -1869,30 +617,26 @@ class FftOptionsDialog(QDialog): self._lbl_nfft = QLabel() self._lbl_freq_res = QLabel() - self._lbl_vel_res = QLabel() + self._lbl_vel_res = QLabel() for lbl in (self._lbl_nfft, self._lbl_freq_res, self._lbl_vel_res): - lbl.setStyleSheet("font-size: 11px; color: #aaa;") + lbl.setStyleSheet(_CSS_HINT) zl.addWidget(lbl) layout.addWidget(grp_zp) # ---- Buttons --------------------------------------------------- buttons = QDialogButtonBox() - self._apply_btn = buttons.addButton( - "Apply", QDialogButtonBox.ButtonRole.AcceptRole) - self._cancel_btn = buttons.addButton( - "Cancel", QDialogButtonBox.ButtonRole.RejectRole) - self._apply_btn.clicked.connect(self.accept) - self._cancel_btn.clicked.connect(self.reject) + buttons.addButton("Apply", QDialogButtonBox.ButtonRole.AcceptRole + ).clicked.connect(self.accept) + buttons.addButton("Cancel", QDialogButtonBox.ButtonRole.RejectRole + ).clicked.connect(self.reject) layout.addWidget(buttons) self._update_info() - # ------------------------------------------------------------------ - def _update_info(self): spf = self._samples_per_frame - sr = self._sample_rate_hz + sr = self._sample_rate_hz pad = self._spin_pad.value() if spf is None or sr is None: @@ -1902,13 +646,12 @@ class FftOptionsDialog(QDialog): return n_fft = spf * pad - freq_res_hz = sr / n_fft + freq_res_hz = sr / n_fft freq_res_mhz = freq_res_hz / 1e6 # v (m/s) = freq (MHz) × grating (µm) - vel_res_ms = freq_res_mhz * self._grating_um + vel_res_ms = freq_res_mhz * self._grating_um - self._lbl_nfft.setText( - f"FFT points: {spf} × {pad} = {n_fft:,}") + self._lbl_nfft.setText(f"FFT points: {spf} × {pad} = {n_fft:,}") self._lbl_freq_res.setText( f"Frequency bin: {freq_res_mhz:.4f} MHz ({freq_res_hz / 1e3:.2f} kHz)") self._lbl_vel_res.setText( @@ -1916,7 +659,7 @@ class FftOptionsDialog(QDialog): f"(at grating = {self._grating_um:.2f} µm)") def get_backend(self) -> str: - return "pyfftw" if self._btn_pyfftw.isChecked() and _pyfftw_available else "numpy" + return "pyfftw" if self._btn_pyfftw.isChecked() and PYFFTW_AVAILABLE else "numpy" def get_pad_factor(self) -> int: return max(1, self._spin_pad.value()) @@ -1937,23 +680,21 @@ class SrasViewerWindow(QMainWindow): self._current_image: np.ndarray | None = None self._current_angle: int = 0 self._current_ch: int = 0 - self._load_thread: QThread | None = None - self._compute_thread: QThread | None = None self._pending_angle: int = 0 self._pending_ch: int = 0 self._pending_bg_sub: bool = True self._pending_threshold: float = 50.0 # mV - self._progress_dlg: QProgressDialog | None = None + self._pending_fft_pad_factor: int = 1 + + # Live background jobs, keyed by role — see _run_worker. + self._jobs: dict[str, tuple] = {} + self._progress_dlgs: dict[str, QProgressDialog] = {} # FFT settings (configured via FFT Options dialog) self._fft_pad_factor: int = 1 # 1 = no padding - self._pending_fft_pad_factor: int = 1 # Convert menu: batch DC/FFT compute-and-store (v6 -> v7) - self._batch_thread: QThread | None = None - self._batch_worker: BatchCacheWorker | None = None self._batch_errors: list[str] = [] - self._batch_progress_dlg: QProgressDialog | None = None # Display-only settings (colormap, grating) never trigger a # recompute — they're applied to cached data on redraw. DC images @@ -1966,14 +707,11 @@ class SrasViewerWindow(QMainWindow): # same combination is free. self._dc_cache: dict[tuple[int, int], np.ndarray] = {} self._fft_cache: dict[tuple[int, bool, int | None, float], np.ndarray] = {} - self._dc_precompute_thread: QThread | None = None self._dc_precompute_worker: DcPrecomputeWorker | None = None self._dc_generation: int = 0 # Angle alignment ("Fusion" menu) - self._alignment_result: AlignmentResult | None = None - self._alignment_thread: QThread | None = None - self._alignment_worker: AngleAlignmentWorker | None = None + self._alignment_result = None self._alignment_generation: int = 0 self._aligned_cache: dict[tuple, np.ndarray] = {} @@ -1982,6 +720,57 @@ class SrasViewerWindow(QMainWindow): if initial_path: self._load_file(initial_path) + # ------------------------------------------------------------------ + # Background job plumbing + # ------------------------------------------------------------------ + + def _run_worker(self, key: str, worker: QObject, *, + connect: tuple = (), quit_on: tuple = ("finished",), + on_done=None) -> bool: + """Move *worker* onto its own QThread and start it. Returns False if + a job under *key* is already running. + + Centralises two lifetime hazards that each cost a process abort: + + 1. The job is claimed in self._jobs *before* start() and before + anything below that can pump the Qt event loop (a + QProgressDialog.show() does on first display). If it weren't, a + re-entrant editingFinished could slip past the busy check, start a + second thread, and then have the first call's own assignment + clobber — and destroy while still running — that second QThread. + + 2. thread.finished fires as the thread winds down but does not + guarantee the OS thread has joined. Dropping the last reference to + a QThread whose thread is still running logs "QThread: Destroyed + while thread is still running" and aborts, so wait() first. + """ + if key in self._jobs: + return False + + thread = QThread() + self._jobs[key] = (thread, worker, on_done) # claim before anything pumps + worker.moveToThread(thread) + thread.started.connect(worker.run) + for signal_name, slot in connect: + getattr(worker, signal_name).connect(slot) + for signal_name in quit_on: + getattr(worker, signal_name).connect(thread.quit) + thread.finished.connect(lambda k=key: self._on_job_finished(k)) + thread.start() + return True + + def _on_job_finished(self, key: str): + job = self._jobs.pop(key, None) + if job is None: + return + thread, _worker, on_done = job + thread.wait() # join before releasing our last reference + if on_done is not None: + on_done() + + def _job_running(self, key: str) -> bool: + return key in self._jobs + # ------------------------------------------------------------------ # UI construction # ------------------------------------------------------------------ @@ -1993,27 +782,33 @@ class SrasViewerWindow(QMainWindow): root.setContentsMargins(8, 8, 8, 8) root.setSpacing(8) - # ---- Left control panel ---------------------------------------- + root.addWidget(self._build_left_panel()) + root.addWidget(self._build_canvases(), stretch=1) + root.addWidget(self._build_right_panel()) + + self.statusBar().showMessage("Open an .sras file to begin.") + self._build_menus() + + def _build_left_panel(self) -> QWidget: panel = QWidget() panel.setFixedWidth(260) panel_layout = QVBoxLayout(panel) panel_layout.setContentsMargins(0, 0, 0, 0) panel_layout.setSpacing(6) - root.addWidget(panel) - # File + # ---- File ------------------------------------------------------- grp_file = QGroupBox("File") fl = QVBoxLayout(grp_file) self.btn_open = QPushButton("Open .sras…") self.btn_open.clicked.connect(self._on_open) self.lbl_filename = QLabel("No file loaded") self.lbl_filename.setWordWrap(True) - self.lbl_filename.setStyleSheet("color: #888; font-size: 11px;") + self.lbl_filename.setStyleSheet(_CSS_MUTED) fl.addWidget(self.btn_open) fl.addWidget(self.lbl_filename) panel_layout.addWidget(grp_file) - # Scan info + # ---- Scan info -------------------------------------------------- grp_info = QGroupBox("Scan Info") il = QVBoxLayout(grp_info) self._info = {} @@ -2021,26 +816,25 @@ class SrasViewerWindow(QMainWindow): "Sample rate", "X start", "Pixel Δx", "Laser freq"): lbl = QLabel(f"{key}: —") lbl.setWordWrap(True) - lbl.setStyleSheet("font-size: 11px;") + lbl.setStyleSheet(_CSS_INFO) il.addWidget(lbl) self._info[key] = lbl - # Frame count warning (hidden until needed) - self.lbl_frame_warn = QLabel("") + + self.lbl_frame_warn = QLabel("") # frame-count / format notes self.lbl_frame_warn.setWordWrap(True) - self.lbl_frame_warn.setStyleSheet("color: #e07000; font-size: 11px;") + self.lbl_frame_warn.setStyleSheet(_CSS_WARN) il.addWidget(self.lbl_frame_warn) - # Background DC-precompute progress (hidden until a file is loaded) - self.lbl_dc_precompute = QLabel("") + + self.lbl_dc_precompute = QLabel("") # background DC-precompute progress self.lbl_dc_precompute.setWordWrap(True) - self.lbl_dc_precompute.setStyleSheet("color: #4a90d9; font-size: 11px;") + self.lbl_dc_precompute.setStyleSheet(_CSS_BUSY) il.addWidget(self.lbl_dc_precompute) panel_layout.addWidget(grp_info) - # View settings + # ---- View settings ---------------------------------------------- grp_view = QGroupBox("View Settings") vl = QVBoxLayout(grp_view) - # Angle ar = QHBoxLayout() ar.addWidget(QLabel("Angle:")) self.spin_angle = QSpinBox() @@ -2052,7 +846,6 @@ class SrasViewerWindow(QMainWindow): ar.addWidget(self.lbl_angle_deg) vl.addLayout(ar) - # Channel cr = QHBoxLayout() cr.addWidget(QLabel("Channel:")) self.combo_channel = QComboBox() @@ -2062,12 +855,12 @@ class SrasViewerWindow(QMainWindow): cr.addWidget(self.combo_channel) vl.addLayout(cr) - # DC threshold (for RF / CH1 masking) sep = QFrame() sep.setFrameShape(QFrame.Shape.HLine) sep.setStyleSheet("color: #555;") vl.addWidget(sep) + # DC threshold (for RF / CH1 masking) self.grp_threshold = QGroupBox("RF Mask Threshold (CH1 only)") tl = QVBoxLayout(self.grp_threshold) thr_row = QHBoxLayout() @@ -2082,8 +875,8 @@ class SrasViewerWindow(QMainWindow): self.spin_threshold_mv.editingFinished.connect(self._on_threshold_changed) thr_row.addWidget(self.spin_threshold_mv) tl.addLayout(thr_row) - self.lbl_threshold_adc = QLabel(f"≈ {mv_to_adc(50.0):.1f} ADC counts") # updated on file load - self.lbl_threshold_adc.setStyleSheet("font-size: 11px; color: #888;") + self.lbl_threshold_adc = QLabel(f"≈ {mv_to_adc(50.0):.1f} ADC counts") + self.lbl_threshold_adc.setStyleSheet(_CSS_MUTED) tl.addWidget(self.lbl_threshold_adc) vl.addWidget(self.grp_threshold) @@ -2110,39 +903,16 @@ class SrasViewerWindow(QMainWindow): self.chk_aligned_view.toggled.connect(self._on_aligned_view_toggled) vl.addWidget(self.chk_aligned_view) - # Velocity settings (visible only in velocity mode) - self.grp_velocity = QGroupBox("Velocity Settings (CH1 only)") - vel_l = QVBoxLayout(self.grp_velocity) - grat_row = QHBoxLayout() - grat_row.addWidget(QLabel("Grating size:")) - self.spin_grating_um = QDoubleSpinBox() - self.spin_grating_um.setRange(0.1, 1000.0) - self.spin_grating_um.setDecimals(2) - self.spin_grating_um.setSingleStep(0.5) - self.spin_grating_um.setSuffix(" µm") - self.spin_grating_um.setValue(25) - self.spin_grating_um.setEnabled(False) - self.spin_grating_um.editingFinished.connect(self._on_grating_changed) - grat_row.addWidget(self.spin_grating_um) - vel_l.addLayout(grat_row) - self.lbl_velocity_formula = QLabel("v (m/s) = freq (MHz) × grating (µm)") - self.lbl_velocity_formula.setStyleSheet("font-size: 10px; color: #888;") - vel_l.addWidget(self.lbl_velocity_formula) - self.grp_velocity.setVisible(False) - # (grp_velocity will be added to the right panel below) - - # Export self.btn_export_csv = QPushButton("Export Image as CSV…") self.btn_export_csv.setEnabled(False) self.btn_export_csv.setToolTip( - "Save the current CH1 image (one scan row per CSV line)." - ) + "Save the current CH1 image (one scan row per CSV line).") self.btn_export_csv.clicked.connect(self._on_export_csv) vl.addWidget(self.btn_export_csv) panel_layout.addWidget(grp_view) - # ---- ROI (Region of Interest) --------------------------------- + # ---- ROI --------------------------------------------------------- grp_roi = QGroupBox("ROI (Region of Interest)") rl = QVBoxLayout(grp_roi) @@ -2152,8 +922,7 @@ class SrasViewerWindow(QMainWindow): self.btn_draw_roi.setToolTip( "Arm next click+drag on the image to draw a new ROI\n" "(replaces any existing one). Click again to cancel.\n" - "After drawing, drag inside to move, grab corners to resize,\n" - "or use the handle above the top edge to rotate.\n" + "After drawing, drag inside to move, or grab corners to reshape.\n" "The ROI is persistent across channels / modes / angles." ) self.btn_draw_roi.toggled.connect(self._on_draw_roi_toggled) @@ -2175,16 +944,66 @@ class SrasViewerWindow(QMainWindow): rl.addWidget(self.btn_export_roi) self.lbl_roi_center = QLabel("centroid: —") - self.lbl_roi_size = QLabel("bbox: —") - self.lbl_roi_npix = QLabel("pixels inside: —") + self.lbl_roi_size = QLabel("bbox: —") + self.lbl_roi_npix = QLabel("pixels inside: —") for lbl in (self.lbl_roi_center, self.lbl_roi_size, self.lbl_roi_npix): - lbl.setStyleSheet("font-size: 11px; color: #aaa;") + lbl.setStyleSheet(_CSS_HINT) rl.addWidget(lbl) panel_layout.addWidget(grp_roi) panel_layout.addStretch() + return panel + + def _build_canvases(self) -> QWidget: + splitter = QSplitter(Qt.Orientation.Vertical) + + img_widget = QWidget() + img_vl = QVBoxLayout(img_widget) + img_vl.setContentsMargins(0, 0, 0, 0) + self.image_canvas = ImageCanvas() + self.image_canvas.pixel_clicked.connect(self._on_pixel_clicked) + self.image_canvas.roi_changed.connect(self._update_roi_ui) + self.image_canvas.draw_mode_changed.connect(self._on_draw_mode_changed) + img_vl.addWidget(NavigationToolbar2QT(self.image_canvas, img_widget)) + img_vl.addWidget(self.image_canvas) + splitter.addWidget(img_widget) + + wave_widget = QWidget() + wave_vl = QVBoxLayout(wave_widget) + wave_vl.setContentsMargins(0, 0, 0, 0) + self.lbl_wave_hint = QLabel( + "Click a pixel in the image above to inspect its waveform.") + self.lbl_wave_hint.setAlignment(Qt.AlignmentFlag.AlignCenter) + self.lbl_wave_hint.setStyleSheet(_CSS_MUTED) + self.wave_canvas = WaveformCanvas() + wave_vl.addWidget(self.lbl_wave_hint) + wave_vl.addWidget(self.wave_canvas) + splitter.addWidget(wave_widget) + + splitter.setSizes([580, 250]) + return splitter + + def _build_right_panel(self) -> QWidget: + # Velocity settings (visible only in velocity mode) + self.grp_velocity = QGroupBox("Velocity Settings (CH1 only)") + vel_l = QVBoxLayout(self.grp_velocity) + grat_row = QHBoxLayout() + grat_row.addWidget(QLabel("Grating size:")) + self.spin_grating_um = QDoubleSpinBox() + self.spin_grating_um.setRange(0.1, 1000.0) + self.spin_grating_um.setDecimals(2) + self.spin_grating_um.setSingleStep(0.5) + self.spin_grating_um.setSuffix(" µm") + self.spin_grating_um.setValue(25) + self.spin_grating_um.setEnabled(False) + self.spin_grating_um.editingFinished.connect(self._on_grating_changed) + grat_row.addWidget(self.spin_grating_um) + vel_l.addLayout(grat_row) + lbl_formula = QLabel("v (m/s) = freq (MHz) × grating (µm)") + lbl_formula.setStyleSheet("font-size: 10px; color: #888;") + vel_l.addWidget(lbl_formula) + self.grp_velocity.setVisible(False) - # ---- Display Options group (added to right panel below) ------------ grp_display = QGroupBox("Display Options") dl = QVBoxLayout(grp_display) @@ -2215,54 +1034,19 @@ class SrasViewerWindow(QMainWindow): row.addWidget(spin) dl.addLayout(row) - # ---- Right: image + waveform splitter -------------------------- - splitter = QSplitter(Qt.Orientation.Vertical) - root.addWidget(splitter, stretch=1) - - # Image canvas - img_widget = QWidget() - img_vl = QVBoxLayout(img_widget) - img_vl.setContentsMargins(0, 0, 0, 0) - self.image_canvas = ImageCanvas() - self.image_canvas.pixel_clicked.connect(self._on_pixel_clicked) - self.image_canvas.roi_changed.connect(self._on_roi_changed) - self.image_canvas.draw_mode_changed.connect(self._on_draw_mode_changed) - toolbar = NavigationToolbar2QT(self.image_canvas, img_widget) - img_vl.addWidget(toolbar) - img_vl.addWidget(self.image_canvas) - splitter.addWidget(img_widget) - - # Waveform inspector - wave_widget = QWidget() - wave_vl = QVBoxLayout(wave_widget) - wave_vl.setContentsMargins(0, 0, 0, 0) - self.lbl_wave_hint = QLabel( - "Click a pixel in the image above to inspect its waveform." - ) - self.lbl_wave_hint.setAlignment(Qt.AlignmentFlag.AlignCenter) - self.lbl_wave_hint.setStyleSheet("color: #888; font-size: 11px;") - self.wave_canvas = WaveformCanvas() - wave_vl.addWidget(self.lbl_wave_hint) - wave_vl.addWidget(self.wave_canvas) - splitter.addWidget(wave_widget) - - splitter.setSizes([580, 250]) - - # ---- Right control panel ------------------------------------------- right_panel = QWidget() right_panel.setFixedWidth(270) - right_panel_layout = QVBoxLayout(right_panel) - right_panel_layout.setContentsMargins(0, 0, 0, 0) - right_panel_layout.setSpacing(6) - right_panel_layout.addWidget(self.grp_velocity) - right_panel_layout.addWidget(grp_display) - right_panel_layout.addStretch() - root.addWidget(right_panel) + layout = QVBoxLayout(right_panel) + layout.setContentsMargins(0, 0, 0, 0) + layout.setSpacing(6) + layout.addWidget(self.grp_velocity) + layout.addWidget(grp_display) + layout.addStretch() + return right_panel - self.statusBar().showMessage("Open an .sras file to begin.") - - # ---- Menu bar ---------------------------------------------------------- + def _build_menus(self): menubar = self.menuBar() + fft_menu = menubar.addMenu("&FFT") fft_act = QAction("FFT &Options…", self) fft_act.setStatusTip("Configure FFT backend and zero-padding") @@ -2311,43 +1095,26 @@ class SrasViewerWindow(QMainWindow): def _on_open(self): path, _ = QFileDialog.getOpenFileName( - self, "Open SRAS File", "", "SRAS Files (*.sras);;All Files (*)" - ) + self, "Open SRAS File", "", "SRAS Files (*.sras);;All Files (*)") if path: self._load_file(path) def _load_file(self, path: str): - if self._load_thread is not None: - return - # Claim self._load_thread before any call below that can pump the - # Qt event loop — see the comment in _start_compute for why. - self._load_worker = LoadWorker(path) - self._load_thread = QThread() - self._load_worker.moveToThread(self._load_thread) - self._load_thread.started.connect(self._load_worker.run) - self._load_worker.finished.connect(self._on_load_done) - self._load_worker.error.connect( - lambda msg: self.statusBar().showMessage(f"Error: {msg}") + started = self._run_worker( + "load", LoadWorker(path), + connect=( + ("finished", self._on_load_done), + ("error", lambda msg: self.statusBar().showMessage(f"Error: {msg}")), + ), ) - self._load_worker.finished.connect(self._load_thread.quit) - self._load_thread.finished.connect(self._on_load_thread_finished) - + if not started: + return self.btn_open.setEnabled(False) self.statusBar().showMessage(f"Loading {Path(path).name}…") - self._show_progress(f"Loading {Path(path).name}…") - - self._load_thread.start() - - def _on_load_thread_finished(self): - # See the comment in _on_compute_thread_finished: wait() before - # releasing our reference to avoid destroying a QThread whose OS - # thread hasn't fully joined yet. - if self._load_thread is not None: - self._load_thread.wait() - self._load_thread = None + self._show_progress("main", f"Loading {Path(path).name}…") def _on_load_done(self, sras): - self._close_progress() + self._close_progress("main") self.btn_open.setEnabled(True) if sras is None: return @@ -2355,15 +1122,14 @@ class SrasViewerWindow(QMainWindow): self._current_image = None # Caches (and any in-flight DC precompute) belong to the previous - # file's geometry — discard and start fresh. + # file's geometry — discard and start fresh. Bumping the generation + # counters makes any still-running worker's result get dropped when + # it lands. self._dc_cache = {} self._fft_cache = {} + self._dc_generation += 1 self.lbl_dc_precompute.setText("") - # Same for any alignment result — belongs to the previous file's - # geometry. Bump the generation counter so a still-running - # alignment worker's result is discarded when it lands (see - # _on_alignment_done). self._alignment_result = None self._aligned_cache = {} self._alignment_generation += 1 @@ -2391,13 +1157,8 @@ class SrasViewerWindow(QMainWindow): self.combo_channel.blockSignals(False) self._update_controls_enabled(True) - # Refresh the ADC-count label now that we have file calibration - self._on_threshold_changed() + self._on_threshold_changed() # refresh ADC label with file calibration self._on_view_changed() - - # Precompute DC images for every angle in the background so - # switching angles while viewing a DC channel is instant, and the - # FFT masking step rarely has to wait on a DC4 image either. self._start_dc_precompute() # ------------------------------------------------------------------ @@ -2408,44 +1169,41 @@ class SrasViewerWindow(QMainWindow): s = self._sras if s is None: return - angle_idx = self.spin_angle.value() - self._info["Angles"].setText(f"Angles: {s.n_angles}") - self._info["Rows"].setText(f"Rows: {s.n_rows[angle_idx]}") - self._info["Frames / row"].setText(f"Frames / row: {s.n_frames[angle_idx]}") - self._info["Samples / frame"].setText(f"Samples / frame: {s.samples_per_frame}") - self._info["Sample rate"].setText(f"Sample rate: {s.sample_rate_hz/1e9:.4g} GS/s") - self._info["X start"].setText(f"X start: {s.x_start_mm[angle_idx]:.4g} mm") - self._info["Pixel Δx"].setText(f"Pixel Δx: {s.pixel_x_mm*1e3:.3g} µm") - self._info["Laser freq"].setText(f"Laser freq: {s.laser_freq_hz/1e3:.4g} kHz") + a = self.spin_angle.value() + for key, text in ( + ("Angles", f"{s.n_angles}"), + ("Rows", f"{s.n_rows[a]}"), + ("Frames / row", f"{s.n_frames[a]}"), + ("Samples / frame", f"{s.samples_per_frame}"), + ("Sample rate", f"{s.sample_rate_hz / 1e9:.4g} GS/s"), + ("X start", f"{s.x_start_mm[a]:.4g} mm"), + ("Pixel Δx", f"{s.pixel_x_mm * 1e3:.3g} µm"), + ("Laser freq", f"{s.laser_freq_hz / 1e3:.4g} kHz"), + ): + self._info[key].setText(f"{key}: {text}") notes = [] if s.frame_count_mismatch: - notes.append( - f"! Header n_frames={s.n_frames_header}, " - f"actual={s.n_frames[angle_idx]} (scanner bug — corrected)" - ) + notes.append(f"! Header n_frames={s.n_frames_header}, " + f"actual={s.n_frames[a]} (scanner bug — corrected)") if s.scan_aborted: - notes.append( - f"! Scan aborted: {s.n_angles}/{s.n_angles_declared} angles complete" - ) + notes.append(f"! Scan aborted: {s.n_angles}/{s.n_angles_declared} " + "angles complete") if s.background is not None: notes.append(f"Background waveform: {len(s.background)} samples") - if isinstance(s.precomputed_freq_mhz, np.ndarray): - bg_note = " (bg-sub)" if s.precomputed_bg_sub else " (no bg-sub)" - notes.append(f"v5: precomputed images present{bg_note} — display is instant") if s.version in (6, 7): notes.append("v6/v7 format: rows / frames / x_start are per-angle") - if s.version == 7: - n_dc = sum(1 for x in s.precomputed_dc4_mv if x is not None) - n_fft = sum(1 for x in s.precomputed_freq_mhz if x is not None) - if n_dc or n_fft: - bg_note = " (bg-sub)" if s.precomputed_bg_sub else " (no bg-sub)" - notes.append( - f"v7: cached DC ({n_dc}/{s.n_angles} angles), " - f"FFT ({n_fft}/{s.n_angles} angles{bg_note if n_fft else ''}) " - "— display is instant for cached angles") - else: - notes.append("v7 format: no cache blocks stored yet") + + n_dc = sum(1 for x in s.precomputed_dc4_mv if x is not None) + n_fft = sum(1 for x in s.precomputed_freq_mhz if x is not None) + if n_dc or n_fft: + bg_note = " (bg-sub)" if s.precomputed_bg_sub else " (no bg-sub)" + notes.append( + f"Cached images: DC {n_dc}/{s.n_angles} angles, " + f"FFT {n_fft}/{s.n_angles} angles{bg_note if n_fft else ''} " + "— display is instant for cached angles") + elif s.version == 7: + notes.append("v7 format: no cache blocks stored yet") self.lbl_frame_warn.setText("\n".join(notes)) # ------------------------------------------------------------------ @@ -2454,61 +1212,50 @@ class SrasViewerWindow(QMainWindow): def _update_controls_enabled(self, enabled: bool): s = self._sras - self.spin_angle.setEnabled(enabled and s is not None and s.n_angles > 1) + has_file = enabled and s is not None + ch_idx = self.combo_channel.currentIndex() + is_ch1 = enabled and ch_idx in CH1_DERIVED_MODES + is_vel = enabled and ch_idx == VELOCITY_MODE_IDX + + self.spin_angle.setEnabled(has_file and s.n_angles > 1) self.combo_channel.setEnabled(enabled) self.combo_cmap.setEnabled(enabled) self.chk_auto.setEnabled(enabled) manual = enabled and not self.chk_auto.isChecked() self.spin_vmin.setEnabled(manual) self.spin_vmax.setEnabled(manual) - ch_idx = self.combo_channel.currentIndex() - is_ch1 = enabled and ch_idx in CH1_DERIVED_MODES + # Threshold and bg-sub apply to all CH1 modes self.spin_threshold_mv.setEnabled(is_ch1) - has_bg = enabled and s is not None and s.background is not None - self.chk_bg_sub.setEnabled(has_bg and is_ch1) - # Velocity grating spinbox - is_vel = enabled and ch_idx == VELOCITY_MODE_IDX + self.chk_bg_sub.setEnabled(has_file and s.background is not None and is_ch1) self.spin_grating_um.setEnabled(is_vel) self.grp_velocity.setVisible(is_vel) - # CSV export: enabled when a CH1-derived image is displayed + self.btn_export_csv.setEnabled(is_ch1 and self._current_image is not None) # ROI: always usable once a file is loaded (independent of channel) - self.btn_draw_roi.setEnabled(enabled and s is not None) - # Batch Convert actions: pick their own files, independent of - # whatever's currently open — only gated on no batch already running - can_batch = self._batch_thread is None + self.btn_draw_roi.setEnabled(has_file) + + # Batch Convert actions pick their own files, independent of + # whatever's currently open — only gated on no batch already running. + can_batch = not self._job_running("batch") self._batch_dc_act.setEnabled(can_batch) self._batch_fft_act.setEnabled(can_batch) - # Angle alignment: needs >1 angle and no alignment already running - can_align = (enabled and s is not None and s.n_angles > 1 - and self._alignment_thread is None) - self._alignment_act.setEnabled(can_align) - self.chk_aligned_view.setEnabled( - enabled and self._alignment_result is not None) + + self._alignment_act.setEnabled( + has_file and s.n_angles > 1 and not self._job_running("align")) + self.chk_aligned_view.setEnabled(enabled and self._alignment_result is not None) self._update_roi_ui() def _on_channel_changed(self): - ch_idx = self.combo_channel.currentIndex() - has_file = self._sras is not None - is_ch1 = ch_idx in CH1_DERIVED_MODES - self.spin_threshold_mv.setEnabled(is_ch1 and has_file) - has_bg = has_file and self._sras.background is not None - self.chk_bg_sub.setEnabled(has_bg and is_ch1) - is_vel = ch_idx == VELOCITY_MODE_IDX - self.spin_grating_um.setEnabled(is_vel and has_file) - self.grp_velocity.setVisible(is_vel) - self.btn_export_csv.setEnabled(is_ch1 and has_file and self._current_image is not None) + self._update_controls_enabled(self._sras is not None) self._on_view_changed() def _on_bg_sub_toggled(self): # Background subtraction changes the FFT input, so it genuinely # invalidates the cached raw FFT (the cache key includes it) — - # _refresh_display() will recompute only if there's no cache hit - # for the new bg-sub state. - if self._sras is not None: - if self.combo_channel.currentIndex() in CH1_DERIVED_MODES: - self._refresh_display() + # _refresh_display() recomputes only on a miss for the new state. + if self._sras is not None and self.combo_channel.currentIndex() in CH1_DERIVED_MODES: + self._refresh_display() def _on_grating_changed(self): # Grating is a pure post-multiply on the cached frequency image — @@ -2516,27 +1263,119 @@ class SrasViewerWindow(QMainWindow): if self._sras is not None and self.combo_channel.currentIndex() == VELOCITY_MODE_IDX: self._refresh_display() + def _on_threshold_changed(self): + mv = self.spin_threshold_mv.value() + cal = (self._sras.cal(CH4_IDX) if self._sras is not None + else (_FALLBACK_YMULT_MV, _FALLBACK_YOFF_ADC, 0.0)) + self.lbl_threshold_adc.setText(f"≈ {mv_to_adc(mv, *cal):.1f} ADC counts") + # Threshold decides which pixels get an FFT at all, so changing it is + # a genuine cache-key change — but the recompute reuses the cached DC4 + # image to skip masked-out pixels. + if self._sras is not None and self.combo_channel.currentIndex() in CH1_DERIVED_MODES: + self._refresh_display() + + def _on_autoscale_toggled(self, checked: bool): + manual = not checked + self.spin_vmin.setEnabled(manual and self._sras is not None) + self.spin_vmax.setEnabled(manual and self._sras is not None) + if self._sras is not None and self._current_image is not None: + self._redraw_image(self._current_image) + + def _on_manual_range_changed(self): + if not self.chk_auto.isChecked() and self._current_image is not None: + self._redraw_image(self._current_image) + + def _on_cmap_changed(self): + # Colormap is purely how the existing image is rendered. + if self._current_image is not None: + self._redraw_image(self._current_image) + + def _on_view_changed(self): + if self._sras is None: + return + idx = self.spin_angle.value() + self.lbl_angle_deg.setText(f"({self._sras.angles_deg[idx]:.1f}°)") + self._update_scan_info_labels() + self._refresh_display() + + def _on_aligned_view_toggled(self, checked: bool): + if self._current_image is not None: + self._redraw_image(self._current_image) + + # ------------------------------------------------------------------ + # CSV export + # ------------------------------------------------------------------ + def _on_export_csv(self): if self._current_image is None or self._sras is None: return - ch_idx = self._current_ch - angle = self._current_angle - ch_name = CH_NAMES[ch_idx] - default_name = ( - f"{self._sras.path.stem}_angle{angle}_{ch_name}.csv" - ) + default_name = (f"{self._sras.path.stem}_angle{self._current_angle}" + f"_{CH_NAMES[self._current_ch]}.csv") path, _ = QFileDialog.getSaveFileName( self, "Export Image as CSV", str(self._sras.path.parent / default_name), - "CSV files (*.csv);;All files (*)", - ) + "CSV files (*.csv);;All files (*)") if not path: return np.savetxt(path, self._current_image, delimiter=",", fmt="%.6g") self.statusBar().showMessage(f"Exported {Path(path).name}") + def _on_export_roi_csv(self): + if self._current_image is None or self._sras is None: + return + roi = self.image_canvas.get_roi() + if roi is None: + self.statusBar().showMessage("No ROI — draw one first") + return + s = self._sras + x_axis = s.x_axis_mm(self._current_angle) + y_axis = s.y_positions_mm(self._current_angle) + mask = roi.mask_for_grid(x_axis, y_axis) + if not mask.any(): + self.statusBar().showMessage("ROI does not overlap any pixel") + return + img = self._current_image + if img.shape != mask.shape: + self.statusBar().showMessage( + f"ROI shape {mask.shape} does not match image {img.shape}") + return + + X, Y = np.meshgrid(np.asarray(x_axis, dtype=np.float64), + np.asarray(y_axis, dtype=np.float64)) + rows_idx, frames_idx = np.where(mask) + n_pix = int(mask.sum()) + + ch_name = CH_NAMES[self._current_ch] + angle = self._current_angle + default_name = f"{s.path.stem}_angle{angle}_{ch_name}_ROI.csv" + path, _ = QFileDialog.getSaveFileName( + self, "Export ROI as CSV", + str(s.path.parent / default_name), + "CSV files (*.csv);;All files (*)") + if not path: + return + + corners_str = " ".join(f"({p[0]:.6g},{p[1]:.6g})" for p in roi.corners()) + header = ( + f"# ROI quad corners (BL BR TR TL) mm: {corners_str}\n" + f"# source: {s.path.name}, channel={ch_name}, " + f"angle_idx={angle}, angle_deg={s.angles_deg[angle]:.4g}\n" + f"# n_pixels={n_pix}\n" + "row,frame,x_mm,y_mm,value" + ) + data = np.column_stack([ + rows_idx.astype(np.int64), frames_idx.astype(np.int64), + X[mask], Y[mask], img[mask].astype(np.float64), + ]) + # integer columns first, floats after — use a per-column format list + np.savetxt(path, data, delimiter=",", + fmt=["%d", "%d", "%.6g", "%.6g", "%.6g"], + header=header, comments="") + self.statusBar().showMessage( + f"Exported ROI ({n_pix} pixels) to {Path(path).name}") + # ------------------------------------------------------------------ - # ROI (rectangle on the image) + # ROI # ------------------------------------------------------------------ def _on_draw_roi_toggled(self, checked: bool): @@ -2553,8 +1392,9 @@ class SrasViewerWindow(QMainWindow): self.btn_draw_roi.setChecked(active) self.btn_draw_roi.blockSignals(False) - def _on_roi_changed(self): - self._update_roi_ui() + def _on_clear_roi(self): + self.image_canvas.clear_roi() + self.statusBar().showMessage("ROI cleared") def _update_roi_ui(self): roi = self.image_canvas.get_roi() @@ -2565,12 +1405,12 @@ class SrasViewerWindow(QMainWindow): self.btn_clear_roi.setEnabled(False) self.btn_export_roi.setEnabled(False) return - cen = roi.centroid() + + cen = roi.centroid() bbox = roi.bbox_size() - self.lbl_roi_center.setText( - f"centroid: ({cen[0]:.3f}, {cen[1]:.3f}) mm") - self.lbl_roi_size.setText( - f"bbox: {bbox[0]:.3f} × {bbox[1]:.3f} mm") + self.lbl_roi_center.setText(f"centroid: ({cen[0]:.3f}, {cen[1]:.3f}) mm") + self.lbl_roi_size.setText(f"bbox: {bbox[0]:.3f} × {bbox[1]:.3f} mm") + npix = 0 if self._sras is not None: try: @@ -2578,133 +1418,24 @@ class SrasViewerWindow(QMainWindow): # Aligned View is on: _on_export_roi_csv also exports on the # raw grid (never synthetically-resampled pixels), so this # readout must match what Export ROI actually writes. - mask = roi.mask_for_grid(self._sras.x_axis_mm(self._current_angle), - self._sras.y_positions_mm(self._current_angle)) + mask = roi.mask_for_grid( + self._sras.x_axis_mm(self._current_angle), + self._sras.y_positions_mm(self._current_angle)) npix = int(mask.sum()) except Exception: npix = 0 self.lbl_roi_npix.setText(f"pixels inside: {npix}") self.btn_clear_roi.setEnabled(True) - self.btn_export_roi.setEnabled( - self._current_image is not None and npix > 0) - - def _on_clear_roi(self): - self.image_canvas.clear_roi() - self.statusBar().showMessage("ROI cleared") - - def _on_roi_angle_edited(self): - pass # rotation control removed — corners are dragged individually - - def _on_export_roi_csv(self): - if self._current_image is None or self._sras is None: - return - roi = self.image_canvas.get_roi() - if roi is None: - self.statusBar().showMessage("No ROI — draw one first") - return - s = self._sras - x_axis = s.x_axis_mm(self._current_angle) - y_axis = s.y_positions_mm(self._current_angle) - X, Y = np.meshgrid(np.asarray(x_axis, dtype=np.float64), - np.asarray(y_axis, dtype=np.float64)) - mask = roi.mask_for_grid(x_axis, y_axis) - if not mask.any(): - self.statusBar().showMessage("ROI does not overlap any pixel") - return - img = self._current_image - if img.shape != mask.shape: - self.statusBar().showMessage( - f"ROI shape {mask.shape} does not match image {img.shape}") - return - - rows_idx, frames_idx = np.where(mask) - xs = X[mask] - ys = Y[mask] - vals = img[mask] - - ch_idx = self._current_ch - ch_name = CH_NAMES[ch_idx] - angle = self._current_angle - default_name = (f"{s.path.stem}_angle{angle}_{ch_name}_ROI.csv") - path, _ = QFileDialog.getSaveFileName( - self, "Export ROI as CSV", - str(s.path.parent / default_name), - "CSV files (*.csv);;All files (*)", - ) - if not path: - return - - pts = roi.corners() - corners_str = " ".join(f"({p[0]:.6g},{p[1]:.6g})" for p in pts) - header = ( - f"# ROI quad corners (BL BR TR TL) mm: {corners_str}\n" - f"# source: {s.path.name}, channel={ch_name}, " - f"angle_idx={angle}, angle_deg={s.angles_deg[angle]:.4g}\n" - f"# n_pixels={int(mask.sum())}\n" - "row,frame,x_mm,y_mm,value" - ) - data = np.column_stack([ - rows_idx.astype(np.int64), - frames_idx.astype(np.int64), - xs, ys, vals.astype(np.float64), - ]) - # integer columns first, floats after — use a per-column format list - np.savetxt(path, data, delimiter=",", - fmt=["%d", "%d", "%.6g", "%.6g", "%.6g"], - header=header, comments="") - self.statusBar().showMessage( - f"Exported ROI ({int(mask.sum())} pixels) to {Path(path).name}") - - def _on_threshold_changed(self): - mv = self.spin_threshold_mv.value() - if self._sras is not None: - ymult = self._sras.ch_ymult_mv[CH4_IDX] - yoff = self._sras.ch_yoff_adc[CH4_IDX] - yzero = self._sras.ch_yzero_mv[CH4_IDX] - else: - ymult, yoff, yzero = _FALLBACK_YMULT_MV, _FALLBACK_YOFF_ADC, 0.0 - self.lbl_threshold_adc.setText(f"≈ {mv_to_adc(mv, ymult, yoff, yzero):.1f} ADC counts") - # Threshold decides which pixels get an FFT at all (see - # ComputeWorker), so changing it is a genuine cache key change — - # _refresh_display() recomputes only on a miss, and that recompute - # reuses the cached DC4 image to skip masked-out pixels. - if self._sras is not None and self.combo_channel.currentIndex() in CH1_DERIVED_MODES: - self._refresh_display() - - def _on_autoscale_toggled(self, checked: bool): - manual = not checked - self.spin_vmin.setEnabled(manual and self._sras is not None) - self.spin_vmax.setEnabled(manual and self._sras is not None) - if self._sras is not None and self._current_image is not None: - self._redraw_image(self._current_image) - - def _on_manual_range_changed(self): - if not self.chk_auto.isChecked() and self._current_image is not None: - self._redraw_image(self._current_image) - - def _on_cmap_changed(self): - # Colormap is purely how the existing image is rendered — never - # needs a recompute. - if self._current_image is not None: - self._redraw_image(self._current_image) - - def _on_view_changed(self): - if self._sras is None: - return - idx = self.spin_angle.value() - self.lbl_angle_deg.setText(f"({self._sras.angles_deg[idx]:.1f}°)") - self._update_scan_info_labels() - self._refresh_display() + self.btn_export_roi.setEnabled(self._current_image is not None and npix > 0) # ------------------------------------------------------------------ - # Computation + # Display # ------------------------------------------------------------------ def _current_n_fft(self) -> int | None: - pad_factor = self._fft_pad_factor - if pad_factor <= 1 or self._sras is None: + if self._fft_pad_factor <= 1 or self._sras is None: return None - return self._sras.samples_per_frame * pad_factor + return self._sras.samples_per_frame * self._fft_pad_factor def _scale_for_display(self, freq_mhz: np.ndarray, ch_idx: int) -> np.ndarray: """Velocity is a pure post-multiply of the (already DC-masked) @@ -2713,28 +1444,26 @@ class SrasViewerWindow(QMainWindow): return freq_mhz * self.spin_grating_um.value() return freq_mhz - # ------------------------------------------------------------------ - # Aligned View (Fusion) display helpers - # ------------------------------------------------------------------ + def _fft_cache_key(self, angle_idx: int) -> tuple: + return (angle_idx, self.chk_bg_sub.isChecked(), self._current_n_fft(), + self.spin_threshold_mv.value()) def _aligned_cache_key(self, angle_idx: int, ch_idx: int) -> tuple: """Mirrors _fft_cache's key granularity so a stale aligned image is never shown after bg_sub/threshold/pad/grating changes.""" if ch_idx in CH1_DERIVED_MODES: - return (angle_idx, ch_idx, self.chk_bg_sub.isChecked(), - self._current_n_fft(), self.spin_threshold_mv.value(), + return (*self._fft_cache_key(angle_idx), ch_idx, self.spin_grating_um.value() if ch_idx == VELOCITY_MODE_IDX else None) return (angle_idx, ch_idx) def _aligned_canvas_axes(self) -> tuple[np.ndarray, np.ndarray]: r = self._alignment_result n_rows, n_cols = r.canvas_shape - x_axis = r.canvas_origin_mm[0] + np.arange(n_cols) * r.canvas_dx_mm - y_axis = r.canvas_origin_mm[1] + np.arange(n_rows) * r.canvas_dy_mm - return x_axis, y_axis + return (r.canvas_origin_mm[0] + np.arange(n_cols) * r.canvas_dx_mm, + r.canvas_origin_mm[1] + np.arange(n_rows) * r.canvas_dy_mm) def _get_aligned_display_image(self, raw_img: np.ndarray, angle_idx: int, - ch_idx: int) -> np.ndarray: + ch_idx: int) -> np.ndarray: key = self._aligned_cache_key(angle_idx, ch_idx) cached = self._aligned_cache.get(key) if cached is None: @@ -2744,138 +1473,164 @@ class SrasViewerWindow(QMainWindow): def _refresh_display(self): """Show the image for the current angle/channel/threshold, using - cached data whenever possible and only falling back to a - background compute (with progress popup) when genuinely nothing - is cached yet for these settings.""" + cached data whenever possible and only falling back to a background + compute (with progress popup) when genuinely nothing is cached yet.""" if self._sras is None: return angle_idx = self.spin_angle.value() - ch_idx = self.combo_channel.currentIndex() + ch_idx = self.combo_channel.currentIndex() - if ch_idx in (CH3_IDX, CH4_IDX): + if ch_idx in CH1_DERIVED_MODES: + raw = self._fft_cache.get(self._fft_cache_key(angle_idx)) + if raw is not None: + self._show_image_now(self._scale_for_display(raw, ch_idx), + angle_idx, ch_idx) + return + else: cached = self._dc_cache.get((angle_idx, ch_idx)) if cached is not None: self._show_image_now(cached, angle_idx, ch_idx) return - elif ch_idx in (CH1_IDX, VELOCITY_MODE_IDX): - apply_bg_sub = self.chk_bg_sub.isChecked() - threshold_mv = self.spin_threshold_mv.value() - key = (angle_idx, apply_bg_sub, self._current_n_fft(), threshold_mv) - raw = self._fft_cache.get(key) - if raw is not None: - img = self._scale_for_display(raw, ch_idx) - self._show_image_now(img, angle_idx, ch_idx) - return # Nothing cached for these settings — need a real compute. Changing # the DC threshold changes *which* pixels get an FFT at all, so it # can't be satisfied from the cache — but with the DC map already - # known, the recompute skips the FFT for masked-out pixels and is - # much cheaper than a full-image compute would be. + # known, the recompute skips the FFT for masked-out pixels. self._start_compute() def _show_image_now(self, img: np.ndarray, angle_idx: int, ch_idx: int): """Display an already-available image with no compute involved.""" self._current_image = img self._current_angle = angle_idx - self._current_ch = ch_idx + self._current_ch = ch_idx self.btn_export_csv.setEnabled(ch_idx in CH1_DERIVED_MODES) self._redraw_image(img) self._update_roi_ui() + def _redraw_image(self, img: np.ndarray): + s = self._sras + angle_idx = self._current_angle + ch_idx = self._current_ch + + aligned = (self.chk_aligned_view.isChecked() + and self._alignment_result is not None + and angle_idx in self._alignment_result.per_angle) + if aligned: + display_img = self._get_aligned_display_image(img, angle_idx, ch_idx) + x_axis, y_axis = self._aligned_canvas_axes() + else: + display_img = img + x_axis = s.x_axis_mm(angle_idx) + y_axis = s.y_positions_mm(angle_idx) + + dx = x_axis[1] - x_axis[0] if len(x_axis) > 1 else s.pixel_x_mm + dy = float(y_axis[1] - y_axis[0]) if len(y_axis) > 1 else 1.0 + extent = [x_axis[0] - dx / 2, x_axis[-1] + dx / 2, + y_axis[-1] + dy / 2, y_axis[0] - dy / 2] + + if self.chk_auto.isChecked(): + vmin, vmax = float(display_img.min()), float(display_img.max()) + for spin, val in ((self.spin_vmin, vmin), (self.spin_vmax, vmax)): + spin.blockSignals(True) + spin.setValue(val) + spin.blockSignals(False) + else: + vmin, vmax = self.spin_vmin.value(), self.spin_vmax.value() + + angle_deg = s.angles_deg[angle_idx] + mode_str, unit, colorbar_label = _CHANNEL_DISPLAY[ch_idx] + if ch_idx == VELOCITY_MODE_IDX: + ch_label = f"Velocity [grating={self.spin_grating_um.value():.2f} µm]" + else: + ch_label = CH_LABELS[ch_idx] + + title = f"{CH_NAMES[ch_idx]} | {mode_str} | {angle_deg:.1f}°" + if aligned: + title += " [Aligned]" + + self.image_canvas.show_image( + display_img, extent, + cmap=self.combo_cmap.currentText(), + vmin=vmin, vmax=vmax, + xlabel="X (mm)", ylabel="Y (mm)", + title=title, colorbar_label=colorbar_label, + ) + self.statusBar().showMessage( + f"{s.path.name} | {ch_label} @ {angle_deg:.1f}° " + f"| {display_img.shape[1]} × {display_img.shape[0]} px | {unit}" + f"{' | Aligned' if aligned else ''}" + ) + # ------------------------------------------------------------------ # Background compute (only reached on a genuine cache miss) # ------------------------------------------------------------------ def _start_compute(self): - if self._sras is None: + if self._sras is None or self._job_running("compute"): + return # re-checked when the running compute finishes + + angle_idx = self.spin_angle.value() + ch_idx = self.combo_channel.currentIndex() + is_fft = ch_idx in CH1_DERIVED_MODES + + self._pending_angle = angle_idx + self._pending_ch = ch_idx + self._pending_bg_sub = self.chk_bg_sub.isChecked() + self._pending_threshold = self.spin_threshold_mv.value() + self._pending_fft_pad_factor = self._fft_pad_factor + + worker = ComputeWorker( + self._sras, angle_idx, ch_idx, + apply_bg_sub=self._pending_bg_sub, + n_fft=self._current_n_fft(), + dc_threshold_mv=self._pending_threshold, + # Reuse the cached DC4 image (if the precompute has reached this + # angle) so the FFT skips masked-out pixels entirely and doesn't + # need to re-read the CH4 channel from disk. + dc4_mv=self._dc_cache.get((angle_idx, CH4_IDX)), + is_fft_mode=is_fft, + ) + if not self._run_worker( + "compute", worker, + connect=( + ("finished", self._on_compute_done), + ("error", lambda msg: self.statusBar().showMessage( + f"Compute error: {msg}")), + ), + on_done=self._after_compute): return - if self._compute_thread is not None: - return # re-check in _on_compute_thread_finished - angle_idx = self.spin_angle.value() - ch_idx = self.combo_channel.currentIndex() - apply_bg_sub = self.chk_bg_sub.isChecked() - threshold_mv = self.spin_threshold_mv.value() - pad_factor = self._fft_pad_factor - n_fft = self._current_n_fft() - # Reuse the cached DC4 image (if the precompute has reached this - # angle) so the FFT compute skips masked-out pixels entirely and - # doesn't need to re-read the CH4 channel from disk. - dc4_mv = self._dc_cache.get((angle_idx, CH4_IDX)) - - self._pending_angle = angle_idx - self._pending_ch = ch_idx - self._pending_bg_sub = apply_bg_sub - self._pending_threshold = threshold_mv - self._pending_fft_pad_factor = pad_factor - - # Claim self._compute_thread *before* anything below that can pump - # the Qt event loop (e.g. QProgressDialog.show() on first display). - # If that happened first, a re-entrant editingFinished/signal could - # slip past the guard above, start a second thread, and then have - # this call's own assignment clobber (and destroy while still - # running) that second thread's QThread object — which aborts the - # process. Assigning immediately closes that window. - self._compute_worker = ComputeWorker( - self._sras, angle_idx, ch_idx, apply_bg_sub, n_fft=n_fft, - dc_threshold_mv=threshold_mv, dc4_mv=dc4_mv, - ) - self._compute_thread = QThread() - self._compute_worker.moveToThread(self._compute_thread) - self._compute_thread.started.connect(self._compute_worker.run) - self._compute_worker.finished.connect(self._on_compute_done) - self._compute_worker.error.connect( - lambda msg: self.statusBar().showMessage(f"Compute error: {msg}") - ) - self._compute_worker.finished.connect(self._compute_thread.quit) - self._compute_thread.finished.connect(self._on_compute_thread_finished) - - if ch_idx in (CH1_IDX, VELOCITY_MODE_IDX): + if is_fft: self.statusBar().showMessage("Computing FFT…") self._show_progress( + "main", f"Computing FFT for angle {angle_idx}…\n" "This can take a while on a large scan — result is cached " - "so revisiting this angle/mode/threshold will be instant." - ) + "so revisiting this angle/mode/threshold will be instant.") else: self.statusBar().showMessage("Computing DC image…") - self._show_progress(f"Computing DC image for angle {angle_idx}…") + self._show_progress("main", f"Computing DC image for angle {angle_idx}…") - self._compute_thread.start() - - def _on_compute_thread_finished(self): - # Block until the OS thread has actually joined before dropping our - # last reference — deallocating a QThread whose thread hasn't fully - # terminated yet logs "QThread: Destroyed while thread is still - # running" and aborts the process. The finished() signal fires as - # the thread is winding down but does not guarantee it has joined. - if self._compute_thread is not None: - self._compute_thread.wait() - self._compute_thread = None - angle_idx = self.spin_angle.value() - ch_idx = self.combo_channel.currentIndex() - apply_bg_sub = self.chk_bg_sub.isChecked() - threshold_mv = self.spin_threshold_mv.value() - if (angle_idx, ch_idx, apply_bg_sub, threshold_mv, self._fft_pad_factor) != ( + def _after_compute(self): + """If settings changed while the compute was running, re-dispatch + through the cache-aware path — the now-current combination may + already be cached.""" + if (self.spin_angle.value(), self.combo_channel.currentIndex(), + self.chk_bg_sub.isChecked(), self.spin_threshold_mv.value(), + self._fft_pad_factor) != ( self._pending_angle, self._pending_ch, self._pending_bg_sub, self._pending_threshold, self._pending_fft_pad_factor): - # Settings changed while this compute was running — re-dispatch - # through the cache-aware path in case that now-current - # combination happens to already be cached. self._refresh_display() def _on_compute_done(self, result): - self._close_progress() + self._close_progress("main") angle_idx = self._pending_angle - ch_idx = self._pending_ch + ch_idx = self._pending_ch - if ch_idx in (CH1_IDX, VELOCITY_MODE_IDX): - masked_freq_mhz = result - key = (angle_idx, self._pending_bg_sub, self._current_n_fft(), - self._pending_threshold) - self._fft_cache[key] = masked_freq_mhz - img = self._scale_for_display(masked_freq_mhz, ch_idx) + if ch_idx in CH1_DERIVED_MODES: + self._fft_cache[(angle_idx, self._pending_bg_sub, + self._current_n_fft(), self._pending_threshold)] = result + img = self._scale_for_display(result, ch_idx) else: img = result self._dc_cache[(angle_idx, ch_idx)] = img @@ -2889,135 +1644,48 @@ class SrasViewerWindow(QMainWindow): def _start_dc_precompute(self): if self._sras is None: return - self._dc_generation += 1 generation = self._dc_generation - n_angles = self._sras.n_angles + n_angles = self._sras.n_angles worker = DcPrecomputeWorker(self._sras) - thread = QThread() - worker.moveToThread(thread) - thread.started.connect(worker.run) - worker.angle_done.connect( - lambda a, dc3, dc4, g=generation: self._on_dc_precompute_angle_done( - g, a, dc3, dc4, n_angles) - ) - worker.error.connect( - lambda msg: self.statusBar().showMessage(f"DC precompute error: {msg}", 5000) - ) - worker.finished.connect(thread.quit) - worker.error.connect(thread.quit) - thread.finished.connect(lambda: self._on_dc_precompute_thread_finished(thread)) - - # See _start_compute for why the thread/worker are claimed on self - # before .start() rather than after. self._dc_precompute_worker = worker - self._dc_precompute_thread = thread - thread.start() + started = self._run_worker( + "dc_precompute", worker, + connect=( + ("angle_done", lambda a, dc3, dc4, g=generation: + self._on_dc_precompute_angle_done(g, a, dc3, dc4, n_angles)), + ("error", lambda msg: self.statusBar().showMessage( + f"DC precompute error: {msg}", 5000)), + ), + quit_on=("finished", "error"), + on_done=lambda: setattr(self, "_dc_precompute_worker", None), + ) + if not started: + self._dc_precompute_worker = None def _on_dc_precompute_angle_done(self, generation: int, angle_idx: int, - dc3_mv: np.ndarray, dc4_mv: np.ndarray, - n_angles: int): + dc3_mv: np.ndarray, dc4_mv: np.ndarray, + n_angles: int): if generation != self._dc_generation: return # stale result from a previously-loaded file — discard self._dc_cache[(angle_idx, CH3_IDX)] = dc3_mv self._dc_cache[(angle_idx, CH4_IDX)] = dc4_mv + done = sum(1 for a in range(n_angles) if (a, CH4_IDX) in self._dc_cache) - if done < n_angles: - self.lbl_dc_precompute.setText( - f"Precomputing DC images: {done}/{n_angles} angles ready…") - else: - self.lbl_dc_precompute.setText("DC images ready for all angles.") - # If we just finished the angle/channel the user is currently - # looking at and it wasn't shown yet (e.g. they switched here - # before precompute caught up and are still waiting), show it now. - if (self._sras is not None and angle_idx == self.spin_angle.value() - and self._compute_thread is None - and self.combo_channel.currentIndex() in (CH3_IDX, CH4_IDX) - and (self._current_angle != angle_idx - or self._current_ch != self.combo_channel.currentIndex())): + self.lbl_dc_precompute.setText( + f"Precomputing DC images: {done}/{n_angles} angles ready…" + if done < n_angles else "DC images ready for all angles.") + + # If we just finished the angle/channel the user is currently looking + # at and it wasn't shown yet (they switched here before the precompute + # caught up and are still waiting), show it now. + current_ch = self.combo_channel.currentIndex() + if (angle_idx == self.spin_angle.value() + and not self._job_running("compute") + and current_ch in (CH3_IDX, CH4_IDX) + and (self._current_angle != angle_idx or self._current_ch != current_ch)): self._refresh_display() - def _on_dc_precompute_thread_finished(self, thread: QThread): - # See the comment in _on_compute_thread_finished: wait() before - # releasing the reference to avoid destroying a QThread whose OS - # thread hasn't fully joined yet. - thread.wait() - if self._dc_precompute_thread is thread: - self._dc_precompute_thread = None - self._dc_precompute_worker = None - - def _redraw_image(self, img: np.ndarray): - s = self._sras - angle_idx = self._current_angle - ch_idx = self._current_ch - - aligned = (self.chk_aligned_view.isChecked() - and self._alignment_result is not None - and angle_idx in self._alignment_result.per_angle) - if aligned: - display_img = self._get_aligned_display_image(img, angle_idx, ch_idx) - x_axis, y_axis = self._aligned_canvas_axes() - else: - display_img = img - x_axis = s.x_axis_mm(angle_idx) - y_axis = s.y_positions_mm(angle_idx) - - dx = x_axis[1] - x_axis[0] if len(x_axis) > 1 else s.pixel_x_mm - dy = float(y_axis[1] - y_axis[0]) if len(y_axis) > 1 else 1.0 - - extent = [ - x_axis[0] - dx / 2, - x_axis[-1] + dx / 2, - y_axis[-1] + dy / 2, - y_axis[0] - dy / 2, - ] - - if self.chk_auto.isChecked(): - vmin, vmax = float(display_img.min()), float(display_img.max()) - for spin, val in ((self.spin_vmin, vmin), (self.spin_vmax, vmax)): - spin.blockSignals(True) - spin.setValue(val) - spin.blockSignals(False) - else: - vmin = self.spin_vmin.value() - vmax = self.spin_vmax.value() - - angle_deg = s.angles_deg[self._current_angle] - ch_label = CH_LABELS[ch_idx] - - if ch_idx == CH1_IDX: - mode_str = "RF" - unit = "Peak frequency (MHz)" - colorbar_label = "MHz" - elif ch_idx == VELOCITY_MODE_IDX: - grating = self.spin_grating_um.value() - mode_str = "Velocity" - unit = "Velocity (m/s)" - colorbar_label = "m/s" - ch_label = f"Velocity [grating={grating:.2f} µm]" - else: - mode_str = "DC" - unit = "DC mean (mV)" - colorbar_label = "mV" - - title = f"{CH_NAMES[ch_idx]} | {mode_str} | {angle_deg:.1f}°" - if aligned: - title += " [Aligned]" - - self.image_canvas.show_image( - display_img, extent, - cmap=self.combo_cmap.currentText(), - vmin=vmin, vmax=vmax, - xlabel="X (mm)", ylabel="Y (mm)", - title=title, - colorbar_label=colorbar_label, - ) - self.statusBar().showMessage( - f"{s.path.name} | {ch_label} @ {angle_deg:.1f}° " - f"| {display_img.shape[1]} × {display_img.shape[0]} px | {unit}" - f"{' | Aligned' if aligned else ''}" - ) - # ------------------------------------------------------------------ # Pixel inspector # ------------------------------------------------------------------ @@ -3028,110 +1696,97 @@ class SrasViewerWindow(QMainWindow): angle_idx = self._current_angle if (self.chk_aligned_view.isChecked() and self._alignment_result is not None and angle_idx in self._alignment_result.per_angle): - # The click landed on the shared aligned canvas — invert the - # same canvas->raw affine used to display it back to a raw - # (row, frame) index before looking up the waveform. + # The click landed on the shared aligned canvas — invert the same + # canvas->raw affine used to display it back to a raw (row, frame) + # index before looking up the waveform. t = self._alignment_result.per_angle[angle_idx] raw = t.matrix @ np.array([row_idx, frame_idx], dtype=np.float64) + t.offset row_idx, frame_idx = int(round(raw[0])), int(round(raw[1])) - n_rows_a = int(self._sras.n_rows[angle_idx]) - n_frames_a = int(self._sras.n_frames[angle_idx]) + n_rows_a, n_frames_a = self._sras.image_shape(angle_idx) if not (0 <= row_idx < n_rows_a and 0 <= frame_idx < n_frames_a): self.statusBar().showMessage( "No source waveform here (padding region of the aligned canvas).") return + self.lbl_wave_hint.hide() - ch_idx = self._current_ch - if ch_idx in CH1_DERIVED_MODES: + if self._current_ch in CH1_DERIVED_MODES: self.wave_canvas.show_rf_waveform( self._sras, angle_idx, row_idx, frame_idx, - apply_bg_sub=self.chk_bg_sub.isChecked(), - ) + apply_bg_sub=self.chk_bg_sub.isChecked()) else: self.wave_canvas.show_dc_waveform( - self._sras, angle_idx, ch_idx, row_idx, frame_idx - ) + self._sras, angle_idx, self._current_ch, row_idx, frame_idx) # ------------------------------------------------------------------ - # Progress dialog helpers + # Progress dialogs # ------------------------------------------------------------------ - def _show_progress(self, message: str): - if self._progress_dlg is not None: - self._progress_dlg.setLabelText(message) + def _show_progress(self, key: str, message: str, maximum: int = 0): + """Show (or relabel) the progress dialog under *key*. maximum=0 gives + an indeterminate busy indicator.""" + dlg = self._progress_dlgs.get(key) + if dlg is not None: + dlg.setLabelText(message) return - dlg = QProgressDialog(message, "", 0, 0, self) + dlg = QProgressDialog(message, "", 0, maximum, self) dlg.setWindowTitle("Please wait…") dlg.setCancelButton(None) dlg.setWindowModality(Qt.WindowModality.WindowModal) - dlg.setMinimumDuration(300) # only appears if operation takes > 300 ms + dlg.setMinimumDuration(300) # only appears if it takes > 300 ms dlg.show() - self._progress_dlg = dlg + self._progress_dlgs[key] = dlg - def _close_progress(self): - if self._progress_dlg is not None: - self._progress_dlg.close() - self._progress_dlg = None + def _set_progress(self, key: str, pct: int): + dlg = self._progress_dlgs.get(key) + if dlg is not None: + dlg.setValue(pct) + + def _close_progress(self, key: str): + dlg = self._progress_dlgs.pop(key, None) + if dlg is not None: + dlg.close() # ------------------------------------------------------------------ # Convert menu: batch DC/FFT compute-and-store (v6 -> v7) # ------------------------------------------------------------------ def _on_batch_compute(self, mode: str): - if self._batch_thread is not None: + if self._job_running("batch"): return label = "DC" if mode == "dc" else "FFT" paths, _ = QFileDialog.getOpenFileNames( self, f"Select .sras files to batch-compute {label}", "", - "SRAS files (*.sras);;All files (*)", - ) + "SRAS files (*.sras);;All files (*)") if not paths: return - if self._batch_thread is not None: - return # a second trigger snuck in while the file dialog was open - apply_bg = self.chk_bg_sub.isChecked() self._batch_errors = [] - - # Claim self._batch_thread before anything below that can pump the - # Qt event loop — see the comment in _start_compute for why. - self._batch_worker = BatchCacheWorker(paths, mode, apply_bg) - self._batch_thread = QThread() - self._batch_worker.moveToThread(self._batch_thread) - self._batch_thread.started.connect(self._batch_worker.run) - self._batch_worker.progress.connect(self._on_batch_progress) - self._batch_worker.file_done.connect(self._on_batch_file_done) - self._batch_worker.finished.connect(lambda paths=paths: self._on_batch_finished(paths)) - self._batch_worker.finished.connect(self._batch_thread.quit) - self._batch_thread.finished.connect(self._on_batch_thread_finished) + worker = BatchCacheWorker(paths, mode, self.chk_bg_sub.isChecked()) + started = self._run_worker( + "batch", worker, + connect=( + ("progress", lambda pct: self._set_progress("batch", pct)), + ("file_done", self._on_batch_file_done), + ("finished", lambda p=paths: self._on_batch_finished(p)), + ), + on_done=self._after_batch, + ) + if not started: + return # a second trigger snuck in while the file dialog was open self._batch_dc_act.setEnabled(False) self._batch_fft_act.setEnabled(False) - - self._batch_progress_dlg = QProgressDialog( - f"Batch computing {label} for {len(paths)} file(s)…", "", 0, 100, self) - self._batch_progress_dlg.setWindowTitle("Please wait…") - self._batch_progress_dlg.setCancelButton(None) - self._batch_progress_dlg.setWindowModality(Qt.WindowModality.WindowModal) - self._batch_progress_dlg.setMinimumDuration(300) - self._batch_progress_dlg.show() - - self._batch_thread.start() - - def _on_batch_progress(self, pct: int): - if self._batch_progress_dlg is not None: - self._batch_progress_dlg.setValue(pct) + self._show_progress( + "batch", f"Batch computing {label} for {len(paths)} file(s)…", + maximum=100) def _on_batch_file_done(self, path: str, err: str): if err: self._batch_errors.append(f"{Path(path).name} — {err}") - if self._batch_progress_dlg is not None: - self._batch_progress_dlg.setLabelText(f"Processed {Path(path).name}…") + self._show_progress("batch", f"Processed {Path(path).name}…") def _on_batch_finished(self, paths: list[str]): - if self._batch_progress_dlg is not None: - self._batch_progress_dlg.close() - self._batch_progress_dlg = None + self._close_progress("batch") n_total = len(paths) n_failed = len(self._batch_errors) @@ -3144,18 +1799,12 @@ class SrasViewerWindow(QMainWindow): self.statusBar().showMessage(summary) self._batch_errors = [] - # If the currently-open file was in this batch, reload it so the - # GUI picks up the newly-written v7 cache instead of stale state. + # If the currently-open file was in this batch, reload it so the GUI + # picks up the newly-written v7 cache instead of stale state. if self._sras is not None and str(self._sras.path) in paths: self._load_file(str(self._sras.path)) - def _on_batch_thread_finished(self): - # See the comment in _on_compute_thread_finished: wait() before - # releasing our reference to avoid destroying a QThread whose OS - # thread hasn't fully joined yet. - if self._batch_thread is not None: - self._batch_thread.wait() - self._batch_thread = None + def _after_batch(self): self._batch_dc_act.setEnabled(True) self._batch_fft_act.setEnabled(True) @@ -3166,46 +1815,31 @@ class SrasViewerWindow(QMainWindow): def _on_angle_alignment(self): if self._sras is None or self._sras.n_angles <= 1: return - if self._alignment_thread is not None: - return - ref_idx = 0 + ref_idx = 0 threshold_mv = self.spin_threshold_mv.value() - generation = self._alignment_generation + generation = self._alignment_generation - # Claim self._alignment_thread before anything below that can pump - # the Qt event loop — see the comment in _start_compute for why. - self._alignment_worker = AngleAlignmentWorker(self._sras, ref_idx, threshold_mv) - self._alignment_thread = QThread() - self._alignment_worker.moveToThread(self._alignment_thread) - self._alignment_thread.started.connect(self._alignment_worker.run) - self._alignment_worker.progress.connect(self._on_alignment_progress) - self._alignment_worker.finished.connect( - lambda result, err, g=generation: self._on_alignment_done(g, result, err)) - self._alignment_worker.finished.connect(self._alignment_thread.quit) - self._alignment_thread.finished.connect(self._on_alignment_thread_finished) + started = self._run_worker( + "align", AngleAlignmentWorker(self._sras, ref_idx, threshold_mv), + connect=( + ("progress", lambda pct: self._set_progress("main", pct)), + ("finished", lambda result, err, g=generation: + self._on_alignment_done(g, result, err)), + ), + on_done=lambda: self._update_controls_enabled(self._sras is not None), + ) + if not started: + return self._alignment_act.setEnabled(False) self._show_progress( + "main", f"Computing angle alignment ({self._sras.n_angles} angles, " - f"ref=angle 0, CH4 mask ≥ {threshold_mv:.3f} mV)…" - ) - self._alignment_thread.start() - - def _on_alignment_progress(self, pct: int): - if self._progress_dlg is not None: - self._progress_dlg.setValue(pct) - - def _on_alignment_thread_finished(self): - # See the comment in _on_compute_thread_finished: wait() before - # releasing our reference to avoid destroying a QThread whose OS - # thread hasn't fully joined yet. - if self._alignment_thread is not None: - self._alignment_thread.wait() - self._alignment_thread = None - self._update_controls_enabled(self._sras is not None) + f"ref=angle 0, CH4 mask ≥ {threshold_mv:.3f} mV)…", + maximum=100) def _on_alignment_done(self, generation: int, result, error_msg: str): - self._close_progress() + self._close_progress("main") if generation != self._alignment_generation: return # a new file was loaded while this was computing — discard if error_msg: @@ -3223,47 +1857,37 @@ class SrasViewerWindow(QMainWindow): f"canvas {nc}×{nr} px).") self._refresh_display() - def _on_aligned_view_toggled(self, checked: bool): - if self._current_image is not None: - self._redraw_image(self._current_image) - # ------------------------------------------------------------------ # FFT Options # ------------------------------------------------------------------ def _on_fft_options(self): - global _fft_backend - spf = self._sras.samples_per_frame if self._sras is not None else None - sr = self._sras.sample_rate_hz if self._sras is not None else None dlg = FftOptionsDialog( self, - current_backend=_fft_backend, + current_backend=compute.get_fft_backend(), current_pad_factor=self._fft_pad_factor, - samples_per_frame=spf, - sample_rate_hz=sr, + samples_per_frame=self._sras.samples_per_frame if self._sras else None, + sample_rate_hz=self._sras.sample_rate_hz if self._sras else None, grating_um=self.spin_grating_um.value(), ) - if dlg.exec() == QDialog.DialogCode.Accepted: - _fft_backend = dlg.get_backend() - self._fft_pad_factor = dlg.get_pad_factor() - # Pad factor changes the FFT bin count, so it genuinely - # invalidates the cached raw FFT (part of the cache key) — - # _refresh_display() will recompute only on a cache miss. - if self._sras is not None and self.combo_channel.currentIndex() in ( - CH1_IDX, VELOCITY_MODE_IDX): - self._refresh_display() + if dlg.exec() != QDialog.DialogCode.Accepted: + return + compute.set_fft_backend(dlg.get_backend()) + self._fft_pad_factor = dlg.get_pad_factor() + # Pad factor changes the FFT bin count, so it genuinely invalidates + # the cached raw FFT (part of the cache key) — _refresh_display() + # recomputes only on a cache miss. + if self._sras is not None and self.combo_channel.currentIndex() in CH1_DERIVED_MODES: + self._refresh_display() # ------------------------------------------------------------------ def closeEvent(self, event): if self._dc_precompute_worker is not None: self._dc_precompute_worker.stop() - for attr in ("_load_thread", "_compute_thread", "_batch_thread", - "_dc_precompute_thread", "_alignment_thread"): - t = getattr(self, attr, None) - if t is not None: - t.quit() - t.wait(2000) + for thread, _worker, _on_done in list(self._jobs.values()): + thread.quit() + thread.wait(2000) super().closeEvent(event) @@ -3271,8 +1895,8 @@ class SrasViewerWindow(QMainWindow): def main(): app = QApplication(sys.argv) - initial = sys.argv[1] if len(sys.argv) > 1 else None - window = SrasViewerWindow(initial_path=initial) + window = SrasViewerWindow( + initial_path=sys.argv[1] if len(sys.argv) > 1 else None) window.show() sys.exit(app.exec()) diff --git a/sras_workers.py b/sras_workers.py new file mode 100644 index 0000000..1445f9e --- /dev/null +++ b/sras_workers.py @@ -0,0 +1,224 @@ +#!/usr/bin/env python3 +"""Background workers for the SRAS viewer. + +Every worker is a plain QObject moved onto its own QThread by +SrasViewerWindow._run_worker, exposing signals only. Workers must never touch +GUI-thread-owned state (the display caches in particular) — they take +everything they need through their constructor and hand results back by signal. +""" + +import os +from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed + +import numpy as np +from PyQt6.QtCore import QObject, pyqtSignal + +import sras_compute as compute +from sras_compute import ( + cache_file, compute_angle_alignment, compute_rf_image, dc_image_mv, +) +from sras_format import CH3_IDX, CH4_IDX, SrasFile + +# Concurrency caps. Batch conversion runs one process per file, and each of +# those processes threads internally, so the two must be divided rather than +# both set to the core count. Files also commonly sit on one external drive, +# where a dozen concurrent readers is slower than a few — hence the low +# default, overridable from the environment. +_BATCH_MAX_PROCS = int(os.environ.get("SRAS_BATCH_PROCS", 0)) or min( + 4, os.cpu_count() or 2) + + +class LoadWorker(QObject): + finished = pyqtSignal(object) # SrasFile | None + error = pyqtSignal(str) + + def __init__(self, path: str): + super().__init__() + self._path = path + + def run(self): + try: + self.finished.emit(SrasFile(self._path)) + except Exception as exc: + self.error.emit(str(exc)) + self.finished.emit(None) + + +class ComputeWorker(QObject): + """Computes one displayable image for (angle, channel). + + For CH1/Velocity (FFT-derived) channels, the FFT is only run for pixels + whose DC4 (Bias B) mean is at or above dc_threshold_mv — masked pixels are + left at 0 MHz without ever being FFT'd, since that's the expensive part of + a scan. If the DC4 image for this angle is already known, pass it in as + *dc4_mv* to skip re-reading the CH4 channel from disk entirely. + + Emits a plain ``np.ndarray`` already in display units. + """ + finished = pyqtSignal(object) + error = pyqtSignal(str) + + def __init__(self, sras: SrasFile, angle_idx: int, ch_idx: int, + apply_bg_sub: bool = True, n_fft: int | None = None, + dc_threshold_mv: float = 0.0, + dc4_mv: np.ndarray | None = None, + is_fft_mode: bool = False): + super().__init__() + self._sras = sras + self._angle = angle_idx + self._ch = ch_idx + self._apply_bg_sub = apply_bg_sub + self._n_fft = n_fft + self._dc_threshold = dc_threshold_mv + self._dc4_mv = dc4_mv + self._is_fft_mode = is_fft_mode + + def run(self): + try: + if self._is_fft_mode: + img = compute_rf_image( + self._sras, self._angle, dc_threshold_mv=self._dc_threshold, + apply_bg_sub=self._apply_bg_sub, n_fft=self._n_fft, + dc4_mv=self._dc4_mv) + else: + img = dc_image_mv(self._sras, self._angle, self._ch) + self.finished.emit(img) + except Exception as exc: + self.error.emit(str(exc)) + + +class DcPrecomputeWorker(QObject): + """Computes CH3/CH4 DC images for every angle in the background. + + DC images are cheap (a per-waveform mean, no FFT) compared to the + CH1/Velocity FFT, so precomputing them for the whole file right after load + makes switching angles instant while on a DC channel, and also means the + FFT masking step (which needs a DC4 image) rarely has to wait on anything. + + Angles are computed on a thread pool — the work is a pure mean over the + waveform block, so it is I/O- and bandwidth-bound and embarrassingly + parallel. Results are emitted one at a time as they land (out of angle + order), and always from this worker's own thread: nothing emits a Qt + signal from a pool thread. + """ + angle_done = pyqtSignal(int, np.ndarray, np.ndarray) # angle_idx, dc3_mv, dc4_mv + finished = pyqtSignal() + error = pyqtSignal(str) + + def __init__(self, sras: SrasFile): + super().__init__() + self._sras = sras + self._stop = False + + def stop(self): + self._stop = True + + def _one_angle(self, a: int) -> tuple[int, np.ndarray, np.ndarray]: + return a, dc_image_mv(self._sras, a, CH3_IDX), dc_image_mv(self._sras, a, CH4_IDX) + + def run(self): + try: + n = self._sras.n_angles + n_workers = max(1, min(compute._MAX_WORKERS, n)) + with ThreadPoolExecutor(max_workers=n_workers) as pool: + futures = {pool.submit(self._one_angle, a): a for a in range(n)} + try: + for fut in as_completed(futures): + if self._stop: + break + a, dc3, dc4 = fut.result() + self.angle_done.emit(a, dc3, dc4) + finally: + if self._stop: + for fut in futures: + fut.cancel() + self.finished.emit() + except Exception as exc: + self.error.emit(str(exc)) + + +class BatchCacheWorker(QObject): + """Batch-computes and stores DC or FFT images into each of *paths*'s v7 + CACH tail, in place — converting v6 sources to v7 on first use, or updating + an existing v7 file's cache blocks without disturbing whatever the other + block already holds. + + *mode* is ``"dc"`` (CH3/CH4 mean images) or ``"fft"`` (CH1 peak-frequency + images, unmasked — masking is applied at display time, same as v5's PREC + convention). + + Files are processed one per subprocess: they are fully independent, each + opens its own memmap and writes only its own bytes, and only path strings + and scalars cross the process boundary. Emits ``progress(int)`` (0–100 by + files completed), ``file_done(str, str)`` (path, error message or "") so + one file's failure doesn't abort the batch, and ``finished()``. + """ + progress = pyqtSignal(int) + file_done = pyqtSignal(str, str) + finished = pyqtSignal() + + def __init__(self, paths: list[str], mode: str, apply_bg_sub: bool): + super().__init__() + self._paths = paths + self._mode = mode + self._apply_bg_sub = apply_bg_sub + + def _pool(self, n_files: int): + """A process pool, falling back to threads if the platform refuses to + spawn (still a win: the compute releases the GIL).""" + n = max(1, min(_BATCH_MAX_PROCS, n_files)) + try: + return ProcessPoolExecutor(max_workers=n), n + except (OSError, ValueError): + return ThreadPoolExecutor(max_workers=n), n + + def run(self): + paths = self._paths + executor, n_procs = self._pool(len(paths)) + # Each child threads internally; divide the machine rather than + # letting every process claim every core. + per_proc_workers = max(1, (os.cpu_count() or 4) // n_procs) + + done = 0 + with executor: + futures = { + executor.submit(cache_file, p, self._mode, self._apply_bg_sub, + compute.get_fft_backend(), per_proc_workers): p + for p in paths + } + for fut in as_completed(futures): + path = futures[fut] + try: + err = fut.result() + except Exception as exc: + err = str(exc) + done += 1 + self.file_done.emit(path, err) + self.progress.emit(int(done / max(1, len(paths)) * 100)) + + self.finished.emit() + + +class AngleAlignmentWorker(QObject): + """Computes rotation+translation alignment for every angle in *sras*, + referenced to *ref_angle_idx*, from each angle's binarized CH4 mask. + Rotation is analytic (from sras.angles_deg); only translation is found by + phase correlation. + """ + progress = pyqtSignal(int) # 0–100 + finished = pyqtSignal(object, str) # AlignmentResult|None, error ("" = success) + + def __init__(self, sras: SrasFile, ref_angle_idx: int, dc_threshold_mv: float): + super().__init__() + self._sras = sras + self._ref = ref_angle_idx + self._threshold = dc_threshold_mv + + def run(self): + try: + result = compute_angle_alignment( + self._sras, self._ref, self._threshold, + progress_cb=self.progress.emit) + self.finished.emit(result, "") + except Exception as exc: + self.finished.emit(None, str(exc))