pre merge cleanup
This commit is contained in:
Thomas Ales
+76
-27
@@ -18,7 +18,9 @@ import numpy as np
|
||||
import scipy.fft as scipy_fft
|
||||
import scipy.ndimage as scipy_ndimage
|
||||
|
||||
from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, SrasFile, adc_to_mv
|
||||
from sras_format import (
|
||||
CH1_IDX, CH3_IDX, CH4_IDX, PAD_FACTOR_MAX, SrasFile, adc_to_mv,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# FFT backend
|
||||
@@ -420,6 +422,56 @@ def dc_image_mv(sras: SrasFile, angle_idx: int, ch_idx: int,
|
||||
*sras.cal(ch_idx))
|
||||
|
||||
|
||||
def cached_rf_image(sras: SrasFile, angle_idx: int,
|
||||
dc_threshold_mv: float | None,
|
||||
apply_bg_sub: bool = True,
|
||||
n_fft: int | None = None,
|
||||
dc4_mv: np.ndarray | None = None,
|
||||
allow_dc_recompute: bool = True) -> np.ndarray | None:
|
||||
"""The stored peak-frequency image for one angle (v5 PREC or v7 CACH),
|
||||
masked and ready to display — or None if no stored image can serve these
|
||||
settings and a real FFT is needed.
|
||||
|
||||
Stored images are unmasked, so they serve any threshold — but they carry
|
||||
one bg-sub state and one zero-padding, and a view asking for either of the
|
||||
others is asking for different numbers, not a different rendering of the
|
||||
same ones. Padding is the sharp edge: a pad-10 view resolves peaks on a
|
||||
10x finer bin grid, so handing back a pad-1 image would quietly undo the
|
||||
setting. v5 PREC and CACH v1 tails record no pad and are natural
|
||||
resolution by construction.
|
||||
|
||||
Cheap enough to call on the GUI thread — a copy plus a comparison —
|
||||
*except* when the CH4 mask has to be read from disk to build it. Pass
|
||||
*allow_dc_recompute* False to return None in that case instead, and leave
|
||||
the whole-channel read to a worker.
|
||||
"""
|
||||
cached_freq = sras.precomputed_freq_mhz[angle_idx]
|
||||
# n_fft None means "no padding", i.e. exactly samples_per_frame points.
|
||||
want_n_fft = n_fft if n_fft is not None else sras.samples_per_frame
|
||||
if (cached_freq is None
|
||||
or want_n_fft != sras.precomputed_pad_factor * sras.samples_per_frame
|
||||
or sras.precomputed_bg_sub != (apply_bg_sub and sras.background is not None)):
|
||||
return None
|
||||
|
||||
if dc_threshold_mv is None:
|
||||
return cached_freq.copy()
|
||||
|
||||
# DC4 mask, in priority order: stored DC block, caller-supplied image,
|
||||
# or a fresh (cheap — no FFT) recompute.
|
||||
dc4_img = sras.cached_dc_mv(angle_idx, CH4_IDX)
|
||||
if dc4_img is None:
|
||||
dc4_img = dc4_mv
|
||||
if dc4_img is None:
|
||||
if not allow_dc_recompute:
|
||||
return None
|
||||
dc4_img = adc_to_mv(compute_dc_image(sras, angle_idx, CH4_IDX),
|
||||
*sras.cal(CH4_IDX))
|
||||
|
||||
freq_img = cached_freq.copy()
|
||||
freq_img[dc4_img < dc_threshold_mv] = 0.0
|
||||
return freq_img
|
||||
|
||||
|
||||
def compute_rf_image(sras: SrasFile, angle_idx: int,
|
||||
dc_threshold_mv: float | None,
|
||||
apply_bg_sub: bool = True,
|
||||
@@ -455,28 +507,20 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
|
||||
|
||||
Fast path: if the file has a precomputed peak-frequency image for this
|
||||
angle (v5 PREC or v7 CACH), zero-padding is off, and the bg-sub flag
|
||||
matches, the stored image is used directly — no FFT is run.
|
||||
matches, the stored image is used directly — no FFT is run. That test is
|
||||
cached_rf_image, which the viewer also applies before it ever dispatches
|
||||
a compute at all.
|
||||
"""
|
||||
n_rows, n_frames = sras.image_shape(angle_idx)
|
||||
data = sras.data[angle_idx]
|
||||
|
||||
# ---- Fast path: precomputed image (v5 PREC or v7 CACH) ----------------
|
||||
cached_freq = sras.precomputed_freq_mhz[angle_idx]
|
||||
if (cached_freq is not None
|
||||
and n_fft is None # no custom zero-padding
|
||||
and sras.precomputed_bg_sub == (apply_bg_sub and sras.background is not None)):
|
||||
freq_img = cached_freq.copy()
|
||||
if dc_threshold_mv is not None:
|
||||
# DC4 mask, in priority order: already-cached DC block, caller-
|
||||
# supplied image, or a fresh (cheap — no FFT) recompute.
|
||||
dc4_img = sras.cached_dc_mv(angle_idx, CH4_IDX)
|
||||
if dc4_img is None:
|
||||
dc4_img = dc4_mv if dc4_mv is not None else adc_to_mv(
|
||||
compute_dc_image(sras, angle_idx, CH4_IDX), *sras.cal(CH4_IDX))
|
||||
freq_img[dc4_img < dc_threshold_mv] = 0.0
|
||||
return freq_img
|
||||
cached = cached_rf_image(sras, angle_idx, dc_threshold_mv,
|
||||
apply_bg_sub=apply_bg_sub, n_fft=n_fft,
|
||||
dc4_mv=dc4_mv)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
# ---- Chunked FFT path --------------------------------------------------
|
||||
n_rows, n_frames = sras.image_shape(angle_idx)
|
||||
data = sras.data[angle_idx]
|
||||
exact = exact or _FFT_EXACT_ENV
|
||||
spf = sras.samples_per_frame
|
||||
n_len = n_fft if n_fft is not None else spf
|
||||
@@ -589,17 +633,19 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def cache_file(path: str, mode: str, apply_bg_sub: bool,
|
||||
fft_backend: str = "scipy", max_workers: int = 0) -> str:
|
||||
fft_backend: str = "scipy", max_workers: int = 0,
|
||||
pad_factor: int = 1) -> str:
|
||||
"""Compute and store DC or FFT images for every angle of one file,
|
||||
converting v6 → v7 in place. Returns "" on success or an error message.
|
||||
|
||||
Module-level and picklable so it can run in a ProcessPoolExecutor. The
|
||||
FFT backend and worker cap are passed explicitly because module globals
|
||||
do not survive a spawn.
|
||||
FFT backend, worker cap and pad factor are passed explicitly because
|
||||
module globals do not survive a spawn.
|
||||
|
||||
The stored FFT cache is always natural-resolution (pad 1): the v7 SFFT
|
||||
block records no pad factor, and padded views compute live fast enough
|
||||
(see _peak_bins_zoom) that caching them is not worth a format change.
|
||||
*pad_factor* is stored alongside the images: a cache is only usable by a
|
||||
view asking for the same padding, so caching at the viewer's own setting
|
||||
is the difference between a batch that pays off and one that can never be
|
||||
read back (see cached_rf_image).
|
||||
"""
|
||||
global _MAX_WORKERS
|
||||
try:
|
||||
@@ -630,15 +676,18 @@ def cache_file(path: str, mode: str, apply_bg_sub: bool,
|
||||
sras.write_v7_cache(new_dc3_mv=dc3, new_dc4_mv=dc4)
|
||||
else:
|
||||
effective_bg = apply_bg_sub and sras.background is not None
|
||||
pad = max(1, min(PAD_FACTOR_MAX, int(pad_factor)))
|
||||
n_fft = None if pad <= 1 else sras.samples_per_frame * pad
|
||||
# dc_threshold_mv=None: store unmasked images and mask at display
|
||||
# time (same convention as v5's PREC block). Skipping the mask
|
||||
# also skips reading CH4 entirely. The FFT path parallelises
|
||||
# internally over blocks, so angles run one at a time with the
|
||||
# full budget.
|
||||
freq = [compute_rf_image(sras, a, dc_threshold_mv=None,
|
||||
apply_bg_sub=effective_bg)
|
||||
apply_bg_sub=effective_bg, n_fft=n_fft)
|
||||
for a in range(n)]
|
||||
sras.write_v7_cache(new_freq_mhz=freq, new_bg_sub=effective_bg)
|
||||
sras.write_v7_cache(new_freq_mhz=freq, new_bg_sub=effective_bg,
|
||||
new_pad_factor=pad)
|
||||
return ""
|
||||
except Exception as exc:
|
||||
return str(exc)
|
||||
|
||||
Reference in New Issue
Block a user