diff --git a/core/sras_analysis.py b/core/sras_analysis.py new file mode 100644 index 0000000..2aa9b95 --- /dev/null +++ b/core/sras_analysis.py @@ -0,0 +1,296 @@ +"""SRAS analysis: scope calibration, image reducers, and the SAW +matched-filter pipeline. Qt-free; operates on the per-angle arrays +returned by core.sras_format.SrasFile.load_angle(). + +Channel semantics (fixed by the acquisition app): + CH1 — RF Acoustic Packet: FFT → peak frequency + CH3 — Bias A (DC): waveform mean + CH4 — Bias B (DC): waveform mean — also the RF valid-pixel mask source +""" +from __future__ import annotations + +import os +import re +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass + +import numpy as np +from scipy.signal import butter, hilbert, sosfiltfilt + +# Channel indices into the on-disk channel axis (order fixed by SCAN_CHANNELS) +CH1_IDX, CH3_IDX, CH4_IDX = 0, 1, 2 + +# Fallback scope calibration for preambles missing YMULT/YOFF/YZERO: +# 50 mV/div, 8 div full-scale, int8 ADC, position = -2.72 div +FALLBACK_YMULT_MV = 1.5625 # mV per ADC count +FALLBACK_YOFF_ADC = -87.04 # ADC count that represents 0 V + + +def parse_preamble(preamble: str) -> dict[str, float]: + """Extract YMULT, YOFF, YZERO from a Tektronix WFMOutpre string.""" + 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 + + +@dataclass +class ChannelCalibration: + """Per-channel ADC↔mV conversion, parsed from the file's preambles.""" + ymult_mv: list[float] + yoff_adc: list[float] + yzero_mv: list[float] + + @classmethod + def from_preambles(cls, preambles: list[str]) -> "ChannelCalibration": + cal = cls([], [], []) + for p in preambles: + vals = parse_preamble(p) + # YMULT/YZERO from the scope are in V; stored here in mV + cal.ymult_mv.append(vals.get("YMULT", FALLBACK_YMULT_MV / 1000) * 1000) + cal.yoff_adc.append(vals.get("YOFF", FALLBACK_YOFF_ADC)) + cal.yzero_mv.append(vals.get("YZERO", 0.0) * 1000) + return cal + + def adc_to_mv(self, adc, ch: int): + return (adc - self.yoff_adc[ch]) * self.ymult_mv[ch] + self.yzero_mv[ch] + + def mv_to_adc(self, mv, ch: int): + return (mv - self.yzero_mv[ch]) / self.ymult_mv[ch] + self.yoff_adc[ch] + + +def power_spectrum(waveform: np.ndarray) -> np.ndarray: + """FFT power with the DC bin suppressed.""" + power = np.abs(np.fft.rfft(waveform)) ** 2 + power[..., 0] = 0.0 + return power + + +# --------------------------------------------------------------------------- +# Image reducers — all take the (n_rows, n_ch, n_frames, spf) angle view +# --------------------------------------------------------------------------- + +def compute_dc_image(angle_view: np.ndarray, ch_idx: int) -> np.ndarray: + """Mean of each waveform → (n_rows, n_frames) float32. + + Computed directly on the int8 view — no float32 copy of the block. + """ + return angle_view[:, ch_idx].mean(axis=-1, dtype=np.float32) + + +def _valid_ch1_waveforms(angle_view: np.ndarray, calib: ChannelCalibration, + dc_threshold_mv: float, + background: np.ndarray | None, + ) -> tuple[np.ndarray, np.ndarray]: + """CH4-DC mask + float32 CH1 waveforms for only the pixels that pass. + + Materializes float32 for the valid pixels alone (fancy-index on the int8 + view first), so a mostly-masked angle costs almost nothing. + """ + dc4_mv = calib.adc_to_mv(compute_dc_image(angle_view, CH4_IDX), CH4_IDX) + valid = dc4_mv >= dc_threshold_mv + if not valid.any(): + return valid, np.empty((0, angle_view.shape[-1]), dtype=np.float32) + waves = angle_view[:, CH1_IDX][valid].astype(np.float32) # (n_valid, spf) + if background is not None: + waves -= background + return valid, waves + + +def compute_rf_image(angle_view: np.ndarray, calib: ChannelCalibration, + freq_axis_mhz: np.ndarray, + dc_threshold_mv: float, + background: np.ndarray | None = None, + gate_start_ns: float | None = None, + gate_end_ns: float | None = None, + time_axis_ns: np.ndarray | None = None) -> np.ndarray: + """FFT of each CH1 waveform; pixel = peak frequency in MHz. + + Pixels whose CH4 DC mean (in mV) is below dc_threshold_mv are 0 and the + FFT is skipped for them. Optional time gate zeroes samples outside + [gate_start_ns, gate_end_ns] before the FFT. + """ + valid, waves = _valid_ch1_waveforms(angle_view, calib, dc_threshold_mv, background) + img = np.zeros(valid.shape, dtype=np.float32) + if len(waves): + if (gate_start_ns is not None or gate_end_ns is not None) and time_axis_ns is not None: + keep = np.ones(len(time_axis_ns), dtype=bool) + if gate_start_ns is not None: + keep &= time_axis_ns >= gate_start_ns + if gate_end_ns is not None: + keep &= time_axis_ns <= gate_end_ns + waves[:, ~keep] = 0.0 + peak_bins = np.argmax(power_spectrum(waves), axis=-1) + img[valid] = freq_axis_mhz[peak_bins] + return img + + +def compute_saw_image(angle_view: np.ndarray, calib: ChannelCalibration, + dc_threshold_mv: float, + pipeline: "SawPipeline", mode: str, + background: np.ndarray | None = None) -> np.ndarray: + """Matched-filter pipeline over every valid pixel. + + mode : "amplitude" → MF envelope peak in the SAW window + "tof" → arrival time (ns) of that peak + Only the requested scalar is kept per pixel — the per-shot intermediate + arrays are dropped inside the worker instead of being accumulated. + """ + valid, waves = _valid_ch1_waveforms(angle_view, calib, dc_threshold_mv, background) + img = np.zeros(valid.shape, dtype=np.float32) + if len(waves): + key = "peak_amplitude" if mode == "amplitude" else "peak_time_ns" + + def _scalar(w): + return pipeline.process_shot_metrics(w)[key] + + n_workers = min(os.cpu_count() or 4, len(waves)) + with ThreadPoolExecutor(max_workers=n_workers) as executor: + img[valid] = np.fromiter(executor.map(_scalar, waves), + dtype=np.float32, count=len(waves)) + return img + + +# --------------------------------------------------------------------------- +# SAW signal processing pipeline +# --------------------------------------------------------------------------- + +class SawPipeline: + """EMI-cleaning and SAW extraction pipeline. + + Stages (each independently bypassable): + 1. EMI gate — cosine-taper the first `emi_gate_ns` ns to suppress the + laser-firing burst at t≈0; leaves the SAW packet alone. + 2. Bandpass — 6th-order Butterworth zero-phase (sosfiltfilt). + 3. Matched filter — FFT cross-correlation with a Hann-windowed template + built from the average of N clean shots. + 4. Analytic — Hilbert transform of MF output → amplitude envelope. + """ + + def __init__(self, sample_rate_hz: float, + emi_gate_ns: float = 50.0, + bp_lo_mhz: float = 85.0, + bp_hi_mhz: float = 200.0, + saw_window_ns: tuple[float, float] = (80.0, 350.0)): + self.sample_rate_hz = float(sample_rate_hz) + self.emi_gate_ns = float(emi_gate_ns) + self.bp_lo_mhz = float(bp_lo_mhz) + self.bp_hi_mhz = float(bp_hi_mhz) + self.saw_window_ns = (float(saw_window_ns[0]), float(saw_window_ns[1])) + self.template: np.ndarray | None = None + self._template_fft: dict[int, np.ndarray] = {} # nfft → rfft(template) + self._emi_gate_samples = max(1, int(round( + self.emi_gate_ns * 1e-9 * self.sample_rate_hz))) + nyq = self.sample_rate_hz / 2.0 + lo = np.clip(self.bp_lo_mhz * 1e6 / nyq, 1e-6, 0.999) + hi = np.clip(self.bp_hi_mhz * 1e6 / nyq, lo + 1e-6, 0.9999) + # 6th-order Butterworth → 12th-order bandpass; ~120 dB/decade rolloff + self._sos = butter(6, [lo, hi], btype='bandpass', output='sos') + + def gate_emi(self, signal: np.ndarray) -> np.ndarray: + """Cosine-taper (raised cosine 0→1) the first emi_gate samples. + + The taper rolls up smoothly from zero so the abrupt EMI burst is + suppressed without introducing a step discontinuity at the gate edge. + """ + n = min(self._emi_gate_samples, len(signal)) + out = signal.copy() + out[:n] *= 0.5 * (1.0 - np.cos(np.pi * np.arange(n) / n)) + return out + + def bandpass(self, signal: np.ndarray) -> np.ndarray: + """Zero-phase IIR Butterworth bandpass (sosfiltfilt), float32 in/out.""" + return sosfiltfilt(self._sos, signal).astype(np.float32, copy=False) + + def build_template(self, waveforms: np.ndarray) -> None: + """Average N shots (EMI-gated + bandpassed), Hann-windowed to the + declared SAW window, to form the matched-filter template.""" + processed = np.stack([ + self.bandpass(self.gate_emi(np.asarray(w, dtype=np.float32))) + for w in waveforms + ]) + avg = processed.mean(axis=0) + + n = len(avg) + t_ns = np.arange(n) / self.sample_rate_hz * 1e9 + i0 = max(0, int(np.searchsorted(t_ns, self.saw_window_ns[0]))) + i1 = min(n, int(np.searchsorted(t_ns, self.saw_window_ns[1]))) + windowed = np.zeros(n, dtype=np.float32) + if i1 > i0: + windowed[i0:i1] = avg[i0:i1] * np.hanning(i1 - i0) + self.template = windowed + self._template_fft.clear() + + def matched_filter(self, signal: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """FFT cross-correlation with the template → (mf_output, envelope).""" + if self.template is None: + raise RuntimeError("No template — call build_template() first") + n = len(signal) + nfft = 1 << (n + len(self.template) - 1).bit_length() + T = self._template_fft.get(nfft) + if T is None: + T = np.conj(np.fft.rfft(self.template, nfft)) + self._template_fft[nfft] = T + S = np.fft.rfft(signal, nfft) + mf = np.fft.irfft(S * T, nfft)[:n] + env = np.abs(hilbert(mf)) + return mf.astype(np.float32, copy=False), env.astype(np.float32, copy=False) + + def process_shot(self, signal: np.ndarray) -> dict: + """EMI gate → bandpass → matched filter on one shot; returns every + stage plus metrics (for diagnostics displays).""" + raw = np.asarray(signal, dtype=np.float32) + gated = self.gate_emi(raw) + filtered = self.bandpass(gated) + + if self.template is not None: + mf_out, env = self.matched_filter(filtered) + else: + mf_out = filtered.copy() + env = np.abs(hilbert(filtered)).astype(np.float32) + + metrics = self._envelope_metrics(env, filtered) + return { + "raw": raw, + "gated": gated, + "filtered": filtered, + "mf_output": mf_out, + "envelope": env, + "sample_rate_hz": self.sample_rate_hz, + **metrics, + } + + def process_shot_metrics(self, signal: np.ndarray) -> dict: + """Like process_shot but returns only the scalar metrics — used for + whole-image sweeps where retaining per-shot arrays would multiply + memory by the pixel count.""" + filtered = self.bandpass(self.gate_emi(np.asarray(signal, dtype=np.float32))) + if self.template is not None: + _, env = self.matched_filter(filtered) + else: + env = np.abs(hilbert(filtered)) + return self._envelope_metrics(env, filtered) + + def _envelope_metrics(self, env: np.ndarray, filtered: np.ndarray) -> dict: + sr = self.sample_rate_hz + t_ns = np.arange(len(env)) / sr * 1e9 + s0, s1 = self.saw_window_ns + roi = (t_ns >= s0) & (t_ns <= s1) + if roi.any(): + peak_sample = int(np.where(roi)[0][np.argmax(env[roi])]) + else: + peak_sample = int(np.argmax(env)) + peak_amplitude = float(env[peak_sample]) + peak_time_ns = float(peak_sample / sr * 1e9) + + # SNR: peak / RMS of the noise floor inside the gated EMI region + noise_seg = filtered[:self._emi_gate_samples] + noise_rms = float(np.sqrt(np.mean(noise_seg ** 2))) if len(noise_seg) else 1.0 + return { + "peak_amplitude": peak_amplitude, + "peak_sample": peak_sample, + "peak_time_ns": peak_time_ns, + "snr": peak_amplitude / noise_rms if noise_rms > 0 else 0.0, + } diff --git a/sras_viewer.py b/sras_viewer.py index 7c6954f..a16fb6c 100755 --- a/sras_viewer.py +++ b/sras_viewer.py @@ -1,249 +1,42 @@ #!/usr/bin/env python3 """ SRAS Scan File Viewer -PyQt6 application for visualizing channel data from .sras binary scan files. +PyQt6 application for visualizing channel data from v6 .sras scan files. 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 + CH1 — RF Acoustic Packet: FFT → peak frequency + CH3 — Bias A (DC): waveform mean + CH4 — Bias B (DC): waveform mean 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 data is memory-mapped (core.sras_format.SrasFile), so opening a +multi-GB file costs only the pages actually touched by the current view. +Incomplete (aborted/resumed) files are handled via the format's frontier +analysis: only the rows actually on disk are shown. """ -import re import sys -import struct -import numpy as np from pathlib import Path -from concurrent.futures import ThreadPoolExecutor -import os -from scipy.signal import butter, sosfiltfilt, decimate as sp_decimate, hilbert + +import numpy as np from PyQt6.QtWidgets import ( QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, QGroupBox, QLabel, QPushButton, QComboBox, QSpinBox, QDoubleSpinBox, QFileDialog, QSizePolicy, QSplitter, QCheckBox, QFrame, QProgressDialog, ) -from PyQt6.QtCore import Qt, QThread, pyqtSignal, QObject +from PyQt6.QtCore import Qt, QThread, QTimer, pyqtSignal, QObject from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg, NavigationToolbar2QT from matplotlib.figure import Figure -# --------------------------------------------------------------------------- -# SAW signal processing pipeline -# --------------------------------------------------------------------------- +sys.path.insert(0, str(Path(__file__).resolve().parent)) -class SawPipeline: - """Modular EMI-cleaning and SAW extraction pipeline. - - Stages (each independently bypassable): - 1. EMI gate — cosine-taper first `emi_gate_ns` ns to suppress the - laser-firing burst at t≈0; leaves SAW packet untouched. - 2. Bandpass — 6th-order Butterworth zero-phase (sosfiltfilt), default - 85–200 MHz. Matches hardware bandpass already applied. - 3. Decimate — optional; reduces to ~781 MS/s (factor-8) before the - matched filter without losing SAW information. - 4. Matched filter — FFT cross-correlation with a Hann-windowed template - built from the average of N clean shots. - 5. Analytic — Hilbert transform of MF output → amplitude envelope and - instantaneous phase. - - Typical usage:: - - pipe = SawPipeline(sras.sample_rate_hz) - waveforms = sras.data[angle, :, CH1_IDX, :, :].reshape(-1, spf) - pipe.build_template(waveforms[:50]) - result = pipe.process_shot(waveform) - # result["envelope"], result["peak_amplitude"], result["peak_time_ns"] - """ - - DECIMATE_FACTOR = 8 # 6250 MS/s → 781.25 MS/s (~4× SAW BW of 200 MHz) - - def __init__(self, sample_rate_hz: float, - emi_gate_ns: float = 50.0, - bp_lo_mhz: float = 85.0, - bp_hi_mhz: float = 200.0, - saw_window_ns: tuple[float, float] = (80.0, 350.0), - decimate_enable: bool = False): - self.sample_rate_hz = float(sample_rate_hz) - self.emi_gate_ns = float(emi_gate_ns) - self.bp_lo_mhz = float(bp_lo_mhz) - self.bp_hi_mhz = float(bp_hi_mhz) - self.saw_window_ns = (float(saw_window_ns[0]), float(saw_window_ns[1])) - self.decimate_enable = decimate_enable - self.template: np.ndarray | None = None - self._emi_gate_samples: int = 0 - self._sos = None - self._effective_sr = self.sample_rate_hz - self._build_filter() - - # ------------------------------------------------------------------ - # Setup - # ------------------------------------------------------------------ - - def _build_filter(self): - self._emi_gate_samples = max(1, int(round( - self.emi_gate_ns * 1e-9 * self.sample_rate_hz))) - nyq = self.sample_rate_hz / 2.0 - lo = np.clip(self.bp_lo_mhz * 1e6 / nyq, 1e-6, 0.999) - hi = np.clip(self.bp_hi_mhz * 1e6 / nyq, lo + 1e-6, 0.9999) - # 6th-order Butterworth → 12th-order bandpass; ~120 dB/decade rolloff - self._sos = butter(6, [lo, hi], btype='bandpass', output='sos') - self._effective_sr = (self.sample_rate_hz / self.DECIMATE_FACTOR - if self.decimate_enable else self.sample_rate_hz) - - # ------------------------------------------------------------------ - # Individual stages - # ------------------------------------------------------------------ - - def gate_emi(self, signal: np.ndarray) -> np.ndarray: - """Cosine-taper (raised cosine 0→1) the first `emi_gate_samples` samples. - - The taper rolls up smoothly from zero so the abrupt EMI burst is - suppressed without introducing a step discontinuity at the gate edge. - """ - n = min(self._emi_gate_samples, len(signal)) - out = signal.copy() - out[:n] *= 0.5 * (1.0 - np.cos(np.pi * np.arange(n) / n)) - return out - - def bandpass(self, signal: np.ndarray) -> np.ndarray: - """Zero-phase IIR Butterworth bandpass (sosfiltfilt).""" - return sosfiltfilt(self._sos, signal.astype(np.float64)).astype(np.float32) - - def decimate_signal(self, signal: np.ndarray) -> np.ndarray: - """Decimate by DECIMATE_FACTOR with scipy anti-alias filter.""" - return sp_decimate(signal.astype(np.float64), self.DECIMATE_FACTOR, - zero_phase=True).astype(np.float32) - - # ------------------------------------------------------------------ - # Template construction - # ------------------------------------------------------------------ - - def build_template(self, waveforms: np.ndarray) -> None: - """Build Hann-windowed average template. - - Parameters - ---------- - waveforms : ndarray, shape (N, n_samples) - Raw or pre-processed CH1 waveforms. EMI gating + bandpass are - applied here before averaging so the template is clean. - """ - processed = np.stack([ - self.bandpass(self.gate_emi(w.astype(np.float32))) - for w in waveforms - ]) - avg = processed.mean(axis=0) - - # Hann window restricted to the declared SAW window region - n = len(avg) - t_ns = np.arange(n) / self.sample_rate_hz * 1e9 - i0 = max(0, int(np.searchsorted(t_ns, self.saw_window_ns[0]))) - i1 = min(n, int(np.searchsorted(t_ns, self.saw_window_ns[1]))) - windowed = np.zeros(n, dtype=np.float32) - win_len = i1 - i0 - if win_len > 0: - windowed[i0:i1] = avg[i0:i1] * np.hanning(win_len) - self.template = windowed - - # ------------------------------------------------------------------ - # Matched filter - # ------------------------------------------------------------------ - - def matched_filter(self, signal: np.ndarray) -> tuple[np.ndarray, np.ndarray]: - """FFT cross-correlation with template. - - Returns - ------- - mf_output : float32 ndarray, length = len(signal) - envelope : float32 ndarray, Hilbert amplitude envelope of mf_output - """ - if self.template is None: - raise RuntimeError("No template — call build_template() first") - n = len(signal) - nfft = 1 << (n + len(self.template) - 1).bit_length() - S = np.fft.rfft(signal.astype(np.float64), nfft) - T = np.fft.rfft(self.template.astype(np.float64), nfft) - mf = np.fft.irfft(S * np.conj(T), nfft)[:n] - env = np.abs(hilbert(mf)) - return mf.astype(np.float32), env.astype(np.float32) - - # ------------------------------------------------------------------ - # Full pipeline for a single waveform - # ------------------------------------------------------------------ - - def process_shot(self, signal: np.ndarray) -> dict: - """Run EMI gate → bandpass → (decimate) → matched filter on one shot. - - Returns a dict with keys: - raw, gated, filtered, [decimated], mf_output, envelope, - peak_amplitude (float), peak_sample (int), peak_time_ns (float), - snr (float), sample_rate_hz (float). - """ - raw = signal.astype(np.float32) - gated = self.gate_emi(raw) - filtered = self.bandpass(gated) - - if self.decimate_enable: - proc = self.decimate_signal(filtered) - sr = self._effective_sr - else: - proc = filtered - sr = self.sample_rate_hz - - if self.template is not None: - mf_out, env = self.matched_filter(proc) - else: - mf_out = proc.copy() - env = np.abs(hilbert(proc)).astype(np.float32) - - t_ns = np.arange(len(env)) / sr * 1e9 - - # Peak within SAW window - s0, s1 = self.saw_window_ns - roi = (t_ns >= s0) & (t_ns <= s1) - if roi.any(): - idx_in_roi = np.argmax(env[roi]) - peak_sample = int(np.where(roi)[0][idx_in_roi]) - else: - peak_sample = int(np.argmax(env)) - peak_amplitude = float(env[peak_sample]) - peak_time_ns = float(peak_sample / sr * 1e9) - - # SNR: peak / RMS of noise floor in the gated EMI region (after bandpass) - noise_seg = filtered[:self._emi_gate_samples] - noise_rms = float(np.sqrt(np.mean(noise_seg ** 2))) if len(noise_seg) > 0 else 1.0 - snr = peak_amplitude / noise_rms if noise_rms > 0 else 0.0 - - return { - "raw": raw, - "gated": gated, - "filtered": filtered, - "mf_output": mf_out, - "envelope": env, - "peak_amplitude": peak_amplitude, - "peak_sample": peak_sample, - "peak_time_ns": peak_time_ns, - "snr": snr, - "sample_rate_hz": sr, - } - - -# --------------------------------------------------------------------------- -# SRAS format -# --------------------------------------------------------------------------- - -HDR_FMT = ">4sBHHffffIIdBB" -HDR_SIZE = struct.calcsize(HDR_FMT) # 43 bytes - -# Fixed-order channels in the file: index 0=CH1, 1=CH3, 2=CH4 -# Fixed channel indices into the .sras data array (CH1=RF, CH3/CH4=Bias DC) -CH1_IDX, CH3_IDX, CH4_IDX = 0, 1, 2 +from core.sras_format import SrasFile +from core.sras_analysis import ( + CH1_IDX, CH3_IDX, CH4_IDX, ChannelCalibration, SawPipeline, + compute_dc_image, compute_rf_image, compute_saw_image, power_spectrum, +) CH_LABELS = [ "CH1 — RF (FFT peak freq)", @@ -253,359 +46,102 @@ CH_LABELS = [ "CH1 — SAW Amplitude (matched filter)", "CH1 — SAW Arrival time (matched filter)", ] -CH_NAMES = ["CH1", "CH3", "CH4", "VEL", "SAW-AMP", "SAW-TOF"] +CH_NAMES = ["CH1", "CH3", "CH4", "VEL", "SAW-AMP", "SAW-TOF"] # Combo indices for derived modes (all use CH1_IDX data) -VELOCITY_MODE_IDX = 3 -SAW_MODE_AMP_IDX = 4 -SAW_MODE_TOF_IDX = 5 -SAW_MODES = (SAW_MODE_AMP_IDX, SAW_MODE_TOF_IDX) -# All modes that operate on CH1 waveforms -CH1_DERIVED_MODES = (CH1_IDX, VELOCITY_MODE_IDX) + SAW_MODES +VELOCITY_MODE_IDX = 3 +SAW_MODE_AMP_IDX = 4 +SAW_MODE_TOF_IDX = 5 +SAW_MODES = (SAW_MODE_AMP_IDX, SAW_MODE_TOF_IDX) +CH1_DERIVED_MODES = (CH1_IDX, VELOCITY_MODE_IDX) + SAW_MODES -# 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 +# Per-mode display strings: (mode_str, unit, colorbar_label) +MODE_INFO = { + CH1_IDX: ("RF", "Peak frequency (MHz)", "MHz"), + VELOCITY_MODE_IDX: ("Velocity", "Velocity (m/s)", "m/s"), + SAW_MODE_AMP_IDX: ("SAW-AMP", "MF envelope peak (arb.)", "amplitude"), + SAW_MODE_TOF_IDX: ("SAW-TOF", "SAW arrival time (ns)", "ns"), + CH3_IDX: ("DC", "DC mean (mV)", "mV"), + CH4_IDX: ("DC", "DC mean (mV)", "mV"), +} 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 +RECOMPUTE_DEBOUNCE_MS = 250 -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/v3/v4 .sras file.""" +class LoadedScan: + """A parsed scan plus everything the viewer derives from it once.""" def __init__(self, path: str): - self.path = Path(path) - self._parse() + self.sras = SrasFile(Path(path)) + self.calib = ChannelCalibration.from_preambles(self.sras.preambles) + bg = np.frombuffer(self.sras.background, dtype=np.int8) + self.background = bg.astype(np.float32) if len(bg) else None + # Rows actually on disk per angle (aborted/resumed scans) + self.rows_available = [s.n_rows_available for s in self.sras.angle_status()] - def _parse(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 + def angle_view(self, angle_idx: int) -> np.ndarray: + """Available rows of one angle: (rows, n_ch, n_frames, spf) int8 view.""" + return self.sras.load_angle(angle_idx, n_rows=self.rows_available[angle_idx]) - if magic != b"SRAS": - raise ValueError(f"Bad magic bytes: {magic!r}") - if ver not in (2, 3, 4): - raise ValueError(f"Unsupported version: {ver}") - - self.n_angles = n_angles - self.n_rows = n_rows - self.x_start_mm = float(x_start) - self.x_delta_mm = float(x_delta) - 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 - - 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 - - raw = f.read() - - total_samples = len(raw) // bps - samples_per_row_per_ch = n_ch * spf - - # Compute actual frames per channel from the file size — the scanner - # writes the configured frame count in the header before acquisition - # begins, but ACQuire:NUMFRAMESACQuired may be lower. - 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.n_frames = actual_n_frames # actual, use this for all indexing - self.n_frames_header = n_frames_hdr - self.frame_count_mismatch = (actual_n_frames != n_frames_hdr) - self.n_frames_remainder = remainder # partial last-row samples - - # Reshape using the actual count; discard any fractional last row - dtype = np.int8 if bps == 1 else ">i2" - good = n_angles * n_rows * n_ch * actual_n_frames * spf - data = np.frombuffer(raw[:good * bps], dtype=dtype) - data = data.reshape(n_angles, n_rows, n_ch, actual_n_frames, spf) - self.data = data.astype(np.int16 if bps == 2 else np.int8) - - self.angles_deg = angles - self.y_positions_mm = y_pos - - # ------------------------------------------------------------------ - # Axes helpers - # ------------------------------------------------------------------ - - @property - def pixel_x_mm(self) -> float: - return self.velocity_mm_s / self.laser_freq_hz - - def x_axis_mm(self) -> np.ndarray: - return self.x_start_mm + np.arange(self.n_frames) * self.pixel_x_mm - - def time_axis_ns(self) -> np.ndarray: - return np.arange(self.samples_per_frame) / self.sample_rate_hz * 1e9 - - def freq_axis_mhz(self) -> np.ndarray: - return np.fft.rfftfreq(self.samples_per_frame, d=1.0 / self.sample_rate_hz) / 1e6 - - -# --------------------------------------------------------------------------- -# Image computation (vectorised) -# --------------------------------------------------------------------------- - - - -def compute_dc_image(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray: - """Mean of each waveform → (n_rows, n_frames) float32.""" - return sras.data[angle_idx, :, ch_idx, :, :].astype(np.float32).mean(axis=-1) - - -def compute_rf_image(sras: SrasFile, angle_idx: int, - dc_threshold_mv: float, - apply_bg_sub: bool = True, - gate_start_ns: float | None = None, - gate_end_ns: float | 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; FFT is skipped for - those pixels entirely. DC is always computed before any FFT work. - The threshold and DC mean are both in mV, using per-channel calibration - from the file (or fallback constants for v2 files). - - If apply_bg_sub is True and the file contains a background waveform - (v4+), each CH1 waveform has the background subtracted before the FFT. - - gate_start_ns / gate_end_ns: when either is set, samples outside the - [start, end] time window are zeroed before the FFT (time-domain gating). - """ - # --- Step 1: compute CH4 DC mask before any FFT work --- - dc4_mv = 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]) - mask = dc4_mv < dc_threshold_mv # True = below threshold = skip FFT - valid = ~mask # pixels that require FFT - - img = np.zeros(mask.shape, dtype=np.float32) - - if valid.any(): - # --- Step 2: FFT only on pixels that passed the DC threshold --- - waveforms = sras.data[angle_idx, :, CH1_IDX, :, :].astype(np.float32) - # shape: (n_rows, n_frames, samples_per_frame) - - if apply_bg_sub and sras.background is not None: - waveforms = waveforms - sras.background[np.newaxis, np.newaxis, :] - - if gate_start_ns is not None or gate_end_ns is not None: - t_ns = sras.time_axis_ns() - keep = np.ones(len(t_ns), dtype=bool) - if gate_start_ns is not None: - keep &= t_ns >= gate_start_ns - if gate_end_ns is not None: - keep &= t_ns <= gate_end_ns - waveforms = waveforms.copy() - waveforms[..., ~keep] = 0.0 - - valid_waves = waveforms[valid] # (n_valid, spf) - fft_pow = np.abs(np.fft.rfft(valid_waves, axis=-1)) ** 2 - fft_pow[:, 0] = 0.0 # suppress DC bin - peak_bins = np.argmax(fft_pow, axis=-1) # (n_valid,) - img[valid] = sras.freq_axis_mhz()[peak_bins] - - return img - - -def compute_saw_image(sras: SrasFile, angle_idx: int, dc_threshold_mv: float, - pipeline: SawPipeline, mode: str, - apply_bg_sub: bool = True) -> np.ndarray: - """Run the SAW matched-filter pipeline over every pixel. - - mode : "amplitude" → MF envelope peak in SAW window - "tof" → arrival time (ns) of that peak - Returns (n_rows, n_frames) float32, DC-masked. - """ - # --- Step 1: compute CH4 DC mask before running the pipeline --- - dc4_mv = 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]) - mask = dc4_mv < dc_threshold_mv # True = below threshold = skip pipeline - valid = ~mask - - img = np.zeros(mask.shape, dtype=np.float32) - - if valid.any(): - # --- Step 2: run pipeline only on pixels that passed the DC threshold --- - waveforms = sras.data[angle_idx, :, CH1_IDX, :, :].astype(np.float32) - if apply_bg_sub and sras.background is not None: - waveforms = waveforms - sras.background[np.newaxis, np.newaxis, :] - - valid_waves = waveforms[valid] # (n_valid, spf) - n_workers = min(os.cpu_count() or 4, len(valid_waves)) - with ThreadPoolExecutor(max_workers=n_workers) as executor: - results = list(executor.map(pipeline.process_shot, valid_waves)) - - if mode == "amplitude": - vals = np.array([r["peak_amplitude"] for r in results], dtype=np.float32) - else: - vals = np.array([r["peak_time_ns"] for r in results], dtype=np.float32) - - img[valid] = vals - - return img + def close(self): + self.sras.close() # --------------------------------------------------------------------------- # Background workers # --------------------------------------------------------------------------- -class LoadWorker(QObject): - finished = pyqtSignal(object) # SrasFile | None - error = pyqtSignal(str) +class FnWorker(QObject): + """Runs a callable on a QThread; emits its return value or the error.""" + finished = pyqtSignal(object) + error = pyqtSignal(str) - def __init__(self, path: str): + def __init__(self, fn): super().__init__() - self._path = path + self._fn = fn 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): - finished = pyqtSignal(np.ndarray) - error = pyqtSignal(str) - - def __init__(self, sras: SrasFile, angle_idx: int, - ch_idx: int, dc_threshold_mv: float, - grating_um: float = 12.5, - apply_bg_sub: bool = True, - gate_start_ns: float | None = None, - gate_end_ns: float | None = None, - saw_pipeline: "SawPipeline | None" = None): - super().__init__() - self._sras = sras - self._angle = angle_idx - self._ch = ch_idx - self._threshold = dc_threshold_mv - self._grating_um = grating_um - self._apply_bg_sub = apply_bg_sub - self._gate_start = gate_start_ns - self._gate_end = gate_end_ns - self._saw_pipeline = saw_pipeline - - def run(self): - try: - if self._ch == CH1_IDX: - img = compute_rf_image(self._sras, self._angle, self._threshold, - self._apply_bg_sub, - self._gate_start, self._gate_end) - elif self._ch == VELOCITY_MODE_IDX: - # velocity (m/s) = freq (MHz) × grating (µm) [units cancel to m/s] - img = compute_rf_image(self._sras, self._angle, self._threshold, - self._apply_bg_sub, - self._gate_start, self._gate_end) - img = img * self._grating_um - elif self._ch in SAW_MODES: - if self._saw_pipeline is None or self._saw_pipeline.template is None: - raise RuntimeError( - "SAW pipeline: no template built yet.\n" - "Use \"Build Template\" in the SAW Pipeline panel first.") - mode = "amplitude" if self._ch == SAW_MODE_AMP_IDX else "tof" - img = compute_saw_image( - self._sras, self._angle, self._threshold, - self._saw_pipeline, mode, self._apply_bg_sub) - else: - # DC channels: convert ADC counts → mV - adc_img = compute_dc_image(self._sras, self._angle, self._ch) - img = adc_to_mv(adc_img, - self._sras.ch_ymult_mv[self._ch], - self._sras.ch_yoff_adc[self._ch], - self._sras.ch_yzero_mv[self._ch]) - self.finished.emit(img) + self.finished.emit(self._fn()) except Exception as exc: self.error.emit(str(exc)) -class TemplateBuildWorker(QObject): - """Background thread worker that calls SawPipeline.build_template().""" - finished = pyqtSignal() - error = pyqtSignal(str) +def compute_image(scan: LoadedScan, angle_idx: int, ch_idx: int, + dc_threshold_mv: float, apply_bg_sub: bool, + gate_start_ns: float | None, gate_end_ns: float | None, + saw_pipeline: SawPipeline | None) -> np.ndarray: + """Channel-mode dispatch for the image computation worker. - def __init__(self, pipeline: SawPipeline, waveforms: np.ndarray): - super().__init__() - self._pipeline = pipeline - self._waveforms = waveforms + VELOCITY mode returns the plain RF image (MHz) — the grating scaling is + a display-time scalar multiply, so changing the grating never triggers a + recompute. + """ + view = scan.angle_view(angle_idx) + if view.shape[0] == 0: + return np.zeros((0, 0), dtype=np.float32) + sras = scan.sras + bg = scan.background if apply_bg_sub else None - def run(self): - try: - self._pipeline.build_template(self._waveforms) - self.finished.emit() - except Exception as exc: - self.error.emit(str(exc)) + if ch_idx in (CH1_IDX, VELOCITY_MODE_IDX): + return compute_rf_image( + view, scan.calib, sras.freq_axis_mhz(sras.header.samples_per_frame), + dc_threshold_mv, background=bg, + gate_start_ns=gate_start_ns, gate_end_ns=gate_end_ns, + time_axis_ns=sras.time_axis_ns(), + ) + if ch_idx in SAW_MODES: + if saw_pipeline is None or saw_pipeline.template is None: + raise RuntimeError( + "SAW pipeline: no template built yet.\n" + "Use \"Build Template\" in the SAW Pipeline panel first.") + mode = "amplitude" if ch_idx == SAW_MODE_AMP_IDX else "tof" + return compute_saw_image(view, scan.calib, dc_threshold_mv, + saw_pipeline, mode, background=bg) + # DC channels: ADC counts → mV + return scan.calib.adc_to_mv(compute_dc_image(view, ch_idx), ch_idx) # --------------------------------------------------------------------------- @@ -621,38 +157,50 @@ 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 + self._im = None + self._cbar = None self.mpl_connect("button_press_event", self._on_click) def show_image(self, img: np.ndarray, extent: list[float], cmap: str, vmin: float, vmax: float, xlabel: str, ylabel: str, title: str, colorbar_label: str = ""): - self.figure.clf() - self.ax = self.figure.add_subplot(111) - - self._extent = extent + # Rebuild the artist tree only when the geometry changes; colormap, + # limits, and data updates reuse the existing AxesImage + colorbar. + rebuild = (self._im is None or img.shape != self._img_shape + or extent != self._extent) + self._extent = extent self._img_shape = img.shape - im = self.ax.imshow( - img, aspect="auto", origin="upper", - extent=extent, cmap=cmap, vmin=vmin, vmax=vmax, - interpolation="nearest", - ) - cb = self.figure.colorbar(im, ax=self.ax, fraction=0.046, pad=0.04) - if colorbar_label: - cb.set_label(colorbar_label) + if rebuild: + self.figure.clf() + self.ax = self.figure.add_subplot(111) + self._im = self.ax.imshow( + img, aspect="auto", origin="upper", + extent=extent, cmap=cmap, vmin=vmin, vmax=vmax, + interpolation="nearest", + ) + self._cbar = self.figure.colorbar(self._im, ax=self.ax, + fraction=0.046, pad=0.04) + else: + self._im.set_data(img) + self._im.set_cmap(cmap) + self._im.set_clim(vmin, vmax) + self._cbar.set_label(colorbar_label or "") self.ax.set_xlabel(xlabel) self.ax.set_ylabel(ylabel) self.ax.set_title(title) - self.draw() + self.draw_idle() def _on_click(self, event): if event.inaxes is not self.ax or self._extent is None: return x0, x1, y_bot, y_top = self._extent n_rows, n_frames = self._img_shape + if n_rows == 0 or n_frames == 0: + return 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)) @@ -669,25 +217,22 @@ class WaveformCanvas(FigureCanvasQTAgg): self.setParent(parent) self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding) - def show_rf_waveform(self, sras: SrasFile, angle_idx: int, + def show_rf_waveform(self, scan: LoadedScan, angle_idx: int, row_idx: int, frame_idx: int, apply_bg_sub: bool = True, gate_start_ns: float | None = None, gate_end_ns: float | None = None): - """CH1 RF: time-domain + FFT spectrum. + """CH1 RF: time-domain + FFT spectrum (background overlay if present).""" + sras = scan.sras + waveform = sras.load_row(angle_idx, row_idx, CH1_IDX)[frame_idx].astype(np.float32) + t_ns = sras.time_axis_ns() + f_mhz = sras.freq_axis_mhz(sras.header.samples_per_frame) + dc3_val = float(sras.load_row(angle_idx, row_idx, CH3_IDX)[frame_idx] + .mean(dtype=np.float32)) + dc4_val = float(sras.load_row(angle_idx, row_idx, CH4_IDX)[frame_idx] + .mean(dtype=np.float32)) - If apply_bg_sub is True and sras.background is not None, the background - waveform is overlaid on the time-domain plot and the FFT is computed - on the subtracted signal. The unsubtracted FFT is also shown faintly - for comparison. - """ - waveform = sras.data[angle_idx, row_idx, CH1_IDX, frame_idx, :].astype(np.float32) - t_ns = sras.time_axis_ns() - f_mhz = sras.freq_axis_mhz() - dc3_val = sras.data[angle_idx, row_idx, CH3_IDX, frame_idx, :].astype(np.float32).mean() - dc4_val = sras.data[angle_idx, 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 + bg = scan.background if apply_bg_sub else None waveform_plot = waveform - bg if bg is not None else waveform self.ax_wave.cla() @@ -704,7 +249,6 @@ class WaveformCanvas(FigureCanvasQTAgg): else: self.ax_wave.plot(t_ns, waveform, linewidth=0.7, color="#4488cc") - # Draw gate boundaries if active if gate_start_ns is not None: self.ax_wave.axvline(gate_start_ns, color="#22cc44", linestyle="--", linewidth=1.0, label=f"gate start {gate_start_ns:.0f} ns") @@ -712,33 +256,29 @@ class WaveformCanvas(FigureCanvasQTAgg): self.ax_wave.axvline(gate_end_ns, color="#cc4422", linestyle="--", linewidth=1.0, label=f"gate end {gate_end_ns:.0f} ns") if gate_start_ns is not None or gate_end_ns is not None: - t_ns = sras.time_axis_ns() lo = gate_start_ns if gate_start_ns is not None else t_ns[0] - hi = gate_end_ns if gate_end_ns is not None else t_ns[-1] + hi = gate_end_ns if gate_end_ns is not None else t_ns[-1] self.ax_wave.axvspan(t_ns[0], lo, alpha=0.10, color="#cc4422") self.ax_wave.axvspan(hi, t_ns[-1], alpha=0.10, color="#cc4422") + calib = scan.calib self.ax_wave.set_xlabel("Time (ns)") self.ax_wave.set_ylabel("ADC counts") bg_tag = " [bg sub]" if bg is not None else "" 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"({calib.adc_to_mv(dc3_val, CH3_IDX):.2f} / " + f"{calib.adc_to_mv(dc4_val, CH4_IDX):.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 + power_sub = power_spectrum(waveform_plot) peak_idx = int(np.argmax(power_sub)) peak_mhz = f_mhz[peak_idx] if bg is not None: - # Also show the unsubtracted FFT for reference - power_raw = np.abs(np.fft.rfft(waveform)) ** 2 - power_raw[0] = 0.0 + power_raw = power_spectrum(waveform) self.ax_right.plot(f_mhz, power_raw, linewidth=0.5, color="#aaaaaa", label="raw FFT", zorder=1) @@ -752,16 +292,16 @@ class WaveformCanvas(FigureCanvasQTAgg): self.ax_right.set_xlim(0, 500) self.ax_right.legend(fontsize=8) - self.draw() + self.draw_idle() - def show_dc_waveform(self, sras: SrasFile, angle_idx: int, ch_idx: int, + def show_dc_waveform(self, scan: LoadedScan, angle_idx: int, ch_idx: int, 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() + sras = scan.sras + waveform = sras.load_row(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 = scan.calib.adc_to_mv(mean_val, ch_idx) self.ax_wave.cla() self.ax_right.cla() @@ -786,7 +326,7 @@ class WaveformCanvas(FigureCanvasQTAgg): ) self.ax_right.set_axis_off() - self.draw() + self.draw_idle() # --------------------------------------------------------------------------- @@ -796,21 +336,16 @@ class WaveformCanvas(FigureCanvasQTAgg): class SawDiagnosticWindow(QMainWindow): """6-panel matplotlib window showing every SAW pipeline stage for one pixel. - Panels: - 1. Raw signal with zone shading (EMI gate / noise region / SAW window) - 2. After EMI gating (cosine taper) — same zone shading - 3. After bandpass filter (time domain) — zone shading - 4. Frequency spectrum of bandpass output — passband shading + raw PSD - 5. Matched filter output + Hilbert envelope + metrics — zone shading - 6. Shot-to-shot overlay (up to 20 frames from the same row) + Panels: raw / EMI-gated / bandpassed / PSD / matched filter / shot-to-shot + overlay. The pipeline runs (once for the pixel + up to 20 overlay shots) + happen on a worker thread; only the plotting touches the GUI thread. """ - # Zone colour constants (all panels use the same palette) - _C_EMI = "#e05030" # red — EMI gate + _C_EMI = "#e05030" # red — EMI gate _C_NOISE = "#ccaa00" # amber — noise / inter-packet region - _C_SAW = "#30c060" # green — SAW window + _C_SAW = "#30c060" # green — SAW window - def __init__(self, sras: SrasFile, pipeline: SawPipeline, + def __init__(self, scan: LoadedScan, pipeline: SawPipeline, angle_idx: int, row_idx: int, frame_idx: int, parent=None): super().__init__(parent) @@ -819,12 +354,12 @@ class SawDiagnosticWindow(QMainWindow): f"SAW Diagnostics angle={angle_idx} row={row_idx} frame={frame_idx}") self.resize(1400, 940) - # Store references for refresh - self._sras = sras - self._pipeline = pipeline + self._scan = scan + self._pipeline = pipeline self._angle_idx = angle_idx - self._row_idx = row_idx + self._row_idx = row_idx self._frame_idx = frame_idx + self._worker_thread: QThread | None = None central = QWidget() self.setCentralWidget(central) @@ -832,33 +367,66 @@ class SawDiagnosticWindow(QMainWindow): vl.setContentsMargins(4, 4, 4, 4) vl.setSpacing(4) - # Toolbar row: matplotlib toolbar + refresh button toolbar_row = QHBoxLayout() fig = Figure(figsize=(14, 9), tight_layout=True) self._canvas = FigureCanvasQTAgg(fig) mpl_toolbar = NavigationToolbar2QT(self._canvas, central) toolbar_row.addWidget(mpl_toolbar, stretch=1) - btn_refresh = QPushButton("Re-apply filter & refresh PSDs") - btn_refresh.setToolTip( + self._btn_refresh = QPushButton("Re-apply filter & refresh PSDs") + self._btn_refresh.setToolTip( "Re-run the current pipeline on this pixel and redraw all panels.\n" "Use after rebuilding the template or changing pipeline parameters.") - btn_refresh.clicked.connect(self._on_refresh) - toolbar_row.addWidget(btn_refresh) + self._btn_refresh.clicked.connect(self._start_compute) + toolbar_row.addWidget(self._btn_refresh) vl.addLayout(toolbar_row) vl.addWidget(self._canvas) - self._plot() + self._start_compute() - # ------------------------------------------------------------------ - # Zone shading helper — call on any time-domain axes - # ------------------------------------------------------------------ + def _compute(self): + """Worker-thread body: pixel result + overlay envelopes.""" + sras = self._scan.sras + row = sras.load_row(self._angle_idx, self._row_idx, CH1_IDX) + raw_adc = row[self._frame_idx].astype(np.float32) + result = self._pipeline.process_shot(raw_adc) + + n_frames = row.shape[0] + n_overlay = min(20, n_frames) + overlay_idxs = np.linspace(0, n_frames - 1, n_overlay, dtype=int) + overlays = [ + self._pipeline.process_shot(row[fi].astype(np.float32))["envelope"] + for fi in overlay_idxs + ] + return raw_adc, result, overlays + + def _start_compute(self): + if self._worker_thread is not None: + return + self._btn_refresh.setEnabled(False) + worker = FnWorker(self._compute) + thread = QThread() + worker.moveToThread(thread) + thread.started.connect(worker.run) + worker.finished.connect(self._on_computed) + worker.error.connect(lambda msg: self.statusBar().showMessage(f"Error: {msg}")) + worker.finished.connect(thread.quit) + worker.error.connect(thread.quit) + thread.finished.connect(self._on_thread_finished) + self._worker = worker + self._worker_thread = thread + thread.start() + + def _on_thread_finished(self): + self._worker_thread = None + self._worker = None + self._btn_refresh.setEnabled(True) def _shade_time_zones(self, ax, t_full: np.ndarray, emi_end_ns: float, s0: float, s1: float, show_legend: bool = False): """Shade EMI gate, noise region, and SAW window on a time-domain axes.""" t0, t_end = float(t_full[0]), float(t_full[-1]) - ax.axvspan(t0, emi_end_ns, alpha=0.15, color=self._C_EMI, + ax.axvspan(t0, emi_end_ns, alpha=0.15, color=self._C_EMI, label=f"EMI gate 0–{emi_end_ns:.0f} ns") if emi_end_ns < s0: ax.axvspan(emi_end_ns, s0, alpha=0.08, color=self._C_NOISE, @@ -868,47 +436,29 @@ class SawDiagnosticWindow(QMainWindow): if show_legend: ax.legend(fontsize=7, loc="upper right") - # ------------------------------------------------------------------ + def _on_computed(self, payload): + raw_adc, result, overlays = payload + sras = self._scan.sras + pipeline = self._pipeline + row_idx, frame_idx = self._row_idx, self._frame_idx - def _on_refresh(self): - self._plot() + t_full = sras.time_axis_ns() + sr = result["sample_rate_hz"] + t_proc = np.arange(len(result["envelope"])) / sr * 1e9 - def _plot(self): - sras = self._sras - pipeline = self._pipeline - angle_idx = self._angle_idx - row_idx = self._row_idx - frame_idx = self._frame_idx + peak_amps = [float(e.max()) if len(e) > 0 else 0.0 for e in overlays] + peak_var = float(np.var(peak_amps)) + n_overlay = len(overlays) - raw_adc = sras.data[angle_idx, row_idx, CH1_IDX, frame_idx, :].astype(np.float32) - t_full = sras.time_axis_ns() - result = pipeline.process_shot(raw_adc) - sr = result["sample_rate_hz"] - t_proc = np.arange(len(result["envelope"])) / sr * 1e9 - - # Collect up to 20 frames from the same row for the overlay panel - n_overlay = min(20, sras.n_frames) - overlay_idxs = np.linspace(0, sras.n_frames - 1, n_overlay, dtype=int) - overlays = [] - for fi in overlay_idxs: - sig = sras.data[angle_idx, row_idx, CH1_IDX, fi, :].astype(np.float32) - r = pipeline.process_shot(sig) - overlays.append(r["envelope"]) - - # Shot-to-shot peak amplitude variance - peak_amps = [float(e[np.argmax(e)]) if len(e) > 0 else 0.0 - for e in overlays] - peak_var = float(np.var(peak_amps)) - - fig = self._canvas.figure + fig = self._canvas.figure fig.clf() axes = fig.subplots(3, 2) - ax_raw, ax_gated = axes[0] - ax_filt, ax_spec = axes[1] - ax_mf, ax_over = axes[2] + ax_raw, ax_gated = axes[0] + ax_filt, ax_spec = axes[1] + ax_mf, ax_over = axes[2] emi_end_ns = pipeline.emi_gate_ns - s0, s1 = pipeline.saw_window_ns + s0, s1 = pipeline.saw_window_ns # --- 1. Raw signal --- ax_raw.plot(t_full, raw_adc, lw=0.6, color="#4488cc", zorder=3) @@ -936,15 +486,13 @@ class SawDiagnosticWindow(QMainWindow): f"6th-order Butterworth, zero-phase)") # --- 4. Frequency spectrum — raw PSD + post-bandpass PSD --- - raw_sig = raw_adc.astype(np.float64) filt_sig = result["filtered"] - f_hz = np.fft.rfftfreq(len(filt_sig), d=1.0 / sras.sample_rate_hz) - f_mhz = f_hz / 1e6 - spec_raw = np.abs(np.fft.rfft(raw_sig, n=len(filt_sig))) ** 2 + f_mhz = np.fft.rfftfreq(len(filt_sig), d=1.0 / sras.header.sample_rate) / 1e6 + spec_raw = np.abs(np.fft.rfft(raw_adc, n=len(filt_sig))) ** 2 spec_filt = np.abs(np.fft.rfft(filt_sig)) ** 2 spec_raw[0] = 0.0 spec_filt[0] = 0.0 - ax_spec.plot(f_mhz, spec_raw, lw=0.5, color="#aaaaaa", alpha=0.7, + ax_spec.plot(f_mhz, spec_raw, lw=0.5, color="#aaaaaa", alpha=0.7, label="raw PSD", zorder=1) ax_spec.plot(f_mhz, spec_filt, lw=0.8, color="#44aa44", label="bandpass PSD", zorder=2) @@ -953,7 +501,7 @@ class SawDiagnosticWindow(QMainWindow): ax_spec.set_xlabel("Frequency (MHz)") ax_spec.set_ylabel("Power (arb.)") ax_spec.set_title("4 — FFT PSD: raw vs. after bandpass") - ax_spec.set_xlim(0, min(600.0, sras.sample_rate_hz / 2e6)) + ax_spec.set_xlim(0, min(600.0, sras.header.sample_rate / 2e6)) ax_spec.legend(fontsize=7) # --- 5. Matched filter output + envelope --- @@ -986,7 +534,6 @@ class SawDiagnosticWindow(QMainWindow): f"6 — Shot-to-shot overlay (row {row_idx}, {n_overlay} frames) | " f"peak-amp variance = {peak_var:.4g}") - # Bottom metrics bar fig.text( 0.5, 0.005, f"Peak amplitude: {result['peak_amplitude']:.4g} | " @@ -997,7 +544,13 @@ class SawDiagnosticWindow(QMainWindow): bbox=dict(boxstyle="round,pad=0.3", facecolor="#2a2a2a", alpha=0.85), color="#e8e8e8", ) - self._canvas.draw() + self._canvas.draw_idle() + + def closeEvent(self, event): + if self._worker_thread is not None: + self._worker_thread.quit() + self._worker_thread.wait(2000) + super().closeEvent(event) # --------------------------------------------------------------------------- @@ -1011,22 +564,21 @@ class SrasViewerWindow(QMainWindow): self.resize(1560, 840) self.setAcceptDrops(True) - self._sras: SrasFile | None = None - self._current_image: np.ndarray | None = None + self._scan: LoadedScan | None = None + self._current_image: np.ndarray | None = None # unscaled (RF in MHz for VEL mode) 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_threshold: float = 50.0 # mV - self._pending_grating_um: float = 12.5 # µm - self._pending_bg_sub: bool = True - self._pending_gate_enabled: bool = False - self._pending_gate_start: float = 0.0 - self._pending_gate_end: float = 200.0 + self._compute_pending = False self._progress_dlg: QProgressDialog | None = None + # Debounce: rapid spinbox changes coalesce into one recompute + self._recompute_timer = QTimer(self) + self._recompute_timer.setSingleShot(True) + self._recompute_timer.setInterval(RECOMPUTE_DEBOUNCE_MS) + self._recompute_timer.timeout.connect(self._start_compute) + # SAW pipeline state self._saw_pipeline: SawPipeline | None = None self._template_thread: QThread | None = None @@ -1081,7 +633,7 @@ class SrasViewerWindow(QMainWindow): lbl.setStyleSheet("font-size: 11px;") il.addWidget(lbl) self._info[key] = lbl - # Frame count warning (hidden until needed) + # Truncation / background notes (hidden until needed) self.lbl_frame_warn = QLabel("") self.lbl_frame_warn.setWordWrap(True) self.lbl_frame_warn.setStyleSheet("color: #e07000; font-size: 11px;") @@ -1134,18 +686,18 @@ class SrasViewerWindow(QMainWindow): self.spin_threshold_mv.valueChanged.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 = QLabel("") self.lbl_threshold_adc.setStyleSheet("font-size: 11px; color: #888;") tl.addWidget(self.lbl_threshold_adc) vl.addWidget(self.grp_threshold) - # Background subtraction (v4 files only) + # Background subtraction self.chk_bg_sub = QCheckBox("Background subtraction (CH1 only)") self.chk_bg_sub.setChecked(True) self.chk_bg_sub.setEnabled(False) self.chk_bg_sub.setToolTip( "Subtract the stored background waveform from each CH1 frame\n" - "before computing the FFT (v4 files only)." + "before computing the FFT." ) self.chk_bg_sub.toggled.connect(self._on_bg_sub_toggled) vl.addWidget(self.chk_bg_sub) @@ -1231,7 +783,6 @@ class SrasViewerWindow(QMainWindow): grp_saw = QGroupBox("SAW Pipeline (CH1 only)") sl = QVBoxLayout(grp_saw) - # EMI gate end emi_row = QHBoxLayout() emi_row.addWidget(QLabel("EMI gate end:")) self.spin_saw_emi_ns = QDoubleSpinBox() @@ -1243,7 +794,6 @@ class SrasViewerWindow(QMainWindow): emi_row.addWidget(self.spin_saw_emi_ns) sl.addLayout(emi_row) - # SAW window sl.addWidget(QLabel("SAW window (ns):")) saw_win_row = QHBoxLayout() self.spin_saw_win_start = QDoubleSpinBox() @@ -1261,7 +811,6 @@ class SrasViewerWindow(QMainWindow): saw_win_row.addWidget(self.spin_saw_win_end) sl.addLayout(saw_win_row) - # Bandpass limits sl.addWidget(QLabel("Bandpass (MHz):")) bp_row = QHBoxLayout() self.spin_saw_bp_lo = QDoubleSpinBox() @@ -1279,7 +828,6 @@ class SrasViewerWindow(QMainWindow): bp_row.addWidget(self.spin_saw_bp_hi) sl.addLayout(bp_row) - # Template shots tmpl_row = QHBoxLayout() tmpl_row.addWidget(QLabel("Template shots:")) self.spin_saw_n_shots = QSpinBox() @@ -1301,7 +849,6 @@ class SrasViewerWindow(QMainWindow): self.lbl_saw_status.setWordWrap(True) sl.addWidget(self.lbl_saw_status) - # ---- Apply matched filter controls ---- sep_mf = QFrame() sep_mf.setFrameShape(QFrame.Shape.HLine) sep_mf.setStyleSheet("color: #555;") @@ -1328,10 +875,10 @@ class SrasViewerWindow(QMainWindow): panel_layout.addStretch() # Colormap - sep2 = QFrame() - sep2.setFrameShape(QFrame.Shape.HLine) - sep2.setStyleSheet("color: #555;") - vl.addWidget(sep2) + sep3 = QFrame() + sep3.setFrameShape(QFrame.Shape.HLine) + sep3.setStyleSheet("color: #555;") + vl.addWidget(sep3) cmr = QHBoxLayout() cmr.addWidget(QLabel("Colormap:")) @@ -1339,7 +886,7 @@ class SrasViewerWindow(QMainWindow): self.combo_cmap.addItems(CMAPS) self.combo_cmap.setCurrentText("gray") self.combo_cmap.setEnabled(False) - self.combo_cmap.currentIndexChanged.connect(self._on_view_changed) + self.combo_cmap.currentIndexChanged.connect(self._on_display_changed) cmr.addWidget(self.combo_cmap) vl.addLayout(cmr) @@ -1365,7 +912,6 @@ class SrasViewerWindow(QMainWindow): 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) @@ -1376,7 +922,6 @@ class SrasViewerWindow(QMainWindow): 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) @@ -1399,7 +944,7 @@ class SrasViewerWindow(QMainWindow): splitter.setSizes([580, 250]) - # ---- Right control panel (SAW pipeline) ---------------------------- + # ---- Right control panel (SAW pipeline) ------------------------ right_panel = QWidget() right_panel.setFixedWidth(260) right_panel_layout = QVBoxLayout(right_panel) @@ -1410,6 +955,7 @@ class SrasViewerWindow(QMainWindow): root.addWidget(right_panel) self.statusBar().showMessage("Open an .sras file to begin.") + self._update_threshold_label() # ------------------------------------------------------------------ # Drag-and-drop @@ -1436,279 +982,324 @@ class SrasViewerWindow(QMainWindow): def _load_file(self, path: str): if self._load_thread is not None: + self.statusBar().showMessage("Still loading previous file…") return self.btn_open.setEnabled(False) self.statusBar().showMessage(f"Loading {Path(path).name}…") self._show_progress(f"Loading {Path(path).name}…") - 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}") - ) - self._load_worker.finished.connect(self._load_thread.quit) - self._load_thread.finished.connect(lambda: setattr(self, "_load_thread", None)) - self._load_thread.start() + worker = FnWorker(lambda: LoadedScan(path)) + thread = QThread() + worker.moveToThread(thread) + thread.started.connect(worker.run) + worker.finished.connect(self._on_load_done) + worker.error.connect(self._on_load_error) + worker.finished.connect(thread.quit) + worker.error.connect(thread.quit) + thread.finished.connect(lambda: setattr(self, "_load_thread", None)) + self._load_worker = worker + self._load_thread = thread + thread.start() - def _on_load_done(self, sras): + def _on_load_error(self, msg: str): self._close_progress() self.btn_open.setEnabled(True) - if sras is None: - return - self._sras = sras + self.statusBar().showMessage(f"Error: {msg}") + + def _on_load_done(self, scan: LoadedScan): + self._close_progress() + self.btn_open.setEnabled(True) + + # Drop every reference to the previous file so its mapping (and any + # multi-GB views derived from it) can actually be freed. + if self._diag_window is not None: + self._diag_window.close() + self._diag_window = None + self._load_worker = None self._current_image = None + self._last_row = None + self._last_frame = None + old = self._scan + self._scan = scan + if old is not None: + old.close() - s = sras - self.lbl_filename.setText(s.path.name) - self._info["Angles"].setText(f"Angles: {s.n_angles}") - self._info["Rows"].setText(f"Rows: {s.n_rows}") - self._info["Frames / row"].setText(f"Frames / row: {s.n_frames}") - 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:.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") - - notes = [] - if s.frame_count_mismatch: - notes.append( - f"! Header n_frames={s.n_frames_header}, " - f"actual={s.n_frames} (scanner bug — corrected)" - ) - if s.background is not None: - notes.append(f"Background waveform: {len(s.background)} samples") - self.lbl_frame_warn.setText("\n".join(notes)) + sras = scan.sras + h = sras.header + self.lbl_filename.setText(sras.path.name) + self._info["Angles"].setText(f"Angles: {h.n_angles}") + self._info["Samples / frame"].setText(f"Samples / frame: {h.samples_per_frame}") + self._info["Sample rate"].setText(f"Sample rate: {h.sample_rate/1e9:.4g} GS/s") + self._info["Pixel Δx"].setText(f"Pixel Δx: {sras.pixel_pitch_x_mm()*1e3:.3g} µm") + self._info["Laser freq"].setText(f"Laser freq: {h.laser_freq/1e3:.4g} kHz") self.spin_angle.blockSignals(True) - self.spin_angle.setRange(0, max(0, s.n_angles - 1)) + self.spin_angle.setRange(0, max(0, h.n_angles - 1)) self.spin_angle.setValue(0) self.spin_angle.blockSignals(False) self._update_controls_enabled(True) - # Refresh the ADC-count label now that we have file calibration - self._on_threshold_changed(self.spin_threshold_mv.value()) + self._update_threshold_label() self._on_view_changed() + def _update_angle_info(self): + """Per-angle info fields (v6 geometry is ragged across angles).""" + scan = self._scan + if scan is None: + return + ai = self.spin_angle.value() + pa = scan.sras.per_angle[ai] + avail = scan.rows_available[ai] + rows_txt = f"Rows: {pa.n_rows}" if avail == pa.n_rows \ + else f"Rows: {avail} of {pa.n_rows} on disk" + self._info["Rows"].setText(rows_txt) + self._info["Frames / row"].setText(f"Frames / row: {pa.n_frames}") + self._info["X start"].setText(f"X start: {pa.x_start:.4g} mm") + + notes = [] + if avail < pa.n_rows: + notes.append(f"! Angle {ai}: only {avail}/{pa.n_rows} rows on disk " + "(scan was aborted or is still running)") + if scan.background is not None: + notes.append(f"Background waveform: {len(scan.background)} samples") + self.lbl_frame_warn.setText("\n".join(notes)) + # ------------------------------------------------------------------ # Controls # ------------------------------------------------------------------ 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) + scan = self._scan + has_file = enabled and scan is not None + self.spin_angle.setEnabled(has_file and scan.sras.header.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 - is_fft = enabled and ch_idx in (CH1_IDX, VELOCITY_MODE_IDX) - is_saw = enabled and ch_idx in SAW_MODES - # Threshold and bg-sub apply to all CH1 modes + + ch_idx = self.combo_channel.currentIndex() + is_ch1 = enabled and ch_idx in CH1_DERIVED_MODES + is_fft = enabled and ch_idx in (CH1_IDX, VELOCITY_MODE_IDX) self.spin_threshold_mv.setEnabled(is_ch1) - has_bg = enabled and s is not None and s.background is not None + has_bg = has_file and scan.background is not None self.chk_bg_sub.setEnabled(has_bg and is_ch1) - # Time gate only for legacy FFT modes (SAW pipeline has its own gating) + # Time gate only for FFT modes (the SAW pipeline has its own gating) self.chk_gate.setEnabled(is_fft) gate_active = is_fft and self.chk_gate.isChecked() self.spin_gate_start.setEnabled(gate_active) self.spin_gate_end.setEnabled(gate_active) - # Velocity grating spinbox + is_vel = enabled and ch_idx == VELOCITY_MODE_IDX self.spin_grating_um.setEnabled(is_vel) self.grp_velocity.setVisible(is_vel) - # SAW pipeline build button - has_file = enabled and s is not None + self.btn_build_template.setEnabled(has_file) - # Diagnostics button: need template + a clicked pixel - has_template = self._saw_pipeline is not None and self._saw_pipeline.template is not None - has_pixel = self._last_row is not None + has_template = (self._saw_pipeline is not None + and self._saw_pipeline.template is not None) + has_pixel = self._last_row is not None self.btn_saw_diag.setEnabled(has_file and has_template and has_pixel) - # Apply filter button: need file + template self.btn_apply_mf.setEnabled(has_file and has_template) self.combo_mf_mode.setEnabled(has_file and has_template) - # CSV export: enabled when a CH1-derived image is displayed self.btn_export_csv.setEnabled(is_ch1 and self._current_image is not None) 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 - is_fft = ch_idx in (CH1_IDX, VELOCITY_MODE_IDX) - 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) - self.chk_gate.setEnabled(is_fft and has_file) - gate_active = is_fft and has_file and self.chk_gate.isChecked() - self.spin_gate_start.setEnabled(gate_active) - self.spin_gate_end.setEnabled(gate_active) - 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._on_view_changed() + self._update_controls_enabled(True) + self._request_compute() def _on_bg_sub_toggled(self): - if self._sras is not None: - if self.combo_channel.currentIndex() in CH1_DERIVED_MODES: - self._start_compute() + if self._scan is not None and self.combo_channel.currentIndex() in CH1_DERIVED_MODES: + self._request_compute() def _on_gate_toggled(self, checked: bool): self.spin_gate_start.setEnabled(checked) self.spin_gate_end.setEnabled(checked) - if self._sras is not None: - ch_idx = self.combo_channel.currentIndex() - if ch_idx in (CH1_IDX, VELOCITY_MODE_IDX): - self._start_compute() + if self._scan is not None and \ + self.combo_channel.currentIndex() in (CH1_IDX, VELOCITY_MODE_IDX): + self._request_compute() def _on_gate_changed(self): - if self._sras is not None and self.chk_gate.isChecked(): - ch_idx = self.combo_channel.currentIndex() - if ch_idx in (CH1_IDX, VELOCITY_MODE_IDX): - self._start_compute() + if self._scan is not None and self.chk_gate.isChecked() and \ + self.combo_channel.currentIndex() in (CH1_IDX, VELOCITY_MODE_IDX): + self._request_compute() def _on_grating_changed(self): - if self._sras is not None and self.combo_channel.currentIndex() == VELOCITY_MODE_IDX: - self._start_compute() + # Velocity = RF (MHz) × grating (µm): a display-time scalar multiply, + # so no recompute — just redraw. + if self._scan is not None and self._current_image is not None \ + and self._current_ch == VELOCITY_MODE_IDX: + self._redraw_image(self._current_image) def _on_export_csv(self): - if self._current_image is None or self._sras is None: + if self._current_image is None or self._scan is None: return - ch_idx = self._current_ch - angle = self._current_angle + 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" - ) + stem = self._scan.sras.path.stem + default_name = f"{stem}_angle{angle}_{ch_name}.csv" path, _ = QFileDialog.getSaveFileName( self, "Export Image as CSV", - str(self._sras.path.parent / default_name), + str(self._scan.sras.path.parent / default_name), "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}") + img = self._display_image(self._current_image) + + worker = FnWorker(lambda: np.savetxt(path, img, delimiter=",", fmt="%.6g")) + thread = QThread() + worker.moveToThread(thread) + thread.started.connect(worker.run) + worker.finished.connect( + lambda _: self.statusBar().showMessage(f"Exported {Path(path).name}")) + worker.error.connect( + lambda msg: self.statusBar().showMessage(f"Export failed: {msg}")) + worker.finished.connect(thread.quit) + worker.error.connect(thread.quit) + thread.finished.connect(lambda: setattr(self, "_csv_thread", None)) + self._csv_worker = worker + self._csv_thread = thread + thread.start() + + def _update_threshold_label(self): + mv = self.spin_threshold_mv.value() + if self._scan is not None: + adc = self._scan.calib.mv_to_adc(mv, CH4_IDX) + else: + adc = ChannelCalibration.from_preambles([""] * 3).mv_to_adc(mv, CH4_IDX) + self.lbl_threshold_adc.setText(f"≈ {adc:.1f} ADC counts") def _on_threshold_changed(self, mv: float): - 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") - if self._sras is not None and self.combo_channel.currentIndex() in CH1_DERIVED_MODES: - self._start_compute() + self._update_threshold_label() + if self._scan is not None and self.combo_channel.currentIndex() in CH1_DERIVED_MODES: + self._request_compute() 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) + self.spin_vmin.setEnabled(manual and self._scan is not None) + self.spin_vmax.setEnabled(manual and self._scan is not None) + self._on_display_changed() def _on_manual_range_changed(self): - if not self.chk_auto.isChecked() and self._current_image is not None: + if not self.chk_auto.isChecked(): + self._on_display_changed() + + def _on_display_changed(self): + """Colormap/scale changes: redraw only, no recompute.""" + if self._scan is not None and self._current_image is not None: self._redraw_image(self._current_image) def _on_view_changed(self): - if self._sras is None: + if self._scan is None: return idx = self.spin_angle.value() - self.lbl_angle_deg.setText(f"({self._sras.angles_deg[idx]:.1f}°)") - self._start_compute() + self.lbl_angle_deg.setText(f"({self._scan.sras.per_angle[idx].angle_deg:.1f}°)") + self._update_angle_info() + self._request_compute() # ------------------------------------------------------------------ # Computation # ------------------------------------------------------------------ + def _request_compute(self): + """Debounced entry point — rapid widget changes coalesce into one run.""" + if self._scan is None: + return + self._recompute_timer.start() + def _start_compute(self): - if self._sras is None: + if self._scan is None: return if self._compute_thread is not None: - return # re-check in _on_compute_thread_finished + # A run is in flight; run again when it finishes. + self._compute_pending = True + return - angle_idx = self.spin_angle.value() - ch_idx = self.combo_channel.currentIndex() + scan = self._scan + angle_idx = self.spin_angle.value() + ch_idx = self.combo_channel.currentIndex() threshold_mv = self.spin_threshold_mv.value() - grating_um = self.spin_grating_um.value() apply_bg_sub = self.chk_bg_sub.isChecked() gate_enabled = self.chk_gate.isChecked() - gate_start = self.spin_gate_start.value() if gate_enabled else None - gate_end = self.spin_gate_end.value() if gate_enabled else None + gate_start = self.spin_gate_start.value() if gate_enabled else None + gate_end = self.spin_gate_end.value() if gate_enabled else None + pipeline = self._saw_pipeline - self._pending_angle = angle_idx - self._pending_ch = ch_idx - self._pending_threshold = threshold_mv - self._pending_grating_um = grating_um - self._pending_bg_sub = apply_bg_sub - self._pending_gate_enabled = gate_enabled - self._pending_gate_start = self.spin_gate_start.value() - self._pending_gate_end = self.spin_gate_end.value() + self._computing_angle = angle_idx + self._computing_ch = ch_idx self.statusBar().showMessage("Computing image…") self._show_progress("Computing image…") - self._compute_worker = ComputeWorker( - self._sras, angle_idx, ch_idx, threshold_mv, grating_um, apply_bg_sub, - gate_start_ns=gate_start, gate_end_ns=gate_end, - saw_pipeline=self._saw_pipeline, - ) - 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) - self._compute_thread.start() + worker = FnWorker(lambda: compute_image( + scan, angle_idx, ch_idx, threshold_mv, apply_bg_sub, + gate_start, gate_end, pipeline)) + thread = QThread() + worker.moveToThread(thread) + thread.started.connect(worker.run) + worker.finished.connect(self._on_compute_done) + worker.error.connect(self._on_compute_error) + worker.finished.connect(thread.quit) + worker.error.connect(thread.quit) + thread.finished.connect(self._on_compute_thread_finished) + self._compute_worker = worker + self._compute_thread = thread + thread.start() def _on_compute_thread_finished(self): self._compute_thread = None - angle_idx = self.spin_angle.value() - ch_idx = self.combo_channel.currentIndex() - threshold_mv = self.spin_threshold_mv.value() - grating_um = self.spin_grating_um.value() - apply_bg_sub = self.chk_bg_sub.isChecked() - gate_enabled = self.chk_gate.isChecked() - if (angle_idx, ch_idx, threshold_mv, grating_um, apply_bg_sub, - gate_enabled, - self.spin_gate_start.value(), self.spin_gate_end.value()) != ( - self._pending_angle, self._pending_ch, - self._pending_threshold, self._pending_grating_um, - self._pending_bg_sub, - self._pending_gate_enabled, - self._pending_gate_start, self._pending_gate_end): + self._compute_worker = None + if self._compute_pending: + self._compute_pending = False self._start_compute() + def _on_compute_error(self, msg: str): + self._close_progress() + self.statusBar().showMessage(f"Compute error: {msg}") + def _on_compute_done(self, img: np.ndarray): self._close_progress() self._current_image = img - self._current_angle = self._pending_angle - self._current_ch = self._pending_ch - self.btn_export_csv.setEnabled(self._pending_ch in CH1_DERIVED_MODES) + self._current_angle = self._computing_angle + self._current_ch = self._computing_ch + self.btn_export_csv.setEnabled(self._current_ch in CH1_DERIVED_MODES) self._redraw_image(img) + # ------------------------------------------------------------------ + # Drawing + # ------------------------------------------------------------------ + + def _display_image(self, img: np.ndarray) -> np.ndarray: + """The image as displayed (velocity mode scales RF MHz by grating µm).""" + if self._current_ch == VELOCITY_MODE_IDX: + return img * self.spin_grating_um.value() + return img + def _redraw_image(self, img: np.ndarray): - s = self._sras - x_axis = s.x_axis_mm() - y_axis = s.y_positions_mm - dx = x_axis[1] - x_axis[0] if len(x_axis) > 1 else s.pixel_x_mm + scan = self._scan + sras = scan.sras + ai = self._current_angle + pa = sras.per_angle[ai] + + if img.size == 0: + self.statusBar().showMessage( + f"{sras.path.name} | angle {ai}: no data on disk") + return + + img = self._display_image(img) + + x_axis = sras.x_axis_mm(ai) + y_axis = np.asarray(pa.y_positions[:img.shape[0]]) + dx = x_axis[1] - x_axis[0] if len(x_axis) > 1 else sras.pixel_pitch_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[0] - dx / 2, x_axis[-1] + dx / 2, y_axis[-1] + dy / 2, - y_axis[0] - dy / 2, + y_axis[0] - dy / 2, ] if self.chk_auto.isChecked(): @@ -1721,32 +1312,12 @@ class SrasViewerWindow(QMainWindow): vmin = self.spin_vmin.value() vmax = self.spin_vmax.value() - ch_idx = self._current_ch - 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]" - elif ch_idx == SAW_MODE_AMP_IDX: - mode_str = "SAW-AMP" - unit = "MF envelope peak (arb.)" - colorbar_label = "amplitude" - elif ch_idx == SAW_MODE_TOF_IDX: - mode_str = "SAW-TOF" - unit = "SAW arrival time (ns)" - colorbar_label = "ns" - else: - mode_str = "DC" - unit = "DC mean (mV)" - colorbar_label = "mV" + ch_idx = self._current_ch + angle_deg = pa.angle_deg + mode_str, unit, colorbar_label = MODE_INFO[ch_idx] + ch_label = CH_LABELS[ch_idx] + if ch_idx == VELOCITY_MODE_IDX: + ch_label = f"Velocity [grating={self.spin_grating_um.value():.2f} µm]" title = f"{CH_NAMES[ch_idx]} | {mode_str} | {angle_deg:.1f}°" @@ -1759,7 +1330,7 @@ class SrasViewerWindow(QMainWindow): colorbar_label=colorbar_label, ) self.statusBar().showMessage( - f"{s.path.name} | {ch_label} @ {angle_deg:.1f}° " + f"{sras.path.name} | {ch_label} @ {angle_deg:.1f}° " f"| {img.shape[1]} × {img.shape[0]} px | {unit}" ) @@ -1768,103 +1339,106 @@ class SrasViewerWindow(QMainWindow): # ------------------------------------------------------------------ def _on_pixel_clicked(self, row_idx: int, frame_idx: int): - if self._sras is None or self._current_image is None: + if self._scan is None or self._current_image is None: return - self._last_row = row_idx + self._last_row = row_idx self._last_frame = frame_idx self.lbl_wave_hint.hide() ch_idx = self._current_ch if ch_idx in CH1_DERIVED_MODES: gate_enabled = self.chk_gate.isChecked() and ch_idx in (CH1_IDX, VELOCITY_MODE_IDX) self.wave_canvas.show_rf_waveform( - self._sras, self._current_angle, row_idx, frame_idx, + self._scan, self._current_angle, row_idx, frame_idx, apply_bg_sub=self.chk_bg_sub.isChecked(), gate_start_ns=self.spin_gate_start.value() if gate_enabled else None, - gate_end_ns=self.spin_gate_end.value() if gate_enabled else None, + gate_end_ns=self.spin_gate_end.value() if gate_enabled else None, ) else: self.wave_canvas.show_dc_waveform( - self._sras, self._current_angle, ch_idx, row_idx, frame_idx + self._scan, self._current_angle, ch_idx, row_idx, frame_idx ) - # Update diagnostics button availability has_template = (self._saw_pipeline is not None and self._saw_pipeline.template is not None) - self.btn_saw_diag.setEnabled( - self._sras is not None and has_template and True) + self.btn_saw_diag.setEnabled(has_template) # ------------------------------------------------------------------ # SAW pipeline management # ------------------------------------------------------------------ def _on_build_template_clicked(self): - if self._sras is None or self._template_thread is not None: + if self._scan is None or self._template_thread is not None: return - # (Re-)create pipeline with current settings - sr = self._sras.sample_rate_hz + scan = self._scan self._saw_pipeline = SawPipeline( - sample_rate_hz = sr, - emi_gate_ns = self.spin_saw_emi_ns.value(), - bp_lo_mhz = self.spin_saw_bp_lo.value(), - bp_hi_mhz = self.spin_saw_bp_hi.value(), - saw_window_ns = (self.spin_saw_win_start.value(), - self.spin_saw_win_end.value()), - decimate_enable = False, + sample_rate_hz=scan.sras.header.sample_rate, + emi_gate_ns=self.spin_saw_emi_ns.value(), + bp_lo_mhz=self.spin_saw_bp_lo.value(), + bp_hi_mhz=self.spin_saw_bp_hi.value(), + saw_window_ns=(self.spin_saw_win_start.value(), + self.spin_saw_win_end.value()), ) - + pipeline = self._saw_pipeline angle_idx = self.spin_angle.value() + n_shots_req = self.spin_saw_n_shots.value() + row_sel, frame_sel = self._last_row, self._last_frame + apply_bg = scan.background is not None and self.chk_bg_sub.isChecked() - if self._last_row is not None and self._last_frame is not None: - # Build from the single selected pixel's waveform - waveforms = self._sras.data[ - angle_idx, self._last_row, CH1_IDX, - self._last_frame:self._last_frame + 1, : - ].astype(np.float32) - n_shots = 1 - src_desc = f"selected pixel (row={self._last_row}, frame={self._last_frame})" + if row_sel is not None and frame_sel is not None: + src_desc = f"selected pixel (row={row_sel}, frame={frame_sel})" else: - # No pixel selected — sample N shots spread across the whole scan - n_shots = min(self.spin_saw_n_shots.value(), - self._sras.n_frames * self._sras.n_rows) - waveforms_all = self._sras.data[angle_idx, :, CH1_IDX, :, :].astype(np.float32) - waveforms_flat = waveforms_all.reshape(-1, waveforms_all.shape[-1]) - indices = np.linspace(0, len(waveforms_flat) - 1, n_shots, dtype=int) - waveforms = waveforms_flat[indices] - src_desc = f"{n_shots} shots (full scan)" + src_desc = f"{n_shots_req} shots (full scan)" + self._template_src_desc = src_desc - if (self._sras.background is not None and self.chk_bg_sub.isChecked()): - waveforms = waveforms - self._sras.background[np.newaxis, :] + def _build(): + view = scan.angle_view(angle_idx) + n_rows, _, n_frames, _ = view.shape + if n_rows == 0: + raise RuntimeError("This angle has no data on disk.") + if row_sel is not None and frame_sel is not None: + shots = view[row_sel, CH1_IDX, frame_sel:frame_sel + 1].astype(np.float32) + else: + # Fancy-index only the selected shots — never materialize the + # whole angle in float32. + n = min(n_shots_req, n_rows * n_frames) + flat = np.linspace(0, n_rows * n_frames - 1, n, dtype=int) + rows, frames = np.divmod(flat, n_frames) + shots = view[rows, CH1_IDX, frames].astype(np.float32) + if apply_bg: + shots -= scan.background + pipeline.build_template(shots) + return None self.btn_build_template.setEnabled(False) self.lbl_saw_status.setText(f"Building template from {src_desc}…") self._show_progress("Building SAW template…") - self._template_worker = TemplateBuildWorker(self._saw_pipeline, waveforms) - self._template_thread = QThread() - self._template_worker.moveToThread(self._template_thread) - self._template_thread.started.connect(self._template_worker.run) - self._template_worker.finished.connect(self._on_template_built) - self._template_worker.error.connect(self._on_template_error) - self._template_worker.finished.connect(self._template_thread.quit) - self._template_worker.error.connect(self._template_thread.quit) - self._template_thread.finished.connect( - lambda: setattr(self, "_template_thread", None)) - self._template_thread.start() + worker = FnWorker(_build) + thread = QThread() + worker.moveToThread(thread) + thread.started.connect(worker.run) + worker.finished.connect(self._on_template_built) + worker.error.connect(self._on_template_error) + worker.finished.connect(thread.quit) + worker.error.connect(thread.quit) + thread.finished.connect(self._on_template_thread_finished) + self._template_worker = worker + self._template_thread = thread + thread.start() - def _on_template_built(self): + def _on_template_thread_finished(self): + self._template_thread = None + self._template_worker = None + + def _on_template_built(self, _result=None): self._close_progress() self.btn_build_template.setEnabled(True) - if self._last_row is not None and self._last_frame is not None: - src = f"pixel row={self._last_row} frame={self._last_frame}" - else: - src = f"{self.spin_saw_n_shots.value()} shots" self.lbl_saw_status.setText( - f"Template ready ({src})\n" + f"Template ready ({self._template_src_desc})\n" f"EMI gate: {self.spin_saw_emi_ns.value():.0f} ns " f"BP: {self.spin_saw_bp_lo.value():.0f}–{self.spin_saw_bp_hi.value():.0f} MHz") self.lbl_saw_status.setStyleSheet("font-size: 11px; color: #44cc66;") - has_pixel = self._last_row is not None - self.btn_saw_diag.setEnabled(self._sras is not None and has_pixel) + self.btn_saw_diag.setEnabled(self._last_row is not None) self.btn_apply_mf.setEnabled(True) self.combo_mf_mode.setEnabled(True) @@ -1875,7 +1449,7 @@ class SrasViewerWindow(QMainWindow): self.lbl_saw_status.setStyleSheet("font-size: 11px; color: #e05030;") def _on_open_diagnostics(self): - if (self._sras is None or self._saw_pipeline is None or + if (self._scan is None or self._saw_pipeline is None or self._last_row is None): return if self._diag_window is not None: @@ -1885,7 +1459,7 @@ class SrasViewerWindow(QMainWindow): pass # C++ object already deleted (user closed the window) self._diag_window = None self._diag_window = SawDiagnosticWindow( - self._sras, self._saw_pipeline, + self._scan, self._saw_pipeline, self._current_angle, self._last_row, self._last_frame, parent=None, # free-floating window ) @@ -1896,7 +1470,7 @@ class SrasViewerWindow(QMainWindow): self._diag_window.show() def _on_apply_mf_clicked(self): - if self._sras is None or self._saw_pipeline is None: + if self._scan is None or self._saw_pipeline is None: return if self._saw_pipeline.template is None: self.statusBar().showMessage( @@ -1934,7 +1508,8 @@ class SrasViewerWindow(QMainWindow): # ------------------------------------------------------------------ def closeEvent(self, event): - for attr in ("_load_thread", "_compute_thread", "_template_thread"): + for attr in ("_load_thread", "_compute_thread", "_template_thread", + "_csv_thread"): t = getattr(self, attr, None) if t is not None: t.quit() diff --git a/sras_viewer_requirements.txt b/sras_viewer_requirements.txt index 05adf9a..d9e2e82 100755 --- a/sras_viewer_requirements.txt +++ b/sras_viewer_requirements.txt @@ -1,3 +1,4 @@ PyQt6==6.10.2 numpy==2.4.1 matplotlib==3.10.8 +scipy==1.16.3 diff --git a/tests/test_sras_analysis.py b/tests/test_sras_analysis.py new file mode 100644 index 0000000..17b3e7f --- /dev/null +++ b/tests/test_sras_analysis.py @@ -0,0 +1,142 @@ +"""core.sras_analysis + the viewer's LoadedScan/compute path, run headlessly +over the golden fixtures.""" +from pathlib import Path + +import numpy as np +import pytest + +from core.sras_analysis import ( + CH1_IDX, CH4_IDX, ChannelCalibration, FALLBACK_YMULT_MV, + SawPipeline, compute_dc_image, compute_rf_image, power_spectrum, +) + +GOLDEN = Path(__file__).parent / "golden" + + +def _loaded_scan(name="complete.sras"): + # Imported lazily: sras_viewer pulls in PyQt6/matplotlib. + from sras_viewer import LoadedScan + return LoadedScan(str(GOLDEN / name)) + + +def test_loaded_scan_basics(): + scan = _loaded_scan() + assert scan.rows_available == [3, 3] + assert scan.background is not None and len(scan.background) == 8 + assert len(scan.calib.ymult_mv) == 3 + view = scan.angle_view(0) + assert view.shape == (3, 3, 4, 8) + scan.close() + + +def test_loaded_scan_truncated_rows(): + scan = _loaded_scan("trunc_rowboundary_a1.sras") + assert scan.rows_available == [3, 2] + assert scan.angle_view(1).shape[0] == 2 + scan = _loaded_scan("header_only.sras") + assert scan.rows_available == [0, 0] + assert scan.angle_view(0).shape[0] == 0 + + +def test_compute_dc_image_matches_manual(): + scan = _loaded_scan() + view = scan.angle_view(0) + img = compute_dc_image(view, CH4_IDX) + manual = view[:, CH4_IDX].astype(np.float64).mean(axis=-1) + assert img.shape == (3, 4) + assert img.dtype == np.float32 + np.testing.assert_allclose(img, manual, rtol=1e-6) + + +def test_compute_rf_image_mask_and_values(): + scan = _loaded_scan() + view = scan.angle_view(0) + sras = scan.sras + f_mhz = sras.freq_axis_mhz(sras.header.samples_per_frame) + + # Threshold below everything: all pixels valid, values from the freq axis + img_all = compute_rf_image(view, scan.calib, f_mhz, dc_threshold_mv=-1e9) + assert img_all.shape == (3, 4) + assert set(np.unique(img_all)).issubset(set(f_mhz)) + + # Threshold above everything: fully masked, zero image, no FFT work + img_none = compute_rf_image(view, scan.calib, f_mhz, dc_threshold_mv=1e9) + assert not img_none.any() + + +def test_compute_rf_image_gate_zeroes_samples(): + scan = _loaded_scan() + view = scan.angle_view(0) + sras = scan.sras + f_mhz = sras.freq_axis_mhz(sras.header.samples_per_frame) + t_ns = sras.time_axis_ns() + img = compute_rf_image(view, scan.calib, f_mhz, dc_threshold_mv=-1e9, + gate_start_ns=float(t_ns[2]), gate_end_ns=float(t_ns[5]), + time_axis_ns=t_ns) + assert img.shape == (3, 4) + + +def test_calibration_roundtrip_and_fallback(): + cal = ChannelCalibration.from_preambles(["", "", ""]) + assert cal.ymult_mv == [FALLBACK_YMULT_MV] * 3 + assert cal.mv_to_adc(cal.adc_to_mv(42.0, 0), 0) == pytest.approx(42.0) + + cal2 = ChannelCalibration.from_preambles( + ["YMULT 1.0E-3;YOFF -10.0;YZERO 2.0E-3"]) + assert cal2.ymult_mv[0] == pytest.approx(1.0) + assert cal2.yoff_adc[0] == pytest.approx(-10.0) + assert cal2.yzero_mv[0] == pytest.approx(2.0) + + +def test_power_spectrum_dc_suppressed(): + x = np.ones(64, dtype=np.float32) * 5.0 + p = power_spectrum(x) + assert p[0] == 0.0 + assert not p[1:].any() + + +def test_saw_pipeline_finds_injected_packet(): + sr = 6.25e9 + n = 1250 + t = np.arange(n) / sr + rng = np.random.default_rng(42) + + def make_shot(delay_ns=200.0): + sig = rng.normal(0, 0.05, n).astype(np.float32) + packet = np.exp(-((t * 1e9 - delay_ns) / 25.0) ** 2) \ + * np.sin(2 * np.pi * 140e6 * t) + return sig + 3.0 * packet.astype(np.float32) + + pipe = SawPipeline(sr, emi_gate_ns=50.0, bp_lo_mhz=85.0, bp_hi_mhz=200.0, + saw_window_ns=(80.0, 350.0)) + pipe.build_template(np.stack([make_shot() for _ in range(10)])) + assert pipe.template is not None + + metrics = pipe.process_shot_metrics(make_shot()) + assert metrics["peak_time_ns"] == pytest.approx(200.0, abs=15.0) + assert metrics["snr"] > 3.0 + + # process_shot returns the same metrics plus the stage arrays + full = pipe.process_shot(make_shot()) + for key in ("raw", "gated", "filtered", "mf_output", "envelope"): + assert isinstance(full[key], np.ndarray) + + +def test_compute_saw_image_scalars_only(): + # Synthetic angle view: golden frames (8 samples) are too short for the + # 6th-order zero-phase bandpass; real records are >1000 samples. + from core.sras_analysis import compute_saw_image + rng = np.random.default_rng(1) + view = rng.integers(-40, 40, size=(3, 3, 4, 512), dtype=np.int8) + view[:, CH4_IDX] = 100 # every pixel passes the DC mask + calib = ChannelCalibration.from_preambles(["", "", ""]) + + pipe = SawPipeline(6.25e9, saw_window_ns=(10.0, 70.0)) + pipe.build_template(view[0, CH1_IDX, :4].astype(np.float32)) + img = compute_saw_image(view, calib, -1e9, pipe, "amplitude") + assert img.shape == (3, 4) + assert img.dtype == np.float32 + assert (img > 0).all() + + tof = compute_saw_image(view, calib, -1e9, pipe, "tof") + assert tof.shape == (3, 4)