Thomas Ales
2954 lines
120 KiB
Python
2954 lines
120 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
SRAS Scan File Viewer
|
||
PyQt6 application for visualizing channel data from .sras binary 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
|
||
|
||
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.
|
||
"""
|
||
|
||
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 scipy.fft as scipy_fft
|
||
|
||
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
|
||
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
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# FFT backend
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_pyfftw_available = False
|
||
try:
|
||
import pyfftw
|
||
pyfftw.interfaces.cache.enable()
|
||
_pyfftw_available = True
|
||
except ImportError:
|
||
pass
|
||
|
||
# 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)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# SAW signal processing pipeline
|
||
# ---------------------------------------------------------------------------
|
||
|
||
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
|
||
# ------------------------------------------------------------------
|
||
# 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:
|
||
n = len(proc)
|
||
nfft = 1 << (n + len(self.template) - 1).bit_length()
|
||
S = scipy_fft.rfft(proc.astype(np.float64), n=nfft, workers=1)
|
||
T = scipy_fft.rfft(self.template.astype(np.float64), n=nfft, workers=1)
|
||
mf_out = scipy_fft.irfft(S * np.conj(T), n=nfft, workers=1)[:n].astype(np.float32)
|
||
env = np.abs(hilbert(mf_out)).astype(np.float32)
|
||
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
|
||
|
||
CH_LABELS = [
|
||
"CH1 — RF (FFT peak freq)",
|
||
"CH3 — Bias A (DC mean)",
|
||
"CH4 — Bias B (DC mean)",
|
||
"CH1 — Velocity (SRAS)",
|
||
"CH1 — SAW Amplitude (matched filter)",
|
||
"CH1 — SAW Arrival time (matched filter)",
|
||
]
|
||
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
|
||
|
||
# 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
|
||
|
||
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/v3/v4 .sras file."""
|
||
|
||
def __init__(self, path: str):
|
||
self.path = Path(path)
|
||
self._parse()
|
||
|
||
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
|
||
|
||
if magic != b"SRAS":
|
||
raise ValueError(f"Bad magic bytes: {magic!r}")
|
||
if ver not in (2, 3, 4, 5):
|
||
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
|
||
|
||
# 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
|
||
|
||
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
|
||
|
||
# Detect v5 PREC tail: if ver == 5, the tail size is known once we
|
||
# know n_frames. We resolve this by computing n_frames from the
|
||
# waveform-only region. For a v5 file the writer stores the exact
|
||
# actual frame count in n_frames_hdr, so we trust it.
|
||
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.n_frames = actual_n_frames
|
||
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)
|
||
self.data = np.memmap(
|
||
str(self.path),
|
||
dtype=waveform_dtype,
|
||
mode="r",
|
||
offset=data_offset,
|
||
shape=waveform_shape,
|
||
)
|
||
|
||
self.angles_deg = angles
|
||
self.y_positions_mm = y_pos
|
||
|
||
# ---- 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
|
||
|
||
# ------------------------------------------------------------------
|
||
# v5 writer
|
||
# ------------------------------------------------------------------
|
||
|
||
def write_v5(self, dest_path: str,
|
||
freq_images: np.ndarray,
|
||
dc4_images: np.ndarray,
|
||
dc3_images: np.ndarray,
|
||
bg_sub_applied: bool,
|
||
progress_cb=None):
|
||
"""Write a v5 .sras file to *dest_path*.
|
||
|
||
Copies the raw waveform bytes verbatim from the current file,
|
||
bumps the version byte to 5, patches n_frames_hdr to the actual
|
||
frame count, then appends the PREC section.
|
||
|
||
*freq_images* / *dc4_images* / *dc3_images*:
|
||
shape (n_angles, n_rows, n_frames) float32.
|
||
|
||
*progress_cb*: optional callable(fraction: float) for UI updates.
|
||
"""
|
||
import shutil
|
||
|
||
dest = Path(dest_path)
|
||
src = self.path
|
||
|
||
# --- Copy the source file verbatim, then patch the header -------
|
||
shutil.copy2(str(src), str(dest))
|
||
|
||
waveform_bytes = (self.n_angles * self.n_rows * self.n_channels
|
||
* self.n_frames * self.samples_per_frame
|
||
* self.bytes_per_sample)
|
||
|
||
with open(str(dest), "r+b") as f:
|
||
# Patch version byte (offset 4 in the header struct)
|
||
f.seek(4)
|
||
f.write(struct.pack("B", 5))
|
||
|
||
# Patch n_frames_hdr (uint32, big-endian) with the actual count.
|
||
# Locate its offset: magic(4) + ver(1) + n_angles(2) + n_rows(2) = 9
|
||
# then x_start(4)+x_delta(4)+vel(4)+freq(4) = 16, total = 25
|
||
# then n_frames_hdr is at offset 25 as ">I" (4 bytes)
|
||
f.seek(25)
|
||
f.write(struct.pack(">I", self.n_frames))
|
||
|
||
# Truncate anything after the waveform data (e.g. old PREC tail)
|
||
# and seek to the append position.
|
||
waveform_end = self._data_offset_for_write()
|
||
f.seek(waveform_end + waveform_bytes)
|
||
f.truncate()
|
||
|
||
# --- Write PREC section -------------------------------------
|
||
n_stored = self.n_angles
|
||
flags = 0x01 if bg_sub_applied else 0x00
|
||
f.write(b"PREC")
|
||
f.write(struct.pack("BB", 1, flags))
|
||
f.write(struct.pack(">H", n_stored))
|
||
|
||
for aidx in range(n_stored):
|
||
if progress_cb is not None:
|
||
progress_cb(aidx / n_stored)
|
||
f.write(struct.pack(">H", aidx))
|
||
f.write(freq_images[aidx].astype(">f4").tobytes())
|
||
f.write(dc4_images[aidx].astype(">f4").tobytes())
|
||
f.write(dc3_images[aidx].astype(">f4").tobytes())
|
||
|
||
if progress_cb is not None:
|
||
progress_cb(1.0)
|
||
|
||
def _data_offset_for_write(self) -> int:
|
||
"""Return the file offset where waveform data starts (used by write_v5)."""
|
||
# Re-derive the offset by walking the header fields, since we do not
|
||
# persist data_offset as an attribute from _parse.
|
||
with open(self.path, "rb") as f:
|
||
fields = struct.unpack(HDR_FMT, f.read(HDR_SIZE))
|
||
ver = fields[1]
|
||
n_ch = fields[12]
|
||
n_bg_bytes = 0
|
||
|
||
with open(self.path, "rb") as f:
|
||
f.seek(HDR_SIZE)
|
||
f.read(self.n_angles * 4) # angles
|
||
f.read(self.n_rows * 4) # y_pos
|
||
if ver >= 3:
|
||
for _ in range(n_ch):
|
||
(length,) = struct.unpack(">H", f.read(2))
|
||
f.read(length)
|
||
if ver >= 4:
|
||
(n_bg,) = struct.unpack(">I", f.read(4))
|
||
f.read(n_bg)
|
||
return f.tell()
|
||
|
||
# ------------------------------------------------------------------
|
||
# Axes helpers
|
||
# ------------------------------------------------------------------
|
||
|
||
@property
|
||
def pixel_x_mm(self) -> float:
|
||
return self.velocity_mm_s / self.laser_freq_hz
|
||
|
||
def x_axis_mm(self) -> 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, 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)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
|
||
_CHUNK_ROWS = 32 # rows processed per FFT/DC chunk; tune for RAM vs. throughput
|
||
|
||
|
||
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 so the float32 working buffer is bounded to
|
||
``_CHUNK_ROWS × n_frames × spf × 4`` bytes regardless of scan size.
|
||
"""
|
||
n_rows, n_frames = sras.n_rows, sras.n_frames
|
||
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] = (
|
||
sras.data[angle_idx, r0:r1, ch_idx, :, :]
|
||
.astype(np.float32)
|
||
.mean(axis=-1)
|
||
)
|
||
return img
|
||
|
||
|
||
def compute_rf_image(sras: SrasFile, angle_idx: int,
|
||
dc_threshold_mv: float,
|
||
apply_bg_sub: bool = True,
|
||
gate_start_ns: float | None = None,
|
||
gate_end_ns: float | None = None,
|
||
n_fft: int | 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.
|
||
|
||
Fast path: if the file contains v5 precomputed peak-frequency images,
|
||
and time-gating and zero-padding are not active, and the bg-sub flag
|
||
matches, the stored images are used directly — no FFT is run.
|
||
|
||
Otherwise, data is processed in row chunks of ``_CHUNK_ROWS`` to bound
|
||
peak RAM to roughly ``_CHUNK_ROWS × n_frames × max(spf, n_fft) × 24``
|
||
bytes regardless of scan size.
|
||
"""
|
||
n_rows, n_frames = sras.n_rows, sras.n_frames
|
||
|
||
# ---- Fast path: v5 precomputed images ----------------------------------
|
||
can_use_precomputed = (
|
||
sras.precomputed_freq_mhz is not None
|
||
and gate_start_ns is None
|
||
and gate_end_ns is None
|
||
and n_fft is None # no custom zero-padding
|
||
and sras.precomputed_bg_sub == (apply_bg_sub and sras.background is not None)
|
||
)
|
||
if can_use_precomputed:
|
||
freq_img = sras.precomputed_freq_mhz[angle_idx].copy()
|
||
dc4_img = sras.precomputed_dc4_mv[angle_idx]
|
||
freq_img[dc4_img < dc_threshold_mv] = 0.0
|
||
return freq_img
|
||
|
||
# ---- Chunked FFT path --------------------------------------------------
|
||
freq_axis = sras.freq_axis_mhz(n_fft)
|
||
img = np.zeros((n_rows, n_frames), dtype=np.float32)
|
||
_n_workers = os.cpu_count() or 4
|
||
|
||
gate_keep: np.ndarray | None = None
|
||
if gate_start_ns is not None or gate_end_ns is not None:
|
||
t_ns = sras.time_axis_ns()
|
||
gate_keep = np.ones(len(t_ns), dtype=bool)
|
||
if gate_start_ns is not None:
|
||
gate_keep &= t_ns >= gate_start_ns
|
||
if gate_end_ns is not None:
|
||
gate_keep &= t_ns <= gate_end_ns
|
||
|
||
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)
|
||
dc4_raw = sras.data[angle_idx, r0:r1, CH4_IDX, :, :].astype(np.float32)
|
||
dc4_mv = 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_mv >= dc_threshold_mv # True = above threshold = run FFT
|
||
|
||
if not valid.any():
|
||
continue
|
||
|
||
waveforms = sras.data[angle_idx, r0:r1, CH1_IDX, :, :].astype(np.float32)
|
||
|
||
if apply_bg_sub and sras.background is not None:
|
||
waveforms -= sras.background # background is 1-D (spf,)
|
||
|
||
if gate_keep is not None:
|
||
waveforms[:, :, ~gate_keep] = 0.0
|
||
|
||
valid_waves = waveforms[valid] # (n_valid, spf)
|
||
del waveforms
|
||
|
||
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
|
||
|
||
|
||
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
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 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):
|
||
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 = 25,
|
||
apply_bg_sub: bool = True,
|
||
gate_start_ns: float | None = None,
|
||
gate_end_ns: float | None = None,
|
||
saw_pipeline: "SawPipeline | None" = None,
|
||
n_fft: int | 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
|
||
self._n_fft = n_fft
|
||
|
||
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,
|
||
n_fft=self._n_fft)
|
||
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,
|
||
n_fft=self._n_fft)
|
||
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)
|
||
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 __init__(self, pipeline: SawPipeline, waveforms: np.ndarray):
|
||
super().__init__()
|
||
self._pipeline = pipeline
|
||
self._waveforms = waveforms
|
||
|
||
def run(self):
|
||
try:
|
||
self._pipeline.build_template(self._waveforms)
|
||
self.finished.emit()
|
||
except Exception as exc:
|
||
self.error.emit(str(exc))
|
||
|
||
|
||
class PreprocessWorker(QObject):
|
||
"""Compute all-angle FFT and DC images and write a v5 file.
|
||
|
||
Emits ``progress(int)`` (0–100) as each angle completes and
|
||
``finished(str)`` with an empty string on success or an error message
|
||
on failure.
|
||
"""
|
||
progress = pyqtSignal(int) # 0–100
|
||
finished = pyqtSignal(str) # empty = success, else error message
|
||
|
||
def __init__(self, sras: SrasFile, dest_path: str, apply_bg_sub: bool):
|
||
super().__init__()
|
||
self._sras = sras
|
||
self._dest_path = dest_path
|
||
self._apply_bg_sub = apply_bg_sub
|
||
|
||
def run(self):
|
||
try:
|
||
sras = self._sras
|
||
n = sras.n_angles
|
||
freq_images = np.empty((n, sras.n_rows, sras.n_frames), dtype=np.float32)
|
||
dc4_images = np.empty_like(freq_images)
|
||
dc3_images = np.empty_like(freq_images)
|
||
|
||
for aidx in range(n):
|
||
# Compute raw peak-frequency (no DC-threshold masking yet)
|
||
freq_images[aidx] = compute_rf_image(
|
||
sras, aidx, dc_threshold_mv=-1e9, # mask nothing
|
||
apply_bg_sub=self._apply_bg_sub)
|
||
|
||
# DC images (ADC counts → mV)
|
||
dc4_images[aidx] = adc_to_mv(
|
||
compute_dc_image(sras, aidx, CH4_IDX),
|
||
sras.ch_ymult_mv[CH4_IDX],
|
||
sras.ch_yoff_adc[CH4_IDX],
|
||
sras.ch_yzero_mv[CH4_IDX])
|
||
dc3_images[aidx] = adc_to_mv(
|
||
compute_dc_image(sras, aidx, CH3_IDX),
|
||
sras.ch_ymult_mv[CH3_IDX],
|
||
sras.ch_yoff_adc[CH3_IDX],
|
||
sras.ch_yzero_mv[CH3_IDX])
|
||
|
||
self.progress.emit(int((aidx + 1) / n * 90))
|
||
|
||
bg_sub_flag = self._apply_bg_sub and sras.background is not None
|
||
sras.write_v5(
|
||
self._dest_path,
|
||
freq_images, dc4_images, dc3_images,
|
||
bg_sub_applied=bg_sub_flag,
|
||
progress_cb=lambda frac: self.progress.emit(90 + int(frac * 10)),
|
||
)
|
||
self.finished.emit("")
|
||
except Exception as exc:
|
||
self.finished.emit(str(exc))
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# ROI (free quadrilateral in data coordinates)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class RoiQuad:
|
||
"""Free quadrilateral defined in data coordinates (mm).
|
||
|
||
Stored as 4 corner points (shape (4, 2)) in CCW order: BL, BR, TR, TL.
|
||
Each corner can be positioned independently, allowing skewed /
|
||
non-orthogonal regions of interest. Because it lives in scan/data
|
||
coords it persists unchanged when the displayed channel/mode switches.
|
||
"""
|
||
|
||
def __init__(self, pts: np.ndarray):
|
||
"""pts : array-like, shape (4, 2)."""
|
||
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":
|
||
"""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)
|
||
|
||
def copy(self) -> "RoiQuad":
|
||
return RoiQuad(self._pts.copy())
|
||
|
||
def corners(self) -> np.ndarray:
|
||
"""World-coord corners, shape (4, 2), CCW: BL, BR, TR, TL."""
|
||
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:
|
||
"""Width and height of the axis-aligned bounding box, shape (2,)."""
|
||
return self._pts.max(axis=0) - self._pts.min(axis=0)
|
||
|
||
def contains(self, x: float, y: float) -> bool:
|
||
return bool(MplPath(self._pts).contains_point((x, y)))
|
||
|
||
def mask_for_grid(self, x_axis: np.ndarray,
|
||
y_axis: np.ndarray) -> np.ndarray:
|
||
"""Boolean mask (n_rows, n_frames) of pixels whose centres lie
|
||
inside the quadrilateral.
|
||
"""
|
||
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)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Matplotlib canvases
|
||
# ---------------------------------------------------------------------------
|
||
|
||
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
|
||
|
||
# Interaction state values
|
||
_IDLE = "idle"
|
||
_DRAW_NEW = "draw_new"
|
||
_MOVE = "move"
|
||
_DRAG_CORNER = "drag_corner"
|
||
|
||
# Hit tolerance (display pixels) for handles.
|
||
_HANDLE_PX = 12
|
||
_CLICK_THRESH_PX = 4 # releases within this of press count as a click
|
||
|
||
def __init__(self, parent=None):
|
||
fig = Figure(figsize=(7, 5), tight_layout=True)
|
||
self.ax = fig.add_subplot(111)
|
||
super().__init__(fig)
|
||
self.setParent(parent)
|
||
self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
|
||
self._extent = None
|
||
self._img_shape = None
|
||
|
||
# ROI state
|
||
self._roi: RoiQuad | None = None
|
||
self._roi_artists: list = []
|
||
self._state = self._IDLE
|
||
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.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)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Public API
|
||
# ------------------------------------------------------------------
|
||
|
||
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)
|
||
# Patches and lines are destroyed by figure.clf(); drop stale refs.
|
||
self._roi_artists = []
|
||
|
||
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)
|
||
|
||
self.ax.set_xlabel(xlabel)
|
||
self.ax.set_ylabel(ylabel)
|
||
self.ax.set_title(title)
|
||
|
||
# Re-draw the ROI (if any) on top of the fresh image so it persists
|
||
# unchanged across mode / angle / channel switches.
|
||
self._draw_roi()
|
||
self.draw()
|
||
|
||
def get_roi(self) -> RoiQuad | None:
|
||
return self._roi
|
||
|
||
def set_roi(self, roi: RoiQuad | None):
|
||
self._roi = roi.copy() if roi is not None else None
|
||
self._draw_roi()
|
||
self.draw_idle()
|
||
self.roi_changed.emit()
|
||
|
||
def clear_roi(self):
|
||
self._roi = None
|
||
self._remove_roi_artists()
|
||
self.draw_idle()
|
||
self.roi_changed.emit()
|
||
|
||
def start_drawing(self):
|
||
"""Arm the next click+drag on the image to create a new ROI,
|
||
replacing any existing one."""
|
||
self._draw_mode = True
|
||
self.setCursor(Qt.CursorShape.CrossCursor)
|
||
self.draw_mode_changed.emit(True)
|
||
|
||
def cancel_drawing(self):
|
||
if self._draw_mode:
|
||
self._draw_mode = False
|
||
self.setCursor(Qt.CursorShape.ArrowCursor)
|
||
self.draw_mode_changed.emit(False)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Rendering
|
||
# ------------------------------------------------------------------
|
||
|
||
def _remove_roi_artists(self):
|
||
for a in self._roi_artists:
|
||
try:
|
||
a.remove()
|
||
except (ValueError, AttributeError, NotImplementedError):
|
||
pass
|
||
self._roi_artists = []
|
||
|
||
def _draw_roi(self):
|
||
self._remove_roi_artists()
|
||
if self._roi is None or self.ax is None:
|
||
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)
|
||
|
||
# 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)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Hit testing (uses display pixels for handles, data coords for "inside")
|
||
# ------------------------------------------------------------------
|
||
|
||
def _hit_test(self, event) -> tuple[str, int | None] | None:
|
||
if self._roi is None or self.ax is None:
|
||
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])
|
||
|
||
for i in range(4):
|
||
if np.hypot(*(corners_disp[i] - click)) <= self._HANDLE_PX:
|
||
return ("corner", i)
|
||
|
||
if event.xdata is not None and event.ydata is not None:
|
||
if self._roi.contains(event.xdata, event.ydata):
|
||
return ("inside", None)
|
||
return None
|
||
|
||
# ------------------------------------------------------------------
|
||
# Mouse event handlers
|
||
# ------------------------------------------------------------------
|
||
|
||
def _on_press(self, event):
|
||
if event.inaxes is not self.ax or self._extent is None:
|
||
return
|
||
if event.button != 1: # only left mouse button
|
||
return
|
||
# If the matplotlib toolbar is in pan / zoom mode, let it handle
|
||
# the interaction instead of starting a ROI manipulation.
|
||
tb = getattr(self, "toolbar", None)
|
||
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_button = event.button
|
||
|
||
if self._draw_mode:
|
||
self._draw_previous = self._roi.copy() if self._roi else None
|
||
self._roi = RoiQuad.from_bbox(event.xdata, event.ydata,
|
||
event.xdata, event.ydata)
|
||
self._state = self._DRAW_NEW
|
||
self._draw_roi()
|
||
self.draw_idle()
|
||
return
|
||
|
||
hit = self._hit_test(event)
|
||
if hit is None:
|
||
self._state = self._IDLE
|
||
return
|
||
|
||
kind, idx = hit
|
||
self._snapshot = self._roi.copy()
|
||
|
||
if kind == "corner":
|
||
self._state = self._DRAG_CORNER
|
||
self._drag_corner_idx = idx
|
||
elif kind == "inside":
|
||
self._state = self._MOVE
|
||
self._move_anchor = (event.xdata, event.ydata)
|
||
|
||
def _on_motion(self, event):
|
||
if self._state == self._IDLE:
|
||
return
|
||
if event.xdata is None or event.ydata is None:
|
||
return
|
||
if event.inaxes is not self.ax:
|
||
return
|
||
|
||
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)
|
||
|
||
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])
|
||
|
||
elif self._state == self._DRAG_CORNER:
|
||
self._roi._pts[self._drag_corner_idx] = [event.xdata, event.ydata]
|
||
|
||
self._draw_roi()
|
||
self.draw_idle()
|
||
|
||
def _on_release(self, event):
|
||
if event.button != 1 and self._press_button != 1:
|
||
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
|
||
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._press_xy = self._press_pixel = None
|
||
self._press_button = None
|
||
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
|
||
|
||
|
||
class WaveformCanvas(FigureCanvasQTAgg):
|
||
def __init__(self, parent=None):
|
||
fig = Figure(figsize=(8, 3), tight_layout=True)
|
||
self.ax_wave = fig.add_subplot(121)
|
||
self.ax_right = fig.add_subplot(122)
|
||
super().__init__(fig)
|
||
self.setParent(parent)
|
||
self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
|
||
|
||
def show_rf_waveform(self, sras: SrasFile, 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.
|
||
|
||
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
|
||
waveform_plot = waveform - bg if bg is not None else waveform
|
||
|
||
self.ax_wave.cla()
|
||
self.ax_right.cla()
|
||
|
||
if bg is not None:
|
||
self.ax_wave.plot(t_ns, waveform, linewidth=0.5, color="#aaaaaa",
|
||
label="raw", zorder=1)
|
||
self.ax_wave.plot(t_ns, bg, linewidth=0.5, color="#e07030",
|
||
linestyle="--", label="background", zorder=2)
|
||
self.ax_wave.plot(t_ns, waveform_plot, linewidth=0.7, color="#4488cc",
|
||
label="subtracted", zorder=3)
|
||
self.ax_wave.legend(fontsize=7, loc="upper right")
|
||
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")
|
||
if gate_end_ns is not None:
|
||
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]
|
||
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")
|
||
|
||
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)",
|
||
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]
|
||
|
||
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
|
||
self.ax_right.plot(f_mhz, power_raw, linewidth=0.5, color="#aaaaaa",
|
||
label="raw FFT", zorder=1)
|
||
|
||
self.ax_right.plot(f_mhz, power_sub, linewidth=0.7, color="#4488cc",
|
||
label="subtracted FFT" if bg is not None else None, zorder=2)
|
||
self.ax_right.axvline(peak_mhz, color="tomato", linestyle="--",
|
||
linewidth=1.2, label=f"peak = {peak_mhz:.1f} MHz")
|
||
self.ax_right.set_xlabel("Frequency (MHz)")
|
||
self.ax_right.set_ylabel("Power (arb.)")
|
||
self.ax_right.set_title("FFT Power Spectrum")
|
||
self.ax_right.set_xlim(0, 500)
|
||
self.ax_right.legend(fontsize=8)
|
||
|
||
self.draw()
|
||
|
||
def show_dc_waveform(self, sras: SrasFile, 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()
|
||
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])
|
||
|
||
self.ax_wave.cla()
|
||
self.ax_right.cla()
|
||
|
||
self.ax_wave.plot(t_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}"
|
||
)
|
||
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",
|
||
ha="center", va="center",
|
||
transform=self.ax_right.transAxes, fontsize=11,
|
||
)
|
||
self.ax_right.set_axis_off()
|
||
|
||
self.draw()
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# SAW diagnostic window
|
||
# ---------------------------------------------------------------------------
|
||
|
||
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)
|
||
"""
|
||
|
||
# Zone colour constants (all panels use the same palette)
|
||
_C_EMI = "#e05030" # red — EMI gate
|
||
_C_NOISE = "#ccaa00" # amber — noise / inter-packet region
|
||
_C_SAW = "#30c060" # green — SAW window
|
||
|
||
def __init__(self, sras: SrasFile, pipeline: SawPipeline,
|
||
angle_idx: int, row_idx: int, frame_idx: int,
|
||
parent=None):
|
||
super().__init__(parent)
|
||
self.setAttribute(Qt.WidgetAttribute.WA_DeleteOnClose)
|
||
self.setWindowTitle(
|
||
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._angle_idx = angle_idx
|
||
self._row_idx = row_idx
|
||
self._frame_idx = frame_idx
|
||
|
||
central = QWidget()
|
||
self.setCentralWidget(central)
|
||
vl = QVBoxLayout(central)
|
||
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(
|
||
"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)
|
||
vl.addLayout(toolbar_row)
|
||
vl.addWidget(self._canvas)
|
||
|
||
self._plot()
|
||
|
||
# ------------------------------------------------------------------
|
||
# Zone shading helper — call on any time-domain axes
|
||
# ------------------------------------------------------------------
|
||
|
||
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,
|
||
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,
|
||
label=f"noise {emi_end_ns:.0f}–{s0:.0f} ns")
|
||
ax.axvspan(s0, min(s1, t_end), alpha=0.10, color=self._C_SAW,
|
||
label=f"SAW {s0:.0f}–{s1:.0f} ns")
|
||
if show_legend:
|
||
ax.legend(fontsize=7, loc="upper right")
|
||
|
||
# ------------------------------------------------------------------
|
||
|
||
def _on_refresh(self):
|
||
self._plot()
|
||
|
||
def _plot(self):
|
||
sras = self._sras
|
||
pipeline = self._pipeline
|
||
angle_idx = self._angle_idx
|
||
row_idx = self._row_idx
|
||
frame_idx = self._frame_idx
|
||
|
||
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.clf()
|
||
axes = fig.subplots(3, 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
|
||
|
||
# --- 1. Raw signal ---
|
||
ax_raw.plot(t_full, raw_adc, lw=0.6, color="#4488cc", zorder=3)
|
||
self._shade_time_zones(ax_raw, t_full, emi_end_ns, s0, s1, show_legend=True)
|
||
ax_raw.set_xlabel("Time (ns)")
|
||
ax_raw.set_ylabel("ADC counts")
|
||
ax_raw.set_title(f"1 — Raw signal (row={row_idx}, frame={frame_idx})")
|
||
|
||
# --- 2. After EMI gating ---
|
||
ax_gated.plot(t_full[:len(result["gated"])], result["gated"],
|
||
lw=0.6, color="#cc8833", zorder=3)
|
||
self._shade_time_zones(ax_gated, t_full, emi_end_ns, s0, s1)
|
||
ax_gated.set_xlabel("Time (ns)")
|
||
ax_gated.set_ylabel("Amplitude")
|
||
ax_gated.set_title("2 — After EMI gating (cosine taper)")
|
||
|
||
# --- 3. After bandpass ---
|
||
ax_filt.plot(t_full[:len(result["filtered"])], result["filtered"],
|
||
lw=0.6, color="#44aa44", zorder=3)
|
||
self._shade_time_zones(ax_filt, t_full, emi_end_ns, s0, s1)
|
||
ax_filt.set_xlabel("Time (ns)")
|
||
ax_filt.set_ylabel("Amplitude")
|
||
ax_filt.set_title(
|
||
f"3 — After bandpass ({pipeline.bp_lo_mhz:.0f}–{pipeline.bp_hi_mhz:.0f} MHz, "
|
||
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
|
||
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,
|
||
label="raw PSD", zorder=1)
|
||
ax_spec.plot(f_mhz, spec_filt, lw=0.8, color="#44aa44",
|
||
label="bandpass PSD", zorder=2)
|
||
ax_spec.axvspan(pipeline.bp_lo_mhz, pipeline.bp_hi_mhz,
|
||
alpha=0.14, color=self._C_SAW, label="passband", zorder=0)
|
||
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.legend(fontsize=7)
|
||
|
||
# --- 5. Matched filter output + envelope ---
|
||
ax_mf.plot(t_proc, result["mf_output"], lw=0.5, color="#8855cc",
|
||
alpha=0.55, label="MF output", zorder=3)
|
||
ax_mf.plot(t_proc, result["envelope"], lw=1.3, color="#cc4488",
|
||
label="envelope (Hilbert)", zorder=4)
|
||
if pipeline.template is not None:
|
||
ax_mf.axvline(result["peak_time_ns"], color="#ffaa00",
|
||
linestyle="--", lw=1.2, zorder=5,
|
||
label=f"peak {result['peak_time_ns']:.1f} ns")
|
||
self._shade_time_zones(ax_mf, t_proc, emi_end_ns, s0, s1)
|
||
ax_mf.set_xlabel("Time (ns)")
|
||
ax_mf.set_ylabel("Amplitude")
|
||
ax_mf.set_title(
|
||
f"5 — Matched filter | "
|
||
f"A = {result['peak_amplitude']:.3f} | "
|
||
f"SNR = {result['snr']:.1f} | "
|
||
f"t = {result['peak_time_ns']:.1f} ns")
|
||
ax_mf.legend(fontsize=7)
|
||
|
||
# --- 6. Shot-to-shot overlay ---
|
||
for env in overlays:
|
||
t_ov = np.arange(len(env)) / sr * 1e9
|
||
ax_over.plot(t_ov, env, lw=0.5, alpha=0.45, color="#cc4488", zorder=3)
|
||
self._shade_time_zones(ax_over, t_proc, emi_end_ns, s0, s1)
|
||
ax_over.set_xlabel("Time (ns)")
|
||
ax_over.set_ylabel("MF envelope amplitude")
|
||
ax_over.set_title(
|
||
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} | "
|
||
f"Arrival time: {result['peak_time_ns']:.2f} ns | "
|
||
f"SNR: {result['snr']:.1f} | "
|
||
f"Shot-to-shot peak-amp variance ({n_overlay} shots): {peak_var:.4g}",
|
||
ha="center", va="bottom", fontsize=9,
|
||
bbox=dict(boxstyle="round,pad=0.3", facecolor="#2a2a2a", alpha=0.85),
|
||
color="#e8e8e8",
|
||
)
|
||
self._canvas.draw()
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Main window
|
||
# ---------------------------------------------------------------------------
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# FFT Options dialog
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class FftOptionsDialog(QDialog):
|
||
"""Configure FFT backend and zero-padding.
|
||
|
||
Changes take effect only when the user clicks Apply. Cancel discards
|
||
all pending edits. The live 'frequency resolution' label updates as
|
||
the user adjusts the pad factor so they can see the trade-off before
|
||
committing.
|
||
"""
|
||
|
||
def __init__(self, parent=None, *,
|
||
current_backend: str,
|
||
current_pad_factor: int,
|
||
samples_per_frame: int | None,
|
||
sample_rate_hz: float | None,
|
||
grating_um: float):
|
||
super().__init__(parent)
|
||
self.setWindowTitle("FFT Options")
|
||
self.setModal(True)
|
||
self.setMinimumWidth(380)
|
||
|
||
self._samples_per_frame = samples_per_frame
|
||
self._sample_rate_hz = sample_rate_hz
|
||
self._grating_um = grating_um
|
||
|
||
layout = QVBoxLayout(self)
|
||
|
||
# ---- Backend ---------------------------------------------------
|
||
grp_backend = QGroupBox("FFT Backend")
|
||
bl = QVBoxLayout(grp_backend)
|
||
|
||
self._btn_numpy = QRadioButton(
|
||
"NumPy FFT (always available)")
|
||
self._btn_pyfftw = QRadioButton(
|
||
"pyFFTW (faster for large arrays)" if _pyfftw_available
|
||
else "pyFFTW (not installed — run: pip install pyfftw)")
|
||
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_pyfftw, id=1)
|
||
|
||
if current_backend == "pyfftw" and _pyfftw_available:
|
||
self._btn_pyfftw.setChecked(True)
|
||
else:
|
||
self._btn_numpy.setChecked(True)
|
||
|
||
bl.addWidget(self._btn_numpy)
|
||
bl.addWidget(self._btn_pyfftw)
|
||
layout.addWidget(grp_backend)
|
||
|
||
# ---- Zero-padding ----------------------------------------------
|
||
grp_zp = QGroupBox("Zero-Padding")
|
||
zl = QVBoxLayout(grp_zp)
|
||
|
||
pad_row = QHBoxLayout()
|
||
pad_row.addWidget(QLabel("Pad factor:"))
|
||
self._spin_pad = QSpinBox()
|
||
self._spin_pad.setRange(1, 256)
|
||
self._spin_pad.setValue(max(1, current_pad_factor))
|
||
self._spin_pad.setToolTip(
|
||
"Multiply the waveform length by this factor via zero-padding\n"
|
||
"before computing the FFT.\n"
|
||
"1 = no padding (natural length).\n"
|
||
"Powers of 2 (2, 4, 8 …) give the best performance."
|
||
)
|
||
self._spin_pad.valueChanged.connect(self._update_info)
|
||
pad_row.addWidget(self._spin_pad)
|
||
zl.addLayout(pad_row)
|
||
|
||
self._lbl_nfft = QLabel()
|
||
self._lbl_freq_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;")
|
||
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)
|
||
layout.addWidget(buttons)
|
||
|
||
self._update_info()
|
||
|
||
# ------------------------------------------------------------------
|
||
|
||
def _update_info(self):
|
||
spf = self._samples_per_frame
|
||
sr = self._sample_rate_hz
|
||
pad = self._spin_pad.value()
|
||
|
||
if spf is None or sr is None:
|
||
self._lbl_nfft.setText("Load a file to preview FFT parameters.")
|
||
self._lbl_freq_res.setText("")
|
||
self._lbl_vel_res.setText("")
|
||
return
|
||
|
||
n_fft = spf * pad
|
||
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
|
||
|
||
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(
|
||
f"Velocity bin: {vel_res_ms:.3f} m/s "
|
||
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"
|
||
|
||
def get_pad_factor(self) -> int:
|
||
return max(1, self._spin_pad.value())
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class SrasViewerWindow(QMainWindow):
|
||
def __init__(self, initial_path: str | None = None):
|
||
super().__init__()
|
||
self.setWindowTitle("SRAS Scan Viewer")
|
||
self.resize(1560, 840)
|
||
self.setAcceptDrops(True)
|
||
|
||
self._sras: SrasFile | None = None
|
||
self._current_image: np.ndarray | None = None
|
||
self._current_angle: int = 0
|
||
self._current_ch: int = 0
|
||
self._load_thread: QThread | None = None
|
||
self._compute_thread: QThread | None = None
|
||
self._pending_angle: int = 0
|
||
self._pending_ch: int = 0
|
||
self._pending_threshold: float = 50.0 # mV
|
||
self._pending_grating_um: float = 25 # µ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._progress_dlg: QProgressDialog | None = None
|
||
|
||
# FFT settings (configured via FFT Options dialog)
|
||
self._fft_pad_factor: int = 1 # 1 = no padding
|
||
self._pending_fft_pad_factor: int = 1
|
||
|
||
self._preprocess_thread: QThread | None = None
|
||
|
||
# SAW pipeline state
|
||
self._saw_pipeline: SawPipeline | None = None
|
||
self._template_thread: QThread | None = None
|
||
self._diag_window: SawDiagnosticWindow | None = None
|
||
self._last_row: int | None = None
|
||
self._last_frame: int | None = None
|
||
|
||
self._build_ui()
|
||
|
||
if initial_path:
|
||
self._load_file(initial_path)
|
||
|
||
# ------------------------------------------------------------------
|
||
# UI construction
|
||
# ------------------------------------------------------------------
|
||
|
||
def _build_ui(self):
|
||
central = QWidget()
|
||
self.setCentralWidget(central)
|
||
root = QHBoxLayout(central)
|
||
root.setContentsMargins(8, 8, 8, 8)
|
||
root.setSpacing(8)
|
||
|
||
# ---- Left control panel ----------------------------------------
|
||
panel = QWidget()
|
||
panel.setFixedWidth(260)
|
||
panel_layout = QVBoxLayout(panel)
|
||
panel_layout.setContentsMargins(0, 0, 0, 0)
|
||
panel_layout.setSpacing(6)
|
||
root.addWidget(panel)
|
||
|
||
# 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;")
|
||
fl.addWidget(self.btn_open)
|
||
fl.addWidget(self.lbl_filename)
|
||
panel_layout.addWidget(grp_file)
|
||
|
||
# Scan info
|
||
grp_info = QGroupBox("Scan Info")
|
||
il = QVBoxLayout(grp_info)
|
||
self._info = {}
|
||
for key in ("Angles", "Rows", "Frames / row", "Samples / frame",
|
||
"Sample rate", "X start", "Pixel Δx", "Laser freq"):
|
||
lbl = QLabel(f"{key}: —")
|
||
lbl.setWordWrap(True)
|
||
lbl.setStyleSheet("font-size: 11px;")
|
||
il.addWidget(lbl)
|
||
self._info[key] = lbl
|
||
# Frame count warning (hidden until needed)
|
||
self.lbl_frame_warn = QLabel("")
|
||
self.lbl_frame_warn.setWordWrap(True)
|
||
self.lbl_frame_warn.setStyleSheet("color: #e07000; font-size: 11px;")
|
||
il.addWidget(self.lbl_frame_warn)
|
||
panel_layout.addWidget(grp_info)
|
||
|
||
# View settings
|
||
grp_view = QGroupBox("View Settings")
|
||
vl = QVBoxLayout(grp_view)
|
||
|
||
# Angle
|
||
ar = QHBoxLayout()
|
||
ar.addWidget(QLabel("Angle:"))
|
||
self.spin_angle = QSpinBox()
|
||
self.spin_angle.setRange(0, 0)
|
||
self.spin_angle.setEnabled(False)
|
||
self.spin_angle.editingFinished.connect(self._on_view_changed)
|
||
self.lbl_angle_deg = QLabel("—")
|
||
ar.addWidget(self.spin_angle)
|
||
ar.addWidget(self.lbl_angle_deg)
|
||
vl.addLayout(ar)
|
||
|
||
# Channel
|
||
cr = QHBoxLayout()
|
||
cr.addWidget(QLabel("Channel:"))
|
||
self.combo_channel = QComboBox()
|
||
self.combo_channel.addItems(CH_LABELS)
|
||
self.combo_channel.setEnabled(False)
|
||
self.combo_channel.currentIndexChanged.connect(self._on_channel_changed)
|
||
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)
|
||
|
||
self.grp_threshold = QGroupBox("RF Mask Threshold (CH1 only)")
|
||
tl = QVBoxLayout(self.grp_threshold)
|
||
thr_row = QHBoxLayout()
|
||
thr_row.addWidget(QLabel("DC threshold:"))
|
||
self.spin_threshold_mv = QDoubleSpinBox()
|
||
self.spin_threshold_mv.setRange(-500.0, 500.0)
|
||
self.spin_threshold_mv.setDecimals(3)
|
||
self.spin_threshold_mv.setSingleStep(0.025)
|
||
self.spin_threshold_mv.setSuffix(" mV")
|
||
self.spin_threshold_mv.setValue(50.0)
|
||
self.spin_threshold_mv.setEnabled(False)
|
||
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;")
|
||
tl.addWidget(self.lbl_threshold_adc)
|
||
vl.addWidget(self.grp_threshold)
|
||
|
||
# Background subtraction (v4 files only)
|
||
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)."
|
||
)
|
||
self.chk_bg_sub.toggled.connect(self._on_bg_sub_toggled)
|
||
vl.addWidget(self.chk_bg_sub)
|
||
|
||
# Time gate (for FFT; CH1/velocity only)
|
||
self.grp_gate = QGroupBox("Time Gate (CH1 only)")
|
||
gl = QVBoxLayout(self.grp_gate)
|
||
self.chk_gate = QCheckBox("Enable time gate")
|
||
self.chk_gate.setChecked(False)
|
||
self.chk_gate.setEnabled(False)
|
||
self.chk_gate.setToolTip(
|
||
"Zero-out samples outside the specified time window before\n"
|
||
"computing the FFT (useful for isolating a specific acoustic packet)."
|
||
)
|
||
self.chk_gate.toggled.connect(self._on_gate_toggled)
|
||
gl.addWidget(self.chk_gate)
|
||
|
||
gate_start_row = QHBoxLayout()
|
||
gate_start_row.addWidget(QLabel("Start:"))
|
||
self.spin_gate_start = QDoubleSpinBox()
|
||
self.spin_gate_start.setRange(0.0, 100000.0)
|
||
self.spin_gate_start.setDecimals(1)
|
||
self.spin_gate_start.setSingleStep(10.0)
|
||
self.spin_gate_start.setSuffix(" ns")
|
||
self.spin_gate_start.setValue(50.0)
|
||
self.spin_gate_start.setEnabled(False)
|
||
self.spin_gate_start.editingFinished.connect(self._on_gate_changed)
|
||
gate_start_row.addWidget(self.spin_gate_start)
|
||
gl.addLayout(gate_start_row)
|
||
|
||
gate_end_row = QHBoxLayout()
|
||
gate_end_row.addWidget(QLabel("End:"))
|
||
self.spin_gate_end = QDoubleSpinBox()
|
||
self.spin_gate_end.setRange(0.0, 100000.0)
|
||
self.spin_gate_end.setDecimals(1)
|
||
self.spin_gate_end.setSingleStep(10.0)
|
||
self.spin_gate_end.setSuffix(" ns")
|
||
self.spin_gate_end.setValue(200.0)
|
||
self.spin_gate_end.setEnabled(False)
|
||
self.spin_gate_end.editingFinished.connect(self._on_gate_changed)
|
||
gate_end_row.addWidget(self.spin_gate_end)
|
||
gl.addLayout(gate_end_row)
|
||
|
||
# (grp_gate will be added to the right panel below)
|
||
|
||
# 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)."
|
||
)
|
||
self.btn_export_csv.clicked.connect(self._on_export_csv)
|
||
vl.addWidget(self.btn_export_csv)
|
||
|
||
panel_layout.addWidget(grp_view)
|
||
|
||
# ---- ROI (Region of Interest) ---------------------------------
|
||
grp_roi = QGroupBox("ROI (Region of Interest)")
|
||
rl = QVBoxLayout(grp_roi)
|
||
|
||
self.btn_draw_roi = QPushButton("Draw ROI")
|
||
self.btn_draw_roi.setCheckable(True)
|
||
self.btn_draw_roi.setEnabled(False)
|
||
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"
|
||
"The ROI is persistent across channels / modes / angles."
|
||
)
|
||
self.btn_draw_roi.toggled.connect(self._on_draw_roi_toggled)
|
||
rl.addWidget(self.btn_draw_roi)
|
||
|
||
self.btn_clear_roi = QPushButton("Clear ROI")
|
||
self.btn_clear_roi.setEnabled(False)
|
||
self.btn_clear_roi.clicked.connect(self._on_clear_roi)
|
||
rl.addWidget(self.btn_clear_roi)
|
||
|
||
self.btn_export_roi = QPushButton("Export ROI as CSV…")
|
||
self.btn_export_roi.setEnabled(False)
|
||
self.btn_export_roi.setToolTip(
|
||
"Save every pixel whose centre lies inside the ROI as CSV.\n"
|
||
"Columns: row, frame, x_mm, y_mm, value.\n"
|
||
"Corner coordinates of the quad are written in the file header."
|
||
)
|
||
self.btn_export_roi.clicked.connect(self._on_export_roi_csv)
|
||
rl.addWidget(self.btn_export_roi)
|
||
|
||
self.lbl_roi_center = QLabel("centroid: —")
|
||
self.lbl_roi_size = QLabel("bbox: —")
|
||
self.lbl_roi_npix = QLabel("pixels inside: —")
|
||
for lbl in (self.lbl_roi_center, self.lbl_roi_size, self.lbl_roi_npix):
|
||
lbl.setStyleSheet("font-size: 11px; color: #aaa;")
|
||
rl.addWidget(lbl)
|
||
|
||
panel_layout.addWidget(grp_roi)
|
||
|
||
# ---- SAW Pipeline panel ----------------------------------------
|
||
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()
|
||
self.spin_saw_emi_ns.setRange(1.0, 10000.0)
|
||
self.spin_saw_emi_ns.setDecimals(1)
|
||
self.spin_saw_emi_ns.setSingleStep(5.0)
|
||
self.spin_saw_emi_ns.setSuffix(" ns")
|
||
self.spin_saw_emi_ns.setValue(50.0)
|
||
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()
|
||
self.spin_saw_win_start.setRange(0.0, 100000.0)
|
||
self.spin_saw_win_start.setDecimals(1)
|
||
self.spin_saw_win_start.setSuffix(" ns")
|
||
self.spin_saw_win_start.setValue(80.0)
|
||
saw_win_row.addWidget(self.spin_saw_win_start)
|
||
saw_win_row.addWidget(QLabel("–"))
|
||
self.spin_saw_win_end = QDoubleSpinBox()
|
||
self.spin_saw_win_end.setRange(0.0, 100000.0)
|
||
self.spin_saw_win_end.setDecimals(1)
|
||
self.spin_saw_win_end.setSuffix(" ns")
|
||
self.spin_saw_win_end.setValue(350.0)
|
||
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()
|
||
self.spin_saw_bp_lo.setRange(1.0, 3000.0)
|
||
self.spin_saw_bp_lo.setDecimals(1)
|
||
self.spin_saw_bp_lo.setSuffix(" MHz")
|
||
self.spin_saw_bp_lo.setValue(85.0)
|
||
bp_row.addWidget(self.spin_saw_bp_lo)
|
||
bp_row.addWidget(QLabel("–"))
|
||
self.spin_saw_bp_hi = QDoubleSpinBox()
|
||
self.spin_saw_bp_hi.setRange(1.0, 3000.0)
|
||
self.spin_saw_bp_hi.setDecimals(1)
|
||
self.spin_saw_bp_hi.setSuffix(" MHz")
|
||
self.spin_saw_bp_hi.setValue(200.0)
|
||
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()
|
||
self.spin_saw_n_shots.setRange(1, 10000)
|
||
self.spin_saw_n_shots.setValue(50)
|
||
tmpl_row.addWidget(self.spin_saw_n_shots)
|
||
sl.addLayout(tmpl_row)
|
||
|
||
self.btn_build_template = QPushButton("Build Template")
|
||
self.btn_build_template.setEnabled(False)
|
||
self.btn_build_template.setToolTip(
|
||
"Average N shots (EMI-gated + bandpass-filtered) to form a\n"
|
||
"Hann-windowed template for the matched filter.")
|
||
self.btn_build_template.clicked.connect(self._on_build_template_clicked)
|
||
sl.addWidget(self.btn_build_template)
|
||
|
||
self.lbl_saw_status = QLabel("No template")
|
||
self.lbl_saw_status.setStyleSheet("font-size: 11px; color: #888;")
|
||
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;")
|
||
sl.addWidget(sep_mf)
|
||
|
||
sl.addWidget(QLabel("Apply matched filter:"))
|
||
|
||
mf_mode_row = QHBoxLayout()
|
||
mf_mode_row.addWidget(QLabel("Output:"))
|
||
self.combo_mf_mode = QComboBox()
|
||
self.combo_mf_mode.addItems(["Amplitude", "Time-of-Flight"])
|
||
self.combo_mf_mode.setEnabled(False)
|
||
mf_mode_row.addWidget(self.combo_mf_mode)
|
||
sl.addLayout(mf_mode_row)
|
||
|
||
self.btn_apply_mf = QPushButton("Apply Filter → Image")
|
||
self.btn_apply_mf.setEnabled(False)
|
||
self.btn_apply_mf.setToolTip(
|
||
"Switch to the SAW matched-filter channel and compute the image.\n"
|
||
"Requires a template to be built first.")
|
||
self.btn_apply_mf.clicked.connect(self._on_apply_mf_clicked)
|
||
sl.addWidget(self.btn_apply_mf)
|
||
|
||
panel_layout.addStretch()
|
||
|
||
# ---- Display Options group (added to right panel below) ------------
|
||
grp_display = QGroupBox("Display Options")
|
||
dl = QVBoxLayout(grp_display)
|
||
|
||
cmr = QHBoxLayout()
|
||
cmr.addWidget(QLabel("Colormap:"))
|
||
self.combo_cmap = QComboBox()
|
||
self.combo_cmap.addItems(CMAPS)
|
||
self.combo_cmap.setCurrentText("gray")
|
||
self.combo_cmap.setEnabled(False)
|
||
self.combo_cmap.currentIndexChanged.connect(self._on_view_changed)
|
||
cmr.addWidget(self.combo_cmap)
|
||
dl.addLayout(cmr)
|
||
|
||
self.chk_auto = QCheckBox("Auto-scale colormap")
|
||
self.chk_auto.setChecked(True)
|
||
self.chk_auto.toggled.connect(self._on_autoscale_toggled)
|
||
dl.addWidget(self.chk_auto)
|
||
|
||
for label, attr in (("min:", "spin_vmin"), ("max:", "spin_vmax")):
|
||
row = QHBoxLayout()
|
||
row.addWidget(QLabel(label))
|
||
spin = QDoubleSpinBox()
|
||
spin.setRange(-1e9, 1e9)
|
||
spin.setDecimals(4)
|
||
spin.setEnabled(False)
|
||
spin.editingFinished.connect(self._on_manual_range_changed)
|
||
setattr(self, attr, spin)
|
||
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()
|
||
self.btn_saw_diag = QPushButton("Open SAW Diagnostics…")
|
||
self.btn_saw_diag.setEnabled(False)
|
||
self.btn_saw_diag.setToolTip(
|
||
"Show the 6-panel SAW pipeline diagnostic for the clicked pixel.\n"
|
||
"Requires a SAW template to be built first.")
|
||
self.btn_saw_diag.clicked.connect(self._on_open_diagnostics)
|
||
wave_vl.addWidget(self.lbl_wave_hint)
|
||
wave_vl.addWidget(self.btn_saw_diag)
|
||
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_gate)
|
||
right_panel_layout.addWidget(self.grp_velocity)
|
||
right_panel_layout.addWidget(grp_display)
|
||
right_panel_layout.addWidget(grp_saw)
|
||
right_panel_layout.addStretch()
|
||
root.addWidget(right_panel)
|
||
|
||
self.statusBar().showMessage("Open an .sras file to begin.")
|
||
|
||
# ---- Menu bar ----------------------------------------------------------
|
||
menubar = self.menuBar()
|
||
fft_menu = menubar.addMenu("&FFT")
|
||
fft_act = QAction("FFT &Options…", self)
|
||
fft_act.setStatusTip("Configure FFT backend and zero-padding")
|
||
fft_act.triggered.connect(self._on_fft_options)
|
||
fft_menu.addAction(fft_act)
|
||
|
||
fft_menu.addSeparator()
|
||
self._preprocess_act = QAction("&Pre-process and Save as v5…", self)
|
||
self._preprocess_act.setStatusTip(
|
||
"Compute FFT and DC images for all angles and save to a v5 file "
|
||
"for instant re-opening (no FFT on load).")
|
||
self._preprocess_act.setEnabled(False)
|
||
self._preprocess_act.triggered.connect(self._on_preprocess)
|
||
fft_menu.addAction(self._preprocess_act)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Drag-and-drop
|
||
# ------------------------------------------------------------------
|
||
|
||
def dragEnterEvent(self, event):
|
||
urls = event.mimeData().urls()
|
||
if urls and urls[0].toLocalFile().lower().endswith(".sras"):
|
||
event.acceptProposedAction()
|
||
|
||
def dropEvent(self, event):
|
||
self._load_file(event.mimeData().urls()[0].toLocalFile())
|
||
|
||
# ------------------------------------------------------------------
|
||
# File loading
|
||
# ------------------------------------------------------------------
|
||
|
||
def _on_open(self):
|
||
path, _ = QFileDialog.getOpenFileName(
|
||
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
|
||
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()
|
||
|
||
def _on_load_done(self, sras):
|
||
self._close_progress()
|
||
self.btn_open.setEnabled(True)
|
||
if sras is None:
|
||
return
|
||
self._sras = sras
|
||
self._current_image = None
|
||
|
||
# A ROI from the previous file no longer matches the new scan's
|
||
# geometry, so discard it on every load.
|
||
self.image_canvas.clear_roi()
|
||
|
||
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")
|
||
if s.precomputed_freq_mhz is not None:
|
||
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")
|
||
self.lbl_frame_warn.setText("\n".join(notes))
|
||
|
||
self.spin_angle.blockSignals(True)
|
||
self.spin_angle.setRange(0, max(0, s.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._on_view_changed()
|
||
|
||
# ------------------------------------------------------------------
|
||
# 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)
|
||
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
|
||
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)
|
||
# Time gate only for legacy FFT modes (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
|
||
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)
|
||
# ROI: always usable once a file is loaded (independent of channel)
|
||
self.btn_draw_roi.setEnabled(has_file)
|
||
# Pre-process: available whenever a file is loaded and not already running
|
||
self._preprocess_act.setEnabled(has_file and self._preprocess_thread is 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
|
||
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()
|
||
|
||
def _on_bg_sub_toggled(self):
|
||
if self._sras is not None:
|
||
if self.combo_channel.currentIndex() in CH1_DERIVED_MODES:
|
||
self._start_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()
|
||
|
||
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()
|
||
|
||
def _on_grating_changed(self):
|
||
if self._sras is not None and self.combo_channel.currentIndex() == VELOCITY_MODE_IDX:
|
||
self._start_compute()
|
||
|
||
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"
|
||
)
|
||
path, _ = QFileDialog.getSaveFileName(
|
||
self, "Export Image as CSV",
|
||
str(self._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}")
|
||
|
||
# ------------------------------------------------------------------
|
||
# ROI (rectangle on the image)
|
||
# ------------------------------------------------------------------
|
||
|
||
def _on_draw_roi_toggled(self, checked: bool):
|
||
if checked:
|
||
self.image_canvas.start_drawing()
|
||
self.statusBar().showMessage(
|
||
"Click and drag on the image to draw a new rectangle.")
|
||
else:
|
||
self.image_canvas.cancel_drawing()
|
||
|
||
def _on_draw_mode_changed(self, active: bool):
|
||
# Keep the toggle button's visual state in sync with the canvas.
|
||
self.btn_draw_roi.blockSignals(True)
|
||
self.btn_draw_roi.setChecked(active)
|
||
self.btn_draw_roi.blockSignals(False)
|
||
|
||
def _on_roi_changed(self):
|
||
self._update_roi_ui()
|
||
|
||
def _update_roi_ui(self):
|
||
roi = self.image_canvas.get_roi()
|
||
if roi is None:
|
||
self.lbl_roi_center.setText("centroid: —")
|
||
self.lbl_roi_size.setText("bbox: —")
|
||
self.lbl_roi_npix.setText("pixels inside: —")
|
||
self.btn_clear_roi.setEnabled(False)
|
||
self.btn_export_roi.setEnabled(False)
|
||
return
|
||
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")
|
||
npix = 0
|
||
if self._sras is not None:
|
||
try:
|
||
mask = roi.mask_for_grid(self._sras.x_axis_mm(),
|
||
self._sras.y_positions_mm)
|
||
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()
|
||
y_axis = s.y_positions_mm
|
||
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")
|
||
if self._sras is not None and self.combo_channel.currentIndex() in CH1_DERIVED_MODES:
|
||
self._start_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)
|
||
|
||
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_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._start_compute()
|
||
|
||
# ------------------------------------------------------------------
|
||
# Computation
|
||
# ------------------------------------------------------------------
|
||
|
||
def _start_compute(self):
|
||
if self._sras is None:
|
||
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()
|
||
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
|
||
pad_factor = self._fft_pad_factor
|
||
n_fft = (self._sras.samples_per_frame * pad_factor
|
||
if pad_factor > 1 else None)
|
||
|
||
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._pending_fft_pad_factor = pad_factor
|
||
|
||
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,
|
||
n_fft=n_fft,
|
||
)
|
||
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()
|
||
|
||
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._fft_pad_factor) != (
|
||
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._pending_fft_pad_factor):
|
||
self._start_compute()
|
||
|
||
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._redraw_image(img)
|
||
self._update_roi_ui()
|
||
|
||
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
|
||
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(img.min()), float(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()
|
||
|
||
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"
|
||
|
||
title = f"{CH_NAMES[ch_idx]} | {mode_str} | {angle_deg:.1f}°"
|
||
|
||
self.image_canvas.show_image(
|
||
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"| {img.shape[1]} × {img.shape[0]} px | {unit}"
|
||
)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Pixel inspector
|
||
# ------------------------------------------------------------------
|
||
|
||
def _on_pixel_clicked(self, row_idx: int, frame_idx: int):
|
||
if self._sras is None or self._current_image is None:
|
||
return
|
||
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,
|
||
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,
|
||
)
|
||
else:
|
||
self.wave_canvas.show_dc_waveform(
|
||
self._sras, 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)
|
||
|
||
# ------------------------------------------------------------------
|
||
# SAW pipeline management
|
||
# ------------------------------------------------------------------
|
||
|
||
def _on_build_template_clicked(self):
|
||
if self._sras is None or self._template_thread is not None:
|
||
return
|
||
|
||
# (Re-)create pipeline with current settings
|
||
sr = self._sras.sample_rate_hz
|
||
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,
|
||
)
|
||
|
||
angle_idx = self.spin_angle.value()
|
||
|
||
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})"
|
||
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)"
|
||
|
||
if (self._sras.background is not None and self.chk_bg_sub.isChecked()):
|
||
waveforms = waveforms - self._sras.background[np.newaxis, :]
|
||
|
||
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()
|
||
|
||
def _on_template_built(self):
|
||
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"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_apply_mf.setEnabled(True)
|
||
self.combo_mf_mode.setEnabled(True)
|
||
|
||
def _on_template_error(self, msg: str):
|
||
self._close_progress()
|
||
self.btn_build_template.setEnabled(True)
|
||
self.lbl_saw_status.setText(f"Error: {msg}")
|
||
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
|
||
self._last_row is None):
|
||
return
|
||
if self._diag_window is not None:
|
||
try:
|
||
self._diag_window.close()
|
||
except RuntimeError:
|
||
pass # C++ object already deleted (user closed the window)
|
||
self._diag_window = None
|
||
self._diag_window = SawDiagnosticWindow(
|
||
self._sras, self._saw_pipeline,
|
||
self._current_angle, self._last_row, self._last_frame,
|
||
parent=None, # free-floating window
|
||
)
|
||
# Clear our reference when the user closes the window so we never
|
||
# call into a deleted C++ object again.
|
||
self._diag_window.destroyed.connect(
|
||
lambda: setattr(self, "_diag_window", None))
|
||
self._diag_window.show()
|
||
|
||
def _on_apply_mf_clicked(self):
|
||
if self._sras is None or self._saw_pipeline is None:
|
||
return
|
||
if self._saw_pipeline.template is None:
|
||
self.statusBar().showMessage(
|
||
"No template built yet — use 'Build Template' first.")
|
||
return
|
||
target_ch = (SAW_MODE_AMP_IDX
|
||
if self.combo_mf_mode.currentIndex() == 0
|
||
else SAW_MODE_TOF_IDX)
|
||
self.combo_channel.blockSignals(True)
|
||
self.combo_channel.setCurrentIndex(target_ch)
|
||
self.combo_channel.blockSignals(False)
|
||
self._on_channel_changed()
|
||
|
||
# ------------------------------------------------------------------
|
||
# Progress dialog helpers
|
||
# ------------------------------------------------------------------
|
||
|
||
def _show_progress(self, message: str):
|
||
if self._progress_dlg is not None:
|
||
self._progress_dlg.setLabelText(message)
|
||
return
|
||
dlg = QProgressDialog(message, "", 0, 0, self)
|
||
dlg.setWindowTitle("Please wait…")
|
||
dlg.setCancelButton(None)
|
||
dlg.setWindowModality(Qt.WindowModality.WindowModal)
|
||
dlg.setMinimumDuration(300) # only appears if operation takes > 300 ms
|
||
dlg.show()
|
||
self._progress_dlg = dlg
|
||
|
||
def _close_progress(self):
|
||
if self._progress_dlg is not None:
|
||
self._progress_dlg.close()
|
||
self._progress_dlg = None
|
||
|
||
# ------------------------------------------------------------------
|
||
# Pre-process → save v5
|
||
# ------------------------------------------------------------------
|
||
|
||
def _on_preprocess(self):
|
||
if self._sras is None:
|
||
return
|
||
s = self._sras
|
||
default_name = s.path.stem + "_v5" + s.path.suffix
|
||
dest, _ = QFileDialog.getSaveFileName(
|
||
self, "Save Pre-processed v5 File",
|
||
str(s.path.parent / default_name),
|
||
"SRAS files (*.sras);;All files (*)",
|
||
)
|
||
if not dest:
|
||
return
|
||
|
||
apply_bg = self.chk_bg_sub.isChecked() and s.background is not None
|
||
|
||
n_total = s.n_angles
|
||
n_px = s.n_rows * s.n_frames
|
||
approx_mb = n_total * n_px * 3 * 4 / 1e6
|
||
self._show_progress(
|
||
f"Pre-processing {n_total} angle(s) "
|
||
f"({s.n_rows}×{s.n_frames} px each, ~{approx_mb:.0f} MB output)…"
|
||
)
|
||
|
||
self._preprocess_worker = PreprocessWorker(s, dest, apply_bg)
|
||
self._preprocess_thread = QThread()
|
||
self._preprocess_worker.moveToThread(self._preprocess_thread)
|
||
self._preprocess_thread.started.connect(self._preprocess_worker.run)
|
||
self._preprocess_worker.progress.connect(self._on_preprocess_progress)
|
||
self._preprocess_worker.finished.connect(self._on_preprocess_done)
|
||
self._preprocess_worker.finished.connect(self._preprocess_thread.quit)
|
||
self._preprocess_thread.finished.connect(self._on_preprocess_thread_finished)
|
||
self._preprocess_act.setEnabled(False)
|
||
self._preprocess_thread.start()
|
||
|
||
def _on_preprocess_progress(self, pct: int):
|
||
if self._progress_dlg is not None:
|
||
self._progress_dlg.setValue(pct)
|
||
|
||
def _on_preprocess_done(self, error_msg: str):
|
||
self._close_progress()
|
||
if error_msg:
|
||
self.statusBar().showMessage(f"Pre-process failed: {error_msg}")
|
||
else:
|
||
self.statusBar().showMessage("v5 file written — re-open it for instant display.")
|
||
|
||
def _on_preprocess_thread_finished(self):
|
||
self._preprocess_thread = None
|
||
self._preprocess_act.setEnabled(self._sras is not None)
|
||
|
||
# ------------------------------------------------------------------
|
||
# 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_pad_factor=self._fft_pad_factor,
|
||
samples_per_frame=spf,
|
||
sample_rate_hz=sr,
|
||
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()
|
||
if self._sras is not None and self.combo_channel.currentIndex() in (
|
||
CH1_IDX, VELOCITY_MODE_IDX):
|
||
self._start_compute()
|
||
|
||
# ------------------------------------------------------------------
|
||
|
||
def closeEvent(self, event):
|
||
for attr in ("_load_thread", "_compute_thread",
|
||
"_template_thread", "_preprocess_thread"):
|
||
t = getattr(self, attr, None)
|
||
if t is not None:
|
||
t.quit()
|
||
t.wait(2000)
|
||
if self._diag_window is not None:
|
||
try:
|
||
self._diag_window.close()
|
||
except RuntimeError:
|
||
pass
|
||
super().closeEvent(event)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def main():
|
||
app = QApplication(sys.argv)
|
||
initial = sys.argv[1] if len(sys.argv) > 1 else None
|
||
window = SrasViewerWindow(initial_path=initial)
|
||
window.show()
|
||
sys.exit(app.exec())
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|