Thomas Ales [M S E]
1caf6373cb
Drop the coarse+fine zoom refinement, the SciPy FFT backend, and the exact= audit path in favor of a single always-on full-transform peak search (_peak_bins). Block size is now derived from a per-thread memory budget (_fft_block_for/SRAS_FFT_PLAN_BUDGET_MB) instead of a fixed constant, so the existing block-parallel PyFFTW pool stays memory-safe at high pad factors without the zoom algorithm's bookkeeping. Also removes the now-unused threadpoolctl dependency and the FFT backend selector from the UI. Also includes a pre-existing min_freq_mhz peak-search floor (excludes bins below a caller-supplied frequency from the argmax) that was already implemented and tested in the working tree. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
298 lines
13 KiB
Python
298 lines
13 KiB
Python
"""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 CH1_IDX, 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 agree bit-for-bit with an independent
|
|
reference that row-averages the raw waveforms and background-subtracts
|
|
them, then runs a plain scipy rfft + argmax at the same pad factor --
|
|
i.e. row-averaging composes correctly with the padded peak search."""
|
|
import scipy.fft as scipy_fft
|
|
|
|
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
|
|
n_rows, n_frames = sras.image_shape(0)
|
|
n_fft = spf * 40
|
|
|
|
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 = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True,
|
|
row_avg_n=3, n_fft=n_fft)
|
|
|
|
assert avg_natural.shape == raw_natural.shape == avg_padded.shape
|
|
assert np.all(np.isfinite(avg_padded))
|
|
|
|
weights = compute._row_average_weights(3)
|
|
freq32 = sras.freq_axis_mhz(n_fft).astype(np.float32)
|
|
data = sras.data[0]
|
|
expected = np.zeros((n_rows, n_frames), dtype=np.float32)
|
|
for r in range(n_rows):
|
|
v = np.ones(n_frames, dtype=bool)
|
|
avg = compute._row_average_waveforms(
|
|
data[r, CH1_IDX].astype(np.float32), v, weights)
|
|
avg = avg - sras.background
|
|
S = scipy_fft.rfft(avg, n=n_fft, axis=-1, workers=1)
|
|
power = S.real ** 2 + S.imag ** 2
|
|
power[:, 0] = 0.0
|
|
expected[r] = freq32[np.argmax(power, axis=1)]
|
|
|
|
assert np.array_equal(avg_padded, expected), \
|
|
"row-averaged waveforms feed the padded peak search identically " \
|
|
"to an independent reference"
|
|
|
|
|
|
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(raw, spf)
|
|
avg_bins = compute._peak_bins(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_MAX", 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"
|