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, def compute_rf_image(sras: SrasFile, angle_idx: int,
dc_threshold_mv: float, dc_threshold_mv: float,
apply_bg_sub: bool = True, 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. """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, 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 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) r1 = min(r0 + chunk_rows, n_rows)
# DC mask for this chunk (float32 expansion is only chunk-sized) # 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_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_ymult_mv[CH4_IDX],
sras.ch_yoff_adc[CH4_IDX], sras.ch_yoff_adc[CH4_IDX],
sras.ch_yzero_mv[CH4_IDX]) sras.ch_yzero_mv[CH4_IDX])
del dc4_raw 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(): if not valid.any():
continue 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: if apply_bg_sub and sras.background is not None:
waveforms -= sras.background # background is 1-D (spf,) valid_waves -= sras.background # background is 1-D (spf,)
valid_waves = waveforms[valid] # (n_valid, spf)
del waveforms
fft_pow = np.abs(_do_rfft(valid_waves, n=n_fft, axis=-1, workers=_n_workers)) ** 2 fft_pow = np.abs(_do_rfft(valid_waves, n=n_fft, axis=-1, workers=_n_workers)) ** 2
del valid_waves del valid_waves
@@ -715,16 +730,15 @@ class LoadWorker(QObject):
class ComputeWorker(QObject): class ComputeWorker(QObject):
"""Computes one displayable image for (angle, channel). """Computes one displayable image for (angle, channel).
For CH1/Velocity (FFT-derived) channels, the FFT is run *unmasked* For CH1/Velocity (FFT-derived) channels, the FFT is only run for
(dc_threshold_mv=-1e9 → every pixel is computed) so the result is pixels whose DC4 (Bias B) mean is at or above dc_threshold_mv — masked
reusable across threshold/grating/colormap changes without re-running pixels are left at 0 MHz without ever being FFT'd, since that's the
the FFT — those are then pure post-processing on the cached array. expensive part of a scan. If the DC4 image for this angle is already
The DC4 image needed to apply that mask at display time is computed known (e.g. from the DC-channel precompute cache), pass it in as
alongside it (cheap: a per-waveform mean, no FFT) and emitted with it *dc4_mv* to skip re-reading the CH4 channel from disk entirely.
so the caller can cache both.
Emits a plain ``np.ndarray`` (already in display units) for DC Emits a plain ``np.ndarray`` (already in display units, masked for
channels, or a ``(raw_freq_mhz, dc4_mv)`` tuple for CH1/Velocity. CH1/Velocity) for both DC and FFT-derived channels.
""" """
finished = pyqtSignal(object) finished = pyqtSignal(object)
error = pyqtSignal(str) error = pyqtSignal(str)
@@ -732,26 +746,26 @@ class ComputeWorker(QObject):
def __init__(self, sras: SrasFile, angle_idx: int, def __init__(self, sras: SrasFile, angle_idx: int,
ch_idx: int, ch_idx: int,
apply_bg_sub: bool = True, 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__() super().__init__()
self._sras = sras self._sras = sras
self._angle = angle_idx self._angle = angle_idx
self._ch = ch_idx self._ch = ch_idx
self._apply_bg_sub = apply_bg_sub self._apply_bg_sub = apply_bg_sub
self._n_fft = n_fft self._n_fft = n_fft
self._dc_threshold = dc_threshold_mv
self._dc4_mv = dc4_mv
def run(self): def run(self):
try: try:
if self._ch in (CH1_IDX, VELOCITY_MODE_IDX): if self._ch in (CH1_IDX, VELOCITY_MODE_IDX):
raw_freq_mhz = compute_rf_image( img = compute_rf_image(
self._sras, self._angle, dc_threshold_mv=-1e9, self._sras, self._angle, dc_threshold_mv=self._dc_threshold,
apply_bg_sub=self._apply_bg_sub, n_fft=self._n_fft) 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=self._dc4_mv)
dc4_mv = adc_to_mv(dc4_adc, self.finished.emit(img)
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))
else: else:
# DC channels: convert ADC counts → mV # DC channels: convert ADC counts → mV
adc_img = compute_dc_image(self._sras, self._angle, self._ch) 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_angle: int = 0
self._pending_ch: int = 0 self._pending_ch: int = 0
self._pending_bg_sub: bool = True self._pending_bg_sub: bool = True
self._pending_threshold: float = 50.0 # mV
self._progress_dlg: QProgressDialog | None = None self._progress_dlg: QProgressDialog | None = None
# FFT settings (configured via FFT Options dialog) # FFT settings (configured via FFT Options dialog)
@@ -1486,15 +1501,17 @@ class SrasViewerWindow(QMainWindow):
self._preprocess_thread: QThread | None = None self._preprocess_thread: QThread | None = None
# Display-only settings (colormap, threshold, grating) no longer # Display-only settings (colormap, grating) never trigger a
# trigger a recompute — they're applied to cached data on redraw. # recompute — they're applied to cached data on redraw. DC images
# DC images (CH3/CH4) are cheap and precomputed for every angle in # (CH3/CH4) are cheap and precomputed for every angle in the
# the background right after load; CH1/Velocity FFT images are # background right after load. CH1/Velocity FFT images are
# computed lazily (with a progress popup) the first time an angle # computed lazily (with a progress popup) the first time an
# is viewed in that mode, cached unmasked, and re-masked/re-scaled # angle/threshold combination is viewed — using the cached DC4
# for free afterward. # 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._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_thread: QThread | None = None
self._dc_precompute_worker: DcPrecomputeWorker | None = None self._dc_precompute_worker: DcPrecomputeWorker | None = None
self._dc_generation: int = 0 self._dc_generation: int = 0
@@ -2125,8 +2142,10 @@ class SrasViewerWindow(QMainWindow):
else: else:
ymult, yoff, yzero = _FALLBACK_YMULT_MV, _FALLBACK_YOFF_ADC, 0.0 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") 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 # Threshold decides which pixels get an FFT at all (see
# image, not baked into the compute — never needs a recompute. # 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: if self._sras is not None and self.combo_channel.currentIndex() in CH1_DERIVED_MODES:
self._refresh_display() self._refresh_display()
@@ -2165,22 +2184,18 @@ class SrasViewerWindow(QMainWindow):
return None return None
return self._sras.samples_per_frame * pad_factor return self._sras.samples_per_frame * pad_factor
def _apply_threshold_and_scale(self, raw_freq_mhz: np.ndarray, def _scale_for_display(self, freq_mhz: np.ndarray, ch_idx: int) -> np.ndarray:
dc4_mv: np.ndarray, ch_idx: int) -> np.ndarray: """Velocity is a pure post-multiply of the (already DC-masked)
"""Mask a raw (unmasked) FFT peak-frequency image against the cached frequency image — never worth a recompute on its own."""
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
if ch_idx == VELOCITY_MODE_IDX: if ch_idx == VELOCITY_MODE_IDX:
img = img * self.spin_grating_um.value() return freq_mhz * self.spin_grating_um.value()
return img return freq_mhz
def _refresh_display(self): def _refresh_display(self):
"""Show the image for the current angle/channel, using cached data """Show the image for the current angle/channel/threshold, using
whenever possible and only falling back to a background compute cached data whenever possible and only falling back to a
(with progress popup) when genuinely nothing is cached yet for background compute (with progress popup) when genuinely nothing
these settings.""" is cached yet for these settings."""
if self._sras is None: if self._sras is None:
return return
angle_idx = self.spin_angle.value() angle_idx = self.spin_angle.value()
@@ -2193,15 +2208,19 @@ class SrasViewerWindow(QMainWindow):
return return
elif ch_idx in (CH1_IDX, VELOCITY_MODE_IDX): elif ch_idx in (CH1_IDX, VELOCITY_MODE_IDX):
apply_bg_sub = self.chk_bg_sub.isChecked() 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) raw = self._fft_cache.get(key)
dc4 = self._dc_cache.get((angle_idx, CH4_IDX)) if raw is not None:
if raw is not None and dc4 is not None: img = self._scale_for_display(raw, ch_idx)
img = self._apply_threshold_and_scale(raw, dc4, ch_idx)
self._show_image_now(img, angle_idx, ch_idx) self._show_image_now(img, angle_idx, ch_idx)
return 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() self._start_compute()
def _show_image_now(self, img: np.ndarray, angle_idx: int, ch_idx: int): 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() angle_idx = self.spin_angle.value()
ch_idx = self.combo_channel.currentIndex() ch_idx = self.combo_channel.currentIndex()
apply_bg_sub = self.chk_bg_sub.isChecked() apply_bg_sub = self.chk_bg_sub.isChecked()
threshold_mv = self.spin_threshold_mv.value()
pad_factor = self._fft_pad_factor pad_factor = self._fft_pad_factor
n_fft = self._current_n_fft() 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_angle = angle_idx
self._pending_ch = ch_idx self._pending_ch = ch_idx
self._pending_bg_sub = apply_bg_sub self._pending_bg_sub = apply_bg_sub
self._pending_threshold = threshold_mv
self._pending_fft_pad_factor = pad_factor self._pending_fft_pad_factor = pad_factor
# Claim self._compute_thread *before* anything below that can pump # Claim self._compute_thread *before* anything below that can pump
@@ -2243,6 +2268,7 @@ class SrasViewerWindow(QMainWindow):
# process. Assigning immediately closes that window. # process. Assigning immediately closes that window.
self._compute_worker = ComputeWorker( self._compute_worker = ComputeWorker(
self._sras, angle_idx, ch_idx, apply_bg_sub, n_fft=n_fft, 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_thread = QThread()
self._compute_worker.moveToThread(self._compute_thread) self._compute_worker.moveToThread(self._compute_thread)
@@ -2259,7 +2285,7 @@ class SrasViewerWindow(QMainWindow):
self._show_progress( self._show_progress(
f"Computing FFT for angle {angle_idx}…\n" f"Computing FFT for angle {angle_idx}…\n"
"This can take a while on a large scan — result is cached " "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: else:
self.statusBar().showMessage("Computing DC image…") self.statusBar().showMessage("Computing DC image…")
@@ -2279,9 +2305,10 @@ class SrasViewerWindow(QMainWindow):
angle_idx = self.spin_angle.value() angle_idx = self.spin_angle.value()
ch_idx = self.combo_channel.currentIndex() ch_idx = self.combo_channel.currentIndex()
apply_bg_sub = self.chk_bg_sub.isChecked() apply_bg_sub = self.chk_bg_sub.isChecked()
if (angle_idx, ch_idx, apply_bg_sub, self._fft_pad_factor) != ( threshold_mv = self.spin_threshold_mv.value()
self._pending_angle, self._pending_ch, if (angle_idx, ch_idx, apply_bg_sub, threshold_mv, self._fft_pad_factor) != (
self._pending_bg_sub, self._pending_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 # Settings changed while this compute was running — re-dispatch
# through the cache-aware path in case that now-current # through the cache-aware path in case that now-current
# combination happens to already be cached. # combination happens to already be cached.
@@ -2293,11 +2320,11 @@ class SrasViewerWindow(QMainWindow):
ch_idx = self._pending_ch ch_idx = self._pending_ch
if ch_idx in (CH1_IDX, VELOCITY_MODE_IDX): if ch_idx in (CH1_IDX, VELOCITY_MODE_IDX):
raw_freq_mhz, dc4_mv = result masked_freq_mhz = result
key = (angle_idx, self._pending_bg_sub, self._current_n_fft()) key = (angle_idx, self._pending_bg_sub, self._current_n_fft(),
self._fft_cache[key] = raw_freq_mhz self._pending_threshold)
self._dc_cache[(angle_idx, CH4_IDX)] = dc4_mv self._fft_cache[key] = masked_freq_mhz
img = self._apply_threshold_and_scale(raw_freq_mhz, dc4_mv, ch_idx) img = self._scale_for_display(masked_freq_mhz, ch_idx)
else: else:
img = result img = result
self._dc_cache[(angle_idx, ch_idx)] = img self._dc_cache[(angle_idx, ch_idx)] = img