Merge pull request 'Refactor/parallel and dedup' (#1) from refactor/parallel-and-dedup into main

Reviewed-on: #1
This commit was merged in pull request #1.
This commit is contained in:
2026-07-31 15:20:33 -05:00
11 changed files with 5123 additions and 2161 deletions
+117 -2
View File
@@ -1,4 +1,4 @@
# SRAS Scan Binary Format — Version 6
# SRAS Scan Binary Format — Version 6 / 7
Each `.sras` file contains **one complete scan**: all GR rotation angles and all
Y rows. Files are named `{prefix}.sras`.
@@ -10,6 +10,14 @@ rotated by that specific angle** — not the worst case across all angles — so
fewer rows than a 45° scan of the same ROI, and the file format reflects that
instead of forcing every angle to the largest bounding box.
v7 is byte-identical to v6 (same header, angle table, geometry table, row
table, preamble blocks, background block, waveform data) plus an optional
trailing **Cache Tail** holding precomputed per-angle DC and/or FFT images
(see [Cache Tail (v7)](#cache-tail-v7) below) — only the header's `version`
field and the presence of that trailing section differ. Scans come off the
scope as v6; `sras_viewer.py`'s "Convert" menu batch actions convert a file
to v7 **in place** the first time either cache block is computed and stored.
---
## File Layout
@@ -22,6 +30,7 @@ instead of forcing every angle to the largest bounding box.
[Preamble Blocks — n_channels × (uint16 length + UTF-8 WFMOutpre string)]
[Background Block — uint32 n_bg_samples + n_bg_samples × int8 bytes]
[Waveform Data (ragged) — per angle: n_rows[a] × n_channels × n_frames[a] × samples_per_frame × bps bytes]
[Cache Tail (optional) — "CACH" + DC block (optional) + FFT block (optional); v7 only]
```
All multi-byte integers and floats use **big-endian** byte order
@@ -182,6 +191,111 @@ using that angle's `x_start` from the Per-Angle Geometry Table (not
---
## Cache Tail (v7)
Present iff `version == 7` and `file_size > cache_offset`, where:
```
cache_offset = data_offset + Σ over angles a of:
n_rows[a] × n_channels × n_frames[a] × samples_per_frame × bytes_per_sample
```
i.e. exactly `data_offset + waveform_bytes` — the same "Total data size"
formula as Waveform Data above. This offset is derivable from the header and
Per-Angle Geometry Table alone and does **not** depend on which cache
block(s) are present, so a writer can always seek straight there without
reading or touching any waveform byte before it.
Unlike the (removed) v5 `PREC` section, which stored one omnibus per-angle
entry (FFT + both DC channels together) in a dense, uniform-geometry array,
the v7 Cache Tail splits DC and FFT into two **independent** sub-blocks —
each sized per-angle from the Per-Angle Geometry Table, each independently
present, and each independently updatable in any order, any number of
times, without disturbing the other. This matches `sras_viewer.py`'s
"Convert" menu, which exposes DC and FFT store as two separate batch
actions.
### CACH outer header (6 bytes, `">4sBB"`)
| Offset | Size | Type | Field | Description |
|--------|------|------|-------|-------------|
| 0 | 4 | `char[4]` | `cach_magic` | `CACH` (ASCII). Missing/wrong magic → treat file as having no cache. |
| 4 | 1 | `u8` | `cach_version` | Cache format version. Currently `1`. Readers must treat the file as uncached if this is not a version they understand (unlike v5's `PREC` section, which read but never validated its version byte). |
| 5 | 1 | `u8` | `block_flags` | Bit 0 = DC block (`SDCB`) follows. Bit 1 = FFT block (`SFFT`) follows, immediately after the DC block if both are present. Bits 27 reserved, must be zero on write. |
### DC block `SDCB` (present iff `block_flags & 0x01`)
7-byte block header, format `">4sBH"`:
| Offset (rel) | Size | Type | Field | Description |
|--------------|------|------|-------|-------------|
| 0 | 4 | `char[4]` | `magic` | `SDCB` |
| 4 | 1 | `u8` | `reserved` | `0`, reserved for future use |
| 5 | 2 | `u16` | `n_stored` | Number of angle entries that follow, `0 ≤ n_stored ≤ n_angles` |
followed by `n_stored` entries, each:
```
u16 angle_idx — index into the angle table (0-based)
f32[n_rows[angle_idx] × n_frames[angle_idx]] dc3_mv — CH3 waveform mean, mV, row-major
f32[n_rows[angle_idx] × n_frames[angle_idx]] dc4_mv — CH4 waveform mean, mV, row-major
```
Entries may appear in any order and need not be contiguous from angle 0 —
this supports storing (or re-storing) a subset of angles, or an
interrupted batch run leaving only some angles cached. Readers bounds-check
`angle_idx < n_angles` on each entry and stop parsing on an out-of-range
value, same as v5's `PREC` section.
### FFT block `SFFT` (present iff `block_flags & 0x02`)
7-byte block header, format `">4sBH"`:
| Offset (rel) | Size | Type | Field | Description |
|--------------|------|------|-------|-------------|
| 0 | 4 | `char[4]` | `magic` | `SFFT` |
| 4 | 1 | `u8` | `flags` | Bit 0 = `bg_sub_applied` — background waveform was subtracted from CH1 before the FFT when these images were computed. Bits 17 reserved. |
| 5 | 2 | `u16` | `n_stored` | Number of angle entries that follow |
followed by `n_stored` entries, each:
```
u16 angle_idx — index into the angle table (0-based)
f32[n_rows[angle_idx] × n_frames[angle_idx]] peak_freq_mhz — CH1 FFT peak frequency, MHz, row-major
```
**`peak_freq_mhz`** is computed without any DC-threshold masking (i.e. the
FFT is run on every pixel unconditionally, same as v5's `PREC` convention).
Readers apply the DC4 threshold at display time:
```
pixel is valid ⟺ dc4_mv[r][f] ≥ threshold_mv
display_value = peak_freq_mhz[r][f] if valid, else 0
```
using the DC4 image from the DC block if that angle is also cached there,
else computed on demand.
Readers must fall back to real-time FFT computation (ignoring stored
`peak_freq_mhz`) under the same conditions as v5's PREC fast path: time-domain
gating is active, zero-padding (`n_fft ≠ samples_per_frame`) is requested, or
the reader's background-subtraction setting doesn't match `flags.bg_sub_applied`.
### In-place write ordering
A writer updating a file's Cache Tail must write the payload (CACH header +
whichever block(s) are present) **before** flipping the header's `version`
byte to `7`, and `truncate()` the file to the new payload's end immediately
after writing it. If the process is interrupted between the payload write
and the version-byte flip, the file is still valid v6 — v6 parsing only
bounds-checks each angle's `offset + nbytes ≤ file_size`, it never asserts
exactly how many bytes follow the last angle's waveform block — so the
interrupted write leaves harmless trailing bytes rather than a corrupt file,
and the next successful write overwrites them via the same deterministic
`cache_offset`.
---
## Acquisition Settings (fixed by sc3_aui_app.py)
| Parameter | Value |
@@ -206,5 +320,6 @@ using that angle's `x_start` from the Per-Angle Geometry Table (not
| 3 | Added preamble blocks (WFMOutpre strings) after the row table, one length-prefixed UTF-8 block per channel. |
| 4 | Added background waveform block (CH1, Helios ON / Genesis OFF) after the preamble blocks; stored as `uint32` sample count followed by raw `int8` ADC bytes. |
| 5 | (skipped) |
| 6 | Each angle now scans only the bounding box of the nominal ROI rotated by that angle instead of the AABB-expanded worst case across all angles. Header no longer carries a single global `x_start`/`x_delta`/`n_rows` — replaced with `*_nominal` reference fields plus a new Per-Angle Geometry Table (`x_start`, `x_delta`, `n_frames`, `n_rows` per angle) and a ragged Row Table / Waveform Data block sized per angle. **Not compatible with pre-v6 readers** that assume uniform geometry. `sras_viewer.py` reads v6 natively (per-angle `n_rows`/`n_frames`/`x_start`); the "Pre-process and Save as v5" fast-path export is not offered for v6 files since the flat v5 layout cannot represent per-angle geometry. |
| 6 | Each angle now scans only the bounding box of the nominal ROI rotated by that angle instead of the AABB-expanded worst case across all angles. Header no longer carries a single global `x_start`/`x_delta`/`n_rows` — replaced with `*_nominal` reference fields plus a new Per-Angle Geometry Table (`x_start`, `x_delta`, `n_frames`, `n_rows` per angle) and a ragged Row Table / Waveform Data block sized per angle. **Not compatible with pre-v6 readers** that assume uniform geometry. `sras_viewer.py` reads v6 natively (per-angle `n_rows`/`n_frames`/`x_start`). |
| 7 | Adds an optional trailing **Cache Tail** (`CACH` section, see [Cache Tail (v7)](#cache-tail-v7)) after the ragged waveform data, holding independently-present, independently-updatable per-angle DC (`SDCB`: dc3_mv + dc4_mv) and FFT (`SFFT`: peak_freq_mhz) blocks, so previously-computed images redisplay instantly instead of being recomputed. Header / angle table / geometry table / row table / preamble blocks / background block / waveform data are byte-identical to v6 — only the version byte and the optional Cache Tail differ. Scans still come off the scope as v6; `sras_viewer.py`'s "Convert" menu ("Batch Compute DC and Store" / "Batch Compute FFT and Store") converts a file to v7 **in place** on first use, or updates an existing v7 file's cache blocks, without rewriting any waveform bytes. Supersedes the removed "Pre-process and Save as v5" workflow, which was never available for v6 sources since the flat v5 `PREC` layout can't represent per-angle geometry. |
+78 -153
View File
@@ -15,25 +15,24 @@ Options:
Default: include a partial average for the last group.
"""
import sys
import struct
import argparse
import numpy as np
import shutil
import struct
import sys
from pathlib import Path
# ---------------------------------------------------------------------------
# Header format — must match sras_viewer.py exactly
# ---------------------------------------------------------------------------
import numpy as np
HDR_FMT = ">4sBHHffffIIdBB"
HDR_SIZE = struct.calcsize(HDR_FMT) # 43 bytes
from sras_format import HDR_FMT, HDR_SIZE, SrasFile
_SUPPORTED = (2, 3, 4)
def parse_args():
p = argparse.ArgumentParser(
description="Average every N waveforms in a .sras file and write a new file."
)
p.add_argument("input", help="Input .sras file")
p.add_argument("input", help="Input .sras file")
p.add_argument("output", help="Output .sras file")
p.add_argument("--n", type=int, required=True, metavar="N",
help="Number of consecutive frames to average into one")
@@ -42,143 +41,68 @@ def parse_args():
return p.parse_args()
def read_sras(path: Path):
"""Read all sections of a .sras file and return them as a dict."""
with open(path, "rb") as f:
header_bytes = f.read(HDR_SIZE)
fields = struct.unpack(HDR_FMT, header_bytes)
(magic, ver, n_angles, n_rows, x_start, x_delta, vel, freq,
n_frames_hdr, spf, sr, bps, n_ch) = fields
if magic != b"SRAS":
raise ValueError(f"Not a .sras file (bad magic: {magic!r})")
if ver not in (2, 3, 4):
raise ValueError(f"Unsupported .sras version: {ver}")
with open(path, "rb") as f:
def read_header_sections(sras: SrasFile) -> bytes:
"""The raw bytes between the header and the waveform data (angle table,
row table, preambles, background), copied through verbatim so nothing is
lost in a re-encode."""
with open(sras.path, "rb") as f:
f.seek(HDR_SIZE)
angles = f.read(n_angles * 4) # big-endian float32 array, raw bytes
y_pos = f.read(n_rows * 4) # big-endian float32 array, raw bytes
preambles = [] # list of raw bytes (length-prefixed strings)
if ver >= 3:
for _ in range(n_ch):
(length,) = struct.unpack(">H", f.read(2))
preambles.append(f.read(length))
background = b"" # raw bytes for the v4 background block
if ver >= 4:
(n_bg,) = struct.unpack(">I", f.read(4))
background = f.read(n_bg)
raw = f.read() # all waveform data
# -----------------------------------------------------------------------
# Determine actual frame count from file size (the header value can be
# wrong — the viewer does the same correction)
# -----------------------------------------------------------------------
total_samples = len(raw) // bps
samples_per_pixel = n_ch * spf # samples in one (angle, row, frame) cell
samples_per_full = n_angles * n_rows * samples_per_pixel
actual_n_frames = total_samples // (n_angles * n_rows * samples_per_pixel)
good_bytes = actual_n_frames * samples_per_full * bps
# Decode waveform data
dtype = np.int8 if bps == 1 else ">i2"
data = np.frombuffer(raw[:good_bytes], dtype=dtype)
data = data.reshape(n_angles, n_rows, n_ch, actual_n_frames, spf)
# Work in int16 (safe intermediate for both int8 and int16 inputs)
data = data.astype(np.int16)
return {
"ver": ver, "n_angles": n_angles, "n_rows": n_rows,
"x_start": x_start, "x_delta": x_delta, "vel": vel, "freq": freq,
"n_frames_hdr": n_frames_hdr, "spf": spf, "sr": sr,
"bps": bps, "n_ch": n_ch,
"angles_raw": angles, "y_pos_raw": y_pos,
"preambles": preambles, "background": background,
"data": data, # shape: (n_angles, n_rows, n_ch, n_frames, spf), int16
}
return f.read(sras._data_offset - HDR_SIZE)
def average_frames(data: np.ndarray, n: int, discard_remainder: bool) -> np.ndarray:
"""Average every N frames along axis 3.
def average_rows(block: np.ndarray, n: int, discard_remainder: bool) -> np.ndarray:
"""Average every N frames of one angle's (n_rows, n_ch, n_frames, spf)
block. Returns int16 of shape (n_rows, n_ch, n_out, spf).
data shape: (n_angles, n_rows, n_ch, n_frames, spf)
Returns array of shape (n_angles, n_rows, n_ch, n_out, spf).
Averaging is done in float32 and rounded on cast, matching numpy's mean
followed by an int16 cast in the original implementation.
"""
n_frames = data.shape[3]
n_full = n_frames // n
n_frames = block.shape[2]
n_full = n_frames // n
remainder = n_frames % n
# Average full groups using reshape-trick (no Python loop)
if n_full > 0:
full = data[:, :, :, :n_full * n, :] # trim to full groups
full = full.reshape(data.shape[0], data.shape[1], data.shape[2],
n_full, n, data.shape[4]) # (..., n_out, n, spf)
averaged = full.mean(axis=4).astype(np.int16) # (..., n_out, spf)
else:
averaged = np.empty((*data.shape[:3], 0, data.shape[4]), dtype=np.int16)
parts = []
if n_full:
full = block[:, :, :n_full * n, :].astype(np.float32)
full = full.reshape(block.shape[0], block.shape[1], n_full, n, block.shape[3])
parts.append(full.mean(axis=3).astype(np.int16))
if remainder and not discard_remainder:
tail = block[:, :, n_full * n:, :].astype(np.float32)
parts.append(tail.mean(axis=2, keepdims=True).astype(np.int16))
if remainder > 0 and not discard_remainder:
tail = data[:, :, :, n_full * n:, :] # shape (..., remainder, spf)
tail_avg = tail.mean(axis=3, keepdims=True).astype(np.int16)
averaged = np.concatenate([averaged, tail_avg], axis=3)
return averaged
if not parts:
return np.empty((*block.shape[:2], 0, block.shape[3]), dtype=np.int16)
return parts[0] if len(parts) == 1 else np.concatenate(parts, axis=2)
def write_sras(path: Path, src: dict, data_out: np.ndarray):
"""Write a new .sras file with the averaged waveform data."""
ver = src["ver"]
bps = src["bps"]
n_out = data_out.shape[3]
def write_averaged(out_path: Path, sras: SrasFile, mid_sections: bytes,
n: int, discard_remainder: bool) -> int:
"""Stream each angle through the averager, writing as we go so peak RAM
stays at one angle's block rather than the whole file."""
bps = sras.bytes_per_sample
n_frames_in = int(sras.n_frames[0])
n_out = n_frames_in // n
if n_frames_in % n and not discard_remainder:
n_out += 1
# Pack header — update only n_frames_hdr; everything else stays the same
header = struct.pack(
HDR_FMT,
b"SRAS",
ver,
src["n_angles"],
src["n_rows"],
src["x_start"],
src["x_delta"],
src["vel"],
src["freq"],
n_out, # updated frame count
src["spf"],
src["sr"],
bps,
src["n_ch"],
HDR_FMT, b"SRAS", sras.version, sras.n_angles, int(sras.n_rows[0]),
float(sras.x_start_mm[0]), float(sras.x_delta_mm),
sras.velocity_mm_s, sras.laser_freq_hz,
n_out, # updated frame count
sras.samples_per_frame, sras.sample_rate_hz, bps, sras.n_channels,
)
# Encode waveform data back to original dtype
if bps == 1:
raw_out = np.clip(data_out, -128, 127).astype(np.int8).tobytes()
else:
# big-endian int16
raw_out = data_out.astype(">i2").tobytes()
with open(path, "wb") as f:
with open(out_path, "wb") as f:
f.write(header)
f.write(src["angles_raw"])
f.write(src["y_pos_raw"])
if ver >= 3:
for preamble_bytes in src["preambles"]:
f.write(struct.pack(">H", len(preamble_bytes)))
f.write(preamble_bytes)
if ver >= 4:
bg = src["background"]
f.write(struct.pack(">I", len(bg)))
f.write(bg)
f.write(raw_out)
f.write(mid_sections)
for a in range(sras.n_angles):
averaged = average_rows(sras.data[a], n, discard_remainder)
if bps == 1:
f.write(np.clip(averaged, -128, 127).astype(np.int8).tobytes())
else:
f.write(averaged.astype(">i2").tobytes())
return n_out
def main():
@@ -188,32 +112,35 @@ def main():
print("Error: --n must be at least 1.", file=sys.stderr)
sys.exit(1)
in_path = Path(args.input)
in_path = Path(args.input)
out_path = Path(args.output)
if not in_path.exists():
print(f"Error: input file not found: {in_path}", file=sys.stderr)
sys.exit(1)
if out_path.resolve() == in_path.resolve():
print("Error: output path must differ from input path.", file=sys.stderr)
sys.exit(1)
print(f"Reading {in_path} ...", flush=True)
src = read_sras(in_path)
sras = SrasFile(str(in_path))
if sras.version not in _SUPPORTED:
print(f"Error: unsupported .sras version: {sras.version} "
f"(this tool handles v{'/v'.join(map(str, _SUPPORTED))})",
file=sys.stderr)
sys.exit(1)
n_frames_in = src["data"].shape[3]
print(f" Version : v{src['ver']}")
print(f" Angles : {src['n_angles']}")
print(f" Rows : {src['n_rows']}")
n_frames_in = int(sras.n_frames[0])
print(f" Version : v{sras.version}")
print(f" Angles : {sras.n_angles}")
print(f" Rows : {int(sras.n_rows[0])}")
print(f" Frames (actual): {n_frames_in}")
print(f" Channels : {src['n_ch']}")
print(f" Samples/frame : {src['spf']}")
print(f" Bytes/sample : {src['bps']}")
print(f" Channels : {sras.n_channels}")
print(f" Samples/frame : {sras.samples_per_frame}")
print(f" Bytes/sample : {sras.bytes_per_sample}")
if args.n == 1:
print("--n 1: no averaging needed; copying file as-is.")
import shutil
shutil.copy2(in_path, out_path)
print(f"Wrote {out_path}")
return
@@ -223,27 +150,25 @@ def main():
"The entire dataset will be averaged into a single frame.")
print(f"\nAveraging every {args.n} frames ...", flush=True)
data_out = average_frames(src["data"], args.n, args.discard_remainder)
n_frames_out = data_out.shape[3]
print(f"\nWriting {out_path} ...", flush=True)
mid_sections = read_header_sections(sras)
n_frames_out = write_averaged(out_path, sras, mid_sections,
args.n, args.discard_remainder)
n_full = n_frames_in // args.n
remainder = n_frames_in % args.n
n_full = n_frames_in // args.n
remainder = n_frames_in % args.n
if remainder and not args.discard_remainder:
status = f"({n_full} full groups + 1 partial group of {remainder})"
elif remainder and args.discard_remainder:
elif remainder:
status = f"({n_full} full groups, {remainder} trailing frames discarded)"
else:
status = f"({n_full} full groups)"
print(f" {n_frames_in} frames -> {n_frames_out} frames {status}")
print(f"\nWriting {out_path} ...", flush=True)
write_sras(out_path, src, data_out)
in_mb = in_path.stat().st_size / 1024**2
in_mb = in_path.stat().st_size / 1024**2
out_mb = out_path.stat().st_size / 1024**2
print(f" Input size : {in_mb:.1f} MB")
print(f" Output size: {out_mb:.1f} MB ({out_mb/in_mb*100:.1f}% of input)")
print(f" Output size: {out_mb:.1f} MB ({out_mb / in_mb * 100:.1f}% of input)")
print("Done.")
+1055
View File
File diff suppressed because it is too large Load Diff
+593
View File
@@ -0,0 +1,593 @@
#!/usr/bin/env python3
"""SRAS binary scan file format — parsing and writing.
Reads v2v7 .sras files. Depends only on numpy + struct, so compute workers
(including multiprocessing children) can import it without pulling in Qt or
matplotlib. See scan_format.md for the full v6/v7 spec.
Channel semantics (fixed by sc3_aui_app.py acquisition settings):
CH1 — RF Acoustic Packet (AC-coupled, 100 mV/div): FFT → peak frequency
CH3 — Bias A (DC-coupled, 50 mV/div): waveform mean
CH4 — Bias B (DC-coupled, 50 mV/div): waveform mean
Frame-count correction: the scanner writes the *configured* frame count in the
header before acquisition, but the scope may acquire fewer frames. The actual
count is computed from the file size and used for the reshape so channels are
correctly aligned.
Scan geometry: v6/v7 files scan a different bounding box per angle (x_start,
x_delta, n_frames, n_rows all vary by angle), so geometry is exposed per-angle
via SrasFile.n_rows / n_frames / x_start_mm arrays and the x_axis_mm() /
y_positions_mm() methods. v2v5 files have uniform geometry across angles, so
those arrays simply repeat the same value n_angles times.
v7 files are v6 files with an optional trailing cache section holding
precomputed per-angle DC and/or FFT images, so display never has to recompute
them after the "Convert" menu's batch actions have stored them once.
"""
import os
import re
import struct
from pathlib import Path
import numpy as np
# ---------------------------------------------------------------------------
# Format constants
# ---------------------------------------------------------------------------
# v2v5: fixed header, uniform geometry across angles (43 bytes)
HDR_FMT = ">4sBHHffffIIdBB"
HDR_SIZE = struct.calcsize(HDR_FMT)
# v6/v7: fixed header, per-angle geometry in a separate table (49 bytes)
HDR_FMT_V6 = ">4sBHfffffffIdBB"
HDR_SIZE_V6 = struct.calcsize(HDR_FMT_V6)
# v6/v7: per-angle geometry record (x_start, x_delta, n_frames, n_rows)
GEO_FMT_V6 = ">ffIH"
GEO_SIZE_V6 = struct.calcsize(GEO_FMT_V6)
# v5 precomputed-image tail
PREC_MAGIC = b"PREC"
PREC_FLAG_BG_SUB = 0x01
# v7 cache tail. v7 is byte-identical to v6 (the version byte is the only
# header difference) plus this optional trailing section. Sub-block sizes are
# per-angle (n_rows[a] * n_frames[a]), taken from the Per-Angle Geometry Table
# already parsed for v6 — no new geometry fields are needed.
CACH_MAGIC = b"CACH"
CACH_HDR_FMT = ">4sBB" # magic, cach_version, block_flags
CACH_HDR_SIZE = struct.calcsize(CACH_HDR_FMT)
CACH_VERSION = 1
CACH_FLAG_DC = 0x01
CACH_FLAG_FFT = 0x02
SDCB_MAGIC = b"SDCB"
SDCB_HDR_FMT = ">4sBH" # magic, reserved, n_stored
SDCB_HDR_SIZE = struct.calcsize(SDCB_HDR_FMT)
SFFT_MAGIC = b"SFFT"
SFFT_HDR_FMT = ">4sBH" # magic, flags, n_stored
SFFT_HDR_SIZE = struct.calcsize(SFFT_HDR_FMT)
SFFT_FLAG_BG_SUB = 0x01
# Fixed channel indices into the .sras data array (CH1=RF, CH3/CH4=Bias DC)
CH1_IDX, CH3_IDX, CH4_IDX = 0, 1, 2
CH_NAMES = ["CH1", "CH3", "CH4", "VEL"]
# Fallback scope calibration used only when reading v2 files without embedded
# preambles. v3+ files carry the WFMOutpre string so these are not used.
# 50 mV/div, 8 div full-scale, int8 ADC, position = -2.72 div
# ymult = 50 mV × 8 / 256 = 1.5625 mV/count
# yoff = position × (256/8) = -2.72 × 32 = -87.04 (ADC count for 0 V)
_FALLBACK_YMULT_MV = 1.5625 # mV per ADC count
_FALLBACK_YOFF_ADC = -87.04 # ADC count that represents 0 V
# ---------------------------------------------------------------------------
# Calibration
# ---------------------------------------------------------------------------
def _parse_preamble(preamble: str) -> dict[str, float]:
"""Extract YMULT, YOFF, YZERO from a Tektronix WFMOutpre string.
Returns a dict with float values for whichever keys are present.
YMULT is left in V/count as the scope reports it.
"""
result = {}
for key in ("YMULT", "YOFF", "YZERO"):
m = re.search(rf'\b{key}\s+([-+]?\d*\.?\d+(?:[Ee][+-]?\d+)?)', preamble)
if m:
result[key] = float(m.group(1))
return result
def mv_to_adc(mv: float, ymult_mv: float = _FALLBACK_YMULT_MV,
yoff_adc: float = _FALLBACK_YOFF_ADC,
yzero_mv: float = 0.0) -> float:
return (mv - yzero_mv) / ymult_mv + yoff_adc
def adc_to_mv(adc, ymult_mv: float = _FALLBACK_YMULT_MV,
yoff_adc: float = _FALLBACK_YOFF_ADC,
yzero_mv: float = 0.0):
return (adc - yoff_adc) * ymult_mv + yzero_mv
# ---------------------------------------------------------------------------
# Binary read helpers
# ---------------------------------------------------------------------------
def _read_struct(f, fmt: str) -> tuple:
return struct.unpack(fmt, f.read(struct.calcsize(fmt)))
def _read_preambles(f, n_ch: int) -> list[str]:
"""n_ch length-prefixed UTF-8 WFMOutpre strings."""
out = []
for _ in range(n_ch):
(length,) = _read_struct(f, ">H")
out.append(f.read(length).decode("utf-8"))
return out
def _read_background(f) -> np.ndarray:
"""uint32 sample count followed by that many int8 samples."""
(n_bg,) = _read_struct(f, ">I")
return np.frombuffer(f.read(n_bg), dtype=np.int8).astype(np.float32)
def _read_f32_image(f, shape: tuple[int, int]) -> np.ndarray:
"""One big-endian float32 image, converted to native float32.
The conversion matters: np.frombuffer hands back a read-only big-endian
view, and these arrays flow straight into the display caches and every
downstream arithmetic op.
"""
n_bytes = shape[0] * shape[1] * 4
return np.frombuffer(f.read(n_bytes), dtype=">f4").reshape(shape).astype(np.float32)
# ---------------------------------------------------------------------------
# File parser
# ---------------------------------------------------------------------------
class SrasFile:
"""Parsed in-memory representation of a v2v7 .sras file.
Scan geometry (rows, frames, x_start) is exposed per-angle via the
``n_rows`` / ``n_frames`` / ``x_start_mm`` arrays and the ``x_axis_mm()``
/ ``y_positions_mm()`` methods, since v6/v7 files scan a different
bounding box per angle. v2v5 files have uniform geometry, so these
arrays just repeat the same value ``n_angles`` times. Waveform data is
likewise exposed as ``data[angle_idx]``, an array of shape
``(n_rows[a], n_channels, n_frames[a], samples_per_frame)``.
Precomputed images (v5's PREC tail or v7's CACH tail) are exposed as
``precomputed_dc3_mv`` / ``precomputed_dc4_mv`` / ``precomputed_freq_mhz``,
always as ragged per-angle lists (``list[np.ndarray | None]``, one entry
per angle, ``None`` where that angle was never stored) regardless of
source version.
"""
def __init__(self, path: str):
self.path = Path(path)
self._parse()
def _parse(self):
with open(self.path, "rb") as f:
magic = f.read(4)
if magic != b"SRAS":
raise ValueError(f"Bad magic bytes: {magic!r}")
(version,) = struct.unpack(">B", f.read(1))
self.version = version
if version in (2, 3, 4, 5):
self._parse_legacy()
elif version in (6, 7):
self._parse_v6()
else:
raise ValueError(f"Unsupported version: {version}")
# ------------------------------------------------------------------
# Calibration
# ------------------------------------------------------------------
def _set_calibration(self, preambles: list[str] | None, n_ch: int):
"""Populate the per-channel ymult/yoff/yzero lists from preamble
strings, falling back to the hardcoded scope constants for v2 files
that carry no preambles."""
if preambles is None:
self.preambles = None
self.ch_ymult_mv = [_FALLBACK_YMULT_MV] * n_ch
self.ch_yoff_adc = [_FALLBACK_YOFF_ADC] * n_ch
self.ch_yzero_mv = [0.0] * n_ch
return
self.preambles = preambles
self.ch_ymult_mv, self.ch_yoff_adc, self.ch_yzero_mv = [], [], []
for p in preambles:
cal = _parse_preamble(p)
# The scope reports YMULT and YZERO in volts; store both as mV.
self.ch_ymult_mv.append(cal.get("YMULT", _FALLBACK_YMULT_MV / 1000) * 1000)
self.ch_yoff_adc.append(cal.get("YOFF", _FALLBACK_YOFF_ADC))
self.ch_yzero_mv.append(cal.get("YZERO", 0.0) * 1000)
def cal(self, ch_idx: int) -> tuple[float, float, float]:
"""(ymult_mv, yoff_adc, yzero_mv) for one channel — splat straight
into adc_to_mv / mv_to_adc."""
return (self.ch_ymult_mv[ch_idx], self.ch_yoff_adc[ch_idx],
self.ch_yzero_mv[ch_idx])
def _init_precomputed(self, n_angles: int):
self.precomputed_freq_mhz: list[np.ndarray | None] = [None] * n_angles
self.precomputed_dc4_mv: list[np.ndarray | None] = [None] * n_angles
self.precomputed_dc3_mv: list[np.ndarray | None] = [None] * n_angles
self.precomputed_bg_sub: bool = False
def cached_dc_mv(self, angle_idx: int, ch_idx: int) -> np.ndarray | None:
"""A stored DC image (already in mV) for (angle, channel), or None."""
store = self.precomputed_dc3_mv if ch_idx == CH3_IDX else self.precomputed_dc4_mv
return store[angle_idx] if angle_idx < len(store) else None
def image_shape(self, angle_idx: int) -> tuple[int, int]:
return int(self.n_rows[angle_idx]), int(self.n_frames[angle_idx])
# ------------------------------------------------------------------
# v2v5 parsing (uniform geometry, flat waveform block)
# ------------------------------------------------------------------
def _parse_legacy(self):
with open(self.path, "rb") as f:
(magic, ver, n_angles, n_rows, x_start, x_delta, vel, freq,
n_frames_hdr, spf, sr, bps, n_ch) = _read_struct(f, HDR_FMT)
self.n_angles = n_angles
self.velocity_mm_s = float(vel)
self.laser_freq_hz = float(freq)
self.n_frames_header = n_frames_hdr # configured count (may be wrong)
self.samples_per_frame = spf
self.sample_rate_hz = float(sr)
self.bytes_per_sample = bps
self.n_channels = n_ch
self._init_precomputed(n_angles)
self.scan_aborted = False
self.n_angles_declared = n_angles
angles = np.frombuffer(f.read(n_angles * 4), dtype=">f4").astype(np.float32)
y_pos = np.frombuffer(f.read(n_rows * 4), dtype=">f4").astype(np.float32)
self._set_calibration(_read_preambles(f, n_ch) if ver >= 3 else None, n_ch)
self.background = _read_background(f) if ver >= 4 else None
# Record where raw waveform data begins; np.memmap maps from here.
data_offset = f.tell()
# ---- Determine actual frame count from file size ---------------
# For v4 and earlier the header n_frames may be the *configured*
# count before acquisition; the actual count is derived from the
# bytes on disk. For v5 files a PREC tail follows the waveform
# data, so we must not include those extra bytes in the frame count.
file_size = self.path.stat().st_size
samples_per_row_per_ch = n_ch * spf
available_bytes = file_size - data_offset
if ver == 5:
actual_n_frames = n_frames_hdr
remainder = 0
else:
total_samples = available_bytes // bps
per_frame = n_angles * n_rows * samples_per_row_per_ch
actual_n_frames = total_samples // per_frame
remainder = total_samples % per_frame
self.frame_count_mismatch = (actual_n_frames != n_frames_hdr)
self.n_frames_remainder = remainder
# ---- Memory-map the waveform data (zero RAM cost) --------------
# Instead of f.read() → astype() (which peaks at 2× file size),
# memmap lets the OS page only the bytes that are actually touched.
data5d = np.memmap(
str(self.path),
dtype=np.int8 if bps == 1 else ">i2",
mode="r",
offset=data_offset,
shape=(n_angles, n_rows, n_ch, actual_n_frames, spf),
)
# Expose as a list of per-angle views so downstream code shares one
# indexing convention with v6: sras.data[a][row, ch, frame, sample]
self.data = [data5d[a] for a in range(n_angles)]
# Uniform per-angle geometry, repeated so callers don't need to
# special-case legacy vs. v6 files.
self.n_rows = np.full(n_angles, n_rows, dtype=np.int64)
self.n_frames = np.full(n_angles, actual_n_frames, dtype=np.int64)
self.x_start_mm = np.full(n_angles, float(x_start), dtype=np.float64)
self.x_delta_mm = float(x_delta) # reference only; kept for re-encode
self._y_pos_per_angle = [y_pos] * n_angles
self.angles_deg = angles
self._data_offset = data_offset
if ver >= 5:
waveform_bytes = actual_n_frames * n_angles * n_rows * n_ch * spf * bps
prec_offset = data_offset + waveform_bytes
if file_size > prec_offset:
self._parse_prec_section(prec_offset)
def _parse_prec_section(self, offset: int):
"""Parse the v5 PREC tail that holds precomputed images.
Stored per-angle as (freq, dc4, dc3), each a full-image float32
block prefixed by its uint16 angle index.
"""
with open(self.path, "rb") as f:
f.seek(offset)
header_raw = f.read(6) # magic(4) + fmt_ver(1) + flags(1)
if len(header_raw) < 6 or header_raw[:4] != PREC_MAGIC:
return
self.precomputed_bg_sub = bool(header_raw[5] & PREC_FLAG_BG_SUB)
(n_stored,) = _read_struct(f, ">H")
for _ in range(n_stored):
(aidx,) = _read_struct(f, ">H")
if aidx >= self.n_angles:
break
shape = self.image_shape(aidx)
self.precomputed_freq_mhz[aidx] = _read_f32_image(f, shape)
self.precomputed_dc4_mv[aidx] = _read_f32_image(f, shape)
self.precomputed_dc3_mv[aidx] = _read_f32_image(f, shape)
# ------------------------------------------------------------------
# v6/v7 parsing (per-angle geometry, ragged waveform blocks)
# ------------------------------------------------------------------
def _parse_v6(self):
with open(self.path, "rb") as f:
(magic, ver, n_angles, x_start_nom, y_start_nom, x_delta_nom,
y_delta_nom, row_spacing, vel, freq, spf, sr, bps,
n_ch) = _read_struct(f, HDR_FMT_V6)
n_angles_declared = n_angles
self.velocity_mm_s = float(vel)
self.laser_freq_hz = float(freq)
self.samples_per_frame = spf
self.sample_rate_hz = float(sr)
self.bytes_per_sample = bps
self.n_channels = n_ch
# Reference-only fields: the ROI as entered before per-angle
# bounding-box expansion. Actual per-angle geometry used for
# rendering comes from the Per-Angle Geometry Table below.
self.x_start_nominal_mm = float(x_start_nom)
self.y_start_nominal_mm = float(y_start_nom)
self.x_delta_nominal_mm = float(x_delta_nom)
self.y_delta_nominal_mm = float(y_delta_nom)
self.row_spacing_mm = float(row_spacing)
self.n_frames_header = None
self.frame_count_mismatch = False
self.n_frames_remainder = 0
angles = np.frombuffer(f.read(n_angles * 4), dtype=">f4").astype(np.float32)
x_start = np.empty(n_angles, dtype=np.float64)
n_frames = np.empty(n_angles, dtype=np.int64)
n_rows = np.empty(n_angles, dtype=np.int64)
for a in range(n_angles):
xs, xd, nf, nr = _read_struct(f, GEO_FMT_V6)
x_start[a], n_frames[a], n_rows[a] = xs, nf, nr
y_pos_per_angle = [
np.frombuffer(f.read(int(n_rows[a]) * 4), dtype=">f4").astype(np.float32)
for a in range(n_angles)
]
self._set_calibration(_read_preambles(f, n_ch), n_ch)
self.background = _read_background(f)
data_offset = f.tell()
self._data_offset = data_offset
# ---- Memory-map each angle's ragged waveform block -------------
# v6 gives each angle its own row/frame count, so waveform data is
# no longer one uniform (n_angles, n_rows, ...) block — each angle's
# block sits at a different offset with its own shape. An aborted
# scan truncates the file mid-angle; per the format spec we keep
# whatever complete angles are present rather than refusing to open
# the file.
file_size = self.path.stat().st_size
waveform_dtype = np.int8 if bps == 1 else ">i2"
data = []
offset = data_offset
for a in range(n_angles):
nr, nf = int(n_rows[a]), int(n_frames[a])
nbytes = nr * n_ch * nf * spf * bps
if offset + nbytes > file_size:
break
data.append(np.memmap(
str(self.path), dtype=waveform_dtype, mode="r",
offset=offset, shape=(nr, n_ch, nf, spf),
))
offset += nbytes
n_complete = len(data)
if n_complete == 0:
raise ValueError(
"v6 file has no complete angle blocks — scan was aborted "
"before the first angle finished.")
self.data = data
self.n_angles = n_complete
self.n_angles_declared = n_angles_declared
self.scan_aborted = n_complete < n_angles_declared
self.angles_deg = angles[:n_complete]
self.x_start_mm = x_start[:n_complete]
self.n_frames = n_frames[:n_complete]
self.n_rows = n_rows[:n_complete]
self._y_pos_per_angle = y_pos_per_angle[:n_complete]
self._init_precomputed(n_complete)
if self.version == 7 and offset < file_size:
self._parse_cach_section(offset)
# ------------------------------------------------------------------
# v7 cache tail (CACH section: precomputed DC / FFT images)
# ------------------------------------------------------------------
def _cache_tail_offset(self) -> int:
"""Deterministic file offset where the CACH tail starts (or would
start), derived purely from the header + Per-Angle Geometry Table —
independent of whether a cache tail is actually present. Used by
both the parser and the in-place writer."""
waveform_bytes = sum(
int(self.n_rows[a]) * self.n_channels * int(self.n_frames[a])
* self.samples_per_frame * self.bytes_per_sample
for a in range(self.n_angles)
)
return self._data_offset + int(waveform_bytes)
def _read_cache_block(self, f, hdr_fmt: str, magic: bytes,
stores: list[list]) -> int | None:
"""Read one CACH sub-block header, then its per-angle image entries
into *stores* (one list per image the block stores per angle).
Returns the header's flags byte, or None if the block is malformed.
"""
raw = f.read(struct.calcsize(hdr_fmt))
if len(raw) < struct.calcsize(hdr_fmt):
return None
block_magic, flags, n_stored = struct.unpack(hdr_fmt, raw)
if block_magic != magic:
return None
for _ in range(n_stored):
(angle_idx,) = _read_struct(f, ">H")
if angle_idx >= self.n_angles:
break
shape = self.image_shape(angle_idx)
for store in stores:
store[angle_idx] = _read_f32_image(f, shape)
return flags
def _parse_cach_section(self, offset: int):
"""Parse the v7 CACH tail that holds precomputed DC/FFT images."""
with open(self.path, "rb") as f:
f.seek(offset)
header_raw = f.read(CACH_HDR_SIZE)
if len(header_raw) < CACH_HDR_SIZE:
return
magic, cach_version, block_flags = struct.unpack(CACH_HDR_FMT, header_raw)
if magic != CACH_MAGIC or cach_version != CACH_VERSION:
return
if block_flags & CACH_FLAG_DC:
if self._read_cache_block(
f, SDCB_HDR_FMT, SDCB_MAGIC,
[self.precomputed_dc3_mv, self.precomputed_dc4_mv]) is None:
return
if block_flags & CACH_FLAG_FFT:
flags = self._read_cache_block(
f, SFFT_HDR_FMT, SFFT_MAGIC, [self.precomputed_freq_mhz])
if flags is None:
return
self.precomputed_bg_sub = bool(flags & SFFT_FLAG_BG_SUB)
def write_v7_cache(self, *,
new_dc3_mv: list[np.ndarray | None] | None = None,
new_dc4_mv: list[np.ndarray | None] | None = None,
new_freq_mhz: list[np.ndarray | None] | None = None,
new_bg_sub: bool | None = None):
"""Store computed DC and/or FFT images into this file's CACH tail,
in place, converting a v6 source to v7 (or updating an existing v7
file). Only the block(s) passed in are recomputed; whichever block
isn't passed is carried forward unchanged from whatever this
``SrasFile`` already has in memory (from parsing, or a prior write
in this same session) — its bytes are never re-read from disk.
The waveform data itself is never touched: the cache tail always
starts at ``_cache_tail_offset()``, a fixed offset derived from the
header and geometry table alone.
"""
if self.version not in (6, 7):
raise ValueError(
f"write_v7_cache only supports v6/v7 source files, got v{self.version}")
final_dc3 = new_dc3_mv if new_dc3_mv is not None else self.precomputed_dc3_mv
final_dc4 = new_dc4_mv if new_dc4_mv is not None else self.precomputed_dc4_mv
final_freq = new_freq_mhz if new_freq_mhz is not None else self.precomputed_freq_mhz
final_bg_sub = new_bg_sub if new_bg_sub is not None else self.precomputed_bg_sub
dc_entries = [a for a in range(self.n_angles) if final_dc3[a] is not None]
fft_entries = [a for a in range(self.n_angles) if final_freq[a] is not None]
block_flags = ((CACH_FLAG_DC if dc_entries else 0)
| (CACH_FLAG_FFT if fft_entries else 0))
payload = bytearray()
payload += struct.pack(CACH_HDR_FMT, CACH_MAGIC, CACH_VERSION, block_flags)
if dc_entries:
payload += struct.pack(SDCB_HDR_FMT, SDCB_MAGIC, 0, len(dc_entries))
for a in dc_entries:
payload += struct.pack(">H", a)
payload += final_dc3[a].astype(">f4").tobytes()
payload += final_dc4[a].astype(">f4").tobytes()
if fft_entries:
fft_flags = SFFT_FLAG_BG_SUB if final_bg_sub else 0
payload += struct.pack(SFFT_HDR_FMT, SFFT_MAGIC, fft_flags, len(fft_entries))
for a in fft_entries:
payload += struct.pack(">H", a)
payload += final_freq[a].astype(">f4").tobytes()
with open(self.path, "r+b") as f:
f.seek(self._cache_tail_offset())
f.write(payload)
f.truncate()
f.flush()
os.fsync(f.fileno())
# Version-byte flip last: if the process dies before this point,
# the file is still readable as plain v6 (v6 parsing only
# bounds-checks per-angle offset+nbytes <= file_size, it never
# asserts exactly how many bytes follow the last angle) — so an
# interrupted write can never corrupt the file, only leave
# harmless trailing bytes that the next successful write
# overwrites via this same deterministic cache offset.
f.seek(4)
f.write(struct.pack("B", 7))
f.flush()
os.fsync(f.fileno())
self.version = 7
self.precomputed_dc3_mv = final_dc3
self.precomputed_dc4_mv = final_dc4
self.precomputed_freq_mhz = final_freq
self.precomputed_bg_sub = final_bg_sub
# ------------------------------------------------------------------
# Axes helpers
# ------------------------------------------------------------------
@property
def pixel_x_mm(self) -> float:
return self.velocity_mm_s / self.laser_freq_hz
def x_axis_mm(self, angle_idx: int) -> np.ndarray:
n = int(self.n_frames[angle_idx])
return self.x_start_mm[angle_idx] + np.arange(n) * self.pixel_x_mm
def y_positions_mm(self, angle_idx: int) -> np.ndarray:
return self._y_pos_per_angle[angle_idx]
def time_axis_ns(self) -> np.ndarray:
return np.arange(self.samples_per_frame) / self.sample_rate_hz * 1e9
def freq_axis_mhz(self, n_fft: int | None = None) -> np.ndarray:
n = n_fft if n_fft is not None else self.samples_per_frame
return np.fft.rfftfreq(n, d=1.0 / self.sample_rate_hz) / 1e6
+1647 -2006
View File
File diff suppressed because it is too large Load Diff
+402
View File
@@ -0,0 +1,402 @@
#!/usr/bin/env python3
"""Background workers for the SRAS viewer.
Every worker is a plain QObject moved onto its own QThread by
SrasViewerWindow._run_worker, exposing signals only. Workers must never touch
GUI-thread-owned state (the display caches in particular) — they take
everything they need through their constructor and hand results back by signal.
"""
import os
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
from concurrent.futures.process import BrokenProcessPool
import numpy as np
from PyQt6.QtCore import QObject, pyqtSignal
import sras_compute as compute
from sras_compute import (
cache_file, compute_angle_alignment, compute_rf_image, dc_image_mv,
)
from sras_format import CH3_IDX, CH4_IDX, SrasFile
# Concurrency caps. Batch conversion runs one process per file, and each of
# those processes threads internally, so the two must be divided rather than
# both set to the core count. Files also commonly sit on one external drive,
# where a dozen concurrent readers is slower than a few — hence the low
# default, overridable from the environment.
_BATCH_MAX_PROCS = int(os.environ.get("SRAS_BATCH_PROCS", 0)) or min(
4, os.cpu_count() or 2)
# Spawning a pool costs roughly a second of interpreter startup (each child
# re-imports the entry module). That is noise against a multi-GB scan but
# dominates a batch of small files, where it would make the job *slower* —
# so below this total size the batch just runs in the worker thread.
_BATCH_POOL_MIN_BYTES = int(os.environ.get("SRAS_BATCH_POOL_MIN_MB", 512)) * 1024 * 1024
class CancellableWorker(QObject):
"""A worker whose compute polls stop() between row chunks.
Without this a shutdown has to wait out whatever is in flight, and on a
large scan a single angle is ~40 s — far too long to block closing the
window. Chunk-level polling bounds the wait to one chunk instead.
"""
def __init__(self):
super().__init__()
self._stop = False
def stop(self):
self._stop = True
def _stopped(self) -> bool:
return self._stop
class LoadWorker(QObject):
finished = pyqtSignal(object) # SrasFile | None
error = pyqtSignal(str)
def __init__(self, path: str):
super().__init__()
self._path = path
def run(self):
try:
self.finished.emit(SrasFile(self._path))
except Exception as exc:
self.error.emit(str(exc))
self.finished.emit(None)
class ComputeWorker(CancellableWorker):
"""Computes one displayable image for (angle, channel).
For CH1/Velocity (FFT-derived) channels, the FFT is only run for pixels
whose DC4 (Bias B) mean is at or above dc_threshold_mv — masked pixels are
left at 0 MHz without ever being FFT'd, since that's the expensive part of
a scan. If the DC4 image for this angle is already known, pass it in as
*dc4_mv* to skip re-reading the CH4 channel from disk entirely.
Emits a plain ``np.ndarray`` already in display units.
"""
finished = pyqtSignal(object)
error = pyqtSignal(str)
def __init__(self, sras: SrasFile, angle_idx: int, ch_idx: int,
apply_bg_sub: bool = True, n_fft: int | None = None,
dc_threshold_mv: float = 0.0,
dc4_mv: np.ndarray | None = None,
is_fft_mode: bool = False):
super().__init__()
self._sras = sras
self._angle = angle_idx
self._ch = ch_idx
self._apply_bg_sub = apply_bg_sub
self._n_fft = n_fft
self._dc_threshold = dc_threshold_mv
self._dc4_mv = dc4_mv
self._is_fft_mode = is_fft_mode
def run(self):
try:
if self._is_fft_mode:
img = compute_rf_image(
self._sras, self._angle, dc_threshold_mv=self._dc_threshold,
apply_bg_sub=self._apply_bg_sub, n_fft=self._n_fft,
dc4_mv=self._dc4_mv, should_stop=self._stopped)
else:
img = dc_image_mv(self._sras, self._angle, self._ch,
should_stop=self._stopped)
# On cancellation the image is only partly filled, so hand back
# None rather than something that would be cached as real. The
# signal still fires either way — it is what quits the thread.
self.finished.emit(None if self._stop else img)
except Exception as exc:
self.error.emit(str(exc))
class DcPrecomputeWorker(CancellableWorker):
"""Computes CH3/CH4 DC images for every angle in the background.
DC images are cheap (a per-waveform mean, no FFT) compared to the
CH1/Velocity FFT, so precomputing them for the whole file right after load
makes switching angles instant while on a DC channel, and also means the
FFT masking step (which needs a DC4 image) rarely has to wait on anything.
Angles are computed on a thread pool — the work is a pure mean over the
waveform block, so it is I/O- and bandwidth-bound and embarrassingly
parallel. Results are emitted one at a time as they land (out of angle
order), and always from this worker's own thread: nothing emits a Qt
signal from a pool thread.
"""
angle_done = pyqtSignal(int, np.ndarray, np.ndarray) # angle_idx, dc3_mv, dc4_mv
finished = pyqtSignal()
error = pyqtSignal(str)
def __init__(self, sras: SrasFile):
super().__init__()
self._sras = sras
def _one_angle(self, a: int) -> tuple[int, np.ndarray, np.ndarray]:
# max_workers=1 *and* a budget share: this call is one of several
# concurrent angles, and both the thread count and the buffer size
# have to be divided (see compute.plan_angle_level).
kw = dict(max_workers=1, budget=self._angle_budget,
should_stop=self._stopped)
return (a,
dc_image_mv(self._sras, a, CH3_IDX, **kw),
dc_image_mv(self._sras, a, CH4_IDX, **kw))
def run(self):
try:
n = self._sras.n_angles
n_workers, self._angle_budget = compute.plan_angle_level(self._sras)
pool = ThreadPoolExecutor(max_workers=n_workers)
try:
futures = {pool.submit(self._one_angle, a): a for a in range(n)}
for fut in as_completed(futures):
if self._stop:
break
a, dc3, dc4 = fut.result()
self.angle_done.emit(a, dc3, dc4)
finally:
# cancel_futures drops the queued angles; should_stop lets the
# in-flight ones bail within a chunk. Not waiting here is what
# keeps closing the window responsive on a large scan.
pool.shutdown(wait=not self._stop, cancel_futures=True)
self.finished.emit()
except Exception as exc:
self.error.emit(str(exc))
class BatchCacheWorker(QObject):
"""Batch-computes and stores DC or FFT images into each of *paths*'s v7
CACH tail, in place — converting v6 sources to v7 on first use, or updating
an existing v7 file's cache blocks without disturbing whatever the other
block already holds.
*mode* is ``"dc"`` (CH3/CH4 mean images) or ``"fft"`` (CH1 peak-frequency
images, unmasked — masking is applied at display time, same as v5's PREC
convention).
Files are processed one per subprocess: they are fully independent, each
opens its own memmap and writes only its own bytes, and only path strings
and scalars cross the process boundary. Emits ``progress(int)`` (0100 by
files completed), ``file_done(str, str)`` (path, error message or "") so
one file's failure doesn't abort the batch, and ``finished()``.
"""
progress = pyqtSignal(int)
file_done = pyqtSignal(str, str)
finished = pyqtSignal()
def __init__(self, paths: list[str], mode: str, apply_bg_sub: bool):
super().__init__()
self._paths = paths
self._mode = mode
self._apply_bg_sub = apply_bg_sub
def _report(self, path: str, err: str, done: int, total: int):
self.file_done.emit(path, err)
self.progress.emit(int(done / max(1, total) * 100))
def _run_pooled(self, paths: list[str], n_procs: int) -> list[str]:
"""Process the batch across *n_procs* subprocesses. Returns the paths
that never got a real answer because the pool itself died, so the
caller can retry them in-process.
Under spawn each child re-imports the entry module, so the batch must
survive that going wrong (an unguarded __main__, a frozen build, a
sandbox that forbids subprocesses) rather than reporting every file as
failed — hence the retry list instead of a per-file error.
"""
# Each child threads internally; divide the machine rather than
# letting every process claim every core.
per_proc_workers = max(1, (os.cpu_count() or 4) // n_procs)
unresolved: list[str] = []
done = 0
with ProcessPoolExecutor(max_workers=n_procs) as executor:
futures = {
executor.submit(cache_file, p, self._mode, self._apply_bg_sub,
compute.get_fft_backend(), per_proc_workers): p
for p in paths
}
for fut in as_completed(futures):
path = futures[fut]
try:
err = fut.result()
except BrokenProcessPool:
unresolved.append(path)
continue
except Exception as exc:
err = str(exc)
done += 1
self._report(path, err, done, len(paths))
return unresolved
def _run_inline(self, paths: list[str], done: int, total: int):
"""Fallback / single-file path: compute in this thread. Still uses the
full core count internally, since nothing else is competing."""
for path in paths:
try:
err = cache_file(path, self._mode, self._apply_bg_sub,
compute.get_fft_backend(), compute._MAX_WORKERS)
except Exception as exc:
err = str(exc)
done += 1
self._report(path, err, done, total)
def _worth_pooling(self, paths: list[str]) -> bool:
if len(paths) < 2:
return False
total = 0
for p in paths:
try:
total += os.path.getsize(p)
except OSError:
pass # unreadable files are reported by cache_file
return total >= _BATCH_POOL_MIN_BYTES
def run(self):
paths = self._paths
n_procs = max(1, min(_BATCH_MAX_PROCS, len(paths)))
if not self._worth_pooling(paths):
self._run_inline(paths, 0, len(paths))
self.finished.emit()
return
try:
unresolved = self._run_pooled(paths, n_procs)
except Exception:
# The pool could not be created or collapsed wholesale.
unresolved = list(paths)
if unresolved:
self._run_inline(unresolved, len(paths) - len(unresolved), len(paths))
self.finished.emit()
class AngleAlignmentWorker(QObject):
"""Computes rotation+translation alignment for every angle in *sras*,
referenced to *ref_angle_idx*, from each angle's binarized CH4 mask.
Rotation is analytic (from sras.angles_deg); only translation is found by
phase correlation.
"""
progress = pyqtSignal(int) # 0100
finished = pyqtSignal(object, str) # AlignmentResult|None, error ("" = success)
def __init__(self, sras: SrasFile, ref_angle_idx: int, dc_threshold_mv: float):
super().__init__()
self._sras = sras
self._ref = ref_angle_idx
self._threshold = dc_threshold_mv
def run(self):
try:
result = compute_angle_alignment(
self._sras, self._ref, self._threshold,
progress_cb=self.progress.emit)
self.finished.emit(result, "")
except Exception as exc:
self.finished.emit(None, str(exc))
class Ch4MaskWorker(QObject):
"""Fetches each requested angle's CH4 (Bias B) DC image in mV, for
ManualAlignmentDialog's initial threshold-mask overlay.
Reuses dc_image_mv, which prefers a stored v5/v7 cache over recomputing
from raw waveforms, so this only does real work for a file that hasn't
gone through the v7 "Convert" batch step and for angles the main
window's own DcPrecomputeWorker (which runs automatically right after
every file load) hasn't reached yet. In the common case — the user opens
Fusion -> Manual Alignment after DC precompute has already finished —
*angle_indices* is empty and this worker is never even constructed (see
ManualAlignmentDialog._start_mask_prep).
"""
angle_done = pyqtSignal(int, np.ndarray) # angle_idx, dc4_mv
finished = pyqtSignal()
error = pyqtSignal(str)
def __init__(self, sras: SrasFile, angle_indices: list[int]):
super().__init__()
self._sras = sras
self._angles = angle_indices
def run(self):
try:
n_workers, budget = compute.plan_angle_level(self._sras)
pool = ThreadPoolExecutor(max_workers=n_workers)
try:
futures = {
pool.submit(dc_image_mv, self._sras, a, CH4_IDX,
max_workers=1, budget=budget): a
for a in self._angles
}
for fut in as_completed(futures):
a = futures[fut]
self.angle_done.emit(a, fut.result())
finally:
pool.shutdown(wait=True)
self.finished.emit()
except Exception as exc:
self.error.emit(str(exc))
class CrossCorrelateWorker(QObject):
"""FFT phase-correlation translation for each of *angle_indices* against
*ref_angle_idx*, for ManualAlignmentDialog's Auto Cross-Correlate button.
Runs on a background thread — a real many-angle, high-resolution scan's
correlation (even at its downsampled working resolution) can take long
enough that doing all of them on the GUI thread would visibly freeze the
dialog. Rotation is set to the same analytic scan-angle delta Auto
De-rotate uses alongside the correlated shift, since a translation
search is only meaningful once both angles' content is already oriented
the same way. dc4_mv/pivot_mm are the dialog's own already-in-memory
per-angle images/pivots — this worker does no fetching of its own.
"""
angle_done = pyqtSignal(int, float, float, float) # angle_idx, rotation_deg, shift_x_mm, shift_y_mm
finished = pyqtSignal()
error = pyqtSignal(str)
def __init__(self, sras: SrasFile, ref_angle_idx: int, angle_indices: list[int],
dc4_mv: dict[int, np.ndarray], pivot_mm: dict[int, tuple[float, float]],
*, use_mask: bool, dc_threshold_mv: float, margin_frac: float):
super().__init__()
self._sras = sras
self._ref = ref_angle_idx
self._angles = angle_indices
self._dc4_mv = dc4_mv
self._pivot_mm = pivot_mm
self._use_mask = use_mask
self._threshold = dc_threshold_mv
self._margin = margin_frac
def _one(self, a: int) -> tuple[int, float, float, float]:
theta = compute._theta_deg(self._sras, a, self._ref)
dx, dy = compute.correlate_translation_mm(
self._sras, a, self._ref, self._dc4_mv, self._pivot_mm,
use_mask=self._use_mask, dc_threshold_mv=self._threshold,
margin_frac=self._margin)
return a, theta, dx, dy
def run(self):
try:
n_workers, _budget = compute.plan_angle_level(self._sras)
pool = ThreadPoolExecutor(max_workers=max(1, n_workers))
try:
futures = [pool.submit(self._one, a) for a in self._angles]
for fut in as_completed(futures):
a, theta, dx, dy = fut.result()
self.angle_done.emit(a, theta, dx, dy)
finally:
pool.shutdown(wait=True)
self.finished.emit()
except Exception as exc:
self.error.emit(str(exc))
View File
+197
View File
@@ -0,0 +1,197 @@
#!/usr/bin/env python3
"""Golden-output equivalence harness for the sras-viewer refactor.
Computes a battery of DC / FFT / alignment outputs and prints a stable hash
for each. Run it on the pre-refactor commit to capture a baseline, then again
after the refactor and diff the two reports — every line must match.
Imports work against both the pre-refactor monolith (`sras_viewer`) and the
post-refactor split (`sras_format` + `sras_compute`), so the *same* script
produces both sides of the comparison.
Hashes canonicalise to native little-endian float64 before hashing, so a
deliberate dtype/byte-order change that preserves values does not show up as
a false mismatch.
Usage:
python tools/check_equivalence.py --out baseline.txt
python tools/check_equivalence.py --out after.txt --real /path/to/big.sras
diff baseline.txt after.txt
"""
import argparse
import copy
import hashlib
import sys
from pathlib import Path
import numpy as np
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
# --- Import shim: split modules if present, else the monolith --------------
try:
from sras_format import SrasFile, CH1_IDX, CH3_IDX, CH4_IDX, adc_to_mv
import sras_compute as C
_LAYOUT = "split"
except ImportError:
import sras_viewer as _V
from sras_viewer import SrasFile, CH1_IDX, CH3_IDX, CH4_IDX, adc_to_mv
C = _V
_LAYOUT = "monolith"
compute_dc_image = C.compute_dc_image
compute_rf_image = C.compute_rf_image
compute_alignment = C._compute_angle_alignment
apply_alignment = C.apply_alignment
import tools.make_test_sras as gen # noqa: E402
def h(arr) -> str:
"""Stable hash of an array's *values*, independent of dtype/byte order."""
a = np.ascontiguousarray(np.asarray(arr, dtype=np.float64))
return hashlib.sha256(a.tobytes()).hexdigest()[:16]
def row_slice(sras: SrasFile, angle_idx: int, n_rows: int) -> SrasFile:
"""A shallow view of *sras* restricted to the first *n_rows* rows of
*angle_idx*, so the huge real file can be exercised in seconds."""
view = copy.copy(sras)
n = min(int(n_rows), int(sras.n_rows[angle_idx]))
view.n_rows = np.array(sras.n_rows, copy=True)
view.n_rows[angle_idx] = n
view.data = list(sras.data)
view.data[angle_idx] = sras.data[angle_idx][:n]
view._y_pos_per_angle = list(sras._y_pos_per_angle)
view._y_pos_per_angle[angle_idx] = sras._y_pos_per_angle[angle_idx][:n]
return view
def report(lines: list[str], label: str, value: str):
lines.append(f"{label:<58} {value}")
def check_file(path: Path, lines: list[str], tag: str,
angles: list[int], n_rows: int | None):
sras = SrasFile(str(path))
report(lines, f"[{tag}] version/n_angles",
f"v{sras.version} n={sras.n_angles}")
report(lines, f"[{tag}] geometry",
f"rows={list(map(int, sras.n_rows))} frames={list(map(int, sras.n_frames))}")
report(lines, f"[{tag}] calibration",
" ".join(f"{m:.6g}/{o:.6g}/{z:.6g}" for m, o, z in
zip(sras.ch_ymult_mv, sras.ch_yoff_adc, sras.ch_yzero_mv)))
report(lines, f"[{tag}] angles_deg", h(sras.angles_deg))
if sras.background is not None:
report(lines, f"[{tag}] background", h(sras.background))
for a in angles:
if a >= sras.n_angles:
continue
s = row_slice(sras, a, n_rows) if n_rows else sras
report(lines, f"[{tag}] x_axis_mm(a={a})", h(s.x_axis_mm(a)))
report(lines, f"[{tag}] y_positions_mm(a={a})", h(s.y_positions_mm(a)))
for ch, name in ((CH3_IDX, "CH3"), (CH4_IDX, "CH4")):
raw = compute_dc_image(s, a, ch)
report(lines, f"[{tag}] dc_adc(a={a},{name})", h(raw))
mv = adc_to_mv(raw, s.ch_ymult_mv[ch], s.ch_yoff_adc[ch],
s.ch_yzero_mv[ch])
report(lines, f"[{tag}] dc_mv(a={a},{name})", h(mv))
dc4 = adc_to_mv(compute_dc_image(s, a, CH4_IDX),
s.ch_ymult_mv[CH4_IDX], s.ch_yoff_adc[CH4_IDX],
s.ch_yzero_mv[CH4_IDX])
thresholds = [-1e9, float(np.median(dc4))]
for bg in (False, True):
if bg and s.background is None:
continue
for pad in (1, 2):
n_fft = s.samples_per_frame * pad if pad > 1 else None
for ti, thr in enumerate(thresholds):
img = compute_rf_image(s, a, dc_threshold_mv=thr,
apply_bg_sub=bg, n_fft=n_fft)
report(lines,
f"[{tag}] rf(a={a},bg={int(bg)},pad={pad},thr{ti})",
h(img))
# dc4_mv passthrough must give the identical result
img2 = compute_rf_image(s, a, dc_threshold_mv=thr,
apply_bg_sub=bg, n_fft=n_fft,
dc4_mv=dc4)
same = "SAME" if h(img) == h(img2) else "DIFFER"
report(lines,
f"[{tag}] rf-dc4arg(a={a},bg={int(bg)},pad={pad},thr{ti})",
same)
def check_alignment(path: Path, lines: list[str], tag: str):
sras = SrasFile(str(path))
if sras.n_angles < 2:
return
dc4 = adc_to_mv(compute_dc_image(sras, 0, CH4_IDX),
sras.ch_ymult_mv[CH4_IDX], sras.ch_yoff_adc[CH4_IDX],
sras.ch_yzero_mv[CH4_IDX])
thr = float(np.median(dc4))
res = compute_alignment(sras, 0, thr)
report(lines, f"[{tag}] align canvas_shape", str(res.canvas_shape))
report(lines, f"[{tag}] align canvas_origin",
f"{res.canvas_origin_mm[0]:.9g},{res.canvas_origin_mm[1]:.9g}")
report(lines, f"[{tag}] align canvas_pitch",
f"{res.canvas_dx_mm:.9g},{res.canvas_dy_mm:.9g}")
for a in sorted(res.per_angle):
t = res.per_angle[a]
report(lines, f"[{tag}] align shift_mm(a={a})",
f"{t.shift_mm[0]:.9g},{t.shift_mm[1]:.9g}")
report(lines, f"[{tag}] align rot(a={a})", f"{t.rotation_deg:.9g}")
report(lines, f"[{tag}] align matrix(a={a})", h(t.matrix))
report(lines, f"[{tag}] align offset(a={a})", h(t.offset))
img = compute_dc_image(sras, a, CH4_IDX)
report(lines, f"[{tag}] align resampled(a={a})",
h(apply_alignment(res, a, img)))
def main():
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--out", required=True, help="report file to write")
p.add_argument("--real", help="optional path to a real .sras file")
p.add_argument("--real-rows", type=int, default=2,
help="rows per angle to sample from the real file")
p.add_argument("--real-angles", type=int, default=2,
help="how many angles to sample from the real file")
p.add_argument("--scratch", default=".",
help="directory for generated synthetic files")
args = p.parse_args()
lines = [f"# layout: {_LAYOUT}", f"# numpy: {np.__version__}"]
scratch = Path(args.scratch)
synth = scratch / "equiv_synth.sras"
gen.write(synth, n_angles=4, seed=0, samples_per_frame=64)
check_file(synth, lines, "synth", angles=[0, 1, 2, 3], n_rows=None)
check_alignment(synth, lines, "synth")
# A second synthetic with an odd sample count, to catch off-by-one in
# rfft bin handling and chunk-boundary arithmetic.
synth_odd = scratch / "equiv_synth_odd.sras"
gen.write(synth_odd, n_angles=2, seed=7, samples_per_frame=37)
check_file(synth_odd, lines, "odd", angles=[0, 1], n_rows=None)
if args.real:
real = Path(args.real)
if real.exists():
check_file(real, lines, "real",
angles=list(range(args.real_angles)),
n_rows=args.real_rows)
else:
lines.append(f"# real file not found: {real}")
Path(args.out).write_text("\n".join(lines) + "\n")
print("\n".join(lines))
print(f"\nWrote {args.out} ({len(lines)} lines)")
if __name__ == "__main__":
main()
+191
View File
@@ -0,0 +1,191 @@
#!/usr/bin/env python3
"""Generate small synthetic .sras files for testing.
Writes v6 files (per-angle geometry, ragged waveform blocks) matching
scan_format.md, with deterministic pseudo-random waveform content so a test
can compute expected DC/FFT images independently of the reader under test.
Usage:
python tools/make_test_sras.py out.sras [--angles 3] [--seed 0]
"""
import argparse
import struct
from pathlib import Path
import numpy as np
HDR_FMT_V6 = ">4sBHfffffffIdBB"
GEO_FMT_V6 = ">ffIH"
# Per-angle (n_rows, n_frames) — deliberately different per angle so ragged
# geometry handling is actually exercised.
_GEOMETRY = [(5, 7), (4, 11), (6, 9), (3, 13), (7, 6)]
_SAMPLE_RATE_HZ = 6.25e9
_VELOCITY_MM_S = 20.0
_LASER_FREQ_HZ = 1000.0
_ROW_SPACING_MM = 0.05
def _preamble(ymult_v: float, yoff_adc: float, yzero_v: float) -> bytes:
"""A Tektronix WFMOutpre string in verbose (keyword) form — the reader
pulls YMULT/YOFF/YZERO out of it by name, so the keywords must be
present literally. YMULT/YZERO are in volts, as the scope reports them."""
return (
":WFMOUTPRE:BYT_NR 1;BIT_NR 8;ENCDG BIN;BN_FMT RI;BYT_OR MSB;"
'WFID "Ch1, DC coupling";NR_PT 2500;PT_FMT Y;'
"XINCR 1.6000E-10;XZERO 0.0E0;XUNIT \"s\";"
f"YMULT {ymult_v:.6E};YOFF {yoff_adc:.6E};YZERO {yzero_v:.6E};"
'YUNIT "V"'
).encode("utf-8")
def build(n_angles: int, seed: int, samples_per_frame: int,
geometry: list[tuple[int, int]] | None = None) -> tuple[bytes, dict]:
rng = np.random.default_rng(seed)
src_geom = geometry or _GEOMETRY
geom = [src_geom[a % len(src_geom)] for a in range(n_angles)]
n_ch = 3
bps = 1
angles_deg = np.linspace(0.0, 60.0, n_angles, dtype=np.float32)
# Distinct calibration per channel so a swapped-channel bug is visible.
cal = [
(1.5625e-3, -87.04, 0.0),
(2.0000e-3, -60.00, 1.0e-3),
(2.5000e-3, -40.00, -2.0e-3),
]
out = bytearray()
out += struct.pack(
HDR_FMT_V6, b"SRAS", 6, n_angles,
0.0, 0.0, 1.0, 1.0, _ROW_SPACING_MM,
_VELOCITY_MM_S, _LASER_FREQ_HZ,
samples_per_frame, _SAMPLE_RATE_HZ, bps, n_ch,
)
out += angles_deg.astype(">f4").tobytes()
x_starts = []
for a, (n_rows, n_frames) in enumerate(geom):
x_start = -0.5 + 0.1 * a
x_starts.append(x_start)
out += struct.pack(GEO_FMT_V6, x_start, 1.0, n_frames, n_rows)
y_positions = []
for a, (n_rows, _) in enumerate(geom):
y = (0.2 * a + np.arange(n_rows) * _ROW_SPACING_MM).astype(np.float32)
y_positions.append(y)
out += y.astype(">f4").tobytes()
for ymult_v, yoff, yzero_v in cal:
p = _preamble(ymult_v, yoff, yzero_v)
out += struct.pack(">H", len(p)) + p
background = rng.integers(-8, 9, size=samples_per_frame, dtype=np.int8)
out += struct.pack(">I", samples_per_frame) + background.tobytes()
# Waveform data. CH1 gets a sinusoid at a per-pixel frequency so the FFT
# peak is predictable; CH3/CH4 get per-pixel DC levels so the mean is too.
t = np.arange(samples_per_frame)
waveforms = []
for a, (n_rows, n_frames) in enumerate(geom):
block = np.empty((n_rows, n_ch, n_frames, samples_per_frame), dtype=np.int8)
for r in range(n_rows):
for f in range(n_frames):
bin_idx = 3 + ((a + r + f) % 17)
phase = 2 * np.pi * bin_idx * t / samples_per_frame
block[r, 0, f] = np.clip(
np.round(60 * np.sin(phase)), -128, 127).astype(np.int8)
block[r, 1, f] = np.int8((a * 7 + r * 3 + f) % 100 - 50)
block[r, 2, f] = np.int8((a * 5 + r * 11 + f * 2) % 120 - 60)
waveforms.append(block)
out += block.tobytes()
meta = {
"n_angles": n_angles,
"geometry": geom,
"angles_deg": angles_deg,
"x_starts": x_starts,
"y_positions": y_positions,
"cal": cal,
"background": background,
"waveforms": waveforms,
"samples_per_frame": samples_per_frame,
"sample_rate_hz": _SAMPLE_RATE_HZ,
}
return bytes(out), meta
def write(path: Path, n_angles: int = 3, seed: int = 0,
samples_per_frame: int = 64,
geometry: list[tuple[int, int]] | None = None) -> dict:
payload, meta = build(n_angles, seed, samples_per_frame, geometry)
path.write_bytes(payload)
return meta
HDR_FMT_LEGACY = ">4sBHHffffIIdBB"
def write_legacy(path: Path, version: int = 4, n_angles: int = 2,
n_rows: int = 4, n_frames: int = 10,
samples_per_frame: int = 32, seed: int = 0) -> dict:
"""Write a v2/v3/v4 file: uniform geometry, one flat waveform block.
Used to exercise sras_average.py, which only handles the legacy formats.
"""
rng = np.random.default_rng(seed)
n_ch, bps = 3, 1
out = bytearray()
out += struct.pack(
HDR_FMT_LEGACY, b"SRAS", version, n_angles, n_rows,
-0.5, 1.0, _VELOCITY_MM_S, _LASER_FREQ_HZ,
n_frames, samples_per_frame, _SAMPLE_RATE_HZ, bps, n_ch,
)
angles = np.linspace(0.0, 45.0, n_angles, dtype=np.float32)
out += angles.astype(">f4").tobytes()
y = (np.arange(n_rows) * _ROW_SPACING_MM).astype(np.float32)
out += y.astype(">f4").tobytes()
if version >= 3:
for ymult_v, yoff, yzero_v in ((1.5625e-3, -87.04, 0.0),
(2.0e-3, -60.0, 1.0e-3),
(2.5e-3, -40.0, -2.0e-3))[:n_ch]:
p = _preamble(ymult_v, yoff, yzero_v)
out += struct.pack(">H", len(p)) + p
background = rng.integers(-8, 9, size=samples_per_frame, dtype=np.int8)
if version >= 4:
out += struct.pack(">I", samples_per_frame) + background.tobytes()
data = rng.integers(-100, 101,
size=(n_angles, n_rows, n_ch, n_frames, samples_per_frame),
dtype=np.int8)
out += data.tobytes()
path.write_bytes(bytes(out))
return {"version": version, "n_angles": n_angles, "n_rows": n_rows,
"n_frames": n_frames, "samples_per_frame": samples_per_frame,
"n_channels": n_ch, "data": data, "angles_deg": angles,
"y_positions": y, "background": background}
def main():
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("output")
p.add_argument("--angles", type=int, default=3)
p.add_argument("--seed", type=int, default=0)
p.add_argument("--spf", type=int, default=64, help="samples per frame")
args = p.parse_args()
out = Path(args.output)
meta = write(out, args.angles, args.seed, args.spf)
print(f"Wrote {out} ({out.stat().st_size:,} bytes)")
print(f" angles : {meta['n_angles']}")
print(f" geometry : {meta['geometry']}")
print(f" spf : {meta['samples_per_frame']}")
if __name__ == "__main__":
main()
+485
View File
@@ -0,0 +1,485 @@
#!/usr/bin/env python3
"""Headless GUI test: drives SrasViewerWindow through the real Qt widgets,
signals and worker threads under the offscreen platform plugin.
Covers the interactions a manual smoke test would: load, switch angles and
channels, background DC precompute, lazy FFT compute, threshold and bg-sub
changes, angle alignment, manual angle alignment, aligned view, ROI
draw/move, and CSV export.
Usage: QT_QPA_PLATFORM=offscreen python tools/test_gui.py [file.sras]
"""
import json
import os
import sys
import tempfile
from pathlib import Path
from unittest.mock import patch
os.environ.setdefault("QT_QPA_PLATFORM", "offscreen")
import numpy as np # noqa: E402
from PyQt6.QtCore import QEventLoop, Qt, QTimer # noqa: E402
from PyQt6.QtTest import QTest # noqa: E402
from PyQt6.QtWidgets import QApplication, QMessageBox # noqa: E402
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
import sras_compute as compute # noqa: E402
from sras_format import CH1_IDX, CH3_IDX, CH4_IDX # noqa: E402
from sras_viewer import RoiQuad, SrasViewerWindow, VELOCITY_MODE_IDX # 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 pump(ms: int = 250):
"""Run the event loop for a while so queued signals and worker threads
make progress."""
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()
def main():
app = QApplication(sys.argv)
errors: list[str] = []
tmpdir = Path(tempfile.mkdtemp(prefix="sras_gui_"))
path = Path(sys.argv[1]) if len(sys.argv) > 1 else tmpdir / "gui.sras"
if len(sys.argv) <= 1:
gen.write(path, n_angles=4, seed=11, samples_per_frame=256)
print(f"\nloading {path.name}")
win = SrasViewerWindow()
win.show()
# Capture anything the app reports as an error via the status bar.
win.statusBar().messageChanged.connect(
lambda m: errors.append(m) if m and "error" in m.lower() else None)
win._load_file(str(path))
check("file loaded", wait_until(lambda: win._sras is not None))
s = win._sras
check("parsed as v6", s.version == 6, f"v{s.version}")
check("defaults to CH4", win.combo_channel.currentIndex() == CH4_IDX)
check("image displayed", win._current_image is not None)
check("angle spinbox ranges over all angles",
win.spin_angle.maximum() == s.n_angles - 1)
check("scan info populated",
win._info["Angles"].text() == f"Angles: {s.n_angles}",
win._info["Angles"].text())
print("\nbackground DC precompute (all angles)")
ok = wait_until(lambda: all((a, CH4_IDX) in win._dc_cache
and (a, CH3_IDX) in win._dc_cache
for a in range(s.n_angles)))
check("every angle cached for CH3 and CH4", ok,
f"{len(win._dc_cache)} entries")
check("status label reports completion",
"ready for all angles" in win.lbl_dc_precompute.text(),
win.lbl_dc_precompute.text())
print("\nangle switching (DC, should be served from cache)")
for a in range(s.n_angles):
win.spin_angle.setValue(a)
win._on_view_changed()
pump(60)
expected = win._sras.image_shape(a)
check(f"angle {a} shows its own geometry {expected}",
win._current_image.shape == expected,
str(win._current_image.shape))
check("no compute job needed for cached DC angles",
not win._job_running("compute"))
print("\nchannel switching")
win.spin_angle.setValue(0)
win._on_view_changed()
pump(60)
win.combo_channel.setCurrentIndex(CH3_IDX)
check("CH3 displayed", wait_until(lambda: win._current_ch == CH3_IDX))
win.combo_channel.setCurrentIndex(CH1_IDX)
check("CH1 (FFT) computed", wait_until(
lambda: win._current_ch == CH1_IDX and not win._job_running("compute")))
check("FFT result cached", len(win._fft_cache) > 0, f"{len(win._fft_cache)} keys")
rf_img = win._current_image
check("FFT image is non-degenerate", len(np.unique(rf_img)) > 1,
f"{len(np.unique(rf_img))} unique values")
print("\nvelocity mode (pure post-multiply, no recompute)")
n_fft_before = len(win._fft_cache)
win.combo_channel.setCurrentIndex(VELOCITY_MODE_IDX)
check("velocity displayed", wait_until(
lambda: win._current_ch == VELOCITY_MODE_IDX and not win._job_running("compute")))
grating = win.spin_grating_um.value()
check("velocity == freq x grating",
np.allclose(win._current_image, rf_img * grating, atol=1e-3))
check("velocity reused the cached FFT", len(win._fft_cache) == n_fft_before,
f"{n_fft_before} -> {len(win._fft_cache)}")
check("grating spinbox visible in velocity mode", win.grp_velocity.isVisible())
print("\nthreshold change (genuine cache-key change)")
win.combo_channel.setCurrentIndex(CH1_IDX)
wait_until(lambda: not win._job_running("compute"))
dc4 = win._dc_cache[(0, CH4_IDX)]
win.spin_threshold_mv.setValue(float(np.median(dc4)))
win._on_threshold_changed()
check("recomputed at new threshold", wait_until(
lambda: not win._job_running("compute") and len(win._fft_cache) > n_fft_before))
check("masking zeroed some pixels",
int((win._current_image == 0).sum()) > 0,
f"{int((win._current_image == 0).sum())} of {win._current_image.size}")
print("\nbackground subtraction toggle")
n_before = len(win._fft_cache)
win.chk_bg_sub.setChecked(False)
check("recomputed without bg-sub", wait_until(
lambda: not win._job_running("compute") and len(win._fft_cache) > n_before))
win.chk_bg_sub.setChecked(True)
pump(200)
check("returning to bg-sub was a cache hit (no recompute)",
not win._job_running("compute"))
print("\nROI")
x = s.x_axis_mm(0)
y = s.y_positions_mm(0)
roi = RoiQuad.from_bbox(float(x[1]), float(y[1]),
float(x[-2]), float(y[-2]))
win.image_canvas.set_roi(roi)
pump(120)
check("ROI registered", win.image_canvas.get_roi() is not None)
check("pixel count reported",
"pixels inside" in win.lbl_roi_npix.text()
and win.lbl_roi_npix.text() != "pixels inside: —",
win.lbl_roi_npix.text())
npix = int(win.lbl_roi_npix.text().split(":")[1])
check("ROI pixel count is plausible",
0 < npix <= win._current_image.size, f"{npix}")
check("Export ROI enabled", win.btn_export_roi.isEnabled())
csv_path = tmpdir / "roi.csv"
with patch("sras_viewer.QFileDialog.getSaveFileName",
return_value=(str(csv_path), "")):
win._on_export_roi_csv()
check("ROI CSV written", csv_path.exists())
if csv_path.exists():
body = [l for l in csv_path.read_text().splitlines() if not l.startswith("#")]
check("ROI CSV has header + one line per pixel",
len(body) == npix + 1, f"{len(body)} lines for {npix} pixels")
img_csv = tmpdir / "img.csv"
with patch("sras_viewer.QFileDialog.getSaveFileName",
return_value=(str(img_csv), "")):
win._on_export_csv()
check("image CSV written", img_csv.exists())
if img_csv.exists():
arr = np.loadtxt(img_csv, delimiter=",")
check("image CSV round-trips the displayed image",
arr.shape == win._current_image.shape
and np.allclose(arr, win._current_image, rtol=1e-5, atol=1e-4))
print("\nROI survives angle and channel switches")
win.spin_angle.setValue(1)
win._on_view_changed()
wait_until(lambda: not win._job_running("compute"))
check("ROI still present after angle switch",
win.image_canvas.get_roi() is not None)
win.combo_channel.setCurrentIndex(CH4_IDX)
wait_until(lambda: win._current_ch == CH4_IDX)
check("ROI still present after channel switch",
win.image_canvas.get_roi() is not None)
print("\nangle alignment (Fusion)")
win.spin_angle.setValue(0)
win._on_view_changed()
wait_until(lambda: not win._job_running("compute"))
check("alignment action enabled", win._alignment_act.isEnabled())
win._on_angle_alignment()
check("alignment completed", wait_until(
lambda: win._alignment_result is not None and not win._job_running("align"),
timeout_ms=60000))
if win._alignment_result is not None:
r = win._alignment_result
check("transform for every angle", len(r.per_angle) == s.n_angles)
check("canvas is at least as large as any single angle",
all(r.canvas_shape[0] >= int(s.n_rows[a])
and r.canvas_shape[1] >= int(s.n_frames[a])
for a in range(s.n_angles)), str(r.canvas_shape))
check("reference angle has zero shift",
r.per_angle[r.ref_angle_idx].shift_mm == (0.0, 0.0))
check("Aligned View auto-enabled and checked",
win.chk_aligned_view.isEnabled() and win.chk_aligned_view.isChecked())
pump(200)
check("displayed image is on the alignment canvas",
win.image_canvas._img_shape == r.canvas_shape,
f"{win.image_canvas._img_shape} vs {r.canvas_shape}")
win.chk_aligned_view.setChecked(False)
pump(200)
check("unchecking returns to the raw per-angle grid",
win.image_canvas._img_shape == s.image_shape(0),
str(win.image_canvas._img_shape))
print("\nmanual alignment (Fusion)")
check("manual alignment action enabled", win._manual_align_act.isEnabled())
# --- Alignment pivot is a signal-weighted centroid, not the raw bbox --
# center, and is independent of any DC threshold (so a threshold that
# happens to leave a real angle's binary mask empty can't silently
# degrade the pivot back to the bbox center).
corner_signal = np.zeros(s.image_shape(0), dtype=np.float32)
corner_signal[0, 0] = 1.0 # single spike -> weighted centroid is exact
expected_corner = (float(s.x_axis_mm(0)[0]), float(s.y_positions_mm(0)[0]))
centroid = compute._signal_centroid_mm(s, 0, corner_signal)
check("signal-weighted centroid of a single spike pixel is that pixel exactly",
np.allclose(centroid, expected_corner), f"{centroid} vs {expected_corner}")
bbox_center = compute._bbox_center_mm(s, 0)
check("signal centroid differs from the raw scan-window bbox center",
not np.allclose(centroid, bbox_center),
f"centroid {centroid} vs bbox center {bbox_center}")
# compute_pivot_points_mm should reuse a pre-computed dc4_mv dict rather
# than recomputing from the real DC4 image (which has no such spike and
# would give a different answer if silently recomputed).
reused_pivot = compute.compute_pivot_points_mm(s, dc4_mv={0: corner_signal})[0]
check("compute_pivot_points_mm reuses a pre-computed dc4_mv dict",
np.allclose(reused_pivot, expected_corner))
# A perfectly flat signal carries no information to weight by, so it
# falls back to the bbox center rather than producing a NaN/degenerate
# centroid.
flat_signal = np.full(s.image_shape(0), 5.0, dtype=np.float32)
flat_centroid = compute._signal_centroid_mm(s, 0, flat_signal)
check("a perfectly flat signal falls back to the bbox center",
np.allclose(flat_centroid, bbox_center))
# --- Rotation sign convention: negative of the raw angles_deg delta ----
check("_theta_deg negates the raw angles_deg delta (GR stage's positive "
"angle is the opposite rotational sense from this module's CCW "
"math convention)",
all(np.isclose(compute._theta_deg(s, a, 0),
-(float(s.angles_deg[a]) - float(s.angles_deg[0])))
for a in range(s.n_angles)))
# --- FFT phase correlation recovers a known synthetic pixel shift ------
rng = np.random.default_rng(0)
corr_ref = np.zeros((40, 50), dtype=np.float32)
corr_ref[10:25, 15:35] = 1.0
corr_ref += 0.05 * rng.standard_normal(corr_ref.shape).astype(np.float32)
corr_mov = np.roll(corr_ref, shift=(4, -7), axis=(0, 1))
dr, dc = compute._phase_correlate_shift(corr_ref, corr_mov)
check("phase correlation recovers the shift that aligns mov onto ref",
(dr, dc) == (-4, 7), f"got (dr, dc)={(dr, dc)}")
# --- Open: must NOT seed from the still-live automatic AlignmentResult --
# The automatic result's translation comes from FFT phase correlation --
# the very thing manual mode exists to work around -- so manual mode
# must start from identity (centroids coincide, zero shift) regardless
# of whatever the automatic run last computed. Only a previously *saved
# manual* alignment (sidecar) should ever seed this dialog.
win._on_manual_alignment()
check("dialog opened", win._manual_align_dialog is not None)
dlg = win._manual_align_dialog
check("mask prep needed no background worker (already DC-cached)",
not win._job_running("manual_align_masks"))
check("no manual sidecar yet -> dialog starts at identity, not the "
"automatic result",
all(dlg._angle_params[a] == compute.ManualAngleParams()
for a in range(s.n_angles)))
# --- Reference angle is locked -------------------------------------------
dlg.combo_active_angle.setCurrentIndex(dlg._ref_angle_idx)
pump(30)
before_ref = dlg._angle_params[dlg._ref_angle_idx]
dlg._on_nudge_translate(1, 0, False)
dlg._on_nudge_rotate(1, False)
check("reference angle group disabled", not dlg.grp_manual_adjust.isEnabled())
check("reference angle untouched by nudge attempts",
dlg._angle_params[dlg._ref_angle_idx] == before_ref)
# --- Nudging a real angle (fine + coarse, translate + rotate) -----------
active = 1 if s.n_angles > 1 else 0
dlg.combo_active_angle.setCurrentIndex(active)
pump(30)
before = dlg._angle_params[active].shift_mm
dlg._on_nudge_translate(1, 0, False) # fine +X
fine_step = dlg.spin_step_translate_mm.value()
check("fine translate nudge moved shift_x by exactly one fine step",
abs(dlg._angle_params[active].shift_mm[0] - (before[0] + fine_step)) < 1e-9)
before = dlg._angle_params[active].shift_mm
dlg._on_nudge_translate(0, -1, True) # coarse -Y
coarse_step = fine_step * dlg.spin_step_multiplier.value()
check("coarse translate nudge uses the multiplier",
abs(dlg._angle_params[active].shift_mm[1] - (before[1] - coarse_step)) < 1e-9)
before_rot = dlg._angle_params[active].rotation_deg
dlg._on_nudge_rotate(1, False)
check("rotate nudge changed rotation_deg",
dlg._angle_params[active].rotation_deg != before_rot)
check("preview canvas rebuilt for every angle after a rotation nudge",
len(dlg._preview_layers) == s.n_angles)
# --- Real key-event wiring (proves keyPressEvent -> signal -> slot) -----
before = dlg._angle_params[active].shift_mm
QTest.keyClick(dlg.canvas, Qt.Key.Key_Right)
check("a real Right-arrow key event nudged shift_x",
dlg._angle_params[active].shift_mm[0] > before[0])
# --- Auto De-rotate: rotation only, translation untouched ---------------
shift_before_derotate = dlg._angle_params[active].shift_mm
dlg._on_auto_derotate()
expected_theta = compute._theta_deg(s, active, dlg._ref_angle_idx)
check("auto de-rotate set the known analytic angle",
abs(dlg._angle_params[active].rotation_deg - expected_theta) < 1e-6)
check("auto de-rotate left translation untouched",
dlg._angle_params[active].shift_mm == shift_before_derotate)
check("reference angle stays identity after auto de-rotate",
dlg._angle_params[dlg._ref_angle_idx].rotation_deg == 0.0)
# --- Auto Cross-Correlate: rotation + FFT-correlated shift, backgrounded -
check("cross-correlate action enabled once masks are ready",
dlg.btn_auto_correlate.isEnabled())
dlg._on_auto_correlate()
check("auto cross-correlate completed", wait_until(
lambda: not win._job_running("manual_align_correlate"), timeout_ms=30000))
check("auto cross-correlate set the known analytic angle for every angle",
all(abs(dlg._angle_params[a].rotation_deg
- compute._theta_deg(s, a, dlg._ref_angle_idx)) < 1e-6
for a in range(s.n_angles) if a != dlg._ref_angle_idx))
check("auto cross-correlate reference angle stays identity",
dlg._angle_params[dlg._ref_angle_idx] == compute.ManualAngleParams())
check("auto cross-correlate re-enabled controls when done",
dlg.grp_correlate.isEnabled() and dlg.btn_save.isEnabled())
check("preview canvas rebuilt after cross-correlate",
len(dlg._preview_layers) == s.n_angles)
# The thresholded-mask option should also work end to end.
dlg.combo_correlate_source.setCurrentIndex(1) # thresholded mask
dlg._on_auto_correlate()
check("auto cross-correlate (thresholded-mask option) completed", wait_until(
lambda: not win._job_running("manual_align_correlate"), timeout_ms=30000))
# --- Save -----------------------------------------------------------------
dlg._on_save()
sidecar = compute.sidecar_path(s.path)
check("sidecar file written", sidecar.exists())
raw = json.loads(sidecar.read_text()) if sidecar.exists() else {}
check("sidecar schema_version is current",
raw.get("schema_version") == compute._SIDECAR_SCHEMA_VERSION)
check("sidecar per_angle round-trips the dialog's resolved params",
all(raw.get("per_angle", {}).get(str(a), {}).get("rotation_deg")
== dlg._angle_params[a].rotation_deg for a in range(s.n_angles)))
check("main window's alignment_result replaced by the manual build",
win._alignment_result is not None
and win._alignment_result.per_angle[active].rotation_deg
== dlg._angle_params[active].rotation_deg)
check("Aligned View auto-enabled after Save",
win.chk_aligned_view.isEnabled() and win.chk_aligned_view.isChecked())
# --- An old-schema sidecar (pre-pivot/sign fix) is treated as absent ------
stale = dict(raw)
stale["schema_version"] = compute._SIDECAR_SCHEMA_VERSION - 1
sidecar.write_text(json.dumps(stale))
check("a sidecar with an old schema_version is not loaded",
compute.load_manual_alignment(s) is None)
sidecar.write_text(json.dumps(raw)) # restore for the rest of this section
# --- Clear (with confirmation) --------------------------------------------
with patch("sras_viewer.QMessageBox.question",
return_value=QMessageBox.StandardButton.Yes):
dlg._on_clear()
check("sidecar file deleted", not sidecar.exists())
check("dialog params reset to identity",
all(dlg._angle_params[a] == compute.ManualAngleParams()
for a in range(s.n_angles)))
check("main window alignment_result cleared", win._alignment_result is None)
check("Aligned View disabled after Clear",
not win.chk_aligned_view.isEnabled() and not win.chk_aligned_view.isChecked())
dlg.close()
pump(150)
check("dialog reference released on close", win._manual_align_dialog is None)
# --- Sidecar auto-restore on next load ------------------------------------
win._on_manual_alignment()
dlg = win._manual_align_dialog
dlg.combo_active_angle.setCurrentIndex(active)
pump(30)
dlg._on_auto_derotate()
dlg._on_nudge_translate(1, 1, True)
saved_rotation = dlg._angle_params[active].rotation_deg
saved_shift = dlg._angle_params[active].shift_mm
dlg._on_save()
dlg.close()
pump(150)
old_sras_id = id(win._sras)
win._load_file(str(path)) # reload the same file fresh
check("file reloaded", wait_until(
lambda: win._sras is not None and id(win._sras) != old_sras_id))
s = win._sras
check("manual dialog force-closed by a reload", win._manual_align_dialog is None)
check("reload restores the saved manual alignment automatically",
win._alignment_result is not None)
if win._alignment_result is not None:
check("restored rotation matches what was saved",
abs(win._alignment_result.per_angle[active].rotation_deg
- saved_rotation) < 1e-9)
check("restored shift matches what was saved",
win._alignment_result.per_angle[active].shift_mm == saved_shift)
check("Aligned View auto-checked after restoring a saved alignment",
win.chk_aligned_view.isChecked())
print("\npixel inspector")
win.chk_aligned_view.setChecked(False)
pump(100)
win._on_pixel_clicked(0, 0)
pump(150)
check("waveform hint hidden after a click", win.lbl_wave_hint.isHidden())
win.combo_channel.setCurrentIndex(CH1_IDX)
wait_until(lambda: not win._job_running("compute"))
win._on_pixel_clicked(1, 1)
pump(150)
check("RF waveform panel rendered",
len(win.wave_canvas.ax_wave.lines) > 0,
f"{len(win.wave_canvas.ax_wave.lines)} lines")
print("\nshutdown")
win.close()
pump(400)
check("all background jobs released", len(win._jobs) == 0,
f"{list(win._jobs)}")
print()
unexpected = [e for e in errors if e]
if unexpected:
print(f"status-bar errors seen: {unexpected}")
_failures.append("status-bar errors")
if _failures:
print(f"{len(_failures)} FAILURE(S): " + ", ".join(_failures))
return 1
print("All GUI checks passed.")
return 0
if __name__ == "__main__":
sys.exit(main())
+358
View File
@@ -0,0 +1,358 @@
#!/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"]
expect0 = src_data[0][:, :, 0:4, :].astype(np.float32).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"
subprocess.run([sys.executable, str(repo / "sras_average.py"),
str(src), str(dst2), "--n", "5"],
capture_output=True, text=True, cwd=repo)
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"
subprocess.run([sys.executable, str(repo / "sras_average.py"),
str(src), str(dst3), "--n", "5", "--discard-remainder"],
capture_output=True, text=True, cwd=repo)
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()