tka/sras-viewer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Skip FFT compute and disk I/O for below-threshold pixels in RF images

Browse Source 
Previously the FFT was always run unmasked (dc_threshold_mv=-1e9) so the
result could be cached independent of the threshold, with masking applied
afterward as a display-time step. That meant every pixel's CH1 waveform
was read and FFT'd even when most of the scan is background — wasted work
now that the DC map (and its threshold mask) is already known ahead of
time via the DC precompute cache.

compute_rf_image now takes the real threshold and an optional precomputed
dc4_mv array (reused from the DC cache, avoiding a redundant CH4 read),
and skips the FFT for masked-out pixels entirely, same as before this
threshold-caching detour was introduced. The masked-out CH1 samples are
also never read from disk: the boolean valid-pixel 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.

Threshold is therefore back to being part of the FFT cache key (changing
it now decides which pixels get computed at all, so it can't be satisfied
from a cache built for a different threshold) — but the recompute it
triggers reuses the cached DC4 image and skips both the FFT and the I/O
for masked pixels, so it's much cheaper than the original full-image
compute. Grating and colormap remain pure post-processing with no
recompute.

Verified bit-identical output against the un-optimized reference path on
synthetic data (including scattered, non-row-aligned masking), and timed
on the real 532 GB / 17-angle file: angle 0's FFT (75.2% of pixels above
the default 50 mV threshold) went from ~121s (fully unmasked) to 114s
(FFT skipped, I/O not skipped) to 90.2s with this change (both skipped).

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
Thomas Ales
2026-07-30 20:59:02 -05:00
parent 158f60d2a7
commit f745efe7e0
1 changed files with 102 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sras_viewer.py
+92 -65
View File
@@ -622,10 +622,21 @@ def compute_dc_image(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray:
def compute_rf_image(sras: SrasFile, angle_idx: int,
dc_threshold_mv: float,
apply_bg_sub: bool = True,
n_fft: int | None = None) -> np.ndarray:
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.
Pixels where CH4_dc < dc_threshold_mv are set to 0 — and the FFT is
never run for them, since that's the expensive part. The masked-out
CH1 samples are also never *read*: the boolean mask is applied to the
raw memmap slice before any dtype conversion, so numpy only pages in
the bytes for pixels that pass the threshold (an untouched memmap
page is never read from disk). If the DC4 image for this angle is
already known (e.g. from the DC-channel precompute cache), pass it as
*dc4_mv* (mV, shape (n_rows, n_frames)) to reuse it directly instead
of re-reading/re-averaging the CH4 channel here; otherwise it's
computed chunk-by-chunk internally (which does need every pixel's
CH4 data, since that's what determines validity in the first place).
Fast path: if the file contains v5 precomputed peak-frequency images,
and zero-padding is not active, and the bg-sub flag matches, the stored
@@ -662,24 +673,28 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
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_mv = adc_to_mv(dc4_raw.mean(axis=-1),
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_mv >= dc_threshold_mv # True = above threshold = run FFT
valid = dc4_chunk >= dc_threshold_mv # True = above threshold = run FFT
if not valid.any():
continue
waveforms = data[r0:r1, CH1_IDX, :, :].astype(np.float32)
# 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:
waveforms -= sras.background # background is 1-D (spf,)
valid_waves = waveforms[valid] # (n_valid, spf)
del waveforms
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
@@ -715,16 +730,15 @@ class LoadWorker(QObject):
class ComputeWorker(QObject):
"""Computes one displayable image for (angle, channel).
For CH1/Velocity (FFT-derived) channels, the FFT is run *unmasked*
(dc_threshold_mv=-1e9 → every pixel is computed) so the result is
reusable across threshold/grating/colormap changes without re-running
the FFT — those are then pure post-processing on the cached array.
The DC4 image needed to apply that mask at display time is computed
alongside it (cheap: a per-waveform mean, no FFT) and emitted with it
so the caller can cache both.
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) for DC
channels, or a ``(raw_freq_mhz, dc4_mv)`` tuple for CH1/Velocity.
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)
@@ -732,26 +746,26 @@ class ComputeWorker(QObject):
def __init__(self, sras: SrasFile, angle_idx: int,
ch_idx: int,
apply_bg_sub: bool = True,
n_fft: int | None = None):
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):
raw_freq_mhz = compute_rf_image(
self._sras, self._angle, dc_threshold_mv=-1e9,
apply_bg_sub=self._apply_bg_sub, n_fft=self._n_fft)
dc4_adc = compute_dc_image(self._sras, self._angle, CH4_IDX)
dc4_mv = adc_to_mv(dc4_adc,
self._sras.ch_ymult_mv[CH4_IDX],
self._sras.ch_yoff_adc[CH4_IDX],
self._sras.ch_yzero_mv[CH4_IDX])
self.finished.emit((raw_freq_mhz, dc4_mv))
img = compute_rf_image(
self._sras, self._angle, dc_threshold_mv=self._dc_threshold,
apply_bg_sub=self._apply_bg_sub, n_fft=self._n_fft,
dc4_mv=self._dc4_mv)
self.finished.emit(img)
else:
# DC channels: convert ADC counts → mV
adc_img = compute_dc_image(self._sras, self._angle, self._ch)
@@ -1478,6 +1492,7 @@ class SrasViewerWindow(QMainWindow):
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)
@@ -1486,15 +1501,17 @@ class SrasViewerWindow(QMainWindow):
self._preprocess_thread: QThread | None = None
# Display-only settings (colormap, threshold, grating) no longer
# 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
# is viewed in that mode, cached unmasked, and re-masked/re-scaled
# for free afterward.
# 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], 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
@@ -2125,8 +2142,10 @@ class SrasViewerWindow(QMainWindow):
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 is applied as a mask against the cached raw FFT + DC4
# image, not baked into the compute — never needs a recompute.
# 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()
@@ -2165,22 +2184,18 @@ class SrasViewerWindow(QMainWindow):
return None
return self._sras.samples_per_frame * pad_factor
def _apply_threshold_and_scale(self, raw_freq_mhz: np.ndarray,
dc4_mv: np.ndarray, ch_idx: int) -> np.ndarray:
"""Mask a raw (unmasked) FFT peak-frequency image against the
current DC threshold, and scale to velocity if needed. Both are
cheap array ops — never worth a recompute on their own."""
img = raw_freq_mhz.copy()
img[dc4_mv < self.spin_threshold_mv.value()] = 0.0
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:
img = img * self.spin_grating_um.value()
return img
return freq_mhz * self.spin_grating_um.value()
return freq_mhz
def _refresh_display(self):
"""Show the image for the current angle/channel, 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."""
"""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()
@@ -2193,15 +2208,19 @@ class SrasViewerWindow(QMainWindow):
return
elif ch_idx in (CH1_IDX, VELOCITY_MODE_IDX):
apply_bg_sub = self.chk_bg_sub.isChecked()
key = (angle_idx, apply_bg_sub, self._current_n_fft())
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)
dc4 = self._dc_cache.get((angle_idx, CH4_IDX))
if raw is not None and dc4 is not None:
img = self._apply_threshold_and_scale(raw, dc4, ch_idx)
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.
# 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):
@@ -2226,12 +2245,18 @@ class SrasViewerWindow(QMainWindow):
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
@@ -2243,6 +2268,7 @@ class SrasViewerWindow(QMainWindow):
# 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)
@@ -2259,7 +2285,7 @@ class SrasViewerWindow(QMainWindow):
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 will be instant."
"so revisiting this angle/mode/threshold will be instant."
)
else:
self.statusBar().showMessage("Computing DC image…")
@@ -2279,9 +2305,10 @@ class SrasViewerWindow(QMainWindow):
angle_idx = self.spin_angle.value()
ch_idx = self.combo_channel.currentIndex()
apply_bg_sub = self.chk_bg_sub.isChecked()
if (angle_idx, ch_idx, apply_bg_sub, self._fft_pad_factor) != (
self._pending_angle, self._pending_ch,
self._pending_bg_sub, self._pending_fft_pad_factor):
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.
@@ -2293,11 +2320,11 @@ class SrasViewerWindow(QMainWindow):
ch_idx = self._pending_ch
if ch_idx in (CH1_IDX, VELOCITY_MODE_IDX):
raw_freq_mhz, dc4_mv = result
key = (angle_idx, self._pending_bg_sub, self._current_n_fft())
self._fft_cache[key] = raw_freq_mhz
self._dc_cache[(angle_idx, CH4_IDX)] = dc4_mv
img = self._apply_threshold_and_scale(raw_freq_mhz, dc4_mv, ch_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