Thomas K Ales [MSE]
145 lines
5.1 KiB
Python
Executable File
145 lines
5.1 KiB
Python
Executable File
"""core.sras_analysis + the viewer's LoadedScan/compute path, run headlessly
|
|
over the golden fixtures."""
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from core.sras_analysis import (
|
|
CH1_IDX, CH4_IDX, ChannelCalibration, FALLBACK_YMULT_MV,
|
|
SawPipeline, compute_dc_image, compute_rf_image, power_spectrum,
|
|
)
|
|
|
|
GOLDEN = Path(__file__).parent / "golden"
|
|
|
|
|
|
def _loaded_scan(name="complete.sras"):
|
|
# Imported lazily: sras_viewer pulls in PyQt6/matplotlib.
|
|
from sras_viewer import LoadedScan
|
|
return LoadedScan(str(GOLDEN / name))
|
|
|
|
|
|
def test_loaded_scan_basics():
|
|
scan = _loaded_scan()
|
|
assert scan.rows_available == [3, 3]
|
|
# A legacy v6 fixture: its one background stands in for every angle.
|
|
assert all(scan.background(ai) is not None and len(scan.background(ai)) == 8
|
|
for ai in (0, 1))
|
|
assert len(scan.calib.ymult_mv) == 3
|
|
view = scan.angle_view(0)
|
|
assert view.shape == (3, 3, 4, 8)
|
|
scan.close()
|
|
|
|
|
|
def test_loaded_scan_truncated_rows():
|
|
scan = _loaded_scan("trunc_rowboundary_a1.sras")
|
|
assert scan.rows_available == [3, 2]
|
|
assert scan.angle_view(1).shape[0] == 2
|
|
scan = _loaded_scan("header_only.sras")
|
|
assert scan.rows_available == [0, 0]
|
|
assert scan.angle_view(0).shape[0] == 0
|
|
|
|
|
|
def test_compute_dc_image_matches_manual():
|
|
scan = _loaded_scan()
|
|
view = scan.angle_view(0)
|
|
img = compute_dc_image(view, CH4_IDX)
|
|
manual = view[:, CH4_IDX].astype(np.float64).mean(axis=-1)
|
|
assert img.shape == (3, 4)
|
|
assert img.dtype == np.float32
|
|
np.testing.assert_allclose(img, manual, rtol=1e-6)
|
|
|
|
|
|
def test_compute_rf_image_mask_and_values():
|
|
scan = _loaded_scan()
|
|
view = scan.angle_view(0)
|
|
sras = scan.sras
|
|
f_mhz = sras.freq_axis_mhz(sras.header.samples_per_frame)
|
|
|
|
# Threshold below everything: all pixels valid, values from the freq axis
|
|
img_all = compute_rf_image(view, scan.calib, f_mhz, dc_threshold_mv=-1e9)
|
|
assert img_all.shape == (3, 4)
|
|
assert set(np.unique(img_all)).issubset(set(f_mhz))
|
|
|
|
# Threshold above everything: fully masked, zero image, no FFT work
|
|
img_none = compute_rf_image(view, scan.calib, f_mhz, dc_threshold_mv=1e9)
|
|
assert not img_none.any()
|
|
|
|
|
|
def test_compute_rf_image_gate_zeroes_samples():
|
|
scan = _loaded_scan()
|
|
view = scan.angle_view(0)
|
|
sras = scan.sras
|
|
f_mhz = sras.freq_axis_mhz(sras.header.samples_per_frame)
|
|
t_ns = sras.time_axis_ns()
|
|
img = compute_rf_image(view, scan.calib, f_mhz, dc_threshold_mv=-1e9,
|
|
gate_start_ns=float(t_ns[2]), gate_end_ns=float(t_ns[5]),
|
|
time_axis_ns=t_ns)
|
|
assert img.shape == (3, 4)
|
|
|
|
|
|
def test_calibration_roundtrip_and_fallback():
|
|
cal = ChannelCalibration.from_preambles(["", "", ""])
|
|
assert cal.ymult_mv == [FALLBACK_YMULT_MV] * 3
|
|
assert cal.mv_to_adc(cal.adc_to_mv(42.0, 0), 0) == pytest.approx(42.0)
|
|
|
|
cal2 = ChannelCalibration.from_preambles(
|
|
["YMULT 1.0E-3;YOFF -10.0;YZERO 2.0E-3"])
|
|
assert cal2.ymult_mv[0] == pytest.approx(1.0)
|
|
assert cal2.yoff_adc[0] == pytest.approx(-10.0)
|
|
assert cal2.yzero_mv[0] == pytest.approx(2.0)
|
|
|
|
|
|
def test_power_spectrum_dc_suppressed():
|
|
x = np.ones(64, dtype=np.float32) * 5.0
|
|
p = power_spectrum(x)
|
|
assert p[0] == 0.0
|
|
assert not p[1:].any()
|
|
|
|
|
|
def test_saw_pipeline_finds_injected_packet():
|
|
sr = 6.25e9
|
|
n = 1250
|
|
t = np.arange(n) / sr
|
|
rng = np.random.default_rng(42)
|
|
|
|
def make_shot(delay_ns=200.0):
|
|
sig = rng.normal(0, 0.05, n).astype(np.float32)
|
|
packet = np.exp(-((t * 1e9 - delay_ns) / 25.0) ** 2) \
|
|
* np.sin(2 * np.pi * 140e6 * t)
|
|
return sig + 3.0 * packet.astype(np.float32)
|
|
|
|
pipe = SawPipeline(sr, emi_gate_ns=50.0, bp_lo_mhz=85.0, bp_hi_mhz=200.0,
|
|
saw_window_ns=(80.0, 350.0))
|
|
pipe.build_template(np.stack([make_shot() for _ in range(10)]))
|
|
assert pipe.template is not None
|
|
|
|
metrics = pipe.process_shot_metrics(make_shot())
|
|
assert metrics["peak_time_ns"] == pytest.approx(200.0, abs=15.0)
|
|
assert metrics["snr"] > 3.0
|
|
|
|
# process_shot returns the same metrics plus the stage arrays
|
|
full = pipe.process_shot(make_shot())
|
|
for key in ("raw", "gated", "filtered", "mf_output", "envelope"):
|
|
assert isinstance(full[key], np.ndarray)
|
|
|
|
|
|
def test_compute_saw_image_scalars_only():
|
|
# Synthetic angle view: golden frames (8 samples) are too short for the
|
|
# 6th-order zero-phase bandpass; real records are >1000 samples.
|
|
from core.sras_analysis import compute_saw_image
|
|
rng = np.random.default_rng(1)
|
|
view = rng.integers(-40, 40, size=(3, 3, 4, 512), dtype=np.int8)
|
|
view[:, CH4_IDX] = 100 # every pixel passes the DC mask
|
|
calib = ChannelCalibration.from_preambles(["", "", ""])
|
|
|
|
pipe = SawPipeline(6.25e9, saw_window_ns=(10.0, 70.0))
|
|
pipe.build_template(view[0, CH1_IDX, :4].astype(np.float32))
|
|
img = compute_saw_image(view, calib, -1e9, pipe, "amplitude")
|
|
assert img.shape == (3, 4)
|
|
assert img.dtype == np.float32
|
|
assert (img > 0).all()
|
|
|
|
tof = compute_saw_image(view, calib, -1e9, pipe, "tof")
|
|
assert tof.shape == (3, 4)
|