"""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"