Thomas Ales [M S E]
6976cbf767
Add Export -> Batch Export Images..., which renders and saves one PNG per angle for whichever of CH1 (RF), CH3 (DC-A), CH4 (DC-B), and Velocity the user selects, for the currently open file. Each channel gets its own fixed min/max colorbar range, entered once in the dialog and held constant across every exported angle, so the resulting images are directly comparable to each other instead of each auto-scaling to its own data (today's live-view default). BatchExportDialog (sras_viewer.py) collects the output folder, file prefix, and per-channel checkbox + range, seeded from whatever's already cached (session DC/FFT caches, or the file's own v7 cache blocks) so opening the dialog never triggers a fresh compute. Export runs on a background thread via a new BatchExportWorker (sras_workers.py), which computes each angle's channel image with the existing dc_image_mv/compute_rf_image helpers (reusing one FFT per angle for both RF and Velocity) and renders it with a headless matplotlib Agg canvas. Also includes several small correctness/robustness fixes that were already staged in the working tree: an int16-vs-float32 accumulator mismatch in sras_average.py's row averaging, a DC4-mask cache reuse and atomic sidecar write in sras_compute.py, a dc3/dc4 pairing guard in sras_format.py's v7 cache writer, and matching updates to the tools/ test fixtures. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
366 lines
16 KiB
Python
366 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
"""Behavioural tests for the sras-viewer refactor.
|
|
|
|
Covers what the golden-hash harness can't: the v6->v7 cache round-trip
|
|
(including block carry-forward), parallel-vs-serial identity, the no-mask
|
|
fast path, and the ROI bounding-box mask optimisation.
|
|
|
|
Usage: python tools/test_refactor.py [--scratch DIR]
|
|
"""
|
|
|
|
import argparse
|
|
import shutil
|
|
import sys
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
|
|
|
import sras_compute as compute # noqa: E402
|
|
from sras_compute import ( # noqa: E402
|
|
cache_file, compute_dc_image, compute_rf_image, dc_image_mv,
|
|
)
|
|
from sras_format import CH3_IDX, CH4_IDX, SrasFile, adc_to_mv # noqa: E402
|
|
import tools.make_test_sras as gen # noqa: E402
|
|
|
|
_failures: list[str] = []
|
|
|
|
|
|
def check(name: str, ok: bool, detail: str = ""):
|
|
print(f" {'PASS' if ok else 'FAIL'} {name}" + (f" — {detail}" if detail else ""))
|
|
if not ok:
|
|
_failures.append(name)
|
|
|
|
|
|
def test_cache_roundtrip(scratch: Path):
|
|
"""v6 -> v7 for DC, then FFT, asserting the first block survives the
|
|
second write (the carry-forward path in write_v7_cache)."""
|
|
print("\ncache round-trip (v6 -> v7, both blocks)")
|
|
path = scratch / "roundtrip.sras"
|
|
gen.write(path, n_angles=3, seed=1, samples_per_frame=64)
|
|
|
|
src = SrasFile(str(path))
|
|
check("source is v6", src.version == 6, f"got v{src.version}")
|
|
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)]
|
|
expect_fft = [compute_rf_image(src, a, dc_threshold_mv=None, apply_bg_sub=True)
|
|
for a in range(src.n_angles)]
|
|
|
|
err = cache_file(str(path), "dc", True)
|
|
check("dc cache_file succeeded", err == "", err)
|
|
|
|
after_dc = SrasFile(str(path))
|
|
check("version flipped to 7", after_dc.version == 7, f"got v{after_dc.version}")
|
|
check("dc3 stored for every angle",
|
|
all(x is not None for x in after_dc.precomputed_dc3_mv))
|
|
check("dc3 values round-trip",
|
|
all(np.allclose(after_dc.precomputed_dc3_mv[a], expect_dc3[a], atol=1e-4)
|
|
for a in range(after_dc.n_angles)))
|
|
check("dc4 values round-trip",
|
|
all(np.allclose(after_dc.precomputed_dc4_mv[a], expect_dc4[a], atol=1e-4)
|
|
for a in range(after_dc.n_angles)))
|
|
check("no fft block yet",
|
|
all(x is None for x in after_dc.precomputed_freq_mhz))
|
|
check("cached images are native float32",
|
|
after_dc.precomputed_dc3_mv[0].dtype == np.float32
|
|
and after_dc.precomputed_dc3_mv[0].dtype.byteorder in ("=", "|"),
|
|
str(after_dc.precomputed_dc3_mv[0].dtype.byteorder))
|
|
check("cached images are writable",
|
|
after_dc.precomputed_dc3_mv[0].flags.writeable)
|
|
|
|
err = cache_file(str(path), "fft", True)
|
|
check("fft cache_file succeeded", err == "", err)
|
|
|
|
both = SrasFile(str(path))
|
|
check("fft stored for every angle",
|
|
all(x is not None for x in both.precomputed_freq_mhz))
|
|
check("fft values round-trip",
|
|
all(np.allclose(both.precomputed_freq_mhz[a], expect_fft[a], atol=1e-3)
|
|
for a in range(both.n_angles)))
|
|
check("DC block carried forward through the FFT write",
|
|
all(np.allclose(both.precomputed_dc3_mv[a], expect_dc3[a], atol=1e-4)
|
|
for a in range(both.n_angles)))
|
|
check("bg_sub flag persisted", both.precomputed_bg_sub is True)
|
|
|
|
# The fast path must reproduce a fresh compute, and masking must still
|
|
# apply on top of a cached (unmasked) image.
|
|
fresh = SrasFile(str(path))
|
|
fresh.precomputed_freq_mhz = [None] * fresh.n_angles
|
|
dc4 = dc_image_mv(both, 0, CH4_IDX)
|
|
thr = float(np.median(dc4))
|
|
check("cached fast path == fresh compute (unmasked)",
|
|
np.allclose(compute_rf_image(both, 0, dc_threshold_mv=None, apply_bg_sub=True),
|
|
compute_rf_image(fresh, 0, dc_threshold_mv=None, apply_bg_sub=True),
|
|
atol=1e-3))
|
|
check("cached fast path == fresh compute (masked)",
|
|
np.allclose(compute_rf_image(both, 0, dc_threshold_mv=thr, apply_bg_sub=True),
|
|
compute_rf_image(fresh, 0, dc_threshold_mv=thr, apply_bg_sub=True),
|
|
atol=1e-3))
|
|
|
|
# Waveform data must be byte-identical to the pre-cache file.
|
|
orig = scratch / "roundtrip_orig.sras"
|
|
gen.write(orig, n_angles=3, seed=1, samples_per_frame=64)
|
|
o, n = SrasFile(str(orig)), SrasFile(str(path))
|
|
check("waveform data untouched by the cache write",
|
|
all(np.array_equal(np.asarray(o.data[a]), np.asarray(n.data[a]))
|
|
for a in range(o.n_angles)))
|
|
|
|
|
|
def test_partial_v7_cache(scratch: Path):
|
|
"""Only some angles cached: uncached angles must compute, not read zeros.
|
|
This is the v5 bug the ragged normalisation fixed, checked via v7."""
|
|
print("\npartial cache (only some angles stored)")
|
|
path = scratch / "partial.sras"
|
|
gen.write(path, n_angles=3, seed=2, samples_per_frame=64)
|
|
|
|
src = SrasFile(str(path))
|
|
expected = [compute_rf_image(src, a, dc_threshold_mv=None, apply_bg_sub=True)
|
|
for a in range(src.n_angles)]
|
|
partial = [expected[0], None, expected[2]] # angle 1 deliberately absent
|
|
src.write_v7_cache(new_freq_mhz=partial, new_bg_sub=True)
|
|
|
|
reread = SrasFile(str(path))
|
|
check("angle 1 is not cached", reread.precomputed_freq_mhz[1] is None)
|
|
check("angles 0 and 2 are cached",
|
|
reread.precomputed_freq_mhz[0] is not None
|
|
and reread.precomputed_freq_mhz[2] is not None)
|
|
img1 = compute_rf_image(reread, 1, dc_threshold_mv=None, apply_bg_sub=True)
|
|
check("uncached angle computes rather than returning zeros",
|
|
np.any(img1 != 0) and np.allclose(img1, expected[1], atol=1e-3))
|
|
|
|
|
|
def test_parallel_identity(scratch: Path):
|
|
"""Forcing 1 worker vs many must give identical output — catches
|
|
chunk-boundary and race bugs."""
|
|
print("\nparallel vs serial identity")
|
|
path = scratch / "parallel.sras"
|
|
# Many rows, so the row loop actually splits into several chunks.
|
|
n_rows, n_frames, spf = 48, 9, 256
|
|
gen.write(path, n_angles=1, seed=3, samples_per_frame=spf,
|
|
geometry=[(n_rows, n_frames)])
|
|
sras = SrasFile(str(path))
|
|
|
|
saved_budget, saved_workers = compute._TOTAL_BYTES_BUDGET, compute._MAX_WORKERS
|
|
try:
|
|
# Shrink the budget so chunk_rows collapses to 1 and every row is
|
|
# its own chunk — the worst case for boundary bugs.
|
|
compute._TOTAL_BYTES_BUDGET = 8 * n_frames * spf * 4
|
|
|
|
compute._MAX_WORKERS = 1
|
|
dc_serial = compute_dc_image(sras, 0, CH4_IDX)
|
|
rf_serial = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True)
|
|
dc4 = adc_to_mv(dc_serial, *sras.cal(CH4_IDX))
|
|
thr = float(np.median(dc4))
|
|
rf_masked_serial = compute_rf_image(sras, 0, dc_threshold_mv=thr,
|
|
apply_bg_sub=True)
|
|
|
|
chunk_rows, n_workers = compute._plan_chunks(
|
|
n_rows, n_frames, spf, live_multiplier=compute._FFT_LIVE_MULTIPLIER)
|
|
check("serial plan uses 1 worker", n_workers == 1, f"chunk_rows={chunk_rows}")
|
|
check("work actually splits into multiple chunks", chunk_rows < n_rows,
|
|
f"chunk_rows={chunk_rows} of {n_rows} rows")
|
|
|
|
compute._MAX_WORKERS = 8
|
|
chunk_rows, n_workers = compute._plan_chunks(
|
|
n_rows, n_frames, spf, live_multiplier=compute._FFT_LIVE_MULTIPLIER)
|
|
check("parallel plan uses >1 worker", n_workers > 1,
|
|
f"chunk_rows={chunk_rows} workers={n_workers}")
|
|
|
|
dc_par = compute_dc_image(sras, 0, CH4_IDX)
|
|
rf_par = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True)
|
|
rf_masked_par = compute_rf_image(sras, 0, dc_threshold_mv=thr, apply_bg_sub=True)
|
|
|
|
check("dc image identical", np.array_equal(dc_serial, dc_par))
|
|
check("rf image identical (unmasked)", np.array_equal(rf_serial, rf_par))
|
|
check("rf image identical (masked)",
|
|
np.array_equal(rf_masked_serial, rf_masked_par))
|
|
finally:
|
|
compute._TOTAL_BYTES_BUDGET, compute._MAX_WORKERS = saved_budget, saved_workers
|
|
|
|
|
|
def test_nomask_equals_low_threshold(scratch: Path):
|
|
"""dc_threshold_mv=None must equal a threshold below every pixel, while
|
|
skipping the CH4 read."""
|
|
print("\nno-mask path")
|
|
path = scratch / "nomask.sras"
|
|
gen.write(path, n_angles=2, seed=4, samples_per_frame=128)
|
|
sras = SrasFile(str(path))
|
|
for a in range(sras.n_angles):
|
|
none_img = compute_rf_image(sras, a, dc_threshold_mv=None, apply_bg_sub=True)
|
|
low_img = compute_rf_image(sras, a, dc_threshold_mv=-1e9, apply_bg_sub=True)
|
|
check(f"angle {a}: None == -1e9 threshold",
|
|
np.array_equal(none_img, low_img))
|
|
check(f"angle {a}: image is non-degenerate",
|
|
len(np.unique(none_img)) > 1, f"{len(np.unique(none_img))} unique")
|
|
|
|
|
|
def test_roi_mask():
|
|
"""The bbox-restricted mask must equal a full-grid point-in-polygon test."""
|
|
print("\nROI mask (bbox fast path vs full grid)")
|
|
from matplotlib.path import Path as MplPath
|
|
from sras_viewer import RoiQuad
|
|
|
|
rng = np.random.default_rng(0)
|
|
x = np.linspace(-2.0, 3.0, 137)
|
|
y = np.linspace(1.0, 4.0, 91)
|
|
|
|
cases = {
|
|
"axis-aligned rect": np.array([[0.0, 1.5], [1.0, 1.5], [1.0, 3.0], [0.0, 3.0]]),
|
|
"skewed quad": np.array([[-0.5, 1.2], [1.7, 1.9], [1.2, 3.4], [-1.0, 2.6]]),
|
|
"entirely outside": np.array([[8.0, 8.0], [9.0, 8.0], [9.0, 9.0], [8.0, 9.0]]),
|
|
"covers whole grid": np.array([[-9.0, -9.0], [9.0, -9.0], [9.0, 9.0], [-9.0, 9.0]]),
|
|
"straddles left edge": np.array([[-4.0, 2.0], [0.5, 2.0], [0.5, 3.0], [-4.0, 3.0]]),
|
|
}
|
|
for _ in range(5):
|
|
cases[f"random {_}"] = rng.uniform([-2.5, 0.5], [3.5, 4.5], size=(4, 2))
|
|
|
|
for name, pts in cases.items():
|
|
roi = RoiQuad(pts)
|
|
fast = roi.mask_for_grid(x, y)
|
|
X, Y = np.meshgrid(x.astype(np.float64), y.astype(np.float64))
|
|
slow = MplPath(pts).contains_points(
|
|
np.column_stack([X.ravel(), Y.ravel()])).reshape(X.shape)
|
|
check(f"{name} ({int(slow.sum())} px inside)", np.array_equal(fast, slow))
|
|
|
|
# Descending y axis (images are stored top-down in some scans).
|
|
roi = RoiQuad(cases["skewed quad"])
|
|
y_desc = y[::-1]
|
|
fast = roi.mask_for_grid(x, y_desc)
|
|
X, Y = np.meshgrid(x.astype(np.float64), y_desc.astype(np.float64))
|
|
slow = MplPath(cases["skewed quad"]).contains_points(
|
|
np.column_stack([X.ravel(), Y.ravel()])).reshape(X.shape)
|
|
check("descending y axis", np.array_equal(fast, slow))
|
|
|
|
|
|
def test_legacy_parse_and_average(scratch: Path):
|
|
"""v2-v4 parsing plus the sras_average.py rewrite (which now streams via
|
|
SrasFile rather than slurping the whole file)."""
|
|
import subprocess
|
|
print("\nlegacy formats (v2-v4) and sras_average")
|
|
repo = Path(__file__).resolve().parent.parent
|
|
|
|
for version in (2, 3, 4):
|
|
path = scratch / f"legacy_v{version}.sras"
|
|
meta = gen.write_legacy(path, version=version, n_angles=2, n_rows=4,
|
|
n_frames=12, samples_per_frame=32, seed=version)
|
|
s = SrasFile(str(path))
|
|
check(f"v{version} parses", s.version == version, f"got v{s.version}")
|
|
check(f"v{version} geometry uniform across angles",
|
|
list(s.n_rows) == [4, 4] and list(s.n_frames) == [12, 12],
|
|
f"rows={list(s.n_rows)} frames={list(s.n_frames)}")
|
|
check(f"v{version} waveform data matches what was written",
|
|
all(np.array_equal(np.asarray(s.data[a]), meta["data"][a])
|
|
for a in range(s.n_angles)))
|
|
check(f"v{version} background {'present' if version >= 4 else 'absent'}",
|
|
(s.background is not None) == (version >= 4))
|
|
check(f"v{version} precomputed stores are ragged lists",
|
|
isinstance(s.precomputed_freq_mhz, list)
|
|
and len(s.precomputed_freq_mhz) == s.n_angles)
|
|
# DC image must equal a direct mean of the known input.
|
|
expect = meta["data"][0][:, CH3_IDX, :, :].astype(np.float64).mean(axis=-1)
|
|
check(f"v{version} DC image equals a direct mean",
|
|
np.allclose(compute_dc_image(s, 0, CH3_IDX), expect, atol=1e-3))
|
|
|
|
src = scratch / "legacy_v4.sras"
|
|
meta = gen.write_legacy(src, version=4, n_angles=2, n_rows=4, n_frames=12,
|
|
samples_per_frame=32, seed=4)
|
|
dst = scratch / "legacy_v4_avg.sras"
|
|
if dst.exists():
|
|
dst.unlink()
|
|
proc = subprocess.run(
|
|
[sys.executable, str(repo / "sras_average.py"), str(src), str(dst), "--n", "4"],
|
|
capture_output=True, text=True, cwd=repo)
|
|
check("sras_average ran", proc.returncode == 0,
|
|
(proc.stderr or proc.stdout).strip()[-200:])
|
|
|
|
if dst.exists():
|
|
avg = SrasFile(str(dst))
|
|
check("averaged file parses", avg.version == 4)
|
|
check("frame count divided by 4",
|
|
list(avg.n_frames) == [3, 3], f"{list(avg.n_frames)}")
|
|
check("angles/rows/channels unchanged",
|
|
avg.n_angles == 2 and list(avg.n_rows) == [4, 4]
|
|
and avg.n_channels == meta["n_channels"])
|
|
check("calibration preserved",
|
|
np.allclose(avg.ch_ymult_mv, SrasFile(str(src)).ch_ymult_mv))
|
|
check("background preserved",
|
|
np.array_equal(avg.background, SrasFile(str(src)).background))
|
|
src_data = meta["data"]
|
|
# int16 (not float32) before .mean(): matches average_rows' own
|
|
# float64-accumulator behavior for integer input, so this doesn't
|
|
# drift from what average_rows actually guarantees.
|
|
expect0 = src_data[0][:, :, 0:4, :].astype(np.int16).mean(axis=2).astype(np.int16)
|
|
check("first averaged group equals the mean of its 4 source frames",
|
|
np.array_equal(np.asarray(avg.data[0])[:, :, 0, :], expect0))
|
|
|
|
# Remainder handling: 12 frames / 5 -> 2 full groups + 1 partial.
|
|
dst2 = scratch / "legacy_v4_avg5.sras"
|
|
proc2 = subprocess.run([sys.executable, str(repo / "sras_average.py"),
|
|
str(src), str(dst2), "--n", "5"],
|
|
capture_output=True, text=True, cwd=repo)
|
|
check("sras_average ran (--n 5)", proc2.returncode == 0,
|
|
(proc2.stderr or proc2.stdout).strip()[-200:])
|
|
if dst2.exists():
|
|
check("partial trailing group kept by default",
|
|
list(SrasFile(str(dst2)).n_frames) == [3, 3],
|
|
f"{list(SrasFile(str(dst2)).n_frames)}")
|
|
dst3 = scratch / "legacy_v4_avg5d.sras"
|
|
proc3 = subprocess.run([sys.executable, str(repo / "sras_average.py"),
|
|
str(src), str(dst3), "--n", "5", "--discard-remainder"],
|
|
capture_output=True, text=True, cwd=repo)
|
|
check("sras_average ran (--n 5 --discard-remainder)", proc3.returncode == 0,
|
|
(proc3.stderr or proc3.stdout).strip()[-200:])
|
|
if dst3.exists():
|
|
check("--discard-remainder drops the partial group",
|
|
list(SrasFile(str(dst3)).n_frames) == [2, 2],
|
|
f"{list(SrasFile(str(dst3)).n_frames)}")
|
|
|
|
|
|
def test_unsupported_version_reported(scratch: Path):
|
|
"""cache_file must report, not raise, for a file it can't handle."""
|
|
print("\nerror reporting")
|
|
bogus = scratch / "bogus.sras"
|
|
bogus.write_bytes(b"SRAS" + bytes([99]) + b"\x00" * 200)
|
|
err = cache_file(str(bogus), "dc", True)
|
|
check("bad version returns an error string", bool(err), err)
|
|
missing = cache_file(str(scratch / "does_not_exist.sras"), "dc", True)
|
|
check("missing file returns an error string", bool(missing), missing)
|
|
|
|
|
|
def main():
|
|
p = argparse.ArgumentParser(description=__doc__)
|
|
p.add_argument("--scratch")
|
|
args = p.parse_args()
|
|
|
|
tmp = None
|
|
if args.scratch:
|
|
scratch = Path(args.scratch)
|
|
scratch.mkdir(parents=True, exist_ok=True)
|
|
else:
|
|
tmp = tempfile.mkdtemp(prefix="sras_test_")
|
|
scratch = Path(tmp)
|
|
|
|
try:
|
|
test_cache_roundtrip(scratch)
|
|
test_partial_v7_cache(scratch)
|
|
test_parallel_identity(scratch)
|
|
test_nomask_equals_low_threshold(scratch)
|
|
test_roi_mask()
|
|
test_legacy_parse_and_average(scratch)
|
|
test_unsupported_version_reported(scratch)
|
|
finally:
|
|
if tmp:
|
|
shutil.rmtree(tmp, ignore_errors=True)
|
|
|
|
print()
|
|
if _failures:
|
|
print(f"{len(_failures)} FAILURE(S): " + ", ".join(_failures))
|
|
sys.exit(1)
|
|
print("All checks passed.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|