Thomas Ales
ed0eba4bae
Scans come off the scope as v6; a new trailing CACH section (independent SDCB/SFFT blocks, sized per-angle from the existing geometry table) lets computed DC and FFT images be cached in place, bumping the file to v7 the first time either is stored. The Convert menu's "Batch Compute DC and Store"/"Batch Compute FFT and Store" actions run this across multiple files in the background. compute_rf_image and the DC compute workers now reuse cached v7 data instead of recomputing it. Removes "Pre-process and Save as v5", which never supported v6 sources and is superseded by in-place v7 caching. Legacy v2-v5 reading, including v5's PREC section, is unaffected. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
3282 lines
140 KiB
Python
3282 lines
140 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.
|
||
|
||
Scan geometry: v6/v7 files scan a different bounding box per angle (x_start,
|
||
x_delta, n_frames, n_rows all vary by angle), so geometry is exposed per-angle
|
||
via SrasFile.n_rows / n_frames / x_start_mm arrays and the x_axis_mm() /
|
||
y_positions_mm() methods. v2–v5 files have uniform geometry across angles, so
|
||
those arrays simply repeat the same value n_angles times.
|
||
|
||
v7 files are v6 files with an optional trailing cache section holding
|
||
precomputed per-angle DC and/or FFT images (see scan_format.md), so display
|
||
never has to recompute them after the "Convert" menu's batch actions have
|
||
stored them once.
|
||
"""
|
||
|
||
import re
|
||
import sys
|
||
import struct
|
||
import faulthandler
|
||
import numpy as np
|
||
from pathlib import Path
|
||
from dataclasses import dataclass
|
||
import os
|
||
|
||
faulthandler.enable() # print a native stack trace on SIGSEGV/SIGABRT/etc.
|
||
|
||
from PyQt6.QtWidgets import (
|
||
QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
|
||
QGroupBox, QLabel, QPushButton, QComboBox, QSpinBox, QDoubleSpinBox,
|
||
QFileDialog, QSizePolicy, QSplitter, QCheckBox, QFrame, QProgressDialog,
|
||
QDialog, QDialogButtonBox, QRadioButton, QButtonGroup,
|
||
)
|
||
from PyQt6.QtGui import QAction
|
||
from PyQt6.QtCore import Qt, QThread, pyqtSignal, QObject
|
||
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
|
||
|
||
import scipy.fft as scipy_fft
|
||
import scipy.ndimage as scipy_ndimage
|
||
|
||
# Runtime-mutable settings changed via FftOptionsDialog
|
||
_fft_backend = "numpy" # "numpy" or "pyfftw"
|
||
|
||
|
||
def _do_rfft(x: np.ndarray, n: int | None = None, axis: int = -1,
|
||
workers: int = 1) -> np.ndarray:
|
||
"""Dispatch rfft to the selected backend with optional multithreading."""
|
||
if _fft_backend == "pyfftw" and _pyfftw_available:
|
||
return pyfftw.interfaces.numpy_fft.rfft(x, n=n, axis=axis, threads=workers)
|
||
return scipy_fft.rfft(x, n=n, axis=axis, workers=workers)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# SRAS format
|
||
# ---------------------------------------------------------------------------
|
||
|
||
# v2–v5: fixed header, uniform geometry across angles (43 bytes)
|
||
HDR_FMT = ">4sBHHffffIIdBB"
|
||
HDR_SIZE = struct.calcsize(HDR_FMT) # 43 bytes
|
||
|
||
# v6: fixed header, per-angle geometry in a separate table (49 bytes)
|
||
HDR_FMT_V6 = ">4sBHfffffffIdBB"
|
||
HDR_SIZE_V6 = struct.calcsize(HDR_FMT_V6) # 49 bytes
|
||
|
||
# v6: per-angle geometry table record (x_start, x_delta, n_frames, n_rows)
|
||
GEO_FMT_V6 = ">ffIH"
|
||
GEO_SIZE_V6 = struct.calcsize(GEO_FMT_V6) # 14 bytes
|
||
|
||
# v7: identical to v6 (same header/geometry/waveform layout, version byte
|
||
# is the only header difference) plus an optional trailing cache section
|
||
# ("CACH") holding precomputed per-angle DC and/or FFT images so a v7 file
|
||
# never needs to recompute them on open. See scan_format.md for the full
|
||
# spec. Sub-block sizes are per-angle (n_rows[a] * n_frames[a]), taken from
|
||
# the Per-Angle Geometry Table already parsed for v6 — no new geometry
|
||
# fields are needed.
|
||
CACH_MAGIC = b"CACH"
|
||
CACH_HDR_FMT = ">4sBB" # magic, cach_version, block_flags
|
||
CACH_HDR_SIZE = struct.calcsize(CACH_HDR_FMT) # 6 bytes
|
||
CACH_VERSION = 1
|
||
CACH_FLAG_DC = 0x01
|
||
CACH_FLAG_FFT = 0x02
|
||
|
||
SDCB_MAGIC = b"SDCB"
|
||
SDCB_HDR_FMT = ">4sBH" # magic, reserved, n_stored
|
||
SDCB_HDR_SIZE = struct.calcsize(SDCB_HDR_FMT) # 7 bytes
|
||
|
||
SFFT_MAGIC = b"SFFT"
|
||
SFFT_HDR_FMT = ">4sBH" # magic, flags, n_stored
|
||
SFFT_HDR_SIZE = struct.calcsize(SFFT_HDR_FMT) # 7 bytes
|
||
SFFT_FLAG_BG_SUB = 0x01
|
||
|
||
# Fixed-order channels in the file: index 0=CH1, 1=CH3, 2=CH4
|
||
# Fixed channel indices into the .sras data array (CH1=RF, CH3/CH4=Bias DC)
|
||
CH1_IDX, CH3_IDX, CH4_IDX = 0, 1, 2
|
||
|
||
CH_LABELS = [
|
||
"CH1 — RF (FFT peak freq)",
|
||
"CH3 — Bias A (DC mean)",
|
||
"CH4 — Bias B (DC mean)",
|
||
"CH1 — Velocity (SRAS)",
|
||
]
|
||
CH_NAMES = ["CH1", "CH3", "CH4", "VEL"]
|
||
|
||
# Combo index for the derived velocity mode (uses CH1_IDX data)
|
||
VELOCITY_MODE_IDX = 3
|
||
# All modes that operate on CH1 waveforms
|
||
CH1_DERIVED_MODES = (CH1_IDX, VELOCITY_MODE_IDX)
|
||
|
||
# Fallback scope calibration used only when reading v2 files without embedded
|
||
# preambles. v3+ files carry the WFMOutpre string so these are not used.
|
||
# 50 mV/div, 8 div full-scale, int8 ADC, position = -2.72 div
|
||
# ymult = 50 mV × 8 / 256 = 1.5625 mV/count
|
||
# yoff = position × (256/8) = -2.72 × 32 = -87.04 (ADC count for 0 V)
|
||
_FALLBACK_YMULT_MV = 1.5625 # mV per ADC count
|
||
_FALLBACK_YOFF_ADC = -87.04 # ADC count that represents 0 V
|
||
|
||
CMAPS = ["gray", "viridis", "plasma", "inferno", "hot", "jet", "RdBu_r", "seismic"]
|
||
|
||
|
||
def _parse_preamble(preamble: str) -> dict[str, float]:
|
||
"""Extract YMULT, YOFF, YZERO from a Tektronix WFMOutpre string.
|
||
|
||
Returns a dict with float values for whichever keys are present.
|
||
YMULT is left in V/count as the scope reports it.
|
||
"""
|
||
result = {}
|
||
for key in ("YMULT", "YOFF", "YZERO"):
|
||
m = re.search(rf'\b{key}\s+([-+]?\d*\.?\d+(?:[Ee][+-]?\d+)?)', preamble)
|
||
if m:
|
||
result[key] = float(m.group(1))
|
||
return result
|
||
|
||
|
||
def mv_to_adc(mv: float, ymult_mv: float = _FALLBACK_YMULT_MV,
|
||
yoff_adc: float = _FALLBACK_YOFF_ADC,
|
||
yzero_mv: float = 0.0) -> float:
|
||
return (mv - yzero_mv) / ymult_mv + yoff_adc
|
||
|
||
|
||
def adc_to_mv(adc: float, ymult_mv: float = _FALLBACK_YMULT_MV,
|
||
yoff_adc: float = _FALLBACK_YOFF_ADC,
|
||
yzero_mv: float = 0.0) -> float:
|
||
return (adc - yoff_adc) * ymult_mv + yzero_mv
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# File parser
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class SrasFile:
|
||
"""Parsed in-memory representation of a v2–v7 .sras file.
|
||
|
||
Scan geometry (rows, frames, x_start) is exposed per-angle via the
|
||
``n_rows`` / ``n_frames`` / ``x_start_mm`` arrays and the ``x_axis_mm()``
|
||
/ ``y_positions_mm()`` methods, since v6/v7 files scan a different
|
||
bounding box per angle. v2–v5 files have uniform geometry, so these
|
||
arrays just repeat the same value ``n_angles`` times. Waveform data is
|
||
likewise exposed as ``data[angle_idx]``, an array of shape
|
||
``(n_rows[a], n_channels, n_frames[a], samples_per_frame)``.
|
||
|
||
v7 files (v6 plus an optional trailing cache section) expose any stored
|
||
precomputed images as ``precomputed_dc3_mv`` / ``precomputed_dc4_mv`` /
|
||
``precomputed_freq_mhz`` — for v6/v7 these are ragged per-angle lists
|
||
(``list[np.ndarray | None]``, one entry per angle); for v5 they are dense
|
||
``(n_angles, n_rows, n_frames)`` arrays, since v5 geometry is uniform.
|
||
"""
|
||
|
||
def __init__(self, path: str):
|
||
self.path = Path(path)
|
||
self._parse()
|
||
|
||
def _parse(self):
|
||
with open(self.path, "rb") as f:
|
||
magic = f.read(4)
|
||
if magic != b"SRAS":
|
||
raise ValueError(f"Bad magic bytes: {magic!r}")
|
||
(version,) = struct.unpack(">B", f.read(1))
|
||
|
||
self.version = version
|
||
if version in (2, 3, 4, 5):
|
||
self._parse_legacy()
|
||
elif version in (6, 7):
|
||
self._parse_v6()
|
||
else:
|
||
raise ValueError(f"Unsupported version: {version}")
|
||
|
||
# ------------------------------------------------------------------
|
||
# v2–v5 parsing (uniform geometry, flat waveform block)
|
||
# ------------------------------------------------------------------
|
||
|
||
def _parse_legacy(self):
|
||
with open(self.path, "rb") as f:
|
||
fields = struct.unpack(HDR_FMT, f.read(HDR_SIZE))
|
||
(magic, ver, n_angles, n_rows, x_start, x_delta, vel, freq,
|
||
n_frames_hdr, spf, sr, bps, n_ch) = fields
|
||
|
||
self.n_angles = n_angles
|
||
self.velocity_mm_s = float(vel)
|
||
self.laser_freq_hz = float(freq)
|
||
self.n_frames_header = n_frames_hdr # configured count (may be wrong)
|
||
self.samples_per_frame = spf
|
||
self.sample_rate_hz = float(sr)
|
||
self.bytes_per_sample = bps
|
||
self.n_channels = n_ch
|
||
|
||
# Precomputed-image cache (populated when reading a v5 file).
|
||
# These are (n_angles, n_rows, n_frames) float32 arrays or None.
|
||
self.precomputed_freq_mhz: np.ndarray | None = None
|
||
self.precomputed_dc4_mv: np.ndarray | None = None
|
||
self.precomputed_dc3_mv: np.ndarray | None = None
|
||
self.precomputed_bg_sub: bool = False
|
||
self.scan_aborted = False
|
||
self.n_angles_declared = n_angles
|
||
|
||
with open(self.path, "rb") as f:
|
||
f.seek(HDR_SIZE)
|
||
angles = np.frombuffer(f.read(n_angles * 4), dtype=">f4").astype(np.float32)
|
||
y_pos = np.frombuffer(f.read(n_rows * 4), dtype=">f4").astype(np.float32)
|
||
|
||
if ver >= 3:
|
||
preambles = []
|
||
for _ in range(n_ch):
|
||
(length,) = struct.unpack(">H", f.read(2))
|
||
preambles.append(f.read(length).decode("utf-8"))
|
||
self.preambles = preambles
|
||
self.ch_ymult_mv = []
|
||
self.ch_yoff_adc = []
|
||
self.ch_yzero_mv = []
|
||
for p in preambles:
|
||
cal = _parse_preamble(p)
|
||
# YMULT from scope is V/count; store as mV/count
|
||
self.ch_ymult_mv.append(cal.get("YMULT", _FALLBACK_YMULT_MV / 1000) * 1000)
|
||
self.ch_yoff_adc.append(cal.get("YOFF", _FALLBACK_YOFF_ADC))
|
||
# YZERO from scope is in V; store as mV
|
||
self.ch_yzero_mv.append(cal.get("YZERO", 0.0) * 1000)
|
||
else:
|
||
self.preambles = None
|
||
self.ch_ymult_mv = [_FALLBACK_YMULT_MV] * n_ch
|
||
self.ch_yoff_adc = [_FALLBACK_YOFF_ADC] * n_ch
|
||
self.ch_yzero_mv = [0.0] * n_ch
|
||
|
||
if ver >= 4:
|
||
(n_bg,) = struct.unpack(">I", f.read(4))
|
||
self.background = np.frombuffer(f.read(n_bg), dtype=np.int8).astype(np.float32)
|
||
else:
|
||
self.background = None
|
||
|
||
# Record the byte offset where raw waveform data begins.
|
||
# np.memmap will use this to map only the waveform section.
|
||
data_offset = f.tell()
|
||
|
||
# ---- Determine actual frame count from file size ---------------
|
||
# For v4 and earlier the header n_frames may be the *configured*
|
||
# count before acquisition; the actual count is derived from the
|
||
# bytes on disk. For v5 files a PREC tail follows the waveform
|
||
# data, so we must not include those extra bytes in the frame count.
|
||
file_size = self.path.stat().st_size
|
||
samples_per_row_per_ch = n_ch * spf
|
||
|
||
# Upper bound: bytes from data_offset to end of file
|
||
available_bytes = file_size - data_offset
|
||
|
||
if ver == 5:
|
||
actual_n_frames = n_frames_hdr
|
||
remainder = 0
|
||
else:
|
||
total_samples = available_bytes // bps
|
||
actual_n_frames = total_samples // (n_angles * n_rows * samples_per_row_per_ch)
|
||
remainder = total_samples % (n_angles * n_rows * samples_per_row_per_ch)
|
||
|
||
self.frame_count_mismatch = (actual_n_frames != n_frames_hdr)
|
||
self.n_frames_remainder = remainder
|
||
|
||
# ---- Memory-map the waveform data (zero RAM cost) --------------
|
||
# Instead of f.read() → astype() (which peaks at 2× file size),
|
||
# memmap lets the OS page only the bytes that are actually touched.
|
||
waveform_dtype = np.int8 if bps == 1 else ">i2"
|
||
waveform_shape = (n_angles, n_rows, n_ch, actual_n_frames, spf)
|
||
data5d = np.memmap(
|
||
str(self.path),
|
||
dtype=waveform_dtype,
|
||
mode="r",
|
||
offset=data_offset,
|
||
shape=waveform_shape,
|
||
)
|
||
# Expose as a list of per-angle views so downstream code shares one
|
||
# indexing convention with v6: sras.data[a][row, ch, frame, sample]
|
||
self.data = [data5d[a] for a in range(n_angles)]
|
||
|
||
# Uniform per-angle geometry, repeated so callers don't need to
|
||
# special-case legacy vs. v6 files.
|
||
self.n_rows = np.full(n_angles, n_rows, dtype=np.int64)
|
||
self.n_frames = np.full(n_angles, actual_n_frames, dtype=np.int64)
|
||
self.x_start_mm = np.full(n_angles, float(x_start), dtype=np.float64)
|
||
self._y_pos_per_angle = [y_pos] * n_angles
|
||
|
||
self.angles_deg = angles
|
||
|
||
# ---- Read v5 precomputed section if present --------------------
|
||
if ver >= 5:
|
||
waveform_bytes = actual_n_frames * n_angles * n_rows * n_ch * spf * bps
|
||
prec_offset = data_offset + waveform_bytes
|
||
if file_size > prec_offset:
|
||
self._parse_prec_section(prec_offset, n_angles, n_rows, actual_n_frames)
|
||
|
||
def _parse_prec_section(self, offset: int,
|
||
n_angles: int, n_rows: int, n_frames: int):
|
||
"""Parse the v5 PREC tail that holds precomputed images."""
|
||
_PREC_MAGIC = b"PREC"
|
||
px = n_rows * n_frames # pixels per angle image
|
||
img_bytes = px * 4 # float32
|
||
|
||
with open(self.path, "rb") as f:
|
||
f.seek(offset)
|
||
header_raw = f.read(6) # magic(4) + fmt_ver(1) + flags(1)
|
||
if len(header_raw) < 6 or header_raw[:4] != _PREC_MAGIC:
|
||
return
|
||
flags = header_raw[5]
|
||
self.precomputed_bg_sub = bool(flags & 0x01)
|
||
|
||
(n_stored,) = struct.unpack(">H", f.read(2))
|
||
if n_stored == 0:
|
||
return
|
||
|
||
freq_buf = np.zeros((n_angles, n_rows, n_frames), dtype=np.float32)
|
||
dc4_buf = np.zeros((n_angles, n_rows, n_frames), dtype=np.float32)
|
||
dc3_buf = np.zeros((n_angles, n_rows, n_frames), dtype=np.float32)
|
||
|
||
for _ in range(n_stored):
|
||
(aidx,) = struct.unpack(">H", f.read(2))
|
||
if aidx >= n_angles:
|
||
break
|
||
freq_buf[aidx] = np.frombuffer(
|
||
f.read(img_bytes), dtype=">f4").reshape(n_rows, n_frames)
|
||
dc4_buf[aidx] = np.frombuffer(
|
||
f.read(img_bytes), dtype=">f4").reshape(n_rows, n_frames)
|
||
dc3_buf[aidx] = np.frombuffer(
|
||
f.read(img_bytes), dtype=">f4").reshape(n_rows, n_frames)
|
||
|
||
self.precomputed_freq_mhz = freq_buf
|
||
self.precomputed_dc4_mv = dc4_buf
|
||
self.precomputed_dc3_mv = dc3_buf
|
||
|
||
# ------------------------------------------------------------------
|
||
# v6 parsing (per-angle geometry, ragged waveform blocks)
|
||
# ------------------------------------------------------------------
|
||
|
||
def _parse_v6(self):
|
||
with open(self.path, "rb") as f:
|
||
fields = struct.unpack(HDR_FMT_V6, f.read(HDR_SIZE_V6))
|
||
(magic, ver, n_angles, x_start_nom, y_start_nom, x_delta_nom,
|
||
y_delta_nom, row_spacing, vel, freq, spf, sr, bps, n_ch) = fields
|
||
|
||
n_angles_declared = n_angles
|
||
|
||
self.velocity_mm_s = float(vel)
|
||
self.laser_freq_hz = float(freq)
|
||
self.samples_per_frame = spf
|
||
self.sample_rate_hz = float(sr)
|
||
self.bytes_per_sample = bps
|
||
self.n_channels = n_ch
|
||
|
||
# Reference-only fields: the ROI as entered before per-angle
|
||
# bounding-box expansion. Actual per-angle geometry used for
|
||
# rendering comes from the Per-Angle Geometry Table below.
|
||
self.x_start_nominal_mm = float(x_start_nom)
|
||
self.y_start_nominal_mm = float(y_start_nom)
|
||
self.x_delta_nominal_mm = float(x_delta_nom)
|
||
self.y_delta_nominal_mm = float(y_delta_nom)
|
||
self.row_spacing_mm = float(row_spacing)
|
||
|
||
self.n_frames_header = None
|
||
self.frame_count_mismatch = False
|
||
self.n_frames_remainder = 0
|
||
|
||
angles = np.frombuffer(f.read(n_angles * 4), dtype=">f4").astype(np.float32)
|
||
|
||
x_start = np.empty(n_angles, dtype=np.float64)
|
||
n_frames = np.empty(n_angles, dtype=np.int64)
|
||
n_rows = np.empty(n_angles, dtype=np.int64)
|
||
for a in range(n_angles):
|
||
xs, xd, nf, nr = struct.unpack(GEO_FMT_V6, f.read(GEO_SIZE_V6))
|
||
x_start[a] = xs
|
||
n_frames[a] = nf
|
||
n_rows[a] = nr
|
||
|
||
y_pos_per_angle = []
|
||
for a in range(n_angles):
|
||
nr = int(n_rows[a])
|
||
y_pos_per_angle.append(
|
||
np.frombuffer(f.read(nr * 4), dtype=">f4").astype(np.float32))
|
||
|
||
preambles = []
|
||
for _ in range(n_ch):
|
||
(length,) = struct.unpack(">H", f.read(2))
|
||
preambles.append(f.read(length).decode("utf-8"))
|
||
self.preambles = preambles
|
||
self.ch_ymult_mv = []
|
||
self.ch_yoff_adc = []
|
||
self.ch_yzero_mv = []
|
||
for p in preambles:
|
||
cal = _parse_preamble(p)
|
||
self.ch_ymult_mv.append(cal.get("YMULT", _FALLBACK_YMULT_MV / 1000) * 1000)
|
||
self.ch_yoff_adc.append(cal.get("YOFF", _FALLBACK_YOFF_ADC))
|
||
self.ch_yzero_mv.append(cal.get("YZERO", 0.0) * 1000)
|
||
|
||
(n_bg,) = struct.unpack(">I", f.read(4))
|
||
self.background = np.frombuffer(f.read(n_bg), dtype=np.int8).astype(np.float32)
|
||
|
||
data_offset = f.tell()
|
||
|
||
self._data_offset = data_offset
|
||
|
||
# ---- Memory-map each angle's ragged waveform block -------------
|
||
# v6 gives each angle its own row/frame count, so waveform data is
|
||
# no longer one uniform (n_angles, n_rows, ...) block — each angle's
|
||
# block sits at a different offset with its own shape. An aborted
|
||
# scan truncates the file mid-angle; per the format spec we keep
|
||
# whatever complete angles are present rather than refusing to open
|
||
# the file.
|
||
file_size = self.path.stat().st_size
|
||
waveform_dtype = np.int8 if bps == 1 else ">i2"
|
||
|
||
data = []
|
||
offset = data_offset
|
||
n_complete = 0
|
||
for a in range(n_angles):
|
||
nr = int(n_rows[a])
|
||
nf = int(n_frames[a])
|
||
nbytes = nr * n_ch * nf * spf * bps
|
||
if offset + nbytes > file_size:
|
||
break
|
||
data.append(np.memmap(
|
||
str(self.path), dtype=waveform_dtype, mode="r",
|
||
offset=offset, shape=(nr, n_ch, nf, spf),
|
||
))
|
||
offset += nbytes
|
||
n_complete += 1
|
||
|
||
if n_complete == 0:
|
||
raise ValueError(
|
||
"v6 file has no complete angle blocks — scan was aborted "
|
||
"before the first angle finished.")
|
||
|
||
self.data = data
|
||
self.n_angles = n_complete
|
||
self.n_angles_declared = n_angles_declared
|
||
self.scan_aborted = n_complete < n_angles_declared
|
||
self.angles_deg = angles[:n_complete]
|
||
self.x_start_mm = x_start[:n_complete]
|
||
self.n_frames = n_frames[:n_complete]
|
||
self.n_rows = n_rows[:n_complete]
|
||
self._y_pos_per_angle = y_pos_per_angle[:n_complete]
|
||
|
||
# Precomputed-image cache (v7 only). Ragged per-angle lists — unlike
|
||
# v5's dense (n_angles, n_rows, n_frames) arrays, v6/v7 geometry
|
||
# varies per angle, so each entry is its own (n_rows[a], n_frames[a])
|
||
# array or None if that angle isn't cached yet.
|
||
self.precomputed_freq_mhz: list[np.ndarray | None] = [None] * n_complete
|
||
self.precomputed_dc4_mv: list[np.ndarray | None] = [None] * n_complete
|
||
self.precomputed_dc3_mv: list[np.ndarray | None] = [None] * n_complete
|
||
self.precomputed_bg_sub: bool = False
|
||
|
||
if self.version == 7 and offset < file_size:
|
||
self._parse_cach_section(offset)
|
||
|
||
# ------------------------------------------------------------------
|
||
# v7 cache tail (CACH section: precomputed DC / FFT images)
|
||
# ------------------------------------------------------------------
|
||
|
||
def _cache_tail_offset(self) -> int:
|
||
"""Deterministic file offset where the CACH tail starts (or would
|
||
start), derived purely from the header + Per-Angle Geometry Table —
|
||
independent of whether a cache tail is actually present. Used by
|
||
both the parser and the in-place writer."""
|
||
n_ch = self.n_channels
|
||
spf = self.samples_per_frame
|
||
bps = self.bytes_per_sample
|
||
waveform_bytes = int(sum(
|
||
int(self.n_rows[a]) * n_ch * int(self.n_frames[a]) * spf * bps
|
||
for a in range(self.n_angles)
|
||
))
|
||
return self._data_offset + waveform_bytes
|
||
|
||
def _parse_cach_section(self, offset: int):
|
||
"""Parse the v7 CACH tail that holds precomputed DC/FFT images."""
|
||
with open(self.path, "rb") as f:
|
||
f.seek(offset)
|
||
header_raw = f.read(CACH_HDR_SIZE)
|
||
if len(header_raw) < CACH_HDR_SIZE:
|
||
return
|
||
magic, cach_version, block_flags = struct.unpack(CACH_HDR_FMT, header_raw)
|
||
if magic != CACH_MAGIC or cach_version != CACH_VERSION:
|
||
return
|
||
|
||
if block_flags & CACH_FLAG_DC:
|
||
sdcb_raw = f.read(SDCB_HDR_SIZE)
|
||
if len(sdcb_raw) < SDCB_HDR_SIZE:
|
||
return
|
||
sdcb_magic, _reserved, n_stored = struct.unpack(SDCB_HDR_FMT, sdcb_raw)
|
||
if sdcb_magic != SDCB_MAGIC:
|
||
return
|
||
for _ in range(n_stored):
|
||
(angle_idx,) = struct.unpack(">H", f.read(2))
|
||
if angle_idx >= self.n_angles:
|
||
break
|
||
px = int(self.n_rows[angle_idx]) * int(self.n_frames[angle_idx])
|
||
shape = (int(self.n_rows[angle_idx]), int(self.n_frames[angle_idx]))
|
||
self.precomputed_dc3_mv[angle_idx] = np.frombuffer(
|
||
f.read(px * 4), dtype=">f4").reshape(shape)
|
||
self.precomputed_dc4_mv[angle_idx] = np.frombuffer(
|
||
f.read(px * 4), dtype=">f4").reshape(shape)
|
||
|
||
if block_flags & CACH_FLAG_FFT:
|
||
sfft_raw = f.read(SFFT_HDR_SIZE)
|
||
if len(sfft_raw) < SFFT_HDR_SIZE:
|
||
return
|
||
sfft_magic, flags, n_stored = struct.unpack(SFFT_HDR_FMT, sfft_raw)
|
||
if sfft_magic != SFFT_MAGIC:
|
||
return
|
||
self.precomputed_bg_sub = bool(flags & SFFT_FLAG_BG_SUB)
|
||
for _ in range(n_stored):
|
||
(angle_idx,) = struct.unpack(">H", f.read(2))
|
||
if angle_idx >= self.n_angles:
|
||
break
|
||
px = int(self.n_rows[angle_idx]) * int(self.n_frames[angle_idx])
|
||
shape = (int(self.n_rows[angle_idx]), int(self.n_frames[angle_idx]))
|
||
self.precomputed_freq_mhz[angle_idx] = np.frombuffer(
|
||
f.read(px * 4), dtype=">f4").reshape(shape)
|
||
|
||
def write_v7_cache(self, *,
|
||
new_dc3_mv: list[np.ndarray | None] | None = None,
|
||
new_dc4_mv: list[np.ndarray | None] | None = None,
|
||
new_freq_mhz: list[np.ndarray | None] | None = None,
|
||
new_bg_sub: bool | None = None):
|
||
"""Store computed DC and/or FFT images into this file's CACH tail,
|
||
in place, converting a v6 source to v7 (or updating an existing v7
|
||
file). Only the block(s) passed in are recomputed; whichever block
|
||
isn't passed is carried forward unchanged from whatever this
|
||
``SrasFile`` already has in memory (from parsing, or a prior write
|
||
in this same session) — its bytes are never re-read from disk.
|
||
|
||
The waveform data itself is never touched: the cache tail always
|
||
starts at ``_cache_tail_offset()``, a fixed offset derived from the
|
||
header and geometry table alone.
|
||
"""
|
||
if self.version not in (6, 7):
|
||
raise ValueError(
|
||
f"write_v7_cache only supports v6/v7 source files, got v{self.version}")
|
||
|
||
final_dc3 = new_dc3_mv if new_dc3_mv is not None else self.precomputed_dc3_mv
|
||
final_dc4 = new_dc4_mv if new_dc4_mv is not None else self.precomputed_dc4_mv
|
||
final_freq = new_freq_mhz if new_freq_mhz is not None else self.precomputed_freq_mhz
|
||
final_bg_sub = new_bg_sub if new_bg_sub is not None else self.precomputed_bg_sub
|
||
|
||
dc_entries = [a for a in range(self.n_angles) if final_dc3[a] is not None]
|
||
fft_entries = [a for a in range(self.n_angles) if final_freq[a] is not None]
|
||
|
||
block_flags = 0
|
||
if dc_entries:
|
||
block_flags |= CACH_FLAG_DC
|
||
if fft_entries:
|
||
block_flags |= CACH_FLAG_FFT
|
||
|
||
payload = bytearray()
|
||
payload += struct.pack(CACH_HDR_FMT, CACH_MAGIC, CACH_VERSION, block_flags)
|
||
|
||
if dc_entries:
|
||
payload += struct.pack(SDCB_HDR_FMT, SDCB_MAGIC, 0, len(dc_entries))
|
||
for a in dc_entries:
|
||
payload += struct.pack(">H", a)
|
||
payload += final_dc3[a].astype(">f4").tobytes()
|
||
payload += final_dc4[a].astype(">f4").tobytes()
|
||
|
||
if fft_entries:
|
||
fft_flags = SFFT_FLAG_BG_SUB if final_bg_sub else 0
|
||
payload += struct.pack(SFFT_HDR_FMT, SFFT_MAGIC, fft_flags, len(fft_entries))
|
||
for a in fft_entries:
|
||
payload += struct.pack(">H", a)
|
||
payload += final_freq[a].astype(">f4").tobytes()
|
||
|
||
cache_offset = self._cache_tail_offset()
|
||
with open(self.path, "r+b") as f:
|
||
f.seek(cache_offset)
|
||
f.write(payload)
|
||
f.truncate()
|
||
f.flush()
|
||
os.fsync(f.fileno())
|
||
# Version-byte flip last: if the process dies before this point,
|
||
# the file is still readable as plain v6 (v6 parsing only
|
||
# bounds-checks per-angle offset+nbytes <= file_size, it never
|
||
# asserts exactly how many bytes follow the last angle) — so an
|
||
# interrupted write can never corrupt the file, only leave
|
||
# harmless trailing bytes that the next successful write
|
||
# overwrites via this same deterministic cache_offset.
|
||
f.seek(4)
|
||
f.write(struct.pack("B", 7))
|
||
f.flush()
|
||
os.fsync(f.fileno())
|
||
|
||
self.version = 7
|
||
self.precomputed_dc3_mv = final_dc3
|
||
self.precomputed_dc4_mv = final_dc4
|
||
self.precomputed_freq_mhz = final_freq
|
||
self.precomputed_bg_sub = final_bg_sub
|
||
|
||
# ------------------------------------------------------------------
|
||
# Axes helpers
|
||
# ------------------------------------------------------------------
|
||
|
||
@property
|
||
def pixel_x_mm(self) -> float:
|
||
return self.velocity_mm_s / self.laser_freq_hz
|
||
|
||
def x_axis_mm(self, angle_idx: int) -> np.ndarray:
|
||
n = int(self.n_frames[angle_idx])
|
||
return self.x_start_mm[angle_idx] + np.arange(n) * self.pixel_x_mm
|
||
|
||
def y_positions_mm(self, angle_idx: int) -> np.ndarray:
|
||
return self._y_pos_per_angle[angle_idx]
|
||
|
||
def time_axis_ns(self) -> np.ndarray:
|
||
return np.arange(self.samples_per_frame) / self.sample_rate_hz * 1e9
|
||
|
||
def freq_axis_mhz(self, n_fft: int | None = None) -> np.ndarray:
|
||
n = n_fft if n_fft is not None else self.samples_per_frame
|
||
return np.fft.rfftfreq(n, d=1.0 / self.sample_rate_hz) / 1e6
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Image computation (vectorised)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
# Rows are batched so the float32 working buffer for one channel's chunk
|
||
# (chunk_rows × n_frames × spf × 4 bytes) stays under this budget. A fixed
|
||
# row count (the original design) works fine for small legacy scans but is
|
||
# catastrophic for a v6 scan with a large per-angle frame/sample count —
|
||
# e.g. a 7500-frame × 2500-sample angle needs ~2.4 GB for a single 32-row
|
||
# chunk, times several such buffers alive at once for the FFT step, which
|
||
# can exceed physical RAM entirely on its own. Sizing the chunk to the
|
||
# actual dimensions keeps peak RAM bounded regardless of scan size.
|
||
_CHUNK_BYTES_BUDGET = 128 * 1024 * 1024 # ~128 MB per channel-buffer chunk
|
||
_CHUNK_ROWS_MAX = 32 # cap for small scans (old behavior)
|
||
|
||
|
||
def _chunk_rows_for(n_frames: int, samples_per_frame: int) -> int:
|
||
bytes_per_row = max(1, n_frames * samples_per_frame * 4) # float32
|
||
rows = _CHUNK_BYTES_BUDGET // bytes_per_row
|
||
return int(max(1, min(_CHUNK_ROWS_MAX, rows)))
|
||
|
||
|
||
def compute_dc_image(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray:
|
||
"""Mean of each waveform → (n_rows, n_frames) float32.
|
||
|
||
Processes in row chunks sized to a fixed memory budget (see
|
||
``_chunk_rows_for``) so the float32 working buffer stays bounded
|
||
regardless of scan size.
|
||
"""
|
||
n_rows = int(sras.n_rows[angle_idx])
|
||
n_frames = int(sras.n_frames[angle_idx])
|
||
data = sras.data[angle_idx]
|
||
chunk_rows = _chunk_rows_for(n_frames, sras.samples_per_frame)
|
||
img = np.empty((n_rows, n_frames), dtype=np.float32)
|
||
for r0 in range(0, n_rows, chunk_rows):
|
||
r1 = min(r0 + chunk_rows, n_rows)
|
||
img[r0:r1] = (
|
||
data[r0:r1, ch_idx, :, :]
|
||
.astype(np.float32)
|
||
.mean(axis=-1)
|
||
)
|
||
return img
|
||
|
||
|
||
def _cached_dc_mv(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray | None:
|
||
"""Return a v7-cached DC image (mV, already converted) for
|
||
(angle_idx, ch_idx) if the file's CACH section has it, else None.
|
||
|
||
Only applies to the ragged v6/v7 list representation of
|
||
``precomputed_dc3_mv``/``precomputed_dc4_mv`` — v5's dense PREC arrays
|
||
don't track per-angle validity (unstored angles are left at 0.0 rather
|
||
than a sentinel), so they're intentionally not consulted here.
|
||
"""
|
||
store = sras.precomputed_dc3_mv if ch_idx == CH3_IDX else sras.precomputed_dc4_mv
|
||
if isinstance(store, list) and angle_idx < len(store):
|
||
return store[angle_idx]
|
||
return None
|
||
|
||
|
||
def compute_rf_image(sras: SrasFile, angle_idx: int,
|
||
dc_threshold_mv: float,
|
||
apply_bg_sub: bool = True,
|
||
n_fft: int | None = None,
|
||
dc4_mv: np.ndarray | None = None) -> np.ndarray:
|
||
"""FFT of each CH1 waveform; pixel = peak frequency in MHz.
|
||
|
||
Pixels where CH4_dc < dc_threshold_mv are set to 0 — and the FFT is
|
||
never run for them, since that's the expensive part. The masked-out
|
||
CH1 samples are also never *read*: the boolean mask is applied to the
|
||
raw memmap slice before any dtype conversion, so numpy only pages in
|
||
the bytes for pixels that pass the threshold (an untouched memmap
|
||
page is never read from disk). If the DC4 image for this angle is
|
||
already known (e.g. from the DC-channel precompute cache), pass it as
|
||
*dc4_mv* (mV, shape (n_rows, n_frames)) to reuse it directly instead
|
||
of re-reading/re-averaging the CH4 channel here; otherwise it's
|
||
computed chunk-by-chunk internally (which does need every pixel's
|
||
CH4 data, since that's what determines validity in the first place).
|
||
|
||
Fast path: if the file contains precomputed peak-frequency images for
|
||
this angle (v5's dense per-file PREC section, or v7's ragged per-angle
|
||
CACH section), and zero-padding is not active, and the bg-sub flag
|
||
matches, the stored image is used directly — no FFT is run.
|
||
|
||
Otherwise, data is processed in row chunks sized to a fixed memory
|
||
budget (see ``_chunk_rows_for``) to bound peak RAM regardless of scan
|
||
size.
|
||
"""
|
||
n_rows = int(sras.n_rows[angle_idx])
|
||
n_frames = int(sras.n_frames[angle_idx])
|
||
data = sras.data[angle_idx]
|
||
|
||
# ---- Fast path: precomputed images (v5 PREC or v7 CACH) ----------------
|
||
# v5 stores a dense (n_angles, n_rows, n_frames) ndarray (uniform
|
||
# geometry); v6/v7 store a ragged list, one entry per angle (None where
|
||
# that angle hasn't been cached), since geometry varies per angle.
|
||
freq_store = sras.precomputed_freq_mhz
|
||
if isinstance(freq_store, list):
|
||
angle_cached = angle_idx < len(freq_store) and freq_store[angle_idx] is not None
|
||
else:
|
||
angle_cached = freq_store is not None
|
||
|
||
can_use_precomputed = (
|
||
angle_cached
|
||
and n_fft is None # no custom zero-padding
|
||
and sras.precomputed_bg_sub == (apply_bg_sub and sras.background is not None)
|
||
)
|
||
if can_use_precomputed:
|
||
freq_img = freq_store[angle_idx].copy()
|
||
|
||
# DC4 mask, in priority order: already-cached DC block, caller-
|
||
# supplied image, or a fresh (cheap — no FFT) recompute.
|
||
dc4_store = sras.precomputed_dc4_mv
|
||
if isinstance(dc4_store, list):
|
||
dc4_img = _cached_dc_mv(sras, angle_idx, CH4_IDX)
|
||
else:
|
||
dc4_img = dc4_store[angle_idx] if dc4_store is not None else None
|
||
if dc4_img is None:
|
||
if dc4_mv is not None:
|
||
dc4_img = dc4_mv
|
||
else:
|
||
dc4_img = adc_to_mv(compute_dc_image(sras, angle_idx, CH4_IDX),
|
||
sras.ch_ymult_mv[CH4_IDX], sras.ch_yoff_adc[CH4_IDX],
|
||
sras.ch_yzero_mv[CH4_IDX])
|
||
|
||
freq_img[dc4_img < dc_threshold_mv] = 0.0
|
||
return freq_img
|
||
|
||
# ---- Chunked FFT path --------------------------------------------------
|
||
freq_axis = sras.freq_axis_mhz(n_fft)
|
||
img = np.zeros((n_rows, n_frames), dtype=np.float32)
|
||
_n_workers = os.cpu_count() or 4
|
||
n_fft_bins = n_fft if n_fft is not None else sras.samples_per_frame
|
||
chunk_rows = _chunk_rows_for(n_frames, max(sras.samples_per_frame, n_fft_bins))
|
||
|
||
for r0 in range(0, n_rows, chunk_rows):
|
||
r1 = min(r0 + chunk_rows, n_rows)
|
||
|
||
# DC mask for this chunk (float32 expansion is only chunk-sized)
|
||
if dc4_mv is not None:
|
||
dc4_chunk = dc4_mv[r0:r1]
|
||
else:
|
||
dc4_raw = data[r0:r1, CH4_IDX, :, :].astype(np.float32)
|
||
dc4_chunk = adc_to_mv(dc4_raw.mean(axis=-1),
|
||
sras.ch_ymult_mv[CH4_IDX],
|
||
sras.ch_yoff_adc[CH4_IDX],
|
||
sras.ch_yzero_mv[CH4_IDX])
|
||
del dc4_raw
|
||
valid = dc4_chunk >= dc_threshold_mv # True = above threshold = run FFT
|
||
|
||
if not valid.any():
|
||
continue
|
||
|
||
# Index the raw memmap slice with the boolean mask *before*
|
||
# converting dtype — this is a lazy view until touched, so only
|
||
# the (n_valid, spf) selected elements are actually read from
|
||
# disk; masked-out pixels' pages are never paged in at all.
|
||
valid_waves = data[r0:r1, CH1_IDX, :, :][valid].astype(np.float32)
|
||
|
||
if apply_bg_sub and sras.background is not None:
|
||
valid_waves -= sras.background # background is 1-D (spf,)
|
||
|
||
fft_pow = np.abs(_do_rfft(valid_waves, n=n_fft, axis=-1, workers=_n_workers)) ** 2
|
||
del valid_waves
|
||
fft_pow[:, 0] = 0.0 # suppress DC bin
|
||
peak_bins = np.argmax(fft_pow, axis=-1)
|
||
del fft_pow
|
||
|
||
img[r0:r1][valid] = freq_axis[peak_bins]
|
||
|
||
return img
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Angle alignment (Fusion menu)
|
||
#
|
||
# Puts every angle's images onto one shared, zero-padded pixel grid using a
|
||
# rigid transform only (rotation + translation, never scale). Rotation for
|
||
# angle `a` is the *known* scan-angle delta relative to a reference angle —
|
||
# never searched. Only the residual translation is found, via FFT phase
|
||
# correlation of each angle's binarized CH4 ("dc-mask") image.
|
||
#
|
||
# Rotation is done in physical mm space rather than on raw pixel indices:
|
||
# the x-pixel pitch (SrasFile.pixel_x_mm) is file-wide constant but the
|
||
# y-pixel pitch (row spacing) can differ from it, and for v6 files can even
|
||
# vary per angle. Rotating the raw index grid directly would implicitly
|
||
# assume square pixels and shear a non-square-pixel image — an unwanted
|
||
# effective anisotropic scale. Instead each angle gets one affine that maps
|
||
# shared-canvas pixel index -> mm -> undo rotation/shift -> that angle's own
|
||
# local mm -> that angle's own raw pixel index, matching the output->input
|
||
# convention scipy.ndimage.affine_transform expects.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@dataclass
|
||
class AngleTransform:
|
||
rotation_deg: float
|
||
shift_mm: tuple[float, float] # (dx_mm, dy_mm) found by phase correlation
|
||
matrix: np.ndarray # (2,2): canvas (row,col) -> this angle's raw (row,col)
|
||
offset: np.ndarray # (2,)
|
||
|
||
|
||
@dataclass
|
||
class AlignmentResult:
|
||
ref_angle_idx: int
|
||
dc_threshold_mv: float
|
||
canvas_shape: tuple[int, int] # (n_rows, n_cols)
|
||
canvas_dx_mm: float
|
||
canvas_dy_mm: float
|
||
canvas_origin_mm: tuple[float, float] # mm at canvas pixel index (0, 0)
|
||
per_angle: dict[int, AngleTransform]
|
||
|
||
|
||
def _pixel_pitch_mm(sras: SrasFile, angle_idx: int) -> tuple[float, float]:
|
||
"""(dx, dy) mm/pixel for one angle: dx is the file-wide constant
|
||
pixel_x_mm; dy is this angle's own row spacing (assumed uniform, same
|
||
assumption _redraw_image already makes when it builds the display
|
||
extent)."""
|
||
dx = sras.pixel_x_mm
|
||
y = sras.y_positions_mm(angle_idx)
|
||
dy = float(y[1] - y[0]) if len(y) > 1 else 1.0
|
||
return dx, dy
|
||
|
||
|
||
def _bbox_center_mm(sras: SrasFile, angle_idx: int) -> tuple[float, float]:
|
||
x = sras.x_axis_mm(angle_idx)
|
||
y = sras.y_positions_mm(angle_idx)
|
||
return float((x[0] + x[-1]) / 2.0), float((y[0] + y[-1]) / 2.0)
|
||
|
||
|
||
def _bbox_corners_mm(sras: SrasFile, angle_idx: int) -> np.ndarray:
|
||
"""4 corners (x, y) of this angle's raw mm bounding box, shape (4, 2)."""
|
||
x = sras.x_axis_mm(angle_idx)
|
||
y = sras.y_positions_mm(angle_idx)
|
||
return np.array([[xx, yy] for xx in (x[0], x[-1]) for yy in (y[0], y[-1])])
|
||
|
||
|
||
def _rotation_matrix(theta_deg: float) -> np.ndarray:
|
||
t = np.radians(theta_deg)
|
||
c, s = np.cos(t), np.sin(t)
|
||
return np.array([[c, -s], [s, c]]) # CCW rotation acting on (x, y)
|
||
|
||
|
||
def _build_affine_canvas_to_raw(sras: SrasFile, angle_idx: int, ref_idx: int,
|
||
shift_mm: tuple[float, float],
|
||
canvas_dx: float, canvas_dy: float,
|
||
canvas_origin_mm: tuple[float, float]
|
||
) -> tuple[np.ndarray, np.ndarray]:
|
||
"""matrix, offset s.t. raw_index = matrix @ [row_out, col_out] + offset,
|
||
matching scipy.ndimage.affine_transform's output->input convention.
|
||
|
||
Pipeline (all mm unless noted):
|
||
[X;Y] = A_out @ [row_out;col_out] + b_out # canvas idx -> ref-frame mm
|
||
[lx;ly] = R(theta)^T @ ([X;Y]-c_ref-shift) + c_a # undo rotation+shift -> angle a's local mm
|
||
[row;col] = D @ ([lx;ly] - [x_start_a; y0_a]) # local mm -> angle a's raw idx
|
||
|
||
where theta = angles_deg[angle_idx] - angles_deg[ref_idx], c_ref/c_a are
|
||
each angle's own raw-bbox mm centroid (the rotation pivot — this keeps
|
||
rotated content centered, minimizing required canvas padding), and
|
||
A_out/D are the index<->mm scaling matrices for the canvas pitch and
|
||
this angle's own native pitch respectively.
|
||
"""
|
||
theta = float(sras.angles_deg[angle_idx] - sras.angles_deg[ref_idx])
|
||
Rinv = _rotation_matrix(theta).T
|
||
cx_a, cy_a = _bbox_center_mm(sras, angle_idx)
|
||
cx_ref, cy_ref = _bbox_center_mm(sras, ref_idx)
|
||
dx_a, dy_a = _pixel_pitch_mm(sras, angle_idx)
|
||
x0_a = float(sras.x_start_mm[angle_idx])
|
||
y0_a = float(sras.y_positions_mm(angle_idx)[0])
|
||
|
||
A_out = np.array([[0.0, canvas_dx], [canvas_dy, 0.0]]) # [row,col] -> [X,Y]
|
||
b_out = np.array(canvas_origin_mm, dtype=np.float64)
|
||
D = np.array([[0.0, 1.0 / dy_a], [1.0 / dx_a, 0.0]]) # [x,y] -> [row,col]
|
||
shift = np.array(shift_mm, dtype=np.float64)
|
||
c_ref_v = np.array([cx_ref, cy_ref])
|
||
c_a_v = np.array([cx_a, cy_a])
|
||
origin_a = np.array([x0_a, y0_a])
|
||
|
||
matrix = D @ Rinv @ A_out
|
||
offset = D @ Rinv @ (b_out - c_ref_v - shift) + D @ (c_a_v - origin_a)
|
||
return matrix, offset
|
||
|
||
|
||
def apply_alignment(result: AlignmentResult, angle_idx: int, img: np.ndarray,
|
||
order: int = 0) -> np.ndarray:
|
||
"""Resample any already-computed 2D image for `angle_idx` (same shape as
|
||
that angle's raw (n_rows, n_frames) — e.g. compute_dc_image /
|
||
compute_rf_image output) onto the shared alignment canvas. order=0
|
||
(nearest) avoids blending real data with zero-padding or with
|
||
masked-out (0-valued) CH1/velocity pixels at mask edges. Channel-
|
||
agnostic: the same per-angle transform (found from the CH4 mask) works
|
||
for any channel's image of that angle."""
|
||
t = result.per_angle[angle_idx]
|
||
return scipy_ndimage.affine_transform(
|
||
img.astype(np.float32, copy=False), t.matrix, offset=t.offset,
|
||
output_shape=result.canvas_shape, order=order,
|
||
mode="constant", cval=0.0)
|
||
|
||
|
||
def _block_mean_downsample(img: np.ndarray, factor: int) -> np.ndarray:
|
||
if factor <= 1:
|
||
return img
|
||
h, w = img.shape
|
||
h2, w2 = (h // factor) * factor, (w // factor) * factor
|
||
trimmed = img[:h2, :w2]
|
||
return trimmed.reshape(h2 // factor, factor, w2 // factor, factor).mean(axis=(1, 3))
|
||
|
||
|
||
def _phase_correlate_shift(ref_img: np.ndarray, mov_img: np.ndarray) -> tuple[int, int]:
|
||
"""FFT normalized cross-power-spectrum phase correlation. Returns the
|
||
integer (dr, dc) pixel shift of mov_img relative to ref_img; both must
|
||
be the same shape. Risk: if the true shift is near +/- half the array
|
||
size, wraparound can bias the peak — mitigated by the generous
|
||
margin_frac padding in _working_canvas_for_pair, which keeps the true
|
||
residual shift small relative to the correlation canvas."""
|
||
F1 = scipy_fft.fft2(ref_img.astype(np.float64))
|
||
F2 = scipy_fft.fft2(mov_img.astype(np.float64))
|
||
R = F1 * np.conj(F2)
|
||
R /= np.maximum(np.abs(R), 1e-12)
|
||
corr = scipy_fft.ifft2(R).real
|
||
dr, dc = np.unravel_index(np.argmax(corr), corr.shape)
|
||
h, w = corr.shape
|
||
if dr > h // 2:
|
||
dr -= h
|
||
if dc > w // 2:
|
||
dc -= w
|
||
return int(dr), int(dc)
|
||
|
||
|
||
def _working_canvas_for_pair(sras: SrasFile, ref_idx: int, a_idx: int,
|
||
dx: float, dy: float, margin_frac: float = 0.3
|
||
) -> tuple[tuple[float, float], tuple[int, int]]:
|
||
"""Union of the reference's own raw bbox and angle a's raw bbox rotated
|
||
(about its own center) into the ref frame with zero shift, padded by
|
||
margin_frac on each side — sized generously so the true phase-
|
||
correlation shift lands well inside the canvas (see
|
||
_phase_correlate_shift's wraparound note)."""
|
||
theta = float(sras.angles_deg[a_idx] - sras.angles_deg[ref_idx])
|
||
R = _rotation_matrix(theta)
|
||
c_a = np.array(_bbox_center_mm(sras, a_idx))
|
||
c_ref = np.array(_bbox_center_mm(sras, ref_idx))
|
||
pts = list(_bbox_corners_mm(sras, ref_idx))
|
||
for corner in _bbox_corners_mm(sras, a_idx):
|
||
pts.append(R @ (corner - c_a) + c_ref)
|
||
pts = np.array(pts)
|
||
x_min, y_min = pts.min(axis=0)
|
||
x_max, y_max = pts.max(axis=0)
|
||
pad_x, pad_y = (x_max - x_min) * margin_frac, (y_max - y_min) * margin_frac
|
||
x_min, x_max = x_min - pad_x, x_max + pad_x
|
||
y_min, y_max = y_min - pad_y, y_max + pad_y
|
||
n_cols = int(np.ceil((x_max - x_min) / dx)) + 1
|
||
n_rows = int(np.ceil((y_max - y_min) / abs(dy))) + 1
|
||
origin = (x_min, y_min if dy > 0 else y_max)
|
||
return origin, (n_rows, n_cols)
|
||
|
||
|
||
def _compute_angle_alignment(sras: SrasFile, ref_angle_idx: int,
|
||
dc_threshold_mv: float,
|
||
progress_cb=None) -> AlignmentResult:
|
||
"""Top-level alignment driver. Runs on a background thread (see
|
||
AngleAlignmentWorker) — deliberately recomputes CH4 DC images from
|
||
scratch via compute_dc_image rather than reading the GUI-thread
|
||
_dc_cache dict, since background-thread workers must not touch
|
||
GUI-thread-owned caches (BatchCacheWorker follows the same rule)."""
|
||
n = sras.n_angles # already the *complete*-angle count for aborted v6 scans
|
||
dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx)
|
||
|
||
# ---- Step 1: binarized CH4 mask per angle, native per-angle grid -----
|
||
masks: dict[int, np.ndarray] = {}
|
||
for a in range(n):
|
||
dc4 = adc_to_mv(compute_dc_image(sras, a, CH4_IDX),
|
||
sras.ch_ymult_mv[CH4_IDX], sras.ch_yoff_adc[CH4_IDX],
|
||
sras.ch_yzero_mv[CH4_IDX])
|
||
masks[a] = (dc4 >= dc_threshold_mv).astype(np.float32)
|
||
if progress_cb:
|
||
progress_cb(int((a + 1) / n * 25))
|
||
|
||
# ---- Step 2: coarse correlation stage (downsample first, then rotate)
|
||
# Downsampling before affine_transform (not after) is what keeps this
|
||
# tractable for a v6 scan with thousands of rows/frames per angle.
|
||
max_dim = max(max(m.shape) for m in masks.values())
|
||
factor = max(1, int(np.ceil(max_dim / 1024)))
|
||
dx_c, dy_c = dx_ref * factor, dy_ref * factor
|
||
masks_small = {a: _block_mean_downsample(m, factor) for a, m in masks.items()}
|
||
|
||
shifts_mm: dict[int, tuple[float, float]] = {ref_angle_idx: (0.0, 0.0)}
|
||
for a in range(n):
|
||
if a == ref_angle_idx:
|
||
continue
|
||
work_origin, work_shape = _working_canvas_for_pair(
|
||
sras, ref_angle_idx, a, dx_c, dy_c, margin_frac=0.3)
|
||
|
||
# Coarse canvas->raw affine at native pitch, then rescale by /factor
|
||
# so it maps coarse-canvas idx -> coarse (downsampled) raw idx —
|
||
# exact for block-mean downsampling (up to the trimmed remainder).
|
||
m_a, o_a = _build_affine_canvas_to_raw(
|
||
sras, a, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin)
|
||
m_ref, o_ref = _build_affine_canvas_to_raw(
|
||
sras, ref_angle_idx, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin)
|
||
rotated_a = scipy_ndimage.affine_transform(
|
||
masks_small[a], m_a / factor, offset=o_a / factor,
|
||
output_shape=work_shape, order=0, mode="constant", cval=0.0)
|
||
embedded_ref = scipy_ndimage.affine_transform(
|
||
masks_small[ref_angle_idx], m_ref / factor, offset=o_ref / factor,
|
||
output_shape=work_shape, order=0, mode="constant", cval=0.0)
|
||
|
||
dr, dc = _phase_correlate_shift(embedded_ref, rotated_a)
|
||
shifts_mm[a] = (dc * dx_c, dr * dy_c)
|
||
if progress_cb:
|
||
progress_cb(25 + int((a + 1) / n * 50))
|
||
|
||
# ---- Step 3: union bounding box over all angles (rotation+shift applied)
|
||
corners_ref_frame = []
|
||
for a in range(n):
|
||
theta = float(sras.angles_deg[a] - sras.angles_deg[ref_angle_idx])
|
||
R = _rotation_matrix(theta)
|
||
c_a = np.array(_bbox_center_mm(sras, a))
|
||
c_ref = np.array(_bbox_center_mm(sras, ref_angle_idx))
|
||
shift = np.array(shifts_mm[a])
|
||
for corner in _bbox_corners_mm(sras, a):
|
||
corners_ref_frame.append(R @ (corner - c_a) + c_ref + shift)
|
||
corners_ref_frame = np.array(corners_ref_frame)
|
||
x_min, y_min = corners_ref_frame.min(axis=0)
|
||
x_max, y_max = corners_ref_frame.max(axis=0)
|
||
n_cols = int(np.ceil((x_max - x_min) / dx_ref)) + 1
|
||
n_rows = int(np.ceil((y_max - y_min) / abs(dy_ref))) + 1
|
||
canvas_origin_mm = (float(x_min), float(y_min if dy_ref > 0 else y_max))
|
||
canvas_shape = (n_rows, n_cols)
|
||
|
||
# ---- Step 4: final per-angle full-resolution affine (canvas -> raw idx)
|
||
per_angle: dict[int, AngleTransform] = {}
|
||
for a in range(n):
|
||
matrix, offset = _build_affine_canvas_to_raw(
|
||
sras, a, ref_angle_idx, shifts_mm[a], dx_ref, dy_ref, canvas_origin_mm)
|
||
theta = float(sras.angles_deg[a] - sras.angles_deg[ref_angle_idx])
|
||
per_angle[a] = AngleTransform(theta, shifts_mm[a], matrix, offset)
|
||
if progress_cb:
|
||
progress_cb(75 + int((a + 1) / n * 25))
|
||
|
||
return AlignmentResult(ref_angle_idx, dc_threshold_mv, canvas_shape,
|
||
dx_ref, dy_ref, canvas_origin_mm, per_angle)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Background workers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class LoadWorker(QObject):
|
||
finished = pyqtSignal(object) # SrasFile | None
|
||
error = pyqtSignal(str)
|
||
|
||
def __init__(self, path: str):
|
||
super().__init__()
|
||
self._path = path
|
||
|
||
def run(self):
|
||
try:
|
||
self.finished.emit(SrasFile(self._path))
|
||
except Exception as exc:
|
||
self.error.emit(str(exc))
|
||
self.finished.emit(None)
|
||
|
||
|
||
class ComputeWorker(QObject):
|
||
"""Computes one displayable image for (angle, channel).
|
||
|
||
For CH1/Velocity (FFT-derived) channels, the FFT is only run for
|
||
pixels whose DC4 (Bias B) mean is at or above dc_threshold_mv — masked
|
||
pixels are left at 0 MHz without ever being FFT'd, since that's the
|
||
expensive part of a scan. If the DC4 image for this angle is already
|
||
known (e.g. from the DC-channel precompute cache), pass it in as
|
||
*dc4_mv* to skip re-reading the CH4 channel from disk entirely.
|
||
|
||
Emits a plain ``np.ndarray`` (already in display units, masked for
|
||
CH1/Velocity) for both DC and FFT-derived channels.
|
||
"""
|
||
finished = pyqtSignal(object)
|
||
error = pyqtSignal(str)
|
||
|
||
def __init__(self, sras: SrasFile, angle_idx: int,
|
||
ch_idx: int,
|
||
apply_bg_sub: bool = True,
|
||
n_fft: int | None = None,
|
||
dc_threshold_mv: float = 0.0,
|
||
dc4_mv: np.ndarray | None = None):
|
||
super().__init__()
|
||
self._sras = sras
|
||
self._angle = angle_idx
|
||
self._ch = ch_idx
|
||
self._apply_bg_sub = apply_bg_sub
|
||
self._n_fft = n_fft
|
||
self._dc_threshold = dc_threshold_mv
|
||
self._dc4_mv = dc4_mv
|
||
|
||
def run(self):
|
||
try:
|
||
if self._ch in (CH1_IDX, VELOCITY_MODE_IDX):
|
||
img = compute_rf_image(
|
||
self._sras, self._angle, dc_threshold_mv=self._dc_threshold,
|
||
apply_bg_sub=self._apply_bg_sub, n_fft=self._n_fft,
|
||
dc4_mv=self._dc4_mv)
|
||
self.finished.emit(img)
|
||
else:
|
||
# DC channels: use the v7 cache if this angle is already
|
||
# stored (already mV), else compute fresh and convert.
|
||
cached = _cached_dc_mv(self._sras, self._angle, self._ch)
|
||
if cached is not None:
|
||
self.finished.emit(cached)
|
||
return
|
||
adc_img = compute_dc_image(self._sras, self._angle, self._ch)
|
||
img = adc_to_mv(adc_img,
|
||
self._sras.ch_ymult_mv[self._ch],
|
||
self._sras.ch_yoff_adc[self._ch],
|
||
self._sras.ch_yzero_mv[self._ch])
|
||
self.finished.emit(img)
|
||
except Exception as exc:
|
||
self.error.emit(str(exc))
|
||
|
||
|
||
class DcPrecomputeWorker(QObject):
|
||
"""Computes CH3/CH4 DC images for every angle in the background.
|
||
|
||
DC images are cheap (a per-waveform mean, no FFT) compared to the
|
||
CH1/Velocity FFT, so precomputing them for the whole file right after
|
||
load makes switching angles instant while on a DC channel, and also
|
||
means the FFT masking step (which needs a DC4 image) rarely has to
|
||
wait on anything. Emits one ``angle_done`` signal per angle as it
|
||
completes rather than waiting for the whole file, so the cache fills
|
||
in progressively.
|
||
"""
|
||
angle_done = pyqtSignal(int, np.ndarray, np.ndarray) # angle_idx, dc3_mv, dc4_mv
|
||
finished = pyqtSignal()
|
||
error = pyqtSignal(str)
|
||
|
||
def __init__(self, sras: SrasFile):
|
||
super().__init__()
|
||
self._sras = sras
|
||
self._stop = False
|
||
|
||
def stop(self):
|
||
self._stop = True
|
||
|
||
def run(self):
|
||
try:
|
||
for a in range(self._sras.n_angles):
|
||
if self._stop:
|
||
break
|
||
cached_dc3 = _cached_dc_mv(self._sras, a, CH3_IDX)
|
||
cached_dc4 = _cached_dc_mv(self._sras, a, CH4_IDX)
|
||
if cached_dc3 is not None and cached_dc4 is not None:
|
||
self.angle_done.emit(a, cached_dc3, cached_dc4)
|
||
continue
|
||
dc3_mv = adc_to_mv(
|
||
compute_dc_image(self._sras, a, CH3_IDX),
|
||
self._sras.ch_ymult_mv[CH3_IDX],
|
||
self._sras.ch_yoff_adc[CH3_IDX],
|
||
self._sras.ch_yzero_mv[CH3_IDX])
|
||
dc4_mv = adc_to_mv(
|
||
compute_dc_image(self._sras, a, CH4_IDX),
|
||
self._sras.ch_ymult_mv[CH4_IDX],
|
||
self._sras.ch_yoff_adc[CH4_IDX],
|
||
self._sras.ch_yzero_mv[CH4_IDX])
|
||
self.angle_done.emit(a, dc3_mv, dc4_mv)
|
||
self.finished.emit()
|
||
except Exception as exc:
|
||
self.error.emit(str(exc))
|
||
|
||
|
||
class BatchCacheWorker(QObject):
|
||
"""Batch-computes and stores DC or FFT images into each of *paths*'s
|
||
v7 CACH tail, in place — converting v6 sources to v7 on first use, or
|
||
updating an existing v7 file's cache blocks without disturbing whatever
|
||
the other block already holds.
|
||
|
||
*mode* is ``"dc"`` (CH3/CH4 mean images) or ``"fft"`` (CH1 peak-frequency
|
||
images, unmasked — masking is applied at display time, same as v5's
|
||
PREC convention).
|
||
|
||
Emits ``progress(int)`` (0–100, weighted by total angle count across the
|
||
whole batch), ``file_done(str, str)`` (path, error message or "" on
|
||
success) after each file so one file's failure doesn't abort the batch,
|
||
and ``finished()`` once every file has been attempted.
|
||
"""
|
||
progress = pyqtSignal(int)
|
||
file_done = pyqtSignal(str, str)
|
||
finished = pyqtSignal()
|
||
|
||
def __init__(self, paths: list[str], mode: str, apply_bg_sub: bool):
|
||
super().__init__()
|
||
self._paths = paths
|
||
self._mode = mode
|
||
self._apply_bg_sub = apply_bg_sub
|
||
|
||
def run(self):
|
||
# Pass 1: quick open of each file just to weight progress by total
|
||
# angle count. Don't hold every file's memmap open at once — reopen
|
||
# fresh per file in pass 2 below.
|
||
total_angles = 0
|
||
for path in self._paths:
|
||
try:
|
||
total_angles += SrasFile(path).n_angles
|
||
except Exception:
|
||
pass # unreadable files are reported properly in pass 2
|
||
total_angles = max(total_angles, 1)
|
||
|
||
done_angles = 0
|
||
for path in self._paths:
|
||
try:
|
||
sras = SrasFile(path)
|
||
if sras.version not in (6, 7):
|
||
self.file_done.emit(
|
||
path, f"unsupported version {sras.version} — only "
|
||
"v6/v7 files can be batch-cached")
|
||
continue
|
||
|
||
n = sras.n_angles
|
||
if self._mode == "dc":
|
||
new_dc3 = [None] * n
|
||
new_dc4 = [None] * n
|
||
for a in range(n):
|
||
new_dc3[a] = adc_to_mv(
|
||
compute_dc_image(sras, a, CH3_IDX),
|
||
sras.ch_ymult_mv[CH3_IDX], sras.ch_yoff_adc[CH3_IDX],
|
||
sras.ch_yzero_mv[CH3_IDX])
|
||
new_dc4[a] = adc_to_mv(
|
||
compute_dc_image(sras, a, CH4_IDX),
|
||
sras.ch_ymult_mv[CH4_IDX], sras.ch_yoff_adc[CH4_IDX],
|
||
sras.ch_yzero_mv[CH4_IDX])
|
||
done_angles += 1
|
||
self.progress.emit(int(done_angles / total_angles * 100))
|
||
sras.write_v7_cache(new_dc3_mv=new_dc3, new_dc4_mv=new_dc4)
|
||
else:
|
||
effective_bg = self._apply_bg_sub and sras.background is not None
|
||
new_freq = [None] * n
|
||
for a in range(n):
|
||
new_freq[a] = compute_rf_image(
|
||
sras, a, dc_threshold_mv=-1e9, # mask nothing
|
||
apply_bg_sub=effective_bg)
|
||
done_angles += 1
|
||
self.progress.emit(int(done_angles / total_angles * 100))
|
||
sras.write_v7_cache(new_freq_mhz=new_freq, new_bg_sub=effective_bg)
|
||
|
||
self.file_done.emit(path, "")
|
||
except Exception as exc:
|
||
self.file_done.emit(path, str(exc))
|
||
|
||
self.finished.emit()
|
||
|
||
|
||
class AngleAlignmentWorker(QObject):
|
||
"""Computes rotation+translation alignment for every angle in *sras*,
|
||
referenced to *ref_angle_idx*, from each angle's binarized CH4 mask.
|
||
Rotation is analytic (from sras.angles_deg); only translation is found
|
||
by phase correlation. Uses the same progress(int)/finished(...) worker
|
||
shape as the other background-thread workers in this file.
|
||
"""
|
||
progress = pyqtSignal(int) # 0–100
|
||
finished = pyqtSignal(object, str) # AlignmentResult|None, error msg ("" = success)
|
||
|
||
def __init__(self, sras: SrasFile, ref_angle_idx: int, dc_threshold_mv: float):
|
||
super().__init__()
|
||
self._sras = sras
|
||
self._ref = ref_angle_idx
|
||
self._threshold = dc_threshold_mv
|
||
|
||
def run(self):
|
||
try:
|
||
result = _compute_angle_alignment(
|
||
self._sras, self._ref, self._threshold,
|
||
progress_cb=lambda pct: self.progress.emit(pct))
|
||
self.finished.emit(result, "")
|
||
except Exception as exc:
|
||
self.finished.emit(None, str(exc))
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 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):
|
||
"""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.
|
||
"""
|
||
data = sras.data[angle_idx]
|
||
waveform = data[row_idx, CH1_IDX, frame_idx, :].astype(np.float32)
|
||
t_ns = sras.time_axis_ns()
|
||
f_mhz = sras.freq_axis_mhz()
|
||
dc3_val = data[row_idx, CH3_IDX, frame_idx, :].astype(np.float32).mean()
|
||
dc4_val = data[row_idx, CH4_IDX, frame_idx, :].astype(np.float32).mean()
|
||
|
||
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")
|
||
|
||
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()
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 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())
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Main window
|
||
# ---------------------------------------------------------------------------
|
||
|
||
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_bg_sub: bool = True
|
||
self._pending_threshold: float = 50.0 # mV
|
||
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
|
||
|
||
# Convert menu: batch DC/FFT compute-and-store (v6 -> v7)
|
||
self._batch_thread: QThread | None = None
|
||
self._batch_worker: BatchCacheWorker | None = None
|
||
self._batch_errors: list[str] = []
|
||
self._batch_progress_dlg: QProgressDialog | None = None
|
||
|
||
# Display-only settings (colormap, grating) never trigger a
|
||
# recompute — they're applied to cached data on redraw. DC images
|
||
# (CH3/CH4) are cheap and precomputed for every angle in the
|
||
# background right after load. CH1/Velocity FFT images are
|
||
# computed lazily (with a progress popup) the first time an
|
||
# angle/threshold combination is viewed — using the cached DC4
|
||
# image to skip the FFT entirely for masked-out pixels — and
|
||
# cached per (angle, bg_sub, n_fft, threshold) so revisiting the
|
||
# same combination is free.
|
||
self._dc_cache: dict[tuple[int, int], np.ndarray] = {}
|
||
self._fft_cache: dict[tuple[int, bool, int | None, float], np.ndarray] = {}
|
||
self._dc_precompute_thread: QThread | None = None
|
||
self._dc_precompute_worker: DcPrecomputeWorker | None = None
|
||
self._dc_generation: int = 0
|
||
|
||
# Angle alignment ("Fusion" menu)
|
||
self._alignment_result: AlignmentResult | None = None
|
||
self._alignment_thread: QThread | None = None
|
||
self._alignment_worker: AngleAlignmentWorker | None = None
|
||
self._alignment_generation: int = 0
|
||
self._aligned_cache: dict[tuple, np.ndarray] = {}
|
||
|
||
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)
|
||
# Background DC-precompute progress (hidden until a file is loaded)
|
||
self.lbl_dc_precompute = QLabel("")
|
||
self.lbl_dc_precompute.setWordWrap(True)
|
||
self.lbl_dc_precompute.setStyleSheet("color: #4a90d9; font-size: 11px;")
|
||
il.addWidget(self.lbl_dc_precompute)
|
||
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)
|
||
|
||
# Aligned View (Fusion → Angle Alignment result)
|
||
self.chk_aligned_view = QCheckBox("Aligned View (Fusion)")
|
||
self.chk_aligned_view.setChecked(False)
|
||
self.chk_aligned_view.setEnabled(False)
|
||
self.chk_aligned_view.setToolTip(
|
||
"Show the current angle/channel resampled onto the shared,\n"
|
||
"rotation+translation-aligned canvas from Fusion → Angle\n"
|
||
"Alignment. Uncheck to see the raw per-angle scan grid."
|
||
)
|
||
self.chk_aligned_view.toggled.connect(self._on_aligned_view_toggled)
|
||
vl.addWidget(self.chk_aligned_view)
|
||
|
||
# Velocity settings (visible only in velocity mode)
|
||
self.grp_velocity = QGroupBox("Velocity Settings (CH1 only)")
|
||
vel_l = QVBoxLayout(self.grp_velocity)
|
||
grat_row = QHBoxLayout()
|
||
grat_row.addWidget(QLabel("Grating size:"))
|
||
self.spin_grating_um = QDoubleSpinBox()
|
||
self.spin_grating_um.setRange(0.1, 1000.0)
|
||
self.spin_grating_um.setDecimals(2)
|
||
self.spin_grating_um.setSingleStep(0.5)
|
||
self.spin_grating_um.setSuffix(" µm")
|
||
self.spin_grating_um.setValue(25)
|
||
self.spin_grating_um.setEnabled(False)
|
||
self.spin_grating_um.editingFinished.connect(self._on_grating_changed)
|
||
grat_row.addWidget(self.spin_grating_um)
|
||
vel_l.addLayout(grat_row)
|
||
self.lbl_velocity_formula = QLabel("v (m/s) = freq (MHz) × grating (µm)")
|
||
self.lbl_velocity_formula.setStyleSheet("font-size: 10px; color: #888;")
|
||
vel_l.addWidget(self.lbl_velocity_formula)
|
||
self.grp_velocity.setVisible(False)
|
||
# (grp_velocity will be added to the right panel below)
|
||
|
||
# Export
|
||
self.btn_export_csv = QPushButton("Export Image as CSV…")
|
||
self.btn_export_csv.setEnabled(False)
|
||
self.btn_export_csv.setToolTip(
|
||
"Save the current CH1 image (one scan row per CSV line)."
|
||
)
|
||
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)
|
||
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_cmap_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()
|
||
wave_vl.addWidget(self.lbl_wave_hint)
|
||
wave_vl.addWidget(self.wave_canvas)
|
||
splitter.addWidget(wave_widget)
|
||
|
||
splitter.setSizes([580, 250])
|
||
|
||
# ---- Right control panel -------------------------------------------
|
||
right_panel = QWidget()
|
||
right_panel.setFixedWidth(270)
|
||
right_panel_layout = QVBoxLayout(right_panel)
|
||
right_panel_layout.setContentsMargins(0, 0, 0, 0)
|
||
right_panel_layout.setSpacing(6)
|
||
right_panel_layout.addWidget(self.grp_velocity)
|
||
right_panel_layout.addWidget(grp_display)
|
||
right_panel_layout.addStretch()
|
||
root.addWidget(right_panel)
|
||
|
||
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)
|
||
|
||
fusion_menu = menubar.addMenu("&Fusion")
|
||
self._alignment_act = QAction("Angle &Alignment", self)
|
||
self._alignment_act.setStatusTip(
|
||
"Compute a rotation+translation alignment across all angles "
|
||
"(from CH4 masks) and enable Aligned View. Requires >1 angle.")
|
||
self._alignment_act.setEnabled(False)
|
||
self._alignment_act.triggered.connect(self._on_angle_alignment)
|
||
fusion_menu.addAction(self._alignment_act)
|
||
|
||
convert_menu = menubar.addMenu("&Convert")
|
||
self._batch_dc_act = QAction("Batch Compute DC and &Store…", self)
|
||
self._batch_dc_act.setStatusTip(
|
||
"Select .sras files and compute+store DC images (CH3/CH4 mean) "
|
||
"for every angle, converting v6 files to v7 in place.")
|
||
self._batch_dc_act.triggered.connect(lambda: self._on_batch_compute("dc"))
|
||
convert_menu.addAction(self._batch_dc_act)
|
||
|
||
self._batch_fft_act = QAction("Batch Compute FFT and Sto&re…", self)
|
||
self._batch_fft_act.setStatusTip(
|
||
"Select .sras files and compute+store FFT peak-frequency images "
|
||
"for every angle, converting v6 files to v7 in place.")
|
||
self._batch_fft_act.triggered.connect(lambda: self._on_batch_compute("fft"))
|
||
convert_menu.addAction(self._batch_fft_act)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Drag-and-drop
|
||
# ------------------------------------------------------------------
|
||
|
||
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
|
||
# Claim self._load_thread before any call below that can pump the
|
||
# Qt event loop — see the comment in _start_compute for why.
|
||
self._load_worker = LoadWorker(path)
|
||
self._load_thread = QThread()
|
||
self._load_worker.moveToThread(self._load_thread)
|
||
self._load_thread.started.connect(self._load_worker.run)
|
||
self._load_worker.finished.connect(self._on_load_done)
|
||
self._load_worker.error.connect(
|
||
lambda msg: self.statusBar().showMessage(f"Error: {msg}")
|
||
)
|
||
self._load_worker.finished.connect(self._load_thread.quit)
|
||
self._load_thread.finished.connect(self._on_load_thread_finished)
|
||
|
||
self.btn_open.setEnabled(False)
|
||
self.statusBar().showMessage(f"Loading {Path(path).name}…")
|
||
self._show_progress(f"Loading {Path(path).name}…")
|
||
|
||
self._load_thread.start()
|
||
|
||
def _on_load_thread_finished(self):
|
||
# See the comment in _on_compute_thread_finished: wait() before
|
||
# releasing our reference to avoid destroying a QThread whose OS
|
||
# thread hasn't fully joined yet.
|
||
if self._load_thread is not None:
|
||
self._load_thread.wait()
|
||
self._load_thread = None
|
||
|
||
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
|
||
|
||
# Caches (and any in-flight DC precompute) belong to the previous
|
||
# file's geometry — discard and start fresh.
|
||
self._dc_cache = {}
|
||
self._fft_cache = {}
|
||
self.lbl_dc_precompute.setText("")
|
||
|
||
# Same for any alignment result — belongs to the previous file's
|
||
# geometry. Bump the generation counter so a still-running
|
||
# alignment worker's result is discarded when it lands (see
|
||
# _on_alignment_done).
|
||
self._alignment_result = None
|
||
self._aligned_cache = {}
|
||
self._alignment_generation += 1
|
||
self.chk_aligned_view.blockSignals(True)
|
||
self.chk_aligned_view.setChecked(False)
|
||
self.chk_aligned_view.setEnabled(False)
|
||
self.chk_aligned_view.blockSignals(False)
|
||
|
||
# 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()
|
||
|
||
self.lbl_filename.setText(sras.path.name)
|
||
|
||
self.spin_angle.blockSignals(True)
|
||
self.spin_angle.setRange(0, max(0, sras.n_angles - 1))
|
||
self.spin_angle.setValue(0)
|
||
self.spin_angle.blockSignals(False)
|
||
|
||
# DC channels are cheap and give an instant, fluid overview of a
|
||
# scan; CH1/Velocity require an FFT per pixel that can take minutes
|
||
# on a large scan, so don't default to it.
|
||
self.combo_channel.blockSignals(True)
|
||
self.combo_channel.setCurrentIndex(CH4_IDX)
|
||
self.combo_channel.blockSignals(False)
|
||
|
||
self._update_controls_enabled(True)
|
||
# Refresh the ADC-count label now that we have file calibration
|
||
self._on_threshold_changed()
|
||
self._on_view_changed()
|
||
|
||
# Precompute DC images for every angle in the background so
|
||
# switching angles while viewing a DC channel is instant, and the
|
||
# FFT masking step rarely has to wait on a DC4 image either.
|
||
self._start_dc_precompute()
|
||
|
||
# ------------------------------------------------------------------
|
||
# Scan info panel
|
||
# ------------------------------------------------------------------
|
||
|
||
def _update_scan_info_labels(self):
|
||
s = self._sras
|
||
if s is None:
|
||
return
|
||
angle_idx = self.spin_angle.value()
|
||
self._info["Angles"].setText(f"Angles: {s.n_angles}")
|
||
self._info["Rows"].setText(f"Rows: {s.n_rows[angle_idx]}")
|
||
self._info["Frames / row"].setText(f"Frames / row: {s.n_frames[angle_idx]}")
|
||
self._info["Samples / frame"].setText(f"Samples / frame: {s.samples_per_frame}")
|
||
self._info["Sample rate"].setText(f"Sample rate: {s.sample_rate_hz/1e9:.4g} GS/s")
|
||
self._info["X start"].setText(f"X start: {s.x_start_mm[angle_idx]:.4g} mm")
|
||
self._info["Pixel Δx"].setText(f"Pixel Δx: {s.pixel_x_mm*1e3:.3g} µm")
|
||
self._info["Laser freq"].setText(f"Laser freq: {s.laser_freq_hz/1e3:.4g} kHz")
|
||
|
||
notes = []
|
||
if s.frame_count_mismatch:
|
||
notes.append(
|
||
f"! Header n_frames={s.n_frames_header}, "
|
||
f"actual={s.n_frames[angle_idx]} (scanner bug — corrected)"
|
||
)
|
||
if s.scan_aborted:
|
||
notes.append(
|
||
f"! Scan aborted: {s.n_angles}/{s.n_angles_declared} angles complete"
|
||
)
|
||
if s.background is not None:
|
||
notes.append(f"Background waveform: {len(s.background)} samples")
|
||
if isinstance(s.precomputed_freq_mhz, np.ndarray):
|
||
bg_note = " (bg-sub)" if s.precomputed_bg_sub else " (no bg-sub)"
|
||
notes.append(f"v5: precomputed images present{bg_note} — display is instant")
|
||
if s.version in (6, 7):
|
||
notes.append("v6/v7 format: rows / frames / x_start are per-angle")
|
||
if s.version == 7:
|
||
n_dc = sum(1 for x in s.precomputed_dc4_mv if x is not None)
|
||
n_fft = sum(1 for x in s.precomputed_freq_mhz if x is not None)
|
||
if n_dc or n_fft:
|
||
bg_note = " (bg-sub)" if s.precomputed_bg_sub else " (no bg-sub)"
|
||
notes.append(
|
||
f"v7: cached DC ({n_dc}/{s.n_angles} angles), "
|
||
f"FFT ({n_fft}/{s.n_angles} angles{bg_note if n_fft else ''}) "
|
||
"— display is instant for cached angles")
|
||
else:
|
||
notes.append("v7 format: no cache blocks stored yet")
|
||
self.lbl_frame_warn.setText("\n".join(notes))
|
||
|
||
# ------------------------------------------------------------------
|
||
# Controls
|
||
# ------------------------------------------------------------------
|
||
|
||
def _update_controls_enabled(self, enabled: bool):
|
||
s = self._sras
|
||
self.spin_angle.setEnabled(enabled and s is not None and s.n_angles > 1)
|
||
self.combo_channel.setEnabled(enabled)
|
||
self.combo_cmap.setEnabled(enabled)
|
||
self.chk_auto.setEnabled(enabled)
|
||
manual = enabled and not self.chk_auto.isChecked()
|
||
self.spin_vmin.setEnabled(manual)
|
||
self.spin_vmax.setEnabled(manual)
|
||
ch_idx = self.combo_channel.currentIndex()
|
||
is_ch1 = enabled and ch_idx in CH1_DERIVED_MODES
|
||
# Threshold and bg-sub apply to all CH1 modes
|
||
self.spin_threshold_mv.setEnabled(is_ch1)
|
||
has_bg = enabled and s is not None and s.background is not None
|
||
self.chk_bg_sub.setEnabled(has_bg and is_ch1)
|
||
# Velocity grating spinbox
|
||
is_vel = enabled and ch_idx == VELOCITY_MODE_IDX
|
||
self.spin_grating_um.setEnabled(is_vel)
|
||
self.grp_velocity.setVisible(is_vel)
|
||
# CSV export: enabled when a CH1-derived image is displayed
|
||
self.btn_export_csv.setEnabled(is_ch1 and self._current_image is not None)
|
||
# ROI: always usable once a file is loaded (independent of channel)
|
||
self.btn_draw_roi.setEnabled(enabled and s is not None)
|
||
# Batch Convert actions: pick their own files, independent of
|
||
# whatever's currently open — only gated on no batch already running
|
||
can_batch = self._batch_thread is None
|
||
self._batch_dc_act.setEnabled(can_batch)
|
||
self._batch_fft_act.setEnabled(can_batch)
|
||
# Angle alignment: needs >1 angle and no alignment already running
|
||
can_align = (enabled and s is not None and s.n_angles > 1
|
||
and self._alignment_thread is None)
|
||
self._alignment_act.setEnabled(can_align)
|
||
self.chk_aligned_view.setEnabled(
|
||
enabled and self._alignment_result is not None)
|
||
self._update_roi_ui()
|
||
|
||
def _on_channel_changed(self):
|
||
ch_idx = self.combo_channel.currentIndex()
|
||
has_file = self._sras is not None
|
||
is_ch1 = ch_idx in CH1_DERIVED_MODES
|
||
self.spin_threshold_mv.setEnabled(is_ch1 and has_file)
|
||
has_bg = has_file and self._sras.background is not None
|
||
self.chk_bg_sub.setEnabled(has_bg and is_ch1)
|
||
is_vel = ch_idx == VELOCITY_MODE_IDX
|
||
self.spin_grating_um.setEnabled(is_vel and has_file)
|
||
self.grp_velocity.setVisible(is_vel)
|
||
self.btn_export_csv.setEnabled(is_ch1 and has_file and self._current_image is not None)
|
||
self._on_view_changed()
|
||
|
||
def _on_bg_sub_toggled(self):
|
||
# Background subtraction changes the FFT input, so it genuinely
|
||
# invalidates the cached raw FFT (the cache key includes it) —
|
||
# _refresh_display() will recompute only if there's no cache hit
|
||
# for the new bg-sub state.
|
||
if self._sras is not None:
|
||
if self.combo_channel.currentIndex() in CH1_DERIVED_MODES:
|
||
self._refresh_display()
|
||
|
||
def _on_grating_changed(self):
|
||
# Grating is a pure post-multiply on the cached frequency image —
|
||
# never needs a recompute.
|
||
if self._sras is not None and self.combo_channel.currentIndex() == VELOCITY_MODE_IDX:
|
||
self._refresh_display()
|
||
|
||
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:
|
||
# Deliberately always the raw per-angle grid, even when
|
||
# Aligned View is on: _on_export_roi_csv also exports on the
|
||
# raw grid (never synthetically-resampled pixels), so this
|
||
# readout must match what Export ROI actually writes.
|
||
mask = roi.mask_for_grid(self._sras.x_axis_mm(self._current_angle),
|
||
self._sras.y_positions_mm(self._current_angle))
|
||
npix = int(mask.sum())
|
||
except Exception:
|
||
npix = 0
|
||
self.lbl_roi_npix.setText(f"pixels inside: {npix}")
|
||
self.btn_clear_roi.setEnabled(True)
|
||
self.btn_export_roi.setEnabled(
|
||
self._current_image is not None and npix > 0)
|
||
|
||
def _on_clear_roi(self):
|
||
self.image_canvas.clear_roi()
|
||
self.statusBar().showMessage("ROI cleared")
|
||
|
||
def _on_roi_angle_edited(self):
|
||
pass # rotation control removed — corners are dragged individually
|
||
|
||
def _on_export_roi_csv(self):
|
||
if self._current_image is None or self._sras is None:
|
||
return
|
||
roi = self.image_canvas.get_roi()
|
||
if roi is None:
|
||
self.statusBar().showMessage("No ROI — draw one first")
|
||
return
|
||
s = self._sras
|
||
x_axis = s.x_axis_mm(self._current_angle)
|
||
y_axis = s.y_positions_mm(self._current_angle)
|
||
X, Y = np.meshgrid(np.asarray(x_axis, dtype=np.float64),
|
||
np.asarray(y_axis, dtype=np.float64))
|
||
mask = roi.mask_for_grid(x_axis, y_axis)
|
||
if not mask.any():
|
||
self.statusBar().showMessage("ROI does not overlap any pixel")
|
||
return
|
||
img = self._current_image
|
||
if img.shape != mask.shape:
|
||
self.statusBar().showMessage(
|
||
f"ROI shape {mask.shape} does not match image {img.shape}")
|
||
return
|
||
|
||
rows_idx, frames_idx = np.where(mask)
|
||
xs = X[mask]
|
||
ys = Y[mask]
|
||
vals = img[mask]
|
||
|
||
ch_idx = self._current_ch
|
||
ch_name = CH_NAMES[ch_idx]
|
||
angle = self._current_angle
|
||
default_name = (f"{s.path.stem}_angle{angle}_{ch_name}_ROI.csv")
|
||
path, _ = QFileDialog.getSaveFileName(
|
||
self, "Export ROI as CSV",
|
||
str(s.path.parent / default_name),
|
||
"CSV files (*.csv);;All files (*)",
|
||
)
|
||
if not path:
|
||
return
|
||
|
||
pts = roi.corners()
|
||
corners_str = " ".join(f"({p[0]:.6g},{p[1]:.6g})" for p in pts)
|
||
header = (
|
||
f"# ROI quad corners (BL BR TR TL) mm: {corners_str}\n"
|
||
f"# source: {s.path.name}, channel={ch_name}, "
|
||
f"angle_idx={angle}, angle_deg={s.angles_deg[angle]:.4g}\n"
|
||
f"# n_pixels={int(mask.sum())}\n"
|
||
"row,frame,x_mm,y_mm,value"
|
||
)
|
||
data = np.column_stack([
|
||
rows_idx.astype(np.int64),
|
||
frames_idx.astype(np.int64),
|
||
xs, ys, vals.astype(np.float64),
|
||
])
|
||
# integer columns first, floats after — use a per-column format list
|
||
np.savetxt(path, data, delimiter=",",
|
||
fmt=["%d", "%d", "%.6g", "%.6g", "%.6g"],
|
||
header=header, comments="")
|
||
self.statusBar().showMessage(
|
||
f"Exported ROI ({int(mask.sum())} pixels) to {Path(path).name}")
|
||
|
||
def _on_threshold_changed(self):
|
||
mv = self.spin_threshold_mv.value()
|
||
if self._sras is not None:
|
||
ymult = self._sras.ch_ymult_mv[CH4_IDX]
|
||
yoff = self._sras.ch_yoff_adc[CH4_IDX]
|
||
yzero = self._sras.ch_yzero_mv[CH4_IDX]
|
||
else:
|
||
ymult, yoff, yzero = _FALLBACK_YMULT_MV, _FALLBACK_YOFF_ADC, 0.0
|
||
self.lbl_threshold_adc.setText(f"≈ {mv_to_adc(mv, ymult, yoff, yzero):.1f} ADC counts")
|
||
# Threshold decides which pixels get an FFT at all (see
|
||
# ComputeWorker), so changing it is a genuine cache key change —
|
||
# _refresh_display() recomputes only on a miss, and that recompute
|
||
# reuses the cached DC4 image to skip masked-out pixels.
|
||
if self._sras is not None and self.combo_channel.currentIndex() in CH1_DERIVED_MODES:
|
||
self._refresh_display()
|
||
|
||
def _on_autoscale_toggled(self, checked: bool):
|
||
manual = not checked
|
||
self.spin_vmin.setEnabled(manual and self._sras is not None)
|
||
self.spin_vmax.setEnabled(manual and self._sras is not None)
|
||
if self._sras is not None and self._current_image is not None:
|
||
self._redraw_image(self._current_image)
|
||
|
||
def _on_manual_range_changed(self):
|
||
if not self.chk_auto.isChecked() and self._current_image is not None:
|
||
self._redraw_image(self._current_image)
|
||
|
||
def _on_cmap_changed(self):
|
||
# Colormap is purely how the existing image is rendered — never
|
||
# needs a recompute.
|
||
if self._current_image is not None:
|
||
self._redraw_image(self._current_image)
|
||
|
||
def _on_view_changed(self):
|
||
if self._sras is None:
|
||
return
|
||
idx = self.spin_angle.value()
|
||
self.lbl_angle_deg.setText(f"({self._sras.angles_deg[idx]:.1f}°)")
|
||
self._update_scan_info_labels()
|
||
self._refresh_display()
|
||
|
||
# ------------------------------------------------------------------
|
||
# Computation
|
||
# ------------------------------------------------------------------
|
||
|
||
def _current_n_fft(self) -> int | None:
|
||
pad_factor = self._fft_pad_factor
|
||
if pad_factor <= 1 or self._sras is None:
|
||
return None
|
||
return self._sras.samples_per_frame * pad_factor
|
||
|
||
def _scale_for_display(self, freq_mhz: np.ndarray, ch_idx: int) -> np.ndarray:
|
||
"""Velocity is a pure post-multiply of the (already DC-masked)
|
||
cached frequency image — never worth a recompute on its own."""
|
||
if ch_idx == VELOCITY_MODE_IDX:
|
||
return freq_mhz * self.spin_grating_um.value()
|
||
return freq_mhz
|
||
|
||
# ------------------------------------------------------------------
|
||
# Aligned View (Fusion) display helpers
|
||
# ------------------------------------------------------------------
|
||
|
||
def _aligned_cache_key(self, angle_idx: int, ch_idx: int) -> tuple:
|
||
"""Mirrors _fft_cache's key granularity so a stale aligned image is
|
||
never shown after bg_sub/threshold/pad/grating changes."""
|
||
if ch_idx in CH1_DERIVED_MODES:
|
||
return (angle_idx, ch_idx, self.chk_bg_sub.isChecked(),
|
||
self._current_n_fft(), self.spin_threshold_mv.value(),
|
||
self.spin_grating_um.value() if ch_idx == VELOCITY_MODE_IDX else None)
|
||
return (angle_idx, ch_idx)
|
||
|
||
def _aligned_canvas_axes(self) -> tuple[np.ndarray, np.ndarray]:
|
||
r = self._alignment_result
|
||
n_rows, n_cols = r.canvas_shape
|
||
x_axis = r.canvas_origin_mm[0] + np.arange(n_cols) * r.canvas_dx_mm
|
||
y_axis = r.canvas_origin_mm[1] + np.arange(n_rows) * r.canvas_dy_mm
|
||
return x_axis, y_axis
|
||
|
||
def _get_aligned_display_image(self, raw_img: np.ndarray, angle_idx: int,
|
||
ch_idx: int) -> np.ndarray:
|
||
key = self._aligned_cache_key(angle_idx, ch_idx)
|
||
cached = self._aligned_cache.get(key)
|
||
if cached is None:
|
||
cached = apply_alignment(self._alignment_result, angle_idx, raw_img)
|
||
self._aligned_cache[key] = cached
|
||
return cached
|
||
|
||
def _refresh_display(self):
|
||
"""Show the image for the current angle/channel/threshold, using
|
||
cached data whenever possible and only falling back to a
|
||
background compute (with progress popup) when genuinely nothing
|
||
is cached yet for these settings."""
|
||
if self._sras is None:
|
||
return
|
||
angle_idx = self.spin_angle.value()
|
||
ch_idx = self.combo_channel.currentIndex()
|
||
|
||
if ch_idx in (CH3_IDX, CH4_IDX):
|
||
cached = self._dc_cache.get((angle_idx, ch_idx))
|
||
if cached is not None:
|
||
self._show_image_now(cached, angle_idx, ch_idx)
|
||
return
|
||
elif ch_idx in (CH1_IDX, VELOCITY_MODE_IDX):
|
||
apply_bg_sub = self.chk_bg_sub.isChecked()
|
||
threshold_mv = self.spin_threshold_mv.value()
|
||
key = (angle_idx, apply_bg_sub, self._current_n_fft(), threshold_mv)
|
||
raw = self._fft_cache.get(key)
|
||
if raw is not None:
|
||
img = self._scale_for_display(raw, ch_idx)
|
||
self._show_image_now(img, angle_idx, ch_idx)
|
||
return
|
||
|
||
# Nothing cached for these settings — need a real compute. Changing
|
||
# the DC threshold changes *which* pixels get an FFT at all, so it
|
||
# can't be satisfied from the cache — but with the DC map already
|
||
# known, the recompute skips the FFT for masked-out pixels and is
|
||
# much cheaper than a full-image compute would be.
|
||
self._start_compute()
|
||
|
||
def _show_image_now(self, img: np.ndarray, angle_idx: int, ch_idx: int):
|
||
"""Display an already-available image with no compute involved."""
|
||
self._current_image = img
|
||
self._current_angle = angle_idx
|
||
self._current_ch = ch_idx
|
||
self.btn_export_csv.setEnabled(ch_idx in CH1_DERIVED_MODES)
|
||
self._redraw_image(img)
|
||
self._update_roi_ui()
|
||
|
||
# ------------------------------------------------------------------
|
||
# Background compute (only reached on a genuine cache miss)
|
||
# ------------------------------------------------------------------
|
||
|
||
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()
|
||
apply_bg_sub = self.chk_bg_sub.isChecked()
|
||
threshold_mv = self.spin_threshold_mv.value()
|
||
pad_factor = self._fft_pad_factor
|
||
n_fft = self._current_n_fft()
|
||
# Reuse the cached DC4 image (if the precompute has reached this
|
||
# angle) so the FFT compute skips masked-out pixels entirely and
|
||
# doesn't need to re-read the CH4 channel from disk.
|
||
dc4_mv = self._dc_cache.get((angle_idx, CH4_IDX))
|
||
|
||
self._pending_angle = angle_idx
|
||
self._pending_ch = ch_idx
|
||
self._pending_bg_sub = apply_bg_sub
|
||
self._pending_threshold = threshold_mv
|
||
self._pending_fft_pad_factor = pad_factor
|
||
|
||
# Claim self._compute_thread *before* anything below that can pump
|
||
# the Qt event loop (e.g. QProgressDialog.show() on first display).
|
||
# If that happened first, a re-entrant editingFinished/signal could
|
||
# slip past the guard above, start a second thread, and then have
|
||
# this call's own assignment clobber (and destroy while still
|
||
# running) that second thread's QThread object — which aborts the
|
||
# process. Assigning immediately closes that window.
|
||
self._compute_worker = ComputeWorker(
|
||
self._sras, angle_idx, ch_idx, apply_bg_sub, n_fft=n_fft,
|
||
dc_threshold_mv=threshold_mv, dc4_mv=dc4_mv,
|
||
)
|
||
self._compute_thread = QThread()
|
||
self._compute_worker.moveToThread(self._compute_thread)
|
||
self._compute_thread.started.connect(self._compute_worker.run)
|
||
self._compute_worker.finished.connect(self._on_compute_done)
|
||
self._compute_worker.error.connect(
|
||
lambda msg: self.statusBar().showMessage(f"Compute error: {msg}")
|
||
)
|
||
self._compute_worker.finished.connect(self._compute_thread.quit)
|
||
self._compute_thread.finished.connect(self._on_compute_thread_finished)
|
||
|
||
if ch_idx in (CH1_IDX, VELOCITY_MODE_IDX):
|
||
self.statusBar().showMessage("Computing FFT…")
|
||
self._show_progress(
|
||
f"Computing FFT for angle {angle_idx}…\n"
|
||
"This can take a while on a large scan — result is cached "
|
||
"so revisiting this angle/mode/threshold will be instant."
|
||
)
|
||
else:
|
||
self.statusBar().showMessage("Computing DC image…")
|
||
self._show_progress(f"Computing DC image for angle {angle_idx}…")
|
||
|
||
self._compute_thread.start()
|
||
|
||
def _on_compute_thread_finished(self):
|
||
# Block until the OS thread has actually joined before dropping our
|
||
# last reference — deallocating a QThread whose thread hasn't fully
|
||
# terminated yet logs "QThread: Destroyed while thread is still
|
||
# running" and aborts the process. The finished() signal fires as
|
||
# the thread is winding down but does not guarantee it has joined.
|
||
if self._compute_thread is not None:
|
||
self._compute_thread.wait()
|
||
self._compute_thread = None
|
||
angle_idx = self.spin_angle.value()
|
||
ch_idx = self.combo_channel.currentIndex()
|
||
apply_bg_sub = self.chk_bg_sub.isChecked()
|
||
threshold_mv = self.spin_threshold_mv.value()
|
||
if (angle_idx, ch_idx, apply_bg_sub, threshold_mv, self._fft_pad_factor) != (
|
||
self._pending_angle, self._pending_ch, self._pending_bg_sub,
|
||
self._pending_threshold, self._pending_fft_pad_factor):
|
||
# Settings changed while this compute was running — re-dispatch
|
||
# through the cache-aware path in case that now-current
|
||
# combination happens to already be cached.
|
||
self._refresh_display()
|
||
|
||
def _on_compute_done(self, result):
|
||
self._close_progress()
|
||
angle_idx = self._pending_angle
|
||
ch_idx = self._pending_ch
|
||
|
||
if ch_idx in (CH1_IDX, VELOCITY_MODE_IDX):
|
||
masked_freq_mhz = result
|
||
key = (angle_idx, self._pending_bg_sub, self._current_n_fft(),
|
||
self._pending_threshold)
|
||
self._fft_cache[key] = masked_freq_mhz
|
||
img = self._scale_for_display(masked_freq_mhz, ch_idx)
|
||
else:
|
||
img = result
|
||
self._dc_cache[(angle_idx, ch_idx)] = img
|
||
|
||
self._show_image_now(img, angle_idx, ch_idx)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Background DC precompute (all angles, so switching is fluid)
|
||
# ------------------------------------------------------------------
|
||
|
||
def _start_dc_precompute(self):
|
||
if self._sras is None:
|
||
return
|
||
self._dc_generation += 1
|
||
generation = self._dc_generation
|
||
n_angles = self._sras.n_angles
|
||
|
||
worker = DcPrecomputeWorker(self._sras)
|
||
thread = QThread()
|
||
worker.moveToThread(thread)
|
||
thread.started.connect(worker.run)
|
||
worker.angle_done.connect(
|
||
lambda a, dc3, dc4, g=generation: self._on_dc_precompute_angle_done(
|
||
g, a, dc3, dc4, n_angles)
|
||
)
|
||
worker.error.connect(
|
||
lambda msg: self.statusBar().showMessage(f"DC precompute error: {msg}", 5000)
|
||
)
|
||
worker.finished.connect(thread.quit)
|
||
worker.error.connect(thread.quit)
|
||
thread.finished.connect(lambda: self._on_dc_precompute_thread_finished(thread))
|
||
|
||
# See _start_compute for why the thread/worker are claimed on self
|
||
# before .start() rather than after.
|
||
self._dc_precompute_worker = worker
|
||
self._dc_precompute_thread = thread
|
||
thread.start()
|
||
|
||
def _on_dc_precompute_angle_done(self, generation: int, angle_idx: int,
|
||
dc3_mv: np.ndarray, dc4_mv: np.ndarray,
|
||
n_angles: int):
|
||
if generation != self._dc_generation:
|
||
return # stale result from a previously-loaded file — discard
|
||
self._dc_cache[(angle_idx, CH3_IDX)] = dc3_mv
|
||
self._dc_cache[(angle_idx, CH4_IDX)] = dc4_mv
|
||
done = sum(1 for a in range(n_angles) if (a, CH4_IDX) in self._dc_cache)
|
||
if done < n_angles:
|
||
self.lbl_dc_precompute.setText(
|
||
f"Precomputing DC images: {done}/{n_angles} angles ready…")
|
||
else:
|
||
self.lbl_dc_precompute.setText("DC images ready for all angles.")
|
||
# If we just finished the angle/channel the user is currently
|
||
# looking at and it wasn't shown yet (e.g. they switched here
|
||
# before precompute caught up and are still waiting), show it now.
|
||
if (self._sras is not None and angle_idx == self.spin_angle.value()
|
||
and self._compute_thread is None
|
||
and self.combo_channel.currentIndex() in (CH3_IDX, CH4_IDX)
|
||
and (self._current_angle != angle_idx
|
||
or self._current_ch != self.combo_channel.currentIndex())):
|
||
self._refresh_display()
|
||
|
||
def _on_dc_precompute_thread_finished(self, thread: QThread):
|
||
# See the comment in _on_compute_thread_finished: wait() before
|
||
# releasing the reference to avoid destroying a QThread whose OS
|
||
# thread hasn't fully joined yet.
|
||
thread.wait()
|
||
if self._dc_precompute_thread is thread:
|
||
self._dc_precompute_thread = None
|
||
self._dc_precompute_worker = None
|
||
|
||
def _redraw_image(self, img: np.ndarray):
|
||
s = self._sras
|
||
angle_idx = self._current_angle
|
||
ch_idx = self._current_ch
|
||
|
||
aligned = (self.chk_aligned_view.isChecked()
|
||
and self._alignment_result is not None
|
||
and angle_idx in self._alignment_result.per_angle)
|
||
if aligned:
|
||
display_img = self._get_aligned_display_image(img, angle_idx, ch_idx)
|
||
x_axis, y_axis = self._aligned_canvas_axes()
|
||
else:
|
||
display_img = img
|
||
x_axis = s.x_axis_mm(angle_idx)
|
||
y_axis = s.y_positions_mm(angle_idx)
|
||
|
||
dx = x_axis[1] - x_axis[0] if len(x_axis) > 1 else s.pixel_x_mm
|
||
dy = float(y_axis[1] - y_axis[0]) if len(y_axis) > 1 else 1.0
|
||
|
||
extent = [
|
||
x_axis[0] - dx / 2,
|
||
x_axis[-1] + dx / 2,
|
||
y_axis[-1] + dy / 2,
|
||
y_axis[0] - dy / 2,
|
||
]
|
||
|
||
if self.chk_auto.isChecked():
|
||
vmin, vmax = float(display_img.min()), float(display_img.max())
|
||
for spin, val in ((self.spin_vmin, vmin), (self.spin_vmax, vmax)):
|
||
spin.blockSignals(True)
|
||
spin.setValue(val)
|
||
spin.blockSignals(False)
|
||
else:
|
||
vmin = self.spin_vmin.value()
|
||
vmax = self.spin_vmax.value()
|
||
|
||
angle_deg = s.angles_deg[self._current_angle]
|
||
ch_label = CH_LABELS[ch_idx]
|
||
|
||
if ch_idx == CH1_IDX:
|
||
mode_str = "RF"
|
||
unit = "Peak frequency (MHz)"
|
||
colorbar_label = "MHz"
|
||
elif ch_idx == VELOCITY_MODE_IDX:
|
||
grating = self.spin_grating_um.value()
|
||
mode_str = "Velocity"
|
||
unit = "Velocity (m/s)"
|
||
colorbar_label = "m/s"
|
||
ch_label = f"Velocity [grating={grating:.2f} µm]"
|
||
else:
|
||
mode_str = "DC"
|
||
unit = "DC mean (mV)"
|
||
colorbar_label = "mV"
|
||
|
||
title = f"{CH_NAMES[ch_idx]} | {mode_str} | {angle_deg:.1f}°"
|
||
if aligned:
|
||
title += " [Aligned]"
|
||
|
||
self.image_canvas.show_image(
|
||
display_img, extent,
|
||
cmap=self.combo_cmap.currentText(),
|
||
vmin=vmin, vmax=vmax,
|
||
xlabel="X (mm)", ylabel="Y (mm)",
|
||
title=title,
|
||
colorbar_label=colorbar_label,
|
||
)
|
||
self.statusBar().showMessage(
|
||
f"{s.path.name} | {ch_label} @ {angle_deg:.1f}° "
|
||
f"| {display_img.shape[1]} × {display_img.shape[0]} px | {unit}"
|
||
f"{' | Aligned' if aligned else ''}"
|
||
)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Pixel inspector
|
||
# ------------------------------------------------------------------
|
||
|
||
def _on_pixel_clicked(self, row_idx: int, frame_idx: int):
|
||
if self._sras is None or self._current_image is None:
|
||
return
|
||
angle_idx = self._current_angle
|
||
if (self.chk_aligned_view.isChecked() and self._alignment_result is not None
|
||
and angle_idx in self._alignment_result.per_angle):
|
||
# The click landed on the shared aligned canvas — invert the
|
||
# same canvas->raw affine used to display it back to a raw
|
||
# (row, frame) index before looking up the waveform.
|
||
t = self._alignment_result.per_angle[angle_idx]
|
||
raw = t.matrix @ np.array([row_idx, frame_idx], dtype=np.float64) + t.offset
|
||
row_idx, frame_idx = int(round(raw[0])), int(round(raw[1]))
|
||
n_rows_a = int(self._sras.n_rows[angle_idx])
|
||
n_frames_a = int(self._sras.n_frames[angle_idx])
|
||
if not (0 <= row_idx < n_rows_a and 0 <= frame_idx < n_frames_a):
|
||
self.statusBar().showMessage(
|
||
"No source waveform here (padding region of the aligned canvas).")
|
||
return
|
||
self.lbl_wave_hint.hide()
|
||
ch_idx = self._current_ch
|
||
if ch_idx in CH1_DERIVED_MODES:
|
||
self.wave_canvas.show_rf_waveform(
|
||
self._sras, angle_idx, row_idx, frame_idx,
|
||
apply_bg_sub=self.chk_bg_sub.isChecked(),
|
||
)
|
||
else:
|
||
self.wave_canvas.show_dc_waveform(
|
||
self._sras, angle_idx, ch_idx, row_idx, frame_idx
|
||
)
|
||
|
||
# ------------------------------------------------------------------
|
||
# 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
|
||
|
||
# ------------------------------------------------------------------
|
||
# Convert menu: batch DC/FFT compute-and-store (v6 -> v7)
|
||
# ------------------------------------------------------------------
|
||
|
||
def _on_batch_compute(self, mode: str):
|
||
if self._batch_thread is not None:
|
||
return
|
||
label = "DC" if mode == "dc" else "FFT"
|
||
paths, _ = QFileDialog.getOpenFileNames(
|
||
self, f"Select .sras files to batch-compute {label}", "",
|
||
"SRAS files (*.sras);;All files (*)",
|
||
)
|
||
if not paths:
|
||
return
|
||
if self._batch_thread is not None:
|
||
return # a second trigger snuck in while the file dialog was open
|
||
|
||
apply_bg = self.chk_bg_sub.isChecked()
|
||
self._batch_errors = []
|
||
|
||
# Claim self._batch_thread before anything below that can pump the
|
||
# Qt event loop — see the comment in _start_compute for why.
|
||
self._batch_worker = BatchCacheWorker(paths, mode, apply_bg)
|
||
self._batch_thread = QThread()
|
||
self._batch_worker.moveToThread(self._batch_thread)
|
||
self._batch_thread.started.connect(self._batch_worker.run)
|
||
self._batch_worker.progress.connect(self._on_batch_progress)
|
||
self._batch_worker.file_done.connect(self._on_batch_file_done)
|
||
self._batch_worker.finished.connect(lambda paths=paths: self._on_batch_finished(paths))
|
||
self._batch_worker.finished.connect(self._batch_thread.quit)
|
||
self._batch_thread.finished.connect(self._on_batch_thread_finished)
|
||
|
||
self._batch_dc_act.setEnabled(False)
|
||
self._batch_fft_act.setEnabled(False)
|
||
|
||
self._batch_progress_dlg = QProgressDialog(
|
||
f"Batch computing {label} for {len(paths)} file(s)…", "", 0, 100, self)
|
||
self._batch_progress_dlg.setWindowTitle("Please wait…")
|
||
self._batch_progress_dlg.setCancelButton(None)
|
||
self._batch_progress_dlg.setWindowModality(Qt.WindowModality.WindowModal)
|
||
self._batch_progress_dlg.setMinimumDuration(300)
|
||
self._batch_progress_dlg.show()
|
||
|
||
self._batch_thread.start()
|
||
|
||
def _on_batch_progress(self, pct: int):
|
||
if self._batch_progress_dlg is not None:
|
||
self._batch_progress_dlg.setValue(pct)
|
||
|
||
def _on_batch_file_done(self, path: str, err: str):
|
||
if err:
|
||
self._batch_errors.append(f"{Path(path).name} — {err}")
|
||
if self._batch_progress_dlg is not None:
|
||
self._batch_progress_dlg.setLabelText(f"Processed {Path(path).name}…")
|
||
|
||
def _on_batch_finished(self, paths: list[str]):
|
||
if self._batch_progress_dlg is not None:
|
||
self._batch_progress_dlg.close()
|
||
self._batch_progress_dlg = None
|
||
|
||
n_total = len(paths)
|
||
n_failed = len(self._batch_errors)
|
||
n_ok = n_total - n_failed
|
||
if n_failed:
|
||
summary = (f"Batch store: {n_ok}/{n_total} file(s) updated, "
|
||
f"{n_failed} failed: {'; '.join(self._batch_errors)}")
|
||
else:
|
||
summary = f"Batch store: {n_ok}/{n_total} file(s) updated."
|
||
self.statusBar().showMessage(summary)
|
||
self._batch_errors = []
|
||
|
||
# If the currently-open file was in this batch, reload it so the
|
||
# GUI picks up the newly-written v7 cache instead of stale state.
|
||
if self._sras is not None and str(self._sras.path) in paths:
|
||
self._load_file(str(self._sras.path))
|
||
|
||
def _on_batch_thread_finished(self):
|
||
# See the comment in _on_compute_thread_finished: wait() before
|
||
# releasing our reference to avoid destroying a QThread whose OS
|
||
# thread hasn't fully joined yet.
|
||
if self._batch_thread is not None:
|
||
self._batch_thread.wait()
|
||
self._batch_thread = None
|
||
self._batch_dc_act.setEnabled(True)
|
||
self._batch_fft_act.setEnabled(True)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Fusion: angle alignment
|
||
# ------------------------------------------------------------------
|
||
|
||
def _on_angle_alignment(self):
|
||
if self._sras is None or self._sras.n_angles <= 1:
|
||
return
|
||
if self._alignment_thread is not None:
|
||
return
|
||
ref_idx = 0
|
||
threshold_mv = self.spin_threshold_mv.value()
|
||
generation = self._alignment_generation
|
||
|
||
# Claim self._alignment_thread before anything below that can pump
|
||
# the Qt event loop — see the comment in _start_compute for why.
|
||
self._alignment_worker = AngleAlignmentWorker(self._sras, ref_idx, threshold_mv)
|
||
self._alignment_thread = QThread()
|
||
self._alignment_worker.moveToThread(self._alignment_thread)
|
||
self._alignment_thread.started.connect(self._alignment_worker.run)
|
||
self._alignment_worker.progress.connect(self._on_alignment_progress)
|
||
self._alignment_worker.finished.connect(
|
||
lambda result, err, g=generation: self._on_alignment_done(g, result, err))
|
||
self._alignment_worker.finished.connect(self._alignment_thread.quit)
|
||
self._alignment_thread.finished.connect(self._on_alignment_thread_finished)
|
||
|
||
self._alignment_act.setEnabled(False)
|
||
self._show_progress(
|
||
f"Computing angle alignment ({self._sras.n_angles} angles, "
|
||
f"ref=angle 0, CH4 mask ≥ {threshold_mv:.3f} mV)…"
|
||
)
|
||
self._alignment_thread.start()
|
||
|
||
def _on_alignment_progress(self, pct: int):
|
||
if self._progress_dlg is not None:
|
||
self._progress_dlg.setValue(pct)
|
||
|
||
def _on_alignment_thread_finished(self):
|
||
# See the comment in _on_compute_thread_finished: wait() before
|
||
# releasing our reference to avoid destroying a QThread whose OS
|
||
# thread hasn't fully joined yet.
|
||
if self._alignment_thread is not None:
|
||
self._alignment_thread.wait()
|
||
self._alignment_thread = None
|
||
self._update_controls_enabled(self._sras is not None)
|
||
|
||
def _on_alignment_done(self, generation: int, result, error_msg: str):
|
||
self._close_progress()
|
||
if generation != self._alignment_generation:
|
||
return # a new file was loaded while this was computing — discard
|
||
if error_msg:
|
||
self.statusBar().showMessage(f"Angle alignment failed: {error_msg}")
|
||
return
|
||
self._alignment_result = result
|
||
self._aligned_cache = {}
|
||
self.chk_aligned_view.setEnabled(True)
|
||
self.chk_aligned_view.blockSignals(True)
|
||
self.chk_aligned_view.setChecked(True)
|
||
self.chk_aligned_view.blockSignals(False)
|
||
nr, nc = result.canvas_shape
|
||
self.statusBar().showMessage(
|
||
f"Angle alignment computed ({self._sras.n_angles} angles, "
|
||
f"canvas {nc}×{nr} px).")
|
||
self._refresh_display()
|
||
|
||
def _on_aligned_view_toggled(self, checked: bool):
|
||
if self._current_image is not None:
|
||
self._redraw_image(self._current_image)
|
||
|
||
# ------------------------------------------------------------------
|
||
# FFT Options
|
||
# ------------------------------------------------------------------
|
||
|
||
def _on_fft_options(self):
|
||
global _fft_backend
|
||
spf = self._sras.samples_per_frame if self._sras is not None else None
|
||
sr = self._sras.sample_rate_hz if self._sras is not None else None
|
||
dlg = FftOptionsDialog(
|
||
self,
|
||
current_backend=_fft_backend,
|
||
current_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()
|
||
# Pad factor changes the FFT bin count, so it genuinely
|
||
# invalidates the cached raw FFT (part of the cache key) —
|
||
# _refresh_display() will recompute only on a cache miss.
|
||
if self._sras is not None and self.combo_channel.currentIndex() in (
|
||
CH1_IDX, VELOCITY_MODE_IDX):
|
||
self._refresh_display()
|
||
|
||
# ------------------------------------------------------------------
|
||
|
||
def closeEvent(self, event):
|
||
if self._dc_precompute_worker is not None:
|
||
self._dc_precompute_worker.stop()
|
||
for attr in ("_load_thread", "_compute_thread", "_batch_thread",
|
||
"_dc_precompute_thread", "_alignment_thread"):
|
||
t = getattr(self, attr, None)
|
||
if t is not None:
|
||
t.quit()
|
||
t.wait(2000)
|
||
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()
|