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