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

Implement row-averaged FFT feature for same-row SNR cleanup

Browse Source 
Add row-averaged FFT feature with configurable window size (row_avg_n
parameter) for improved signal-to-noise ratio on noisy scans. Includes:

- Gaussian-weighted same-row neighbor averaging (never crosses rows)
- Masked/renormalized convolution handling for edge cases and masked samples
- Cache format v2 with row_avg_n tracking to prevent silent cache mismatches
- GUI dialog option for row-average window configuration
- Comprehensive tests validating kernel properties, background subtraction
  invariance, and cache dispatch

Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
This commit is contained in:
Thomas Ales
2026-08-07 21:50:05 -05:00
parent eecb0e3a82
commit 3989c2a1b8
10 changed files with 1363 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_row_average.py
+278
View File
@@ -0,0 +1,278 @@
"""Row-averaged FFT: same-row, distance-weighted CH1 waveform smoothing.
Covers the properties the design depends on: the kernel is symmetric and
n=0 is a true no-op; the masked/renormalized convolution matches an
independent brute-force reference and gives masked neighbors exactly zero
weight regardless of their content; background subtraction after averaging
is algebraically identical to subtracting before; chunking/worker count
never changes the result; and a pixel that's itself masked is never
"rescued" by averaging.
"""
import numpy as np
import pytest
import sras_compute as compute
from sras_compute import compute_rf_image, dc_image_mv
from sras_format import CH4_IDX, SrasFile
import tools.make_test_sras as gen
def _reference_row_average(masked_waves: np.ndarray, valid: np.ndarray,
weights: np.ndarray) -> np.ndarray:
"""Independent, unvectorized reference for _row_average_waveforms: for
each row position, sum weighted valid neighbors within the kernel's
radius and normalize by the actual included weight sum. Same
definition, computed by brute-force nested loops instead of
correlate1d, so it can't share a bug with the implementation."""
n_frames, spf = masked_waves.shape
n = len(weights) // 2
out = np.zeros_like(masked_waves)
for i in range(n_frames):
num = np.zeros(spf, dtype=np.float64)
den = 0.0
for d in range(-n, n + 1):
j = i + d
if 0 <= j < n_frames and valid[j]:
w = float(weights[d + n])
num += w * masked_waves[j].astype(np.float64)
den += w
out[i] = num / den if den > 0 else 0.0
return out
# ---------------------------------------------------------------------------
# _row_average_weights
# ---------------------------------------------------------------------------
def test_row_average_weights_shape_and_symmetry():
w0 = compute._row_average_weights(0)
assert w0.shape == (1,) and w0[0] == 1.0
for n in (1, 2, 5):
w = compute._row_average_weights(n)
assert w.shape == (2 * n + 1,)
assert w[n] == pytest.approx(1.0), "center tap is the peak weight"
assert np.allclose(w, w[::-1]), "symmetric about the center"
half = w[n:]
assert np.all(np.diff(half) < 0), "strictly decreasing away from center"
# ---------------------------------------------------------------------------
# _row_average_waveforms
# ---------------------------------------------------------------------------
def test_row_average_matches_hand_rolled_reference():
rng = np.random.default_rng(0)
n_frames, spf = 15, 6
raw = rng.integers(-50, 51, size=(n_frames, spf)).astype(np.float32)
valid = np.ones(n_frames, dtype=bool)
valid[[2, 3, 9]] = False # a run of two invalid, plus a lone invalid
masked = raw.copy()
masked[~valid] = 0.0
weights = compute._row_average_weights(3)
got = compute._row_average_waveforms(masked, valid, weights)
ref = _reference_row_average(masked, valid, weights)
assert np.allclose(got[valid], ref[valid], atol=1e-4)
def test_row_average_edge_of_row():
"""A window wider than the row itself must still renormalize correctly
at both ends -- mode='constant', cval=0.0 zero-pads both the numerator
and denominator, so this is not a special case, but it's the one most
likely to break if that padding were ever mismatched between the two."""
n_frames, spf = 6, 3
raw = np.arange(n_frames * spf, dtype=np.float32).reshape(n_frames, spf)
valid = np.ones(n_frames, dtype=bool)
weights = compute._row_average_weights(4) # window (9 taps) > n_frames (6)
got = compute._row_average_waveforms(raw, valid, weights)
ref = _reference_row_average(raw, valid, weights)
assert np.allclose(got, ref, atol=1e-4)
def test_row_average_excludes_masked_neighbor_from_normalization():
"""A masked neighbor must contribute zero *weight* to the normalization,
not participate as a legitimate zero-valued sample at full weight --
the two give different answers, and only the former is correct. (Note:
masked_waves must already be 0 at invalid positions per
_row_average_waveforms's contract -- that's what read_row's zero-filled
scratch buffer guarantees in production -- so the only way to vary
"what a masked position looks like" while respecting that contract is
whether its weight is excluded from the denominator at all.)"""
n_frames, spf = 9, 4
weights = compute._row_average_weights(2)
# A: position 4 is masked -- excluded from the weight sum entirely.
valid_a = np.ones(n_frames, dtype=bool)
valid_a[4] = False
masked_a = np.zeros((n_frames, spf), dtype=np.float32)
masked_a[valid_a] = 1.0
got_a = compute._row_average_waveforms(masked_a, valid_a, weights)
# B: position 4 is valid but genuinely zero-valued -- included in the
# weight sum, diluting neighbors' averages.
valid_b = np.ones(n_frames, dtype=bool)
masked_b = np.ones((n_frames, spf), dtype=np.float32)
masked_b[4] = 0.0
got_b = compute._row_average_waveforms(masked_b, valid_b, weights)
# Every position whose window reaches index 4 must average *higher* in
# A (excluded from the denominator) than in B (included as a real zero).
affected = [2, 3, 5, 6]
assert np.all(got_a[affected] > got_b[affected]), \
"masking must exclude a neighbor from normalization, not just zero its value"
# Positions outside the window (radius 2) are unaffected either way.
assert np.allclose(got_a[[0, 1, 7, 8]], got_b[[0, 1, 7, 8]])
def test_background_subtracted_once_equals_subtract_then_average():
"""Algebraic identity the implementation relies on: subtracting a fixed
background from the already-averaged waveform equals subtracting it
from every valid neighbor first, because the denominator is always the
*actual* included weight sum (never a fixed total)."""
rng = np.random.default_rng(1)
n_frames, spf = 11, 8
raw = rng.integers(-40, 41, size=(n_frames, spf)).astype(np.float32)
valid = np.ones(n_frames, dtype=bool)
valid[[1, 7]] = False
masked = raw.copy()
masked[~valid] = 0.0
background = rng.integers(-5, 6, size=spf).astype(np.float32)
weights = compute._row_average_weights(3)
# Order A (what the code does): average first, subtract background once.
order_a = compute._row_average_waveforms(masked, valid, weights) - background
# Order B: subtract background from every valid neighbor first (restoring
# the "0 at invalid positions" contract afterward), then average.
bg_subbed = masked - background
bg_subbed[~valid] = 0.0
order_b = compute._row_average_waveforms(bg_subbed, valid, weights)
assert np.allclose(order_a[valid], order_b[valid], atol=1e-3)
# ---------------------------------------------------------------------------
# compute_rf_image(row_avg_n=...) integration
# ---------------------------------------------------------------------------
def test_row_average_zero_is_identity(tmp_path):
"""row_avg_n=0 must take the exact same code path as before this
feature existed (row_avg_weights stays None), not a single-tap kernel
that merely computes to the same answer."""
path = tmp_path / "zero.sras"
gen.write(path, n_angles=1, seed=10, samples_per_frame=64)
sras = SrasFile(str(path))
plain = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True)
explicit_zero = compute_rf_image(sras, 0, dc_threshold_mv=None,
apply_bg_sub=True, row_avg_n=0)
assert np.array_equal(plain, explicit_zero)
def test_row_average_respects_own_center_mask(tmp_path):
"""A pixel that's itself below threshold stays masked (0) after row
averaging -- averaging never rescues a masked pixel, matching the
'valid neighbors only' design (masked pixels are excluded from other
pixels' averages, and are never themselves smoothed)."""
path = tmp_path / "center_mask.sras"
gen.write(path, n_angles=1, seed=13, samples_per_frame=64)
sras = SrasFile(str(path))
dc4 = dc_image_mv(sras, 0, CH4_IDX)
thr = float(np.percentile(dc4, 50))
mask = dc4 >= thr
assert mask.any() and not mask.all(), "threshold actually splits the image"
img = compute_rf_image(sras, 0, dc_threshold_mv=thr, apply_bg_sub=True,
row_avg_n=4)
assert np.array_equal(img == 0, ~mask), \
"masked pixels stay exactly 0 after row averaging; valid ones don't"
def test_row_average_composes_with_padding(tmp_path):
"""row_avg_n and n_fft (zero-padding) are independent knobs: using them
together must not raise, and must still agree with the exact (non-zoom)
reference path at that pad factor -- i.e. row-averaging composes with
the zoom peak search correctly, not just with the direct one."""
path = tmp_path / "padded_rowavg.sras"
gen.write(path, n_angles=1, seed=12, samples_per_frame=64)
sras = SrasFile(str(path))
spf = sras.samples_per_frame
raw_natural = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True)
avg_natural = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True,
row_avg_n=3)
avg_padded_zoom = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True,
row_avg_n=3, n_fft=spf * 40)
avg_padded_exact = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True,
row_avg_n=3, n_fft=spf * 40, exact=True)
assert avg_natural.shape == raw_natural.shape == avg_padded_zoom.shape
assert np.all(np.isfinite(avg_padded_zoom))
assert np.array_equal(avg_padded_zoom, avg_padded_exact), \
"row-averaged waveforms feed the zoom and exact FFT paths identically"
def test_row_average_improves_snr_recovery():
"""The actual point of the feature: averaging same-row waveforms that
share a true underlying tone but carry independent noise recovers that
tone far more reliably than any single raw (unaveraged) waveform does."""
rng = np.random.default_rng(42)
n_frames, spf = 21, 128
true_bin = 9
t = np.arange(spf)
tone = 15.0 * np.sin(2 * np.pi * true_bin * t / spf) # same true signal
# at every position
noise_sigma = 40.0 # much larger than the tone -- deliberately poor SNR
raw = (tone[None, :] + rng.normal(scale=noise_sigma, size=(n_frames, spf))
).astype(np.float32)
valid = np.ones(n_frames, dtype=bool)
weights = compute._row_average_weights(8) # wide window: lots of averaging
averaged = compute._row_average_waveforms(raw, valid, weights)
raw_bins = compute._peak_bins_direct(raw, spf)
avg_bins = compute._peak_bins_direct(averaged, spf)
raw_hits = int(np.sum(raw_bins == true_bin))
avg_hits = int(np.sum(avg_bins == true_bin))
assert avg_hits > raw_hits, (
f"row averaging should recover the true bin ({true_bin}) more often "
f"than raw per-pixel estimates: raw {raw_hits}/{n_frames}, "
f"averaged {avg_hits}/{n_frames}")
assert avg_hits >= n_frames * 0.7, \
f"averaged recovery should be reliable, not just barely better: {avg_hits}/{n_frames}"
def test_row_average_parallel_identity(tmp_path, monkeypatch):
"""Forcing 1 worker vs many must give an identical row-averaged image --
catches chunk-boundary bugs (there should be none, since averaging never
crosses rows, but this is the empirical proof, not just inspection)."""
path = tmp_path / "parallel_rowavg.sras"
n_rows, n_frames, spf = 40, 13, 128
gen.write(path, n_angles=1, seed=11, samples_per_frame=spf,
geometry=[(n_rows, n_frames)])
sras = SrasFile(str(path))
monkeypatch.setattr(compute, "_TOTAL_BYTES_BUDGET", 8 * n_frames * spf * 4)
monkeypatch.setattr(compute, "_FFT_BLOCK", 4)
# row_avg_n > 0 halves the effective budget before chunk planning.
fft_rows = compute._plan_fft_rows(n_frames, spf, compute._TOTAL_BYTES_BUDGET // 2)
assert fft_rows < n_rows, \
f"row-averaged FFT work actually splits into multiple chunks ({fft_rows} of {n_rows})"
dc4 = dc_image_mv(sras, 0, CH4_IDX)
thr = float(np.median(dc4))
monkeypatch.setattr(compute, "_MAX_WORKERS", 1)
serial = compute_rf_image(sras, 0, dc_threshold_mv=thr, apply_bg_sub=True,
row_avg_n=4)
monkeypatch.setattr(compute, "_MAX_WORKERS", 8)
parallel = compute_rf_image(sras, 0, dc_threshold_mv=thr, apply_bg_sub=True,
row_avg_n=4)
assert np.array_equal(serial, parallel), \
"row-averaged rf image identical regardless of chunking/worker count"
tests/test_stored_cache.py
+558
View File
@@ -0,0 +1,558 @@
"""Does a batch-computed FFT cache actually spare the viewer the FFT?
Storing a peak-frequency image per angle in the file is only worth doing if
displaying it is then free. The regression this module pins down is the
viewer's *dispatch* decision: it used to find the stored image only inside
ComputeWorker, so after a batch every angle change still queued a background
job behind a "Computing FFT…" popup for an image already on disk.
Both layers are covered — cached_rf_image's accept/reject rules, and the
window never reaching _start_compute for a batch-cached angle.
"""
import struct
from types import SimpleNamespace
from unittest.mock import patch
import numpy as np
import pytest
from PyQt6.QtCore import QEventLoop, QTimer
from PyQt6.QtWidgets import QApplication, QDialog
import sras_compute as compute
import sras_format as fmt
from sras_compute import cache_file, cached_rf_image, compute_rf_image, dc_image_mv
from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, SrasFile
from sras_viewer import SrasViewerWindow, VELOCITY_MODE_IDX
import tools.make_test_sras as gen
_THRESHOLD_MV = 50.0 # the viewer's own default
def pump(ms: int = 200):
loop = QEventLoop()
QTimer.singleShot(ms, loop.quit)
loop.exec()
def wait_until(pred, timeout_ms: int = 20000, step: int = 100) -> bool:
waited = 0
while waited < timeout_ms:
if pred():
return True
pump(step)
waited += step
return pred()
@pytest.fixture(scope="module")
def rig(tmp_path_factory):
"""A v6 file, the FFT images a from-scratch compute gives for it, and the
same file after Batch Compute FFT has written them into its v7 cache."""
path = tmp_path_factory.mktemp("stored_cache") / "cached.sras"
gen.write(path, n_angles=4, seed=7, samples_per_frame=256)
src = SrasFile(str(path))
fresh = {a: compute_rf_image(src, a, dc_threshold_mv=_THRESHOLD_MV,
apply_bg_sub=True)
for a in range(src.n_angles)}
assert src.background is not None, "the fixture file must have a background"
err = cache_file(str(path), "fft", True)
assert err == "", err
cached = SrasFile(str(path))
assert all(x is not None for x in cached.precomputed_freq_mhz)
assert all(x is None for x in cached.precomputed_dc4_mv), \
"FFT-only batch: the mask has to come from the viewer, not the file"
return SimpleNamespace(path=path, fresh=fresh, sras=cached,
n_angles=cached.n_angles)
@pytest.fixture(scope="module")
def dc_rig(tmp_path_factory):
"""A file that has been through Batch Compute DC and Store."""
path = tmp_path_factory.mktemp("stored_dc") / "dc_cached.sras"
gen.write(path, n_angles=4, seed=9, samples_per_frame=128)
err = cache_file(str(path), "dc", True)
assert err == "", err
sras = SrasFile(str(path))
assert all(x is not None for x in sras.precomputed_dc3_mv)
assert all(x is not None for x in sras.precomputed_dc4_mv)
assert len(np.unique(sras.precomputed_dc4_mv[0])) > 1, \
"a degenerate DC image would make the comparisons below vacuous"
return SimpleNamespace(path=path, sras=sras, n_angles=sras.n_angles)
@pytest.fixture
def no_fft(monkeypatch):
"""Make any real FFT work loud: returns a list that stays empty unless a
peak search actually runs."""
calls = []
for name in ("_peak_bins_direct", "_peak_bins_zoom"):
original = getattr(compute, name)
def spy(*args, _f=original, **kwargs):
calls.append(_f.__name__)
return _f(*args, **kwargs)
monkeypatch.setattr(compute, name, spy)
return calls
# ---------------------------------------------------------------------------
# cached_rf_image: when may the stored image stand in for a compute?
# ---------------------------------------------------------------------------
def test_stored_image_matches_a_fresh_compute(rig, no_fft):
for a in range(rig.n_angles):
dc4 = dc_image_mv(SrasFile(str(rig.path)), a, CH4_IDX)
img = cached_rf_image(rig.sras, a, dc_threshold_mv=_THRESHOLD_MV,
apply_bg_sub=True, dc4_mv=dc4)
assert img is not None, f"angle {a} is cached in the file"
assert np.allclose(img, rig.fresh[a], atol=1e-3), \
f"angle {a} differs from a from-scratch compute"
assert not no_fft, f"the stored image was used, no FFT ran: {no_fft}"
def test_unmasked_when_no_threshold(rig):
img = cached_rf_image(rig.sras, 0, dc_threshold_mv=None, apply_bg_sub=True)
assert img is not None and np.array_equal(img, rig.sras.precomputed_freq_mhz[0])
img[:] = -1.0
assert not np.any(rig.sras.precomputed_freq_mhz[0] == -1.0), \
"callers get a copy, never the file's own array"
def test_settings_the_stored_image_cannot_serve(rig):
"""A stored image carries one bg-sub state and one padding, so anything
else must fall through to a real compute rather than lie."""
spf = rig.sras.samples_per_frame
assert cached_rf_image(rig.sras, 0, _THRESHOLD_MV, apply_bg_sub=True,
n_fft=spf * 4) is None, \
"cache was written at pad 1, a pad-4 view resolves different peaks"
assert cached_rf_image(rig.sras, 0, _THRESHOLD_MV, apply_bg_sub=False) is None, \
"cache was written with bg-sub on"
uncached = SrasFile(str(rig.path))
uncached.precomputed_freq_mhz[1] = None
assert cached_rf_image(uncached, 1, _THRESHOLD_MV, apply_bg_sub=True) is None
@pytest.mark.parametrize("pad", [2, 10])
def test_cache_is_stored_at_the_configured_pad(tmp_path, pad):
"""Batching at a pad factor must produce a cache that view can read back —
a pad-1-only cache is one the padded viewer can never use."""
path = tmp_path / f"pad{pad}.sras"
gen.write(path, n_angles=2, seed=13, samples_per_frame=256)
spf = SrasFile(str(path)).samples_per_frame
assert cache_file(str(path), "fft", True, "scipy", 0, pad) == ""
sras = SrasFile(str(path))
assert sras.precomputed_pad_factor == pad, "pad factor survives the round trip"
assert cached_rf_image(sras, 0, None, apply_bg_sub=True,
n_fft=spf * pad) is not None, f"usable at pad {pad}"
assert cached_rf_image(sras, 0, None, apply_bg_sub=True) is None, \
"not usable unpadded"
assert cached_rf_image(sras, 0, None, apply_bg_sub=True,
n_fft=spf * (pad + 1)) is None, "not usable at another pad"
# The stored numbers must be the padded ones, not pad-1 relabelled.
fresh = SrasFile(str(path))
fresh.precomputed_freq_mhz = [None] * fresh.n_angles
for a in range(sras.n_angles):
assert np.allclose(
compute_rf_image(sras, a, dc_threshold_mv=None, apply_bg_sub=True,
n_fft=spf * pad),
compute_rf_image(fresh, a, dc_threshold_mv=None, apply_bg_sub=True,
n_fft=spf * pad), atol=1e-3), \
f"angle {a}: stored image is the pad-{pad} answer"
def test_cach_v1_reads_as_natural_resolution(tmp_path):
"""Files cached before the pad factor existed must keep working: a v1 tail
has no pad field and is pad 1 by construction."""
path = tmp_path / "v1.sras"
gen.write(path, n_angles=2, seed=14, samples_per_frame=256)
assert cache_file(str(path), "fft", True) == ""
# Rewrite the tail as a genuine CACH v1 block (old header, no pad field).
v2 = SrasFile(str(path))
freq, entries = v2.precomputed_freq_mhz, list(range(v2.n_angles))
payload = struct.pack(fmt.CACH_HDR_FMT, fmt.CACH_MAGIC, 1, fmt.CACH_FLAG_FFT)
payload += struct.pack(fmt.SFFT_HDR_FMT_V1, fmt.SFFT_MAGIC,
fmt.SFFT_FLAG_BG_SUB, len(entries))
for a in entries:
payload += struct.pack(">H", a) + freq[a].astype(">f4").tobytes()
head = path.read_bytes()[:v2._cache_tail_offset()]
path.write_bytes(head + payload)
v1 = SrasFile(str(path))
assert v1.precomputed_pad_factor == 1
assert v1.precomputed_bg_sub is True
assert all(np.array_equal(v1.precomputed_freq_mhz[a], freq[a]) for a in entries), \
"v1 images read back unchanged"
assert cached_rf_image(v1, 0, None, apply_bg_sub=True) is not None
assert cached_rf_image(v1, 0, None, apply_bg_sub=True,
n_fft=v1.samples_per_frame * 10) is None
def test_mask_read_can_be_refused(rig, no_fft):
"""With no DC4 in hand, building the mask means reading a whole channel —
the GUI thread asks for None instead."""
assert cached_rf_image(rig.sras, 0, _THRESHOLD_MV, apply_bg_sub=True,
allow_dc_recompute=False) is None
dc4 = dc_image_mv(SrasFile(str(rig.path)), 0, CH4_IDX)
img = cached_rf_image(rig.sras, 0, _THRESHOLD_MV, apply_bg_sub=True,
dc4_mv=dc4, allow_dc_recompute=False)
assert img is not None and np.allclose(img, rig.fresh[0], atol=1e-3)
assert not no_fft, f"no FFT on either branch: {no_fft}"
# ---------------------------------------------------------------------------
# The viewer: no compute job at all for a batch-cached angle
# ---------------------------------------------------------------------------
def test_viewer_shows_stored_angles_without_computing(rig, no_fft, monkeypatch):
app = QApplication.instance() or QApplication([]) # noqa: F841
win = SrasViewerWindow()
win.show()
dispatched = []
original_start = type(win)._start_compute
monkeypatch.setattr(type(win), "_start_compute",
lambda self: (dispatched.append(self.spin_angle.value()),
original_start(self))[1])
try:
win._load_file(str(rig.path))
assert wait_until(lambda: win._sras is not None), "file loaded"
# The file carries no DC block, so the mask comes from the window's own
# background precompute — the state a user is in by the time they click.
assert wait_until(lambda: all((a, CH4_IDX) in win._dc_cache
for a in range(rig.n_angles))), \
"DC precompute finished"
dispatched.clear()
win.combo_channel.setCurrentIndex(CH1_IDX)
assert wait_until(lambda: win._current_ch == CH1_IDX), "CH1 displayed"
for a in list(range(rig.n_angles)) + [1, 0]:
win.spin_angle.setValue(a)
win._on_view_changed()
pump(60)
assert win._current_angle == a, f"angle {a} displayed"
assert np.allclose(win._current_image, rig.fresh[a], atol=1e-3), \
f"angle {a} shows the stored image"
assert dispatched == [], \
f"stored angles need no compute job, dispatched for {dispatched}"
assert not no_fft, f"no FFT ran for any stored angle: {no_fft}"
# Velocity is still a post-multiply of the same stored image.
win.combo_channel.setCurrentIndex(VELOCITY_MODE_IDX)
pump(120)
assert np.allclose(win._current_image,
rig.fresh[0] * win.spin_grating_um.value(), atol=1e-3)
assert dispatched == [] and not no_fft
# ...but a setting the stored image cannot serve must still recompute,
# or the fast path would be showing the wrong picture.
win.chk_bg_sub.setChecked(False)
assert wait_until(lambda: not win._job_running("compute") and bool(no_fft)), \
"bg-sub off falls through to a real FFT"
finally:
win.close()
pump(300)
def test_batch_caches_at_the_viewers_pad_factor(tmp_path, no_fft, monkeypatch):
"""The bug a pad-10 user hits: Batch Compute FFT used to store pad-1
images regardless, so the padded view recomputed every angle forever.
"""
path = tmp_path / "padded_gui.sras"
gen.write(path, n_angles=3, seed=15, samples_per_frame=256)
app = QApplication.instance() or QApplication([]) # noqa: F841
win = SrasViewerWindow()
win.show()
dispatched = []
original_start = type(win)._start_compute
monkeypatch.setattr(type(win), "_start_compute",
lambda self: (dispatched.append(self.spin_angle.value()),
original_start(self))[1])
try:
win._fft_pad_factor = 10
win._load_file(str(path))
assert wait_until(lambda: win._sras is not None), "file loaded"
assert wait_until(lambda: all((a, CH4_IDX) in win._dc_cache
for a in range(win._sras.n_angles))), \
"DC precompute finished"
# Convert -> Batch Compute FFT, on the open file, through the real slot.
with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames",
return_value=([str(path)], "")):
win._on_batch_compute("fft")
assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran"
assert wait_until(lambda: win._sras is not None
and win._sras.version == 7), "file reloaded as v7"
pump(200)
assert win._sras.precomputed_pad_factor == 10, \
f"cached at the viewer's pad, got {win._sras.precomputed_pad_factor}"
expected = {a: compute.cached_rf_image(
win._sras, a, dc_threshold_mv=win.spin_threshold_mv.value(),
apply_bg_sub=win.chk_bg_sub.isChecked(),
n_fft=win._current_n_fft(),
dc4_mv=win._dc_cache.get((a, CH4_IDX)))
for a in range(win._sras.n_angles)}
assert all(v is not None for v in expected.values()), "cache is readable at pad 10"
no_fft.clear()
dispatched.clear()
win.combo_channel.setCurrentIndex(CH1_IDX)
assert wait_until(lambda: win._current_ch == CH1_IDX), "CH1 displayed"
for a in range(win._sras.n_angles):
win.spin_angle.setValue(a)
pump(60)
assert np.allclose(win._current_image, expected[a], atol=1e-3), \
f"angle {a} served from the pad-10 cache"
assert dispatched == [] and not no_fft, \
f"no recompute at pad 10 (jobs={dispatched}, fft={no_fft})"
assert "unusable" not in win.lbl_frame_warn.text()
# Change the pad and the cache legitimately stops applying — and the
# info panel has to say so rather than leave it a mystery.
win._fft_pad_factor = 4
win._update_scan_info_labels()
assert "Cached FFT unusable" in win.lbl_frame_warn.text(), \
win.lbl_frame_warn.text()
assert "pad 10x" in win.lbl_frame_warn.text()
win._refresh_display()
assert wait_until(lambda: not win._job_running("compute") and bool(no_fft)), \
"pad 4 recomputes rather than reusing the pad-10 cache"
finally:
win.close()
pump(300)
def test_viewer_shows_stored_dc_without_computing(dc_rig, monkeypatch):
"""Same for the DC half, with the background precompute silenced so the
file's stored block is the only thing that can be carrying the display."""
app = QApplication.instance() or QApplication([]) # noqa: F841
win = SrasViewerWindow()
win.show()
dispatched = []
original_start = type(win)._start_compute
monkeypatch.setattr(type(win), "_start_compute",
lambda self: (dispatched.append(self.spin_angle.value()),
original_start(self))[1])
monkeypatch.setattr(type(win), "_start_dc_precompute", lambda self: None)
try:
win._load_file(str(dc_rig.path))
assert wait_until(lambda: win._sras is not None), "file loaded"
pump(120)
for ch in (CH3_IDX, CH4_IDX):
win.combo_channel.setCurrentIndex(ch)
for a in range(dc_rig.n_angles):
win.spin_angle.setValue(a)
pump(60)
assert (win._current_angle, win._current_ch) == (a, ch), \
f"angle {a} on channel {ch} displayed"
assert np.array_equal(win._current_image,
dc_rig.sras.cached_dc_mv(a, ch)), \
f"angle {a} channel {ch} shows the file's stored DC image"
assert dispatched == [], \
f"stored DC angles need no compute job, dispatched for {dispatched}"
finally:
win.close()
pump(300)
# ---------------------------------------------------------------------------
# Row-averaged FFT: cache_file("fft_rowavg", ...) and its on-disk provenance
# ---------------------------------------------------------------------------
def test_row_average_flag_and_n_round_trip(tmp_path):
path = tmp_path / "rowavg_roundtrip.sras"
gen.write(path, n_angles=2, seed=20, samples_per_frame=128)
err = cache_file(str(path), "fft_rowavg", True, dc_threshold_mv=-1e9, row_avg_n=5)
assert err == "", err
sras = SrasFile(str(path))
assert sras.version == 7
assert sras.precomputed_row_avg_n == 5
assert all(x is not None for x in sras.precomputed_freq_mhz)
def test_fft_rowavg_mode_requires_positive_n_and_threshold(tmp_path):
path = tmp_path / "rowavg_bad_args.sras"
gen.write(path, n_angles=1, seed=27, samples_per_frame=64)
assert cache_file(str(path), "fft_rowavg", True, dc_threshold_mv=0.0, row_avg_n=0)
assert cache_file(str(path), "fft_rowavg", True, dc_threshold_mv=None, row_avg_n=5)
def test_raw_and_row_averaged_caches_never_cross_served(tmp_path):
"""The central regression this feature must never allow: a raw request
served a row-averaged image (or vice versa), or a request at one window
size served a cache stored at a different one."""
path = tmp_path / "cross_serve.sras"
gen.write(path, n_angles=1, seed=21, samples_per_frame=128)
err = cache_file(str(path), "fft_rowavg", True, dc_threshold_mv=-1e9, row_avg_n=5)
assert err == "", err
sras = SrasFile(str(path))
assert cached_rf_image(sras, 0, None, apply_bg_sub=True, row_avg_n=0) is None, \
"a raw request must not be served a row-averaged cache"
assert cached_rf_image(sras, 0, None, apply_bg_sub=True, row_avg_n=3) is None, \
"a request at the wrong window size must not be served either"
served = cached_rf_image(sras, 0, None, apply_bg_sub=True, row_avg_n=5)
assert served is not None
assert np.array_equal(served, sras.precomputed_freq_mhz[0])
def test_write_v7_cache_row_avg_n_carries_forward(tmp_path):
"""A later DC-only write must leave a previously-written row-averaged
FFT block -- including its row_avg_n -- byte-for-byte unchanged."""
path = tmp_path / "carry_forward.sras"
gen.write(path, n_angles=2, seed=22, samples_per_frame=64)
err = cache_file(str(path), "fft_rowavg", True, dc_threshold_mv=-1e9, row_avg_n=7)
assert err == "", err
before = SrasFile(str(path))
assert before.precomputed_row_avg_n == 7
freq_before = [x.copy() for x in before.precomputed_freq_mhz]
err = cache_file(str(path), "dc", True)
assert err == "", err
after = SrasFile(str(path))
assert after.precomputed_row_avg_n == 7, "row_avg_n survives a DC-only write"
assert all(np.array_equal(after.precomputed_freq_mhz[a], freq_before[a])
for a in range(after.n_angles)), \
"the row-averaged FFT block itself is untouched by a DC-only write"
def test_row_average_never_touches_dc_images(tmp_path):
path = tmp_path / "dc_untouched.sras"
gen.write(path, n_angles=2, seed=23, samples_per_frame=64)
src = SrasFile(str(path))
expect_dc3 = [dc_image_mv(src, a, CH3_IDX) for a in range(src.n_angles)]
expect_dc4 = [dc_image_mv(src, a, CH4_IDX) for a in range(src.n_angles)]
assert cache_file(str(path), "dc", True) == ""
assert cache_file(str(path), "fft_rowavg", True,
dc_threshold_mv=-1e9, row_avg_n=6) == ""
after = SrasFile(str(path))
assert all(np.allclose(after.precomputed_dc3_mv[a], expect_dc3[a], atol=1e-4)
for a in range(after.n_angles))
assert all(np.allclose(after.precomputed_dc4_mv[a], expect_dc4[a], atol=1e-4)
for a in range(after.n_angles))
def test_row_average_never_modifies_raw_waveform_data(tmp_path):
path = tmp_path / "waveform_untouched.sras"
gen.write(path, n_angles=2, seed=24, samples_per_frame=64)
orig = tmp_path / "waveform_untouched_orig.sras"
gen.write(orig, n_angles=2, seed=24, samples_per_frame=64)
assert cache_file(str(path), "fft_rowavg", True,
dc_threshold_mv=-1e9, row_avg_n=5) == ""
o, n = SrasFile(str(orig)), SrasFile(str(path))
assert all(np.array_equal(np.asarray(o.data[a]), np.asarray(n.data[a]))
for a in range(o.n_angles)), \
"waveform data untouched by a row-averaged cache write"
def test_cach_v1_backward_compat_defaults_row_avg_n_zero(tmp_path):
"""A v1 CACH tail predates row-averaged FFT caching entirely (no
row_avg_n byte at all) -- readers must still parse it in full, treating
it as row_avg_n=0. This is what protects an existing real-world v7
file's already-stored FFT cache from silently becoming unusable after
this change ships."""
path = tmp_path / "v1_rowavg.sras"
gen.write(path, n_angles=2, seed=25, samples_per_frame=128)
assert cache_file(str(path), "fft", True) == ""
v2 = SrasFile(str(path))
freq, entries = v2.precomputed_freq_mhz, list(range(v2.n_angles))
payload = struct.pack(fmt.CACH_HDR_FMT, fmt.CACH_MAGIC, 1, fmt.CACH_FLAG_FFT)
payload += struct.pack(fmt.SFFT_HDR_FMT_V1, fmt.SFFT_MAGIC,
fmt.SFFT_FLAG_BG_SUB, len(entries))
for a in entries:
payload += struct.pack(">H", a) + freq[a].astype(">f4").tobytes()
head = path.read_bytes()[:v2._cache_tail_offset()]
path.write_bytes(head + payload)
v1 = SrasFile(str(path))
assert v1.precomputed_row_avg_n == 0
assert all(np.array_equal(v1.precomputed_freq_mhz[a], freq[a]) for a in entries), \
"v1 images read back unchanged"
assert cached_rf_image(v1, 0, None, apply_bg_sub=True, row_avg_n=0) is not None
assert cached_rf_image(v1, 0, None, apply_bg_sub=True, row_avg_n=5) is None, \
"a v1 tail (predating this feature) can never satisfy a row-averaged request"
def test_viewer_batch_row_average_dispatch(tmp_path, monkeypatch):
"""Driving the new 'Batch Compute Row-Averaged FFT and Store' action
end-to-end through the real menu handler: dialog values reach the
worker, the worker reaches cache_file, and the written file is
self-describing afterward. Deliberately does not assert anything about
whether viewing an angle afterward dispatches a compute job -- that is
a separate, pre-existing gap in _refresh_display shared with the plain
DC/FFT batch actions (see test_viewer_shows_stored_angles_without_computing
/ test_viewer_shows_stored_dc_without_computing above), not something
row-averaging introduces or is responsible for fixing."""
path = tmp_path / "rowavg_gui.sras"
gen.write(path, n_angles=2, seed=26, samples_per_frame=128)
class _StubDialog:
def __init__(self, *a, **k):
pass
def exec(self):
return QDialog.DialogCode.Accepted
def get_half_width(self):
return 6
def get_threshold_mv(self):
return -1e9 # mask nothing, keep the comparison simple
monkeypatch.setattr("sras_viewer.main_window.RowAverageFftOptionsDialog", _StubDialog)
app = QApplication.instance() or QApplication([]) # noqa: F841
win = SrasViewerWindow()
win.show()
try:
win._load_file(str(path))
assert wait_until(lambda: win._sras is not None), "file loaded"
assert wait_until(lambda: all((a, CH4_IDX) in win._dc_cache
for a in range(win._sras.n_angles))), \
"DC precompute finished"
with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames",
return_value=([str(path)], "")):
win._on_batch_compute_row_avg()
assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran"
assert wait_until(lambda: win._sras is not None
and win._sras.version == 7), "file reloaded as v7"
pump(200)
assert win._sras.precomputed_row_avg_n == 6
assert "row-averaged n=6" in win.lbl_frame_warn.text(), win.lbl_frame_warn.text()
expected = compute.cached_rf_image(win._sras, 0, dc_threshold_mv=None,
apply_bg_sub=win.chk_bg_sub.isChecked(),
row_avg_n=6)
assert expected is not None, "the batch write left a readable row-averaged cache"
finally:
win.close()
pump(300)