From f745efe7e0e6740844f6288b0a44ecc5248c0790 Mon Sep 17 00:00:00 2001 From: Thomas Ales Date: Thu, 30 Jul 2026 20:59:02 -0500 Subject: [PATCH] Skip FFT compute and disk I/O for below-threshold pixels in RF images MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- sras_viewer.py | 177 ++++++++++++++++++++++++++++--------------------- 1 file changed, 102 insertions(+), 75 deletions(-) diff --git a/sras_viewer.py b/sras_viewer.py index 3ec743c..1001ae1 100644 --- a/sras_viewer.py +++ b/sras_viewer.py @@ -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) - dc4_raw = data[r0:r1, CH4_IDX, :, :].astype(np.float32) - dc4_mv = adc_to_mv(dc4_raw.mean(axis=-1), - sras.ch_ymult_mv[CH4_IDX], - sras.ch_yoff_adc[CH4_IDX], - sras.ch_yzero_mv[CH4_IDX]) - del dc4_raw - valid = dc4_mv >= dc_threshold_mv # True = above threshold = run FFT + if 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 - 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._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