An SRAS file stores the raw RF waveforms captured during a Surface-acoustic-wave Resonance And Spectroscopy (SRAS) scan, along with the scan geometry and scope calibration metadata needed to interpret them.
-
A scan consists of one or more angles (rotation positions of the sample), each containing a 2-D raster of rows × frames. At every pixel, n_channels waveforms of samples_per_frame ADC counts are stored. Channel order is fixed:
-
-
-
-
Index
-
Hardware channel
-
Signal
-
-
-
-
-
0
-
CH1
-
RF acoustic packet (AC-coupled)
-
-
-
1
-
CH3
-
Bias A — DC mean used for masking
-
-
-
2
-
CH4
-
Bias B — DC mean used for masking
-
-
-
-
-
Type notation
-
-
-
-
Symbol
-
C type
-
Size
-
Notes
-
-
-
-
-
u8
-
uint8_t
-
1 byte
-
unsigned
-
-
-
u16
-
uint16_t
-
2 bytes
-
big-endian
-
-
-
u32
-
uint32_t
-
4 bytes
-
big-endian
-
-
-
i8
-
int8_t
-
1 byte
-
signed, used for ADC samples when bytes_per_sample == 1
-
-
-
i16
-
int16_t
-
2 bytes
-
big-endian signed, used when bytes_per_sample == 2
-
-
-
f32
-
float
-
4 bytes
-
big-endian IEEE 754 single
-
-
-
f64
-
double
-
8 bytes
-
big-endian IEEE 754 double
-
-
-
char[N]
-
—
-
N bytes
-
raw bytes, no null terminator unless noted
-
-
-
utf8[N]
-
—
-
N bytes
-
UTF-8 string, length-prefixed (see preamble section)
-
-
-
-
-
Version history
-
-
-
-
Version
-
Added
-
-
-
-
-
2
-
Baseline: fixed header, angle table, row table, raw waveform data. No scope calibration (fallback constants used by readers).
Background waveform section: one CH1 reference shot subtracted from each CH1 frame before FFT.
-
-
-
5
-
(this document) Version byte incremented to 5. n_frames_hdr is now the actual acquired frame count (authoritative). PREC section appended after waveform data with precomputed FFT-peak and DC images for instant re-display.
-
-
-
-
-
v2 note: Version 1 is not defined; version 2 is the lowest observed in the field.
-
-
-
File structure
-
┌─────────────────────────────────────────────┐
-│ 1. Fixed header (43 bytes) │ all versions
-├─────────────────────────────────────────────┤
-│ 2. Angle table (n_angles × 4 bytes)│ all versions
-├─────────────────────────────────────────────┤
-│ 3. Row position table (n_rows × 4 bytes)│ all versions
-├─────────────────────────────────────────────┤
-│ 4. Channel preambles (variable) │ v3+
-├─────────────────────────────────────────────┤
-│ 5. Background waveform (variable) │ v4+
-├─────────────────────────────────────────────┤
-│ 6. Waveform data (variable) │ all versions
-├─────────────────────────────────────────────┤
-│ 7. PREC section (variable) │ v5 only
-└─────────────────────────────────────────────┘
-
-
-
1. Fixed header (43 bytes, all versions)
-
-
-
-
Offset
-
Size
-
Type
-
Field
-
Description
-
-
-
-
-
0
-
4
-
char[4]
-
magic
-
SRAS (ASCII, no null terminator). Reject file if this does not match.
-
-
-
4
-
1
-
u8
-
version
-
Format version. This document describes version 5.
-
-
-
5
-
2
-
u16
-
n_angles
-
Number of scan angles (rotation positions). ≥ 1.
-
-
-
7
-
2
-
u16
-
n_rows
-
Number of scan rows (Y positions). ≥ 1.
-
-
-
9
-
4
-
f32
-
x_start_mm
-
X position of the first frame in the first row, in millimetres.
-
-
-
13
-
4
-
f32
-
x_delta_mm
-
Pre-computed pixel pitch in mm (velocity_mm_s / laser_freq_hz). Provided for convenience; readers should prefer the derived value.
-
-
-
17
-
4
-
f32
-
velocity_mm_s
-
Scanner stage velocity, mm s⁻¹. Used together with laser_freq_hz to compute pixel pitch.
-
-
-
21
-
4
-
f32
-
laser_freq_hz
-
Laser repetition rate, Hz.
-
-
-
25
-
4
-
u32
-
n_frames_hdr
-
v2–v4: the configured frame count written before acquisition; may exceed actual frames acquired (use file-size arithmetic to obtain the true count). v5: the actual acquired frame count — authoritative; readers must not re-derive it from file size.
YOFF = −87.04 ADC counts (scope position = −2.72 div)
-
YZERO = 0 mV
-
-
-
5. Background waveform (v4+)
-
u32 n_bg_samples — number of i8 ADC samples that follow
-i8[n_bg] background — one representative CH1 background shot
-
-
The background waveform has the same samples_per_frame length as a normal CH1 waveform. It is subtracted from each CH1 waveform before FFT processing when background subtraction is enabled. When n_bg_samples == 0 the section is present but empty.
-
-
6. Waveform data (all versions)
-
Begins immediately after the fixed header (v2), preambles (v3), or background waveform (v4+). The waveform data is a flat, contiguous array with the following logical shape, stored in row-major (C) order:
v5: use n_frames_hdr directly; do not use file-size arithmetic.
-
v2–v4:n_frames = floor((file_bytes_after_header_sections) / (bytes_per_sample × n_angles × n_rows × n_channels × samples_per_frame)). Any remainder bytes are a partial trailing row and are discarded.
-
-
-
7. PREC section (v5)
-
The PREC section is appended immediately after the waveform data and is present if and only if version == 5 and the file size exceeds waveform_end_offset.
where data_offset is the file offset of the first waveform byte (the byte immediately after the background waveform, or after the angle/row tables for v2 files).
-
PREC header (8 bytes)
-
-
-
-
Offset (relative)
-
Size
-
Type
-
Field
-
Description
-
-
-
-
-
0
-
4
-
char[4]
-
prec_magic
-
PREC (ASCII). Absent or wrong magic → ignore section.
-
-
-
4
-
1
-
u8
-
prec_version
-
PREC format version. Currently 1.
-
-
-
5
-
1
-
u8
-
flags
-
Bitmask (see below).
-
-
-
6
-
2
-
u16
-
n_stored
-
Number of angle entries that follow. 0 ≤ n_stored ≤ n_angles.
-
-
-
-
Flags byte
-
-
-
-
Bit
-
Mask
-
Meaning when set
-
-
-
-
-
0
-
0x01
-
bg_sub_applied — background waveform was subtracted from CH1 before the FFT when these images were computed.
-
-
-
1–7
-
—
-
Reserved, must be zero on write; readers must ignore.
-
-
-
-
PREC angle entries
-
Repeated n_stored times, in arbitrary angle-index order:
-
for each stored angle:
- u16 angle_idx — index into the angle table (0-based)
- f32[n_rows×n_frames] peak_freq_mhz — CH1 FFT peak frequency, MHz, row-major
- f32[n_rows×n_frames] dc4_mv — CH4 waveform mean, mV, row-major
- f32[n_rows×n_frames] dc3_mv — CH3 waveform mean, mV, row-major
-
-
All image arrays are f32 big-endian, stored in row-major order: element [r][f] is at offset (r × n_frames + f) × 4 bytes within the array.
-
peak_freq_mhz is computed without any DC-threshold masking (i.e. the FFT is run on every pixel unconditionally). Readers apply the dc4_mv threshold at display time:
-
pixel is valid ⟺ dc4_mv[r][f] ≥ threshold_mv
-display_value = peak_freq_mhz[r][f] if valid, else 0
-
-
dc4_mv / dc3_mv are the mean of all ADC samples in the respective channel waveform, converted to millivolts using the channel calibration:
-
dc_mv = (adc_mean − YOFF) × YMULT + YZERO
-
-
When readers must bypass the PREC fast path
-
Readers must fall back to real-time FFT computation (ignoring stored peak_freq_mhz) when any of the following are true:
-
-
Time-domain gating is active (zeroing samples outside a time window changes peak frequency).
-
Zero-padding (n_fft ≠ samples_per_frame) is requested (changes bin spacing).
-
The reader's background-subtraction setting does not match flags.bg_sub_applied.
-
-
-
Derived quantities
-
pixel_pitch_mm = velocity_mm_s / laser_freq_hz
-
-x_axis_mm[f] = x_start_mm + f × pixel_pitch_mm (f = 0 … n_frames−1)
-
-time_axis_ns[s] = s / sample_rate_hz × 1e9 (s = 0 … spf−1)
-
-freq_axis_mhz[k] = k × sample_rate_hz / (n_fft × 1e6) (k = 0 … n_fft/2)
- where n_fft = samples_per_frame unless zero-padding is active
-
-velocity_ms[r][f] = peak_freq_mhz[r][f] × grating_um (grating_um user-supplied)
-
A v4 reader that only accepts versions {2, 3, 4} will reject a v5 file with an "unsupported version" error. This is intentional: a v4 reader would derive n_frames from the file size, incorrectly including the PREC bytes in the sample count, producing a silently wrong reshape.
-
Producing v5 files
-
v5 files are produced by the SRAS viewer's "Pre-process and Save as v5" action. The procedure is:
-
-
Copy the source file (any version) verbatim.
-
Set version = 5 at byte offset 4.
-
Set n_frames_hdr at byte offset 25 to the actual acquired frame count.
-
Truncate the copy to data_offset + waveform_bytes (removes any pre-existing stale PREC tail).
-
Compute peak_freq_mhz, dc4_mv, and dc3_mv for every angle using chunked FFT.
-
Append the PREC section.
-
-
Partially-written PREC sections
-
If n_stored < n_angles (e.g. pre-processing was interrupted), the file is still valid. Readers use stored images for the angles present in the PREC section and fall back to real-time FFT for the remainder. Readers must check angle_idx bounds on each entry and stop parsing on an out-of-range value.
-
-
-
-
-
\ No newline at end of file
diff --git a/SRAS_FORMAT.md b/SRAS_FORMAT.md
deleted file mode 100644
index 4bc0346..0000000
--- a/SRAS_FORMAT.md
+++ /dev/null
@@ -1,379 +0,0 @@
-# SRAS File Format Specification
-
-**Format family:** `.sras`
-**Byte order:** Big-endian (network byte order) throughout, unless noted.
-**Version history:** v2 (baseline), v3 (scope calibration), v4 (background waveform), v5 (precomputed images + guaranteed frame count).
-
----
-
-## Table of Contents
-
-1. [Overview](#overview)
-2. [Type notation](#type-notation)
-3. [Version history](#version-history)
-4. [File structure](#file-structure)
- - [Fixed header (all versions)](#1-fixed-header-43-bytes-all-versions)
- - [Angle table (all versions)](#2-angle-table-all-versions)
- - [Row position table (all versions)](#3-row-position-table-all-versions)
- - [Channel preambles (v3+)](#4-channel-preambles-v3)
- - [Background waveform (v4+)](#5-background-waveform-v4)
- - [Waveform data (all versions)](#6-waveform-data-all-versions)
- - [PREC section (v5)](#7-prec-section-v5)
-5. [Derived quantities](#derived-quantities)
-6. [ADC calibration](#adc-calibration)
-7. [Waveform data layout detail](#waveform-data-layout-detail)
-8. [Size reference](#size-reference)
-9. [Compatibility notes](#compatibility-notes)
-
----
-
-## Overview
-
-An SRAS file stores the raw RF waveforms captured during a Surface-acoustic-wave Resonance And Spectroscopy (SRAS) scan, along with the scan geometry and scope calibration metadata needed to interpret them.
-
-A scan consists of one or more **angles** (rotation positions of the sample), each containing a 2-D raster of **rows** × **frames**. At every pixel, `n_channels` waveforms of `samples_per_frame` ADC counts are stored. Channel order is fixed:
-
-| Index | Hardware channel | Signal |
-|-------|-----------------|--------|
-| 0 | CH1 | RF acoustic packet (AC-coupled) |
-| 1 | CH3 | Bias A — DC mean used for masking |
-| 2 | CH4 | Bias B — DC mean used for masking |
-
----
-
-## Type notation
-
-| Symbol | C type | Size | Notes |
-|--------|--------|------|-------|
-| `u8` | `uint8_t` | 1 byte | unsigned |
-| `u16` | `uint16_t` | 2 bytes | big-endian |
-| `u32` | `uint32_t` | 4 bytes | big-endian |
-| `i8` | `int8_t` | 1 byte | signed, used for ADC samples when `bytes_per_sample == 1` |
-| `i16` | `int16_t` | 2 bytes | big-endian signed, used when `bytes_per_sample == 2` |
-| `f32` | `float` | 4 bytes | big-endian IEEE 754 single |
-| `f64` | `double` | 8 bytes | big-endian IEEE 754 double |
-| `char[N]` | — | N bytes | raw bytes, no null terminator unless noted |
-| `utf8[N]` | — | N bytes | UTF-8 string, length-prefixed (see preamble section) |
-
----
-
-## Version history
-
-| Version | Added |
-|---------|-------|
-| 2 | Baseline: fixed header, angle table, row table, raw waveform data. No scope calibration (fallback constants used by readers). |
-| 3 | Per-channel Tektronix WFMOutpre preamble strings carrying YMULT / YOFF / YZERO calibration. |
-| 4 | Background waveform section: one CH1 reference shot subtracted from each CH1 frame before FFT. |
-| 5 | **(this document)** Version byte incremented to 5. `n_frames_hdr` is now the *actual* acquired frame count (authoritative). PREC section appended after waveform data with precomputed FFT-peak and DC images for instant re-display. |
-
-> **v2 note:** Version 1 is not defined; version 2 is the lowest observed in the field.
-
----
-
-## File structure
-
-```
-┌─────────────────────────────────────────────┐
-│ 1. Fixed header (43 bytes) │ all versions
-├─────────────────────────────────────────────┤
-│ 2. Angle table (n_angles × 4 bytes)│ all versions
-├─────────────────────────────────────────────┤
-│ 3. Row position table (n_rows × 4 bytes)│ all versions
-├─────────────────────────────────────────────┤
-│ 4. Channel preambles (variable) │ v3+
-├─────────────────────────────────────────────┤
-│ 5. Background waveform (variable) │ v4+
-├─────────────────────────────────────────────┤
-│ 6. Waveform data (variable) │ all versions
-├─────────────────────────────────────────────┤
-│ 7. PREC section (variable) │ v5 only
-└─────────────────────────────────────────────┘
-```
-
----
-
-### 1. Fixed header (43 bytes, all versions)
-
-| Offset | Size | Type | Field | Description |
-|--------|------|------|-------|-------------|
-| 0 | 4 | `char[4]` | `magic` | `SRAS` (ASCII, no null terminator). Reject file if this does not match. |
-| 4 | 1 | `u8` | `version` | Format version. This document describes version **5**. |
-| 5 | 2 | `u16` | `n_angles` | Number of scan angles (rotation positions). ≥ 1. |
-| 7 | 2 | `u16` | `n_rows` | Number of scan rows (Y positions). ≥ 1. |
-| 9 | 4 | `f32` | `x_start_mm` | X position of the first frame in the first row, in millimetres. |
-| 13 | 4 | `f32` | `x_delta_mm` | Pre-computed pixel pitch in mm (`velocity_mm_s / laser_freq_hz`). Provided for convenience; readers should prefer the derived value. |
-| 17 | 4 | `f32` | `velocity_mm_s` | Scanner stage velocity, mm s⁻¹. Used together with `laser_freq_hz` to compute pixel pitch. |
-| 21 | 4 | `f32` | `laser_freq_hz` | Laser repetition rate, Hz. |
-| 25 | 4 | `u32` | `n_frames_hdr` | **v2–v4:** the *configured* frame count written before acquisition; may exceed actual frames acquired (use file-size arithmetic to obtain the true count). **v5:** the *actual* acquired frame count — authoritative; readers must not re-derive it from file size. |
-| 29 | 4 | `u32` | `samples_per_frame` | ADC samples per waveform (`spf`). |
-| 33 | 8 | `f64` | `sample_rate_hz` | Oscilloscope sample rate, Hz (e.g. 5 × 10⁹ for 5 GS/s). |
-| 41 | 1 | `u8` | `bytes_per_sample` | ADC word size: `1` → `i8`, `2` → `i16` (big-endian). |
-| 42 | 1 | `u8` | `n_channels` | Number of channels per frame. Currently always `3`. |
-
----
-
-### 2. Angle table (all versions)
-
-Immediately follows the fixed header.
-
-```
-n_angles × f32 — scan angle in degrees
-```
-
-Each entry is a big-endian `f32` giving the sample rotation angle in degrees at which that angle index was acquired.
-
----
-
-### 3. Row position table (all versions)
-
-Immediately follows the angle table.
-
-```
-n_rows × f32 — Y position of each row, in millimetres
-```
-
----
-
-### 4. Channel preambles (v3+)
-
-One entry per channel, in channel-index order (CH1 first).
-
-```
-for each channel:
- u16 preamble_length — byte count of the UTF-8 string that follows
- utf8[N] preamble — Tektronix WFMOutpre string
-```
-
-The preamble is the oscilloscope's `WFMOutpre` response string. Readers extract the following keys (case-insensitive, space-separated value):
-
-| Key | Stored unit | Conversion to mV |
-|-----|-------------|-----------------|
-| `YMULT` | V count⁻¹ | multiply by 1000 |
-| `YOFF` | ADC counts | used directly |
-| `YZERO` | V | multiply by 1000 |
-
-**v2 fallback:** when preambles are absent, readers use:
-- `YMULT` = 1.5625 mV count⁻¹ (50 mV/div, 8 div, 8-bit ADC)
-- `YOFF` = −87.04 ADC counts (scope position = −2.72 div)
-- `YZERO` = 0 mV
-
----
-
-### 5. Background waveform (v4+)
-
-```
-u32 n_bg_samples — number of i8 ADC samples that follow
-i8[n_bg] background — one representative CH1 background shot
-```
-
-The background waveform has the same `samples_per_frame` length as a normal CH1 waveform. It is subtracted from each CH1 waveform before FFT processing when background subtraction is enabled. When `n_bg_samples == 0` the section is present but empty.
-
----
-
-### 6. Waveform data (all versions)
-
-Begins immediately after the fixed header (v2), preambles (v3), or background waveform (v4+). The waveform data is a flat, contiguous array with the following logical shape, stored in row-major (C) order:
-
-```
-waveform_data[n_angles][n_rows][n_channels][n_frames][samples_per_frame]
-```
-
-Each element is a signed ADC count of size `bytes_per_sample`:
-- `bytes_per_sample == 1` → `i8`
-- `bytes_per_sample == 2` → `i16` big-endian
-
-**Total byte count:**
-
-```
-waveform_bytes = n_angles × n_rows × n_channels × n_frames × samples_per_frame × bytes_per_sample
-```
-
-#### Index semantics
-
-| Dimension | Range | Meaning |
-|-----------|-------|---------|
-| `[a]` | 0 … n_angles−1 | Scan angle (rotation position) |
-| `[r]` | 0 … n_rows−1 | Row (Y position); row 0 is the first acquired |
-| `[c]` | 0 … n_channels−1 | Channel (0=CH1 RF, 1=CH3 Bias A, 2=CH4 Bias B) |
-| `[f]` | 0 … n_frames−1 | Frame (X position) within the row |
-| `[s]` | 0 … spf−1 | Sample index within the waveform |
-
-#### Frame-count determination
-
-- **v5:** use `n_frames_hdr` directly; do not use file-size arithmetic.
-- **v2–v4:** `n_frames = floor((file_bytes_after_header_sections) / (bytes_per_sample × n_angles × n_rows × n_channels × samples_per_frame))`. Any remainder bytes are a partial trailing row and are discarded.
-
----
-
-### 7. PREC section (v5)
-
-The PREC section is appended immediately after the waveform data and is present if and only if `version == 5` and the file size exceeds `waveform_end_offset`.
-
-```
-waveform_end_offset = data_offset + waveform_bytes
-```
-
-where `data_offset` is the file offset of the first waveform byte (the byte immediately after the background waveform, or after the angle/row tables for v2 files).
-
-#### PREC header (8 bytes)
-
-| Offset (relative) | Size | Type | Field | Description |
-|-------------------|------|------|-------|-------------|
-| 0 | 4 | `char[4]` | `prec_magic` | `PREC` (ASCII). Absent or wrong magic → ignore section. |
-| 4 | 1 | `u8` | `prec_version` | PREC format version. Currently `1`. |
-| 5 | 1 | `u8` | `flags` | Bitmask (see below). |
-| 6 | 2 | `u16` | `n_stored` | Number of angle entries that follow. 0 ≤ `n_stored` ≤ `n_angles`. |
-
-##### Flags byte
-
-| Bit | Mask | Meaning when set |
-|-----|------|-----------------|
-| 0 | `0x01` | `bg_sub_applied` — background waveform was subtracted from CH1 before the FFT when these images were computed. |
-| 1–7 | — | Reserved, must be zero on write; readers must ignore. |
-
-#### PREC angle entries
-
-Repeated `n_stored` times, in arbitrary angle-index order:
-
-```
-for each stored angle:
- u16 angle_idx — index into the angle table (0-based)
- f32[n_rows×n_frames] peak_freq_mhz — CH1 FFT peak frequency, MHz, row-major
- f32[n_rows×n_frames] dc4_mv — CH4 waveform mean, mV, row-major
- f32[n_rows×n_frames] dc3_mv — CH3 waveform mean, mV, row-major
-```
-
-All image arrays are `f32` big-endian, stored in row-major order: element `[r][f]` is at offset `(r × n_frames + f) × 4` bytes within the array.
-
-**`peak_freq_mhz`** is computed without any DC-threshold masking (i.e. the FFT is run on every pixel unconditionally). Readers apply the `dc4_mv` threshold at display time:
-
-```
-pixel is valid ⟺ dc4_mv[r][f] ≥ threshold_mv
-display_value = peak_freq_mhz[r][f] if valid, else 0
-```
-
-**`dc4_mv` / `dc3_mv`** are the mean of all ADC samples in the respective channel waveform, converted to millivolts using the channel calibration:
-
-```
-dc_mv = (adc_mean − YOFF) × YMULT + YZERO
-```
-
-#### When readers must bypass the PREC fast path
-
-Readers must fall back to real-time FFT computation (ignoring stored `peak_freq_mhz`) when any of the following are true:
-
-- Time-domain gating is active (zeroing samples outside a time window changes peak frequency).
-- Zero-padding (`n_fft ≠ samples_per_frame`) is requested (changes bin spacing).
-- The reader's background-subtraction setting does not match `flags.bg_sub_applied`.
-
----
-
-## Derived quantities
-
-```
-pixel_pitch_mm = velocity_mm_s / laser_freq_hz
-
-x_axis_mm[f] = x_start_mm + f × pixel_pitch_mm (f = 0 … n_frames−1)
-
-time_axis_ns[s] = s / sample_rate_hz × 1e9 (s = 0 … spf−1)
-
-freq_axis_mhz[k] = k × sample_rate_hz / (n_fft × 1e6) (k = 0 … n_fft/2)
- where n_fft = samples_per_frame unless zero-padding is active
-
-velocity_ms[r][f] = peak_freq_mhz[r][f] × grating_um (grating_um user-supplied)
-```
-
----
-
-## ADC calibration
-
-Convert raw ADC counts to millivolts:
-
-```
-voltage_mv = (adc_count − YOFF) × YMULT_mv + YZERO_mv
-```
-
-Invert (mV → ADC count):
-
-```
-adc_count = (voltage_mv − YZERO_mv) / YMULT_mv + YOFF
-```
-
-where `YMULT_mv` is YMULT in mV count⁻¹ (= scope YMULT in V count⁻¹ × 1000).
-
----
-
-## Waveform data layout detail
-
-For a scan with `n_angles=2`, `n_rows=3`, `n_channels=3`, `n_frames=4`, `spf=5` the layout is:
-
-```
-angle 0
- row 0
- CH1: [s0 s1 s2 s3 s4] [s0 s1 s2 s3 s4] [s0 s1 s2 s3 s4] [s0 s1 s2 s3 s4]
- frame 0 frame 1 frame 2 frame 3
- CH3: …(same layout)…
- CH4: …(same layout)…
- row 1
- …
- row 2
- …
-angle 1
- …
-```
-
-The flat byte offset of sample `s` of frame `f`, channel `c`, row `r`, angle `a` is:
-
-```
-offset = data_offset
- + (a × n_rows × n_channels × n_frames × spf
- + r × n_channels × n_frames × spf
- + c × n_frames × spf
- + f × spf
- + s)
- × bytes_per_sample
-```
-
----
-
-## Size reference
-
-Approximate sizes for representative scans (`bytes_per_sample = 1`, `n_channels = 3`).
-
-| n_angles | n_rows | n_frames | spf | Waveform data | PREC section |
-|----------|--------|----------|-----|---------------|-------------|
-| 1 | 500 | 500 | 400 | 300 MB | 12 MB |
-| 4 | 500 | 500 | 400 | 1.2 GB | 48 MB |
-| 1 | 2000 | 2000 | 400 | 4.8 GB | 48 MB |
-| 4 | 2000 | 2000 | 400 | 19.2 GB | 192 MB |
-| 8 | 2000 | 2000 | 400 | 38.4 GB | 384 MB |
-| 16 | 2000 | 2000 | 400 | 76.8 GB | 768 MB |
-
-**PREC section size formula:**
-
-```
-prec_bytes = 8 + n_stored × (2 + 3 × n_rows × n_frames × 4)
-```
-
----
-
-## Compatibility notes
-
-### Reading v5 files with a v4 reader
-
-A v4 reader that only accepts versions `{2, 3, 4}` will reject a v5 file with an "unsupported version" error. This is intentional: a v4 reader would derive `n_frames` from the file size, incorrectly including the PREC bytes in the sample count, producing a silently wrong reshape.
-
-### Producing v5 files
-
-v5 files are produced by the SRAS viewer's **"Pre-process and Save as v5"** action. The procedure is:
-
-1. Copy the source file (any version) verbatim.
-2. Set `version = 5` at byte offset 4.
-3. Set `n_frames_hdr` at byte offset 25 to the actual acquired frame count.
-4. Truncate the copy to `data_offset + waveform_bytes` (removes any pre-existing stale PREC tail).
-5. Compute `peak_freq_mhz`, `dc4_mv`, and `dc3_mv` for every angle using chunked FFT.
-6. Append the PREC section.
-
-### Partially-written PREC sections
-
-If `n_stored < n_angles` (e.g. pre-processing was interrupted), the file is still valid. Readers use stored images for the angles present in the PREC section and fall back to real-time FFT for the remainder. Readers must check `angle_idx` bounds on each entry and stop parsing on an out-of-range value.
diff --git a/docs/design.md b/docs/design.md
new file mode 100644
index 0000000..7b0c583
--- /dev/null
+++ b/docs/design.md
@@ -0,0 +1,398 @@
+# sras-viewer design notes
+
+Rationale that outgrew code comments. Each section is referenced by a short
+pointer comment at the relevant definition, so the code stays scannable and
+the reasoning stays findable.
+
+## Memory budget and row chunking (`sras_compute.py`)
+
+DC images are computed over row chunks so the float32 working buffers for one
+chunk stay under a memory budget. A fixed row count (the original design)
+works fine for small legacy scans but is catastrophic for a v6 scan with a
+large per-angle frame/sample count — e.g. a 7500-frame × 2500-sample angle
+needs ~2.4 GB for a single 32-row chunk.
+
+With chunks running concurrently the budget has to cover *all* live chunks at
+once. On a large scan `chunk_rows` is already clamped to its floor of one row
+(one row alone is ~75 MB of float32 at 7507×2500), so shrinking the per-chunk
+size cannot buy more concurrency — the worker count must be derived from the
+budget instead: `_plan_chunks` picks the worker count *first* and sizes the
+chunk to it. Sizing the chunk first is the trap: a single chunk would always
+consume the whole budget and leave room for exactly one worker, precisely on
+the large scans that need concurrency most.
+
+The 1024 MB default (`SRAS_MEM_BUDGET_MB`) is the measured knee on a 16-core
+machine against a 7507-frame × 2500-sample angle: 512 MB left ~20% of the
+speedup on the table, and 1536+ MB cost ~0.4 GB more resident memory for no
+further gain.
+
+A caller that itself runs several computations concurrently (angle-level
+parallelism, `plan_angle_level`) must pass *both* `max_workers=1` and its
+share of the budget. Capping the workers alone is not enough: the chunk would
+still be sized against the whole budget, and N concurrent callers would each
+allocate all of it.
+
+## FFT peak search: block-parallel zoom refinement (`sras_compute.py`)
+
+The displayed RF value per pixel is the argmax of the zero-padded power
+spectrum of that pixel's CH1 waveform. At the pad factor of 40 needed for
+mapping resolution, materialising padded spectra is hopeless: ~9 GB per scan
+row, which is what used to collapse the old row-chunk planner to one worker
+and make synthesis single-threaded.
+
+`_peak_bins_zoom` never materialises the padded spectrum:
+
+1. a coarse rfft at `next_fast_len(2*spf)` — 2× oversampled, so the padded
+ power spectrum (a trig polynomial of degree spf−1) cannot hide its global
+ max between coarse samples;
+2. every coarse bin within `_ZOOM_CAND_RATIO` (0.7) of its row's coarse max
+ becomes a refinement candidate. Quarter-natural-bin scalloping at the 2×
+ grid can understate a peak's power by at most ~19%, so 0.7 keeps a wide
+ margin. The DC-adjacent window is always refined too: the coarse DC bin
+ is zeroed for suppression, which would otherwise blind the scan to fine
+ bins closer to DC than the first coarse sample (where the leakage skirt
+ of an un-subtracted offset peaks);
+3. each candidate window (±`_ZOOM_HALFWIDTH` = 0.75 coarse spacings; every
+ fine bin lies within 0.5 spacings of its nearest coarse bin) is evaluated
+ on the exact `n_fft` grid by one small complex gemm, with np.argmax's
+ lowest-bin tie-break preserved across windows.
+
+The selected bin is bit-identical to the full padded argmax — enforced by
+`tests/test_compute.py::test_zoom_identity`, a fuzz test over adversarial
+spectra, and the golden-hash harness (`tools/check_equivalence.py`), whose
+baseline was captured on the old full-padded path.
+
+Work fans out over a persistent thread pool in `_FFT_BLOCK` = 512-waveform
+tasks: smaller blocks serialise on GIL-held numpy dispatch, larger ones lose
+cache residency and task granularity (measured on a 16-core machine, where
+this path runs ~35× faster than the old serial padded transform at pad 40).
+pyFFTW runs through per-thread `builders` plans (FFTW_MEASURE, wisdom
+persisted under `~/.cache/sras-viewer/`), and `threadpoolctl` clamps BLAS to
+one thread under the pool so the refinement gemm cannot oversubscribe.
+`compute_rf_image(exact=True)` (or `SRAS_FFT_EXACT=1`) keeps the reference
+full-padded path for audits.
+
+## Row-averaged FFT: same-row, distance-weighted SNR cleanup (`sras_compute.py`)
+
+`compute_rf_image`'s `row_avg_n` parameter averages each pixel's CH1
+waveform with its up-to-n same-row neighbors before the FFT peak search, to
+improve SNR on noisy scans. Never crosses rows: pixel pitch is strongly
+anisotropic and varies by scan (5 µm × 50 µm on a typical scan, but as
+stretched as 5 µm × 1 mm on others), so a physically meaningful "neighbor"
+set can't be a fixed-shape 2-D window — but the X pitch *within one row* is
+a single file-wide constant (`SrasFile.pixel_x_mm`), so restricting to the
+row axis sidesteps the anisotropy question entirely rather than solving it
+with an elliptical or physically-scaled 2-D kernel.
+
+`_row_average_weights` is a Gaussian in pixel-index distance, not physical
+mm distance — deliberately: within one row those are the same function up
+to a fixed scale factor (`pixel_x_mm` is constant along a row), so the
+kernel itself needs no pitch at all. `pixel_x_mm` is used for real exactly
+once, in the GUI's options dialog, to show the window's physical width —
+not in the kernel math, where it would only ever cancel out.
+
+`_row_average_waveforms` is a masked/renormalized convolution (two
+`correlate1d` calls, numerator and denominator, divided) rather than a
+single fixed-normalized convolution, because a masked neighbor must
+contribute *zero weight*, not a zero-amplitude sample at full weight — the
+latter would bias every average near a masked run or a row's own edge
+toward zero. The same two-correlation trick handles row-edge truncation for
+free: `mode="constant", cval=0.0` zero-pads both the numerator and the
+denominator beyond a row's own ends, so the output renormalizes by whatever
+weight sum actually landed inside the row, no separate edge case.
+
+Background subtraction stays exactly where it already was (subtracted once
+from the fully-assembled `waves` buffer) rather than being threaded into the
+per-neighbor gather. This is exact, not an approximation: because
+`_row_average_waveforms`'s denominator is always the *actual* sum of
+included, valid weights (never a fixed total), `Σwᵢ·(rawᵢ−bg) / Σwᵢ`
+distributes to `avg − bg·(Σwᵢ/Σwᵢ) = avg − bg` regardless of which or how
+many neighbors were included — subtracting background from the averaged
+waveform is identical to subtracting it from every neighbor first, for any
+window, at any row edge, with any number of masked-out neighbors.
+
+No cross-row halo is needed: `compute_rf_image`'s chunk loop already splits
+on rows only, and `read_row` already reads one row's complete
+`(n_frames, spf)` slice at a time — averaging happens entirely inside that
+one row's own frame axis, so a chunk boundary (which falls between rows)
+can never truncate a window. Only a row's own start/end can, and that's the
+same edge case the masked convolution already handles.
+
+The averaging step doubles the live per-row scratch memory (a full-width
+`(n_frames, spf)` buffer on top of the existing compacted `waves` buffer),
+so `compute_rf_image` halves its byte budget when `row_avg_n > 0` before
+`_plan_fft_rows`/the exact-path sizing runs — see "Memory budget and row
+chunking" above. On the largest real scans `_plan_chunks` is already
+clamped to its floor of one row regardless, so this costs no concurrency
+where it matters most; it mainly protects moderate-sized scans from an
+unexpected regression.
+
+Persistence: `cached_rf_image` (the extracted fast-path check) requires
+`sras.precomputed_row_avg_n == row_avg_n` exactly, so a raw request can
+never be silently served a row-averaged cache or vice versa, and a request
+at one window size can never be served a cache at another — see
+`scan_format.md`'s Cache Tail / CACH tail version history sections for the
+on-disk `row_avg_n` field this depends on.
+
+## Serving a stored cache: provenance, not just presence
+
+A stored `peak_freq_mhz` image is only interchangeable with a live compute
+for the *exact* settings it was computed under. Three of them are baked
+irreversibly into the numbers — background subtraction, row-averaging window,
+and zero-padding — so all three are recorded in the `SFFT` block and checked
+by `cached_rf_image` before it hands the image back. Getting this wrong is
+not a slow display, it is a *wrong* display, which is why the check is a
+single predicate in one place rather than spread across callers.
+
+Padding is the subtlest of the three, because a padded FFT looks like it
+should be a refinement of the unpadded one. It isn't: zero-padding
+interpolates between the natural bins, so it resolves a different peak
+frequency for the same waveform. `precomputed_pad_factor` exists so a padded
+view can be served from a cache computed at *its* pad while still refusing
+one computed at any other, including pad 1. Before it existed the store was
+pad 1 by definition and any `n_fft` was rejected outright — correct, but it
+meant a user working at a pad factor got nothing at all from batch-computing
+a file, which is most of the point of the feature. `pad_factor_for` maps an
+`n_fft` request onto the integer factor a store could have recorded, and
+returns 0 for a request that is not a whole multiple of `samples_per_frame`
+— unmatchable by construction, since only integer factors are representable.
+
+The batch actions therefore have to cache at the *viewer's* current pad
+factor, not a fixed one: a cache stored at a pad nobody is viewing at is
+dead weight. When the two do diverge (the user changes the pad after
+batching), `_cache_mismatch_notes` says so in the scan info panel, because
+the symptom otherwise is just "the file I pre-computed got slow again" with
+no visible cause.
+
+That divergence note is informational only, not a warning of an impending
+recompute. The *display* path (`_stored_fft_image`) never asks
+`cached_rf_image` whether a stored image matches the window's live
+bg-sub/pad/row-averaging controls — it asks whether the image matches its
+*own* recorded settings (`sras.precomputed_bg_sub`/`precomputed_pad_factor`/
+`precomputed_row_avg_n`), which is always true whenever a stored image
+exists. So presence alone decides whether it's shown; the live controls
+never gate it. They still matter for two things: a genuinely never-computed
+angle's first live compute, and an explicit batch recompute — both of which
+read the live controls and produce new stored data, at which point it's the
+new data's *own* settings that get self-matched from then on. This is what
+keeps a view switch (angle, channel, or flipping bg-sub/pad) from ever
+discarding precomputed data — only an explicit batch recompute does, and it
+already reloads the file afterward so the new data displays immediately.
+Row-averaging has no live control to diverge from in the first place (it's
+only ever set inside the batch dialog), so it never appears in the
+divergence note — the "Cached images" line's own `row-averaged n=…` phrase
+already covers it.
+
+### Two caches, in cost order
+
+`_refresh_display` consults this window's in-session `_fft_cache`/`_dc_cache`
+dicts first, then the open file's stored v5/v7 blocks, and only then
+dispatches a `ComputeWorker`. The second tier is what makes a batch-computed
+file worth having; without it every angle change queued a worker and a
+progress popup for an image already sitting on disk — precisely the cost the
+batch was run to avoid. (The stored image *was* reachable before, but only
+from inside `ComputeWorker`, i.e. after paying for the thread and the popup.)
+
+The file tier asks with `allow_dc_recompute=False`. If applying the DC4 mask
+would mean reading a whole CH4 channel, it declines rather than blocking the
+GUI thread, and the fall-through worker reaches the same stored image via
+`compute_rf_image` and pays for the mask off-thread. So the GUI thread never
+does I/O, and the slow path is still a fast path.
+
+## Angle alignment coordinate frames (`sras_compute.py`)
+
+Alignment puts every angle's images onto one shared, zero-padded pixel grid
+using a rigid transform only — rotation + translation, never scale.
+
+Angle 0 (the reference) is the sole coordinate authority: it is the only
+angle whose stage XY (`x_start_mm` / `y_positions_mm`) is ever read, and the
+shared canvas is literally an extension of angle 0's own pixel grid, so the
+aligned view carries angle 0's real X/Y axes. Every *other* angle is placed
+purely by content — its rotation and translation come from cross-correlating
+its CH4 image against angle 0's (`register_angle_to_reference`) — and its own
+stage XY is deliberately never consulted. That is not an oversight: the
+rotation stage moves the sample relative to the scan window, so where a
+window sat in stage coordinates says nothing about where the sample is, and
+an earlier design that pivoted each angle on a signal-weighted centroid of
+its own window put every angle on a ~20 mm circle around the optical center
+instead of stacking them into one shape.
+
+Only two coordinate frames exist:
+
+* **local mm** — one angle's own physical frame: origin at the *center of its
+ own pixel array*, x along +column, y along +row, scaled by that angle's own
+ pitches. Carries no stage position whatsoever.
+* **ref mm** — the reference angle's local mm. A registration result
+ `(rotation_deg, shift_mm)` is exactly the rigid map from an angle's local
+ mm to ref mm: `q = R(rotation_deg) @ l + shift_mm`. Stage coordinates
+ re-enter once, at the very end, when the canvas origin is converted to
+ angle 0's stage mm (`AlignmentResult.canvas_origin_mm`).
+
+Rotation is done in mm, never on raw pixel indices: the x pitch
+(`SrasFile.pixel_x_mm`, 5 µm on a real scan) and the y/row pitch (50 µm)
+differ by 10×, so rotating the raw index grid would shear the image — an
+unwanted anisotropic scale. Registration runs on a resampled *isotropic* grid
+for the same reason, and every affine maps shared-grid index → mm → undo
+rotation/shift → that angle's own local mm → that angle's own raw index,
+matching the output→input convention `scipy.ndimage.affine_transform` wants.
+
+### Cropping the canvas is index translation, not a second transform
+
+`crop_alignment_result` restricts an `AlignmentResult` to a rectangular window
+of its canvas by folding the crop into each angle's existing affine rather than
+composing a new one. From `_affine_out_to_src`, `matrix = D @ Rinv @ A_out`
+depends only on the pitches and the rotation, and `A_out @ [row0, col0]` is
+exactly the mm displacement of the new origin, so
+
+```
+matrix @ [r', c'] + (offset + matrix @ [row0, col0])
+ == matrix @ [r' + row0, c' + col0] + offset
+```
+
+identically. `matrix` is untouched and `offset` — which already absorbs the
+origin — absorbs the crop too.
+
+Two things follow, and both are relied on. `apply_alignment`, `reproject_mask`
+and the aligned exporter all work on a cropped result with no special-casing:
+resampling a cropped result is *exactly* a slice of resampling the full one
+(`tests/test_align_export.py::test_crop_is_a_window_of_the_full_canvas` asserts
+bit equality). And because the crop offset is a whole number of canvas pixels,
+`canvas_for_params`' snap invariant — the reference angle lands on integer
+canvas pixels — survives the crop, which is what keeps the reference exportable
+as a verbatim block.
+
+## Aligned export (`sras_align_export.py`)
+
+`write_aligned_sras` bakes an alignment into a new v6 file: every angle
+resampled onto the cropped shared canvas, so all of them end up with identical
+geometry and the file opens already aligned. It is the only place in the
+codebase that *resamples* waveform data — `sras_edit_scans` and `sras_average`
+copy waveform bytes verbatim — which is why it is its own top-level module
+rather than part of `sras_format` (scoped to the versioned binary spec, per the
+sidecar section's own rule) or `sras_compute` (imported by every
+multiprocessing child).
+
+**Nearest neighbour, never interpolation.** Each output pixel gets exactly one
+source pixel's three waveforms, verbatim. Averaging two neighbouring CH1
+packets would synthesise a waveform the instrument never measured, whose FFT
+peak is the peak of neither — meaningless for a technique whose entire output is
+that peak frequency. The cost is that some source pixels are duplicated and
+others dropped, which is the same trade `apply_alignment`'s `order=0` already
+makes for the display.
+
+**The rounding rule is `floor(x + 0.5)`, not `np.rint`.** `scipy.ndimage`'s
+`order=0` rounds halves away from zero while `np.rint` rounds them to even. The
+canvas is snapped to the reference's own pixel grid, so an angle whose row pitch
+differs from the reference's lands on exact half-integers across whole rows —
+this is the common case, not a corner case. Getting it wrong shifts those rows
+by one source pixel relative to what the Aligned View drew.
+
+**Out-of-bounds is tested on the fractional coordinate, not the rounded index.**
+`scipy`'s `mode="constant"` writes `cval` wherever the coordinate leaves the
+range of sample *centres*, `[0, n-1]` — a coordinate of −0.4 rounds to a
+perfectly valid index 0 and is still padding. Testing the rounded index instead
+puts a one-pixel rim of real data everywhere the preview shows padding.
+
+**...but with a tolerance (`_EDGE_TOL`).** The affine is built from a chain of
+mm-space multiplications, so an exactly-integer transform comes out a few times
+1e-13 off: the reference angle's offset is `-20 - 7e-15`, not `-20`. A bare
+`>= 0.0` therefore rejects that angle's entire first row, and `<= n-1` its last
+column — for the *reference* angle, whose whole job is to pass through as an
+exact integer crop. The tolerance is ~7 orders of magnitude above that noise and
+~7 below the half-pixel scale at which a rounding decision means anything, so it
+can only ever change pixels whose scipy answer was itself decided by noise.
+
+**Padding is the per-channel ADC code nearest 0 mV, not 0.** Zero ADC decodes to
+`(0 - yoff) * ymult + yzero`, which on real calibration is around +100 mV —
+above any sensible CH4 mask threshold, so a zero fill would paint a solid
+rectangle of "valid" pixels around the sample and corrupt every DC image and ROI
+statistic downstream.
+
+**Source rows are served from sliding in-RAM bands** (`_SourceReader`). A
+rotated angle maps one output row to a *diagonal* across the source array, so
+the pixels of a single output row come from hundreds of different source rows —
+~1.4 MB each on a full-size scan. Indexing a memmap pixel-by-pixel in output
+order re-faults nearly the whole angle per output row: terabytes of paging for a
+gigabyte of data. Reading a contiguous band per output chunk, with the band
+advancing monotonically, costs roughly 2× the source size in total reads.
+
+**Writes go to `.part` and are `os.replace`d into position.** Not politeness: a
+truncated .sras is not detectably broken, because `_parse_v6` drops incomplete
+trailing angle blocks and opens what is left as an aborted scan. A half-written
+export left in place would silently look like a real file with fewer angles.
+
+**The Angle Table is carried over unchanged.** Alignment removes the *spatial*
+rotation of the sample; it does not change which acoustic propagation direction
+each angle measured, and that direction is the scientific content of a
+multi-angle scan. Zeroing the table would make the export self-consistent for
+re-registration and useless for anisotropy work. The consequence is that
+re-registering an export needs `seed_deg=0.0` to put 0° inside the coarse sweep,
+since `nominal_delta_deg` is still non-zero — which is exactly what the seed
+parameter exists for.
+
+## Alignment wizard (`sras_viewer/align_wizard.py`)
+
+A `QWizard` rather than another dialog because the three steps are genuinely
+sequential and the last one is destructive: correlate, choose a crop, write a
+file. It replaces both former Fusion actions, so it also absorbs the old
+`ManualAlignmentDialog`'s by-eye nudge editor — otherwise a scan the search
+cannot fit would have no fallback at all.
+
+Shared state lives on the wizard object, not in `registerField`: the pages pass
+numpy arrays, `ManualAngleParams` and an `AlignmentResult` between them, none of
+which are scalar widget properties.
+
+`IndependentPages` is deliberately left **off**. With it set Qt never calls
+`cleanupPage`, and `cleanupPage` is how the ROI page discards a crop when the
+user goes back to re-correlate — a crop is indexed in canvas pixels, and a new
+rotation means a different canvas, so stale indices would silently be
+reinterpreted against the wrong grid. `geometry_generation` is the belt-and-
+braces check for the same hazard.
+
+The mask-stack preview shares the **final** canvas's origin and uses a pitch
+that is an integer multiple of it, unlike the old manual dialog's padded,
+unsnapped preview canvas. That is what lets the crop page convert a rectangle
+drawn in millimetres into an exact integer window of the real canvas, with no
+second coordinate frame to reconcile.
+
+"Fit to full overlap" uses `largest_rect_at_least`, a largest-rectangle sweep,
+not a bounding box of the fully-covered pixels. The full-overlap region of
+several rotated scans is roughly a disc, and its bounding box has corners no
+angle covers — offering that as the crop would hand the user the padding they
+were trying to avoid.
+
+Every background launch follows the two rules `_run_worker`'s docstring
+establishes: disable the trigger *before* the call (so a re-entrant click cannot
+start a second thread over the first), and never ignore the returned bool.
+Progress is an inline `QProgressBar` on the page rather than a `QProgressDialog`
+— a window-modal popup over a wizard both looks wrong and reintroduces the
+event-loop pumping hazard that ordering exists to avoid. `reject()` refuses to
+close while a job is in flight, since the running worker's signals are connected
+to bound methods of the pages Qt would be deleting.
+
+## Manual-alignment sidecar (`sras_compute.py`)
+
+`.sras.align.json` lives next to the scan file. The code lives in
+`sras_compute`, not `sras_format`: `sras_format` is scoped to the versioned
+binary .sras spec itself (see `scan_format.md`), while a manual alignment is
+a viewer-computed *derived* artifact, analogous in kind to `AlignmentResult`
+— so it belongs with the alignment math it serialises. json + pathlib are
+stdlib, so this adds no dependency to a module whose load-bearing constraint
+is staying free of Qt/matplotlib for cheap multiprocessing-child imports.
+
+### Schema history
+
+The stored `rotation_deg`/`shift_mm` are meaningless without the frame they
+were measured in, so `_SIDECAR_SCHEMA_VERSION` is bumped whenever that frame
+changes. Each bump makes older files describe a different (and, for the bugs
+each bump fixed, actively wrong) transform than the same numbers would today;
+loading one unchanged would silently reproduce the very "scans show up
+everywhere" symptom the bump fixed — so older sidecars are treated as absent
+rather than migrated.
+
+* **1 → 2** — pivot moved from the scan-window bbox center to a
+ content-derived centroid, and the rotation sign convention was corrected.
+* **2 → 3** — the content centroid was abandoned entirely: rotation is now
+ about each angle's own array center, mapped onto the reference's array
+ center, with `shift_mm` in the reference's local mm frame. No angle but the
+ reference contributes stage coordinates any more.
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..9f7942e
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,41 @@
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "sras-viewer"
+version = "0.1.0"
+description = "Viewer and processing tools for SRAS .sras scan files"
+requires-python = ">=3.12"
+dependencies = [
+ "PyQt6==6.10.2",
+ "numpy==2.4.1",
+ "matplotlib==3.10.8",
+ "scipy==1.18.0",
+ # Angle alignment only: masked FFT phase correlation (skimage.registration).
+ "scikit-image==0.26.0",
+ # Faster rfft backend; the viewer falls back to scipy.fft without it.
+ "pyFFTW==0.15.1",
+ # Clamps BLAS threading under the FFT worker pool.
+ "threadpoolctl==3.6.0",
+]
+
+[project.optional-dependencies]
+dev = ["pytest"]
+
+[project.scripts]
+sras-viewer = "sras_viewer.main_window:main"
+
+[tool.setuptools]
+py-modules = [
+ "sras_format",
+ "sras_compute",
+ "sras_workers",
+ "sras_align_export",
+ "sras_average",
+ "sras_edit_scans",
+]
+packages = ["sras_viewer"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
diff --git a/scan_format.md b/scan_format.md
index 3c90e16..6181c02 100644
--- a/scan_format.md
+++ b/scan_format.md
@@ -173,6 +173,41 @@ sum over angles a of: n_rows[a] × 3 × n_frames[a] × samples_per_frame × byte
> Table and check `file_size` against the running total before reshaping —
> a fixed `(n_angles, n_rows, ...)` reshape (as in pre-v6 readers) will not
> work since row/frame counts are no longer uniform across angles.
+>
+> A consequence worth stating explicitly: because a short file opens
+> *successfully* as a scan with fewer angles, a truncated file is not
+> detectably broken. Anything that writes a .sras must therefore stage to a
+> temporary name and rename on success — the viewer's aligned export writes
+> `.part` and `os.replace`s it — or an interrupted write leaves behind
+> something that loads without complaint and silently has the wrong angle count.
+
+---
+
+## Files written by the viewer's Alignment Wizard
+
+The acquisition app is not the only producer of this format. The viewer's
+`Fusion → Alignment Wizard…` writes a **v6** file holding the aligned, cropped
+stack, with these properties:
+
+* Every angle shares one grid — the cropped alignment canvas — so the
+ Per-Angle Geometry Table is `n_angles` identical records and the ragged Row
+ Table is `n_angles` identical spans. The raggedness v6 exists for is still
+ *expressible*, just unused, so any v6 reader works unchanged.
+* `x_delta` is the reference angle's own pitch, which is exactly
+ `velocity_mm_s / laser_freq_hz`, so the derived X axis stays consistent with
+ the header.
+* The `*_nominal` header fields describe the crop. Uniquely for these files they
+ coincide with the actual per-angle geometry, since after alignment every angle
+ really does scan the same box.
+* The **Angle Table is unchanged**. Alignment removes the sample's spatial
+ rotation, not the acoustic propagation direction each angle measured — that
+ direction is the point of a multi-angle scan, so it is preserved.
+* Output pixels with no corresponding source pixel (the canvas corners a rotated
+ scan cannot reach) hold the per-channel ADC code nearest **0 mV**, not zero.
+ Zero ADC decodes to roughly +100 mV on real calibration and would read as
+ signal.
+* No Cache Tail is written: any cached DC/FFT is indexed by the source's grid
+ and would be meaningless on the new one.
---
@@ -220,7 +255,7 @@ actions.
| 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). |
+| 4 | 1 | `u8` | `cach_version` | Cache format version. Currently `3`; readers also accept `1` and `2` (each older tail simply lacks the fields added since — see the `SFFT` block below and [CACH tail version history](#cach-tail-version-history)). Any other value → treat the file as uncached (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 2–7 reserved, must be zero on write. |
### DC block `SDCB` (present iff `block_flags & 0x01`)
@@ -249,13 +284,29 @@ value, same as v5's `PREC` section.
### FFT block `SFFT` (present iff `block_flags & 0x02`)
-7-byte block header, format `">4sBH"`:
+Block header layout depends on `cach_version`:
+
+Each `cach_version` appended one trailing field, so the header grows but
+never shifts an existing offset:
+
+- **`cach_version` 1**: 7 bytes, format `">4sBH"` — magic, flags, n_stored.
+- **`cach_version` 2**: 8 bytes, format `">4sBHB"` — + `row_avg_n`.
+- **`cach_version` 3**: 10 bytes, format `">4sBHBH"` — + `pad_factor`.
+ Always written by current code.
+
+An older tail is read with its absent fields taken as the only value such a
+tail can describe: `row_avg_n = 0` for a `cach_version` 1 tail, which
+predates row-averaged FFT caching, and `pad_factor = 1` for `cach_version`
+1 or 2, which predate padded caching and are therefore natural-resolution.
+Files cached before either change keep working with no recompute.
| 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 1–7 reserved. |
+| 4 | 1 | `u8` | `flags` | Bit 0 = `bg_sub_applied` — background waveform was subtracted from CH1 before the FFT when these images were computed. Bit 1 = `row_averaged` — `peak_freq_mhz` came from same-row, distance-weighted averaged CH1 waveforms rather than raw per-pixel ones; `row_avg_n` (below) is the neighbor half-width used. Bits 2–7 reserved. |
| 5 | 2 | `u16` | `n_stored` | Number of angle entries that follow |
+| 7 | 1 | `u8` | `row_avg_n` | *`cach_version` ≥ 2 only.* Same-row neighbor half-width, in pixels, that `peak_freq_mhz` was averaged over before its FFT; `0` = raw (unaveraged). Meaningful only when `flags` bit 1 is set — a `cach_version` 1 tail has no such byte and is always `row_avg_n = 0`. |
+| 8 | 2 | `u16` | `pad_factor` | *`cach_version` 3 only.* Zero-padding factor the stored `peak_freq_mhz` was resolved at: `n_fft = pad_factor × samples_per_frame`, so `1` = natural resolution. Never `0`; a `cach_version` 1 or 2 tail has no such field and is always `pad_factor = 1`. |
followed by `n_stored` entries, each:
@@ -264,9 +315,10 @@ u16 angle_idx — index into the angle table (0-ba
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:
+**`peak_freq_mhz`** for a raw store (`row_avg_n == 0`) 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
@@ -276,10 +328,27 @@ 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.
+For a row-averaged store (`row_avg_n > 0`), the DC4 threshold is applied
+*during* the store — a pixel below threshold is left at `0` and never
+contributes to any neighbor's average — since neighbor validity can't be
+deferred to display time the way plain masking can. The threshold value
+itself is not recorded, only that averaging happened and at what window
+size. Readers still apply their own live DC4 threshold at display time
+exactly as for a raw store, using whatever mask they currently have.
+
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`.
+`peak_freq_mhz`) whenever the store's recorded provenance doesn't match what
+the reader is asking for: time-domain gating is active, the reader's
+requested `n_fft` doesn't equal `pad_factor × samples_per_frame`, the
+reader's background-subtraction setting doesn't match
+`flags.bg_sub_applied`, or the reader's requested `row_avg_n` doesn't match
+the stored value exactly. A raw request must never be served a row-averaged
+store, or vice versa; a request at one row-averaging window size must never
+be served a store at another; and a request at one padding must never be
+served a store at another, since a padded FFT interpolates between the
+natural bins and so resolves genuinely different peak frequencies. An
+`n_fft` that is not a whole multiple of `samples_per_frame` can never match
+any store, because only an integer `pad_factor` is representable.
### In-place write ordering
@@ -294,6 +363,23 @@ interrupted write leaves harmless trailing bytes rather than a corrupt file,
and the next successful write overwrites them via the same deterministic
`cache_offset`.
+### CACH tail version history
+
+Distinct from the outer `.sras` file `version` byte (top of this document),
+which has stayed `7` since the Cache Tail was introduced — this is the inner
+`cach_version` byte inside the `CACH` header itself.
+
+| cach_version | Change |
+|--------------|--------|
+| 1 | Initial Cache Tail: `SDCB` (DC) and `SFFT` (FFT, 7-byte header) blocks. |
+| 2 | `SFFT` header grows one byte, `row_avg_n` — the same-row neighbor half-width the stored `peak_freq_mhz` was averaged over before its FFT, `0` = raw. Readers still accept a `cach_version` 1 tail, treated as `row_avg_n = 0` for every angle it stores, so files cached before this change keep working without a recompute. |
+| 3 | `SFFT` header grows a `u16` `pad_factor` — the zero-padding factor the stored `peak_freq_mhz` was resolved at, `1` = natural resolution. Before this, a padded view could never use a stored cache at all (the store was pad 1 by definition and readers rejected any `n_fft ≠ samples_per_frame`), so a user working at a pad factor got no benefit from batch-computing a file. Recording the factor lets such a view be served, while still refusing a store resolved at a *different* pad. Readers accept `cach_version` 1 and 2 tails as `pad_factor = 1`. |
+
+A reader that does not know a `cach_version` must treat the file as
+uncached — not attempt a partial parse — and the file still reads as an
+ordinary v7 (byte-identical to v6) scan, so a forward-dated tail costs a
+recompute and never correctness.
+
---
## Acquisition Settings (fixed by sc3_aui_app.py)
diff --git a/sras_align_export.py b/sras_align_export.py
new file mode 100644
index 0000000..4e4311b
--- /dev/null
+++ b/sras_align_export.py
@@ -0,0 +1,456 @@
+#!/usr/bin/env python3
+"""Write an aligned, cropped .sras file from an AlignmentResult.
+
+The alignment machinery in sras_compute never modifies a scan: it produces an
+AlignmentResult, and every consumer resamples on the fly (apply_alignment for
+the display, reproject_mask for the overlay). That is right for a viewer, but it
+means the aligned stack cannot leave the process — no other tool can read it,
+and re-opening the scan re-does the registration.
+
+This module bakes an alignment into a new file. Each angle is resampled onto the
+shared canvas that AlignmentResult already defines, cropped to the caller's
+window, so every output angle ends up with *identical* geometry: same rows, same
+frames, same X/Y coordinates. Rotation and translation are gone, absorbed into
+where each waveform sits. The result is an ordinary v6 file that opens already
+aligned, and registering it against itself returns identity.
+
+Two deliberate choices, both about not inventing data:
+
+ * The resample is a nearest-neighbour **gather** of whole waveforms, never an
+ interpolation. Averaging two neighbouring pixels' CH1 packets would produce
+ a waveform the instrument never measured, whose FFT peak is not the peak of
+ either — meaningless for a technique whose entire output is that peak
+ frequency. So each output pixel gets exactly one source pixel's three
+ waveforms, verbatim, and the cost is that some source pixels are duplicated
+ and others dropped. This matches apply_alignment's order=0 for the same
+ reason.
+ * Output pixels with no source pixel (the canvas corners a rotated scan cannot
+ reach, and anything outside the crop's coverage) are filled with the ADC
+ code for 0 mV, not with zero. See _fill_row.
+
+Depends only on numpy/sras_format/sras_compute — no Qt — so it is directly
+unit-testable and importable from a worker thread.
+"""
+
+import os
+import struct
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import numpy as np
+
+import sras_compute as compute
+from sras_compute import AlignmentResult
+from sras_format import GEO_FMT_V6, HDR_FMT_V6, SrasFile, mv_to_adc
+
+# GEO_FMT_V6 stores n_rows as ">H" and n_frames as ">I". A canvas that overflows
+# either is not representable, and silently truncating would write a file whose
+# geometry table disagrees with its waveform block.
+_MAX_ROWS = 0xFFFF
+_MAX_FRAMES = 0xFFFFFFFF
+
+# Slack, in source pixels, on the in-bounds test at the very edge of a source
+# array. Absorbs the ~1e-13 of float noise an exactly-integer affine picks up
+# from being built in mm space; see _in_bounds.
+_EDGE_TOL = 1e-6
+
+# Output rows per write() call. One row is n_channels * n_cols * spf bytes —
+# ~1.5 MB on a full-size scan — so a handful of rows keeps the peak buffer in
+# the low tens of MB no matter how large the scan is.
+_ROW_CHUNK = 8
+
+
+@dataclass
+class ExportPlan:
+ """What write_aligned_sras would produce, without producing it.
+
+ Derived from the affine transforms alone — no waveform bytes are read — so
+ the wizard can call it on every ROI edit to keep a live size estimate and
+ per-angle coverage readout in front of the user *before* they commit to a
+ multi-gigabyte write.
+ """
+ n_rows: int
+ n_frames: int
+ n_angles: int
+ bytes_per_angle: int
+ total_bytes: int
+ valid_px: dict[int, int] # output pixels with a source pixel
+ warnings: list[str] = field(default_factory=list)
+
+ def coverage_frac(self, angle_idx: int) -> float:
+ px = self.n_rows * self.n_frames
+ return (self.valid_px.get(angle_idx, 0) / px) if px else 0.0
+
+ def empty_angles(self) -> list[int]:
+ """Angles that would be written as pure padding — no output pixel of
+ theirs has a source pixel."""
+ return [a for a in range(self.n_angles) if not self.valid_px.get(a, 0)]
+
+
+def _src_coords(t, rows, n_cols: int) -> tuple[np.ndarray, np.ndarray]:
+ """Fractional source (row, col) coordinates for whole output rows.
+
+ *rows* is an array of output row indices; both results are shaped
+ (len(rows), n_cols).
+ """
+ cols = np.arange(n_cols, dtype=np.float64)
+ r = np.asarray(rows, dtype=np.float64)[:, None]
+ sr = t.matrix[0, 0] * r + t.matrix[0, 1] * cols + t.offset[0]
+ sc = t.matrix[1, 0] * r + t.matrix[1, 1] * cols + t.offset[1]
+ return sr, sc
+
+
+def _round_idx(coord: np.ndarray) -> np.ndarray:
+ """Nearest source index, rounding halves away from zero.
+
+ floor(x + 0.5), not np.rint: scipy.ndimage's order=0 rounds halves away
+ from zero while np.rint rounds them to even, and these have to be the same
+ source pixels the apply_alignment(order=0) preview drew. Exact halves are
+ not a corner case here — the canvas is snapped to the reference angle's own
+ pixel grid (see canvas_for_params), so an unrotated angle lands on
+ half-integers wherever its row pitch differs from the reference's.
+ """
+ return np.floor(coord + 0.5).astype(np.int64)
+
+
+def _in_bounds(sr: np.ndarray, sc: np.ndarray,
+ n_rows: int, n_frames: int) -> np.ndarray:
+ """Which output pixels have a source pixel, by scipy's mode="constant" rule
+ plus a tolerance at the edge.
+
+ Tested on the *fractional* coordinate against the range of sample centres,
+ [0, n-1] inclusive — deliberately not on the rounded index. The two differ
+ around the whole rim: a coordinate of -0.4 rounds to a perfectly valid index
+ 0, but scipy calls it out of bounds and writes cval there, so testing the
+ rounded index would put a one-pixel rim of real data everywhere the Aligned
+ View shows padding.
+
+ _EDGE_TOL is why this is not literally scipy's test. The affine is built
+ from a chain of mm-space multiplications, so an exactly-integer transform
+ comes out a few times 1e-13 off (the reference angle's offset lands on
+ -20 - 7e-15 rather than -20). Bare >= 0.0 then rejects that angle's entire
+ first row, and <= n-1 its last column — for the *reference* angle, whose
+ whole role is to pass through as an exact integer crop. The tolerance is
+ seven orders of magnitude above that noise and seven below the half-pixel
+ scale at which a rounding decision is ever meaningful, so it can only ever
+ change pixels whose scipy answer was itself decided by rounding noise.
+ """
+ return ((sr >= -_EDGE_TOL) & (sr <= n_rows - 1 + _EDGE_TOL)
+ & (sc >= -_EDGE_TOL) & (sc <= n_frames - 1 + _EDGE_TOL))
+
+
+# Output rows evaluated per numpy call when counting coverage. Counting row by
+# row costs one small matmul per row (hundreds of milliseconds per angle on a
+# full-size scan, on every ROI edit); counting the whole canvas at once needs
+# hundreds of MB of index arrays. Blocking gets both: ~10 numpy calls per angle
+# against ~30 MB of live index arrays.
+_COUNT_BLOCK = 128
+
+
+def _count_in_bounds(t, n_rows: int, n_cols: int,
+ src_rows: int, src_frames: int) -> int:
+ """How many of the n_rows x n_cols output pixels have a source pixel."""
+ total = 0
+ for start in range(0, n_rows, _COUNT_BLOCK):
+ rows = np.arange(start, min(start + _COUNT_BLOCK, n_rows))
+ sr, sc = _src_coords(t, rows, n_cols)
+ total += int(np.count_nonzero(_in_bounds(sr, sc, src_rows, src_frames)))
+ return total
+
+
+def plan_export(sras: SrasFile, result: AlignmentResult) -> ExportPlan:
+ """Geometry, size and per-angle coverage of the file *result* would export.
+
+ Coverage is counted from the actual per-pixel index arrays rather than
+ approximated by the footprint parallelogram's area, because the two differ
+ exactly where it matters — a crop that clips one angle's scan window — and
+ this number is what tells the user an angle will come out mostly empty. No
+ waveform bytes are read, so it stays fast enough to call on every ROI edit.
+ """
+ n_rows, n_cols = result.canvas_shape
+ n_angles = sras.n_angles
+ warnings: list[str] = []
+
+ valid_px: dict[int, int] = {}
+ for a in range(n_angles):
+ t = result.per_angle.get(a)
+ if t is None:
+ valid_px[a] = 0
+ warnings.append(f"Angle {a} has no transform and will be all padding.")
+ continue
+ src_rows, src_frames = sras.image_shape(a)
+ valid_px[a] = _count_in_bounds(t, n_rows, n_cols, src_rows, src_frames)
+
+ bytes_per_angle = (n_rows * sras.n_channels * n_cols
+ * sras.samples_per_frame * sras.bytes_per_sample)
+
+ # Built before the remaining warnings so they can be phrased with the
+ # plan's own coverage_frac rather than a second copy of the same division.
+ plan = ExportPlan(n_rows=n_rows, n_frames=n_cols, n_angles=n_angles,
+ bytes_per_angle=bytes_per_angle,
+ total_bytes=bytes_per_angle * n_angles,
+ valid_px=valid_px, warnings=warnings)
+
+ if n_rows > _MAX_ROWS:
+ warnings.append(
+ f"Crop is {n_rows} rows; the .sras geometry table caps rows at "
+ f"{_MAX_ROWS}. Narrow the ROI in Y.")
+ if n_cols > _MAX_FRAMES:
+ warnings.append(f"Crop is {n_cols} frames; the cap is {_MAX_FRAMES}.")
+ for a in range(n_angles):
+ frac = plan.coverage_frac(a)
+ if frac == 0.0:
+ warnings.append(
+ f"Angle {a} has no data inside this crop — it will be written "
+ f"as all padding.")
+ elif frac < 0.10:
+ warnings.append(
+ f"Angle {a} covers only {frac * 100:.1f}% of the crop.")
+ if sras.background is None:
+ warnings.append(
+ "Input has no background waveform (pre-v4 scan); a zero background "
+ "is written, which makes background subtraction a no-op.")
+ if sras.version != 6:
+ warnings.append(f"Input is v{sras.version}; the export is written as v6.")
+ if sras.scan_aborted:
+ warnings.append(
+ f"Input scan was aborted: only its {n_angles} complete angle(s) "
+ f"are exported.")
+
+ return plan
+
+
+def _fill_row(sras: SrasFile, n_cols: int, dtype) -> np.ndarray:
+ """One output row of pure padding, shape (n_channels, n_cols, spf).
+
+ Filled per channel with the ADC code for 0 mV, not with 0. Zero ADC decodes
+ to (0 - yoff) * ymult + yzero, which for a real scope preamble is a long way
+ from 0 mV — often far enough to sit above the CH4 mask threshold, which
+ would paint a solid rectangle of "valid" pixels around the sample and make
+ every DC image and every ROI statistic wrong. Rounding to the integer code
+ lands within half an ADC step of 0 mV, which is as close as the format can
+ represent.
+ """
+ info = np.iinfo(dtype)
+ codes = [int(np.clip(round(mv_to_adc(0.0, *sras.cal(ch))), info.min, info.max))
+ for ch in range(sras.n_channels)]
+ row = np.empty((sras.n_channels, n_cols, sras.samples_per_frame), dtype=dtype)
+ for ch, code in enumerate(codes):
+ row[ch] = code
+ return row
+
+
+class _SourceReader:
+ """Gives the gather source rows without ever reading one twice.
+
+ This is the difference between a usable export and an unusable one, and it
+ is entirely about read amplification. A rotated angle maps one output row to
+ a *diagonal* line across the source array, so the pixels of a single output
+ row come from hundreds of different source rows — on a full-size scan, a
+ ~1.4 MB source row each. Indexing a memmap pixel by pixel in output order
+ therefore re-faults nearly the whole angle for every output row: terabytes
+ of paging for a gigabyte of data.
+
+ So rows are served from a contiguous *band* held in RAM. The band for a
+ chunk of output rows is read in one sequential slice, and because output
+ rows advance monotonically through the source, consecutive chunks' bands
+ barely overlap: each source row is read about once, and the whole job costs
+ roughly 2x the source size in reads rather than a thousand times it.
+
+ Small angles skip the machinery — if the whole block fits the budget it is
+ materialized once and every band is a view of it.
+
+ One honest caveat: a chunk whose diagonal spans more source rows than the
+ budget allows still gets the band it asked for, so the budget can be
+ overshot. The overshoot is bounded by the span of _ROW_CHUNK output rows,
+ and in the worst case (an extreme rotation on a huge scan) that is the whole
+ angle — i.e. no worse than the in-RAM path above. Accepted deliberately:
+ correctness of the gather is not negotiable, and the alternative is the
+ memmap thrashing this class exists to avoid.
+ """
+
+ def __init__(self, sras: SrasFile, angle_idx: int, budget: int):
+ self._src = sras.data[angle_idx]
+ self._n_rows = self._src.shape[0]
+ self._row_bytes = max(1, self._src[0].nbytes)
+ self._whole = np.asarray(self._src) if self._src.nbytes <= budget else None
+ # Leave room for the output buffer and the index arrays alongside.
+ self._max_band = max(1, int(budget * 0.5) // self._row_bytes)
+ self._band = None
+ self._lo = self._hi = 0
+
+ def band(self, lo: int, hi: int) -> tuple[np.ndarray, int]:
+ """Rows [lo, hi) as an in-RAM array, plus the index its row 0 holds."""
+ lo = max(0, min(lo, self._n_rows))
+ hi = max(lo + 1, min(hi, self._n_rows))
+ if self._whole is not None:
+ return self._whole, 0
+ if self._band is None or lo < self._lo or hi > self._hi:
+ # Read a little more than asked so a chunk whose band creeps
+ # forward by a few rows does not re-read the whole span.
+ span = min(self._max_band, max(hi - lo, self._max_band // 2))
+ self._lo = lo
+ self._hi = min(self._n_rows, lo + span)
+ if self._hi < hi: # band cannot cover the ask
+ self._hi = hi
+ self._band = np.asarray(self._src[self._lo:self._hi])
+ return self._band, self._lo
+
+ def close(self):
+ self._whole = None
+ self._band = None
+
+
+def write_aligned_sras(sras: SrasFile, result: AlignmentResult, out_path,
+ *, progress_cb=None, should_stop=None,
+ budget: int | None = None) -> Path:
+ """Write *sras*, aligned per *result* and cropped to its canvas, to a new
+ v6 .sras file. Returns the path written.
+
+ *result* is used exactly as given: crop the canvas first with
+ compute.crop_alignment_result, whose offset shift makes the cropped result
+ resample precisely the window the user selected.
+
+ No cache tail is written. Any DC/FFT the input had cached is indexed by the
+ input's grid and is meaningless on the new one, so — like sras_edit_scans —
+ the export drops it and lets the viewer recompute.
+
+ Writes to a sibling ".part" file and os.replace()s it into position on
+ success, unlinking it on error or cancellation: a half-written .sras is not
+ detectably broken (the v6 parser treats a short file as an aborted scan and
+ opens it happily), so it must never be left where the user might load it.
+
+ *should_stop* is polled once per output row chunk; returning True aborts and
+ raises nothing — the partial file is removed and the returned path will not
+ exist, so callers must check.
+ """
+ out_path = Path(out_path)
+ n_rows, n_cols = result.canvas_shape
+ n_ch, spf = sras.n_channels, sras.samples_per_frame
+ n_angles = sras.n_angles
+
+ if n_rows <= 0 or n_cols <= 0:
+ raise ValueError(f"empty canvas: {n_rows} x {n_cols}")
+ if n_rows > _MAX_ROWS:
+ raise ValueError(
+ f"{n_rows} rows exceeds the .sras per-angle geometry limit of "
+ f"{_MAX_ROWS}; crop further in Y")
+ if n_cols > _MAX_FRAMES:
+ raise ValueError(f"{n_cols} frames exceeds the limit of {_MAX_FRAMES}")
+ missing = [a for a in range(n_angles) if a not in result.per_angle]
+ if missing:
+ raise ValueError(f"alignment result has no transform for angle(s) {missing}")
+ # The source's waveform blocks are live read-only memmaps into sras.path,
+ # so writing over it would corrupt the very reads the gather is making.
+ if out_path.exists() and out_path.samefile(sras.path):
+ raise ValueError(
+ "refusing to export onto the source scan; choose another filename")
+
+ dtype = np.dtype(np.int8 if sras.bytes_per_sample == 1 else ">i2")
+ x0_mm, y0_mm = result.canvas_origin_mm
+ y_rows = (y0_mm + np.arange(n_rows) * result.canvas_dy_mm).astype(">f4")
+
+ # Reference-only header fields. v6/v7 inputs have real ones to carry over;
+ # for a legacy input describe the canvas we are actually writing.
+ if sras.x_start_nominal_mm is not None:
+ nominal = (sras.x_start_nominal_mm, sras.y_start_nominal_mm,
+ sras.x_delta_nominal_mm, sras.y_delta_nominal_mm,
+ sras.row_spacing_mm)
+ else:
+ nominal = (x0_mm, y0_mm,
+ n_cols * sras.pixel_x_mm, n_rows * result.canvas_dy_mm,
+ result.canvas_dy_mm)
+
+ header = struct.pack(
+ HDR_FMT_V6, b"SRAS", 6, n_angles,
+ float(nominal[0]), float(nominal[1]), float(nominal[2]),
+ float(nominal[3]), float(nominal[4]),
+ sras.velocity_mm_s, sras.laser_freq_hz, spf, sras.sample_rate_hz,
+ sras.bytes_per_sample, n_ch)
+
+ # Every angle now shares one grid, so the ragged v6 tables collapse to
+ # n_angles copies of the same record. x_delta is the reference angle's own
+ # pitch (the canvas is its grid extended), which is velocity/laser_freq
+ # exactly, so x_axis_mm() stays self-consistent on re-read.
+ geo = struct.pack(GEO_FMT_V6, float(x0_mm), float(sras.pixel_x_mm),
+ int(n_cols), int(n_rows)) * n_angles
+
+ budget = compute.memory_budget_bytes() if budget is None else max(1, budget)
+ total_chunks = max(1, n_angles * ((n_rows + _ROW_CHUNK - 1) // _ROW_CHUNK))
+ done_chunks = 0
+ cancelled = False
+
+ part_path = out_path.with_name(out_path.name + ".part")
+ try:
+ with open(part_path, "wb") as fout:
+ fout.write(header)
+ fout.write(sras.angles_deg.astype(">f4").tobytes())
+ fout.write(geo)
+ fout.write(y_rows.tobytes() * n_angles)
+ fout.write(sras.encoded_preambles())
+ fout.write(sras.encoded_background())
+
+ pad = _fill_row(sras, n_cols, dtype)
+ for a in range(n_angles):
+ t = result.per_angle[a]
+ reader = _SourceReader(sras, a, budget)
+ src_rows, src_frames = sras.image_shape(a)
+ try:
+ for chunk_start in range(0, n_rows, _ROW_CHUNK):
+ if should_stop is not None and should_stop():
+ cancelled = True
+ break
+ chunk = np.arange(chunk_start,
+ min(chunk_start + _ROW_CHUNK, n_rows))
+ sr, sc = _src_coords(t, chunk, n_cols)
+ ok = _in_bounds(sr, sc, src_rows, src_frames)
+ # Clip rather than trust: _EDGE_TOL admits coordinates a
+ # hair outside the array, and an index off the end here
+ # would silently read the wrong row of the band.
+ idx_r = np.clip(_round_idx(sr), 0, src_rows - 1)
+ idx_c = np.clip(_round_idx(sc), 0, src_frames - 1)
+
+ # One band read covers the whole chunk: every source row
+ # any of these output rows touches, in one sequential
+ # slice. See _SourceReader.
+ if ok.any():
+ band, base = reader.band(int(idx_r[ok].min()),
+ int(idx_r[ok].max()) + 1)
+ else:
+ band, base = None, 0
+
+ for i in range(len(chunk)):
+ out = pad.copy()
+ keep = ok[i]
+ if keep.any():
+ # The two advanced indices are separated by a
+ # slice, so numpy puts the gathered axis first:
+ # (n_sel, n_ch, spf). Move it behind channels.
+ out[:, keep, :] = band[
+ idx_r[i][keep] - base, :, idx_c[i][keep], :
+ ].transpose(1, 0, 2)
+ # out is C-contiguous, so the buffer protocol
+ # writes it straight out — .tobytes() would
+ # copy a full row per row written.
+ fout.write(out)
+ done_chunks += 1
+ if progress_cb is not None:
+ progress_cb(int(done_chunks / total_chunks * 100))
+ finally:
+ reader.close()
+ if cancelled:
+ break
+ if not cancelled:
+ fout.flush()
+ os.fsync(fout.fileno())
+ if cancelled:
+ part_path.unlink(missing_ok=True)
+ return out_path
+ os.replace(part_path, out_path)
+ except BaseException:
+ part_path.unlink(missing_ok=True)
+ raise
+
+ if progress_cb is not None:
+ progress_cb(100)
+ return out_path
diff --git a/sras_average.py b/sras_average.py
index 0c8fc77..aed403f 100644
--- a/sras_average.py
+++ b/sras_average.py
@@ -47,7 +47,7 @@ def read_header_sections(sras: SrasFile) -> bytes:
lost in a re-encode."""
with open(sras.path, "rb") as f:
f.seek(HDR_SIZE)
- return f.read(sras._data_offset - HDR_SIZE)
+ return f.read(sras.data_offset - HDR_SIZE)
def average_rows(block: np.ndarray, n: int, discard_remainder: bool) -> np.ndarray:
diff --git a/sras_compute.py b/sras_compute.py
index a0966d8..ba931c6 100644
--- a/sras_compute.py
+++ b/sras_compute.py
@@ -6,8 +6,10 @@ multiprocessing child can import it without loading Qt or matplotlib —
which matters because Python 3.14 on macOS spawns rather than forks.
"""
+import atexit
import json
import os
+import threading
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from pathlib import Path
@@ -16,7 +18,8 @@ import numpy as np
import scipy.fft as scipy_fft
import scipy.ndimage as scipy_ndimage
-from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, SrasFile, adc_to_mv
+from sras_format import (CH1_IDX, CH3_IDX, CH4_IDX, MAX_PAD_FACTOR, SrasFile,
+ adc_to_mv)
# ---------------------------------------------------------------------------
# FFT backend
@@ -24,20 +27,26 @@ from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, SrasFile, adc_to_mv
try:
import pyfftw
- pyfftw.interfaces.cache.enable()
PYFFTW_AVAILABLE = True
except ImportError:
PYFFTW_AVAILABLE = False
-_fft_backend = "numpy" # "numpy" or "pyfftw"; set via set_fft_backend()
+try:
+ from threadpoolctl import threadpool_limits
+except ImportError:
+ threadpool_limits = None
+
+_fft_backend = "scipy" # "scipy" or "pyfftw"; set via set_fft_backend()
def set_fft_backend(name: str):
"""Select the rfft implementation. Module-level state, so it must be set
explicitly inside each multiprocessing child — it does not survive a
- spawn."""
+ spawn. "numpy" is accepted as a legacy alias for "scipy"."""
global _fft_backend
- _fft_backend = name if (name != "pyfftw" or PYFFTW_AVAILABLE) else "numpy"
+ if name == "numpy":
+ name = "scipy"
+ _fft_backend = name if (name == "pyfftw" and PYFFTW_AVAILABLE) else "scipy"
def get_fft_backend() -> str:
@@ -52,33 +61,235 @@ def _do_rfft(x: np.ndarray, n: int | None = None, axis: int = -1,
return scipy_fft.rfft(x, n=n, axis=axis, workers=workers)
+# ---------------------------------------------------------------------------
+# FFT worker pool and per-thread pyFFTW plans
+# ---------------------------------------------------------------------------
+
+_FFT_BLOCK = 512 # waveforms per FFT task. The knee on a 16-core machine:
+ # smaller blocks serialise on GIL-held numpy dispatch,
+ # larger ones lose cache residency and task granularity.
+_ZOOM_MIN_PAD = 4 # zoom refinement engages at n_fft >= _ZOOM_MIN_PAD * spf
+_FFT_EXACT_ENV = bool(int(os.environ.get("SRAS_FFT_EXACT", "0") or 0))
+
+_pool_lock = threading.Lock()
+_pool: ThreadPoolExecutor | None = None
+
+
+def _fft_pool() -> ThreadPoolExecutor:
+ """The persistent process-wide pool for FFT block tasks."""
+ global _pool
+ with _pool_lock:
+ if _pool is None:
+ _pool = ThreadPoolExecutor(max_workers=_MAX_WORKERS,
+ thread_name_prefix="sras-fft")
+ atexit.register(_pool.shutdown, wait=False, cancel_futures=True)
+ return _pool
+
+
+_WISDOM_PATH = Path.home() / ".cache" / "sras-viewer" / "fftw_wisdom"
+_wisdom_lock = threading.Lock()
+_wisdom_loaded = False
+_fftw_local = threading.local()
+
+
+def _load_wisdom_once():
+ """Import saved FFTW wisdom so FFTW_MEASURE planning is a one-time cost
+ per machine. Purely an optimisation: failures are ignored."""
+ global _wisdom_loaded
+ with _wisdom_lock:
+ if _wisdom_loaded:
+ return
+ _wisdom_loaded = True
+ try:
+ pyfftw.import_wisdom(_WISDOM_PATH.read_bytes().split(b"\x00\n"))
+ except Exception:
+ pass
+
+
+def _save_wisdom():
+ with _wisdom_lock:
+ try:
+ _WISDOM_PATH.parent.mkdir(parents=True, exist_ok=True)
+ _WISDOM_PATH.write_bytes(b"\x00\n".join(pyfftw.export_wisdom()))
+ except Exception:
+ pass
+
+
+def _fftw_block_rfft(waves: np.ndarray, n: int) -> np.ndarray:
+ """rfft of a (B, spf) float32 block via a cached per-thread FFTW plan.
+
+ Plans have a fixed (_FFT_BLOCK, spf) input shape so each worker thread
+ plans once per transform length; a remainder block runs through the same
+ plan with its tail rows ignored. The returned array is the plan's output
+ buffer — consume it before the next call on the same thread.
+ """
+ n_wf, spf = waves.shape
+ plans = getattr(_fftw_local, "plans", None)
+ if plans is None:
+ plans = _fftw_local.plans = {}
+ key = (_FFT_BLOCK, spf, n)
+ plan = plans.get(key)
+ if plan is None:
+ _load_wisdom_once()
+ buf = pyfftw.empty_aligned((_FFT_BLOCK, spf), dtype="float32")
+ # No overwrite_input: FFTW must not scribble on input_array, whose
+ # zero-padded tail (columns spf..n) is zeroed exactly once here.
+ plan = pyfftw.builders.rfft(buf, n=n, axis=-1, threads=1,
+ planner_effort="FFTW_MEASURE")
+ plan.input_array[:] = 0.0
+ plans[key] = plan
+ _save_wisdom()
+ inp = plan.input_array
+ inp[:n_wf, :spf] = waves
+ return plan()[:n_wf]
+
+
+def _block_rfft(waves: np.ndarray, n: int) -> np.ndarray:
+ """Single-threaded rfft of one block; outer parallelism comes from the
+ pool, so the transform itself must not spin up threads."""
+ if _fft_backend == "pyfftw" and PYFFTW_AVAILABLE:
+ return _fftw_block_rfft(waves, n)
+ return scipy_fft.rfft(waves, n=n, axis=-1, workers=1)
+
+
+# ---------------------------------------------------------------------------
+# Zoom peak search: coarse rfft + local fine DFT around the winning bin
+# ---------------------------------------------------------------------------
+
+_ZOOM_HALFWIDTH = 0.75 # refinement window half-width, in coarse spacings.
+ # Every fine bin lies within 0.5 spacings of its
+ # nearest coarse bin, and that bin is guaranteed to
+ # be a candidate (see _ZOOM_CAND_RATIO), so 0.5
+ # suffices; 0.75 adds rounding margin.
+_ZOOM_CAND_RATIO = 0.7 # refine every coarse bin within this power ratio of
+ # its row's coarse maximum. Quarter-natural-bin
+ # scalloping at the 2x-oversampled coarse grid can
+ # understate a peak by at most ~19% in power, so 0.7
+ # keeps a wide margin — near-contenders are resolved
+ # on the fine grid, never ranked from coarse samples.
+
+
+@dataclass
+class _ZoomPlan:
+ """Constants for the coarse+refine peak search, built once per
+ compute_rf_image call and shared across worker threads. *phases* is a
+ lazily-filled window-start -> phase-vector cache; a benign duplicate
+ compute under concurrency is cheaper than locking."""
+ n_fft: int
+ n_coarse: int
+ m: int # fine bins per refinement window
+ n_bins_fine: int
+ E: np.ndarray # (spf, m) complex64 fine-DFT matrix, relative bins
+ j: np.ndarray # (spf,) float64 sample indices
+ phases: dict
+
+
+def _zoom_plan(spf: int, n_fft: int) -> _ZoomPlan:
+ # 2x-oversampled coarse grid: the padded power spectrum is a trig
+ # polynomial of degree spf-1, so at 2x sampling its global max cannot
+ # hide between coarse bins.
+ n_coarse = scipy_fft.next_fast_len(2 * spf, real=True)
+ n_bins_fine = n_fft // 2 + 1
+ m = int(np.ceil(2 * _ZOOM_HALFWIDTH * n_fft / n_coarse)) + 1
+ m = min(m, n_bins_fine - 1)
+ j = np.arange(spf, dtype=np.float64)
+ E = np.exp((-2j * np.pi / n_fft) * np.outer(j, np.arange(m))).astype(np.complex64)
+ return _ZoomPlan(n_fft, n_coarse, m, n_bins_fine, E, j, {})
+
+
+def _window_start(k_c: np.ndarray, zp: _ZoomPlan) -> np.ndarray:
+ """First fine bin of the refinement window around each coarse bin.
+ Clipped to [1, ...] so fine bin 0 stays excluded (DC suppression)."""
+ k0 = np.floor((k_c - _ZOOM_HALFWIDTH) * (zp.n_fft / zp.n_coarse)).astype(np.int64)
+ return np.clip(k0, 1, max(1, zp.n_bins_fine - zp.m))
+
+
+def _refine_window(waves: np.ndarray, rows: np.ndarray, k0: int, zp: _ZoomPlan,
+ best_pow: np.ndarray, best_bin: np.ndarray):
+ """Evaluate fine bins k0..k0+m-1 for the given rows and fold the result
+ into the per-row best (power, bin), preserving np.argmax's lowest-bin
+ tie-break."""
+ ph = zp.phases.get(k0)
+ if ph is None:
+ ph = np.exp((-2j * np.pi * k0 / zp.n_fft) * zp.j).astype(np.complex64)
+ zp.phases[k0] = ph
+ F = (waves[rows] * ph) @ zp.E
+ q = F.real ** 2
+ q += F.imag ** 2
+ i = np.argmax(q, axis=1)
+ p = q[np.arange(len(rows)), i]
+ b = k0 + i
+ upd = (p > best_pow[rows]) | ((p == best_pow[rows]) & (b < best_bin[rows]))
+ ridx = rows[upd]
+ best_pow[ridx] = p[upd]
+ best_bin[ridx] = b[upd]
+
+
+def _peak_bins_zoom(waves: np.ndarray, zp: _ZoomPlan) -> np.ndarray:
+ """Peak fine-bin per waveform without materialising the padded spectrum:
+ coarse rfft, then a small fine DFT (one gemm per shared window) on the
+ exact n_fft bin grid. Identity with the full padded argmax is enforced
+ by tests/test_compute.py::test_zoom_identity and the golden-hash sweep."""
+ n_wf = waves.shape[0]
+ S = _block_rfft(waves, zp.n_coarse)
+ P = S.real ** 2
+ P += S.imag ** 2
+ P[:, 0] = 0.0
+ p_c = np.max(P, axis=1)
+
+ best_pow = np.full(n_wf, -1.0, dtype=np.float32)
+ best_bin = np.full(n_wf, np.iinfo(np.int64).max, dtype=np.int64)
+
+ # For a clean signal this yields one or two windows; a noise spectrum
+ # (many near-equal peaks) yields a dozen or so — still a tiny fraction
+ # of the padded grid.
+ thr = np.where(p_c > 0, np.float32(_ZOOM_CAND_RATIO) * p_c,
+ np.float32(np.inf))
+ rows_c, bins_c = np.nonzero(P >= thr[:, None])
+ k0_c = _window_start(bins_c, zp)
+ # Zeroing the coarse DC bin (suppression) blinds the candidate scan to
+ # fine bins closer to DC than the first coarse sample — where the
+ # DC-leakage skirt of an un-subtracted offset peaks. Always refine the
+ # DC-adjacent window too.
+ rows_c = np.concatenate([rows_c, np.arange(n_wf)])
+ k0_c = np.concatenate([k0_c, np.ones(n_wf, dtype=np.int64)])
+ pair = np.unique(np.stack([rows_c, k0_c], axis=1), axis=0)
+ for k0 in np.unique(pair[:, 1]):
+ _refine_window(waves, pair[pair[:, 1] == k0, 0], int(k0), zp,
+ best_pow, best_bin)
+
+ # An all-zero spectrum must reproduce argmax-of-zeros = bin 0.
+ best_bin[p_c == 0.0] = 0
+ return best_bin
+
+
+def _peak_bins_direct(waves: np.ndarray, n_len: int) -> np.ndarray:
+ """Peak bin per waveform via the full transform — for pad factors too
+ small for the zoom search to pay."""
+ S = _block_rfft(waves, n_len)
+ power = S.real ** 2
+ power += S.imag ** 2
+ power[:, 0] = 0.0
+ return np.argmax(power, axis=1)
+
+
# ---------------------------------------------------------------------------
# Chunking / parallel budget
# ---------------------------------------------------------------------------
#
-# Rows are batched so the float32 working buffers for one chunk stay under a
-# memory budget. A fixed row count (the original design) works fine for small
-# legacy scans but is catastrophic for a v6 scan with a large per-angle
-# frame/sample count — e.g. a 7500-frame x 2500-sample angle needs ~2.4 GB for
-# a single 32-row chunk.
-#
-# With chunks running concurrently the budget has to cover *all* live chunks at
-# once. Note that on a large scan chunk_rows is already clamped to its floor of
-# 1 row (one row alone is ~75 MB of float32 at 7507x2500), so shrinking the
-# per-chunk size cannot buy more concurrency — the worker count must be derived
-# from the budget instead. See _plan_chunks.
-
-# 1024 MB is the measured knee on a 16-core machine against a 7507-frame x
-# 2500-sample angle: 512 MB leaves ~20% of the FFT speedup on the table, and
-# 1536+ MB costs ~0.4 GB more resident for no further gain. Override with
-# SRAS_MEM_BUDGET_MB on a smaller machine.
+# Row chunks are budgeted so all concurrently-live working buffers fit in
+# memory; the worker count is derived from the budget, not vice versa.
+# Rationale and the measured 1024 MB default: docs/design.md ("Memory budget
+# and row chunking").
_TOTAL_BYTES_BUDGET = int(os.environ.get("SRAS_MEM_BUDGET_MB", 1024)) * 1024 * 1024
_CHUNK_ROWS_MAX = 32 # cap for small scans (original behavior)
_MAX_WORKERS = int(os.environ.get("SRAS_MAX_WORKERS", 0)) or (os.cpu_count() or 4)
-# An rfft chunk holds, live at once: the float32 input, the complex64
-# transform, and the float32 power spectrum — roughly 3x the input buffer.
-_FFT_LIVE_MULTIPLIER = 3
+def memory_budget_bytes() -> int:
+ """The module-wide ceiling on concurrently-live working buffers, for
+ callers outside this module that read the same data (the aligned exporter
+ sizes its source-band reader against it)."""
+ return _TOTAL_BYTES_BUDGET
def _chunk_rows_for(n_frames: int, samples_per_frame: int,
@@ -87,6 +298,15 @@ def _chunk_rows_for(n_frames: int, samples_per_frame: int,
return int(max(1, min(_CHUNK_ROWS_MAX, budget // bytes_per_row)))
+def _plan_fft_rows(n_frames: int, samples_per_frame: int, budget: int) -> int:
+ """Rows per outer chunk for the block-FFT path. Only the float32
+ waveform buffer scales with the chunk (in-flight block spectra total a
+ few MB across the whole pool), so budget it with 2x slack and let the
+ block fan-out saturate the pool regardless of pad factor."""
+ bytes_per_row = max(1, n_frames * samples_per_frame * 4 * 2)
+ return int(max(1, min(_CHUNK_ROWS_MAX, budget // bytes_per_row)))
+
+
def _plan_chunks(n_rows: int, n_frames: int, samples_per_frame: int,
live_multiplier: int = 1,
max_workers: int | None = None,
@@ -208,6 +428,175 @@ def dc_image_mv(sras: SrasFile, angle_idx: int, ch_idx: int,
*sras.cal(ch_idx))
+
+# ---------------------------------------------------------------------------
+# Row-averaged FFT: same-row, distance-weighted CH1 waveform smoothing
+#
+# Averages a pixel's CH1 waveform with its same-row neighbors *before* the
+# FFT peak search, to improve SNR on noisy scans. Never crosses rows: pixel
+# pitch is strongly anisotropic and varies by scan, but the X pitch within
+# one row is a single file-wide constant (SrasFile.pixel_x_mm), so a
+# Gaussian in pixel-index distance along a row and one in true physical mm
+# distance are the same function up to that constant scale factor — the
+# kernel itself needs no pitch, only the GUI's physical-width hint label
+# does. Rationale and the masked-average/background-subtraction proofs:
+# docs/design.md ("Row-averaged FFT").
+# ---------------------------------------------------------------------------
+
+_ROW_AVG_SIGMA_FRAC = 0.5 # sigma = n * this; edge weight (distance n) is
+ # exp(-1/(2*frac**2)) ~= 0.135 of the center tap
+
+
+def _row_average_weights(row_avg_n: int) -> np.ndarray:
+ """(2n+1,) float32 Gaussian weights for distance-weighted row averaging,
+ symmetric around the center tap. Not pre-normalized to sum to 1 —
+ _row_average_waveforms renormalizes per output pixel by the actual sum
+ of included, valid, in-window neighbor weights, not a fixed total."""
+ n = max(0, int(row_avg_n))
+ if n == 0:
+ return np.ones(1, dtype=np.float32)
+ d = np.arange(-n, n + 1, dtype=np.float64)
+ sigma = n * _ROW_AVG_SIGMA_FRAC
+ return np.exp(-0.5 * (d / sigma) ** 2).astype(np.float32)
+
+
+def _row_average_waveforms(masked_waves: np.ndarray, valid: np.ndarray,
+ weights: np.ndarray) -> np.ndarray:
+ """Distance-weighted mean of each row position's CH1 waveform with its
+ same-row neighbors, counting only neighbors where *valid* is True.
+
+ masked_waves: (n_frames, spf) float32 — CH1 samples at valid[f]
+ positions; must already be 0.0 elsewhere (the caller must never
+ have read the raw memmap at an invalid position).
+ valid: (n_frames,) bool.
+ weights: (2n+1,) float32 from _row_average_weights.
+
+ Returns (n_frames, spf) float32, meaningful only where valid is True
+ (the caller compacts by that same mask right after, matching
+ compute_rf_image's existing masked-compaction contract).
+
+ Two 1-D correlations along the frame axis — the numerator against the
+ raw masked-zero waveforms, the denominator against the validity mask
+ itself — so a masked neighbor contributes *zero weight* rather than a
+ zero-amplitude sample at full weight, and a window truncated at a row's
+ edge renormalizes correctly with no separate edge case: mode="constant",
+ cval=0.0 pads both convolutions with zero beyond the row's own ends.
+ """
+ num = scipy_ndimage.correlate1d(masked_waves, weights, axis=0,
+ mode="constant", cval=0.0)
+ den = scipy_ndimage.correlate1d(valid.astype(np.float32), weights, axis=0,
+ mode="constant", cval=0.0)
+ den_safe = np.where(valid, den, np.float32(1.0))
+ return num / den_safe[:, None]
+
+
+def pad_factor_for(sras: SrasFile, n_fft: int | None) -> int:
+ """The integer zero-padding factor an *n_fft* request represents, or 0
+ if it represents none — i.e. it is not a whole multiple of this file's
+ samples_per_frame, so no stored image (which only ever records an
+ integer factor) can answer it.
+
+ 0 rather than None so callers can compare it straight against
+ ``sras.precomputed_pad_factor``, which is never 0.
+ """
+ if n_fft is None:
+ return 1
+ spf = sras.samples_per_frame
+ if spf <= 0 or n_fft % spf:
+ return 0
+ return max(1, n_fft // spf)
+
+
+def cache_mismatch_reasons(sras: SrasFile, *, n_fft: int | None,
+ apply_bg_sub: bool, row_avg_n: int) -> list[str]:
+ """Why this file's stored FFT images can't answer a request, as human-
+ readable phrases; empty means they can.
+
+ Padding, background subtraction and row-averaging are all baked into
+ the stored numbers, so a request that differs in any of them has to go
+ back through a real FFT. This is the single accept rule: cached_rf_image
+ serves a stored image iff this returns nothing, and callers that want to
+ *explain* the miss (rather than silently recompute) format these same
+ strings, so the two can't drift apart.
+ """
+ reasons = []
+ pad = pad_factor_for(sras, n_fft)
+ if pad != sras.precomputed_pad_factor:
+ want = f"pad {pad}x" if pad else "a ragged n_fft"
+ reasons.append(f"stored at pad {sras.precomputed_pad_factor}x, "
+ f"requested at {want}")
+ if sras.precomputed_bg_sub != (apply_bg_sub and sras.background is not None):
+ reasons.append("stored with background subtraction "
+ f"{'on' if sras.precomputed_bg_sub else 'off'}")
+ if sras.precomputed_row_avg_n != row_avg_n:
+ have = (f"row-averaged (n={sras.precomputed_row_avg_n})"
+ if sras.precomputed_row_avg_n else "raw per-pixel")
+ want = (f"row-averaged (n={row_avg_n})"
+ if row_avg_n else "raw per-pixel")
+ reasons.append(f"stored {have}, requested {want}")
+ return reasons
+
+
+def cached_rf_image(sras: SrasFile, angle_idx: int,
+ dc_threshold_mv: float | None,
+ apply_bg_sub: bool = True,
+ n_fft: int | None = None,
+ dc4_mv: np.ndarray | None = None,
+ row_avg_n: int = 0,
+ allow_dc_recompute: bool = True) -> np.ndarray | None:
+ """The precomputed-cache fast path for compute_rf_image: a ready-to-
+ display peak-frequency image if this file already has one matching
+ every setting the caller cares about, else None (caller must run a
+ real FFT).
+
+ A stored image stands in only when *all* hold: this angle actually has
+ a stored image (v5 PREC or v7 CACH); the requested zero-padding matches
+ what the store was computed at (see below); the stored bg-sub flag
+ matches what the caller wants; and sras.precomputed_row_avg_n ==
+ row_avg_n exactly (0 == "raw") — this last check is what stops a raw
+ request from ever being silently served a row-averaged image, or vice
+ versa, or a request at one window size being served a cache stored at a
+ different one.
+
+ Padding is checked the same way, and for the same reason: a padded FFT
+ interpolates between the natural bins, so it resolves genuinely
+ different peak frequencies. *n_fft* of None means natural resolution,
+ i.e. pad 1; anything else must be an exact whole multiple of
+ samples_per_frame equal to sras.precomputed_pad_factor. A ragged n_fft
+ that is not such a multiple can never match a stored image, since the
+ store only ever records an integer pad factor.
+
+ If *allow_dc_recompute* is False and no DC4 image is already cached or
+ supplied via *dc4_mv*, applying the mask would mean reading a whole
+ channel on the caller's behalf; this returns None instead so a caller
+ that wants to stay off the I/O path (e.g. a GUI thread) can choose to
+ fall through to a real compute rather than block.
+ """
+ cached_freq = sras.precomputed_freq_mhz[angle_idx]
+ if cached_freq is None or cache_mismatch_reasons(
+ sras, n_fft=n_fft, apply_bg_sub=apply_bg_sub, row_avg_n=row_avg_n):
+ return None
+ dc4_img = None
+ if dc_threshold_mv is not None:
+ # DC4 mask, in priority order: already-cached DC block, caller-
+ # supplied image, or a fresh (cheap — no FFT) recompute. Resolved
+ # before the copy below so the allow_dc_recompute bail-out doesn't
+ # allocate a full image it is about to throw away.
+ dc4_img = sras.cached_dc_mv(angle_idx, CH4_IDX)
+ if dc4_img is None:
+ if dc4_mv is not None:
+ dc4_img = dc4_mv
+ elif allow_dc_recompute:
+ dc4_img = adc_to_mv(compute_dc_image(sras, angle_idx, CH4_IDX),
+ *sras.cal(CH4_IDX))
+ else:
+ return None
+ freq_img = cached_freq.copy()
+ if dc4_img is not None:
+ freq_img[dc4_img < dc_threshold_mv] = 0.0
+ return freq_img
+
+
def compute_rf_image(sras: SrasFile, angle_idx: int,
dc_threshold_mv: float | None,
apply_bg_sub: bool = True,
@@ -215,7 +604,9 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
dc4_mv: np.ndarray | None = None,
max_workers: int | None = None,
budget: int | None = None,
- should_stop=None) -> np.ndarray:
+ should_stop=None,
+ exact: bool = False,
+ row_avg_n: int = 0) -> np.ndarray:
"""FFT of each CH1 waveform; pixel = peak frequency in MHz.
Pixels where CH4_dc < dc_threshold_mv are set to 0 — and the FFT is
@@ -233,37 +624,37 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
DC-channel precompute cache), pass it as *dc4_mv* (mV, shape
(n_rows, n_frames)) to reuse it instead of re-reading CH4 here.
+ At pad factors >= _ZOOM_MIN_PAD the padded spectrum is never
+ materialised: a coarse rfft finds each peak and a local fine DFT
+ resolves it on the exact n_fft bin grid (_peak_bins_zoom). *exact*
+ forces the reference full-padded transform instead — it exists for
+ tests, audits, and the SRAS_FFT_EXACT=1 escape hatch, and is slow and
+ memory-hungry at high pad.
+
+ *row_avg_n* > 0 averages each pixel's CH1 waveform with its up-to-n
+ same-row neighbors (distance-weighted, valid-neighbors-only per the
+ same dc_threshold_mv mask) before the FFT runs — see
+ _row_average_waveforms. 0 (default) is the raw, unaveraged behavior.
+
Fast path: if the file has a precomputed peak-frequency image for this
- angle (v5 PREC or v7 CACH), zero-padding is off, and the bg-sub flag
- matches, the stored image is used directly — no FFT is run.
+ angle (v5 PREC or v7 CACH) matching every one of the caller's settings
+ — including row_avg_n exactly — the stored image is used directly, no
+ FFT is run. See cached_rf_image.
"""
n_rows, n_frames = sras.image_shape(angle_idx)
data = sras.data[angle_idx]
- # ---- Fast path: precomputed image (v5 PREC or v7 CACH) ----------------
- cached_freq = sras.precomputed_freq_mhz[angle_idx]
- if (cached_freq is not None
- and n_fft is None # no custom zero-padding
- and sras.precomputed_bg_sub == (apply_bg_sub and sras.background is not None)):
- freq_img = cached_freq.copy()
- if dc_threshold_mv is not None:
- # DC4 mask, in priority order: already-cached DC block, caller-
- # supplied image, or a fresh (cheap — no FFT) recompute.
- dc4_img = sras.cached_dc_mv(angle_idx, CH4_IDX)
- if dc4_img is None:
- dc4_img = dc4_mv if dc4_mv is not None else adc_to_mv(
- compute_dc_image(sras, angle_idx, CH4_IDX), *sras.cal(CH4_IDX))
- freq_img[dc4_img < dc_threshold_mv] = 0.0
- return freq_img
+ fast = cached_rf_image(sras, angle_idx, dc_threshold_mv, apply_bg_sub=apply_bg_sub,
+ n_fft=n_fft, dc4_mv=dc4_mv, row_avg_n=row_avg_n)
+ if fast is not None:
+ return fast
# ---- Chunked FFT path --------------------------------------------------
- freq_axis = sras.freq_axis_mhz(n_fft)
+ exact = exact or _FFT_EXACT_ENV
+ spf = sras.samples_per_frame
+ n_len = n_fft if n_fft is not None else spf
+ freq32 = sras.freq_axis_mhz(n_fft).astype(np.float32)
img = np.zeros((n_rows, n_frames), dtype=np.float32)
- n_fft_bins = n_fft if n_fft is not None else sras.samples_per_frame
- chunk_rows, n_workers = _plan_chunks(
- n_rows, n_frames, max(sras.samples_per_frame, n_fft_bins),
- live_multiplier=_FFT_LIVE_MULTIPLIER, max_workers=max_workers,
- budget=budget)
background = sras.background if (apply_bg_sub and sras.background is not None) else None
cal4 = sras.cal(CH4_IDX)
# DC4 mask source, in the same priority order as the fast path above:
@@ -272,8 +663,31 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
dc4_full = sras.cached_dc_mv(angle_idx, CH4_IDX) if dc_threshold_mv is not None else None
if dc4_full is None:
dc4_full = dc4_mv
+ row_avg_weights = _row_average_weights(row_avg_n) if row_avg_n > 0 else None
- def chunk(r0: int, r1: int):
+ zp = (_zoom_plan(spf, n_fft)
+ if not exact and n_fft is not None and n_fft >= _ZOOM_MIN_PAD * spf
+ else None)
+ total = _TOTAL_BYTES_BUDGET if budget is None else max(1, budget)
+ if row_avg_n > 0:
+ # One extra same-sized transient buffer (the pre-averaging full-row
+ # scratch array) is live per in-flight row; halve the budget so
+ # _plan_fft_rows/the exact-path sizing accounts for it rather than
+ # relying on _plan_fft_rows's existing 2x slack to happen to cover it.
+ total = max(1, total // 2)
+ cap = max_workers if max_workers is not None else _MAX_WORKERS
+ if exact:
+ # The reference path materialises the full padded spectrum, so rows
+ # are budgeted against it (complex64 + power + temp per bin) and the
+ # chunk runs as one serial transform.
+ bytes_per_row = max(1, n_frames * (4 * spf + 16 * (n_len // 2 + 1)))
+ chunk_rows = int(max(1, min(_CHUNK_ROWS_MAX, total // bytes_per_row)))
+ cap = 1
+ else:
+ chunk_rows = _plan_fft_rows(n_frames, spf, total)
+ pool = _fft_pool() if cap > 1 else None
+
+ def process(r0: int, r1: int):
if dc_threshold_mv is None:
valid = None
else:
@@ -286,36 +700,91 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
if not valid.any():
return
- # Index the raw memmap slice with the boolean mask *before*
- # converting dtype — this is a lazy view until touched, so only the
- # selected elements are actually read from disk; masked-out pixels'
- # pages are never paged in at all.
- raw = data[r0:r1, CH1_IDX, :, :]
- waves = (raw[valid] if valid is not None else raw).astype(np.float32)
+ counts = (valid.sum(axis=1) if valid is not None
+ else np.full(r1 - r0, n_frames, dtype=np.int64))
+ offs = np.concatenate(([0], np.cumsum(counts)))
+ n_wf = int(offs[-1])
+ waves = np.empty((n_wf, spf), dtype=np.float32)
- if background is not None:
- waves -= background # background is 1-D (spf,)
+ def read_row(i: int):
+ if should_stop is not None and should_stop():
+ return
+ raw = data[r0 + i, CH1_IDX]
+ dst = waves[offs[i]:offs[i + 1]]
+ row_valid = valid[i] if valid is not None else None
+ if row_avg_weights is not None:
+ v = row_valid if row_valid is not None else np.ones(n_frames, dtype=bool)
+ # A plain cast + broadcast multiply, not a boolean-indexed
+ # scatter into a zeroed buffer: numpy's fancy indexing holds
+ # the GIL for its whole duration (measured zero speedup
+ # across threads, and net negative once threads outnumber
+ # physical cores), while a cast and a multiply are ordinary
+ # ufuncs that release it — this is what lets read_row actually
+ # parallelize across the pool instead of serializing on the
+ # scatter/gather. Trade-off: every sample in the row is read,
+ # valid or not (row averaging needs broad neighbor context
+ # regardless, unlike the plain path below, which still skips
+ # masked-out pixels entirely).
+ full = raw.astype(np.float32) * v[:, None]
+ avg = _row_average_waveforms(full, v, row_avg_weights)
+ dst[:] = avg[v] if row_valid is not None else avg
+ else:
+ # Index the raw memmap slice with the boolean mask *before*
+ # converting dtype — this is a lazy view until touched, so
+ # only the selected elements are actually read from disk;
+ # masked-out pixels' pages are never paged in at all.
+ dst[:] = raw[row_valid] if row_valid is not None else raw
+ if background is not None:
+ dst -= background # background is 1-D (spf,)
- # Parallelism comes from the outer chunk loop, so keep the inner
- # transform single-threaded to avoid oversubscribing the machine.
- spectrum = _do_rfft(waves, n=n_fft, axis=-1, workers=1)
- del waves
- # |z|^2 without np.abs()'s extra full-size temporary.
- power = spectrum.real ** 2
- power += spectrum.imag ** 2
- del spectrum
- # [..., 0] not [:, 0]: the unmasked path keeps the (rows, frames,
- # bins) shape, where [:, 0] would blank a whole frame.
- power[..., 0] = 0.0 # suppress DC bin
- peak_bins = np.argmax(power, axis=-1)
- del power
-
- if valid is not None:
- img[r0:r1][valid] = freq_axis[peak_bins]
+ if pool is None:
+ for i in range(r1 - r0):
+ read_row(i)
else:
- img[r0:r1] = freq_axis[peak_bins]
+ list(pool.map(read_row, range(r1 - r0)))
- _map_row_chunks(n_rows, chunk_rows, n_workers, chunk, should_stop=should_stop)
+ out = np.empty(n_wf, dtype=np.float32)
+
+ def fft_block(b0: int):
+ if should_stop is not None and should_stop():
+ return
+ b1 = min(b0 + _FFT_BLOCK, n_wf)
+ w = waves[b0:b1]
+ bins = (_peak_bins_zoom(w, zp) if zp is not None
+ else _peak_bins_direct(w, n_len))
+ out[b0:b1] = freq32[bins]
+
+ if exact:
+ spectrum = _do_rfft(waves, n=n_fft, axis=-1, workers=1)
+ power = spectrum.real ** 2
+ power += spectrum.imag ** 2
+ power[:, 0] = 0.0 # suppress DC bin
+ out[:] = freq32[np.argmax(power, axis=1)]
+ elif pool is None:
+ for b0 in range(0, n_wf, _FFT_BLOCK):
+ fft_block(b0)
+ else:
+ list(pool.map(fft_block, range(0, n_wf, _FFT_BLOCK)))
+
+ # Boolean scatter is row-major, matching the read pass's
+ # concatenation order.
+ if valid is not None:
+ img[r0:r1][valid] = out
+ else:
+ img[r0:r1] = out.reshape(r1 - r0, n_frames)
+
+ # BLAS must not thread under the pool (the fine-DFT gemm would multiply
+ # against the pool's own workers).
+ limiter = (threadpool_limits(limits=1)
+ if pool is not None and threadpool_limits is not None else None)
+ try:
+ for r0 in range(0, n_rows, chunk_rows):
+ if should_stop is not None and should_stop():
+ break
+ process(r0, min(r0 + chunk_rows, n_rows))
+ finally:
+ if limiter is not None:
+ limiter.unregister()
return img
@@ -324,18 +793,41 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
# ---------------------------------------------------------------------------
def cache_file(path: str, mode: str, apply_bg_sub: bool,
- fft_backend: str = "numpy", max_workers: int = 0) -> str:
+ fft_backend: str = "scipy", max_workers: int = 0,
+ pad_factor: int = 1,
+ dc_threshold_mv: float | None = None,
+ row_avg_n: int = 0) -> str:
"""Compute and store DC or FFT images for every angle of one file,
converting v6 → v7 in place. Returns "" on success or an error message.
Module-level and picklable so it can run in a ProcessPoolExecutor. The
FFT backend and worker cap are passed explicitly because module globals
do not survive a spawn.
+
+ mode is "dc", "fft" (raw per-pixel FFT, unmasked — masking applied at
+ display time), or "fft_rowavg" (same-row, distance-weighted CH1
+ averaging before the FFT — see compute_rf_image's row_avg_n).
+ fft_rowavg needs *dc_threshold_mv* up front, unlike plain "fft":
+ neighbor validity is baked into the stored numbers, so it can't be
+ deferred to display time the way plain masking can.
+
+ *pad_factor* is the zero-padding factor to resolve the peaks at: 1 (the
+ default) is natural resolution, n_fft == samples_per_frame. It is
+ recorded in the SFFT block so a reader knows which views the stored
+ numbers answer for — and it has to be the pad the viewer is *actually*
+ using, since a pad-1 cache is dead weight to a padded view and vice
+ versa (cached_rf_image refuses the mismatch rather than showing peaks
+ resolved at the wrong resolution).
"""
global _MAX_WORKERS
try:
- if mode not in ("dc", "fft"):
- return f"unknown cache mode {mode!r} (expected 'dc' or 'fft')"
+ if mode not in ("dc", "fft", "fft_rowavg"):
+ return f"unknown cache mode {mode!r} (expected 'dc', 'fft', or 'fft_rowavg')"
+ # write_v7_cache enforces this bound too, but only once every angle's
+ # FFT has already been computed. Checking up front is the difference
+ # between a bad argument costing nothing and costing the whole run.
+ if not (1 <= pad_factor <= MAX_PAD_FACTOR):
+ return f"pad_factor must be 1-{MAX_PAD_FACTOR}, got {pad_factor}"
set_fft_backend(fft_backend)
if max_workers:
_MAX_WORKERS = max_workers
@@ -346,21 +838,49 @@ def cache_file(path: str, mode: str, apply_bg_sub: bool,
"can be batch-cached")
n = sras.n_angles
+ n_fft = sras.samples_per_frame * pad_factor if pad_factor > 1 else None
+ n_workers, angle_budget = plan_angle_level(sras)
if mode == "dc":
- dc3 = [adc_to_mv(compute_dc_image(sras, a, CH3_IDX), *sras.cal(CH3_IDX))
- for a in range(n)]
- dc4 = [adc_to_mv(compute_dc_image(sras, a, CH4_IDX), *sras.cal(CH4_IDX))
- for a in range(n)]
+ dc3 = _parallel_map(
+ lambda a: adc_to_mv(
+ compute_dc_image(sras, a, CH3_IDX, max_workers=1,
+ budget=angle_budget), *sras.cal(CH3_IDX)),
+ range(n), n_workers)
+ dc4 = _parallel_map(
+ lambda a: adc_to_mv(
+ compute_dc_image(sras, a, CH4_IDX, max_workers=1,
+ budget=angle_budget), *sras.cal(CH4_IDX)),
+ range(n), n_workers)
sras.write_v7_cache(new_dc3_mv=dc3, new_dc4_mv=dc4)
- else:
+ elif mode == "fft":
effective_bg = apply_bg_sub and sras.background is not None
# dc_threshold_mv=None: store unmasked images and mask at display
# time (same convention as v5's PREC block). Skipping the mask
- # also skips reading CH4 entirely.
+ # also skips reading CH4 entirely. The FFT path parallelises
+ # internally over blocks, so angles run one at a time with the
+ # full budget.
freq = [compute_rf_image(sras, a, dc_threshold_mv=None,
- apply_bg_sub=effective_bg)
+ apply_bg_sub=effective_bg, n_fft=n_fft)
for a in range(n)]
- sras.write_v7_cache(new_freq_mhz=freq, new_bg_sub=effective_bg)
+ # new_row_avg_n=0 explicitly: these images are raw per-pixel FFTs,
+ # and carrying forward a row_avg_n left by an earlier fft_rowavg
+ # write would label them as something they are not.
+ sras.write_v7_cache(new_freq_mhz=freq, new_bg_sub=effective_bg,
+ new_row_avg_n=0, new_pad_factor=pad_factor)
+ else: # "fft_rowavg"
+ if row_avg_n <= 0:
+ return "row_avg_n must be a positive neighbor half-width for fft_rowavg mode"
+ if dc_threshold_mv is None:
+ return ("fft_rowavg mode requires a DC threshold "
+ "(neighbor validity depends on it)")
+ effective_bg = apply_bg_sub and sras.background is not None
+ freq = [compute_rf_image(sras, a, dc_threshold_mv=dc_threshold_mv,
+ apply_bg_sub=effective_bg, n_fft=n_fft,
+ row_avg_n=row_avg_n)
+ for a in range(n)]
+ sras.write_v7_cache(new_freq_mhz=freq, new_bg_sub=effective_bg,
+ new_row_avg_n=row_avg_n,
+ new_pad_factor=pad_factor)
return ""
except Exception as exc:
return str(exc)
@@ -369,29 +889,22 @@ def cache_file(path: str, mode: str, apply_bg_sub: bool,
# ---------------------------------------------------------------------------
# Angle alignment (Fusion menu)
#
-# Puts every angle's images onto one shared, zero-padded pixel grid using a
-# rigid transform only (rotation + translation, never scale). Rotation for
-# angle `a` is the *known* scan-angle delta relative to a reference angle —
-# never searched. Only the residual translation is found, via FFT phase
-# correlation of each angle's binarized CH4 ("dc-mask") image.
-#
-# Rotation is done in physical mm space rather than on raw pixel indices:
-# the x-pixel pitch (SrasFile.pixel_x_mm) is file-wide constant but the
-# y-pixel pitch (row spacing) can differ from it, and for v6 files can even
-# vary per angle. Rotating the raw index grid directly would implicitly
-# assume square pixels and shear a non-square-pixel image — an unwanted
-# effective anisotropic scale. Instead each angle gets one affine that maps
-# shared-canvas pixel index -> mm -> undo rotation/shift -> that angle's own
-# local mm -> that angle's own raw pixel index, matching the output->input
-# convention scipy.ndimage.affine_transform expects.
+# Rigid transforms only (rotation + translation, never scale), computed in mm
+# on two frames: each angle's "local mm" (origin at its own array center) and
+# the reference angle's local mm. Angle 0 is the sole coordinate authority;
+# every other angle is placed purely by image content. Why, and the full
+# frame/affine conventions: docs/design.md ("Angle alignment coordinate
+# frames").
# ---------------------------------------------------------------------------
@dataclass
class AngleTransform:
rotation_deg: float
- shift_mm: tuple[float, float] # (dx_mm, dy_mm) found by phase correlation
+ shift_mm: tuple[float, float] # (dx_mm, dy_mm) in ref mm
matrix: np.ndarray # (2,2): canvas (row,col) -> this angle's raw (row,col)
offset: np.ndarray # (2,)
+ score: float = 1.0 # registration NCC (1.0 = reference/manual)
+ source: str = "" # image that won registration: "signal"/"mask"
@dataclass
@@ -409,36 +922,90 @@ class AlignmentResult:
class ManualAngleParams:
"""One angle's manual-alignment state, independent of any canvas.
- rotation_deg/shift_mm are exactly AngleTransform's non-derived fields —
- the pair a canvas-bound AngleTransform's matrix/offset get built from
- once a canvas is decided (build_manual_alignment). Defaults to identity
- (no rotation, no shift): a fresh angle with no prior alignment is shown
- raw, exactly as scanned — the same "fully unaligned" state Clear
- Alignment resets back to.
+ rotation_deg/shift_mm are exactly AngleTransform's non-derived fields — the
+ rigid map from this angle's local mm to ref mm, the pair a canvas-bound
+ AngleTransform's matrix/offset get built from once a canvas is decided
+ (build_manual_alignment). Defaults to identity: a fresh angle with no prior
+ alignment is shown centered on the reference with no rotation, which is the
+ same "fully unaligned" state Clear Alignment resets back to.
"""
rotation_deg: float = 0.0
shift_mm: tuple[float, float] = (0.0, 0.0)
-def _pixel_pitch_mm(sras: SrasFile, angle_idx: int) -> tuple[float, float]:
+# Side length of the grid the rotation refinement runs on. The whole cost of
+# registration scales with it: ~325 ms per candidate rotation at 640, roughly
+# quadrupling per doubling, against ~200 MB of live masked-correlation buffers
+# (see _registration_workers). 640 puts a full-size scan at ~0.1 mm/px, which
+# resolves rotation on a 20 mm sample to well under a tenth of a degree.
+_DEFAULT_FINE_DIM = 640
+
+
+@dataclass
+class RigidFit:
+ """What register_angle_to_reference found for one angle."""
+ rotation_deg: float
+ shift_mm: tuple[float, float]
+ score: float # zero-mean NCC over the valid overlap
+ source: str # "signal", "mask", or "reference"
+
+
+def _skimage_phase_cross_correlation():
+ """Import skimage.registration lazily and cache it.
+
+ Deliberately not a module-level import: this module is imported by every
+ multiprocessing child (see the module docstring), skimage costs ~0.6 s to
+ import, and no child ever registers anything — registration runs in GUI-
+ process threads.
+ """
+ global _pcc
+ try:
+ return _pcc
+ except NameError:
+ from skimage.registration import phase_cross_correlation as _fn
+ _pcc = _fn
+ return _pcc
+
+
+# ---- Geometry: local mm, ref mm, and the one affine builder ---------------
+
+def pixel_pitch_mm(sras: SrasFile, angle_idx: int) -> tuple[float, float]:
"""(dx, dy) mm/pixel for one angle: dx is the file-wide constant
pixel_x_mm; dy is this angle's own row spacing (assumed uniform, the same
- assumption _redraw_image makes when it builds the display extent)."""
+ assumption _redraw_image makes when it builds the display extent). dy keeps
+ its sign, so +row always means the same physical direction as +y."""
y = sras.y_positions_mm(angle_idx)
return sras.pixel_x_mm, (float(y[1] - y[0]) if len(y) > 1 else 1.0)
-def _bbox_center_mm(sras: SrasFile, angle_idx: int) -> tuple[float, float]:
- x = sras.x_axis_mm(angle_idx)
- y = sras.y_positions_mm(angle_idx)
- return float((x[0] + x[-1]) / 2.0), float((y[0] + y[-1]) / 2.0)
+def _center_idx(sras: SrasFile, angle_idx: int) -> np.ndarray:
+ """(row, col) index of this angle's array center — the origin of its local
+ mm frame. Purely geometric: it depends on the array shape and nothing
+ else, which is what keeps local mm free of stage position."""
+ n_rows, n_frames = sras.image_shape(angle_idx)
+ return np.array([(n_rows - 1) / 2.0, (n_frames - 1) / 2.0])
-def _bbox_corners_mm(sras: SrasFile, angle_idx: int) -> np.ndarray:
- """4 corners (x, y) of this angle's raw mm bounding box, shape (4, 2)."""
- x = sras.x_axis_mm(angle_idx)
- y = sras.y_positions_mm(angle_idx)
- return np.array([[xx, yy] for xx in (x[0], x[-1]) for yy in (y[0], y[-1])])
+def ref_center_mm(sras: SrasFile, ref_angle_idx: int) -> np.ndarray:
+ """Stage mm of the reference angle's array center: ref mm + this == stage
+ mm.
+
+ The single bridge between ref mm and stage mm, and the only place in the
+ whole alignment path where any angle's stage position is read at all —
+ which is why it takes the *reference* index by name rather than an
+ arbitrary angle.
+ """
+ x = sras.x_axis_mm(ref_angle_idx)
+ y = sras.y_positions_mm(ref_angle_idx)
+ return np.array([(x[0] + x[-1]) / 2.0, (y[0] + y[-1]) / 2.0], dtype=np.float64)
+
+
+def _local_half_extent_mm(sras: SrasFile, angle_idx: int) -> tuple[float, float]:
+ """(half width, half height) in mm from this angle's array center to the
+ center of its outermost pixel."""
+ n_rows, n_frames = sras.image_shape(angle_idx)
+ dx, dy = pixel_pitch_mm(sras, angle_idx)
+ return (n_frames - 1) / 2.0 * abs(dx), (n_rows - 1) / 2.0 * abs(dy)
def _rotation_matrix(theta_deg: float) -> np.ndarray:
@@ -447,152 +1014,63 @@ def _rotation_matrix(theta_deg: float) -> np.ndarray:
return np.array([[c, -s], [s, c]]) # CCW rotation acting on (x, y)
-def _theta_deg(sras: SrasFile, angle_idx: int, ref_idx: int) -> float:
- """CCW rotation, in degrees and in _rotation_matrix's convention, that
- maps angle_idx's own local mm frame onto ref_idx's.
+def nominal_delta_deg(sras: SrasFile, angle_idx: int, ref_idx: int) -> float:
+ """The rotation-stage's own reported angle change between two angles.
- This is the *negative* of the raw angles_deg delta: the GR rotation
- stage's reported angle increases in the opposite rotational sense from
- this module's math-positive (CCW, x toward y) convention in scan mm
- space. Rotating by +(angles_deg[a] - angles_deg[ref]) therefore turns
- misalignment the wrong way — confirmed empirically (Auto De-rotate made
- real scans worse, not better, before this negation).
+ Used only to *seed* the rotation search, never as the answer: the stage's
+ sign convention relative to this module's math-positive (CCW, x toward y)
+ convention in scan mm is not knowable from the file, so
+ register_angle_to_reference scores both +this and -this and lets the image
+ content decide (see _rotation_candidates).
"""
- return -float(sras.angles_deg[angle_idx] - sras.angles_deg[ref_idx])
+ return float(sras.angles_deg[angle_idx] - sras.angles_deg[ref_idx])
-def _signal_centroid_mm(sras: SrasFile, angle_idx: int,
- dc4_mv: np.ndarray) -> tuple[float, float]:
- """Intensity-weighted centroid (mean x, mean y, weighted by CH4 signal
- after subtracting this angle's own minimum) in angle_idx's own local mm
- frame — the alignment pivot used in place of the raw scan-window bbox
- center (see compute_pivot_points_mm).
-
- Weighting by the continuous DC signal, rather than a binary >=
- dc_threshold_mv mask, means the pivot never depends on how well one
- shared threshold happens to suit this particular angle: real signal
- levels vary scan to scan, so a threshold tuned for one angle can leave
- another angle's binary mask empty — and a centroid of an empty mask has
- nothing to fall back to *except* the raw bbox center, silently
- reproducing the exact "aligned to the scan window, not the sample"
- problem this pivot exists to avoid. Falls back to the bbox center only
- in the fully-degenerate case of a perfectly flat signal (nothing to
- weight by at all).
- """
- weights = dc4_mv - dc4_mv.min()
- total = float(weights.sum())
- if total <= 0.0:
- return _bbox_center_mm(sras, angle_idx)
- x = sras.x_axis_mm(angle_idx)
- y = sras.y_positions_mm(angle_idx)
- cx = float((x * weights.sum(axis=0)).sum() / total)
- cy = float((y * weights.sum(axis=1)).sum() / total)
- return cx, cy
+def _footprint_corners_ref_mm(sras: SrasFile, angle_idx: int,
+ rotation_deg: float,
+ shift_mm: tuple[float, float]) -> np.ndarray:
+ """This angle's 4 footprint corners mapped into ref mm by its rigid
+ transform, shape (4, 2). Zero-padding cost, not content: it is the scan
+ window's corners, which is what the shared canvas has to cover."""
+ hw, hh = _local_half_extent_mm(sras, angle_idx)
+ corners = np.array([[sx * hw, sy * hh] for sx in (-1.0, 1.0) for sy in (-1.0, 1.0)])
+ R = _rotation_matrix(rotation_deg)
+ return corners @ R.T + np.asarray(shift_mm, dtype=np.float64)
-def compute_pivot_points_mm(sras: SrasFile,
- dc4_mv: dict[int, np.ndarray] | None = None
- ) -> dict[int, tuple[float, float]]:
- """Per-angle alignment pivot, in each angle's own local mm frame: the
- CH4-signal-weighted centroid of its own footprint (see
- _signal_centroid_mm), rather than the raw scan-window bbox center.
-
- Pivoting on each angle's own content — instead of on wherever its
- scanned window happened to sit in microscope/global XY space — is what
- makes alignment purely relative *between scans* rather than to a global
- coordinate system: the rotation stage's true mechanical axis need not
- coincide with the scan window's geometric center, and the sample need
- not be perfectly centered on that axis either, so a bbox-center pivot
- leaves a residual orbital motion between angles that a content-centroid
- pivot does not.
-
- Deliberately independent of dc_threshold_mv (the RF-mask / overlay
- threshold): that value is a display/masking choice and must never
- silently change where alignment pivots.
-
- *dc4_mv* lets a caller that has already computed each angle's CH4 mV
- image (compute_angle_alignment's Step 1, or ManualAlignmentDialog's own
- cache) reuse it instead of recomputing; angles missing from it are
- computed fresh via dc_image_mv, which prefers a stored v5/v7 cache over
- recomputing from raw waveforms.
- """
- dc4_mv = dc4_mv or {}
- pivots: dict[int, tuple[float, float]] = {}
- for a in range(sras.n_angles):
- img = dc4_mv.get(a)
- if img is None:
- img = dc_image_mv(sras, a, CH4_IDX)
- pivots[a] = _signal_centroid_mm(sras, a, img)
- return pivots
-
-
-def _corners_in_ref_frame(sras: SrasFile, angle_idx: int, ref_idx: int,
- pivot_mm: dict[int, tuple[float, float]],
- shift_mm=(0.0, 0.0),
- theta_deg: float | None = None) -> np.ndarray:
- """Angle *angle_idx*'s bbox corners, rotated about its own alignment
- pivot (pivot_mm[angle_idx] — see compute_pivot_points_mm) into the
- reference frame and translated by *shift_mm*. Shape (4, 2).
-
- theta_deg overrides the analytic angles_deg-derived rotation used by
- default — the manual-alignment path (union_canvas_mm) passes a
- user-chosen rotation here (which may differ from the known scan-angle
- delta) without needing a parallel code path.
- """
- theta = _theta_deg(sras, angle_idx, ref_idx) if theta_deg is None else theta_deg
- R = _rotation_matrix(theta)
- c_a = np.array(pivot_mm[angle_idx])
- c_ref = np.array(pivot_mm[ref_idx])
- shift = np.asarray(shift_mm, dtype=np.float64)
- return np.array([R @ (corner - c_a) + c_ref + shift
- for corner in _bbox_corners_mm(sras, angle_idx)])
-
-
-def _build_affine_canvas_to_raw(sras: SrasFile, angle_idx: int, ref_idx: int,
- shift_mm: tuple[float, float],
- canvas_dx: float, canvas_dy: float,
- canvas_origin_mm: tuple[float, float],
- pivot_mm: dict[int, tuple[float, float]],
- theta_deg: float | None = None
- ) -> tuple[np.ndarray, np.ndarray]:
- """matrix, offset s.t. raw_index = matrix @ [row_out, col_out] + offset,
+def _affine_out_to_src(*, out_pitch_mm: tuple[float, float],
+ out_origin_ref_mm, src_dx_mm: float, src_dy_mm: float,
+ src_center_idx, src_center_off_mm=(0.0, 0.0),
+ rotation_deg: float = 0.0,
+ shift_mm: tuple[float, float] = (0.0, 0.0)
+ ) -> tuple[np.ndarray, np.ndarray]:
+ """matrix, offset s.t. src_index = matrix @ [row_out, col_out] + offset,
matching scipy.ndimage.affine_transform's output->input convention.
- Pipeline (all mm unless noted):
- [X;Y] = A_out @ [row_out;col_out] + b_out # canvas idx -> ref-frame mm
- [lx;ly] = R(theta)^T @ ([X;Y]-c_ref-shift) + c_a # undo rotation+shift -> angle a's local mm
- [row;col] = D @ ([lx;ly] - [x_start_a; y0_a]) # local mm -> angle a's raw idx
+ The single affine builder behind every resampling in this module — the
+ registration grid, the manual-alignment preview and the final canvas all
+ differ only in their arguments.
- where theta is _theta_deg(angle_idx, ref_idx) unless *theta_deg*
- overrides it (see _corners_in_ref_frame's note, used by the manual-
- alignment path), c_ref/c_a are each angle's own alignment pivot
- (pivot_mm — see compute_pivot_points_mm; the CH4-signal-weighted
- centroid of its own footprint, not the raw scan-window bbox center —
- this keeps the sample itself centered post-rotation, minimizing
- required canvas padding and, more importantly, keeping alignment
- relative to the sample rather than to wherever the scan window sat in
- microscope/global XY space), and A_out/D are the index<->mm scaling
- matrices for the canvas pitch and this angle's own native pitch
- respectively.
+ Pipeline (mm unless noted):
+ [X;Y] = A_out @ [row_out;col_out] + out_origin_ref_mm # out idx -> ref mm
+ [lx;ly] = R(rotation)^T @ ([X;Y] - shift) # ref mm -> this angle's local mm
+ [row;col] = D @ ([lx;ly] - src_center_off) + src_center_idx
+
+ *src_center_off_mm* is the local-mm position of the source array's own
+ center, nonzero only when the source has been block-mean downsampled (its
+ center can land up to half a block off the full-resolution center — see
+ _prepare_reg_image). Everything is expressed relative to array centers, so
+ no per-angle stage coordinate appears anywhere in here.
"""
- theta = _theta_deg(sras, angle_idx, ref_idx) if theta_deg is None else theta_deg
- Rinv = _rotation_matrix(theta).T
- cx_a, cy_a = pivot_mm[angle_idx]
- cx_ref, cy_ref = pivot_mm[ref_idx]
- dx_a, dy_a = _pixel_pitch_mm(sras, angle_idx)
- x0_a = float(sras.x_start_mm[angle_idx])
- y0_a = float(sras.y_positions_mm(angle_idx)[0])
-
- A_out = np.array([[0.0, canvas_dx], [canvas_dy, 0.0]]) # [row,col] -> [X,Y]
- b_out = np.array(canvas_origin_mm, dtype=np.float64)
- D = np.array([[0.0, 1.0 / dy_a], [1.0 / dx_a, 0.0]]) # [x,y] -> [row,col]
- shift = np.array(shift_mm, dtype=np.float64)
- c_ref_v = np.array([cx_ref, cy_ref])
- c_a_v = np.array([cx_a, cy_a])
- origin_a = np.array([x0_a, y0_a])
+ Rinv = _rotation_matrix(rotation_deg).T
+ A_out = np.array([[0.0, out_pitch_mm[0]], [out_pitch_mm[1], 0.0]]) # [row,col] -> [x,y]
+ D = np.array([[0.0, 1.0 / src_dy_mm], [1.0 / src_dx_mm, 0.0]]) # [x,y] -> [row,col]
+ b_out = np.asarray(out_origin_ref_mm, dtype=np.float64)
+ shift = np.asarray(shift_mm, dtype=np.float64)
+ off = np.asarray(src_center_off_mm, dtype=np.float64)
matrix = D @ Rinv @ A_out
- offset = D @ Rinv @ (b_out - c_ref_v - shift) + D @ (c_a_v - origin_a)
+ offset = D @ (Rinv @ (b_out - shift) - off) + np.asarray(src_center_idx, dtype=np.float64)
return matrix, offset
@@ -602,7 +1080,7 @@ def apply_alignment(result: AlignmentResult, angle_idx: int, img: np.ndarray,
that angle's raw (n_rows, n_frames)) onto the shared alignment canvas.
order=0 (nearest) avoids blending real data with zero-padding or with
masked-out (0-valued) CH1/velocity pixels at mask edges. Channel-
- agnostic: the same per-angle transform (found from the CH4 mask) works
+ agnostic: the same per-angle transform (found from the CH4 image) works
for any channel's image of that angle."""
t = result.per_angle[angle_idx]
return scipy_ndimage.affine_transform(
@@ -611,122 +1089,695 @@ def apply_alignment(result: AlignmentResult, angle_idx: int, img: np.ndarray,
mode="constant", cval=0.0)
-def _block_mean_downsample(img: np.ndarray, factor: int) -> np.ndarray:
- if factor <= 1:
+def block_mean_2d(img: np.ndarray, fy: int, fx: int) -> np.ndarray:
+ """Block-mean by independent row/column factors. Independent factors matter
+ because the raw grid is strongly anisotropic (5 µm along x, 50 µm along y):
+ a single square factor would either alias along x or throw away rows."""
+ fy, fx = max(1, int(fy)), max(1, int(fx))
+ if fy == 1 and fx == 1:
return img
h, w = img.shape
- h2, w2 = (h // factor) * factor, (w // factor) * factor
+ h2, w2 = (h // fy) * fy, (w // fx) * fx
+ if h2 == 0 or w2 == 0:
+ return img
trimmed = img[:h2, :w2]
- return trimmed.reshape(h2 // factor, factor, w2 // factor, factor).mean(axis=(1, 3))
+ return trimmed.reshape(h2 // fy, fy, w2 // fx, fx).mean(axis=(1, 3))
-def _phase_correlate_shift(ref_img: np.ndarray, mov_img: np.ndarray) -> tuple[int, int]:
- """FFT normalized cross-power-spectrum phase correlation. Returns the
- integer (dr, dc) pixel shift of mov_img relative to ref_img; both must
- be the same shape. Risk: if the true shift is near +/- half the array
- size, wraparound can bias the peak — mitigated by the generous
- margin_frac padding in _working_canvas_for_pair, which keeps the true
- residual shift small relative to the correlation canvas."""
- F1 = scipy_fft.fft2(ref_img.astype(np.float64), workers=-1)
- F2 = scipy_fft.fft2(mov_img.astype(np.float64), workers=-1)
- R = F1 * np.conj(F2)
- R /= np.maximum(np.abs(R), 1e-12)
- corr = scipy_fft.ifft2(R, workers=-1).real
- dr, dc = np.unravel_index(np.argmax(corr), corr.shape)
- h, w = corr.shape
- if dr > h // 2:
- dr -= h
- if dc > w // 2:
- dc -= w
- return int(dr), int(dc)
+# ---- Registration: rigid (rotation + translation) fit against the reference
+#
+# Everything below works on _RegImage, an angle's image resampled to a shared
+# *isotropic* grid centered on its own array center. Nothing in here can see a
+# stage coordinate even in principle, which is the point: the fit is decided by
+# image content alone.
+
+@dataclass
+class _RegImage:
+ """One angle's image prepared for registration: block-mean downsampled to
+ roughly the registration pitch, carrying the physical pitch it ended up
+ with and the local-mm offset of its own array center from the
+ full-resolution array center (block-mean trims a partial trailing block, so
+ the two centers can differ by up to half a block)."""
+ img: np.ndarray
+ dx_mm: float
+ dy_mm: float
+ center_off_mm: tuple[float, float]
-def _working_canvas_for_pair(sras: SrasFile, ref_idx: int, a_idx: int,
- dx: float, dy: float,
- pivot_mm: dict[int, tuple[float, float]],
- margin_frac: float = 0.3
- ) -> tuple[tuple[float, float], tuple[int, int]]:
- """Union of the reference's own raw bbox and angle a's raw bbox rotated
- (about its own alignment pivot) into the ref frame with zero shift,
- padded by margin_frac on each side — sized generously so the true
- phase-correlation shift lands well inside the canvas (see
- _phase_correlate_shift's wraparound note)."""
- pts = np.vstack([_bbox_corners_mm(sras, ref_idx),
- _corners_in_ref_frame(sras, a_idx, ref_idx, pivot_mm)])
- x_min, y_min = pts.min(axis=0)
- x_max, y_max = pts.max(axis=0)
- pad_x, pad_y = (x_max - x_min) * margin_frac, (y_max - y_min) * margin_frac
- x_min, x_max = x_min - pad_x, x_max + pad_x
- y_min, y_max = y_min - pad_y, y_max + pad_y
- n_cols = int(np.ceil((x_max - x_min) / dx)) + 1
- n_rows = int(np.ceil((y_max - y_min) / abs(dy))) + 1
- origin = (x_min, y_min if dy > 0 else y_max)
- return origin, (n_rows, n_cols)
+def _block_center_off(n_full: int, n_small: int, factor: int,
+ pitch_mm: float) -> float:
+ """Local-mm offset of a block-mean-downsampled array's own center from the
+ full-resolution array's center along one axis.
+
+ Small pixel k averages full pixels [k*f, k*f + f - 1], so its center sits
+ at full index k*f + (f-1)/2; block-mean also trims a partial trailing
+ block. Both effects together move the small array's center by up to half a
+ block, which has to be accounted for or a downsampled image registers (or
+ previews) at a systematically shifted position.
+ """
+ return (((n_small - 1) / 2.0) * factor + (factor - 1) / 2.0
+ - (n_full - 1) / 2.0) * pitch_mm
-def correlate_translation_mm(sras: SrasFile, angle_idx: int, ref_angle_idx: int,
- signal_mv: dict[int, np.ndarray],
- pivot_mm: dict[int, tuple[float, float]], *,
- use_mask: bool = False,
- dc_threshold_mv: float = 0.0,
- margin_frac: float = 0.3,
- max_corr_dim: int = 1024) -> tuple[float, float]:
- """FFT phase-correlation translation (shift_x_mm, shift_y_mm) that best
- lines up angle_idx's CH4 content onto ref_angle_idx's, given each angle's
- own alignment pivot (pivot_mm — see compute_pivot_points_mm) already
- rotated about that pivot by the analytic scan-angle delta with zero
- shift. Returns (0.0, 0.0) unconditionally for the reference angle.
+def _prepare_reg_image(sras: SrasFile, angle_idx: int, img: np.ndarray,
+ pitch_mm: float) -> _RegImage:
+ """Block-mean an angle's image down to roughly *pitch_mm* before it is
+ resampled onto the registration grid. Pre-averaging matters: the raw grid
+ is 10x finer along x than along y, so sampling it directly at the (much
+ coarser) isotropic registration pitch would alias badly along x."""
+ dx, dy = pixel_pitch_mm(sras, angle_idx)
+ fx = max(1, int(pitch_mm / abs(dx)))
+ fy = max(1, int(pitch_mm / abs(dy)))
+ small = block_mean_2d(np.asarray(img, dtype=np.float32), fy, fx)
+ n_rows, n_frames = img.shape
+ return _RegImage(
+ small, dx * fx, dy * fy,
+ (_block_center_off(n_frames, small.shape[1], fx, dx),
+ _block_center_off(n_rows, small.shape[0], fy, dy)))
- By default (use_mask=False) correlates on each angle's own raw CH4
- signal minus its own minimum — the same weighting _signal_centroid_mm
- uses — rather than a dc_threshold_mv binary mask: a shared threshold
- that doesn't suit every angle's real signal level can empty or distort
- one angle's mask (the same failure mode the alignment pivot was fixed
- to avoid), and correlating on the raw signal also lets phase
- correlation lock onto real internal sample structure rather than just
- the scan window's silhouette. Subtracting each angle's own minimum
- (rather than using the signal as-is) keeps the zero-padding surrounding
- the rotated image from reading as a spurious high-contrast edge against
- a nonzero DC baseline. use_mask=True switches to the binary
- >= dc_threshold_mv mask instead (compute_angle_alignment's original
- behavior), for cases where raw-signal correlation locks onto noise.
+
+def _embed(reg: _RegImage, pitch_mm: float, n: int, rotation_deg: float,
+ order: int) -> np.ndarray:
+ """Resample a _RegImage onto the shared n x n isotropic registration grid,
+ rotated by *rotation_deg* about the grid center and with no translation
+ (translation is what phase correlation then measures)."""
+ half = (n - 1) / 2.0 * pitch_mm
+ matrix, offset = _affine_out_to_src(
+ out_pitch_mm=(pitch_mm, pitch_mm), out_origin_ref_mm=(-half, -half),
+ src_dx_mm=reg.dx_mm, src_dy_mm=reg.dy_mm,
+ src_center_idx=((reg.img.shape[0] - 1) / 2.0, (reg.img.shape[1] - 1) / 2.0),
+ src_center_off_mm=reg.center_off_mm, rotation_deg=rotation_deg)
+ return scipy_ndimage.affine_transform(
+ reg.img, matrix, offset=offset, output_shape=(n, n), order=order,
+ mode="constant", cval=0.0)
+
+
+def _embed_with_valid(reg: _RegImage, pitch_mm: float, n: int,
+ rotation_deg: float) -> tuple[np.ndarray, np.ndarray]:
+ """Embedded image plus the boolean mask of where that angle actually has
+ data. The valid mask is what lets registration ignore each angle's
+ differently-shaped scan window instead of locking onto its silhouette."""
+ img = _embed(reg, pitch_mm, n, rotation_deg, order=1)
+ ones = _RegImage(np.ones_like(reg.img), reg.dx_mm, reg.dy_mm, reg.center_off_mm)
+ valid = _embed(ones, pitch_mm, n, rotation_deg, order=0) > 0.5
+ return img, valid
+
+
+def _shift_into(img: np.ndarray, dr: int, dc: int) -> np.ndarray:
+ """img translated by whole pixels with zero fill (never wrapping, unlike
+ np.roll — wrapped content would score as a spurious match)."""
+ out = np.zeros_like(img)
+ h, w = img.shape
+ sr0, sr1 = max(0, dr), min(h, h + dr)
+ sc0, sc1 = max(0, dc), min(w, w + dc)
+ if sr0 >= sr1 or sc0 >= sc1:
+ return out
+ out[sr0:sr1, sc0:sc1] = img[sr0 - dr:sr1 - dr, sc0 - dc:sc1 - dc]
+ return out
+
+
+def _overlap_ncc(ref: np.ndarray, ref_valid: np.ndarray,
+ mov: np.ndarray, mov_valid: np.ndarray,
+ min_overlap_frac: float = 0.15) -> float:
+ """Zero-mean normalized cross-correlation over the two images' common valid
+ region — the score every rotation candidate is ranked by.
+
+ Computed on the overlap only, and rejected outright (-1) when the overlap
+ is too small a fraction of the smaller footprint: without that floor a
+ candidate that slides the angles almost entirely apart can win on a handful
+ of coincidentally-similar pixels.
+ """
+ both = ref_valid & mov_valid
+ n = int(both.sum())
+ smaller = min(int(ref_valid.sum()), int(mov_valid.sum()))
+ if smaller == 0 or n < min_overlap_frac * smaller or n < 16:
+ return -1.0
+ a = ref[both].astype(np.float64)
+ b = mov[both].astype(np.float64)
+ a -= a.mean()
+ b -= b.mean()
+ denom = np.sqrt((a * a).sum() * (b * b).sum())
+ return float((a * b).sum() / denom) if denom > 0 else -1.0
+
+
+def _masked_shift(ref: np.ndarray, ref_valid: np.ndarray,
+ mov: np.ndarray, mov_valid: np.ndarray) -> tuple[int, int]:
+ """Integer (dr, dc) that best registers *mov* onto *ref*, from skimage's
+ masked FFT phase correlation (Padfield). The masked variant is the whole
+ reason scikit-image is a dependency: plain phase correlation on these
+ images locks onto the scan window's rectangular silhouette, which differs
+ per angle, instead of onto the sample."""
+ pcc = _skimage_phase_cross_correlation()
+ result = pcc(ref, mov, reference_mask=ref_valid, moving_mask=mov_valid)
+ shift = result[0] if isinstance(result, tuple) else result
+ return int(round(float(shift[0]))), int(round(float(shift[1])))
+
+
+def _subpixel_residual(ref: np.ndarray, mov: np.ndarray,
+ both: np.ndarray) -> tuple[float, float]:
+ """Sub-pixel leftover shift between two already integer-aligned images,
+ from upsampled phase correlation over their common valid region.
+
+ A separate pass because skimage's *masked* phase correlation has no
+ upsample_factor; here both images are zeroed outside the shared overlap and
+ mean-subtracted inside it, so the plain upsampled version is well posed.
+ Clamped to ±1 px: this only ever polishes an already-good integer fit, and
+ a larger "residual" means the peak was spurious.
+
+ normalization=None (plain cross-correlation, not phase correlation) is
+ load-bearing. Whitening the spectrum is what makes phase correlation good
+ at finding a large unknown shift, but here the two images are already
+ aligned to within a pixel and the masked-off surroundings put a hard edge
+ in both: whitened, that edge and the high-frequency noise swamp the true
+ sub-pixel peak and the default returns a flat zero every time.
+ """
+ if not both.any():
+ return 0.0, 0.0
+ a = np.zeros_like(ref)
+ b = np.zeros_like(mov)
+ a[both] = ref[both] - ref[both].mean()
+ b[both] = mov[both] - mov[both].mean()
+ if not (np.any(a) and np.any(b)):
+ return 0.0, 0.0
+ pcc = _skimage_phase_cross_correlation()
+ result = pcc(a, b, upsample_factor=20, normalization=None)
+ shift = result[0] if isinstance(result, tuple) else result
+ dr, dc = float(shift[0]), float(shift[1])
+ if abs(dr) > 1.0 or abs(dc) > 1.0:
+ return 0.0, 0.0
+ return dr, dc
+
+
+def _reg_pitch_and_size(sras: SrasFile, max_dim: int,
+ margin: float = 1.25) -> tuple[float, int]:
+ """Isotropic pitch (mm/px) and side length for the shared registration
+ grid: square, big enough for the largest angle's footprint at any rotation
+ (hence its diagonal) plus *margin* headroom for the translation search.
+
+ The floor on pitch is the *geometric mean* of the two native pitches, not
+ the coarser of them. The raw grid is anisotropic (5 µm along x, 50 µm along
+ y): flooring at 50 µm would throw away all the extra x detail, and rotation
+ precision depends directly on it — a feature at radius r moves by r·δθ, so
+ at 50 µm a 20 mm-wide sample can only resolve rotation to a few tenths of a
+ degree. The geometric mean interpolates y up by ~3x rather than discarding
+ x, which costs a little memory and buys real angular precision. On a
+ full-size scan max_dim binds first and this floor never applies at all.
+ """
+ diag = max(float(np.hypot(*(2 * v for v in _local_half_extent_mm(sras, a))))
+ for a in range(sras.n_angles))
+ span = diag * margin
+ native = float(np.sqrt(
+ abs(sras.pixel_x_mm)
+ * max(abs(pixel_pitch_mm(sras, a)[1]) for a in range(sras.n_angles))))
+ pitch = max(span / max_dim, native)
+ n = int(scipy_fft.next_fast_len(max(16, int(np.ceil(span / pitch)))))
+ return pitch, n
+
+
+def _registration_workers(sras: SrasFile, fine_dim: int) -> int:
+ """How many angles may register concurrently.
+
+ Not plan_angle_level: that budgets for waveform chunks, and registration
+ never touches a waveform — it works on already-computed DC images and a few
+ grid-sized arrays. The real limit is the masked phase correlation, which
+ pads to roughly twice the grid and holds several complex128 arrays of that
+ size live at once, so the count is derived from the same
+ _TOTAL_BYTES_BUDGET the rest of the module honours. The estimate below
+ comes out at 262 MB for the default fine_dim=640, against 199 MB measured —
+ deliberately on the pessimistic side, since overshooting the budget costs
+ swapping while undershooting only costs a little wall time.
+ """
+ per_worker = 10 * (2 * fine_dim) ** 2 * 16 # ~10 complex128 grids
+ return int(max(1, min(_MAX_WORKERS, sras.n_angles,
+ _TOTAL_BYTES_BUDGET // max(1, per_worker))))
+
+
+def registration_workers(sras: SrasFile) -> int:
+ """Public: concurrent-registration cap at the default fine grid."""
+ return _registration_workers(sras, _DEFAULT_FINE_DIM)
+
+
+def default_max_workers() -> int:
+ """Public: the module-wide worker cap (SRAS_MAX_WORKERS or cpu count)."""
+ return _MAX_WORKERS
+
+
+def _rotation_candidates(nominal_deg: float, search_deg: float,
+ step_deg: float,
+ signs: tuple[int, ...] = (-1, 1)) -> list[float]:
+ """Coarse rotation candidates: a window around each requested sign of the
+ stage's reported angle change. Scoring both signs (the default) is what
+ makes the stage's sign convention a non-issue — the images decide which way
+ the stage turns, and a file whose stage reports the opposite sense
+ registers just as well.
+
+ *signs* exists only so a user who has established which way their own stage
+ turns can halve the coarse sweep from the alignment wizard. It is an
+ override, not an inference: nothing in the file says which sign is right.
+ """
+ out: list[float] = []
+ step_deg = max(abs(step_deg), 1e-9) # a 0 step would divide by zero below
+ for sign in signs:
+ center = sign * nominal_deg
+ k = int(np.floor(search_deg / step_deg))
+ for i in range(-k, k + 1):
+ out.append(center + i * step_deg)
+ # Dedupe (the two windows coincide when nominal_deg is 0) while keeping order.
+ seen: set[float] = set()
+ return [t for t in out if not (round(t, 6) in seen or seen.add(round(t, 6)))]
+
+
+def _score_rotation(ref_img: np.ndarray, ref_valid: np.ndarray,
+ mov: _RegImage, pitch: float, n: int, theta: float,
+ subpixel: bool = False) -> tuple[float, tuple[float, float]]:
+ """Best score this rotation can reach, and the translation that reaches it:
+ rotate, phase-correlate for the shift, score the overlap.
+
+ *subpixel* also removes the leftover sub-pixel translation before scoring.
+ That matters more than it sounds: without it every candidate is scored at
+ whole-pixel alignment, so rotations that differ by less than one pixel of
+ rim displacement are ranked by quantization noise rather than by fit, and
+ the refinement stalls a degree or so off. Left off for the coarse sweep,
+ which only has to pick a basin, and on for the refinement.
+ """
+ img, valid = _embed_with_valid(mov, pitch, n, theta)
+ dr, dc = _masked_shift(ref_img, ref_valid, img, valid)
+ shifted = _shift_into(img, dr, dc)
+ shifted_valid = _shift_into(valid.astype(np.float32), dr, dc) > 0.5
+ if subpixel:
+ sub_dr, sub_dc = _subpixel_residual(
+ ref_img, shifted, ref_valid & shifted_valid)
+ if sub_dr or sub_dc:
+ shifted = scipy_ndimage.shift(shifted, (sub_dr, sub_dc), order=1,
+ mode="constant", cval=0.0)
+ dr, dc = dr + sub_dr, dc + sub_dc
+ return (_overlap_ncc(ref_img, ref_valid, shifted, shifted_valid),
+ (float(dr), float(dc)))
+
+
+def _refine_rotation(ref_img: np.ndarray, ref_valid: np.ndarray,
+ mov: _RegImage, pitch: float, n: int,
+ theta: float, score: float, shift: tuple[float, float],
+ step_deg: float, min_step_deg: float = 0.05,
+ max_evals: int = 80
+ ) -> tuple[float, float, tuple[float, float]]:
+ """Hill-climb the rotation from the coarse winner: step out while the score
+ improves, halve the step when it doesn't, stop below *min_step_deg*.
+
+ A walking search rather than a fixed grid around the coarse winner, because
+ the coarse stage ranks rotations at a coarse pitch where a degree can be
+ worth less than the translation quantization — its winner can legitimately
+ land a degree or two off, which a fixed ±½° refinement window could never
+ recover from.
+
+ Every candidate here is scored with subpixel=True, matching how the
+ incoming *score* was measured. Mixing the two is silently fatal: the
+ sub-pixel-corrected score is strictly the higher of the two, so a
+ subpixel-scored start compared against un-corrected candidates can never be
+ beaten and the search sits still at whatever the coarse stage handed it.
+ """
+ evals = 0
+ while step_deg >= min_step_deg and evals < max_evals:
+ trials = []
+ for cand in (theta - step_deg, theta + step_deg):
+ s, sh = _score_rotation(ref_img, ref_valid, mov, pitch, n, cand,
+ subpixel=True)
+ evals += 1
+ trials.append((s, cand, sh))
+ best_s, best_cand, best_sh = max(trials, key=lambda t: t[0])
+ if best_s > score:
+ theta, score, shift = best_cand, best_s, best_sh
+ else:
+ step_deg /= 2.0
+ return theta, score, shift
+
+
+def _source_images(sras: SrasFile, angle_idx: int, ref_angle_idx: int,
+ signal_mv: dict[int, np.ndarray], dc_threshold_mv: float,
+ sources) -> list[tuple[str, np.ndarray, np.ndarray]]:
+ """(name, reference image, moving image) per requested registration source.
+
+ "signal" is each angle's own CH4 image minus its own minimum — subtracting
+ per-angle rather than globally keeps a nonzero DC baseline from reading as
+ a high-contrast edge against the zero padding. "mask" is the binarized
+ >= dc_threshold_mv image, the same silhouette the overlay draws. A mask
+ that is empty or completely full for either angle carries no registration
+ information at all, so that source is dropped rather than scored.
+ """
+ out = []
+ for name in sources:
+ pair = []
+ for a in (ref_angle_idx, angle_idx):
+ img = signal_mv[a]
+ if name == "mask":
+ m = img >= dc_threshold_mv
+ if not m.any() or m.all():
+ pair = []
+ break
+ pair.append(m.astype(np.float32))
+ else:
+ pair.append((img - img.min()).astype(np.float32))
+ if pair:
+ out.append((name, pair[0], pair[1]))
+ return out
+
+
+def register_angle_to_reference(
+ sras: SrasFile, angle_idx: int, ref_angle_idx: int,
+ signal_mv: dict[int, np.ndarray], *,
+ dc_threshold_mv: float = 0.0,
+ sources: tuple[str, ...] = ("signal", "mask"),
+ coarse_dim: int = 256, fine_dim: int = _DEFAULT_FINE_DIM,
+ search_deg: float = 6.0, coarse_step_deg: float = 2.0,
+ seed_deg: float | None = None,
+ seed_signs: tuple[int, ...] = (-1, 1),
+ refine: bool = True) -> RigidFit:
+ """Rigid (rotation + translation, never scale) fit of *angle_idx* onto
+ *ref_angle_idx*, found entirely by cross-correlating image content.
+
+ Two stages:
+ 1. Coarse sweep at *coarse_dim* over a ±*search_deg* window around the
+ requested signs of the stage's reported angle change (see
+ _rotation_candidates), for each requested source image, scored by
+ _overlap_ncc. Only has to pick the right basin.
+ 2. Hill-climbing refinement of the winning (source, rotation) at
+ *fine_dim* with sub-pixel translation folded into every score
+ (_refine_rotation, _score_rotation), down to 0.05°.
+
+ The returned rotation_deg/shift_mm are the rigid map from this angle's
+ local mm to ref mm (q = R @ l + shift); score is the final NCC, which the
+ caller can surface so a bad scan is visible rather than silently fused in.
+ Returns identity for the reference angle itself.
+
+ The last three arguments only exist to let the alignment wizard expose the
+ rotation search; every default reproduces the search this function has
+ always done.
+
+ *seed_deg* replaces the stage's reported angle change as the center of the
+ coarse sweep — the pre-rotation the search starts from. None (the default)
+ means the stage angle from the file's Angle Table, which is nearly always
+ what you want: it puts the sweep within a couple of degrees of the answer.
+ Pass 0.0 to search around no rotation at all, which is the honest choice for
+ a file whose stage angles are known to be wrong.
+
+ *refine=False* stops after the coarse sweep, leaving rotation on the
+ *coarse_step_deg* grid. Combined with search_deg=0.0 and a single seed sign
+ it pins rotation to exactly the seed and searches translation only — for a
+ scan whose stage angles are trusted more than its image content.
"""
if angle_idx == ref_angle_idx:
- return (0.0, 0.0)
+ return RigidFit(0.0, (0.0, 0.0), 1.0, "reference")
- def corr_img(a: int) -> np.ndarray:
- img = signal_mv[a]
- if use_mask:
- return (img >= dc_threshold_mv).astype(np.float32)
- return (img - img.min()).astype(np.float32)
+ candidates = _source_images(sras, angle_idx, ref_angle_idx, signal_mv,
+ dc_threshold_mv, sources)
+ if not candidates:
+ return RigidFit(0.0, (0.0, 0.0), -1.0, "none")
- img_a = corr_img(angle_idx)
- img_ref = corr_img(ref_angle_idx)
+ nominal = (nominal_delta_deg(sras, angle_idx, ref_angle_idx)
+ if seed_deg is None else float(seed_deg))
+ thetas = _rotation_candidates(nominal, search_deg, coarse_step_deg,
+ seed_signs)
- dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx)
- max_dim = max(img_a.shape + img_ref.shape)
- factor = max(1, int(np.ceil(max_dim / max_corr_dim)))
- dx_c, dy_c = dx_ref * factor, dy_ref * factor
- small_a = _block_mean_downsample(img_a, factor)
- small_ref = _block_mean_downsample(img_ref, factor)
+ # ---- Stage 1: coarse sweep, every source ------------------------------
+ pitch_c, n_c = _reg_pitch_and_size(sras, coarse_dim)
+ best = (-2.0, 0.0, (0, 0), "none") # score, theta, (dr, dc), source
+ for name, ref_raw, mov_raw in candidates:
+ ref_reg = _prepare_reg_image(sras, ref_angle_idx, ref_raw, pitch_c)
+ mov_reg = _prepare_reg_image(sras, angle_idx, mov_raw, pitch_c)
+ ref_img, ref_valid = _embed_with_valid(ref_reg, pitch_c, n_c, 0.0)
+ for theta in thetas:
+ score, shift = _score_rotation(ref_img, ref_valid, mov_reg,
+ pitch_c, n_c, theta)
+ if score > best[0]:
+ best = (score, theta, shift, name)
- work_origin, work_shape = _working_canvas_for_pair(
- sras, ref_angle_idx, angle_idx, dx_c, dy_c, pivot_mm, margin_frac=margin_frac)
+ if best[3] == "none":
+ return RigidFit(0.0, (0.0, 0.0), -1.0, "none")
- m_a, o_a = _build_affine_canvas_to_raw(
- sras, angle_idx, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin, pivot_mm)
- m_ref, o_ref = _build_affine_canvas_to_raw(
- sras, ref_angle_idx, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin, pivot_mm)
+ # ---- Stage 2: refine the winner at full registration resolution -------
+ name = best[3]
+ ref_raw, mov_raw = next((r, m) for n_, r, m in candidates if n_ == name)
+ pitch_f, n_f = _reg_pitch_and_size(sras, fine_dim)
+ ref_reg = _prepare_reg_image(sras, ref_angle_idx, ref_raw, pitch_f)
+ mov_reg = _prepare_reg_image(sras, angle_idx, mov_raw, pitch_f)
+ ref_img, ref_valid = _embed_with_valid(ref_reg, pitch_f, n_f, 0.0)
- rotated_a = scipy_ndimage.affine_transform(
- small_a, m_a / factor, offset=o_a / factor,
- output_shape=work_shape, order=0, mode="constant", cval=0.0)
- embedded_ref = scipy_ndimage.affine_transform(
- small_ref, m_ref / factor, offset=o_ref / factor,
- output_shape=work_shape, order=0, mode="constant", cval=0.0)
+ theta = best[1]
+ score, shift = _score_rotation(ref_img, ref_valid, mov_reg, pitch_f, n_f,
+ theta, subpixel=True)
+ if refine:
+ theta, score, shift = _refine_rotation(
+ ref_img, ref_valid, mov_reg, pitch_f, n_f, theta, score, shift,
+ step_deg=coarse_step_deg)
- dr, dc = _phase_correlate_shift(embedded_ref, rotated_a)
- return (dc * dx_c, dr * dy_c)
+ dr, dc = shift
+ return RigidFit(float(theta), (float(dc * pitch_f), float(dr * pitch_f)),
+ float(score), name)
+
+
+# ---- Shared canvas: angle 0's own pixel grid, extended --------------------
+
+def canvas_for_params(sras: SrasFile, ref_angle_idx: int,
+ pitch_mm: tuple[float, float],
+ per_angle_params: dict[int, ManualAngleParams],
+ ) -> tuple[tuple[float, float], tuple[int, int]]:
+ """Shared-canvas origin (stage mm) and (n_rows, n_cols) at *pitch_mm* that
+ contains every angle's footprint after its own rigid transform. Angles
+ missing from per_angle_params default to identity (e.g. a sidecar saved
+ before a rescan added more angles).
+
+ The canvas grid is aligned with the reference angle's own pixel grid, so
+ the reference lands on integer canvas pixels and is resampled by an exact
+ integer translation — the concrete meaning of "the canvas carries angle
+ 0's X/Y coordinates". This requires pitch_mm to be the reference's own
+ pitch.
+ """
+ dx, dy = pitch_mm
+ corners = np.vstack([
+ _footprint_corners_ref_mm(
+ sras, a,
+ per_angle_params.get(a, ManualAngleParams()).rotation_deg,
+ per_angle_params.get(a, ManualAngleParams()).shift_mm)
+ for a in range(sras.n_angles)])
+ x_min, y_min = corners.min(axis=0)
+ x_max, y_max = corners.max(axis=0)
+ center = ref_center_mm(sras, ref_angle_idx)
+ # Express the box in the reference's own pixel indices and grow it
+ # outward to whole pixels, so canvas index k lands exactly where the
+ # reference's own pixel (k + const) does.
+ cy, cx = _center_idx(sras, ref_angle_idx)
+ cols = sorted((x_min / dx + cx, x_max / dx + cx))
+ rows = sorted((y_min / dy + cy, y_max / dy + cy))
+ col0, col1 = int(np.floor(cols[0])), int(np.ceil(cols[1]))
+ row0, row1 = int(np.floor(rows[0])), int(np.ceil(rows[1]))
+ origin_ref = np.array([(col0 - cx) * dx, (row0 - cy) * dy])
+ shape = (row1 - row0 + 1, col1 - col0 + 1)
+
+ origin_stage = origin_ref + center
+ return (float(origin_stage[0]), float(origin_stage[1])), shape
+
+
+def build_canvas_affine(sras: SrasFile, angle_idx: int, ref_angle_idx: int,
+ rotation_deg: float, shift_mm: tuple[float, float],
+ pitch_mm: tuple[float, float],
+ canvas_origin_mm: tuple[float, float],
+ *, src_downsample: tuple[int, int] = (1, 1)
+ ) -> tuple[np.ndarray, np.ndarray]:
+ """canvas index -> this angle's raw index, for a canvas whose origin is
+ given in *stage* mm (the reference's frame). *src_downsample* is the
+ (rows, cols) block-mean factor already applied to the image the caller will
+ resample — 1:1 for the raw image, coarser for the manual-alignment
+ preview's downsampled masks."""
+ dx_a, dy_a = pixel_pitch_mm(sras, angle_idx)
+ n_rows, n_frames = sras.image_shape(angle_idx)
+ fy, fx = (max(1, int(v)) for v in src_downsample)
+ if (fy, fx) != (1, 1):
+ # Same block-mean bookkeeping _prepare_reg_image does: the downsampled
+ # array's own center can sit up to half a block off the full-resolution
+ # center, and that offset has to be undone here or every preview layer
+ # lands slightly (and inconsistently) off.
+ nr_s, nf_s = n_rows // fy, n_frames // fx
+ src_dx, src_dy = dx_a * fx, dy_a * fy
+ src_center = ((nr_s - 1) / 2.0, (nf_s - 1) / 2.0)
+ src_off = (_block_center_off(n_frames, nf_s, fx, dx_a),
+ _block_center_off(n_rows, nr_s, fy, dy_a))
+ else:
+ src_dx, src_dy = dx_a, dy_a
+ src_center = _center_idx(sras, angle_idx)
+ src_off = (0.0, 0.0)
+
+ origin_ref = np.asarray(canvas_origin_mm, dtype=np.float64) \
+ - ref_center_mm(sras, ref_angle_idx)
+ return _affine_out_to_src(
+ out_pitch_mm=pitch_mm, out_origin_ref_mm=origin_ref,
+ src_dx_mm=src_dx, src_dy_mm=src_dy, src_center_idx=src_center,
+ src_center_off_mm=src_off, rotation_deg=rotation_deg, shift_mm=shift_mm)
+
+
+def _result_from_params(sras: SrasFile, ref_angle_idx: int,
+ dc_threshold_mv: float,
+ params: dict[int, ManualAngleParams],
+ extra: dict[int, tuple[float, str]] | None = None
+ ) -> AlignmentResult:
+ """Assemble the final AlignmentResult from per-angle rigid parameters: pick
+ the shared canvas, then build each angle's canvas->raw affine. Pure matrix
+ and bbox math, so it is cheap enough to call synchronously on the GUI
+ thread on every manual edit."""
+ pitch = pixel_pitch_mm(sras, ref_angle_idx)
+ canvas_origin_mm, canvas_shape = canvas_for_params(
+ sras, ref_angle_idx, pitch, params)
+
+ extra = extra or {}
+ per_angle: dict[int, AngleTransform] = {}
+ for a in range(sras.n_angles):
+ p = params.get(a, ManualAngleParams())
+ matrix, offset = build_canvas_affine(
+ sras, a, ref_angle_idx, p.rotation_deg, p.shift_mm,
+ pitch, canvas_origin_mm)
+ score, source = extra.get(a, (1.0, ""))
+ per_angle[a] = AngleTransform(p.rotation_deg, p.shift_mm, matrix, offset,
+ score, source)
+
+ return AlignmentResult(ref_angle_idx, dc_threshold_mv, canvas_shape,
+ pitch[0], pitch[1], canvas_origin_mm, per_angle)
+
+
+def crop_alignment_result(result: AlignmentResult, row0: int, col0: int,
+ n_rows: int, n_cols: int) -> AlignmentResult:
+ """The same alignment restricted to a rectangular window of its canvas —
+ canvas pixel (row0, col0) becomes the cropped canvas's (0, 0).
+
+ Cropping folds into each angle's existing affine instead of becoming a
+ second transform, because a canvas crop is *pure index translation*. From
+ _affine_out_to_src, matrix = D @ Rinv @ A_out depends only on pitch and
+ rotation, and A_out @ [row0, col0] is exactly the mm displacement of the new
+ origin, so:
+
+ matrix @ [r', c'] + (offset + matrix @ [row0, col0])
+ == matrix @ [r' + row0, c' + col0] + offset
+
+ i.e. shifting the offset by matrix @ [row0, col0] reproduces the original
+ mapping at the shifted indices, exactly. That equality is what lets
+ apply_alignment, reproject_mask and the aligned .sras exporter all keep
+ working on a cropped result with no special-casing — and it is why the crop
+ preview a user approves is guaranteed to be the same pixels the exporter
+ writes.
+
+ Bounds are the caller's responsibility (the wizard's ROI page clamps to the
+ canvas): an out-of-range window is geometrically well-defined here and
+ simply resamples padding.
+ """
+ if n_rows <= 0 or n_cols <= 0:
+ raise ValueError(f"empty crop: {n_rows} x {n_cols}")
+
+ delta = np.array([float(row0), float(col0)])
+ per_angle = {
+ a: AngleTransform(t.rotation_deg, t.shift_mm, t.matrix,
+ t.offset + t.matrix @ delta, t.score, t.source)
+ for a, t in result.per_angle.items()
+ }
+ origin = (result.canvas_origin_mm[0] + col0 * result.canvas_dx_mm,
+ result.canvas_origin_mm[1] + row0 * result.canvas_dy_mm)
+ return AlignmentResult(result.ref_angle_idx, result.dc_threshold_mv,
+ (int(n_rows), int(n_cols)),
+ result.canvas_dx_mm, result.canvas_dy_mm,
+ origin, per_angle)
+
+
+def coarse_rect_to_canvas(rect: tuple[int, int, int, int],
+ downsample: tuple[int, int],
+ canvas_shape: tuple[int, int],
+ *, inset_blocks: int = 0
+ ) -> tuple[int, int, int, int]:
+ """A rectangle in coarse (block-mean preview) indices as canvas pixels,
+ both as (row0, col0, n_rows, n_cols) and clamped to the canvas.
+
+ The single place the preview grid's relation to the real canvas is
+ written down: the coarse grid samples canvas pixels 0, f, 2f, …, so a
+ coarse pixel stands for a whole block and every caller has to agree on
+ which canvas pixels that block means, or a crop the user approved on the
+ preview lands a few pixels off in the export.
+
+ *inset_blocks* shrinks the rectangle by that many coarse blocks on each
+ side. A coarse pixel reported as fully covered stands for a block whose
+ far edge may not be, so "fit to full overlap" insets by 1 to stay honestly
+ inside the overlap region; a union rectangle insets by 0 because it wants
+ to contain the region rather than fit inside it.
+ """
+ row0, col0, nr, nc = rect
+ fy, fx = downsample
+ n_rows, n_cols = canvas_shape
+ r0 = min(n_rows - 1, (row0 + inset_blocks) * fy)
+ c0 = min(n_cols - 1, (col0 + inset_blocks) * fx)
+ r1 = min(n_rows, (row0 + nr - inset_blocks) * fy)
+ c1 = min(n_cols, (col0 + nc - inset_blocks) * fx)
+ return r0, c0, max(1, r1 - r0), max(1, c1 - c0)
+
+
+def overlap_stats(counts: np.ndarray, n_angles: int) -> dict:
+ """Summarize a per-pixel "how many angles cover this pixel" image — the
+ number the alignment wizard's mask-stack view is colored by.
+
+ Separated from the drawing code because it is the actual judgement the user
+ makes on that screen ("do the angles land on top of each other?"), and a
+ plain array-in/dict-out function can be tested without Qt.
+ """
+ counts = np.asarray(counts)
+ union = int(np.count_nonzero(counts))
+ full = int(np.count_nonzero(counts >= n_angles))
+ return {
+ "union_px": union,
+ "full_px": full,
+ "full_frac": (full / union) if union else 0.0,
+ "mean_count": float(counts[counts > 0].mean()) if union else 0.0,
+ "max_count": int(counts.max()) if counts.size else 0,
+ "empty": union == 0,
+ }
+
+
+def largest_rect_at_least(counts: np.ndarray, min_count: int
+ ) -> tuple[int, int, int, int] | None:
+ """Largest axis-aligned rectangle whose every pixel has counts >= min_count,
+ as (row0, col0, n_rows, n_cols), or None if no pixel qualifies.
+
+ Backs the wizard's "fit to full overlap" button. A *bounding box* of the
+ qualifying pixels would be the obvious thing and is wrong: the full-overlap
+ region of several rotated scans is roughly a disc, whose bounding box has
+ corners no angle covers at all. Offering that as the crop would hand the
+ user padding they explicitly asked to avoid, so this finds a rectangle that
+ is entirely inside the region.
+
+ Standard largest-rectangle-in-a-histogram sweep — O(rows * cols) on a
+ preview-sized array, so it is instant at interactive rates.
+ """
+ good = np.asarray(counts) >= min_count
+ if not good.any():
+ return None
+
+ n_rows, n_cols = good.shape
+ best = (0, 0, 0, 0) # area, row0, col0, ...
+ best_area = 0
+ heights = np.zeros(n_cols, dtype=np.int64)
+
+ for r in range(n_rows):
+ heights = np.where(good[r], heights + 1, 0)
+ # Sentinel column of height 0 flushes the stack at the end of the row.
+ stack: list[tuple[int, int]] = [] # (start col, height)
+ for c in range(n_cols + 1):
+ h = int(heights[c]) if c < n_cols else 0
+ start = c
+ while stack and stack[-1][1] >= h:
+ s, sh = stack.pop()
+ area = sh * (c - s)
+ if area > best_area:
+ best_area = area
+ best = (s, sh, c - s, r)
+ start = s
+ if h:
+ stack.append((start, h))
+
+ col0, height, width, row_end = best
+ return (row_end - height + 1, col0, height, width)
def _parallel_map(fn, items, n_workers: int) -> list:
@@ -739,14 +1790,17 @@ def _parallel_map(fn, items, n_workers: int) -> list:
def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int,
- dc_threshold_mv: float,
- progress_cb=None) -> AlignmentResult:
- """Top-level alignment driver. Runs on a background thread (see
- AngleAlignmentWorker) — deliberately recomputes CH4 DC images from
- scratch rather than reading the GUI-thread _dc_cache dict, since
- background-thread workers must not touch GUI-thread-owned caches."""
+ dc_threshold_mv: float, progress_cb=None,
+ fine_dim: int = _DEFAULT_FINE_DIM) -> AlignmentResult:
+ """Top-level alignment driver: register every angle onto *ref_angle_idx* by
+ content, then lay them all out on that angle's own coordinate grid.
+
+ Meant for a background thread — deliberately
+ recomputes CH4 DC images from scratch rather than reading the GUI-thread
+ _dc_cache dict, since background-thread workers must not touch
+ GUI-thread-owned caches.
+ """
n = sras.n_angles # already the *complete*-angle count for aborted v6 scans
- dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx)
# Parallel over angles, serial within each — see plan_angle_level.
n_workers, angle_budget = plan_angle_level(sras)
@@ -759,7 +1813,7 @@ def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int,
_ticks.append(1)
progress_cb(base + int(min(len(_ticks), n) / n * span))
- # ---- Step 1: CH4 DC image + binarized mask per angle, native grid -----
+ # ---- Step 1: CH4 DC image per angle, native grid ----------------------
def dc4_for(a: int) -> np.ndarray:
dc4 = adc_to_mv(
compute_dc_image(sras, a, CH4_IDX, max_workers=1, budget=angle_budget),
@@ -769,134 +1823,62 @@ def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int,
dc4_mv = dict(enumerate(_parallel_map(dc4_for, range(n), n_workers)))
- # Each angle's own alignment pivot — the CH4-signal-weighted centroid,
- # not the raw scan-window bbox center (see compute_pivot_points_mm).
- # Reuses the dc4_mv images just computed above, so this is free, and is
- # deliberately independent of dc_threshold_mv (see that function's
- # docstring) so a threshold that happens to leave some angle's binary
- # mask empty can't silently degrade the pivot back to the bbox center.
- pivot_mm = compute_pivot_points_mm(sras, dc4_mv=dc4_mv)
+ # ---- Step 2: rigid registration of every angle against the reference --
+ # Its own worker count: registration is bounded by grid-sized FFT buffers,
+ # not by the waveform chunking plan_angle_level budgets for.
+ _ticks.clear()
- # ---- Step 2: coarse translation via FFT phase correlation, on each
- # angle's binarized CH4 mask against the reference's (see
- # correlate_translation_mm — also the engine behind ManualAlignmentDialog's
- # Auto Cross-Correlate button, there defaulting to the raw signal instead
- # of a mask).
- def shift_for(a: int) -> tuple[float, float]:
- shift = correlate_translation_mm(
- sras, a, ref_angle_idx, dc4_mv, pivot_mm,
- use_mask=True, dc_threshold_mv=dc_threshold_mv)
- tick(25, 50)
- return shift
+ def fit_for(a: int) -> RigidFit:
+ fit = register_angle_to_reference(
+ sras, a, ref_angle_idx, dc4_mv, dc_threshold_mv=dc_threshold_mv,
+ fine_dim=fine_dim)
+ tick(25, 65)
+ return fit
- shifts_mm = dict(enumerate(_parallel_map(shift_for, range(n), n_workers)))
+ fits = dict(enumerate(_parallel_map(
+ fit_for, range(n), _registration_workers(sras, fine_dim))))
- # ---- Step 3: union bounding box over all angles (rotation+shift applied)
- corners = np.vstack([
- _corners_in_ref_frame(sras, a, ref_angle_idx, pivot_mm, shifts_mm[a])
- for a in range(n)])
- x_min, y_min = corners.min(axis=0)
- x_max, y_max = corners.max(axis=0)
- n_cols = int(np.ceil((x_max - x_min) / dx_ref)) + 1
- n_rows = int(np.ceil((y_max - y_min) / abs(dy_ref))) + 1
- canvas_origin_mm = (float(x_min), float(y_min if dy_ref > 0 else y_max))
-
- # ---- Step 4: final per-angle full-resolution affine (canvas -> raw idx)
- per_angle: dict[int, AngleTransform] = {}
- for a in range(n):
- matrix, offset = _build_affine_canvas_to_raw(
- sras, a, ref_angle_idx, shifts_mm[a], dx_ref, dy_ref, canvas_origin_mm,
- pivot_mm)
- per_angle[a] = AngleTransform(
- _theta_deg(sras, a, ref_angle_idx), shifts_mm[a], matrix, offset)
- if progress_cb:
- progress_cb(75 + int((a + 1) / n * 25))
-
- return AlignmentResult(ref_angle_idx, dc_threshold_mv, (n_rows, n_cols),
- dx_ref, dy_ref, canvas_origin_mm, per_angle)
-
-
-# Back-compat alias for the pre-split private name (used by tooling).
-_compute_angle_alignment = compute_angle_alignment
+ # ---- Step 3/4: shared canvas on the reference's grid, per-angle affines
+ params = {a: ManualAngleParams(f.rotation_deg, f.shift_mm) for a, f in fits.items()}
+ extra = {a: (f.score, f.source) for a, f in fits.items()}
+ result = _result_from_params(sras, ref_angle_idx, dc_threshold_mv, params, extra)
+ if progress_cb:
+ progress_cb(100)
+ return result
# ---------------------------------------------------------------------------
# Manual alignment (Fusion menu -> Manual Alignment... dialog)
#
-# Skips compute_angle_alignment's mask + phase-correlation search entirely:
-# every angle's rotation_deg/shift_mm is supplied directly by the caller
-# (nudged by eye against a live multi-angle mask overlay, or pre-seeded from
-# a previous compute_angle_alignment run or a saved sidecar). Building the
-# final AlignmentResult from already-known per-angle parameters is pure
-# closed-form matrix math (union_canvas_mm + _build_affine_canvas_to_raw)
-# with no per-pixel image work at all, so build_manual_alignment is cheap
-# enough to call synchronously on the GUI thread on every edit. The only
-# genuinely expensive per-pixel operation anywhere in this flow is
-# reproject_mask, and only ManualAlignmentDialog's own downsampled preview
-# calls that per keystroke — see that class's docstring for how it limits
-# each nudge to reprojecting only the actively-edited angle.
+# Skips register_angle_to_reference's search entirely: every angle's
+# rotation_deg/shift_mm is supplied directly by the caller (nudged by eye
+# against a live multi-angle mask overlay, or pre-seeded from a registration
+# run or a saved sidecar). Building the final AlignmentResult from already-known
+# per-angle parameters is pure closed-form matrix math (_result_from_params)
+# with no per-pixel image work at all, so build_manual_alignment is cheap enough
+# to call synchronously on the GUI thread on every edit. The only genuinely
+# expensive per-pixel operation anywhere in this flow is reproject_mask, and
+# only the wizard's own downsampled mask stack calls that per keystroke
+# — see AlignmentWizard.rebuild_stack for how it limits a nudge to reprojecting
+# only the actively-edited angle.
# ---------------------------------------------------------------------------
-def union_canvas_mm(sras: SrasFile, ref_angle_idx: int, dx: float, dy: float,
- per_angle_params: dict[int, ManualAngleParams],
- pivot_mm: dict[int, tuple[float, float]],
- margin_frac: float = 0.0
- ) -> tuple[tuple[float, float], tuple[int, int]]:
- """Shared-canvas origin (mm) and (n_rows, n_cols) at pitch (dx, dy) that
- contains every angle's footprint after applying its own rotation+shift —
- the generalisation of compute_angle_alignment's Step 3 to arbitrary (not
- just _theta_deg-analytic) per-angle rotation. Angles missing from
- per_angle_params default to identity (e.g. a sidecar saved before a
- rescan added more angles).
-
- margin_frac pads the box on every side: 0 for a final canvas (this then
- reproduces compute_angle_alignment's own Step-3 math exactly, when every
- angle's rotation_deg equals the analytic delta and shift_mm matches);
- nonzero for ManualAlignmentDialog's downsampled preview canvas, which
- needs headroom so an ordinary translation nudge of the active angle never
- has to trigger a full canvas resize (see that class's docstring — an
- extreme nudge can still, in principle, push content past this padding;
- accepted as a known edge case, same as _phase_correlate_shift's
- wraparound risk note).
- """
- n = sras.n_angles
- corners = np.vstack([
- _corners_in_ref_frame(
- sras, a, ref_angle_idx, pivot_mm,
- per_angle_params.get(a, ManualAngleParams()).shift_mm,
- theta_deg=per_angle_params.get(a, ManualAngleParams()).rotation_deg)
- for a in range(n)])
- x_min, y_min = corners.min(axis=0)
- x_max, y_max = corners.max(axis=0)
- if margin_frac:
- pad_x, pad_y = (x_max - x_min) * margin_frac, (y_max - y_min) * margin_frac
- x_min, x_max = x_min - pad_x, x_max + pad_x
- y_min, y_max = y_min - pad_y, y_max + pad_y
- n_cols = int(np.ceil((x_max - x_min) / dx)) + 1
- n_rows = int(np.ceil((y_max - y_min) / abs(dy))) + 1
- origin = (float(x_min), float(y_min if dy > 0 else y_max))
- return origin, (n_rows, n_cols)
-
-
def reproject_mask(sras: SrasFile, angle_idx: int, ref_angle_idx: int,
mask: np.ndarray, rotation_deg: float,
shift_mm: tuple[float, float],
- canvas_dx: float, canvas_dy: float,
+ canvas_pitch_mm: tuple[float, float],
canvas_origin_mm: tuple[float, float],
canvas_shape: tuple[int, int],
- pivot_mm: dict[int, tuple[float, float]]) -> np.ndarray:
+ *, src_downsample: tuple[int, int] = (1, 1)) -> np.ndarray:
"""Resample one angle's binary/float mask onto an arbitrary canvas via an
- explicit rotation+shift — the single building block
- ManualAlignmentDialog's live preview repeatedly calls (once per
- keystroke, for only the actively-nudged angle), since it bypasses
- compute_angle_alignment's phase-correlation search entirely and just
- takes rotation_deg/shift_mm as given. order=0 (nearest) matches
- apply_alignment's own reasoning: a binary mask must never be blended with
- zero-padding.
- """
- matrix, offset = _build_affine_canvas_to_raw(
- sras, angle_idx, ref_angle_idx, shift_mm, canvas_dx, canvas_dy,
- canvas_origin_mm, pivot_mm, theta_deg=rotation_deg)
+ explicit rotation+shift — the single building block the alignment wizard's
+ live mask stack repeatedly calls (once per angle per edit). *src_downsample* must match the (rows, cols)
+ block-mean factor already applied to *mask*, or the reprojection lands at
+ the wrong scale. order=0 (nearest) matches apply_alignment's own reasoning:
+ a binary mask must never be blended with zero-padding."""
+ matrix, offset = build_canvas_affine(
+ sras, angle_idx, ref_angle_idx, rotation_deg, shift_mm,
+ canvas_pitch_mm, canvas_origin_mm, src_downsample=src_downsample)
return scipy_ndimage.affine_transform(
mask.astype(np.float32, copy=False), matrix, offset=offset,
output_shape=canvas_shape, order=0, mode="constant", cval=0.0)
@@ -904,59 +1886,29 @@ def reproject_mask(sras: SrasFile, angle_idx: int, ref_angle_idx: int,
def build_manual_alignment(sras: SrasFile, ref_angle_idx: int,
dc_threshold_mv: float,
- per_angle_params: dict[int, ManualAngleParams],
- pivot_mm: dict[int, tuple[float, float]] | None = None
+ per_angle_params: dict[int, ManualAngleParams]
) -> AlignmentResult:
"""Build a full, full-resolution AlignmentResult from user-supplied
per-angle rotation+shift — the Manual Alignment counterpart to
- compute_angle_alignment, skipping its mask/phase-correlation search
- entirely (every angle's transform here is exactly what the caller
- supplied). Pure matrix/bbox math once pivot_mm is known, so it is cheap
- enough to call synchronously on the GUI thread. The reference angle's
- params are always forced to identity, regardless of what
- per_angle_params holds for it — it defines the shared origin and must
- never be transformed.
+ compute_angle_alignment, skipping its registration search entirely (every
+ angle's transform here is exactly what the caller supplied). The reference
+ angle's params are always forced to identity, regardless of what
+ per_angle_params holds for it — it defines the shared origin and must never
+ be transformed.
- Pass an already-computed *pivot_mm* (e.g. ManualAlignmentDialog's own
- cache, built once from its live CH4 images) to skip recomputing every
- angle's DC image here; otherwise it's computed fresh via
- compute_pivot_points_mm, which is still fine for a one-off Save or
- sidecar restore (just not free on a very large, not-yet-cached scan).
- dc_threshold_mv itself plays no part in the pivot (see that function's
- docstring) — it's stored on the returned AlignmentResult purely as a
- record of the RF-mask threshold in effect at the time.
+ dc_threshold_mv plays no part in the geometry; it's stored on the returned
+ AlignmentResult purely as a record of the RF-mask threshold in effect at
+ the time.
"""
- n = sras.n_angles
- dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx)
- params = {a: per_angle_params.get(a, ManualAngleParams()) for a in range(n)}
+ params = {a: per_angle_params.get(a, ManualAngleParams())
+ for a in range(sras.n_angles)}
params[ref_angle_idx] = ManualAngleParams()
- if pivot_mm is None:
- pivot_mm = compute_pivot_points_mm(sras)
-
- canvas_origin_mm, canvas_shape = union_canvas_mm(
- sras, ref_angle_idx, dx_ref, dy_ref, params, pivot_mm, margin_frac=0.0)
-
- per_angle: dict[int, AngleTransform] = {}
- for a in range(n):
- p = params[a]
- matrix, offset = _build_affine_canvas_to_raw(
- sras, a, ref_angle_idx, p.shift_mm, dx_ref, dy_ref,
- canvas_origin_mm, pivot_mm, theta_deg=p.rotation_deg)
- per_angle[a] = AngleTransform(p.rotation_deg, p.shift_mm, matrix, offset)
-
- return AlignmentResult(ref_angle_idx, dc_threshold_mv, canvas_shape,
- dx_ref, dy_ref, canvas_origin_mm, per_angle)
+ return _result_from_params(sras, ref_angle_idx, dc_threshold_mv, params)
# ---- Sidecar persistence (.sras.align.json) -------------------------
-#
-# Lives here, not sras_format.py: sras_format.py is scoped to the versioned
-# binary .sras spec itself (see scan_format.md); a manual alignment is a
-# viewer-computed *derived* artifact, analogous in kind to AlignmentResult —
-# so it belongs with the alignment math it serialises, which already lives
-# in this module. json + pathlib are both stdlib, so this doesn't add a new
-# dependency to a module whose only load-bearing constraint is staying free
-# of Qt/matplotlib for cheap multiprocessing-child imports.
+# A viewer-computed derived artifact, so it lives with the alignment math
+# rather than in sras_format (see docs/design.md, "Manual-alignment sidecar").
@dataclass
class ManualAlignmentSidecar:
@@ -973,36 +1925,18 @@ def sidecar_path(sras_path) -> Path:
return p.with_name(p.name + ".align.json")
-# Bumped from 1 -> 2 when the rotation pivot changed from the raw scan-
-# window bbox center to a content-derived centroid, *and* _theta_deg's sign
-# convention was corrected — either change alone makes a version-1 file's
-# stored rotation_deg/shift_mm numbers describe a different (and, for the
-# pivot bug, actively wrong) transform than they would today. Loading one
-# unchanged would silently reproduce exactly the "scans show up everywhere"
-# symptom these fixes address, so version-1 sidecars are treated as absent
-# rather than migrated.
-_SIDECAR_SCHEMA_VERSION = 2
+# Bumped whenever the frame the stored numbers are measured in changes; older
+# sidecars are treated as absent, never migrated. Bump history:
+# docs/design.md ("Schema history").
+_SIDECAR_SCHEMA_VERSION = 3
def save_manual_alignment(sras: SrasFile, ref_angle_idx: int,
dc_threshold_mv: float,
per_angle: dict[int, ManualAngleParams]) -> Path:
"""Write the sidecar JSON for sras.path (overwriting any existing one)
- and return the path written.
-
- Schema (schema_version 2):
- {
- "schema_version": 2,
- "ref_angle_idx": ,
- "dc_threshold_mv": ,
- "per_angle": {
- "": {"rotation_deg": , "shift_mm": [, ]},
- ...
- }
- }
- Angle indices are JSON object keys, so they round-trip as strings —
- load_manual_alignment converts them back to int.
- """
+ and return the path written. Angle indices become JSON object keys, so
+ they round-trip as strings — load_manual_alignment converts them back."""
path = sidecar_path(sras.path)
payload = {
"schema_version": _SIDECAR_SCHEMA_VERSION,
@@ -1057,7 +1991,7 @@ def delete_manual_alignment(sras: SrasFile) -> bool:
"""Delete the sidecar if present. Returns whether a file actually existed
to delete, so Clear Alignment's status message can say so. Genuine I/O
errors (permission denied, read-only share) propagate — the caller
- (ManualAlignmentDialog._on_clear) surfaces them rather than silently
+ (the wizard's Clear path) surfaces them rather than silently
pretending the destructive action succeeded."""
path = sidecar_path(sras.path)
try:
diff --git a/sras_edit_scans.py b/sras_edit_scans.py
new file mode 100644
index 0000000..a906480
--- /dev/null
+++ b/sras_edit_scans.py
@@ -0,0 +1,224 @@
+#!/usr/bin/env python3
+"""
+sras_edit_scans.py — Remove one or more angle scans from a .sras file.
+
+A .sras file holds one or more "angles" (rotation positions); the viewer
+cross-correlates each non-reference angle against the reference to align
+them. If one angle's acquisition went wrong (stage glitch, bad trigger,
+laser dropout, ...) it throws off that alignment for the whole file. This
+tool drops the bad angle(s) and renumbers the rest, writing a new .sras file
+with everything else — waveform samples, calibration preambles, background
+waveform, row/geometry tables — carried over byte-for-byte.
+
+Handles v2-v7. Any precomputed FFT/DC cache (v5 PREC tail, v7 CACH tail) is
+dropped on write, since it's indexed by angle and would be stale/misaligned
+after renumbering; the viewer just recomputes it next time the file opens.
+
+This tool only ever *drops* angles — every kept angle's waveform bytes and
+geometry are carried across verbatim. To write a file whose angles have been
+resampled onto one shared aligned grid and cropped, use the viewer's
+Fusion -> Alignment Wizard (sras_align_export.py) instead.
+
+Usage:
+ python sras_edit_scans.py input.sras --list
+ python sras_edit_scans.py input.sras output.sras --drop 2,5
+ python sras_edit_scans.py input.sras output.sras --keep 0,1,3,4,6
+"""
+
+import argparse
+import struct
+import sys
+from pathlib import Path
+
+from sras_format import GEO_FMT_V6, HDR_FMT, HDR_FMT_V6, HDR_SIZE, SrasFile
+
+_LEGACY_VERSIONS = (2, 3, 4, 5)
+_V6_VERSIONS = (6, 7)
+
+
+def _die(msg: str):
+ print(f"Error: {msg}", file=sys.stderr)
+ sys.exit(1)
+
+
+def parse_args():
+ p = argparse.ArgumentParser(
+ description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+ p.add_argument("input", help="Input .sras file")
+ p.add_argument("output", nargs="?", help="Output .sras file (omit with --list)")
+ p.add_argument("--list", action="store_true",
+ help="Print each angle's index/degrees/geometry and exit")
+ g = p.add_mutually_exclusive_group()
+ g.add_argument("--drop", metavar="I,J,...",
+ help="Comma-separated angle indices to remove")
+ g.add_argument("--keep", metavar="I,J,...",
+ help="Comma-separated angle indices to keep (all others dropped)")
+ return p.parse_args()
+
+
+def _parse_index_list(s: str, n_angles: int) -> set[int]:
+ out = set()
+ for piece in s.split(","):
+ piece = piece.strip()
+ if not piece:
+ continue
+ i = int(piece)
+ if not (0 <= i < n_angles):
+ raise ValueError(f"angle index {i} out of range [0, {n_angles - 1}]")
+ out.add(i)
+ return out
+
+
+def print_listing(sras: SrasFile):
+ print(f"\n{'idx':>4} {'angle_deg':>10} {'x_start_mm':>11} {'rows':>6} {'frames':>7}")
+ for a in range(sras.n_angles):
+ print(f"{a:>4} {sras.angles_deg[a]:>10.4f} {sras.x_start_mm[a]:>11.4f} "
+ f"{int(sras.n_rows[a]):>6} {int(sras.n_frames[a]):>7}")
+
+
+def _copy_range(fin, fout, offset: int, nbytes: int, chunk: int = 64 * 1024 * 1024):
+ """Stream *nbytes* raw bytes from *fin* at *offset* into *fout*, without
+ ever holding more than one chunk in memory (waveform blocks can be
+ hundreds of MB to low GB each)."""
+ fin.seek(offset)
+ remaining = nbytes
+ while remaining:
+ buf = fin.read(min(chunk, remaining))
+ if not buf:
+ raise IOError("unexpected EOF while copying waveform data")
+ fout.write(buf)
+ remaining -= len(buf)
+
+
+# ---------------------------------------------------------------------------
+# Legacy (v2-v5): uniform geometry across angles, one flat waveform block
+# ---------------------------------------------------------------------------
+
+def _write_legacy(sras: SrasFile, keep: list[int], out_path: Path):
+ n_rows = int(sras.n_rows[0])
+ n_frames = int(sras.n_frames[0]) # uniform across angles for v2-v5
+ n_ch = sras.n_channels
+ spf = sras.samples_per_frame
+ bps = sras.bytes_per_sample
+
+ header = struct.pack(
+ HDR_FMT, b"SRAS", sras.version, len(keep), n_rows,
+ float(sras.x_start_mm[0]), float(sras.x_delta_mm),
+ sras.velocity_mm_s, sras.laser_freq_hz,
+ n_frames, spf, sras.sample_rate_hz, bps, n_ch,
+ )
+
+ # Row table + preambles + background sit right after the angle table and
+ # don't vary per angle — copy that whole span through unmodified.
+ angle_table_size = sras.n_angles * 4
+ with open(sras.path, "rb") as f:
+ f.seek(HDR_SIZE + angle_table_size)
+ shared_mid = f.read(sras.data_offset - (HDR_SIZE + angle_table_size))
+
+ angle_bytes = n_rows * n_ch * n_frames * spf * bps
+
+ with open(sras.path, "rb") as fin, open(out_path, "wb") as fout:
+ fout.write(header)
+ fout.write(sras.angles_deg[keep].astype(">f4").tobytes())
+ fout.write(shared_mid)
+ for a in keep:
+ _copy_range(fin, fout, sras.data_offset + a * angle_bytes, angle_bytes)
+
+
+# ---------------------------------------------------------------------------
+# v6/v7: per-angle geometry, ragged waveform blocks
+# ---------------------------------------------------------------------------
+
+def _write_v6(sras: SrasFile, keep: list[int], out_path: Path):
+ header = struct.pack(
+ HDR_FMT_V6, b"SRAS", sras.version, len(keep),
+ sras.x_start_nominal_mm, sras.y_start_nominal_mm,
+ sras.x_delta_nominal_mm, sras.y_delta_nominal_mm,
+ sras.row_spacing_mm, sras.velocity_mm_s, sras.laser_freq_hz,
+ sras.samples_per_frame, sras.sample_rate_hz,
+ sras.bytes_per_sample, sras.n_channels,
+ )
+
+ blocks = {a: (offset, nbytes) for a, offset, nbytes in sras.iter_angle_blocks()}
+ with open(sras.path, "rb") as fin, open(out_path, "wb") as fout:
+ fout.write(header)
+ fout.write(sras.angles_deg[keep].astype(">f4").tobytes())
+ for i in keep:
+ fout.write(struct.pack(
+ GEO_FMT_V6, float(sras.x_start_mm[i]),
+ float(sras.x_delta_mm_per_angle[i]),
+ int(sras.n_frames[i]), int(sras.n_rows[i])))
+ for i in keep:
+ fout.write(sras.y_pos_per_angle[i].astype(">f4").tobytes())
+ fout.write(sras.preambles_raw)
+ fout.write(sras.background_raw)
+ for i in keep:
+ offset, nbytes = blocks[i]
+ _copy_range(fin, fout, offset, nbytes)
+
+
+def main():
+ args = parse_args()
+ in_path = Path(args.input)
+ if not in_path.exists():
+ _die(f"input file not found: {in_path}")
+
+ print(f"Reading {in_path} ...", flush=True)
+ try:
+ sras = SrasFile(str(in_path))
+ except ValueError as e:
+ _die(str(e))
+
+ if sras.version not in (*_LEGACY_VERSIONS, *_V6_VERSIONS):
+ _die(f"unsupported .sras version: {sras.version}")
+
+ aborted_note = " (scan aborted; trailing angle(s) already excluded)" if sras.scan_aborted else ""
+ print(f" Version : v{sras.version}", flush=True)
+ print(f" Angles : {sras.n_angles}{aborted_note}", flush=True)
+
+ if args.list:
+ print_listing(sras)
+ return
+
+ if not args.output:
+ _die("output path required unless --list is given.")
+ if not (args.drop or args.keep):
+ _die("specify --drop or --keep (see --list for indices).")
+
+ out_path = Path(args.output)
+ if out_path.resolve() == in_path.resolve():
+ _die("output path must differ from input path.")
+
+ try:
+ if args.drop:
+ drop = _parse_index_list(args.drop, sras.n_angles)
+ keep = [a for a in range(sras.n_angles) if a not in drop]
+ else:
+ keep = sorted(_parse_index_list(args.keep, sras.n_angles))
+ except ValueError as e:
+ _die(str(e))
+
+ if not keep:
+ _die("at least one angle must remain.")
+
+ dropped = [a for a in range(sras.n_angles) if a not in keep]
+ print(f"\nDropping angle(s): {dropped}")
+ print(f"Keeping angle(s) : {keep} ({len(keep)} of {sras.n_angles})")
+ print(f"\nWriting {out_path} ...", flush=True)
+
+ if sras.version in _LEGACY_VERSIONS:
+ _write_legacy(sras, keep, out_path)
+ else:
+ _write_v6(sras, keep, out_path)
+
+ 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")
+ print("Done.")
+ print("Note: any precomputed FFT/DC cache was dropped (it's indexed by "
+ "angle); the viewer will recompute it next time this file opens.")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/sras_format.py b/sras_format.py
index 5c82f6a..a33a1e5 100644
--- a/sras_format.py
+++ b/sras_format.py
@@ -60,7 +60,15 @@ PREC_FLAG_BG_SUB = 0x01
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_VERSION = 3 # written on every fresh write
+CACH_VERSIONS_READABLE = (1, 2, 3) # accepted on read — see
+ # _read_sfft_block. Each bump only
+ # appended a field, and every older
+ # tail has a well-defined reading:
+ # v1 predates row-averaged FFT
+ # caching (row_avg_n=0) and v1/v2
+ # predate padded caching, so both
+ # are natural-resolution (pad 1).
CACH_FLAG_DC = 0x01
CACH_FLAG_FFT = 0x02
@@ -69,9 +77,19 @@ 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_FMT_V1 = ">4sBH" # magic, flags, n_stored (cach_version 1)
+SFFT_HDR_FMT_V2 = ">4sBHB" # + row_avg_n (cach_version 2)
+SFFT_HDR_FMT = ">4sBHBH" # + pad_factor (cach_version 3)
+SFFT_HDR_SIZE_V1 = struct.calcsize(SFFT_HDR_FMT_V1)
+SFFT_HDR_SIZE_V2 = struct.calcsize(SFFT_HDR_FMT_V2)
SFFT_HDR_SIZE = struct.calcsize(SFFT_HDR_FMT)
+MAX_PAD_FACTOR = 0xFFFF # the H field above
SFFT_FLAG_BG_SUB = 0x01
+SFFT_FLAG_ROW_AVG = 0x02 # peak_freq_mhz came from same-row,
+ # distance-weighted averaged CH1
+ # waveforms, not raw per-pixel ones;
+ # row_avg_n is the neighbor half-width
+ # (pixels) used. Bits 2-7 reserved.
# Fixed channel indices into the .sras data array (CH1=RF, CH3/CH4=Bias DC)
CH1_IDX, CH3_IDX, CH4_IDX = 0, 1, 2
@@ -169,7 +187,10 @@ class SrasFile:
``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.
+ source version. The scalars ``precomputed_bg_sub`` /
+ ``precomputed_row_avg_n`` / ``precomputed_pad_factor`` record the
+ settings the stored FFT images were computed under, so a reader can tell
+ whether they answer the question it is actually asking.
"""
def __init__(self, path: str):
@@ -226,6 +247,54 @@ class SrasFile:
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
+ self.precomputed_row_avg_n: int = 0
+ # Zero-padding factor the stored peak_freq_mhz images were resolved
+ # at: 1 = natural resolution (n_fft == samples_per_frame). A padded
+ # FFT resolves peaks a padded view would, and only such a view can
+ # be served from it — see sras_compute.cached_rf_image.
+ self.precomputed_pad_factor: int = 1
+
+ def encoded_preambles(self) -> bytes:
+ """This file's Preamble Blocks section, as bytes a writer can emit.
+
+ v6/v7 files kept the on-disk span verbatim, which is both cheaper and
+ lossless; legacy files did not keep it, and a v2 file has no preambles
+ at all, so those are re-encoded from the parsed strings (empty ones for
+ v2). Empty is not a silent downgrade: _parse_preamble("") returns {} and
+ _set_calibration falls back to the hardcoded scope constants, which is
+ exactly the calibration a v2 file already gets, so mV values round-trip
+ unchanged.
+
+ Lives here rather than at each writer so the version fan-out sits next
+ to the parser that creates it, and no writer has to probe the object
+ to find out which shape it got.
+ """
+ raw = getattr(self, "preambles_raw", None)
+ if raw is not None:
+ return raw
+ out = bytearray()
+ for s in self.preambles or [""] * self.n_channels:
+ encoded = s.encode("utf-8")
+ out += struct.pack(">H", len(encoded)) + encoded
+ return bytes(out)
+
+ def encoded_background(self) -> bytes:
+ """This file's Background Block, as bytes a writer can emit.
+
+ When there is none (v2/v3), this is samples_per_frame zeros rather than
+ a zero-length block. Every consumer guards on `background is not None`
+ and then subtracts it from a (spf,)-shaped row, so a length-0 array
+ would broadcast-fail at the first background-subtracted FFT; zeros make
+ the subtraction a correct no-op instead.
+ """
+ raw = getattr(self, "background_raw", None)
+ if raw is not None:
+ return raw
+ if self.background is None:
+ samples = np.zeros(self.samples_per_frame, dtype=np.int8)
+ else:
+ samples = np.rint(self.background).astype(np.int8)
+ return struct.pack(">I", samples.size) + samples.tobytes()
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."""
@@ -257,6 +326,15 @@ class SrasFile:
self.scan_aborted = False
self.n_angles_declared = n_angles
+ # Pre-v6 files carry no nominal ROI. Defined as None rather than
+ # left absent so the object's shape does not depend on its version
+ # and writers can ask instead of probing with hasattr.
+ self.x_start_nominal_mm = None
+ self.y_start_nominal_mm = None
+ self.x_delta_nominal_mm = None
+ self.y_delta_nominal_mm = None
+ self.row_spacing_mm = None
+
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)
@@ -375,21 +453,34 @@ class SrasFile:
angles = np.frombuffer(f.read(n_angles * 4), dtype=">f4").astype(np.float32)
x_start = np.empty(n_angles, dtype=np.float64)
+ x_delta = 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
+ x_start[a], x_delta[a], n_frames[a], n_rows[a] = xs, xd, 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)
]
+ # Verbatim on-disk spans of the preamble and background sections,
+ # kept so file-rewriting tools (sras_edit_scans) can carry them
+ # over byte-for-byte without re-parsing.
+ span_start = f.tell()
self._set_calibration(_read_preambles(f, n_ch), n_ch)
- self.background = _read_background(f)
+ span_end = f.tell()
+ f.seek(span_start)
+ self.preambles_raw = f.read(span_end - span_start)
- data_offset = f.tell()
+ span_start = span_end
+ self.background = _read_background(f)
+ span_end = f.tell()
+ f.seek(span_start)
+ self.background_raw = f.read(span_end - span_start)
+
+ data_offset = span_end
self._data_offset = data_offset
@@ -428,6 +519,7 @@ class SrasFile:
self.scan_aborted = n_complete < n_angles_declared
self.angles_deg = angles[:n_complete]
self.x_start_mm = x_start[:n_complete]
+ self.x_delta_mm_per_angle = x_delta[: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]
@@ -436,6 +528,37 @@ class SrasFile:
if self.version == 7 and offset < file_size:
self._parse_cach_section(offset)
+ # ------------------------------------------------------------------
+ # Byte-layout accessors (public: used by file-rewriting tools)
+ # ------------------------------------------------------------------
+
+ @property
+ def data_offset(self) -> int:
+ """File offset where the waveform data begins (headers end)."""
+ return self._data_offset
+
+ @property
+ def y_pos_per_angle(self) -> list[np.ndarray]:
+ """Per-angle Y row positions (mm). The list and its arrays are the
+ live parsed state — tools that reproject may replace entries."""
+ return self._y_pos_per_angle
+
+ @y_pos_per_angle.setter
+ def y_pos_per_angle(self, value: list[np.ndarray]):
+ self._y_pos_per_angle = value
+
+ def iter_angle_blocks(self):
+ """Yields (angle_idx, byte_offset, byte_count) for each complete
+ angle's waveform block. Works for every version: legacy files have
+ uniform per-angle geometry, so the same walk applies."""
+ offset = self._data_offset
+ for a in range(self.n_angles):
+ nbytes = (int(self.n_rows[a]) * self.n_channels
+ * int(self.n_frames[a]) * self.samples_per_frame
+ * self.bytes_per_sample)
+ yield a, offset, nbytes
+ offset += nbytes
+
# ------------------------------------------------------------------
# v7 cache tail (CACH section: precomputed DC / FFT images)
# ------------------------------------------------------------------
@@ -445,12 +568,10 @@ class SrasFile:
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)
+ end = self._data_offset
+ for _, offset, nbytes in self.iter_angle_blocks():
+ end = offset + nbytes
+ return end
def _read_cache_block(self, f, hdr_fmt: str, magic: bytes,
stores: list[list]) -> int | None:
@@ -474,6 +595,40 @@ class SrasFile:
store[angle_idx] = _read_f32_image(f, shape)
return flags
+ def _read_sfft_block(self, f, cach_version: int) -> tuple[int, int, int] | None:
+ """Read the SFFT block header — its layout depends on cach_version,
+ since each bump appended a trailing field (v2 row_avg_n, v3
+ pad_factor) — then n_stored per-angle peak_freq_mhz entries
+ (unchanged across versions).
+
+ Returns (flags, row_avg_n, pad_factor), or None if the block is
+ malformed. The absent fields of an older tail take the value that
+ describes what such a tail can only have been: row_avg_n=0 for v1,
+ which predates row-averaged FFT caching, and pad_factor=1 for v1/v2,
+ which predate padded caching and so are natural-resolution.
+ """
+ hdr_fmt = {1: SFFT_HDR_FMT_V1, 2: SFFT_HDR_FMT_V2}.get(
+ cach_version, SFFT_HDR_FMT)
+ raw = f.read(struct.calcsize(hdr_fmt))
+ if len(raw) < struct.calcsize(hdr_fmt):
+ return None
+ row_avg_n, pad_factor = 0, 1
+ if cach_version == 1:
+ magic, flags, n_stored = struct.unpack(hdr_fmt, raw)
+ elif cach_version == 2:
+ magic, flags, n_stored, row_avg_n = struct.unpack(hdr_fmt, raw)
+ else:
+ magic, flags, n_stored, row_avg_n, pad_factor = struct.unpack(hdr_fmt, raw)
+ if magic != SFFT_MAGIC:
+ return None
+ for _ in range(n_stored):
+ (angle_idx,) = _read_struct(f, ">H")
+ if angle_idx >= self.n_angles:
+ break
+ self.precomputed_freq_mhz[angle_idx] = _read_f32_image(
+ f, self.image_shape(angle_idx))
+ return flags, row_avg_n, pad_factor
+
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:
@@ -482,7 +637,7 @@ class SrasFile:
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:
+ if magic != CACH_MAGIC or cach_version not in CACH_VERSIONS_READABLE:
return
if block_flags & CACH_FLAG_DC:
@@ -492,17 +647,21 @@ class SrasFile:
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:
+ result = self._read_sfft_block(f, cach_version)
+ if result is None:
return
+ flags, row_avg_n, pad_factor = result
self.precomputed_bg_sub = bool(flags & SFFT_FLAG_BG_SUB)
+ self.precomputed_row_avg_n = row_avg_n if (flags & SFFT_FLAG_ROW_AVG) else 0
+ self.precomputed_pad_factor = max(1, pad_factor)
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):
+ new_bg_sub: bool | None = None,
+ new_row_avg_n: int | None = None,
+ new_pad_factor: int | 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
@@ -510,6 +669,18 @@ class SrasFile:
``SrasFile`` already has in memory (from parsing, or a prior write
in this same session) — its bytes are never re-read from disk.
+ *new_row_avg_n* is the same-row neighbor half-width (pixels) the
+ passed *new_freq_mhz* was averaged over before its FFT, 0 for a raw
+ (unaveraged) compute — carried forward like *new_bg_sub* when None.
+ It describes the whole stored FFT block, not per-angle, mirroring
+ how bg-sub has never been tracked per-angle either.
+
+ *new_pad_factor* is the zero-padding factor the passed *new_freq_mhz*
+ was resolved at (1 = natural resolution), carried forward the same
+ way. Like row_avg_n it is provenance, not a hint: a view at a
+ different pad resolves different peaks, so recording it is what lets
+ a reader refuse the cache instead of showing the wrong numbers.
+
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.
@@ -522,6 +693,15 @@ class SrasFile:
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
+ final_row_avg_n = (new_row_avg_n if new_row_avg_n is not None
+ else self.precomputed_row_avg_n)
+ final_pad_factor = (new_pad_factor if new_pad_factor is not None
+ else self.precomputed_pad_factor)
+ if not (0 <= final_row_avg_n <= 255):
+ raise ValueError(f"row_avg_n must fit in a byte (0-255), got {final_row_avg_n}")
+ if not (1 <= final_pad_factor <= MAX_PAD_FACTOR):
+ raise ValueError(
+ f"pad_factor must be 1-{MAX_PAD_FACTOR}, got {final_pad_factor}")
# dc3/dc4 are always populated together by every current caller, but
# guard the per-angle pairing explicitly rather than assume it: an
@@ -547,7 +727,10 @@ class SrasFile:
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))
+ fft_flags |= SFFT_FLAG_ROW_AVG if final_row_avg_n else 0
+ payload += struct.pack(SFFT_HDR_FMT, SFFT_MAGIC, fft_flags,
+ len(fft_entries), final_row_avg_n,
+ final_pad_factor)
for a in fft_entries:
payload += struct.pack(">H", a)
payload += final_freq[a].astype(">f4").tobytes()
@@ -585,6 +768,8 @@ class SrasFile:
self.precomputed_dc4_mv = final_dc4
self.precomputed_freq_mhz = final_freq
self.precomputed_bg_sub = final_bg_sub
+ self.precomputed_row_avg_n = final_row_avg_n
+ self.precomputed_pad_factor = final_pad_factor
# ------------------------------------------------------------------
# Axes helpers
@@ -601,6 +786,24 @@ class SrasFile:
def y_positions_mm(self, angle_idx: int) -> np.ndarray:
return self._y_pos_per_angle[angle_idx]
+ def angles_share_raw_grid(self) -> bool:
+ """True iff every angle's raw (x, y) pixel grid is literally the same
+ array as angle 0's -- the case for a .sras file the viewer's own
+ Alignment Wizard exported (see scan_format.md, "Files written by the
+ viewer's Alignment Wizard"): the writer packs one Per-Angle Geometry
+ record and one Row Table span and repeats those same bytes for every
+ angle, so re-parsed arrays are bit-identical copies rather than
+ independently re-derived numbers -- a bare np.array_equal is the
+ correct test here, no tolerance needed.
+ """
+ if self.n_angles <= 1:
+ return True
+ x0 = self.x_axis_mm(0)
+ y0 = self.y_positions_mm(0)
+ return all(np.array_equal(self.x_axis_mm(a), x0)
+ and np.array_equal(self.y_positions_mm(a), y0)
+ for a in range(1, self.n_angles))
+
def time_axis_ns(self) -> np.ndarray:
return np.arange(self.samples_per_frame) / self.sample_rate_hz * 1e9
diff --git a/sras_viewer.py b/sras_viewer.py
deleted file mode 100644
index a6b53c2..0000000
--- a/sras_viewer.py
+++ /dev/null
@@ -1,3015 +0,0 @@
-#!/usr/bin/env python3
-"""
-SRAS Scan File Viewer
-PyQt6 application for visualizing channel data from .sras binary scan files.
-
-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
-
-RF images are masked: pixels where CH4_dc < dc_threshold show 0.
-
-File parsing lives in sras_format, image/alignment math in sras_compute, and
-background workers in sras_workers — the first two import neither Qt nor
-matplotlib so multiprocessing children can load them cheaply.
-"""
-
-import faulthandler
-import sys
-from pathlib import Path
-
-import matplotlib as mpl
-import numpy as np
-from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg, NavigationToolbar2QT
-from matplotlib.figure import Figure
-from matplotlib.patches import Polygon
-from matplotlib.path import Path as MplPath
-from PyQt6.QtCore import QObject, Qt, QThread, pyqtSignal
-from PyQt6.QtGui import QAction, QKeyEvent
-from PyQt6.QtWidgets import (
- QApplication, QButtonGroup, QCheckBox, QComboBox, QDialog, QDialogButtonBox,
- QDoubleSpinBox, QFileDialog, QFormLayout, QFrame, QGridLayout, QGroupBox,
- QHBoxLayout, QLabel, QLineEdit, QMainWindow, QMessageBox, QProgressDialog,
- QPushButton, QRadioButton, QScrollArea, QSizePolicy, QSpinBox, QSplitter,
- QVBoxLayout, QWidget,
-)
-
-import sras_compute as compute
-from sras_compute import (
- PYFFTW_AVAILABLE, ManualAngleParams, apply_alignment, build_manual_alignment,
- delete_manual_alignment, load_manual_alignment, save_manual_alignment,
- sidecar_path,
-)
-from sras_format import (
- CH1_IDX, CH3_IDX, CH4_IDX, CH_NAMES, SrasFile, adc_to_mv, mv_to_adc,
- _FALLBACK_YMULT_MV, _FALLBACK_YOFF_ADC,
-)
-from sras_workers import (
- AngleAlignmentWorker, BatchCacheWorker, BatchExportWorker, Ch4MaskWorker,
- ComputeWorker, CrossCorrelateWorker, DcPrecomputeWorker, ExportChannel,
- LoadWorker,
-)
-
-faulthandler.enable() # print a native stack trace on SIGSEGV/SIGABRT/etc.
-
-# ---------------------------------------------------------------------------
-# Display constants
-# ---------------------------------------------------------------------------
-
-CH_LABELS = [
- "CH1 — RF (FFT peak freq)",
- "CH3 — Bias A (DC mean)",
- "CH4 — Bias B (DC mean)",
- "CH1 — Velocity (SRAS)",
-]
-
-# Combo index for the derived velocity mode (uses CH1_IDX data)
-VELOCITY_MODE_IDX = 3
-# All modes that operate on CH1 waveforms
-CH1_DERIVED_MODES = (CH1_IDX, VELOCITY_MODE_IDX)
-
-CMAPS = ["gray", "viridis", "plasma", "inferno", "hot", "jet", "RdBu_r", "seismic"]
-
-# (mode_str, status-bar unit, colorbar label) per channel index
-_CHANNEL_DISPLAY = {
- CH1_IDX: ("RF", "Peak frequency (MHz)", "MHz"),
- CH3_IDX: ("DC", "DC mean (mV)", "mV"),
- CH4_IDX: ("DC", "DC mean (mV)", "mV"),
- VELOCITY_MODE_IDX: ("Velocity", "Velocity (m/s)", "m/s"),
-}
-
-_CSS_HINT = "font-size: 11px; color: #aaa;"
-_CSS_INFO = "font-size: 11px;"
-_CSS_MUTED = "color: #888; font-size: 11px;"
-_CSS_WARN = "color: #e07000; font-size: 11px;"
-_CSS_BUSY = "color: #4a90d9; font-size: 11px;"
-
-# Side-panel column widths (the scroll areas that hold the controls).
-_LEFT_PANEL_W = 288
-_RIGHT_PANEL_W = 272
-
-# Minimum width for a spin box so its value + suffix are never clipped.
-_SPIN_MIN_W = 96
-
-
-# ---------------------------------------------------------------------------
-# Small layout helpers
-# ---------------------------------------------------------------------------
-
-def _wrap_label(text: str = "", css: str | None = None) -> QLabel:
- """A word-wrapped QLabel that reports its *wrapped* height to the layout.
-
- A plain word-wrapped QLabel advertises a single-line minimum height, so in a
- fixed-width column the layout happily shrinks it and the extra lines get
- clipped. Enabling height-for-width makes the box layout ask for the real
- height at the column's width instead.
- """
- lbl = QLabel(text)
- lbl.setWordWrap(True)
- sp = lbl.sizePolicy()
- sp.setVerticalPolicy(QSizePolicy.Policy.Minimum)
- sp.setHeightForWidth(True)
- lbl.setSizePolicy(sp)
- if css:
- lbl.setStyleSheet(css)
- return lbl
-
-
-def _group(title: str) -> tuple[QGroupBox, QVBoxLayout]:
- """A group box with consistent, non-cramped internal margins."""
- grp = QGroupBox(title)
- lay = QVBoxLayout(grp)
- lay.setContentsMargins(10, 8, 10, 10)
- lay.setSpacing(6)
- return grp, lay
-
-
-def _form() -> QFormLayout:
- """A label/field form layout for a narrow side panel."""
- form = QFormLayout()
- form.setContentsMargins(0, 0, 0, 0)
- form.setHorizontalSpacing(8)
- form.setVerticalSpacing(6)
- form.setLabelAlignment(Qt.AlignmentFlag.AlignRight
- | Qt.AlignmentFlag.AlignVCenter)
- form.setFormAlignment(Qt.AlignmentFlag.AlignLeft | Qt.AlignmentFlag.AlignTop)
- form.setFieldGrowthPolicy(
- QFormLayout.FieldGrowthPolicy.AllNonFixedFieldsGrow)
- form.setRowWrapPolicy(QFormLayout.RowWrapPolicy.DontWrapRows)
- return form
-
-
-def _scroll_panel(inner: QWidget, width: int) -> QScrollArea:
- """Put a side panel in a fixed-width scroll area.
-
- Without this the panels are sized by the window: a short window squeezes the
- controls past their minimum heights, which is what makes text overlap the
- widget below it. Scrolling keeps every control at its natural size.
- """
- area = QScrollArea()
- area.setWidget(inner)
- area.setWidgetResizable(True)
- area.setFrameShape(QFrame.Shape.NoFrame)
- area.setHorizontalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAlwaysOff)
- area.setVerticalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAsNeeded)
- area.setFixedWidth(width)
- area.viewport().setAutoFillBackground(False)
- inner.setAutoFillBackground(False)
- return area
-
-
-# ---------------------------------------------------------------------------
-# ROI (free quadrilateral in data coordinates)
-# ---------------------------------------------------------------------------
-
-class RoiQuad:
- """Free quadrilateral defined in data coordinates (mm).
-
- Stored as 4 corner points (shape (4, 2)) in CCW order: BL, BR, TR, TL.
- Each corner can be positioned independently, allowing skewed /
- non-orthogonal regions of interest. Because it lives in scan/data
- coords it persists unchanged when the displayed channel/mode switches.
- """
-
- def __init__(self, pts: np.ndarray):
- """pts : array-like, shape (4, 2)."""
- self._pts = np.asarray(pts, dtype=np.float64).reshape(4, 2).copy()
-
- @classmethod
- def from_bbox(cls, x0: float, y0: float, x1: float, y1: float) -> "RoiQuad":
- """Create an axis-aligned rectangle from two opposite corners."""
- lx, rx = min(x0, x1), max(x0, x1)
- by, ty = min(y0, y1), max(y0, y1)
- return cls(np.array([[lx, by], [rx, by], [rx, ty], [lx, ty]]))
-
- def copy(self) -> "RoiQuad":
- return RoiQuad(self._pts.copy())
-
- def corners(self) -> np.ndarray:
- """World-coord corners, shape (4, 2), CCW: BL, BR, TR, TL."""
- return self._pts.copy()
-
- def centroid(self) -> np.ndarray:
- return self._pts.mean(axis=0)
-
- def bbox_size(self) -> np.ndarray:
- """Width and height of the axis-aligned bounding box, shape (2,)."""
- return self._pts.max(axis=0) - self._pts.min(axis=0)
-
- def contains(self, x: float, y: float) -> bool:
- return bool(MplPath(self._pts).contains_point((x, y)))
-
- def mask_for_grid(self, x_axis: np.ndarray,
- y_axis: np.ndarray) -> np.ndarray:
- """Boolean mask (n_rows, n_frames) of pixels whose centres lie
- inside the quadrilateral.
-
- Only the quad's axis-aligned bounding box is tested — meshgrid and
- contains_points over the *whole* grid would be tens of millions of
- point-in-polygon tests (and hundreds of MB of float64 temporaries)
- on a large scan, on every ROI edit.
- """
- x = np.asarray(x_axis, dtype=np.float64)
- y = np.asarray(y_axis, dtype=np.float64)
- mask = np.zeros((y.size, x.size), dtype=bool)
-
- (x0, y0), (x1, y1) = self._pts.min(axis=0), self._pts.max(axis=0)
- cols = np.nonzero((x >= x0) & (x <= x1))[0]
- rows = np.nonzero((y >= y0) & (y <= y1))[0]
- if cols.size == 0 or rows.size == 0:
- return mask
-
- c0, c1 = int(cols[0]), int(cols[-1]) + 1
- r0, r1 = int(rows[0]), int(rows[-1]) + 1
- X, Y = np.meshgrid(x[c0:c1], y[r0:r1])
- inside = MplPath(self._pts).contains_points(
- np.column_stack([X.ravel(), Y.ravel()]))
- mask[r0:r1, c0:c1] = inside.reshape(X.shape)
- return mask
-
-
-# ---------------------------------------------------------------------------
-# Matplotlib canvases
-# ---------------------------------------------------------------------------
-
-class ImageCanvas(FigureCanvasQTAgg):
- pixel_clicked = pyqtSignal(int, int) # row_idx, frame_idx
- roi_changed = pyqtSignal() # ROI created / edited / cleared
- draw_mode_changed = pyqtSignal(bool) # "draw new ROI" arm toggled
-
- # Interaction state values
- _IDLE = "idle"
- _DRAW_NEW = "draw_new"
- _MOVE = "move"
- _DRAG_CORNER = "drag_corner"
-
- # Hit tolerance (display pixels) for handles.
- _HANDLE_PX = 12
- _CLICK_THRESH_PX = 4 # releases within this of press count as a click
-
- def __init__(self, parent=None):
- fig = Figure(figsize=(7, 5), tight_layout=True)
- self.ax = fig.add_subplot(111)
- super().__init__(fig)
- self.setParent(parent)
- self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
- self._extent = None
- self._img_shape = None
-
- # ROI state
- self._roi: RoiQuad | None = None
- self._roi_artists: list = []
- self._state = self._IDLE
- self._draw_mode = False
-
- # Per-interaction snapshots / anchors
- self._press_xy: tuple[float, float] | None = None
- self._press_pixel: tuple[float, float] | None = None
- self._press_button = None
- self._snapshot: RoiQuad | None = None
- self._drag_corner_idx: int = -1
- self._move_anchor = None # press-point in world coords
- self._draw_previous: RoiQuad | None = None
-
- self.mpl_connect("button_press_event", self._on_press)
- self.mpl_connect("motion_notify_event", self._on_motion)
- self.mpl_connect("button_release_event", self._on_release)
-
- # ------------------------------------------------------------------
- # Public API
- # ------------------------------------------------------------------
-
- def show_image(self, img: np.ndarray, extent: list[float], cmap: str,
- vmin: float, vmax: float, xlabel: str, ylabel: str, title: str,
- colorbar_label: str = ""):
- self.figure.clf()
- self.ax = self.figure.add_subplot(111)
- # Patches and lines are destroyed by figure.clf(); drop stale refs.
- self._roi_artists = []
-
- self._extent = extent
- self._img_shape = img.shape
-
- im = self.ax.imshow(
- img, aspect="auto", origin="upper",
- extent=extent, cmap=cmap, vmin=vmin, vmax=vmax,
- interpolation="nearest",
- )
- cb = self.figure.colorbar(im, ax=self.ax, fraction=0.046, pad=0.04)
- if colorbar_label:
- cb.set_label(colorbar_label)
-
- self.ax.set_xlabel(xlabel)
- self.ax.set_ylabel(ylabel)
- self.ax.set_title(title)
-
- # Re-draw the ROI (if any) on top of the fresh image so it persists
- # unchanged across mode / angle / channel switches.
- self._draw_roi()
- self.draw()
-
- def get_roi(self) -> RoiQuad | None:
- return self._roi
-
- def set_roi(self, roi: RoiQuad | None):
- self._roi = roi.copy() if roi is not None else None
- self._draw_roi()
- self.draw_idle()
- self.roi_changed.emit()
-
- def clear_roi(self):
- self._roi = None
- self._remove_roi_artists()
- self.draw_idle()
- self.roi_changed.emit()
-
- def start_drawing(self):
- """Arm the next click+drag on the image to create a new ROI,
- replacing any existing one."""
- self._draw_mode = True
- self.setCursor(Qt.CursorShape.CrossCursor)
- self.draw_mode_changed.emit(True)
-
- def cancel_drawing(self):
- if self._draw_mode:
- self._draw_mode = False
- self.setCursor(Qt.CursorShape.ArrowCursor)
- self.draw_mode_changed.emit(False)
-
- # ------------------------------------------------------------------
- # Rendering
- # ------------------------------------------------------------------
-
- def _remove_roi_artists(self):
- for a in self._roi_artists:
- try:
- a.remove()
- except (ValueError, AttributeError, NotImplementedError):
- pass
- self._roi_artists = []
-
- def _draw_roi(self):
- self._remove_roi_artists()
- if self._roi is None or self.ax is None:
- return
- corners = self._roi.corners()
-
- # Filled quad, then a sharp unfilled edge for visibility over bright
- # images, then draggable corner handles.
- for kwargs in (
- dict(fill=True, facecolor="#ffd93a", edgecolor="#e53935",
- alpha=0.22, linewidth=2.0, zorder=10),
- dict(fill=False, edgecolor="#e53935", linewidth=1.8, zorder=11),
- ):
- patch = Polygon(corners, closed=True, **kwargs)
- self.ax.add_patch(patch)
- self._roi_artists.append(patch)
-
- self._roi_artists.append(self.ax.scatter(
- corners[:, 0], corners[:, 1], s=60, c="white",
- edgecolors="#e53935", linewidths=1.6, zorder=13))
-
- # ------------------------------------------------------------------
- # Hit testing (display pixels for handles, data coords for "inside")
- # ------------------------------------------------------------------
-
- def _hit_test(self, event) -> tuple[str, int | None] | None:
- if self._roi is None or self.ax is None:
- return None
- if event.x is None or event.y is None:
- return None
- corners_disp = self.ax.transData.transform(self._roi.corners())
- click = np.array([event.x, event.y])
-
- for i in range(4):
- if np.hypot(*(corners_disp[i] - click)) <= self._HANDLE_PX:
- return ("corner", i)
-
- if event.xdata is not None and event.ydata is not None:
- if self._roi.contains(event.xdata, event.ydata):
- return ("inside", None)
- return None
-
- # ------------------------------------------------------------------
- # Mouse event handlers
- # ------------------------------------------------------------------
-
- def _on_press(self, event):
- if event.inaxes is not self.ax or self._extent is None:
- return
- if event.button != 1: # only left mouse button
- return
- # If the matplotlib toolbar is in pan / zoom mode, let it handle
- # the interaction instead of starting a ROI manipulation.
- tb = getattr(self, "toolbar", None)
- if tb is not None and getattr(tb, "mode", ""):
- return
-
- self._press_xy = (event.xdata, event.ydata)
- self._press_pixel = (event.x, event.y)
- self._press_button = event.button
-
- if self._draw_mode:
- self._draw_previous = self._roi.copy() if self._roi else None
- self._roi = RoiQuad.from_bbox(event.xdata, event.ydata,
- event.xdata, event.ydata)
- self._state = self._DRAW_NEW
- self._draw_roi()
- self.draw_idle()
- return
-
- hit = self._hit_test(event)
- if hit is None:
- self._state = self._IDLE
- return
-
- kind, idx = hit
- self._snapshot = self._roi.copy()
- if kind == "corner":
- self._state = self._DRAG_CORNER
- self._drag_corner_idx = idx
- else:
- self._state = self._MOVE
- self._move_anchor = (event.xdata, event.ydata)
-
- def _on_motion(self, event):
- if self._state == self._IDLE:
- return
- if event.xdata is None or event.ydata is None:
- return
- if event.inaxes is not self.ax:
- return
-
- if self._state == self._DRAW_NEW:
- x0, y0 = self._press_xy
- self._roi = RoiQuad.from_bbox(x0, y0, event.xdata, event.ydata)
- elif self._state == self._MOVE:
- delta = np.array([event.xdata - self._move_anchor[0],
- event.ydata - self._move_anchor[1]])
- self._roi._pts = self._snapshot.corners() + delta
- elif self._state == self._DRAG_CORNER:
- self._roi._pts[self._drag_corner_idx] = [event.xdata, event.ydata]
-
- self._draw_roi()
- self.draw_idle()
-
- def _on_release(self, event):
- if event.button != 1 and self._press_button != 1:
- return
- prev_state = self._state
- self._state = self._IDLE
- try:
- if prev_state == self._DRAW_NEW:
- self._finish_draw()
- elif prev_state in (self._MOVE, self._DRAG_CORNER):
- self._draw_roi()
- self.draw_idle()
- self.roi_changed.emit()
- else:
- self._maybe_emit_pixel_click(event)
- finally:
- self._press_xy = self._press_pixel = None
- self._press_button = None
-
- def _finish_draw(self):
- """Commit (or reject) a freshly-dragged quad."""
- if self._extent is not None:
- x0, x1, y_bot, y_top = self._extent
- min_w = abs(x1 - x0) * 0.01 # minimum: 1% of each axis range
- min_h = abs(y_bot - y_top) * 0.01
- else:
- min_w = min_h = 1e-6
-
- if self._roi is None:
- too_small = True
- else:
- bbox = self._roi.bbox_size()
- too_small = bbox[0] < min_w or bbox[1] < min_h
- if too_small:
- self._roi = self._draw_previous
-
- self._draw_previous = None
- self.cancel_drawing()
- self._draw_roi()
- self.draw_idle()
- self.roi_changed.emit()
-
- def _maybe_emit_pixel_click(self, event):
- """A release close enough to its press counts as a pixel click."""
- if (self._press_pixel is None or event.x is None or event.y is None
- or self._extent is None or event.inaxes is not self.ax
- or event.xdata is None):
- return
- dx_px = event.x - self._press_pixel[0]
- dy_px = event.y - self._press_pixel[1]
- if dx_px * dx_px + dy_px * dy_px > self._CLICK_THRESH_PX ** 2:
- return
-
- x0, x1, y_bot, y_top = self._extent
- n_rows, n_frames = self._img_shape
- col = int((event.xdata - x0) / (x1 - x0) * n_frames)
- row = int((event.ydata - y_top) / (y_bot - y_top) * n_rows)
- self.pixel_clicked.emit(max(0, min(row, n_rows - 1)),
- max(0, min(col, n_frames - 1)))
-
-
-class WaveformCanvas(FigureCanvasQTAgg):
- def __init__(self, parent=None):
- fig = Figure(figsize=(8, 3), tight_layout=True)
- self.ax_wave = fig.add_subplot(121)
- self.ax_right = fig.add_subplot(122)
- super().__init__(fig)
- self.setParent(parent)
- self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
-
- def show_rf_waveform(self, sras: SrasFile, angle_idx: int,
- row_idx: int, frame_idx: int,
- apply_bg_sub: bool = True):
- """CH1 RF: time-domain + FFT spectrum.
-
- If apply_bg_sub is True and sras.background is not None, the background
- waveform is overlaid on the time-domain plot and the FFT is computed
- on the subtracted signal. The unsubtracted FFT is also shown faintly
- for comparison.
- """
- data = sras.data[angle_idx]
- waveform = data[row_idx, CH1_IDX, frame_idx, :].astype(np.float32)
- t_ns = sras.time_axis_ns()
- f_mhz = sras.freq_axis_mhz()
- dc3_val = data[row_idx, CH3_IDX, frame_idx, :].astype(np.float32).mean()
- dc4_val = data[row_idx, CH4_IDX, frame_idx, :].astype(np.float32).mean()
-
- bg = sras.background if (apply_bg_sub and sras.background is not None) else None
- waveform_plot = waveform - bg if bg is not None else waveform
-
- self.ax_wave.cla()
- self.ax_right.cla()
-
- if bg is not None:
- self.ax_wave.plot(t_ns, waveform, linewidth=0.5, color="#aaaaaa",
- label="raw", zorder=1)
- self.ax_wave.plot(t_ns, bg, linewidth=0.5, color="#e07030",
- linestyle="--", label="background", zorder=2)
- self.ax_wave.plot(t_ns, waveform_plot, linewidth=0.7, color="#4488cc",
- label="subtracted", zorder=3)
- self.ax_wave.legend(fontsize=7, loc="upper right")
- else:
- self.ax_wave.plot(t_ns, waveform, linewidth=0.7, color="#4488cc")
-
- self.ax_wave.set_xlabel("Time (ns)")
- self.ax_wave.set_ylabel("ADC counts")
- bg_tag = " [bg sub]" if bg is not None else ""
- dc3_mv = adc_to_mv(dc3_val, *sras.cal(CH3_IDX))
- dc4_mv = adc_to_mv(dc4_val, *sras.cal(CH4_IDX))
- self.ax_wave.set_title(
- f"CH1 RF row={row_idx} frame={frame_idx}{bg_tag}\n"
- f"CH3={dc3_val:.1f} CH4={dc4_val:.1f} "
- f"({dc3_mv:.2f} / {dc4_mv:.2f} mV)",
- fontsize=8,
- )
-
- # FFT of the (possibly subtracted) waveform
- power_sub = np.abs(np.fft.rfft(waveform_plot)) ** 2
- power_sub[0] = 0.0
- peak_mhz = f_mhz[int(np.argmax(power_sub))]
-
- if bg is not None:
- # Also show the unsubtracted FFT for reference
- power_raw = np.abs(np.fft.rfft(waveform)) ** 2
- power_raw[0] = 0.0
- self.ax_right.plot(f_mhz, power_raw, linewidth=0.5, color="#aaaaaa",
- label="raw FFT", zorder=1)
-
- self.ax_right.plot(f_mhz, power_sub, linewidth=0.7, color="#4488cc",
- label="subtracted FFT" if bg is not None else None, zorder=2)
- self.ax_right.axvline(peak_mhz, color="tomato", linestyle="--",
- linewidth=1.2, label=f"peak = {peak_mhz:.1f} MHz")
- self.ax_right.set_xlabel("Frequency (MHz)")
- self.ax_right.set_ylabel("Power (arb.)")
- self.ax_right.set_title("FFT Power Spectrum")
- self.ax_right.set_xlim(0, 500)
- self.ax_right.legend(fontsize=8)
-
- self.draw()
-
- def show_dc_waveform(self, sras: SrasFile, angle_idx: int, ch_idx: int,
- row_idx: int, frame_idx: int):
- """CH3 or CH4 DC: time-domain + mean annotation."""
- waveform = sras.data[angle_idx][row_idx, ch_idx, frame_idx, :].astype(np.float32)
- mean_val = float(waveform.mean())
- mean_mv = adc_to_mv(mean_val, *sras.cal(ch_idx))
-
- self.ax_wave.cla()
- self.ax_right.cla()
-
- self.ax_wave.plot(sras.time_axis_ns(), waveform, linewidth=0.7, color="#4488cc")
- self.ax_wave.axhline(mean_val, color="tomato", linestyle="--",
- linewidth=1.2, label=f"mean = {mean_val:.2f} ADC")
- self.ax_wave.set_xlabel("Time (ns)")
- self.ax_wave.set_ylabel("ADC counts")
- self.ax_wave.set_title(
- f"{CH_NAMES[ch_idx]} DC row={row_idx} frame={frame_idx}")
- self.ax_wave.legend(fontsize=8)
-
- self.ax_right.text(
- 0.5, 0.5,
- f"DC mode\n\nmean = {mean_val:.3f} ADC\n = {mean_mv:.3f} mV",
- ha="center", va="center",
- transform=self.ax_right.transAxes, fontsize=11,
- )
- self.ax_right.set_axis_off()
-
- self.draw()
-
-
-# ---------------------------------------------------------------------------
-# FFT Options dialog
-# ---------------------------------------------------------------------------
-
-class FftOptionsDialog(QDialog):
- """Configure FFT backend and zero-padding.
-
- Changes take effect only when the user clicks Apply. Cancel discards
- all pending edits. The live 'frequency resolution' label updates as
- the user adjusts the pad factor so they can see the trade-off before
- committing.
- """
-
- def __init__(self, parent=None, *,
- current_backend: str,
- current_pad_factor: int,
- samples_per_frame: int | None,
- sample_rate_hz: float | None,
- grating_um: float):
- super().__init__(parent)
- self.setWindowTitle("FFT Options")
- self.setModal(True)
- self.setMinimumWidth(380)
-
- self._samples_per_frame = samples_per_frame
- self._sample_rate_hz = sample_rate_hz
- self._grating_um = grating_um
-
- layout = QVBoxLayout(self)
-
- # ---- Backend ---------------------------------------------------
- grp_backend = QGroupBox("FFT Backend")
- bl = QVBoxLayout(grp_backend)
-
- self._btn_numpy = QRadioButton("NumPy FFT (always available)")
- self._btn_pyfftw = QRadioButton(
- "pyFFTW (faster for large arrays)" if PYFFTW_AVAILABLE
- else "pyFFTW (not installed — run: pip install pyfftw)")
- self._btn_pyfftw.setEnabled(PYFFTW_AVAILABLE)
-
- self._backend_group = QButtonGroup(self)
- self._backend_group.addButton(self._btn_numpy, id=0)
- self._backend_group.addButton(self._btn_pyfftw, id=1)
-
- if current_backend == "pyfftw" and PYFFTW_AVAILABLE:
- self._btn_pyfftw.setChecked(True)
- else:
- self._btn_numpy.setChecked(True)
-
- bl.addWidget(self._btn_numpy)
- bl.addWidget(self._btn_pyfftw)
- layout.addWidget(grp_backend)
-
- # ---- Zero-padding ----------------------------------------------
- grp_zp = QGroupBox("Zero-Padding")
- zl = QVBoxLayout(grp_zp)
-
- pad_row = QHBoxLayout()
- pad_row.addWidget(QLabel("Pad factor:"))
- self._spin_pad = QSpinBox()
- self._spin_pad.setRange(1, 256)
- self._spin_pad.setValue(max(1, current_pad_factor))
- self._spin_pad.setToolTip(
- "Multiply the waveform length by this factor via zero-padding\n"
- "before computing the FFT.\n"
- "1 = no padding (natural length).\n"
- "Powers of 2 (2, 4, 8 …) give the best performance."
- )
- self._spin_pad.valueChanged.connect(self._update_info)
- pad_row.addWidget(self._spin_pad)
- zl.addLayout(pad_row)
-
- self._lbl_nfft = QLabel()
- self._lbl_freq_res = QLabel()
- self._lbl_vel_res = QLabel()
- for lbl in (self._lbl_nfft, self._lbl_freq_res, self._lbl_vel_res):
- lbl.setStyleSheet(_CSS_HINT)
- zl.addWidget(lbl)
-
- layout.addWidget(grp_zp)
-
- # ---- Buttons ---------------------------------------------------
- buttons = QDialogButtonBox()
- buttons.addButton("Apply", QDialogButtonBox.ButtonRole.AcceptRole
- ).clicked.connect(self.accept)
- buttons.addButton("Cancel", QDialogButtonBox.ButtonRole.RejectRole
- ).clicked.connect(self.reject)
- layout.addWidget(buttons)
-
- self._update_info()
-
- def _update_info(self):
- spf = self._samples_per_frame
- sr = self._sample_rate_hz
- pad = self._spin_pad.value()
-
- if spf is None or sr is None:
- self._lbl_nfft.setText("Load a file to preview FFT parameters.")
- self._lbl_freq_res.setText("")
- self._lbl_vel_res.setText("")
- return
-
- n_fft = spf * pad
- freq_res_hz = sr / n_fft
- freq_res_mhz = freq_res_hz / 1e6
- # v (m/s) = freq (MHz) × grating (µm)
- vel_res_ms = freq_res_mhz * self._grating_um
-
- self._lbl_nfft.setText(f"FFT points: {spf} × {pad} = {n_fft:,}")
- self._lbl_freq_res.setText(
- f"Frequency bin: {freq_res_mhz:.4f} MHz ({freq_res_hz / 1e3:.2f} kHz)")
- self._lbl_vel_res.setText(
- f"Velocity bin: {vel_res_ms:.3f} m/s "
- f"(at grating = {self._grating_um:.2f} µm)")
-
- def get_backend(self) -> str:
- return "pyfftw" if self._btn_pyfftw.isChecked() and PYFFTW_AVAILABLE else "numpy"
-
- def get_pad_factor(self) -> int:
- return max(1, self._spin_pad.value())
-
-
-# ---------------------------------------------------------------------------
-# Batch export dialog (Export -> Batch Export Images...)
-# ---------------------------------------------------------------------------
-
-class BatchExportDialog(QDialog):
- """Configure a batch PNG export of DC/RF/Velocity maps across every angle
- of the currently open file.
-
- Each channel's colorbar range is entered here and held fixed for every
- exported angle (rather than auto-scaled per image, today's live-view
- default) so the exported images are directly comparable to each other.
- """
-
- # (ch_idx, is_velocity, label, colorbar unit, filename tag)
- _ROWS = [
- (CH1_IDX, False, CH_LABELS[CH1_IDX], _CHANNEL_DISPLAY[CH1_IDX][2],
- CH_NAMES[CH1_IDX]),
- (CH3_IDX, False, CH_LABELS[CH3_IDX], _CHANNEL_DISPLAY[CH3_IDX][2],
- CH_NAMES[CH3_IDX]),
- (CH4_IDX, False, CH_LABELS[CH4_IDX], _CHANNEL_DISPLAY[CH4_IDX][2],
- CH_NAMES[CH4_IDX]),
- (CH1_IDX, True, CH_LABELS[VELOCITY_MODE_IDX],
- _CHANNEL_DISPLAY[VELOCITY_MODE_IDX][2], CH_NAMES[VELOCITY_MODE_IDX]),
- ]
- # DC is cheap and precomputed on load; CH1/Velocity need a per-pixel FFT
- # that can take minutes, so don't default to exporting them (same
- # rationale as SrasViewerWindow._on_load_done's channel default).
- _DEFAULT_CHECKED = {CH3_IDX, CH4_IDX}
-
- def __init__(self, parent, *, default_dir: str, default_prefix: str,
- default_ranges: dict[tuple[int, bool], tuple[float, float]]):
- super().__init__(parent)
- self.setWindowTitle("Batch Export Images")
- self.setModal(True)
- self.setMinimumWidth(460)
-
- layout = QVBoxLayout(self)
-
- # ---- Output ------------------------------------------------------
- grp_out, ol = _group("Output")
- dir_row = QHBoxLayout()
- self._edit_dir = QLineEdit(default_dir)
- btn_browse = QPushButton("Browse…")
- btn_browse.clicked.connect(self._on_browse)
- dir_row.addWidget(self._edit_dir)
- dir_row.addWidget(btn_browse)
- out_form = _form()
- out_form.addRow("Folder:", dir_row)
- self._edit_prefix = QLineEdit(default_prefix)
- out_form.addRow("File prefix:", self._edit_prefix)
- ol.addLayout(out_form)
- layout.addWidget(grp_out)
-
- # ---- Channels ------------------------------------------------------
- grp_ch, cl = _group("Channels (fixed range, applied to every angle)")
- grid = QGridLayout()
- grid.setHorizontalSpacing(8)
- grid.setVerticalSpacing(6)
- grid.addWidget(_wrap_label("min:", _CSS_HINT), 0, 1)
- grid.addWidget(_wrap_label("max:", _CSS_HINT), 0, 2)
-
- self._rows: list[tuple[QCheckBox, QDoubleSpinBox, QDoubleSpinBox]] = []
- for i, (ch_idx, is_velocity, label, unit, _tag) in enumerate(self._ROWS, 1):
- chk = QCheckBox(label + (f" [{unit}]" if unit else ""))
- chk.setChecked(not is_velocity and ch_idx in self._DEFAULT_CHECKED)
- vmin, vmax = default_ranges.get((ch_idx, is_velocity), (0.0, 1.0))
- spin_min, spin_max = QDoubleSpinBox(), QDoubleSpinBox()
- for spin, val in ((spin_min, vmin), (spin_max, vmax)):
- spin.setRange(-1e9, 1e9)
- spin.setDecimals(4)
- spin.setMinimumWidth(_SPIN_MIN_W)
- spin.setValue(val)
- grid.addWidget(chk, i, 0)
- grid.addWidget(spin_min, i, 1)
- grid.addWidget(spin_max, i, 2)
- self._rows.append((chk, spin_min, spin_max))
- cl.addLayout(grid)
- layout.addWidget(grp_ch)
-
- # ---- Buttons --------------------------------------------------
- buttons = QDialogButtonBox()
- buttons.addButton("Export", QDialogButtonBox.ButtonRole.AcceptRole
- ).clicked.connect(self.accept)
- buttons.addButton("Cancel", QDialogButtonBox.ButtonRole.RejectRole
- ).clicked.connect(self.reject)
- layout.addWidget(buttons)
-
- def _on_browse(self):
- d = QFileDialog.getExistingDirectory(
- self, "Select Output Folder", self._edit_dir.text())
- if d:
- self._edit_dir.setText(d)
-
- def accept(self):
- """Validate before closing — Cancel bypasses this entirely."""
- if not self.get_prefix():
- QMessageBox.warning(self, "Batch Export", "Enter a file prefix.")
- return
- if not self.get_output_dir():
- QMessageBox.warning(self, "Batch Export", "Choose an output folder.")
- return
- selected = self.get_selected_channels()
- if not selected:
- QMessageBox.warning(
- self, "Batch Export", "Select at least one channel to export.")
- return
- for ch in selected:
- if not ch.vmin < ch.vmax:
- QMessageBox.warning(
- self, "Batch Export", f"{ch.label}: min must be less than max.")
- return
- super().accept()
-
- def get_output_dir(self) -> str:
- return self._edit_dir.text().strip()
-
- def get_prefix(self) -> str:
- return self._edit_prefix.text().strip()
-
- def get_selected_channels(self) -> list[ExportChannel]:
- result = []
- for (chk, spin_min, spin_max), (ch_idx, is_velocity, label, unit, tag) in zip(
- self._rows, self._ROWS):
- if chk.isChecked():
- result.append(ExportChannel(
- ch_idx=ch_idx, is_velocity=is_velocity,
- vmin=spin_min.value(), vmax=spin_max.value(),
- label=label, unit=unit, tag=tag))
- return result
-
-
-# ---------------------------------------------------------------------------
-# Manual alignment dialog (Fusion -> Manual Alignment...)
-# ---------------------------------------------------------------------------
-
-class ManualAlignOverlayCanvas(FigureCanvasQTAgg):
- """Renders ManualAlignmentDialog's multi-angle mask overlay and turns
- keyboard input into translate/rotate nudge requests for whichever angle
- the dialog currently has active.
-
- A pure input+render widget — it holds no alignment state and never
- touches SrasFile itself; ManualAlignmentDialog owns all of that and
- decides, from these signals, whether a cheap single-layer refresh or a
- full preview-canvas rebuild is needed.
-
- FigureCanvasQTAgg is a real QWidget, so keyPressEvent works like on any
- other widget, but Qt only ever delivers key events to whichever widget
- currently has focus — StrongFocus, plus grabbing focus on click and once
- right after the dialog is shown, are both required or arrow keys
- silently do nothing.
-
- Rotate keys are letters (Q/E), not punctuation (comma/period or
- brackets): Shift+letter still reports the same Qt.Key on every platform,
- whereas Shift+comma/bracket can report a different virtual key
- (Key_Less / Key_BraceLeft) depending on platform and keyboard layout —
- which would silently break the "Shift = coarse step" modifier for
- rotation specifically. Arrow keys have no such hazard.
- """
- nudge_translate = pyqtSignal(int, int, bool) # dir_x, dir_y in {-1,0,1}; coarse
- nudge_rotate = pyqtSignal(int, bool) # dir in {-1,1} (CCW/CW); coarse
-
- _TRANSLATE_KEYS = {
- Qt.Key.Key_Left: (-1, 0),
- Qt.Key.Key_Right: (1, 0),
- Qt.Key.Key_Up: (0, -1),
- Qt.Key.Key_Down: (0, 1),
- }
- _ROTATE_KEYS = {Qt.Key.Key_Q: 1, Qt.Key.Key_E: -1} # CCW, CW
-
- def __init__(self, parent=None):
- fig = Figure(figsize=(6, 6), tight_layout=True)
- self.ax = fig.add_subplot(111)
- super().__init__(fig)
- self.setParent(parent)
- self.setFocusPolicy(Qt.FocusPolicy.StrongFocus)
- self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
- self.mpl_connect("button_press_event", lambda _e: self.setFocus())
-
- def show_overlay(self, rgba: np.ndarray, extent: list[float], title: str):
- self.figure.clf()
- self.ax = self.figure.add_subplot(111)
- self.ax.imshow(rgba, extent=extent, origin="upper", aspect="auto")
- self.ax.set_xlabel("X (mm)")
- self.ax.set_ylabel("Y (mm)")
- self.ax.set_title(title)
- self.draw_idle() # coalesces rapid redraws — matters for key-repeat.
-
- def keyPressEvent(self, event: QKeyEvent):
- key = event.key()
- coarse = bool(event.modifiers() & Qt.KeyboardModifier.ShiftModifier)
- if key in self._TRANSLATE_KEYS:
- dx, dy = self._TRANSLATE_KEYS[key]
- self.nudge_translate.emit(dx, dy, coarse)
- event.accept()
- elif key in self._ROTATE_KEYS:
- self.nudge_rotate.emit(self._ROTATE_KEYS[key], coarse)
- event.accept()
- else:
- super().keyPressEvent(event)
-
-
-class ManualAlignmentDialog(QDialog):
- """Non-modal manual angle-alignment editor (Fusion -> Manual Alignment...).
-
- Shows every angle's binarized CH4 (Bias B) mask overlaid in a distinct
- color at partial opacity on one shared canvas, so translation/rotation
- misalignment is visible by eye. Reference angle (always index 0) is
- ground truth and never moves; every other angle is aligned to it. The
- user picks an "active" angle and nudges its rotation+translation with
- the keyboard; Auto De-rotate sets every non-reference angle's rotation to
- the known, analytic scan-angle delta without touching any translation;
- Auto Cross-Correlate does the same rotation and additionally sets
- translation to the FFT-phase-correlation best fit against the reference
- (see compute.correlate_translation_mm) — meant to get every angle roughly
- stacked on top of each other so keyboard nudging only has to make small
- corrections, not find a coarse alignment from scratch. Save writes a
- JSON sidecar next to the .sras file and hands a freshly-built, full-
- resolution AlignmentResult back to the main window — the exact same
- object shape compute_angle_alignment produces, so every existing
- Aligned-View code path (apply_alignment, _aligned_canvas_axes, the
- pixel-inspector inverse-transform) works completely unmodified.
-
- Non-modal by design (shown via .show(), never .exec() or setModal(True))
- so the user can still interact with the main window. Talks back to
- SrasViewerWindow two ways: it reuses parent._run_worker/_jobs directly
- for its background mask-fetch and cross-correlate steps, so the main
- window's existing shutdown/lifecycle plumbing covers both for free, and
- it emits alignment_saved / alignment_cleared signals for the two moments
- that should actually mutate the main window's persistent state —
- everything else (nudging, Auto De-rotate, Auto Cross-Correlate, threshold
- edits) stays purely local to this dialog until Save.
- """
-
- alignment_saved = pyqtSignal(object, str) # AlignmentResult, sidecar path (str)
- alignment_cleared = pyqtSignal()
-
- _PREVIEW_MARGIN_FRAC = 0.15
- _BASE_ALPHA = 0.42
- _ACTIVE_ALPHA = 0.75
- _MAX_PREVIEW_DIM = 1024
-
- def __init__(self, parent: "SrasViewerWindow", sras: SrasFile, *,
- ref_angle_idx: int, dc_threshold_mv: float,
- seed_per_angle: dict[int, ManualAngleParams] | None,
- cached_dc4_mv: dict[int, np.ndarray]):
- super().__init__(parent)
- self._parent = parent
- self._sras = sras
- self._ref_angle_idx = ref_angle_idx
- self._downsample_factor = 1
- self._dc4_mv: dict[int, np.ndarray] = {}
- self._masks_small: dict[int, np.ndarray] = {}
- self._pivot_mm: dict[int, tuple[float, float]] = {}
- self._preview_layers: dict[int, np.ndarray] = {}
- self._preview_origin_mm = (0.0, 0.0)
- self._preview_shape = (1, 1)
- self._preview_dx_mm = self._preview_dy_mm = 1.0
- self._masks_ready = False
-
- self.setWindowTitle(f"Manual Alignment — {sras.path.name}")
- self.resize(1150, 760)
-
- self._seed_initial_params(seed_per_angle)
- n = sras.n_angles
- cmap = mpl.colormaps["tab10"] if n <= 10 else mpl.colormaps["tab20"]
- self._angle_colors = {a: cmap(a % cmap.N)[:3] for a in range(n)}
- self._active_angle = 1 if ref_angle_idx == 0 and n > 1 else 0
-
- self._build_ui(dc_threshold_mv)
- self._set_controls_enabled(False) # re-enabled once masks are ready
- self._start_mask_prep(cached_dc4_mv)
-
- def showEvent(self, event):
- super().showEvent(event)
- self.canvas.setFocus()
-
- # ------------------------------------------------------------------
- # Construction
- # ------------------------------------------------------------------
-
- def _seed_initial_params(self, seed_per_angle: dict[int, ManualAngleParams] | None):
- seed = seed_per_angle or {}
- self._angle_params: dict[int, ManualAngleParams] = {
- a: (ManualAngleParams(seed[a].rotation_deg, seed[a].shift_mm)
- if a in seed else ManualAngleParams())
- for a in range(self._sras.n_angles)
- }
- self._angle_params[self._ref_angle_idx] = ManualAngleParams()
-
- def _build_ui(self, dc_threshold_mv: float):
- root = QHBoxLayout(self)
-
- self.canvas = ManualAlignOverlayCanvas()
- left = QWidget()
- left_l = QVBoxLayout(left)
- left_l.setContentsMargins(0, 0, 0, 0)
- left_l.setSpacing(4)
- left_l.addWidget(NavigationToolbar2QT(self.canvas, left))
- left_l.addWidget(self.canvas)
- root.addWidget(left, stretch=1)
-
- panel = QWidget()
- panel_l = QVBoxLayout(panel)
- panel_l.setContentsMargins(0, 0, 0, 0)
- panel_l.setSpacing(8)
-
- # ---- Active Angle -------------------------------------------------
- grp_angle, al = _group("Active Angle")
- self.combo_active_angle = QComboBox()
- for a in range(self._sras.n_angles):
- label = f"Angle {a} ({self._sras.angles_deg[a]:.1f}°)"
- if a == self._ref_angle_idx:
- label += " [reference]"
- self.combo_active_angle.addItem(label)
- al.addWidget(self.combo_active_angle)
- self.lbl_active_note = _wrap_label("", _CSS_WARN)
- al.addWidget(self.lbl_active_note)
- panel_l.addWidget(grp_angle)
-
- # ---- Manual Adjustment ---------------------------------------------
- self.grp_manual_adjust, mform_box = _group("Manual Adjustment")
- mform = _form()
- self.spin_active_rotation_deg = QDoubleSpinBox()
- self.spin_active_rotation_deg.setRange(-3600.0, 3600.0)
- self.spin_active_rotation_deg.setDecimals(3)
- self.spin_active_rotation_deg.setSuffix(" °")
- self.spin_active_rotation_deg.setMinimumWidth(_SPIN_MIN_W)
- mform.addRow("Rotation:", self.spin_active_rotation_deg)
-
- self.spin_active_shift_x_mm = QDoubleSpinBox()
- self.spin_active_shift_x_mm.setRange(-1e5, 1e5)
- self.spin_active_shift_x_mm.setDecimals(4)
- self.spin_active_shift_x_mm.setSuffix(" mm")
- self.spin_active_shift_x_mm.setMinimumWidth(_SPIN_MIN_W)
- mform.addRow("Shift X:", self.spin_active_shift_x_mm)
-
- self.spin_active_shift_y_mm = QDoubleSpinBox()
- self.spin_active_shift_y_mm.setRange(-1e5, 1e5)
- self.spin_active_shift_y_mm.setDecimals(4)
- self.spin_active_shift_y_mm.setSuffix(" mm")
- self.spin_active_shift_y_mm.setMinimumWidth(_SPIN_MIN_W)
- mform.addRow("Shift Y:", self.spin_active_shift_y_mm)
- mform_box.addLayout(mform)
- panel_l.addWidget(self.grp_manual_adjust)
-
- # ---- Nudge Step Sizes ------------------------------------------------
- self.grp_step_sizes, sl = _group("Nudge Step Sizes")
- sform = _form()
- self.spin_step_translate_mm = QDoubleSpinBox()
- self.spin_step_translate_mm.setRange(0.0001, 1000.0)
- self.spin_step_translate_mm.setDecimals(4)
- self.spin_step_translate_mm.setSuffix(" mm")
- self.spin_step_translate_mm.setValue(0.01)
- self.spin_step_translate_mm.setMinimumWidth(_SPIN_MIN_W)
- sform.addRow("Translate step:", self.spin_step_translate_mm)
-
- self.spin_step_rotate_deg = QDoubleSpinBox()
- self.spin_step_rotate_deg.setRange(0.001, 90.0)
- self.spin_step_rotate_deg.setDecimals(3)
- self.spin_step_rotate_deg.setSuffix(" °")
- self.spin_step_rotate_deg.setValue(0.1)
- self.spin_step_rotate_deg.setMinimumWidth(_SPIN_MIN_W)
- sform.addRow("Rotate step:", self.spin_step_rotate_deg)
-
- self.spin_step_multiplier = QDoubleSpinBox()
- self.spin_step_multiplier.setRange(1.0, 1000.0)
- self.spin_step_multiplier.setDecimals(1)
- self.spin_step_multiplier.setValue(10.0)
- self.spin_step_multiplier.setMinimumWidth(_SPIN_MIN_W)
- sform.addRow("Coarse × (Shift):", self.spin_step_multiplier)
- sl.addLayout(sform)
- sl.addWidget(_wrap_label(
- "Arrow keys nudge X/Y translation; Q/E nudge rotation (CCW/CW). "
- "Hold Shift for the coarse step. Click the image once so it has "
- "keyboard focus.", _CSS_HINT))
- panel_l.addWidget(self.grp_step_sizes)
-
- # ---- Mask Threshold ---------------------------------------------------
- self.grp_mask_threshold, tl = _group("Mask Threshold")
- tform = _form()
- self.spin_mask_threshold_mv = QDoubleSpinBox()
- self.spin_mask_threshold_mv.setRange(-500.0, 500.0)
- self.spin_mask_threshold_mv.setDecimals(3)
- self.spin_mask_threshold_mv.setSuffix(" mV")
- self.spin_mask_threshold_mv.setValue(dc_threshold_mv)
- self.spin_mask_threshold_mv.setMinimumWidth(_SPIN_MIN_W)
- tform.addRow("DC threshold:", self.spin_mask_threshold_mv)
- tl.addLayout(tform)
- panel_l.addWidget(self.grp_mask_threshold)
-
- # ---- Cross-Correlate (FFT) -----------------------------------------
- self.grp_correlate, cl = _group("Cross-Correlate (FFT)")
- cform = _form()
- self.combo_correlate_source = QComboBox()
- self.combo_correlate_source.addItems(
- ["Raw signal (recommended)", "Thresholded mask"])
- cform.addRow("Correlate on:", self.combo_correlate_source)
-
- self.spin_correlate_margin = QDoubleSpinBox()
- self.spin_correlate_margin.setRange(0.05, 2.0)
- self.spin_correlate_margin.setSingleStep(0.05)
- self.spin_correlate_margin.setDecimals(2)
- self.spin_correlate_margin.setValue(0.30)
- self.spin_correlate_margin.setMinimumWidth(_SPIN_MIN_W)
- cform.addRow("Search margin (× extent):", self.spin_correlate_margin)
- cl.addLayout(cform)
- self.btn_auto_correlate = QPushButton("Auto Cross-Correlate (vs Reference)")
- cl.addWidget(self.btn_auto_correlate)
- cl.addWidget(_wrap_label(
- "Sets rotation to the known scan angle and translation to the "
- "FFT-correlated best fit for every non-reference angle. Run this "
- "first, then use manual nudging only for small corrections.",
- _CSS_HINT))
- panel_l.addWidget(self.grp_correlate)
-
- # ---- Actions ------------------------------------------------------
- grp_actions, acl = _group("Actions")
- self.btn_auto_derotate = QPushButton("Auto De-rotate (use known angles)")
- self.btn_save = QPushButton("Save Alignment")
- self.btn_clear = QPushButton("Clear Alignment…")
- self.btn_close = QPushButton("Close")
- for btn in (self.btn_auto_derotate, self.btn_save, self.btn_clear, self.btn_close):
- acl.addWidget(btn)
- panel_l.addWidget(grp_actions)
-
- self.lbl_status = _wrap_label("", _CSS_MUTED)
- panel_l.addWidget(self.lbl_status)
- panel_l.addStretch()
-
- root.addWidget(_scroll_panel(panel, 320))
-
- self.combo_active_angle.currentIndexChanged.connect(self._on_active_angle_changed)
- self.spin_active_rotation_deg.editingFinished.connect(self._on_rotation_spin_edited)
- self.spin_active_shift_x_mm.editingFinished.connect(self._on_shift_spin_edited)
- self.spin_active_shift_y_mm.editingFinished.connect(self._on_shift_spin_edited)
- self.spin_mask_threshold_mv.editingFinished.connect(self._on_mask_threshold_edited)
- self.btn_auto_derotate.clicked.connect(self._on_auto_derotate)
- self.btn_auto_correlate.clicked.connect(self._on_auto_correlate)
- self.btn_save.clicked.connect(self._on_save)
- self.btn_clear.clicked.connect(self._on_clear)
- self.btn_close.clicked.connect(self.close)
- self.canvas.nudge_translate.connect(self._on_nudge_translate)
- self.canvas.nudge_rotate.connect(self._on_nudge_rotate)
-
- self.combo_active_angle.blockSignals(True)
- self.combo_active_angle.setCurrentIndex(self._active_angle)
- self.combo_active_angle.blockSignals(False)
- self._on_active_angle_changed(self._active_angle)
-
- # ------------------------------------------------------------------
- # Mask preparation (initial CH4 fetch + threshold + downsample)
- # ------------------------------------------------------------------
-
- def _start_mask_prep(self, cached_dc4_mv: dict[int, np.ndarray]):
- self._dc4_mv = dict(cached_dc4_mv)
- missing = [a for a in range(self._sras.n_angles) if a not in self._dc4_mv]
- if not missing:
- self._finish_mask_prep()
- return
- self.lbl_status.setText(f"Preparing masks: 0/{len(missing)} angle(s) needed…")
- started = self._parent._run_worker(
- "manual_align_masks", Ch4MaskWorker(self._sras, missing),
- connect=(
- ("angle_done", self._on_mask_angle_done),
- ("error", lambda msg: self.lbl_status.setText(f"Mask prep error: {msg}")),
- ),
- on_done=self._finish_mask_prep)
- if not started:
- self.lbl_status.setText(
- "Could not start mask preparation (busy) — close and reopen.")
-
- def _on_mask_angle_done(self, angle_idx: int, dc4_mv: np.ndarray):
- self._dc4_mv[angle_idx] = dc4_mv
- self.lbl_status.setText(
- f"Preparing masks: {len(self._dc4_mv)}/{self._sras.n_angles} ready…")
-
- def _finish_mask_prep(self):
- if len(self._dc4_mv) < self._sras.n_angles:
- return # a mask-worker error left some angles unfetched
- max_dim = max(max(img.shape) for img in self._dc4_mv.values())
- self._downsample_factor = max(1, int(np.ceil(max_dim / self._MAX_PREVIEW_DIM)))
- self._recompute_masks_small()
- # Alignment pivot: the CH4-signal-weighted centroid of each angle's
- # own footprint (see compute.compute_pivot_points_mm) — computed once
- # from the full-res CH4 images and deliberately independent of the
- # mask threshold, so it never needs recomputing when that changes
- # (unlike _masks_small, which is purely for the overlay's visuals).
- self._pivot_mm = compute.compute_pivot_points_mm(self._sras, self._dc4_mv)
- self._rebuild_preview_canvas()
- self._set_controls_enabled(True)
- self.lbl_status.setText("Ready.")
-
- def _recompute_masks_small(self):
- """Threshold + downsample every angle's already-in-memory full-res
- CH4 mV image. Cheap (a compare + block-mean), so this re-runs in
- full whenever the mask-threshold spin box changes — no re-fetch.
- Purely for the overlay's visuals — the alignment pivot does not
- depend on this threshold (see _pivot_mm / compute_pivot_points_mm)."""
- threshold = self.spin_mask_threshold_mv.value()
- factor = self._downsample_factor
- self._masks_small = {
- a: compute._block_mean_downsample(
- (img >= threshold).astype(np.float32), factor)
- for a, img in self._dc4_mv.items()
- }
-
- # ------------------------------------------------------------------
- # Preview canvas: full rebuild vs. incremental single-layer refresh
- # ------------------------------------------------------------------
-
- def _rebuild_preview_canvas(self):
- """Full geometry rebuild: recomputes the shared preview canvas's
- origin/shape (rotation can grow the union bbox — translation alone
- cannot, per the padding baked in via _PREVIEW_MARGIN_FRAC) and every
- angle's reprojected mask layer. Triggered by: dialog open,
- mask-threshold change, Auto De-rotate, a rotation nudge/edit of the
- active angle. NOT triggered by a translation-only nudge — see
- _refresh_active_preview_layer."""
- dx_ref, dy_ref = compute._pixel_pitch_mm(self._sras, self._ref_angle_idx)
- factor = self._downsample_factor
- dx_c, dy_c = dx_ref * factor, dy_ref * factor
- origin, shape = compute.union_canvas_mm(
- self._sras, self._ref_angle_idx, dx_c, dy_c, self._angle_params,
- self._pivot_mm, margin_frac=self._PREVIEW_MARGIN_FRAC)
- self._preview_origin_mm, self._preview_shape = origin, shape
- self._preview_dx_mm, self._preview_dy_mm = dx_c, dy_c
- self._preview_layers = {
- a: compute.reproject_mask(
- self._sras, a, self._ref_angle_idx, self._masks_small[a],
- self._angle_params[a].rotation_deg, self._angle_params[a].shift_mm,
- dx_c, dy_c, origin, shape, self._pivot_mm)
- for a in range(self._sras.n_angles)
- }
- self._redraw_overlay()
-
- def _refresh_active_preview_layer(self):
- """Cheap path for a translation-only nudge/edit of the active angle:
- reproject just that one angle's downsampled mask onto the *existing*
- preview canvas — every other angle's cached layer is untouched."""
- a = self._active_angle
- self._preview_layers[a] = compute.reproject_mask(
- self._sras, a, self._ref_angle_idx, self._masks_small[a],
- self._angle_params[a].rotation_deg, self._angle_params[a].shift_mm,
- self._preview_dx_mm, self._preview_dy_mm,
- self._preview_origin_mm, self._preview_shape, self._pivot_mm)
- self._redraw_overlay()
-
- def _redraw_overlay(self):
- """Alpha-composite every angle's colored mask layer into one RGBA
- image ("all thresholds overlaid with varying opacity"). Each angle
- keeps a fixed, distinct color regardless of which is active; the
- active angle is drawn last (on top) at a visibly higher alpha so
- it's easy to track while nudging."""
- if not self._preview_layers:
- return # mask prep hasn't finished yet — nothing to draw
- n_rows, n_cols = self._preview_shape
- rgba = np.zeros((n_rows, n_cols, 4), dtype=np.float32)
- order = sorted(range(self._sras.n_angles), key=lambda a: a == self._active_angle)
- for a in order:
- layer = self._preview_layers.get(a)
- if layer is None:
- continue
- alpha = self._ACTIVE_ALPHA if a == self._active_angle else self._BASE_ALPHA
- color = self._angle_colors[a]
- fg_a = layer * alpha
- for c in range(3):
- rgba[..., c] = color[c] * fg_a + rgba[..., c] * rgba[..., 3] * (1 - fg_a)
- rgba[..., 3] = fg_a + rgba[..., 3] * (1 - fg_a)
-
- x0, y0 = self._preview_origin_mm
- dx, dy = self._preview_dx_mm, self._preview_dy_mm
- x_axis = x0 + np.arange(n_cols) * dx
- y_axis = y0 + np.arange(n_rows) * dy
- extent = [x_axis[0] - dx / 2, x_axis[-1] + dx / 2,
- y_axis[-1] + dy / 2, y_axis[0] - dy / 2]
- title = (f"Angle {self._active_angle} active "
- f"({self._sras.angles_deg[self._active_angle]:.1f}°)")
- self.canvas.show_overlay(rgba, extent, title)
-
- # ------------------------------------------------------------------
- # Angle selection / nudge / edit handlers
- # ------------------------------------------------------------------
-
- def _on_active_angle_changed(self, angle_idx: int):
- self._active_angle = angle_idx
- is_ref = angle_idx == self._ref_angle_idx
- self.grp_manual_adjust.setEnabled(self._masks_ready and not is_ref)
- self.lbl_active_note.setText(
- "Reference angle — defines the shared origin, not adjustable." if is_ref else "")
- self._sync_active_spinboxes()
- self._redraw_overlay()
-
- def _sync_active_spinboxes(self):
- p = self._angle_params[self._active_angle]
- for spin, val in ((self.spin_active_rotation_deg, p.rotation_deg),
- (self.spin_active_shift_x_mm, p.shift_mm[0]),
- (self.spin_active_shift_y_mm, p.shift_mm[1])):
- spin.blockSignals(True)
- spin.setValue(val)
- spin.blockSignals(False)
-
- def _on_nudge_translate(self, dir_x: int, dir_y: int, coarse: bool):
- if not self._masks_ready or self._active_angle == self._ref_angle_idx:
- return
- step = self.spin_step_translate_mm.value()
- if coarse:
- step *= self.spin_step_multiplier.value()
- p = self._angle_params[self._active_angle]
- p.shift_mm = (p.shift_mm[0] + dir_x * step, p.shift_mm[1] + dir_y * step)
- self._sync_active_spinboxes()
- self._refresh_active_preview_layer()
-
- def _on_nudge_rotate(self, direction: int, coarse: bool):
- if not self._masks_ready or self._active_angle == self._ref_angle_idx:
- return
- step = self.spin_step_rotate_deg.value()
- if coarse:
- step *= self.spin_step_multiplier.value()
- self._angle_params[self._active_angle].rotation_deg += direction * step
- self._sync_active_spinboxes()
- self._rebuild_preview_canvas()
-
- def _on_rotation_spin_edited(self):
- if self._active_angle == self._ref_angle_idx:
- return
- self._angle_params[self._active_angle].rotation_deg = self.spin_active_rotation_deg.value()
- self._rebuild_preview_canvas()
-
- def _on_shift_spin_edited(self):
- if self._active_angle == self._ref_angle_idx:
- return
- p = self._angle_params[self._active_angle]
- p.shift_mm = (self.spin_active_shift_x_mm.value(), self.spin_active_shift_y_mm.value())
- self._refresh_active_preview_layer()
-
- def _on_mask_threshold_edited(self):
- if not self._masks_ready:
- return
- self._recompute_masks_small()
- self._rebuild_preview_canvas()
-
- # ------------------------------------------------------------------
- # Actions
- # ------------------------------------------------------------------
-
- def _on_auto_derotate(self):
- n_changed = 0
- for a in range(self._sras.n_angles):
- if a == self._ref_angle_idx:
- continue
- self._angle_params[a].rotation_deg = compute._theta_deg(
- self._sras, a, self._ref_angle_idx)
- n_changed += 1
- self._sync_active_spinboxes()
- self._rebuild_preview_canvas()
- self.lbl_status.setText(
- f"Rotation set to the known scan angle for {n_changed} angle(s) "
- "(translation left untouched).")
-
- def _on_auto_correlate(self):
- if not self._masks_ready:
- return
- angles = [a for a in range(self._sras.n_angles) if a != self._ref_angle_idx]
- if not angles:
- return
- use_mask = self.combo_correlate_source.currentIndex() == 1
- worker = CrossCorrelateWorker(
- self._sras, self._ref_angle_idx, angles, self._dc4_mv, self._pivot_mm,
- use_mask=use_mask, dc_threshold_mv=self.spin_mask_threshold_mv.value(),
- margin_frac=self.spin_correlate_margin.value())
- self._correlate_done_count = 0
- self._correlate_total = len(angles)
- self._set_controls_enabled(False)
- self.lbl_status.setText(f"Cross-correlating: 0/{self._correlate_total} angle(s)…")
- started = self._parent._run_worker(
- "manual_align_correlate", worker,
- connect=(
- ("angle_done", self._on_correlate_angle_done),
- ("error", self._on_correlate_error),
- ),
- on_done=self._finish_auto_correlate)
- if not started:
- self._set_controls_enabled(True)
- self.lbl_status.setText("Could not start cross-correlation (busy) — try again.")
-
- def _on_correlate_angle_done(self, angle_idx: int, rotation_deg: float,
- shift_x_mm: float, shift_y_mm: float):
- self._angle_params[angle_idx] = ManualAngleParams(rotation_deg, (shift_x_mm, shift_y_mm))
- self._correlate_done_count += 1
- self.lbl_status.setText(
- f"Cross-correlating: {self._correlate_done_count}/{self._correlate_total} angle(s)…")
-
- def _on_correlate_error(self, msg: str):
- self.lbl_status.setText(f"Cross-correlation error: {msg}")
-
- def _finish_auto_correlate(self):
- self._sync_active_spinboxes()
- self._rebuild_preview_canvas()
- self._set_controls_enabled(True)
- source = "thresholded mask" if self.combo_correlate_source.currentIndex() == 1 \
- else "raw signal"
- self.lbl_status.setText(
- f"Cross-correlated {self._correlate_done_count} angle(s) against "
- f"Angle {self._ref_angle_idx} using the {source}. Nudge from here "
- "for any remaining fine correction.")
-
- def _on_save(self):
- threshold = self.spin_mask_threshold_mv.value()
- resolved = dict(self._angle_params) # already concrete floats
- try:
- path = save_manual_alignment(self._sras, self._ref_angle_idx, threshold, resolved)
- result = build_manual_alignment(self._sras, self._ref_angle_idx, threshold,
- resolved, self._pivot_mm)
- except OSError as exc:
- QMessageBox.warning(self, "Save Alignment Failed", str(exc))
- return
- self.lbl_status.setText(f"Saved to {path.name}.")
- self.alignment_saved.emit(result, str(path))
-
- def _on_clear(self):
- reply = QMessageBox.question(
- self, "Clear Alignment",
- "This resets every angle back to raw/unaligned (0° rotation, no "
- "shift) and deletes the saved alignment file for this scan, if "
- "any. This cannot be undone. Continue?",
- QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No,
- QMessageBox.StandardButton.No)
- if reply != QMessageBox.StandardButton.Yes:
- return
- try:
- existed = delete_manual_alignment(self._sras)
- except OSError as exc:
- QMessageBox.warning(self, "Clear Alignment Failed",
- f"Could not delete the saved alignment file: {exc}")
- return
- self._angle_params = {a: ManualAngleParams() for a in range(self._sras.n_angles)}
- self._sync_active_spinboxes()
- self._rebuild_preview_canvas()
- self.lbl_status.setText(
- "Alignment cleared; saved file removed." if existed
- else "Alignment cleared (there was no saved file).")
- self.alignment_cleared.emit()
-
- def _set_controls_enabled(self, enabled: bool):
- self._masks_ready = enabled
- self.combo_active_angle.setEnabled(enabled)
- self.grp_manual_adjust.setEnabled(enabled and self._active_angle != self._ref_angle_idx)
- self.grp_step_sizes.setEnabled(enabled)
- self.grp_mask_threshold.setEnabled(enabled)
- self.grp_correlate.setEnabled(enabled)
- self.btn_auto_derotate.setEnabled(enabled)
- self.btn_save.setEnabled(enabled)
- self.btn_clear.setEnabled(enabled)
-
-
-# ---------------------------------------------------------------------------
-# Main window
-# ---------------------------------------------------------------------------
-
-class SrasViewerWindow(QMainWindow):
- def __init__(self, initial_path: str | None = None):
- super().__init__()
- self.setWindowTitle("SRAS Scan Viewer")
- self.resize(1560, 840)
- self.setMinimumSize(960, 560)
- self.setAcceptDrops(True)
-
- self._sras: SrasFile | None = None
- self._current_image: np.ndarray | None = None
- self._current_angle: int = 0
- self._current_ch: int = 0
- self._pending_angle: int = 0
- self._pending_ch: int = 0
- self._pending_bg_sub: bool = True
- self._pending_threshold: float = 50.0 # mV
- self._pending_fft_pad_factor: int = 1
-
- # Live background jobs, keyed by role — see _run_worker.
- self._jobs: dict[str, tuple] = {}
- self._progress_dlgs: dict[str, QProgressDialog] = {}
-
- # FFT settings (configured via FFT Options dialog)
- self._fft_pad_factor: int = 1 # 1 = no padding
-
- # Convert menu: batch DC/FFT compute-and-store (v6 -> v7)
- self._batch_errors: list[str] = []
-
- # Export menu: batch image export
- self._export_errors: list[str] = []
- self._export_ok_count: int = 0
-
- # Display-only settings (colormap, grating) never trigger a
- # recompute — they're applied to cached data on redraw. DC images
- # (CH3/CH4) are cheap and precomputed for every angle in the
- # background right after load. CH1/Velocity FFT images are
- # computed lazily (with a progress popup) the first time an
- # angle/threshold combination is viewed — using the cached DC4
- # image to skip the FFT entirely for masked-out pixels — and
- # cached per (angle, bg_sub, n_fft, threshold) so revisiting the
- # same combination is free.
- self._dc_cache: dict[tuple[int, int], np.ndarray] = {}
- self._fft_cache: dict[tuple[int, bool, int | None, float], np.ndarray] = {}
- self._dc_precompute_worker: DcPrecomputeWorker | None = None
- self._dc_generation: int = 0
- self._compute_generation: int = 0
- self._pending_compute_generation: int | None = None
-
- # Angle alignment ("Fusion" menu)
- self._alignment_result = None
- self._alignment_generation: int = 0
- self._aligned_cache: dict[tuple, np.ndarray] = {}
- self._manual_align_dialog: ManualAlignmentDialog | None = None
-
- self._build_ui()
-
- if initial_path:
- self._load_file(initial_path)
-
- # ------------------------------------------------------------------
- # Background job plumbing
- # ------------------------------------------------------------------
-
- def _run_worker(self, key: str, worker: QObject, *,
- connect: tuple = (), quit_on: tuple = ("finished",),
- on_done=None) -> bool:
- """Move *worker* onto its own QThread and start it. Returns False if
- a job under *key* is already running.
-
- Centralises two lifetime hazards that each cost a process abort:
-
- 1. The job is claimed in self._jobs *before* start() and before
- anything below that can pump the Qt event loop (a
- QProgressDialog.show() does on first display). If it weren't, a
- re-entrant editingFinished could slip past the busy check, start a
- second thread, and then have the first call's own assignment
- clobber — and destroy while still running — that second QThread.
-
- 2. thread.finished fires as the thread winds down but does not
- guarantee the OS thread has joined. Dropping the last reference to
- a QThread whose thread is still running logs "QThread: Destroyed
- while thread is still running" and aborts, so wait() first.
-
- 3. The thread must quit once worker.run() returns, regardless of
- which signal (if any) run() emitted on its way out. quit_on exists
- so a worker whose real completion is signalled asynchronously
- from *within* run() (e.g. angle_done fired from pool threads
- before run() itself returns) can still opt into quitting earlier,
- but relying on quit_on alone is fragile: a worker whose run()
- emits `error` instead of `finished` on an exception path, at a
- call site that forgot to add "error" to quit_on, would otherwise
- never quit — silently wedging that job in self._jobs forever.
- Quitting unconditionally after run() returns closes that gap for
- every worker, present and future, without requiring every call
- site to enumerate every signal run() might emit.
- """
- if key in self._jobs:
- return False
-
- thread = QThread()
- self._jobs[key] = (thread, worker, on_done) # claim before anything pumps
- worker.moveToThread(thread)
-
- def _run_then_quit():
- worker.run()
- thread.quit()
- thread.started.connect(_run_then_quit)
- for signal_name, slot in connect:
- getattr(worker, signal_name).connect(slot)
- for signal_name in quit_on:
- getattr(worker, signal_name).connect(thread.quit)
- thread.finished.connect(lambda k=key: self._on_job_finished(k))
- thread.start()
- return True
-
- def _on_job_finished(self, key: str):
- job = self._jobs.pop(key, None)
- if job is None:
- return
- thread, _worker, on_done = job
- thread.wait() # join before releasing our last reference
- if on_done is not None:
- on_done()
-
- def _job_running(self, key: str) -> bool:
- return key in self._jobs
-
- # ------------------------------------------------------------------
- # UI construction
- # ------------------------------------------------------------------
-
- def _build_ui(self):
- central = QWidget()
- self.setCentralWidget(central)
- root = QHBoxLayout(central)
- root.setContentsMargins(8, 8, 8, 8)
- root.setSpacing(8)
-
- root.addWidget(self._build_left_panel())
- root.addWidget(self._build_canvases(), stretch=1)
- root.addWidget(self._build_right_panel())
-
- self.statusBar().showMessage("Open an .sras file to begin.")
- self._build_menus()
-
- def _build_left_panel(self) -> QWidget:
- panel = QWidget()
- panel_layout = QVBoxLayout(panel)
- panel_layout.setContentsMargins(0, 0, 0, 0)
- panel_layout.setSpacing(8)
-
- # ---- File -------------------------------------------------------
- grp_file, fl = _group("File")
- self.btn_open = QPushButton("Open .sras…")
- self.btn_open.clicked.connect(self._on_open)
- self.lbl_filename = _wrap_label("No file loaded", _CSS_MUTED)
- fl.addWidget(self.btn_open)
- fl.addWidget(self.lbl_filename)
- panel_layout.addWidget(grp_file)
-
- # ---- Scan info --------------------------------------------------
- grp_info, il = _group("Scan Info")
- il.setSpacing(3)
- self._info = {}
- for key in ("Angles", "Rows", "Frames / row", "Samples / frame",
- "Sample rate", "X start", "Pixel Δx", "Laser freq"):
- lbl = _wrap_label(f"{key}: —", _CSS_INFO)
- il.addWidget(lbl)
- self._info[key] = lbl
-
- # frame-count / format notes
- self.lbl_frame_warn = _wrap_label("", _CSS_WARN)
- il.addWidget(self.lbl_frame_warn)
-
- # background DC-precompute progress
- self.lbl_dc_precompute = _wrap_label("", _CSS_BUSY)
- il.addWidget(self.lbl_dc_precompute)
- panel_layout.addWidget(grp_info)
-
- # ---- View settings ----------------------------------------------
- grp_view, vl = _group("View Settings")
-
- view_form = _form()
-
- self.spin_angle = QSpinBox()
- self.spin_angle.setRange(0, 0)
- self.spin_angle.setEnabled(False)
- self.spin_angle.setMinimumWidth(64)
- self.spin_angle.editingFinished.connect(self._on_view_changed)
- self.lbl_angle_deg = QLabel("—")
- angle_field = QWidget()
- ar = QHBoxLayout(angle_field)
- ar.setContentsMargins(0, 0, 0, 0)
- ar.setSpacing(6)
- ar.addWidget(self.spin_angle)
- ar.addWidget(self.lbl_angle_deg)
- ar.addStretch()
- view_form.addRow("Angle:", angle_field)
-
- self.combo_channel = QComboBox()
- self.combo_channel.addItems(CH_LABELS)
- self.combo_channel.setEnabled(False)
- self.combo_channel.setSizePolicy(QSizePolicy.Policy.Expanding,
- QSizePolicy.Policy.Fixed)
- self.combo_channel.setSizeAdjustPolicy(
- QComboBox.SizeAdjustPolicy.AdjustToMinimumContentsLengthWithIcon)
- self.combo_channel.setMinimumContentsLength(12)
- self.combo_channel.currentIndexChanged.connect(self._on_channel_changed)
- view_form.addRow("Channel:", self.combo_channel)
- vl.addLayout(view_form)
-
- sep = QFrame()
- sep.setFrameShape(QFrame.Shape.HLine)
- sep.setStyleSheet("color: #555;")
- vl.addWidget(sep)
-
- # DC threshold (for RF / CH1 masking)
- self.grp_threshold, tl = _group("RF Mask Threshold (CH1 only)")
- thr_form = _form()
- self.spin_threshold_mv = QDoubleSpinBox()
- self.spin_threshold_mv.setRange(-500.0, 500.0)
- self.spin_threshold_mv.setDecimals(3)
- self.spin_threshold_mv.setSingleStep(0.025)
- self.spin_threshold_mv.setSuffix(" mV")
- self.spin_threshold_mv.setValue(50.0)
- self.spin_threshold_mv.setEnabled(False)
- self.spin_threshold_mv.setMinimumWidth(_SPIN_MIN_W)
- self.spin_threshold_mv.editingFinished.connect(self._on_threshold_changed)
- thr_form.addRow("DC threshold:", self.spin_threshold_mv)
- tl.addLayout(thr_form)
- self.lbl_threshold_adc = _wrap_label(
- f"≈ {mv_to_adc(50.0):.1f} ADC counts", _CSS_MUTED)
- tl.addWidget(self.lbl_threshold_adc)
- vl.addWidget(self.grp_threshold)
-
- # Background subtraction (v4+ files only)
- self.chk_bg_sub = QCheckBox("Background subtraction (CH1 only)")
- self.chk_bg_sub.setChecked(True)
- self.chk_bg_sub.setEnabled(False)
- self.chk_bg_sub.setToolTip(
- "Subtract the stored background waveform from each CH1 frame\n"
- "before computing the FFT (v4+ files only)."
- )
- self.chk_bg_sub.toggled.connect(self._on_bg_sub_toggled)
- vl.addWidget(self.chk_bg_sub)
-
- # Aligned View (Fusion → Angle Alignment result)
- self.chk_aligned_view = QCheckBox("Aligned View (Fusion)")
- self.chk_aligned_view.setChecked(False)
- self.chk_aligned_view.setEnabled(False)
- self.chk_aligned_view.setToolTip(
- "Show the current angle/channel resampled onto the shared,\n"
- "rotation+translation-aligned canvas from Fusion → Angle\n"
- "Alignment. Uncheck to see the raw per-angle scan grid."
- )
- self.chk_aligned_view.toggled.connect(self._on_aligned_view_toggled)
- vl.addWidget(self.chk_aligned_view)
-
- self.btn_export_csv = QPushButton("Export Image as CSV…")
- self.btn_export_csv.setEnabled(False)
- self.btn_export_csv.setToolTip(
- "Save the current CH1 image (one scan row per CSV line).")
- self.btn_export_csv.clicked.connect(self._on_export_csv)
- vl.addWidget(self.btn_export_csv)
-
- panel_layout.addWidget(grp_view)
-
- # ---- ROI ---------------------------------------------------------
- grp_roi, rl = _group("ROI (Region of Interest)")
-
- self.btn_draw_roi = QPushButton("Draw ROI")
- self.btn_draw_roi.setCheckable(True)
- self.btn_draw_roi.setEnabled(False)
- self.btn_draw_roi.setToolTip(
- "Arm next click+drag on the image to draw a new ROI\n"
- "(replaces any existing one). Click again to cancel.\n"
- "After drawing, drag inside to move, or grab corners to reshape.\n"
- "The ROI is persistent across channels / modes / angles."
- )
- self.btn_draw_roi.toggled.connect(self._on_draw_roi_toggled)
- rl.addWidget(self.btn_draw_roi)
-
- self.btn_clear_roi = QPushButton("Clear ROI")
- self.btn_clear_roi.setEnabled(False)
- self.btn_clear_roi.clicked.connect(self._on_clear_roi)
- rl.addWidget(self.btn_clear_roi)
-
- self.btn_export_roi = QPushButton("Export ROI as CSV…")
- self.btn_export_roi.setEnabled(False)
- self.btn_export_roi.setToolTip(
- "Save every pixel whose centre lies inside the ROI as CSV.\n"
- "Columns: row, frame, x_mm, y_mm, value.\n"
- "Corner coordinates of the quad are written in the file header."
- )
- self.btn_export_roi.clicked.connect(self._on_export_roi_csv)
- rl.addWidget(self.btn_export_roi)
-
- self.lbl_roi_center = _wrap_label("centroid: —", _CSS_HINT)
- self.lbl_roi_size = _wrap_label("bbox: —", _CSS_HINT)
- self.lbl_roi_npix = _wrap_label("pixels inside: —", _CSS_HINT)
- for lbl in (self.lbl_roi_center, self.lbl_roi_size, self.lbl_roi_npix):
- rl.addWidget(lbl)
-
- panel_layout.addWidget(grp_roi)
- panel_layout.addStretch()
- return _scroll_panel(panel, _LEFT_PANEL_W)
-
- def _build_canvases(self) -> QWidget:
- splitter = QSplitter(Qt.Orientation.Vertical)
- splitter.setChildrenCollapsible(False)
-
- img_widget = QWidget()
- img_vl = QVBoxLayout(img_widget)
- img_vl.setContentsMargins(0, 0, 0, 0)
- img_vl.setSpacing(4)
- self.image_canvas = ImageCanvas()
- self.image_canvas.setMinimumHeight(220)
- self.image_canvas.pixel_clicked.connect(self._on_pixel_clicked)
- self.image_canvas.roi_changed.connect(self._update_roi_ui)
- self.image_canvas.draw_mode_changed.connect(self._on_draw_mode_changed)
- img_vl.addWidget(NavigationToolbar2QT(self.image_canvas, img_widget))
- img_vl.addWidget(self.image_canvas)
- splitter.addWidget(img_widget)
-
- wave_widget = QWidget()
- wave_vl = QVBoxLayout(wave_widget)
- wave_vl.setContentsMargins(0, 0, 0, 0)
- wave_vl.setSpacing(4)
- self.lbl_wave_hint = QLabel(
- "Click a pixel in the image above to inspect its waveform.")
- self.lbl_wave_hint.setAlignment(Qt.AlignmentFlag.AlignCenter)
- self.lbl_wave_hint.setStyleSheet(_CSS_MUTED)
- self.wave_canvas = WaveformCanvas()
- self.wave_canvas.setMinimumHeight(150)
- wave_vl.addWidget(self.lbl_wave_hint)
- wave_vl.addWidget(self.wave_canvas)
- splitter.addWidget(wave_widget)
-
- splitter.setStretchFactor(0, 3)
- splitter.setStretchFactor(1, 1)
- splitter.setSizes([580, 250])
- return splitter
-
- def _build_right_panel(self) -> QWidget:
- # Velocity settings (visible only in velocity mode)
- self.grp_velocity, vel_l = _group("Velocity Settings (CH1 only)")
- vel_form = _form()
- self.spin_grating_um = QDoubleSpinBox()
- self.spin_grating_um.setRange(0.1, 1000.0)
- self.spin_grating_um.setDecimals(2)
- self.spin_grating_um.setSingleStep(0.5)
- self.spin_grating_um.setSuffix(" µm")
- self.spin_grating_um.setValue(25)
- self.spin_grating_um.setEnabled(False)
- self.spin_grating_um.setMinimumWidth(_SPIN_MIN_W)
- self.spin_grating_um.editingFinished.connect(self._on_grating_changed)
- vel_form.addRow("Grating size:", self.spin_grating_um)
- vel_l.addLayout(vel_form)
- vel_l.addWidget(_wrap_label("v (m/s) = freq (MHz) × grating (µm)",
- "font-size: 10px; color: #888;"))
- self.grp_velocity.setVisible(False)
-
- grp_display, dl = _group("Display Options")
-
- cmap_form = _form()
- self.combo_cmap = QComboBox()
- self.combo_cmap.addItems(CMAPS)
- self.combo_cmap.setCurrentText("gray")
- self.combo_cmap.setEnabled(False)
- self.combo_cmap.setSizePolicy(QSizePolicy.Policy.Expanding,
- QSizePolicy.Policy.Fixed)
- self.combo_cmap.currentIndexChanged.connect(self._on_cmap_changed)
- cmap_form.addRow("Colormap:", self.combo_cmap)
- dl.addLayout(cmap_form)
-
- self.chk_auto = QCheckBox("Auto-scale colormap")
- self.chk_auto.setChecked(True)
- self.chk_auto.toggled.connect(self._on_autoscale_toggled)
- dl.addWidget(self.chk_auto)
-
- range_form = _form()
- for label, attr in (("min:", "spin_vmin"), ("max:", "spin_vmax")):
- spin = QDoubleSpinBox()
- spin.setRange(-1e9, 1e9)
- spin.setDecimals(4)
- spin.setEnabled(False)
- spin.setMinimumWidth(_SPIN_MIN_W)
- spin.editingFinished.connect(self._on_manual_range_changed)
- setattr(self, attr, spin)
- range_form.addRow(label, spin)
- dl.addLayout(range_form)
-
- right_panel = QWidget()
- layout = QVBoxLayout(right_panel)
- layout.setContentsMargins(0, 0, 0, 0)
- layout.setSpacing(8)
- layout.addWidget(self.grp_velocity)
- layout.addWidget(grp_display)
- layout.addStretch()
- return _scroll_panel(right_panel, _RIGHT_PANEL_W)
-
- def _build_menus(self):
- menubar = self.menuBar()
-
- fft_menu = menubar.addMenu("&FFT")
- fft_act = QAction("FFT &Options…", self)
- fft_act.setStatusTip("Configure FFT backend and zero-padding")
- fft_act.triggered.connect(self._on_fft_options)
- fft_menu.addAction(fft_act)
-
- fusion_menu = menubar.addMenu("&Fusion")
- self._alignment_act = QAction("Angle &Alignment", self)
- self._alignment_act.setStatusTip(
- "Compute a rotation+translation alignment across all angles "
- "(from CH4 masks) and enable Aligned View. Requires >1 angle.")
- self._alignment_act.setEnabled(False)
- self._alignment_act.triggered.connect(self._on_angle_alignment)
- fusion_menu.addAction(self._alignment_act)
-
- self._manual_align_act = QAction("&Manual Alignment…", self)
- self._manual_align_act.setStatusTip(
- "Open an interactive dialog to align angles by eye: overlaid CH4 "
- "threshold masks, keyboard nudge (translate + rotate), auto "
- "de-rotate to the known scan angles, and save/clear a persistent "
- "alignment.")
- self._manual_align_act.setEnabled(False)
- self._manual_align_act.triggered.connect(self._on_manual_alignment)
- fusion_menu.addAction(self._manual_align_act)
-
- convert_menu = menubar.addMenu("&Convert")
- self._batch_dc_act = QAction("Batch Compute DC and &Store…", self)
- self._batch_dc_act.setStatusTip(
- "Select .sras files and compute+store DC images (CH3/CH4 mean) "
- "for every angle, converting v6 files to v7 in place.")
- self._batch_dc_act.triggered.connect(lambda: self._on_batch_compute("dc"))
- convert_menu.addAction(self._batch_dc_act)
-
- self._batch_fft_act = QAction("Batch Compute FFT and Sto&re…", self)
- self._batch_fft_act.setStatusTip(
- "Select .sras files and compute+store FFT peak-frequency images "
- "for every angle, converting v6 files to v7 in place.")
- self._batch_fft_act.triggered.connect(lambda: self._on_batch_compute("fft"))
- convert_menu.addAction(self._batch_fft_act)
-
- export_menu = menubar.addMenu("&Export")
- self._batch_export_act = QAction("&Batch Export Images…", self)
- self._batch_export_act.setStatusTip(
- "Export DC/RF/Velocity map PNGs for every angle of the open "
- "file, with a fixed colorbar range per channel.")
- self._batch_export_act.setEnabled(False)
- self._batch_export_act.triggered.connect(self._on_batch_export)
- export_menu.addAction(self._batch_export_act)
-
- # ------------------------------------------------------------------
- # Drag-and-drop
- # ------------------------------------------------------------------
-
- def dragEnterEvent(self, event):
- urls = event.mimeData().urls()
- if urls and urls[0].toLocalFile().lower().endswith(".sras"):
- event.acceptProposedAction()
-
- def dropEvent(self, event):
- self._load_file(event.mimeData().urls()[0].toLocalFile())
-
- # ------------------------------------------------------------------
- # File loading
- # ------------------------------------------------------------------
-
- def _on_open(self):
- path, _ = QFileDialog.getOpenFileName(
- self, "Open SRAS File", "", "SRAS Files (*.sras);;All Files (*)")
- if path:
- self._load_file(path)
-
- def _load_file(self, path: str):
- started = self._run_worker(
- "load", LoadWorker(path),
- connect=(
- ("finished", self._on_load_done),
- ("error", lambda msg: self.statusBar().showMessage(f"Error: {msg}")),
- ),
- )
- if not started:
- return
- self.btn_open.setEnabled(False)
- self.statusBar().showMessage(f"Loading {Path(path).name}…")
- self._show_progress("load", f"Loading {Path(path).name}…")
-
- def _on_load_done(self, sras):
- self._close_progress("load")
- self.btn_open.setEnabled(True)
- if sras is None:
- return
- self._sras = sras
- self._current_image = None
-
- # A manual-alignment dialog bound to the previous file must not
- # survive a reload — its per-angle state (and the sras it was
- # constructed against) no longer matches the new file's geometry.
- if self._manual_align_dialog is not None:
- self._manual_align_dialog.close()
- self._manual_align_dialog = None
-
- # Caches (and any in-flight DC precompute) belong to the previous
- # file's geometry — discard and start fresh. Bumping the generation
- # counters makes any still-running worker's result get dropped when
- # it lands.
- self._dc_cache = {}
- self._fft_cache = {}
- self._dc_generation += 1
- self._compute_generation += 1
- self.lbl_dc_precompute.setText("")
-
- self._alignment_result = None
- self._aligned_cache = {}
- self._alignment_generation += 1
- self.chk_aligned_view.blockSignals(True)
- self.chk_aligned_view.setChecked(False)
- self.chk_aligned_view.setEnabled(False)
- self.chk_aligned_view.blockSignals(False)
-
- # Silently restore a previously-saved manual alignment, if any, so
- # the work survives closing and reopening the file.
- sidecar = load_manual_alignment(sras)
- if sidecar is not None:
- try:
- self._alignment_result = build_manual_alignment(
- sras, sidecar.ref_angle_idx, sidecar.dc_threshold_mv,
- sidecar.per_angle)
- self.chk_aligned_view.blockSignals(True)
- self.chk_aligned_view.setChecked(True)
- self.chk_aligned_view.blockSignals(False)
- self.statusBar().showMessage(
- f"Restored saved manual alignment from "
- f"{sidecar_path(sras.path).name}")
- except Exception as exc:
- # A corrupt/foreign sidecar or a rescan that shrank n_angles
- # below ref_angle_idx must not block opening the .sras file.
- self.statusBar().showMessage(
- f"Could not restore saved alignment: {exc}")
-
- # A ROI from the previous file no longer matches the new scan's
- # geometry, so discard it on every load.
- self.image_canvas.clear_roi()
-
- self.lbl_filename.setText(sras.path.name)
-
- self.spin_angle.blockSignals(True)
- self.spin_angle.setRange(0, max(0, sras.n_angles - 1))
- self.spin_angle.setValue(0)
- self.spin_angle.blockSignals(False)
-
- # DC channels are cheap and give an instant, fluid overview of a
- # scan; CH1/Velocity require an FFT per pixel that can take minutes
- # on a large scan, so don't default to it.
- self.combo_channel.blockSignals(True)
- self.combo_channel.setCurrentIndex(CH4_IDX)
- self.combo_channel.blockSignals(False)
-
- self._update_controls_enabled(True)
- self._on_threshold_changed() # refresh ADC label with file calibration
- self._on_view_changed()
- self._start_dc_precompute()
-
- # ------------------------------------------------------------------
- # Scan info panel
- # ------------------------------------------------------------------
-
- def _update_scan_info_labels(self):
- s = self._sras
- if s is None:
- return
- a = self.spin_angle.value()
- for key, text in (
- ("Angles", f"{s.n_angles}"),
- ("Rows", f"{s.n_rows[a]}"),
- ("Frames / row", f"{s.n_frames[a]}"),
- ("Samples / frame", f"{s.samples_per_frame}"),
- ("Sample rate", f"{s.sample_rate_hz / 1e9:.4g} GS/s"),
- ("X start", f"{s.x_start_mm[a]:.4g} mm"),
- ("Pixel Δx", f"{s.pixel_x_mm * 1e3:.3g} µm"),
- ("Laser freq", f"{s.laser_freq_hz / 1e3:.4g} kHz"),
- ):
- self._info[key].setText(f"{key}: {text}")
-
- notes = []
- if s.frame_count_mismatch:
- notes.append(f"! Header n_frames={s.n_frames_header}, "
- f"actual={s.n_frames[a]} (scanner bug — corrected)")
- if s.scan_aborted:
- notes.append(f"! Scan aborted: {s.n_angles}/{s.n_angles_declared} "
- "angles complete")
- if s.background is not None:
- notes.append(f"Background waveform: {len(s.background)} samples")
- if s.version in (6, 7):
- notes.append("v6/v7 format: rows / frames / x_start are per-angle")
-
- n_dc = sum(1 for x in s.precomputed_dc4_mv if x is not None)
- n_fft = sum(1 for x in s.precomputed_freq_mhz if x is not None)
- if n_dc or n_fft:
- bg_note = " (bg-sub)" if s.precomputed_bg_sub else " (no bg-sub)"
- notes.append(
- f"Cached images: DC {n_dc}/{s.n_angles} angles, "
- f"FFT {n_fft}/{s.n_angles} angles{bg_note if n_fft else ''} "
- "— display is instant for cached angles")
- elif s.version == 7:
- notes.append("v7 format: no cache blocks stored yet")
- self.lbl_frame_warn.setText("\n".join(notes))
-
- # ------------------------------------------------------------------
- # Controls
- # ------------------------------------------------------------------
-
- def _update_controls_enabled(self, enabled: bool):
- s = self._sras
- has_file = enabled and s is not None
- ch_idx = self.combo_channel.currentIndex()
- is_ch1 = enabled and ch_idx in CH1_DERIVED_MODES
- is_vel = enabled and ch_idx == VELOCITY_MODE_IDX
-
- self.spin_angle.setEnabled(has_file and s.n_angles > 1)
- self.combo_channel.setEnabled(enabled)
- self.combo_cmap.setEnabled(enabled)
- self.chk_auto.setEnabled(enabled)
- manual = enabled and not self.chk_auto.isChecked()
- self.spin_vmin.setEnabled(manual)
- self.spin_vmax.setEnabled(manual)
-
- # Threshold and bg-sub apply to all CH1 modes
- self.spin_threshold_mv.setEnabled(is_ch1)
- self.chk_bg_sub.setEnabled(has_file and s.background is not None and is_ch1)
- self.spin_grating_um.setEnabled(is_vel)
- self.grp_velocity.setVisible(is_vel)
-
- self.btn_export_csv.setEnabled(is_ch1 and self._current_image is not None)
- # ROI: always usable once a file is loaded (independent of channel)
- self.btn_draw_roi.setEnabled(has_file)
-
- # Batch Convert actions pick their own files, independent of
- # whatever's currently open — only gated on no batch already running.
- can_batch = not self._job_running("batch")
- self._batch_dc_act.setEnabled(can_batch)
- self._batch_fft_act.setEnabled(can_batch)
-
- self._alignment_act.setEnabled(
- has_file and s.n_angles > 1 and not self._job_running("align"))
- self._manual_align_act.setEnabled(
- has_file and s.n_angles > 1 and not self._job_running("align"))
- self.chk_aligned_view.setEnabled(enabled and self._alignment_result is not None)
-
- self._batch_export_act.setEnabled(has_file and not self._job_running("export"))
- self._update_roi_ui()
-
- def _on_channel_changed(self):
- self._update_controls_enabled(self._sras is not None)
- self._on_view_changed()
-
- def _on_bg_sub_toggled(self):
- # Background subtraction changes the FFT input, so it genuinely
- # invalidates the cached raw FFT (the cache key includes it) —
- # _refresh_display() recomputes only on a miss for the new state.
- if self._sras is not None and self.combo_channel.currentIndex() in CH1_DERIVED_MODES:
- self._refresh_display()
-
- def _on_grating_changed(self):
- # Grating is a pure post-multiply on the cached frequency image —
- # never needs a recompute.
- if self._sras is not None and self.combo_channel.currentIndex() == VELOCITY_MODE_IDX:
- self._refresh_display()
-
- def _on_threshold_changed(self):
- mv = self.spin_threshold_mv.value()
- cal = (self._sras.cal(CH4_IDX) if self._sras is not None
- else (_FALLBACK_YMULT_MV, _FALLBACK_YOFF_ADC, 0.0))
- self.lbl_threshold_adc.setText(f"≈ {mv_to_adc(mv, *cal):.1f} ADC counts")
- # Threshold decides which pixels get an FFT at all, so changing it is
- # a genuine cache-key change — but the recompute reuses the cached DC4
- # image to skip masked-out pixels.
- if self._sras is not None and self.combo_channel.currentIndex() in CH1_DERIVED_MODES:
- self._refresh_display()
-
- def _on_autoscale_toggled(self, checked: bool):
- manual = not checked
- self.spin_vmin.setEnabled(manual and self._sras is not None)
- self.spin_vmax.setEnabled(manual and self._sras is not None)
- if self._sras is not None and self._current_image is not None:
- self._redraw_image(self._current_image)
-
- def _on_manual_range_changed(self):
- if not self.chk_auto.isChecked() and self._current_image is not None:
- self._redraw_image(self._current_image)
-
- def _on_cmap_changed(self):
- # Colormap is purely how the existing image is rendered.
- if self._current_image is not None:
- self._redraw_image(self._current_image)
-
- def _on_view_changed(self):
- if self._sras is None:
- return
- idx = self.spin_angle.value()
- self.lbl_angle_deg.setText(f"({self._sras.angles_deg[idx]:.1f}°)")
- self._update_scan_info_labels()
- self._refresh_display()
-
- def _on_aligned_view_toggled(self, checked: bool):
- if self._current_image is not None:
- self._redraw_image(self._current_image)
-
- # ------------------------------------------------------------------
- # CSV export
- # ------------------------------------------------------------------
-
- def _on_export_csv(self):
- if self._current_image is None or self._sras is None:
- return
- default_name = (f"{self._sras.path.stem}_angle{self._current_angle}"
- f"_{CH_NAMES[self._current_ch]}.csv")
- path, _ = QFileDialog.getSaveFileName(
- self, "Export Image as CSV",
- str(self._sras.path.parent / default_name),
- "CSV files (*.csv);;All files (*)")
- if not path:
- return
- np.savetxt(path, self._current_image, delimiter=",", fmt="%.6g")
- self.statusBar().showMessage(f"Exported {Path(path).name}")
-
- def _on_export_roi_csv(self):
- if self._current_image is None or self._sras is None:
- return
- roi = self.image_canvas.get_roi()
- if roi is None:
- self.statusBar().showMessage("No ROI — draw one first")
- return
- s = self._sras
- x_axis = s.x_axis_mm(self._current_angle)
- y_axis = s.y_positions_mm(self._current_angle)
- mask = roi.mask_for_grid(x_axis, y_axis)
- if not mask.any():
- self.statusBar().showMessage("ROI does not overlap any pixel")
- return
- img = self._current_image
- if img.shape != mask.shape:
- self.statusBar().showMessage(
- f"ROI shape {mask.shape} does not match image {img.shape}")
- return
-
- X, Y = np.meshgrid(np.asarray(x_axis, dtype=np.float64),
- np.asarray(y_axis, dtype=np.float64))
- rows_idx, frames_idx = np.where(mask)
- n_pix = int(mask.sum())
-
- ch_name = CH_NAMES[self._current_ch]
- angle = self._current_angle
- default_name = f"{s.path.stem}_angle{angle}_{ch_name}_ROI.csv"
- path, _ = QFileDialog.getSaveFileName(
- self, "Export ROI as CSV",
- str(s.path.parent / default_name),
- "CSV files (*.csv);;All files (*)")
- if not path:
- return
-
- corners_str = " ".join(f"({p[0]:.6g},{p[1]:.6g})" for p in roi.corners())
- header = (
- f"# ROI quad corners (BL BR TR TL) mm: {corners_str}\n"
- f"# source: {s.path.name}, channel={ch_name}, "
- f"angle_idx={angle}, angle_deg={s.angles_deg[angle]:.4g}\n"
- f"# n_pixels={n_pix}\n"
- "row,frame,x_mm,y_mm,value"
- )
- data = np.column_stack([
- rows_idx.astype(np.int64), frames_idx.astype(np.int64),
- X[mask], Y[mask], img[mask].astype(np.float64),
- ])
- # integer columns first, floats after — use a per-column format list
- np.savetxt(path, data, delimiter=",",
- fmt=["%d", "%d", "%.6g", "%.6g", "%.6g"],
- header=header, comments="")
- self.statusBar().showMessage(
- f"Exported ROI ({n_pix} pixels) to {Path(path).name}")
-
- # ------------------------------------------------------------------
- # ROI
- # ------------------------------------------------------------------
-
- def _on_draw_roi_toggled(self, checked: bool):
- if checked:
- self.image_canvas.start_drawing()
- self.statusBar().showMessage(
- "Click and drag on the image to draw a new rectangle.")
- else:
- self.image_canvas.cancel_drawing()
-
- def _on_draw_mode_changed(self, active: bool):
- # Keep the toggle button's visual state in sync with the canvas.
- self.btn_draw_roi.blockSignals(True)
- self.btn_draw_roi.setChecked(active)
- self.btn_draw_roi.blockSignals(False)
-
- def _on_clear_roi(self):
- self.image_canvas.clear_roi()
- self.statusBar().showMessage("ROI cleared")
-
- def _update_roi_ui(self):
- roi = self.image_canvas.get_roi()
- if roi is None:
- self.lbl_roi_center.setText("centroid: —")
- self.lbl_roi_size.setText("bbox: —")
- self.lbl_roi_npix.setText("pixels inside: —")
- self.btn_clear_roi.setEnabled(False)
- self.btn_export_roi.setEnabled(False)
- return
-
- cen = roi.centroid()
- bbox = roi.bbox_size()
- self.lbl_roi_center.setText(f"centroid: ({cen[0]:.3f}, {cen[1]:.3f}) mm")
- self.lbl_roi_size.setText(f"bbox: {bbox[0]:.3f} × {bbox[1]:.3f} mm")
-
- npix = 0
- if self._sras is not None:
- try:
- # Deliberately always the raw per-angle grid, even when
- # Aligned View is on: _on_export_roi_csv also exports on the
- # raw grid (never synthetically-resampled pixels), so this
- # readout must match what Export ROI actually writes.
- mask = roi.mask_for_grid(
- self._sras.x_axis_mm(self._current_angle),
- self._sras.y_positions_mm(self._current_angle))
- npix = int(mask.sum())
- except Exception:
- npix = 0
- self.lbl_roi_npix.setText(f"pixels inside: {npix}")
- self.btn_clear_roi.setEnabled(True)
- self.btn_export_roi.setEnabled(self._current_image is not None and npix > 0)
-
- # ------------------------------------------------------------------
- # Display
- # ------------------------------------------------------------------
-
- def _current_n_fft(self) -> int | None:
- if self._fft_pad_factor <= 1 or self._sras is None:
- return None
- return self._sras.samples_per_frame * self._fft_pad_factor
-
- def _scale_for_display(self, freq_mhz: np.ndarray, ch_idx: int) -> np.ndarray:
- """Velocity is a pure post-multiply of the (already DC-masked)
- cached frequency image — never worth a recompute on its own."""
- if ch_idx == VELOCITY_MODE_IDX:
- return freq_mhz * self.spin_grating_um.value()
- return freq_mhz
-
- def _fft_cache_key(self, angle_idx: int) -> tuple:
- return (angle_idx, self.chk_bg_sub.isChecked(), self._current_n_fft(),
- self.spin_threshold_mv.value())
-
- def _aligned_cache_key(self, angle_idx: int, ch_idx: int) -> tuple:
- """Mirrors _fft_cache's key granularity so a stale aligned image is
- never shown after bg_sub/threshold/pad/grating changes."""
- if ch_idx in CH1_DERIVED_MODES:
- return (*self._fft_cache_key(angle_idx), ch_idx,
- self.spin_grating_um.value() if ch_idx == VELOCITY_MODE_IDX else None)
- return (angle_idx, ch_idx)
-
- def _aligned_canvas_axes(self) -> tuple[np.ndarray, np.ndarray]:
- r = self._alignment_result
- n_rows, n_cols = r.canvas_shape
- return (r.canvas_origin_mm[0] + np.arange(n_cols) * r.canvas_dx_mm,
- r.canvas_origin_mm[1] + np.arange(n_rows) * r.canvas_dy_mm)
-
- def _get_aligned_display_image(self, raw_img: np.ndarray, angle_idx: int,
- ch_idx: int) -> np.ndarray:
- key = self._aligned_cache_key(angle_idx, ch_idx)
- cached = self._aligned_cache.get(key)
- if cached is None:
- cached = apply_alignment(self._alignment_result, angle_idx, raw_img)
- self._aligned_cache[key] = cached
- return cached
-
- def _refresh_display(self):
- """Show the image for the current angle/channel/threshold, using
- cached data whenever possible and only falling back to a background
- compute (with progress popup) when genuinely nothing is cached yet."""
- if self._sras is None:
- return
- angle_idx = self.spin_angle.value()
- ch_idx = self.combo_channel.currentIndex()
-
- if ch_idx in CH1_DERIVED_MODES:
- raw = self._fft_cache.get(self._fft_cache_key(angle_idx))
- if raw is not None:
- self._show_image_now(self._scale_for_display(raw, ch_idx),
- angle_idx, ch_idx)
- return
- else:
- cached = self._dc_cache.get((angle_idx, ch_idx))
- if cached is not None:
- self._show_image_now(cached, angle_idx, ch_idx)
- return
-
- # Nothing cached for these settings — need a real compute. Changing
- # the DC threshold changes *which* pixels get an FFT at all, so it
- # can't be satisfied from the cache — but with the DC map already
- # known, the recompute skips the FFT for masked-out pixels.
- self._start_compute()
-
- def _show_image_now(self, img: np.ndarray, angle_idx: int, ch_idx: int):
- """Display an already-available image with no compute involved."""
- self._current_image = img
- self._current_angle = angle_idx
- self._current_ch = ch_idx
- self.btn_export_csv.setEnabled(ch_idx in CH1_DERIVED_MODES)
- self._redraw_image(img)
- self._update_roi_ui()
-
- def _redraw_image(self, img: np.ndarray):
- s = self._sras
- angle_idx = self._current_angle
- ch_idx = self._current_ch
-
- aligned = (self.chk_aligned_view.isChecked()
- and self._alignment_result is not None
- and angle_idx in self._alignment_result.per_angle)
- if aligned:
- display_img = self._get_aligned_display_image(img, angle_idx, ch_idx)
- x_axis, y_axis = self._aligned_canvas_axes()
- else:
- display_img = img
- x_axis = s.x_axis_mm(angle_idx)
- y_axis = s.y_positions_mm(angle_idx)
-
- dx = x_axis[1] - x_axis[0] if len(x_axis) > 1 else s.pixel_x_mm
- dy = float(y_axis[1] - y_axis[0]) if len(y_axis) > 1 else 1.0
- extent = [x_axis[0] - dx / 2, x_axis[-1] + dx / 2,
- y_axis[-1] + dy / 2, y_axis[0] - dy / 2]
-
- if self.chk_auto.isChecked():
- vmin, vmax = float(display_img.min()), float(display_img.max())
- for spin, val in ((self.spin_vmin, vmin), (self.spin_vmax, vmax)):
- spin.blockSignals(True)
- spin.setValue(val)
- spin.blockSignals(False)
- else:
- vmin, vmax = self.spin_vmin.value(), self.spin_vmax.value()
-
- angle_deg = s.angles_deg[angle_idx]
- mode_str, unit, colorbar_label = _CHANNEL_DISPLAY[ch_idx]
- if ch_idx == VELOCITY_MODE_IDX:
- ch_label = f"Velocity [grating={self.spin_grating_um.value():.2f} µm]"
- else:
- ch_label = CH_LABELS[ch_idx]
-
- title = f"{CH_NAMES[ch_idx]} | {mode_str} | {angle_deg:.1f}°"
- if aligned:
- title += " [Aligned]"
-
- self.image_canvas.show_image(
- display_img, extent,
- cmap=self.combo_cmap.currentText(),
- vmin=vmin, vmax=vmax,
- xlabel="X (mm)", ylabel="Y (mm)",
- title=title, colorbar_label=colorbar_label,
- )
- self.statusBar().showMessage(
- f"{s.path.name} | {ch_label} @ {angle_deg:.1f}° "
- f"| {display_img.shape[1]} × {display_img.shape[0]} px | {unit}"
- f"{' | Aligned' if aligned else ''}"
- )
-
- # ------------------------------------------------------------------
- # Background compute (only reached on a genuine cache miss)
- # ------------------------------------------------------------------
-
- def _start_compute(self):
- if self._sras is None or self._job_running("compute"):
- return # re-checked when the running compute finishes
-
- angle_idx = self.spin_angle.value()
- ch_idx = self.combo_channel.currentIndex()
- is_fft = ch_idx in CH1_DERIVED_MODES
-
- self._pending_angle = angle_idx
- self._pending_ch = ch_idx
- self._pending_bg_sub = self.chk_bg_sub.isChecked()
- self._pending_threshold = self.spin_threshold_mv.value()
- self._pending_fft_pad_factor = self._fft_pad_factor
- # Snapshot which file this compute belongs to, so a result that lands
- # after a *different* file has since been loaded (nothing blocks
- # opening a new file while a compute is in flight) gets discarded
- # instead of being cached/displayed against the wrong SrasFile.
- self._pending_compute_generation = self._compute_generation
-
- worker = ComputeWorker(
- self._sras, angle_idx, ch_idx,
- apply_bg_sub=self._pending_bg_sub,
- n_fft=self._current_n_fft(),
- dc_threshold_mv=self._pending_threshold,
- # Reuse the cached DC4 image (if the precompute has reached this
- # angle) so the FFT skips masked-out pixels entirely and doesn't
- # need to re-read the CH4 channel from disk.
- dc4_mv=self._dc_cache.get((angle_idx, CH4_IDX)),
- is_fft_mode=is_fft,
- )
- if not self._run_worker(
- "compute", worker,
- connect=(
- ("finished", self._on_compute_done),
- ("error", lambda msg: self.statusBar().showMessage(
- f"Compute error: {msg}")),
- ),
- on_done=self._after_compute):
- return
-
- if is_fft:
- self.statusBar().showMessage("Computing FFT…")
- self._show_progress(
- "compute",
- f"Computing FFT for angle {angle_idx}…\n"
- "This can take a while on a large scan — result is cached "
- "so revisiting this angle/mode/threshold will be instant.")
- else:
- self.statusBar().showMessage("Computing DC image…")
- self._show_progress("compute", f"Computing DC image for angle {angle_idx}…")
-
- def _after_compute(self):
- """If settings changed while the compute was running, re-dispatch
- through the cache-aware path — the now-current combination may
- already be cached."""
- if (self.spin_angle.value(), self.combo_channel.currentIndex(),
- self.chk_bg_sub.isChecked(), self.spin_threshold_mv.value(),
- self._fft_pad_factor) != (
- self._pending_angle, self._pending_ch, self._pending_bg_sub,
- self._pending_threshold, self._pending_fft_pad_factor):
- self._refresh_display()
-
- def _on_compute_done(self, result):
- self._close_progress("compute")
- generation, self._pending_compute_generation = self._pending_compute_generation, None
- if generation != self._compute_generation:
- return # a different file was loaded while this was computing — discard
- if result is None:
- return # cancelled mid-compute; the partial image must not cache
- angle_idx = self._pending_angle
- ch_idx = self._pending_ch
-
- if ch_idx in CH1_DERIVED_MODES:
- self._fft_cache[(angle_idx, self._pending_bg_sub,
- self._current_n_fft(), self._pending_threshold)] = result
- img = self._scale_for_display(result, ch_idx)
- else:
- img = result
- self._dc_cache[(angle_idx, ch_idx)] = img
-
- self._show_image_now(img, angle_idx, ch_idx)
-
- # ------------------------------------------------------------------
- # Background DC precompute (all angles, so switching is fluid)
- # ------------------------------------------------------------------
-
- def _start_dc_precompute(self):
- if self._sras is None:
- return
- generation = self._dc_generation
- n_angles = self._sras.n_angles
-
- worker = DcPrecomputeWorker(self._sras)
- self._dc_precompute_worker = worker
- started = self._run_worker(
- "dc_precompute", worker,
- connect=(
- ("angle_done", lambda a, dc3, dc4, g=generation:
- self._on_dc_precompute_angle_done(g, a, dc3, dc4, n_angles)),
- ("error", lambda msg: self.statusBar().showMessage(
- f"DC precompute error: {msg}", 5000)),
- ),
- quit_on=("finished", "error"),
- on_done=lambda: setattr(self, "_dc_precompute_worker", None),
- )
- if not started:
- self._dc_precompute_worker = None
-
- def _on_dc_precompute_angle_done(self, generation: int, angle_idx: int,
- dc3_mv: np.ndarray, dc4_mv: np.ndarray,
- n_angles: int):
- if generation != self._dc_generation:
- return # stale result from a previously-loaded file — discard
- self._dc_cache[(angle_idx, CH3_IDX)] = dc3_mv
- self._dc_cache[(angle_idx, CH4_IDX)] = dc4_mv
-
- done = sum(1 for a in range(n_angles) if (a, CH4_IDX) in self._dc_cache)
- self.lbl_dc_precompute.setText(
- f"Precomputing DC images: {done}/{n_angles} angles ready…"
- if done < n_angles else "DC images ready for all angles.")
-
- # If we just finished the angle/channel the user is currently looking
- # at and it wasn't shown yet (they switched here before the precompute
- # caught up and are still waiting), show it now.
- current_ch = self.combo_channel.currentIndex()
- if (angle_idx == self.spin_angle.value()
- and not self._job_running("compute")
- and current_ch in (CH3_IDX, CH4_IDX)
- and (self._current_angle != angle_idx or self._current_ch != current_ch)):
- self._refresh_display()
-
- # ------------------------------------------------------------------
- # Pixel inspector
- # ------------------------------------------------------------------
-
- def _on_pixel_clicked(self, row_idx: int, frame_idx: int):
- if self._sras is None or self._current_image is None:
- return
- angle_idx = self._current_angle
- if (self.chk_aligned_view.isChecked() and self._alignment_result is not None
- and angle_idx in self._alignment_result.per_angle):
- # The click landed on the shared aligned canvas — invert the same
- # canvas->raw affine used to display it back to a raw (row, frame)
- # index before looking up the waveform.
- t = self._alignment_result.per_angle[angle_idx]
- raw = t.matrix @ np.array([row_idx, frame_idx], dtype=np.float64) + t.offset
- row_idx, frame_idx = int(round(raw[0])), int(round(raw[1]))
- n_rows_a, n_frames_a = self._sras.image_shape(angle_idx)
- if not (0 <= row_idx < n_rows_a and 0 <= frame_idx < n_frames_a):
- self.statusBar().showMessage(
- "No source waveform here (padding region of the aligned canvas).")
- return
-
- self.lbl_wave_hint.hide()
- if self._current_ch in CH1_DERIVED_MODES:
- self.wave_canvas.show_rf_waveform(
- self._sras, angle_idx, row_idx, frame_idx,
- apply_bg_sub=self.chk_bg_sub.isChecked())
- else:
- self.wave_canvas.show_dc_waveform(
- self._sras, angle_idx, self._current_ch, row_idx, frame_idx)
-
- # ------------------------------------------------------------------
- # Progress dialogs
- # ------------------------------------------------------------------
-
- def _show_progress(self, key: str, message: str, maximum: int = 0):
- """Show (or relabel) the progress dialog under *key*. maximum=0 gives
- an indeterminate busy indicator."""
- dlg = self._progress_dlgs.get(key)
- if dlg is not None:
- dlg.setLabelText(message)
- return
- dlg = QProgressDialog(message, "", 0, maximum, self)
- dlg.setWindowTitle("Please wait…")
- dlg.setCancelButton(None)
- dlg.setWindowModality(Qt.WindowModality.WindowModal)
- dlg.setMinimumDuration(300) # only appears if it takes > 300 ms
- dlg.show()
- self._progress_dlgs[key] = dlg
-
- def _set_progress(self, key: str, pct: int):
- dlg = self._progress_dlgs.get(key)
- if dlg is not None:
- dlg.setValue(pct)
-
- def _close_progress(self, key: str):
- dlg = self._progress_dlgs.pop(key, None)
- if dlg is not None:
- dlg.close()
-
- # ------------------------------------------------------------------
- # Convert menu: batch DC/FFT compute-and-store (v6 -> v7)
- # ------------------------------------------------------------------
-
- def _on_batch_compute(self, mode: str):
- if self._job_running("batch"):
- return
- label = "DC" if mode == "dc" else "FFT"
- paths, _ = QFileDialog.getOpenFileNames(
- self, f"Select .sras files to batch-compute {label}", "",
- "SRAS files (*.sras);;All files (*)")
- if not paths:
- return
-
- self._batch_errors = []
- worker = BatchCacheWorker(paths, mode, self.chk_bg_sub.isChecked())
- started = self._run_worker(
- "batch", worker,
- connect=(
- ("progress", lambda pct: self._set_progress("batch", pct)),
- ("file_done", self._on_batch_file_done),
- ("finished", lambda p=paths: self._on_batch_finished(p)),
- ),
- on_done=self._after_batch,
- )
- if not started:
- return # a second trigger snuck in while the file dialog was open
-
- self._batch_dc_act.setEnabled(False)
- self._batch_fft_act.setEnabled(False)
- self._show_progress(
- "batch", f"Batch computing {label} for {len(paths)} file(s)…",
- maximum=100)
-
- def _on_batch_file_done(self, path: str, err: str):
- if err:
- self._batch_errors.append(f"{Path(path).name} — {err}")
- self._show_progress("batch", f"Processed {Path(path).name}…")
-
- def _on_batch_finished(self, paths: list[str]):
- self._close_progress("batch")
-
- n_total = len(paths)
- n_failed = len(self._batch_errors)
- n_ok = n_total - n_failed
- if n_failed:
- summary = (f"Batch store: {n_ok}/{n_total} file(s) updated, "
- f"{n_failed} failed: {'; '.join(self._batch_errors)}")
- else:
- summary = f"Batch store: {n_ok}/{n_total} file(s) updated."
- self.statusBar().showMessage(summary)
- self._batch_errors = []
-
- # If the currently-open file was in this batch, reload it so the GUI
- # picks up the newly-written v7 cache instead of stale state.
- if self._sras is not None and str(self._sras.path) in paths:
- self._load_file(str(self._sras.path))
-
- def _after_batch(self):
- self._batch_dc_act.setEnabled(True)
- self._batch_fft_act.setEnabled(True)
-
- # ------------------------------------------------------------------
- # Export menu: batch image export
- # ------------------------------------------------------------------
-
- def _export_default_ranges(self) -> dict[tuple[int, bool], tuple[float, float]]:
- """Best-effort default min/max per exportable channel, from whatever
- is already cached — this must never trigger a fresh compute just to
- seed the dialog (CH1/Velocity in particular can be expensive)."""
- s = self._sras
- ranges: dict[tuple[int, bool], tuple[float, float]] = {}
-
- for ch_idx, precomputed in ((CH3_IDX, s.precomputed_dc3_mv),
- (CH4_IDX, s.precomputed_dc4_mv)):
- arrays = [self._dc_cache.get((a, ch_idx)) for a in range(s.n_angles)]
- arrays = [a for a in arrays if a is not None]
- if not arrays:
- arrays = [a for a in precomputed if a is not None]
- if arrays:
- ranges[(ch_idx, False)] = (
- float(min(a.min() for a in arrays)),
- float(max(a.max() for a in arrays)))
-
- freq_arrays = list(self._fft_cache.values())
- if not freq_arrays:
- freq_arrays = [a for a in s.precomputed_freq_mhz if a is not None]
- if freq_arrays:
- fmin = float(min(a.min() for a in freq_arrays))
- fmax = float(max(a.max() for a in freq_arrays))
- ranges[(CH1_IDX, False)] = (fmin, fmax)
- grating = self.spin_grating_um.value()
- ranges[(CH1_IDX, True)] = (fmin * grating, fmax * grating)
-
- return ranges
-
- def _on_batch_export(self):
- if self._sras is None or self._job_running("export"):
- return
- s = self._sras
-
- dlg = BatchExportDialog(
- self, default_dir=str(s.path.parent), default_prefix=s.path.stem,
- default_ranges=self._export_default_ranges())
- if dlg.exec() != QDialog.DialogCode.Accepted:
- return
-
- output_dir = dlg.get_output_dir()
- try:
- Path(output_dir).mkdir(parents=True, exist_ok=True)
- except OSError as exc:
- QMessageBox.warning(
- self, "Batch Export", f"Could not create output folder:\n{exc}")
- return
-
- channels = dlg.get_selected_channels()
- self._export_errors = []
- self._export_ok_count = 0
- worker = BatchExportWorker(
- s, channels, output_dir, dlg.get_prefix(),
- cmap=self.combo_cmap.currentText(),
- apply_bg_sub=self.chk_bg_sub.isChecked(),
- dc_threshold_mv=self.spin_threshold_mv.value(),
- n_fft=self._current_n_fft(), grating_um=self.spin_grating_um.value())
- started = self._run_worker(
- "export", worker,
- connect=(
- ("progress", lambda pct: self._set_progress("export", pct)),
- ("file_done", self._on_export_file_done),
- ("finished", self._on_export_finished),
- ),
- on_done=self._after_export,
- )
- if not started:
- return # a second trigger snuck in while the dialog was open
-
- self._batch_export_act.setEnabled(False)
- self._show_progress(
- "export", f"Exporting images to {output_dir}…", maximum=100)
-
- def _on_export_file_done(self, path: str, err: str):
- if err:
- self._export_errors.append(f"{Path(path).name if path else '?'} — {err}")
- else:
- self._export_ok_count += 1
- self._show_progress("export", f"Wrote {Path(path).name}…")
-
- def _on_export_finished(self):
- self._close_progress("export")
- n_ok = self._export_ok_count
- n_failed = len(self._export_errors)
- if n_failed:
- summary = (f"Batch export: {n_ok} image(s) written, {n_failed} "
- f"failed: {'; '.join(self._export_errors)}")
- else:
- summary = f"Batch export: {n_ok} image(s) written."
- self.statusBar().showMessage(summary)
- self._export_errors = []
-
- def _after_export(self):
- self._update_controls_enabled(self._sras is not None)
-
- # ------------------------------------------------------------------
- # Fusion: angle alignment
- # ------------------------------------------------------------------
-
- def _on_angle_alignment(self):
- if self._sras is None or self._sras.n_angles <= 1:
- return
- ref_idx = 0
- threshold_mv = self.spin_threshold_mv.value()
- generation = self._alignment_generation
-
- started = self._run_worker(
- "align", AngleAlignmentWorker(self._sras, ref_idx, threshold_mv),
- connect=(
- ("progress", lambda pct: self._set_progress("align", pct)),
- ("finished", lambda result, err, g=generation:
- self._on_alignment_done(g, result, err)),
- ),
- on_done=lambda: self._update_controls_enabled(self._sras is not None),
- )
- if not started:
- return
-
- self._alignment_act.setEnabled(False)
- self._show_progress(
- "align",
- f"Computing angle alignment ({self._sras.n_angles} angles, "
- f"ref=angle 0, CH4 mask ≥ {threshold_mv:.3f} mV)…",
- maximum=100)
-
- def _on_alignment_done(self, generation: int, result, error_msg: str):
- self._close_progress("align")
- if generation != self._alignment_generation:
- return # a new file was loaded while this was computing — discard
- if error_msg:
- self.statusBar().showMessage(f"Angle alignment failed: {error_msg}")
- return
- self._alignment_result = result
- self._aligned_cache = {}
- self.chk_aligned_view.setEnabled(True)
- self.chk_aligned_view.blockSignals(True)
- self.chk_aligned_view.setChecked(True)
- self.chk_aligned_view.blockSignals(False)
- nr, nc = result.canvas_shape
- self.statusBar().showMessage(
- f"Angle alignment computed ({self._sras.n_angles} angles, "
- f"canvas {nc}×{nr} px).")
- self._refresh_display()
-
- # ------------------------------------------------------------------
- # Fusion: manual alignment
- # ------------------------------------------------------------------
-
- def _on_manual_alignment(self):
- if self._sras is None or self._sras.n_angles <= 1:
- return
- if self._manual_align_dialog is not None:
- self._manual_align_dialog.raise_()
- self._manual_align_dialog.activateWindow()
- return
-
- ref_idx = 0
- threshold_mv = self.spin_threshold_mv.value()
- seed: dict[int, ManualAngleParams] = {}
- # Seed only from a previously *saved manual* alignment (this dialog's
- # own Save also writes this sidecar) -- never from self._alignment_result
- # when it holds the automatic Fusion -> Angle Alignment's output. That
- # path's translation comes from FFT phase correlation, which is the
- # very thing manual mode exists to work around; inheriting it here
- # would silently reintroduce the same bad translations under a
- # "manual" label, on top of the (correct) analytic rotation, which is
- # exactly what makes manual mode look like it "still does the same
- # thing" the automatic one does.
- sidecar = load_manual_alignment(self._sras)
- if sidecar is not None and sidecar.ref_angle_idx == ref_idx:
- seed = dict(sidecar.per_angle)
- threshold_mv = sidecar.dc_threshold_mv
-
- cached_dc4 = {a: img for (a, ch), img in self._dc_cache.items() if ch == CH4_IDX}
- dlg = ManualAlignmentDialog(
- self, self._sras, ref_angle_idx=ref_idx, dc_threshold_mv=threshold_mv,
- seed_per_angle=seed, cached_dc4_mv=cached_dc4)
- dlg.alignment_saved.connect(self._on_manual_alignment_saved)
- dlg.alignment_cleared.connect(self._on_manual_alignment_cleared)
- dlg.finished.connect(self._on_manual_align_dialog_closed)
- # Deliberately NOT WA_DeleteOnClose: Ch4MaskWorker/CrossCorrelateWorker
- # have no stop() and are owned by this window's _jobs (not the
- # dialog), so a background mask-fetch/cross-correlate job can still be
- # running after the dialog closes. Its signals stay connected to this
- # dialog's own bound methods (_on_mask_angle_done, _finish_mask_prep,
- # etc.) until it finishes — if Qt had already destroyed the dialog's
- # C++ object, those late callbacks would raise "wrapped C/C++ object
- # ... has been deleted". Leaving the object merely hidden (not
- # destroyed) lets them land harmlessly; it is garbage-collected once
- # the job's own references to it are released.
- self._manual_align_dialog = dlg
- dlg.show()
-
- def _on_manual_align_dialog_closed(self, _result_code: int):
- self._manual_align_dialog = None
-
- def _on_manual_alignment_saved(self, result, sidecar_path_str: str):
- self._alignment_result = result
- self._aligned_cache = {}
- self._alignment_generation += 1
- self.chk_aligned_view.setEnabled(True)
- self.chk_aligned_view.blockSignals(True)
- self.chk_aligned_view.setChecked(True)
- self.chk_aligned_view.blockSignals(False)
- self._update_controls_enabled(self._sras is not None)
- self.statusBar().showMessage(
- f"Manual alignment saved to {Path(sidecar_path_str).name}")
- if self._current_image is not None:
- self._refresh_display()
-
- def _on_manual_alignment_cleared(self):
- self._alignment_result = None
- self._aligned_cache = {}
- self._alignment_generation += 1
- self.chk_aligned_view.blockSignals(True)
- self.chk_aligned_view.setChecked(False)
- self.chk_aligned_view.setEnabled(False)
- self.chk_aligned_view.blockSignals(False)
- self._update_controls_enabled(self._sras is not None)
- self.statusBar().showMessage("Manual alignment cleared.")
- if self._current_image is not None:
- self._refresh_display()
-
- # ------------------------------------------------------------------
- # FFT Options
- # ------------------------------------------------------------------
-
- def _on_fft_options(self):
- dlg = FftOptionsDialog(
- self,
- current_backend=compute.get_fft_backend(),
- current_pad_factor=self._fft_pad_factor,
- samples_per_frame=self._sras.samples_per_frame if self._sras else None,
- sample_rate_hz=self._sras.sample_rate_hz if self._sras else None,
- grating_um=self.spin_grating_um.value(),
- )
- if dlg.exec() != QDialog.DialogCode.Accepted:
- return
- compute.set_fft_backend(dlg.get_backend())
- self._fft_pad_factor = dlg.get_pad_factor()
- # Pad factor changes the FFT bin count, so it genuinely invalidates
- # the cached raw FFT (part of the cache key) — _refresh_display()
- # recomputes only on a cache miss.
- if self._sras is not None and self.combo_channel.currentIndex() in CH1_DERIVED_MODES:
- self._refresh_display()
-
- # ------------------------------------------------------------------
-
- def closeEvent(self, event):
- if self._manual_align_dialog is not None:
- self._manual_align_dialog.close()
-
- # Signal every cancellable worker first, then wait. Waiting without
- # signalling means sitting out whatever is in flight — on a large
- # scan a single angle is ~40 s.
- jobs = list(self._jobs.values())
- for _thread, worker, _on_done in jobs:
- stop = getattr(worker, "stop", None)
- if callable(stop):
- stop()
- for thread, _worker, _on_done in jobs:
- thread.quit()
- thread.wait(5000)
- super().closeEvent(event)
-
-
-# ---------------------------------------------------------------------------
-
-def main():
- app = QApplication(sys.argv)
- window = SrasViewerWindow(
- initial_path=sys.argv[1] if len(sys.argv) > 1 else None)
- window.show()
- sys.exit(app.exec())
-
-
-if __name__ == "__main__":
- main()
diff --git a/sras_viewer/__init__.py b/sras_viewer/__init__.py
new file mode 100644
index 0000000..5cae377
--- /dev/null
+++ b/sras_viewer/__init__.py
@@ -0,0 +1,29 @@
+"""
+SRAS Scan File Viewer
+PyQt6 application for visualizing channel data from .sras binary scan files.
+
+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
+
+RF images are masked: pixels where CH4_dc < dc_threshold show 0.
+
+File parsing lives in sras_format, image/alignment math in sras_compute, and
+background workers in sras_workers — none of which import Qt or matplotlib,
+so multiprocessing children can load them cheaply.
+"""
+
+import faulthandler
+
+faulthandler.enable() # print a native stack trace on SIGSEGV/SIGABRT/etc.
+
+from .align_wizard import AlignmentWizard # noqa: E402,F401
+from .canvases import ( # noqa: E402,F401
+ AlignOverlayCanvas, ImageCanvas, RoiQuad, WaveformCanvas,
+)
+from .common import CH_LABELS, CMAPS, VELOCITY_MODE_IDX # noqa: E402,F401
+from .dialogs import ( # noqa: E402,F401
+ FftOptionsDialog, FusedRoiExportDialog, RowAverageFftOptionsDialog,
+)
+from .main_window import SrasViewerWindow, main # noqa: E402,F401
diff --git a/sras_viewer/__main__.py b/sras_viewer/__main__.py
new file mode 100644
index 0000000..cbd143d
--- /dev/null
+++ b/sras_viewer/__main__.py
@@ -0,0 +1,4 @@
+from .main_window import main
+
+if __name__ == "__main__":
+ main()
diff --git a/sras_viewer/align_wizard.py b/sras_viewer/align_wizard.py
new file mode 100644
index 0000000..2476448
--- /dev/null
+++ b/sras_viewer/align_wizard.py
@@ -0,0 +1,1446 @@
+"""The alignment wizard: correlate, crop, export.
+
+Replaces what used to be two disconnected menu actions. `Angle Alignment` ran
+the registration with every parameter hard-coded and no way to retry, and
+`Manual Alignment…` was a separate dialog that deliberately refused to inherit
+the automatic result — so a bad fit meant starting over by hand, and neither
+path told you whether the alignment was actually any good.
+
+The three steps are the three decisions:
+
+ 1. Correlate. Every parameter that affects the fit is on the page, the run is
+ repeatable, and the verdict is a picture: each angle's DC mask reprojected
+ onto the shared canvas and summed, so the colour of a pixel is *how many
+ angles cover it*. A good alignment is one saturated plateau at N; a bad one
+ is a fan of low-count halos. Angles start pre-rotated by the stage angles
+ stored in the file, so the page is informative before any correlation runs.
+ Manual nudging lives here too, since this page is now the only place to
+ correct an angle the search cannot fit.
+ 2. Crop. Axis-aligned only, because that is the only crop .sras geometry can
+ express — a free quadrilateral would have to be squared off behind the
+ user's back.
+ 3. Export. Writes the aligned, cropped data to a new .sras.
+
+State lives on the wizard rather than in QWizardPage.registerField: the pages
+share numpy arrays, ManualAngleParams and an AlignmentResult, none of which are
+scalar widget properties.
+"""
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import matplotlib as mpl
+import numpy as np
+from matplotlib.backends.backend_qtagg import NavigationToolbar2QT
+from PyQt6.QtCore import QSignalBlocker, Qt, pyqtSignal
+from PyQt6.QtWidgets import (
+ QCheckBox, QFileDialog, QHBoxLayout, QHeaderView, QLabel,
+ QLineEdit, QMessageBox, QProgressBar, QPushButton, QSpinBox,
+ QTableWidget, QTableWidgetItem, QVBoxLayout, QWidget, QWizard, QWizardPage,
+)
+
+import sras_align_export as export
+import sras_compute as compute
+from sras_compute import ManualAngleParams, build_manual_alignment
+from sras_format import SrasFile
+from sras_workers import AlignedExportWorker, Ch4MaskWorker, CrossCorrelateWorker
+
+from .canvases import AlignOverlayCanvas, ImageCanvas, RoiQuad, count_colormap
+from .common import (
+ _CSS_HINT, _CSS_INFO, _CSS_MUTED, _CSS_WARN, Jobs, _axes_extent, _combo,
+ _form, _group, _make_dspin, _scroll_panel, _wrap_label,
+)
+
+if TYPE_CHECKING:
+ from .main_window import SrasViewerWindow
+
+# Longest side of the coarsened preview the mask stack is built on. Same
+# reasoning as the old manual dialog: a full-resolution reprojection per angle
+# per edit is far more detail than an alignment can be judged by eye at.
+_MAX_PREVIEW_DIM = 1024
+
+# Alpha for the per-angle colour view (the active angle sits on top, brighter).
+_BASE_ALPHA = 0.42
+_ACTIVE_ALPHA = 0.75
+
+_PANEL_W = 372
+
+# Rotation this far from the file's stage angle is worth calling out: the stage
+# angles are usually good to within a degree, so more than that means either the
+# metadata or the fit is wrong.
+_DRIFT_WARN_DEG = 1.0
+
+# (label, sources for register_angle_to_reference) — "both" costs roughly double
+# but removes the failure mode where the one chosen source happens to be
+# uninformative for a single angle.
+_CORRELATE_SOURCES = (
+ ("Both, keep best", ("signal", "mask")),
+ ("Raw signal", ("signal",)),
+ ("Thresholded mask", ("mask",)),
+)
+
+# (label, seed_deg override or None for "the file's stage angle", signs to
+# sweep, sign used for the pre-rotation *preview*). The stage's rotational sense
+# relative to this module's math-positive convention is not knowable from the
+# file, which is why "both" exists and is the default; the explicit signs are
+# for a user who has established which way their own stage turns.
+_PREROTATE_MODES = (
+ ("Stage angle, both signs", None, (-1, 1), 1),
+ ("Stage angle, + delta only", None, (1,), 1),
+ ("Stage angle, − delta only", None, (-1,), -1),
+ ("No pre-rotation", 0.0, (1,), 0),
+)
+
+
+@dataclass
+class WizardState:
+ """Everything the three pages share."""
+ dc4_mv: dict[int, np.ndarray] = field(default_factory=dict)
+ masks_small: dict[int, np.ndarray] = field(default_factory=dict)
+ downsample: tuple[int, int] = (1, 1)
+ params: dict[int, ManualAngleParams] = field(default_factory=dict)
+ fits: dict[int, tuple[float, str]] = field(default_factory=dict)
+ ref_angle_idx: int = 0
+ threshold_mv: float = 0.0
+ result: compute.AlignmentResult | None = None # full, uncropped
+ counts: np.ndarray | None = None # coarse overlap counts
+ layers: dict[int, np.ndarray] = field(default_factory=dict)
+ preview_pitch_mm: tuple[float, float] = (1.0, 1.0)
+ preview_shape: tuple[int, int] = (1, 1)
+ geometry_generation: int = 0
+ crop: tuple[int, int, int, int] | None = None # canvas (row0,col0,nr,nc)
+ out_path: str = ""
+ exported_path: str = ""
+
+
+class AlignmentWizard(QWizard):
+ """Three-page alignment + export flow (Fusion → Alignment Wizard…).
+
+ Non-modal by design, like the dialog it replaces: an export can take minutes
+ on a real scan and the user should still be able to look at the data.
+
+ Background work goes through self.run_worker, which inherits the main
+ window's job registry, thread joining and shutdown handling. Two rules that
+ module's docstring establishes and every launch here follows: disable the
+ trigger *before* calling it (a re-entrant click must not be able to start a
+ second thread over the first), and never ignore its return value — False
+ means another job holds the key.
+ """
+
+ PAGE_CORRELATE, PAGE_ROI, PAGE_SAVE = 0, 1, 2
+
+ # cropped AlignmentResult, written path
+ alignment_ready = pyqtSignal(object, str)
+
+ def __init__(self, parent: "SrasViewerWindow", sras: SrasFile, *,
+ ref_angle_idx: int, dc_threshold_mv: float,
+ cached_dc4_mv: dict[int, np.ndarray]):
+ super().__init__(parent)
+ self._parent = parent
+ self._sras = sras
+ self._cached_dc4 = dict(cached_dc4_mv)
+ self._closing = False
+ # (result, crop, cropped result, plan) — see cropped_plan.
+ self._plan_cache: tuple | None = None
+ # Canvas origin the current st.layers were reprojected against; a
+ # change in it invalidates every layer. See rebuild_stack.
+ self._stack_origin_mm: tuple[float, float] | None = None
+
+ self.state = WizardState(
+ ref_angle_idx=ref_angle_idx, threshold_mv=dc_threshold_mv,
+ params={a: ManualAngleParams() for a in range(sras.n_angles)})
+
+ n = sras.n_angles
+ cmap = mpl.colormaps["tab10"] if n <= 10 else mpl.colormaps["tab20"]
+ self.angle_colors = {a: cmap(a % cmap.N)[:3] for a in range(n)}
+
+ self.setWindowTitle(f"Alignment Wizard — {sras.path.name}")
+ self.setWizardStyle(QWizard.WizardStyle.ModernStyle)
+ self.setOption(QWizard.WizardOption.HaveHelpButton, False)
+ self.setOption(QWizard.WizardOption.NoBackButtonOnStartPage, True)
+ # IndependentPages must stay OFF: it suppresses cleanupPage, which is
+ # how the ROI page discards a crop whose canvas is about to change.
+ self.setOption(QWizard.WizardOption.IndependentPages, False)
+ self.resize(1220, 820)
+
+ self.setPage(self.PAGE_CORRELATE, CorrelatePage(self))
+ self.setPage(self.PAGE_ROI, RoiPage(self))
+ self.setPage(self.PAGE_SAVE, SavePage(self))
+ self.setStartId(self.PAGE_CORRELATE)
+
+ # ------------------------------------------------------------------
+ # Shared helpers the pages use
+ # ------------------------------------------------------------------
+
+ @property
+ def sras(self) -> SrasFile:
+ return self._sras
+
+ def rebuild_result(self):
+ """Recompute the full AlignmentResult from the current parameters.
+
+ Closed-form matrix and bbox maths with no per-pixel work, so it is cheap
+ enough to call on every edit — which is what keeps the preview honest
+ about the canvas a rotation change implies.
+ """
+ st = self.state
+ st.result = build_manual_alignment(
+ self._sras, st.ref_angle_idx, st.threshold_mv, st.params)
+
+ def _reproject(self, angle_idx: int) -> np.ndarray:
+ st = self.state
+ p = st.params[angle_idx]
+ return compute.reproject_mask(
+ self._sras, angle_idx, st.ref_angle_idx, st.masks_small[angle_idx],
+ p.rotation_deg, p.shift_mm, st.preview_pitch_mm,
+ st.result.canvas_origin_mm, st.preview_shape,
+ src_downsample=st.downsample)
+
+ def rebuild_stack(self, only_angle: int | None = None):
+ """Reproject every angle's mask onto a coarsened view of the *final*
+ canvas and sum them into an overlap-count image.
+
+ The preview deliberately shares the final canvas's origin and uses a
+ pitch that is an integer multiple of it, rather than the padded,
+ unsnapped preview canvas the old manual dialog used. That is what lets
+ the crop page turn a rectangle drawn in millimetres into an exact
+ integer window of the real canvas, with no second coordinate frame to
+ reconcile.
+
+ *only_angle* says that angle is the only one whose parameters changed.
+ Reprojection is the one genuinely per-pixel operation on the nudge path
+ and it costs the same for every angle, so a full rebuild makes an arrow
+ keypress N times more expensive than it needs to be. Each angle's affine
+ depends on its own parameters plus the shared canvas, so if the canvas
+ came back identical the other layers provably did not move and only the
+ nudged one is redrawn; if the canvas did shift, every angle's affine
+ changed with it and the hint is ignored.
+ """
+ st = self.state
+ if st.result is None or not st.masks_small:
+ return
+ fy, fx = st.downsample
+ n_rows, n_cols = st.result.canvas_shape
+ pitch = (st.result.canvas_dx_mm * fx, st.result.canvas_dy_mm * fy)
+ shape = (max(1, -(-n_rows // fy)), max(1, -(-n_cols // fx)))
+
+ incremental = (only_angle is not None
+ and st.counts is not None
+ and len(st.layers) == self._sras.n_angles
+ and pitch == st.preview_pitch_mm
+ and shape == st.preview_shape
+ and st.result.canvas_origin_mm == self._stack_origin_mm)
+ st.preview_pitch_mm, st.preview_shape = pitch, shape
+ self._stack_origin_mm = st.result.canvas_origin_mm
+
+ if incremental:
+ before = st.layers[only_angle] > 0.5
+ layer = self._reproject(only_angle)
+ st.layers[only_angle] = layer
+ # Exact integer arithmetic on booleans, so this cannot drift from
+ # what a full rebuild would have produced.
+ st.counts += (layer > 0.5).astype(np.int16) - before.astype(np.int16)
+ return
+
+ st.layers = {}
+ counts = np.zeros(shape, dtype=np.int16)
+ for a in range(self._sras.n_angles):
+ layer = self._reproject(a)
+ st.layers[a] = layer
+ counts += (layer > 0.5).astype(np.int16)
+ st.counts = counts
+
+ def cropped_plan(self) -> tuple[compute.AlignmentResult | None,
+ export.ExportPlan | None]:
+ """The current crop's AlignmentResult and ExportPlan, or (None, None)
+ before a crop exists.
+
+ Both are pure functions of state.result and state.crop, so neither is
+ stored as its own piece of state. They are memoized instead, because
+ plan_export counts every output pixel of every angle (hundreds of
+ milliseconds on a full-size scan) and the ROI readout, that page's
+ validate step and the save page's summary all ask for the same crop.
+ """
+ st = self.state
+ if st.result is None or st.crop is None:
+ return None, None
+ cached = self._plan_cache
+ if cached is not None and cached[0] is st.result and cached[1] == st.crop:
+ return cached[2], cached[3]
+ cropped = compute.crop_alignment_result(st.result, *st.crop)
+ plan = export.plan_export(self._sras, cropped)
+ self._plan_cache = (st.result, st.crop, cropped, plan)
+ return cropped, plan
+
+ def preview_extent(self) -> list[float]:
+ st = self.state
+ x0, y0 = st.result.canvas_origin_mm
+ dx, dy = st.preview_pitch_mm
+ n_rows, n_cols = st.preview_shape
+ return _axes_extent(x0 + np.arange(n_cols) * dx,
+ y0 + np.arange(n_rows) * dy, dx, dy)
+
+ def mm_to_canvas(self, x_mm: float, y_mm: float) -> tuple[float, float]:
+ """(col, row) in *final* canvas pixels, fractional."""
+ st = self.state
+ x0, y0 = st.result.canvas_origin_mm
+ return ((x_mm - x0) / st.result.canvas_dx_mm,
+ (y_mm - y0) / st.result.canvas_dy_mm)
+
+ def canvas_to_mm(self, col: float, row: float) -> tuple[float, float]:
+ st = self.state
+ x0, y0 = st.result.canvas_origin_mm
+ return (x0 + col * st.result.canvas_dx_mm,
+ y0 + row * st.result.canvas_dy_mm)
+
+ def run_worker(self, key: str, worker, **kwargs) -> bool:
+ """Start a background job on the main window's registry. Returns False
+ if another job already holds *key*, which callers must not ignore.
+
+ The pages go through here rather than reaching into the parent window
+ themselves, so what the wizard actually needs from its parent — a job
+ runner — is stated in one place instead of at every launch site."""
+ return self._parent._run_worker(key, worker, **kwargs)
+
+ def job_running(self) -> bool:
+ return any(self._parent._job_running(k) for k in
+ (Jobs.ALIGN_MASKS, Jobs.ALIGN_CORRELATE, Jobs.ALIGN_EXPORT))
+
+ # ------------------------------------------------------------------
+ # Lifetime
+ # ------------------------------------------------------------------
+
+ def reject(self):
+ """Refuse to close while a job is in flight.
+
+ The running worker's signals are connected to bound methods of this
+ wizard and its pages; letting Qt delete them under a live thread is a
+ crash, not an inconvenience. Ask the worker to stop and let the job's
+ own completion close us.
+ """
+ if self.job_running():
+ for page_id in (self.PAGE_CORRELATE, self.PAGE_SAVE):
+ page = self.page(page_id)
+ if page is not None:
+ page.request_stop()
+ self._closing = True
+ return
+ super().reject()
+
+ def closeEvent(self, event):
+ """Window-manager close goes through the same deferral as Cancel."""
+ if self.job_running():
+ self.reject()
+ event.ignore()
+ return
+ super().closeEvent(event)
+
+ def accept(self):
+ cropped, _ = self.cropped_plan()
+ self.alignment_ready.emit(cropped, self.state.exported_path)
+ super().accept()
+
+ def maybe_close_after_job(self):
+ """Called by a page when its job finishes; completes a deferred close."""
+ if self._closing and not self.job_running():
+ self._closing = False
+ super().reject()
+
+
+# ---------------------------------------------------------------------------
+# Page 1 — cross-correlation
+# ---------------------------------------------------------------------------
+
+class CorrelatePage(QWizardPage):
+ """Register every angle against the reference, and show whether it worked.
+
+ The stack image is the point of the page. Numbers alone ("score 0.42") do
+ not tell you whether a five-angle fusion is usable; an overlap-count image
+ does, immediately.
+ """
+
+ def __init__(self, wizard: AlignmentWizard):
+ super().__init__(wizard)
+ self._wiz = wizard
+ self._sras = wizard.sras
+ self._busy = False
+ self._masks_ready = False
+ self._active_angle = 0
+ self._worker = None
+ self._done_count = 0
+ self._total = 0
+
+ self.setTitle("Step 1 — Cross-correlate the angles")
+ self.setSubTitle(
+ "Angles start pre-rotated by the stage angles stored in the scan. "
+ "Run the correlation, then check the stack: every pixel is coloured "
+ "by how many angles cover it, so a good alignment is one solid "
+ "plateau.")
+ self._build_ui()
+
+ # ---- construction -------------------------------------------------
+
+ def _build_ui(self):
+ root = QHBoxLayout(self)
+
+ self.canvas = AlignOverlayCanvas()
+ left = QWidget()
+ ll = QVBoxLayout(left)
+ ll.setContentsMargins(0, 0, 0, 0)
+ ll.setSpacing(4)
+ ll.addWidget(NavigationToolbar2QT(self.canvas, left))
+ ll.addWidget(self.canvas, stretch=1)
+ self.lbl_overlap = _wrap_label("", _CSS_INFO)
+ ll.addWidget(self.lbl_overlap)
+ self.lbl_warn = _wrap_label("", _CSS_WARN)
+ ll.addWidget(self.lbl_warn)
+ root.addWidget(left, stretch=1)
+
+ panel = QWidget()
+ pl = QVBoxLayout(panel)
+ pl.setContentsMargins(0, 0, 0, 0)
+ pl.setSpacing(8)
+ pl.addWidget(self._build_view_group())
+ pl.addWidget(self._build_reg_group())
+ pl.addWidget(self._build_run_group())
+ pl.addWidget(self._build_nudge_group())
+ pl.addWidget(self._build_table_group())
+ self.lbl_status = _wrap_label("", _CSS_MUTED)
+ pl.addWidget(self.lbl_status)
+ pl.addStretch()
+ root.addWidget(_scroll_panel(panel, _PANEL_W))
+
+ self._connect()
+
+ def _build_view_group(self) -> QWidget:
+ grp, lay = _group("View")
+ self.combo_view = _combo(["Overlap count", "Per-angle colours"])
+ lay.addWidget(self.combo_view)
+ return grp
+
+ def _build_reg_group(self) -> QWidget:
+ grp, lay = _group("Registration")
+ form = _form()
+
+ self.combo_ref = _combo(
+ f"Angle {a} ({self._sras.angles_deg[a]:.1f}°)"
+ for a in range(self._sras.n_angles))
+ self.combo_ref.setCurrentIndex(self._wiz.state.ref_angle_idx)
+ form.addRow("Reference:", self.combo_ref)
+
+ self.spin_threshold = _make_dspin(-500.0, 500.0, 3, suffix=" mV",
+ value=self._wiz.state.threshold_mv)
+ form.addRow("DC threshold:", self.spin_threshold)
+
+ self.combo_source = _combo(
+ (label, sources) for label, sources in _CORRELATE_SOURCES)
+ form.addRow("Correlate on:", self.combo_source)
+
+ self.combo_prerotate = _combo(
+ (label, (seed, signs, disp))
+ for label, seed, signs, disp in _PREROTATE_MODES)
+ form.addRow("Pre-rotate:", self.combo_prerotate)
+
+ self.chk_lock_rotation = QCheckBox("Lock rotation to the seed")
+ form.addRow("", self.chk_lock_rotation)
+
+ self.spin_search_deg = _make_dspin(0.0, 180.0, 1, suffix=" °",
+ value=6.0, step=1.0)
+ form.addRow("Rotation search (±):", self.spin_search_deg)
+
+ self.spin_coarse_step = _make_dspin(0.1, 30.0, 2, suffix=" °", value=2.0)
+ form.addRow("Coarse step:", self.spin_coarse_step)
+
+ self.combo_fine_dim = _combo((f"{dim} px", dim) for dim in (320, 640, 1024))
+ self.combo_fine_dim.setCurrentIndex(1)
+ form.addRow("Fine grid:", self.combo_fine_dim)
+
+ lay.addLayout(form)
+ lay.addWidget(_wrap_label(
+ "Pre-rotation only seeds the search — the rotation is still found "
+ "from image content, and both signs of the stage's reported angle "
+ "are tried unless you narrow it. Locking instead pins rotation to "
+ "the stage angle and searches translation only.", _CSS_HINT))
+ return grp
+
+ def _build_run_group(self) -> QWidget:
+ grp, lay = _group("Run")
+ self.btn_correlate = QPushButton("Run / Re-run Correlation")
+ lay.addWidget(self.btn_correlate)
+ self.progress = QProgressBar()
+ self.progress.setRange(0, 100)
+ self.progress.setValue(0)
+ lay.addWidget(self.progress)
+ self.btn_reset = QPushButton("Reset to pre-rotation only")
+ lay.addWidget(self.btn_reset)
+ return grp
+
+ def _build_nudge_group(self) -> QWidget:
+ self.grp_nudge, lay = _group("Manual Correction")
+ form = _form()
+ self.combo_active = _combo(
+ f"Angle {a} ({self._sras.angles_deg[a]:.1f}°)"
+ for a in range(self._sras.n_angles))
+ form.addRow("Active angle:", self.combo_active)
+
+ self.spin_rot = _make_dspin(-3600.0, 3600.0, 3, suffix=" °")
+ form.addRow("Rotation:", self.spin_rot)
+ self.spin_shift_x = _make_dspin(-1e5, 1e5, 4, suffix=" mm")
+ form.addRow("Shift X:", self.spin_shift_x)
+ self.spin_shift_y = _make_dspin(-1e5, 1e5, 4, suffix=" mm")
+ form.addRow("Shift Y:", self.spin_shift_y)
+
+ self.spin_step_translate = _make_dspin(0.0001, 1000.0, 4, suffix=" mm",
+ value=0.01)
+ form.addRow("Translate step:", self.spin_step_translate)
+ self.spin_step_rotate = _make_dspin(0.001, 90.0, 3, suffix=" °", value=0.1)
+ form.addRow("Rotate step:", self.spin_step_rotate)
+ self.spin_step_mult = _make_dspin(1.0, 1000.0, 1, value=10.0)
+ form.addRow("Coarse × (Shift):", self.spin_step_mult)
+ lay.addLayout(form)
+ lay.addWidget(_wrap_label(
+ "Arrow keys nudge translation; Q/E nudge rotation (CCW/CW). Hold "
+ "Shift for the coarse step. Click the image once to give it "
+ "keyboard focus. Switch to per-angle colours to see which layer "
+ "you are moving.", _CSS_HINT))
+ self.lbl_active_note = _wrap_label("", _CSS_WARN)
+ lay.addWidget(self.lbl_active_note)
+ return self.grp_nudge
+
+ def _build_table_group(self) -> QWidget:
+ grp, lay = _group("Fit per Angle")
+ self.table = QTableWidget(self._sras.n_angles, 5)
+ self.table.setHorizontalHeaderLabels(
+ ["Angle", "Rot °", "Δ stage °", "Score", "On"])
+ self.table.verticalHeader().setVisible(False)
+ self.table.setEditTriggers(QTableWidget.EditTrigger.NoEditTriggers)
+ self.table.setSelectionMode(QTableWidget.SelectionMode.NoSelection)
+ self.table.horizontalHeader().setSectionResizeMode(
+ QHeaderView.ResizeMode.ResizeToContents)
+ self.table.setMinimumHeight(150)
+ lay.addWidget(self.table)
+ return grp
+
+ def _connect(self):
+ self.combo_view.currentIndexChanged.connect(self._redraw)
+ self.combo_ref.currentIndexChanged.connect(self._on_ref_changed)
+ self.spin_threshold.editingFinished.connect(self._on_threshold_changed)
+ self.combo_prerotate.currentIndexChanged.connect(self._apply_prerotation)
+ self.chk_lock_rotation.toggled.connect(self._on_lock_toggled)
+ self.btn_correlate.clicked.connect(self._on_correlate)
+ self.btn_reset.clicked.connect(self._apply_prerotation)
+ self.combo_active.currentIndexChanged.connect(self._on_active_changed)
+ for spin in (self.spin_rot, self.spin_shift_x, self.spin_shift_y):
+ spin.editingFinished.connect(self._on_manual_edit)
+ self.canvas.nudge_translate.connect(self._on_nudge_translate)
+ self.canvas.nudge_rotate.connect(self._on_nudge_rotate)
+
+ # ---- QWizardPage contract ----------------------------------------
+
+ def initializePage(self):
+ if self._masks_ready or self._busy:
+ return
+ self._set_controls_enabled(False)
+ self._start_mask_prep()
+
+ def isComplete(self) -> bool:
+ """Next is gated on masks being ready and nothing running.
+
+ Returning False while busy also disables Finish, which is what stops a
+ Back/Next/Finish race from reaching the writer with a half-built result.
+ """
+ return self._masks_ready and not self._busy
+
+ def request_stop(self):
+ if self._worker is not None:
+ self._worker.stop()
+
+ # ---- mask preparation --------------------------------------------
+
+ def _start_mask_prep(self):
+ st = self._wiz.state
+ st.dc4_mv = dict(self._wiz._cached_dc4)
+ missing = [a for a in range(self._sras.n_angles) if a not in st.dc4_mv]
+ if not missing:
+ self._finish_mask_prep()
+ return
+ self._busy = True
+ self.completeChanged.emit()
+ self.lbl_status.setText(
+ f"Preparing masks: 0/{len(missing)} angle(s) needed…")
+ worker = Ch4MaskWorker(self._sras, missing)
+ started = self._wiz.run_worker(
+ Jobs.ALIGN_MASKS, worker,
+ connect=(("angle_done", self._on_mask_done),
+ ("error", lambda m: self.lbl_status.setText(
+ f"Mask preparation failed: {m}"))),
+ on_done=self._finish_mask_prep)
+ if not started:
+ self._busy = False
+ self.completeChanged.emit()
+ self.lbl_status.setText(
+ "Another alignment step is still running — close and reopen.")
+ return
+ self._worker = worker
+
+ def _on_mask_done(self, angle_idx: int, dc4_mv: np.ndarray):
+ st = self._wiz.state
+ st.dc4_mv[angle_idx] = dc4_mv
+ self.lbl_status.setText(
+ f"Preparing masks: {len(st.dc4_mv)}/{self._sras.n_angles} ready…")
+
+ def _finish_mask_prep(self):
+ st = self._wiz.state
+ self._worker = None
+ self._busy = False
+ if len(st.dc4_mv) < self._sras.n_angles:
+ self.completeChanged.emit()
+ self.lbl_status.setText(
+ "Mask preparation did not finish for every angle.")
+ self._wiz.maybe_close_after_job()
+ return
+
+ # Rows and columns get independent factors. A real scan is ~7500 frames
+ # wide but only ~750 rows tall, so one shared factor sized for the
+ # frames would throw away 10x more row detail than the preview needs.
+ max_rows = max(img.shape[0] for img in st.dc4_mv.values())
+ max_cols = max(img.shape[1] for img in st.dc4_mv.values())
+ st.downsample = (max(1, -(-max_rows // _MAX_PREVIEW_DIM)),
+ max(1, -(-max_cols // _MAX_PREVIEW_DIM)))
+
+ self._recompute_masks()
+ self._masks_ready = True
+ self._set_controls_enabled(True)
+ # Sets st.params for every angle, so it is the single source of the
+ # starting parameters — nothing needs to seed them beforehand.
+ self._apply_prerotation()
+ self.lbl_status.setText("Ready.")
+ self.completeChanged.emit()
+ self._wiz.maybe_close_after_job()
+
+ def _recompute_masks(self):
+ """Threshold + downsample each angle's in-memory CH4 image. Cheap (a
+ compare and a block-mean), so a threshold change re-runs it in full
+ rather than re-fetching anything."""
+ st = self._wiz.state
+ fy, fx = st.downsample
+ st.masks_small = {
+ a: compute.block_mean_2d((img >= st.threshold_mv).astype(np.float32),
+ fy, fx)
+ for a, img in st.dc4_mv.items()
+ }
+
+ # ---- parameter changes -------------------------------------------
+
+ def _on_ref_changed(self, idx: int):
+ st = self._wiz.state
+ st.ref_angle_idx = idx
+ st.fits = {}
+ self._apply_prerotation()
+
+ def _on_threshold_changed(self):
+ st = self._wiz.state
+ value = self.spin_threshold.value()
+ if value == st.threshold_mv:
+ return
+ st.threshold_mv = value
+ if self._masks_ready:
+ self._recompute_masks()
+ self._refresh()
+
+ def _on_lock_toggled(self, locked: bool):
+ self.spin_search_deg.setEnabled(not locked)
+ self.spin_coarse_step.setEnabled(not locked)
+
+ def _apply_prerotation(self):
+ """Seed every non-reference angle's rotation from the stage angles in
+ the file, and redraw — before any correlation has run.
+
+ This is the cheapest useful thing the page can show: the stage angles
+ are usually within a degree or two of the truth, so the stack already
+ looks close to right, and how close is a fair first read on whether the
+ scan's own metadata can be trusted.
+ """
+ if not self._masks_ready:
+ return
+ st = self._wiz.state
+ _seed, _signs, disp_sign = self.combo_prerotate.currentData()
+ for a in range(self._sras.n_angles):
+ rot = 0.0 if a == st.ref_angle_idx else (
+ disp_sign * compute.nominal_delta_deg(self._sras, a,
+ st.ref_angle_idx))
+ st.params[a] = ManualAngleParams(rot, (0.0, 0.0))
+ st.fits = {}
+ self._refresh()
+
+ def _stage_drift_deg(self, angle_idx: int) -> float:
+ """How far this angle's rotation has ended up from the stage angle the
+ file reports. Both signs are allowed: the stage's rotational sense
+ relative to this module's convention is not knowable from the file,
+ so the closer of the two is the honest comparison."""
+ st = self._wiz.state
+ nominal = compute.nominal_delta_deg(self._sras, angle_idx,
+ st.ref_angle_idx)
+ got = st.params[angle_idx].rotation_deg
+ return min(abs(got - nominal), abs(got + nominal))
+
+ def _refresh(self, only_angle: int | None = None):
+ """Rebuild result, stack and every readout from the current params.
+
+ *only_angle* is passed through to rebuild_stack, which uses it to skip
+ reprojecting angles that cannot have moved; callers that changed more
+ than one angle's parameters must leave it None."""
+ st = self._wiz.state
+ self._wiz.rebuild_result()
+ st.geometry_generation += 1
+ self._wiz.rebuild_stack(only_angle)
+ self._sync_spins()
+ self._update_table()
+ self._redraw()
+
+ # ---- correlation --------------------------------------------------
+
+ def _reg_kwargs(self) -> dict:
+ """Every argument register_angle_to_reference takes from this page, in
+ one dict — so "lock rotation" can express all of what it means in one
+ place instead of half here and half at the call site."""
+ seed, signs, _disp = self.combo_prerotate.currentData()
+ kwargs = {
+ "sources": self.combo_source.currentData(),
+ "dc_threshold_mv": self._wiz.state.threshold_mv,
+ "search_deg": self.spin_search_deg.value(),
+ "seed_deg": seed,
+ "seed_signs": signs,
+ "coarse_step_deg": self.spin_coarse_step.value(),
+ "fine_dim": self.combo_fine_dim.currentData(),
+ }
+ if self.chk_lock_rotation.isChecked():
+ # Exactly one candidate, no hill-climb: rotation is the seed and
+ # only the translation is searched.
+ kwargs["refine"] = False
+ kwargs["search_deg"] = 0.0
+ if len(signs) > 1:
+ kwargs["seed_signs"] = (1,)
+ return kwargs
+
+ def _on_correlate(self):
+ if not self._masks_ready or self._busy:
+ return
+ st = self._wiz.state
+ angles = [a for a in range(self._sras.n_angles) if a != st.ref_angle_idx]
+ if not angles:
+ self.lbl_status.setText("Only one angle — nothing to correlate.")
+ return
+
+ worker = CrossCorrelateWorker(
+ self._sras, st.ref_angle_idx, angles, st.dc4_mv,
+ reg_kwargs=self._reg_kwargs())
+
+ # Claim busy and disable the trigger *before* _run_worker, never after:
+ # anything that pumps the event loop in between could deliver a second
+ # click that starts a thread the first assignment then drops.
+ self._busy = True
+ self._done_count, self._total = 0, len(angles)
+ st.fits = {}
+ self.completeChanged.emit()
+ self._set_controls_enabled(False)
+ self.progress.setValue(0)
+ self.lbl_status.setText(f"Cross-correlating: 0/{self._total} angle(s)…")
+
+ started = self._wiz.run_worker(
+ Jobs.ALIGN_CORRELATE, worker,
+ connect=(("angle_done", self._on_angle_done),
+ ("error", lambda m: self.lbl_status.setText(
+ f"Cross-correlation failed: {m}"))),
+ on_done=self._finish_correlate)
+ if not started:
+ self._busy = False
+ self.completeChanged.emit()
+ self._set_controls_enabled(True)
+ self.lbl_status.setText(
+ "Another alignment step is still running — try again shortly.")
+ return
+ self._worker = worker
+
+ def _on_angle_done(self, angle_idx: int, rot: float, sx: float, sy: float,
+ score: float, source: str):
+ st = self._wiz.state
+ st.params[angle_idx] = ManualAngleParams(rot, (sx, sy))
+ st.fits[angle_idx] = (score, source)
+ self._done_count += 1
+ self.progress.setValue(int(self._done_count / max(1, self._total) * 100))
+ self.lbl_status.setText(
+ f"Cross-correlating: {self._done_count}/{self._total} angle(s)…")
+
+ def _finish_correlate(self):
+ self._worker = None
+ self._busy = False
+ self.progress.setValue(100)
+ self._refresh()
+ self._set_controls_enabled(True)
+ self.completeChanged.emit()
+ self.lbl_status.setText(
+ f"Correlated {self._done_count} angle(s) against angle "
+ f"{self._wiz.state.ref_angle_idx}. " + self._fit_summary())
+ self._wiz.maybe_close_after_job()
+
+ def _fit_summary(self) -> str:
+ """Worst fit and any angle whose rotation disagrees with its stage angle.
+
+ Surfaced rather than buried because a single bad acquisition (stage
+ glitch, laser dropout) registers poorly and would otherwise be fused in
+ silently — knowing *which* angle is what makes dropping it with
+ sras_edit_scans.py actionable.
+ """
+ st = self._wiz.state
+ if not st.fits:
+ return ""
+ rows = sorted(st.fits.items(), key=lambda kv: kv[1][0])
+ worst_a, (worst_score, worst_src) = rows[0]
+ parts = [f"Worst fit: angle {worst_a} (score {worst_score:.3f}, "
+ f"{worst_src or 'n/a'})."]
+ failed = [str(a) for a, (score, src) in rows if score < 0 or src == "none"]
+ if failed:
+ parts.append(
+ "Angle(s) " + ", ".join(failed) + " did not register at all and "
+ "are being treated as unrotated — lower the DC threshold, try "
+ "Raw signal, nudge them by hand, or drop them with "
+ "sras_edit_scans.py.")
+ drifted = [f"{a} ({dev:.2f}°)" for a, _ in rows
+ if (dev := self._stage_drift_deg(a)) > _DRIFT_WARN_DEG]
+ if drifted:
+ parts.append(f"Rotation differs from the stage angle by "
+ f">{_DRIFT_WARN_DEG:g}° for angle(s) "
+ + ", ".join(drifted) + ".")
+ return " ".join(parts)
+
+ # ---- manual nudging ----------------------------------------------
+
+ def _on_active_changed(self, idx: int):
+ self._active_angle = idx
+ is_ref = idx == self._wiz.state.ref_angle_idx
+ for spin in (self.spin_rot, self.spin_shift_x, self.spin_shift_y):
+ spin.setEnabled(self._masks_ready and not is_ref)
+ self.lbl_active_note.setText(
+ "Reference angle — defines the shared origin, not adjustable."
+ if is_ref else "")
+ self._sync_spins()
+ self._redraw()
+
+ def _sync_spins(self):
+ p = self._wiz.state.params.get(self._active_angle, ManualAngleParams())
+ for spin, val in ((self.spin_rot, p.rotation_deg),
+ (self.spin_shift_x, p.shift_mm[0]),
+ (self.spin_shift_y, p.shift_mm[1])):
+ with QSignalBlocker(spin):
+ spin.setValue(val)
+
+ def _on_manual_edit(self):
+ if not self._masks_ready or self._active_angle == self._wiz.state.ref_angle_idx:
+ return
+ self._wiz.state.params[self._active_angle] = ManualAngleParams(
+ self.spin_rot.value(),
+ (self.spin_shift_x.value(), self.spin_shift_y.value()))
+ self._refresh(self._active_angle)
+
+ def _on_nudge_translate(self, dir_x: int, dir_y: int, coarse: bool):
+ if not self._masks_ready or self._active_angle == self._wiz.state.ref_angle_idx:
+ return
+ step = self.spin_step_translate.value()
+ if coarse:
+ step *= self.spin_step_mult.value()
+ p = self._wiz.state.params[self._active_angle]
+ self._wiz.state.params[self._active_angle] = ManualAngleParams(
+ p.rotation_deg,
+ (p.shift_mm[0] + dir_x * step, p.shift_mm[1] + dir_y * step))
+ self._refresh(self._active_angle)
+
+ def _on_nudge_rotate(self, direction: int, coarse: bool):
+ if not self._masks_ready or self._active_angle == self._wiz.state.ref_angle_idx:
+ return
+ step = self.spin_step_rotate.value()
+ if coarse:
+ step *= self.spin_step_mult.value()
+ p = self._wiz.state.params[self._active_angle]
+ self._wiz.state.params[self._active_angle] = ManualAngleParams(
+ p.rotation_deg + direction * step, p.shift_mm)
+ self._refresh(self._active_angle)
+
+ # ---- drawing ------------------------------------------------------
+
+ def _set_controls_enabled(self, enabled: bool):
+ for w in (self.combo_ref, self.spin_threshold, self.combo_source,
+ self.combo_prerotate, self.chk_lock_rotation,
+ self.spin_search_deg, self.spin_coarse_step,
+ self.combo_fine_dim, self.btn_correlate, self.btn_reset,
+ self.grp_nudge):
+ w.setEnabled(enabled)
+ if enabled:
+ self._on_lock_toggled(self.chk_lock_rotation.isChecked())
+ self._on_active_changed(self._active_angle)
+
+ def _redraw(self):
+ st = self._wiz.state
+ if st.counts is None or st.result is None:
+ return
+ extent = self._wiz.preview_extent()
+ if self.combo_view.currentIndex() == 0:
+ self.canvas.show_counts(
+ st.counts, self._sras.n_angles, extent,
+ f"Mask stack — {self._sras.n_angles} angles, "
+ f"ref angle {st.ref_angle_idx}")
+ else:
+ self.canvas.show_overlay(
+ self._rgba(), extent,
+ f"Angle {self._active_angle} active "
+ f"({self._sras.angles_deg[self._active_angle]:.1f}°)")
+ self._update_overlap_text()
+
+ def _rgba(self) -> np.ndarray:
+ """Alpha-composite each angle's mask in its own colour, active on top."""
+ st = self._wiz.state
+ n_rows, n_cols = st.preview_shape
+ rgba = np.zeros((n_rows, n_cols, 4), dtype=np.float32)
+ order = sorted(range(self._sras.n_angles),
+ key=lambda a: a == self._active_angle)
+ for a in order:
+ layer = st.layers.get(a)
+ if layer is None:
+ continue
+ alpha = _ACTIVE_ALPHA if a == self._active_angle else _BASE_ALPHA
+ color = self._wiz.angle_colors[a]
+ fg = layer * alpha
+ for c in range(3):
+ rgba[..., c] = color[c] * fg + rgba[..., c] * rgba[..., 3] * (1 - fg)
+ rgba[..., 3] = fg + rgba[..., 3] * (1 - fg)
+ return rgba
+
+ def _update_overlap_text(self):
+ st = self._wiz.state
+ n = self._sras.n_angles
+ stats = compute.overlap_stats(st.counts, n)
+ px = st.preview_pitch_mm
+ area = abs(px[0] * px[1])
+ self.lbl_overlap.setText(
+ f"Covered by any angle: {stats['union_px']:,} px "
+ f"({stats['union_px'] * area:.2f} mm²) · "
+ f"by all {n}: {stats['full_px']:,} px "
+ f"({stats['full_frac'] * 100:.1f}% of that) · "
+ f"mean overlap {stats['mean_count']:.2f} angles")
+ if stats["empty"]:
+ self.lbl_warn.setText(
+ "No angle covers any pixel — check the DC threshold.")
+ elif stats["max_count"] < n:
+ self.lbl_warn.setText(
+ f"No pixel is covered by all {n} angles (best is "
+ f"{stats['max_count']}). Either the alignment is wrong or these "
+ f"scans genuinely do not all overlap. You can still export.")
+ else:
+ self.lbl_warn.setText("")
+
+ def _update_table(self):
+ st = self._wiz.state
+ for a in range(self._sras.n_angles):
+ score, source = st.fits.get(a, (float("nan"), ""))
+ got = st.params[a].rotation_deg
+ dev = self._stage_drift_deg(a)
+ is_ref = a == st.ref_angle_idx
+ cells = [
+ f"{a}" + (" [ref]" if is_ref else ""),
+ f"{got:.3f}",
+ "—" if is_ref else f"{dev:.2f}",
+ "—" if is_ref or np.isnan(score) else f"{score:.3f}",
+ "ref" if is_ref else (source or "—"),
+ ]
+ bad = (not is_ref) and (score < 0 or source == "none")
+ weak = (not is_ref) and (not np.isnan(score)) and 0 <= score < 0.3
+ for col, text in enumerate(cells):
+ item = QTableWidgetItem(text)
+ item.setTextAlignment(Qt.AlignmentFlag.AlignCenter)
+ if bad:
+ item.setForeground(Qt.GlobalColor.red)
+ elif weak or (not is_ref and dev > 1.0):
+ item.setForeground(Qt.GlobalColor.darkYellow)
+ self.table.setItem(a, col, item)
+
+
+# ---------------------------------------------------------------------------
+# Page 2 — ROI crop
+# ---------------------------------------------------------------------------
+
+class RoiPage(QWizardPage):
+ """Pick the rectangle of the shared canvas to keep.
+
+ Axis-aligned, because v6 geometry is (x_start, x_delta, n_frames, n_rows)
+ plus a Y row table — a rectangle on the canvas grid is the only crop the
+ file can express, and drawing something else would have to be squared off
+ without saying so.
+ """
+
+ def __init__(self, wizard: AlignmentWizard):
+ super().__init__(wizard)
+ self._wiz = wizard
+ self._sras = wizard.sras
+ self._syncing = False
+ self._generation = -1
+
+ self.setTitle("Step 2 — Choose the region to keep")
+ self.setSubTitle(
+ "Drag a rectangle on the stack, or type canvas pixels directly. "
+ "Cropping to the region the angles actually share is usually a big "
+ "size saving over the full canvas.")
+ self._build_ui()
+
+ def _build_ui(self):
+ root = QHBoxLayout(self)
+
+ self.canvas = ImageCanvas(rect_only=True)
+ left = QWidget()
+ ll = QVBoxLayout(left)
+ ll.setContentsMargins(0, 0, 0, 0)
+ ll.setSpacing(4)
+ ll.addWidget(NavigationToolbar2QT(self.canvas, left))
+ ll.addWidget(self.canvas, stretch=1)
+ root.addWidget(left, stretch=1)
+
+ panel = QWidget()
+ pl = QVBoxLayout(panel)
+ pl.setContentsMargins(0, 0, 0, 0)
+ pl.setSpacing(8)
+
+ grp_fit, fl = _group("Preset")
+ self.btn_fit_overlap = QPushButton("Fit to full overlap")
+ self.btn_fit_union = QPushButton("Fit to any coverage")
+ self.btn_whole = QPushButton("Whole canvas (no crop)")
+ for b in (self.btn_fit_overlap, self.btn_fit_union, self.btn_whole):
+ fl.addWidget(b)
+ fl.addWidget(_wrap_label(
+ "“Full overlap” finds the largest rectangle lying entirely inside "
+ "the region every angle covers — not its bounding box, which would "
+ "include corners no angle reaches.", _CSS_HINT))
+ pl.addWidget(grp_fit)
+
+ grp_rect, rl = _group("Crop (canvas pixels)")
+ form = _form()
+ self.spin_col0 = QSpinBox()
+ self.spin_row0 = QSpinBox()
+ self.spin_cols = QSpinBox()
+ self.spin_rows = QSpinBox()
+ for spin, label in ((self.spin_col0, "First column:"),
+ (self.spin_row0, "First row:"),
+ (self.spin_cols, "Columns:"),
+ (self.spin_rows, "Rows:")):
+ spin.setRange(0, 1)
+ spin.setMinimumWidth(96)
+ # Each committed edit recounts coverage over the whole crop, so
+ # only commit on Enter/focus-out rather than on every digit typed
+ # into a four-figure column count.
+ spin.setKeyboardTracking(False)
+ form.addRow(label, spin)
+ rl.addLayout(form)
+ self.btn_draw = QPushButton("Draw a new rectangle")
+ rl.addWidget(self.btn_draw)
+ pl.addWidget(grp_rect)
+
+ grp_info, il = _group("Result")
+ self.lbl_extent = _wrap_label("", _CSS_INFO)
+ self.lbl_size = _wrap_label("", _CSS_INFO)
+ self.lbl_coverage = _wrap_label("", _CSS_MUTED)
+ self.lbl_warn = _wrap_label("", _CSS_WARN)
+ for w in (self.lbl_extent, self.lbl_size, self.lbl_coverage, self.lbl_warn):
+ il.addWidget(w)
+ pl.addWidget(grp_info)
+
+ pl.addStretch()
+ root.addWidget(_scroll_panel(panel, _PANEL_W))
+
+ self.btn_fit_overlap.clicked.connect(self._on_fit_overlap)
+ self.btn_fit_union.clicked.connect(self._on_fit_union)
+ self.btn_whole.clicked.connect(self._on_whole)
+ self.btn_draw.clicked.connect(self.canvas.start_drawing)
+ self.canvas.roi_changed.connect(self._on_roi_changed)
+ for spin in (self.spin_col0, self.spin_row0, self.spin_cols, self.spin_rows):
+ spin.valueChanged.connect(self._on_spin_changed)
+
+ # ---- QWizardPage contract ----------------------------------------
+
+ def initializePage(self):
+ st = self._wiz.state
+ # A crop is indexed in canvas pixels, so it is meaningless against a
+ # canvas built from different rotations. Drop a stale one rather than
+ # silently reinterpreting its indices.
+ if st.crop is not None and self._generation != st.geometry_generation:
+ st.crop = None
+ self._generation = st.geometry_generation
+
+ n_rows, n_cols = st.result.canvas_shape
+ # Blocked: setRange clamps the current value into the new range and
+ # emits valueChanged, which would run _on_spin_changed and commit a
+ # crop built from three not-yet-set spin boxes — landing on 1x1 and
+ # making the crop look already-chosen, so the preset below is skipped.
+ for spin, lo, hi in ((self.spin_col0, 0, n_cols - 1),
+ (self.spin_row0, 0, n_rows - 1),
+ (self.spin_cols, 1, n_cols),
+ (self.spin_rows, 1, n_rows)):
+ with QSignalBlocker(spin):
+ spin.setRange(lo, max(lo, hi))
+ self._draw_counts()
+ # largest_rect_at_least returns None iff no pixel qualifies, so the
+ # .any() answers the enable question without running its O(rows*cols)
+ # Python sweep on the GUI thread just to throw the rectangle away.
+ self.btn_fit_overlap.setEnabled(
+ bool((st.counts >= self._sras.n_angles).any()))
+ if st.crop is None:
+ self._on_fit_union()
+ else:
+ self._set_crop(*st.crop)
+
+ def isComplete(self) -> bool:
+ # _set_crop is the only writer and clamps both extents to >= 1.
+ return self._wiz.state.crop is not None
+
+ def validatePage(self) -> bool:
+ _, plan = self._wiz.cropped_plan()
+ empty = plan.empty_angles()
+ if len(empty) == self._sras.n_angles:
+ QMessageBox.warning(
+ self, "Empty crop",
+ "This rectangle contains no data from any angle. Move or grow "
+ "it before continuing.")
+ return False
+ if empty:
+ answer = QMessageBox.question(
+ self, "Some angles are empty",
+ f"Angle(s) {', '.join(map(str, empty))} have no data inside "
+ f"this crop and would be written as padding.\n\nContinue anyway?")
+ if answer != QMessageBox.StandardButton.Yes:
+ return False
+ return True
+
+ def cleanupPage(self):
+ """Going back means the canvas geometry may change under this crop."""
+ self._wiz.state.crop = None
+
+ # ---- drawing / presets --------------------------------------------
+
+ def _draw_counts(self):
+ st = self._wiz.state
+ n = self._sras.n_angles
+ # Same colormap the previous page used, so a count keeps its colour
+ # across the two pages that show this image on different canvas classes.
+ cmap, norm, ticks = count_colormap(n)
+ self.canvas.show_image(
+ st.counts, self._wiz.preview_extent(), cmap, 0.0, 0.0,
+ "X (mm)", "Y (mm)",
+ f"Overlap count — drag the crop rectangle ({n} angles)",
+ colorbar_label="angles overlapping", cb_ticks=ticks, norm=norm)
+
+ def _coarse_rect_to_canvas(self, rect, *, inset_blocks: int
+ ) -> tuple[int, int, int, int]:
+ st = self._wiz.state
+ return compute.coarse_rect_to_canvas(rect, st.downsample,
+ st.result.canvas_shape,
+ inset_blocks=inset_blocks)
+
+ def _on_fit_overlap(self):
+ st = self._wiz.state
+ rect = compute.largest_rect_at_least(st.counts, self._sras.n_angles)
+ if rect is None:
+ return
+ self._set_crop(*self._coarse_rect_to_canvas(rect, inset_blocks=1))
+
+ def _on_fit_union(self):
+ st = self._wiz.state
+ rows, cols = np.nonzero(st.counts > 0)
+ if rows.size == 0:
+ self._on_whole()
+ return
+ r0, c0 = int(rows.min()), int(cols.min())
+ rect = (r0, c0, int(rows.max()) - r0 + 1, int(cols.max()) - c0 + 1)
+ self._set_crop(*self._coarse_rect_to_canvas(rect, inset_blocks=0))
+
+ def _on_whole(self):
+ n_rows, n_cols = self._wiz.state.result.canvas_shape
+ self._set_crop(0, 0, n_rows, n_cols)
+
+ # ---- two-way sync -------------------------------------------------
+
+ def _set_crop(self, row0: int, col0: int, n_rows: int, n_cols: int):
+ st = self._wiz.state
+ cr, cc = st.result.canvas_shape
+ row0 = int(np.clip(row0, 0, cr - 1))
+ col0 = int(np.clip(col0, 0, cc - 1))
+ n_rows = int(np.clip(n_rows, 1, cr - row0))
+ n_cols = int(np.clip(n_cols, 1, cc - col0))
+ st.crop = (row0, col0, n_rows, n_cols)
+
+ self._syncing = True
+ try:
+ for spin, val in ((self.spin_row0, row0), (self.spin_col0, col0),
+ (self.spin_rows, n_rows), (self.spin_cols, n_cols)):
+ with QSignalBlocker(spin):
+ spin.setValue(val)
+ x0, y0 = self._wiz.canvas_to_mm(col0 - 0.5, row0 - 0.5)
+ x1, y1 = self._wiz.canvas_to_mm(col0 + n_cols - 0.5,
+ row0 + n_rows - 0.5)
+ self.canvas.set_roi(RoiQuad.from_bbox(min(x0, x1), min(y0, y1),
+ max(x0, x1), max(y0, y1)))
+ finally:
+ self._syncing = False
+ self._update_readout()
+ self.completeChanged.emit()
+
+ def _on_spin_changed(self):
+ if self._syncing:
+ return
+ self._set_crop(self.spin_row0.value(), self.spin_col0.value(),
+ self.spin_rows.value(), self.spin_cols.value())
+
+ def _on_roi_changed(self):
+ # set_roi itself emits roi_changed, so without the guard _set_crop
+ # would re-enter through its own canvas update.
+ if self._syncing:
+ return
+ roi = self.canvas.get_roi()
+ if roi is None:
+ return
+ pts = roi.corners()
+ (x0, y0), (x1, y1) = pts.min(axis=0), pts.max(axis=0)
+ c_a, r_a = self._wiz.mm_to_canvas(x0, y0)
+ c_b, r_b = self._wiz.mm_to_canvas(x1, y1)
+ col0, col1 = sorted((c_a, c_b))
+ row0, row1 = sorted((r_a, r_b))
+ # Both edges snap to the nearest pixel boundary with the same rule, so
+ # the numeric boxes -> rectangle -> numeric boxes round trip is exact.
+ # _SNAP_TOL absorbs the float noise of the mm round trip: a boundary
+ # that should land on 11.0 arrives as 11.000000000000002, and a bare
+ # ceil() turns that into 12 — one spurious column per edit.
+ r0, r1 = _snap_edge(row0), _snap_edge(row1)
+ c0, c1 = _snap_edge(col0), _snap_edge(col1)
+ self._set_crop(r0, c0, max(1, r1 - r0), max(1, c1 - c0))
+
+ def _update_readout(self):
+ st = self._wiz.state
+ row0, col0, n_rows, n_cols = st.crop
+ x0, y0 = self._wiz.canvas_to_mm(col0, row0)
+ x1, y1 = self._wiz.canvas_to_mm(col0 + n_cols - 1, row0 + n_rows - 1)
+ self.lbl_extent.setText(
+ f"Output: {n_cols:,} × {n_rows:,} px "
+ f"X {min(x0, x1):.3f} … {max(x0, x1):.3f} mm "
+ f"Y {min(y0, y1):.3f} … {max(y0, y1):.3f} mm")
+
+ _, plan = self._wiz.cropped_plan()
+ self.lbl_size.setText(f"Estimated file size: {_humanize(plan.total_bytes)}"
+ f" ({self._sras.n_angles} angles)")
+ self.lbl_coverage.setText("Real data per angle: " + ", ".join(
+ f"{a}: {plan.coverage_frac(a) * 100:.0f}%"
+ for a in range(self._sras.n_angles)))
+ empty = plan.empty_angles()
+ self.lbl_warn.setText(
+ f"Angle(s) {', '.join(map(str, empty))} have no data here and would "
+ f"be written as padding." if empty else "")
+
+
+# ---------------------------------------------------------------------------
+# Page 3 — save
+# ---------------------------------------------------------------------------
+
+class SavePage(QWizardPage):
+ """Choose a destination and write the aligned, cropped scan.
+
+ The write runs on a worker and is started by a button rather than from
+ validatePage, which must not block the GUI thread for what can be minutes of
+ I/O. Finish only becomes available once a file has actually been written, so
+ the wizard cannot be completed on a failed export.
+ """
+
+ def __init__(self, wizard: AlignmentWizard):
+ super().__init__(wizard)
+ self._wiz = wizard
+ self._sras = wizard.sras
+ self._busy = False
+ self._worker = None
+
+ self.setTitle("Step 3 — Save the aligned scan")
+ self.setSubTitle(
+ "Writes a new v6 .sras in which every angle shares the cropped "
+ "grid, so it opens already aligned. The original is not modified.")
+ self._build_ui()
+
+ def _build_ui(self):
+ root = QVBoxLayout(self)
+
+ row = QHBoxLayout()
+ row.addWidget(QLabel("Save to:"))
+ self.edit_path = QLineEdit()
+ self.edit_path.setReadOnly(True)
+ row.addWidget(self.edit_path, stretch=1)
+ self.btn_browse = QPushButton("Browse…")
+ row.addWidget(self.btn_browse)
+ root.addLayout(row)
+
+ grp, gl = _group("What will be written")
+ self.lbl_summary = _wrap_label("", _CSS_INFO)
+ gl.addWidget(self.lbl_summary)
+ self.lbl_notes = _wrap_label("", _CSS_WARN)
+ gl.addWidget(self.lbl_notes)
+ root.addWidget(grp)
+
+ run = QHBoxLayout()
+ self.btn_export = QPushButton("Write .sras")
+ run.addWidget(self.btn_export)
+ self.btn_cancel = QPushButton("Cancel write")
+ self.btn_cancel.setEnabled(False)
+ run.addWidget(self.btn_cancel)
+ run.addStretch()
+ root.addLayout(run)
+
+ self.progress = QProgressBar()
+ self.progress.setRange(0, 100)
+ root.addWidget(self.progress)
+
+ self.lbl_status = _wrap_label("", _CSS_MUTED)
+ root.addWidget(self.lbl_status)
+ root.addStretch()
+
+ self.btn_browse.clicked.connect(self._on_browse)
+ self.btn_export.clicked.connect(self._on_export)
+ self.btn_cancel.clicked.connect(self.request_stop)
+
+ # ---- QWizardPage contract ----------------------------------------
+
+ def initializePage(self):
+ st = self._wiz.state
+ if not st.out_path:
+ src = Path(self._sras.path)
+ st.out_path = str(src.with_name(f"{src.stem}_aligned.sras"))
+ self.edit_path.setText(st.out_path)
+ st.exported_path = ""
+ self.progress.setValue(0)
+ self.lbl_status.setText("")
+ self._update_summary()
+ self._wiz.setButtonText(QWizard.WizardButton.FinishButton, "Done")
+ self.completeChanged.emit()
+
+ def isComplete(self) -> bool:
+ return bool(self._wiz.state.exported_path) and not self._busy
+
+ def request_stop(self):
+ if self._worker is not None:
+ self._worker.stop()
+ self.lbl_status.setText("Cancelling…")
+
+ # ---- summary ------------------------------------------------------
+
+ def _update_summary(self):
+ cropped, plan = self._wiz.cropped_plan()
+ x0, y0 = cropped.canvas_origin_mm
+ self.lbl_summary.setText(
+ f"Format: .sras v6, no precomputed cache (the viewer recomputes "
+ f"DC/FFT on first open).\n"
+ f"Angles: {plan.n_angles}, all sharing one grid of "
+ f"{plan.n_frames:,} × {plan.n_rows:,} px.\n"
+ f"Origin: X {x0:.4f} mm, Y {y0:.4f} mm · "
+ f"pitch {cropped.canvas_dx_mm * 1000:.2f} × "
+ f"{abs(cropped.canvas_dy_mm) * 1000:.2f} µm.\n"
+ f"Stage angles, calibration preambles and the background waveform "
+ f"are carried over unchanged.\n"
+ f"Size: {_humanize(plan.bytes_per_angle)} per angle, "
+ f"{_humanize(plan.total_bytes)} total.\n"
+ f"Real data per angle: " + ", ".join(
+ f"{a}: {plan.coverage_frac(a) * 100:.0f}%"
+ for a in range(plan.n_angles)))
+ self.lbl_notes.setText("\n".join(plan.warnings))
+
+ # ---- export -------------------------------------------------------
+
+ def _on_browse(self):
+ st = self._wiz.state
+ path, _ = QFileDialog.getSaveFileName(
+ self, "Save aligned .sras", st.out_path,
+ "SRAS scans (*.sras);;All files (*)")
+ if not path:
+ return
+ if not path.lower().endswith(".sras"):
+ path += ".sras"
+ st.out_path = path
+ self.edit_path.setText(path)
+ st.exported_path = ""
+ self.completeChanged.emit()
+
+ def _on_export(self):
+ st = self._wiz.state
+ if self._busy or not st.out_path:
+ return
+
+ cropped, _ = self._wiz.cropped_plan()
+ worker = AlignedExportWorker(self._sras, cropped, st.out_path)
+ self._busy = True
+ self.completeChanged.emit()
+ self.btn_export.setEnabled(False)
+ self.btn_browse.setEnabled(False)
+ self.progress.setValue(0)
+ self.lbl_status.setText(f"Writing {Path(st.out_path).name}…")
+
+ started = self._wiz.run_worker(
+ Jobs.ALIGN_EXPORT, worker,
+ connect=(("progress", self.progress.setValue),
+ ("finished", self._on_export_finished)))
+ if not started:
+ self._busy = False
+ self.completeChanged.emit()
+ self.btn_export.setEnabled(True)
+ self.btn_browse.setEnabled(True)
+ self.lbl_status.setText(
+ "Another alignment step is still running — try again shortly.")
+ return
+ self._worker = worker
+ self.btn_cancel.setEnabled(True)
+
+ def _on_export_finished(self, written: str, error: str):
+ st = self._wiz.state
+ self._worker = None
+ self._busy = False
+ self.btn_export.setEnabled(True)
+ self.btn_browse.setEnabled(True)
+ self.btn_cancel.setEnabled(False)
+
+ if error:
+ self.progress.setValue(0)
+ self.lbl_status.setText(f"Export failed: {error}")
+ elif not written:
+ self.progress.setValue(0)
+ self.lbl_status.setText("Export cancelled; no file was written.")
+ else:
+ st.exported_path = written
+ self.progress.setValue(100)
+ self.lbl_status.setText(
+ f"Wrote {Path(written).name}. Choose Done to apply this "
+ f"alignment to the open scan as well.")
+ self.completeChanged.emit()
+ self._wiz.maybe_close_after_job()
+
+
+# Slack when snapping a rectangle edge, in canvas pixels. The mm round trip is
+# two multiplications and a subtraction, so an exact boundary can come back a
+# few ULPs either side of the integer.
+_SNAP_TOL = 1e-6
+
+
+def _snap_edge(coord: float) -> int:
+ """A fractional pixel-index edge, as the nearest pixel boundary index.
+
+ Pixel k spans [k - 0.5, k + 0.5), so a boundary at fractional coordinate
+ *coord* is boundary index coord + 0.5.
+ """
+ return int(np.floor(coord + 0.5 + _SNAP_TOL))
+
+
+def _humanize(n_bytes: int) -> str:
+ value = float(n_bytes)
+ for unit in ("B", "KB", "MB", "GB", "TB"):
+ if value < 1024 or unit == "TB":
+ return f"{value:.1f} {unit}" if unit != "B" else f"{int(value)} B"
+ value /= 1024
+ return f"{value:.1f} TB"
diff --git a/sras_viewer/canvases.py b/sras_viewer/canvases.py
new file mode 100644
index 0000000..8c0a1e4
--- /dev/null
+++ b/sras_viewer/canvases.py
@@ -0,0 +1,629 @@
+"""Matplotlib canvases and the ROI primitive."""
+
+import matplotlib as mpl
+import numpy as np
+from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg
+from matplotlib.colors import BoundaryNorm, ListedColormap
+from matplotlib.figure import Figure
+from matplotlib.patches import Polygon
+from matplotlib.path import Path as MplPath
+from PyQt6.QtCore import Qt, pyqtSignal
+from PyQt6.QtGui import QKeyEvent
+from PyQt6.QtWidgets import QSizePolicy
+
+from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, CH_NAMES, SrasFile, adc_to_mv
+
+def count_colormap(n_angles: int):
+ """(cmap, norm, ticks) for an integer "how many angles cover this pixel"
+ image, 0..n_angles.
+
+ Discrete, not continuous: the judgement the wizard's stack view exists for
+ is "is this a plateau at N, or a fan of partial overlaps", so a region
+ covered by one angle too few has to read as its own band rather than a
+ slightly darker shade. Count 0 is fully transparent so uncovered canvas
+ cannot be mistaken for a low count.
+
+ Shared by both wizard pages that draw this image — they use different canvas
+ classes, and the same number must not change colour between them.
+ """
+ n = max(1, int(n_angles))
+ base = mpl.colormaps["viridis"].resampled(n)
+ colors = [(0.0, 0.0, 0.0, 0.0)] + [base(i) for i in range(n)]
+ return (ListedColormap(colors),
+ BoundaryNorm(np.arange(-0.5, n + 1), len(colors)),
+ np.arange(0, n + 1))
+
+
+# ---------------------------------------------------------------------------
+# ROI (free quadrilateral in data coordinates)
+# ---------------------------------------------------------------------------
+
+class RoiQuad:
+ """Free quadrilateral defined in data coordinates (mm).
+
+ Stored as 4 corner points (shape (4, 2)) in CCW order: BL, BR, TR, TL.
+ Each corner can be positioned independently, allowing skewed /
+ non-orthogonal regions of interest. Because it lives in scan/data
+ coords it persists unchanged when the displayed channel/mode switches.
+ """
+
+ def __init__(self, pts: np.ndarray):
+ """pts : array-like, shape (4, 2)."""
+ self._pts = np.asarray(pts, dtype=np.float64).reshape(4, 2).copy()
+
+ @classmethod
+ def from_bbox(cls, x0: float, y0: float, x1: float, y1: float) -> "RoiQuad":
+ """Create an axis-aligned rectangle from two opposite corners."""
+ lx, rx = min(x0, x1), max(x0, x1)
+ by, ty = min(y0, y1), max(y0, y1)
+ return cls(np.array([[lx, by], [rx, by], [rx, ty], [lx, ty]]))
+
+ def copy(self) -> "RoiQuad":
+ return RoiQuad(self._pts.copy())
+
+ def corners(self) -> np.ndarray:
+ """World-coord corners, shape (4, 2), CCW: BL, BR, TR, TL."""
+ return self._pts.copy()
+
+ def centroid(self) -> np.ndarray:
+ return self._pts.mean(axis=0)
+
+ def bbox_size(self) -> np.ndarray:
+ """Width and height of the axis-aligned bounding box, shape (2,)."""
+ return self._pts.max(axis=0) - self._pts.min(axis=0)
+
+ def contains(self, x: float, y: float) -> bool:
+ return bool(MplPath(self._pts).contains_point((x, y)))
+
+ def mask_for_grid(self, x_axis: np.ndarray,
+ y_axis: np.ndarray) -> np.ndarray:
+ """Boolean mask (n_rows, n_frames) of pixels whose centres lie
+ inside the quadrilateral.
+
+ Only the quad's axis-aligned bounding box is tested — meshgrid and
+ contains_points over the *whole* grid would be tens of millions of
+ point-in-polygon tests (and hundreds of MB of float64 temporaries)
+ on a large scan, on every ROI edit.
+ """
+ x = np.asarray(x_axis, dtype=np.float64)
+ y = np.asarray(y_axis, dtype=np.float64)
+ mask = np.zeros((y.size, x.size), dtype=bool)
+
+ (x0, y0), (x1, y1) = self._pts.min(axis=0), self._pts.max(axis=0)
+ cols = np.nonzero((x >= x0) & (x <= x1))[0]
+ rows = np.nonzero((y >= y0) & (y <= y1))[0]
+ if cols.size == 0 or rows.size == 0:
+ return mask
+
+ c0, c1 = int(cols[0]), int(cols[-1]) + 1
+ r0, r1 = int(rows[0]), int(rows[-1]) + 1
+ X, Y = np.meshgrid(x[c0:c1], y[r0:r1])
+ inside = MplPath(self._pts).contains_points(
+ np.column_stack([X.ravel(), Y.ravel()]))
+ mask[r0:r1, c0:c1] = inside.reshape(X.shape)
+ return mask
+
+
+# ---------------------------------------------------------------------------
+# Matplotlib canvases
+# ---------------------------------------------------------------------------
+
+class ImageCanvas(FigureCanvasQTAgg):
+ pixel_clicked = pyqtSignal(int, int) # row_idx, frame_idx
+ roi_changed = pyqtSignal() # ROI created / edited / cleared
+ draw_mode_changed = pyqtSignal(bool) # "draw new ROI" arm toggled
+
+ # Interaction state values
+ _IDLE = "idle"
+ _DRAW_NEW = "draw_new"
+ _MOVE = "move"
+ _DRAG_CORNER = "drag_corner"
+
+ # Hit tolerance (display pixels) for handles.
+ _HANDLE_PX = 12
+ _CLICK_THRESH_PX = 4 # releases within this of press count as a click
+
+ def __init__(self, parent=None, *, rect_only: bool = False):
+ """*rect_only* constrains the ROI to an axis-aligned rectangle.
+
+ Used by the alignment wizard's crop page, where a free quadrilateral
+ would be actively misleading: v6 geometry can only express an
+ axis-aligned rectangle, so anything else the user drew would have to be
+ squared off behind their back. Default off, so the main window's
+ free-quad ROI is unaffected.
+ """
+ fig = Figure(figsize=(7, 5), tight_layout=True)
+ self.ax = fig.add_subplot(111)
+ super().__init__(fig)
+ self.setParent(parent)
+ self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
+ self._extent = None
+ self._img_shape = None
+ self._rect_only = rect_only
+
+ # ROI state
+ self._roi: RoiQuad | None = None
+ self._roi_artists: list = []
+ self._state = self._IDLE
+ self._draw_mode = False
+
+ # Per-interaction snapshots / anchors
+ self._press_xy: tuple[float, float] | None = None
+ self._press_pixel: tuple[float, float] | None = None
+ self._press_button = None
+ self._snapshot: RoiQuad | None = None
+ self._drag_corner_idx: int = -1
+ self._move_anchor = None # press-point in world coords
+ self._draw_previous: RoiQuad | None = None
+
+ self.mpl_connect("button_press_event", self._on_press)
+ self.mpl_connect("motion_notify_event", self._on_motion)
+ self.mpl_connect("button_release_event", self._on_release)
+
+ # ------------------------------------------------------------------
+ # Public API
+ # ------------------------------------------------------------------
+
+ def show_image(self, img: np.ndarray, extent: list[float], cmap,
+ vmin: float, vmax: float, xlabel: str, ylabel: str, title: str,
+ colorbar_label: str = "", cb_ticks=None, norm=None):
+ """*cmap* may be a name or a Colormap instance. *norm* (which overrides
+ vmin/vmax) and *cb_ticks* let a caller draw a discrete integer image —
+ the wizard's overlap-count view — with whole-number colorbar bands
+ instead of a continuous shade."""
+ self.figure.clf()
+ self.ax = self.figure.add_subplot(111)
+ # Patches and lines are destroyed by figure.clf(); drop stale refs.
+ self._roi_artists = []
+
+ self._extent = extent
+ self._img_shape = img.shape
+
+ kw = ({"norm": norm} if norm is not None
+ else {"vmin": vmin, "vmax": vmax})
+ im = self.ax.imshow(
+ img, aspect="auto", origin="upper",
+ extent=extent, cmap=cmap, interpolation="nearest", **kw,
+ )
+ cb = self.figure.colorbar(im, ax=self.ax, fraction=0.046, pad=0.04,
+ ticks=cb_ticks)
+ if colorbar_label:
+ cb.set_label(colorbar_label)
+
+ self.ax.set_xlabel(xlabel)
+ self.ax.set_ylabel(ylabel)
+ self.ax.set_title(title)
+
+ # Re-draw the ROI (if any) on top of the fresh image so it persists
+ # unchanged across mode / angle / channel switches.
+ self._draw_roi()
+ self.draw()
+
+ def get_roi(self) -> RoiQuad | None:
+ return self._roi
+
+ def set_roi(self, roi: RoiQuad | None):
+ self._roi = roi.copy() if roi is not None else None
+ self._draw_roi()
+ self.draw_idle()
+ self.roi_changed.emit()
+
+ def clear_roi(self):
+ self._roi = None
+ self._remove_roi_artists()
+ self.draw_idle()
+ self.roi_changed.emit()
+
+ def start_drawing(self):
+ """Arm the next click+drag on the image to create a new ROI,
+ replacing any existing one."""
+ self._draw_mode = True
+ self.setCursor(Qt.CursorShape.CrossCursor)
+ self.draw_mode_changed.emit(True)
+
+ def cancel_drawing(self):
+ if self._draw_mode:
+ self._draw_mode = False
+ self.setCursor(Qt.CursorShape.ArrowCursor)
+ self.draw_mode_changed.emit(False)
+
+ # ------------------------------------------------------------------
+ # Rendering
+ # ------------------------------------------------------------------
+
+ def _remove_roi_artists(self):
+ for a in self._roi_artists:
+ try:
+ a.remove()
+ except (ValueError, AttributeError, NotImplementedError):
+ pass
+ self._roi_artists = []
+
+ def _draw_roi(self):
+ self._remove_roi_artists()
+ if self._roi is None or self.ax is None:
+ return
+ corners = self._roi.corners()
+
+ # Filled quad, then a sharp unfilled edge for visibility over bright
+ # images, then draggable corner handles.
+ for kwargs in (
+ dict(fill=True, facecolor="#ffd93a", edgecolor="#e53935",
+ alpha=0.22, linewidth=2.0, zorder=10),
+ dict(fill=False, edgecolor="#e53935", linewidth=1.8, zorder=11),
+ ):
+ patch = Polygon(corners, closed=True, **kwargs)
+ self.ax.add_patch(patch)
+ self._roi_artists.append(patch)
+
+ self._roi_artists.append(self.ax.scatter(
+ corners[:, 0], corners[:, 1], s=60, c="white",
+ edgecolors="#e53935", linewidths=1.6, zorder=13))
+
+ # ------------------------------------------------------------------
+ # Hit testing (display pixels for handles, data coords for "inside")
+ # ------------------------------------------------------------------
+
+ def _hit_test(self, event) -> tuple[str, int | None] | None:
+ if self._roi is None or self.ax is None:
+ return None
+ if event.x is None or event.y is None:
+ return None
+ corners_disp = self.ax.transData.transform(self._roi.corners())
+ click = np.array([event.x, event.y])
+
+ for i in range(4):
+ if np.hypot(*(corners_disp[i] - click)) <= self._HANDLE_PX:
+ return ("corner", i)
+
+ if event.xdata is not None and event.ydata is not None:
+ if self._roi.contains(event.xdata, event.ydata):
+ return ("inside", None)
+ return None
+
+ # ------------------------------------------------------------------
+ # Mouse event handlers
+ # ------------------------------------------------------------------
+
+ def _on_press(self, event):
+ if event.inaxes is not self.ax or self._extent is None:
+ return
+ if event.button != 1: # only left mouse button
+ return
+ # If the matplotlib toolbar is in pan / zoom mode, let it handle
+ # the interaction instead of starting a ROI manipulation.
+ tb = getattr(self, "toolbar", None)
+ if tb is not None and getattr(tb, "mode", ""):
+ return
+
+ self._press_xy = (event.xdata, event.ydata)
+ self._press_pixel = (event.x, event.y)
+ self._press_button = event.button
+
+ if self._draw_mode:
+ self._draw_previous = self._roi.copy() if self._roi else None
+ self._roi = RoiQuad.from_bbox(event.xdata, event.ydata,
+ event.xdata, event.ydata)
+ self._state = self._DRAW_NEW
+ self._draw_roi()
+ self.draw_idle()
+ return
+
+ hit = self._hit_test(event)
+ if hit is None:
+ self._state = self._IDLE
+ return
+
+ kind, idx = hit
+ self._snapshot = self._roi.copy()
+ if kind == "corner":
+ self._state = self._DRAG_CORNER
+ self._drag_corner_idx = idx
+ else:
+ self._state = self._MOVE
+ self._move_anchor = (event.xdata, event.ydata)
+
+ def _on_motion(self, event):
+ if self._state == self._IDLE:
+ return
+ if event.xdata is None or event.ydata is None:
+ return
+ if event.inaxes is not self.ax:
+ return
+
+ if self._state == self._DRAW_NEW:
+ x0, y0 = self._press_xy
+ self._roi = RoiQuad.from_bbox(x0, y0, event.xdata, event.ydata)
+ elif self._state == self._MOVE:
+ delta = np.array([event.xdata - self._move_anchor[0],
+ event.ydata - self._move_anchor[1]])
+ self._roi._pts = self._snapshot.corners() + delta
+ elif self._state == self._DRAG_CORNER:
+ self._roi._pts[self._drag_corner_idx] = [event.xdata, event.ydata]
+ if self._rect_only:
+ self._rectify_corner(self._drag_corner_idx)
+
+ self._draw_roi()
+ self.draw_idle()
+
+ def _rectify_corner(self, idx: int):
+ """Re-square the quad after a corner drag, anchored on the *opposite*
+ corner.
+
+ Anchoring on the diagonal opposite (idx ^ 2, since corners run
+ BL, BR, TR, TL) rather than taking the bbox of all four points is what
+ lets the rectangle shrink: a bbox over the three stale corners plus the
+ new one is the union of the old rectangle and the new point, so dragging
+ inward would never make it smaller.
+ """
+ pts = self._roi.corners()
+ ax_, ay = pts[idx ^ 2]
+ bx, by = pts[idx]
+ self._roi._pts = RoiQuad.from_bbox(min(ax_, bx), min(ay, by),
+ max(ax_, bx), max(ay, by)).corners()
+
+ def _on_release(self, event):
+ if event.button != 1 and self._press_button != 1:
+ return
+ prev_state = self._state
+ self._state = self._IDLE
+ try:
+ if prev_state == self._DRAW_NEW:
+ self._finish_draw()
+ elif prev_state in (self._MOVE, self._DRAG_CORNER):
+ self._draw_roi()
+ self.draw_idle()
+ self.roi_changed.emit()
+ else:
+ self._maybe_emit_pixel_click(event)
+ finally:
+ self._press_xy = self._press_pixel = None
+ self._press_button = None
+
+ def _finish_draw(self):
+ """Commit (or reject) a freshly-dragged quad."""
+ if self._extent is not None:
+ x0, x1, y_bot, y_top = self._extent
+ min_w = abs(x1 - x0) * 0.01 # minimum: 1% of each axis range
+ min_h = abs(y_bot - y_top) * 0.01
+ else:
+ min_w = min_h = 1e-6
+
+ if self._roi is None:
+ too_small = True
+ else:
+ bbox = self._roi.bbox_size()
+ too_small = bbox[0] < min_w or bbox[1] < min_h
+ if too_small:
+ self._roi = self._draw_previous
+
+ self._draw_previous = None
+ self.cancel_drawing()
+ self._draw_roi()
+ self.draw_idle()
+ self.roi_changed.emit()
+
+ def _maybe_emit_pixel_click(self, event):
+ """A release close enough to its press counts as a pixel click."""
+ if (self._press_pixel is None or event.x is None or event.y is None
+ or self._extent is None or event.inaxes is not self.ax
+ or event.xdata is None):
+ return
+ dx_px = event.x - self._press_pixel[0]
+ dy_px = event.y - self._press_pixel[1]
+ if dx_px * dx_px + dy_px * dy_px > self._CLICK_THRESH_PX ** 2:
+ return
+
+ x0, x1, y_bot, y_top = self._extent
+ n_rows, n_frames = self._img_shape
+ col = int((event.xdata - x0) / (x1 - x0) * n_frames)
+ row = int((event.ydata - y_top) / (y_bot - y_top) * n_rows)
+ self.pixel_clicked.emit(max(0, min(row, n_rows - 1)),
+ max(0, min(col, n_frames - 1)))
+
+
+class WaveformCanvas(FigureCanvasQTAgg):
+ def __init__(self, parent=None):
+ fig = Figure(figsize=(8, 3), tight_layout=True)
+ self.ax_wave = fig.add_subplot(121)
+ self.ax_right = fig.add_subplot(122)
+ super().__init__(fig)
+ self.setParent(parent)
+ self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
+
+ def show_rf_waveform(self, sras: SrasFile, angle_idx: int,
+ row_idx: int, frame_idx: int,
+ apply_bg_sub: bool = True):
+ """CH1 RF: time-domain + FFT spectrum.
+
+ If apply_bg_sub is True and sras.background is not None, the background
+ waveform is overlaid on the time-domain plot and the FFT is computed
+ on the subtracted signal. The unsubtracted FFT is also shown faintly
+ for comparison.
+ """
+ data = sras.data[angle_idx]
+ waveform = data[row_idx, CH1_IDX, frame_idx, :].astype(np.float32)
+ t_ns = sras.time_axis_ns()
+ f_mhz = sras.freq_axis_mhz()
+ dc3_val = data[row_idx, CH3_IDX, frame_idx, :].astype(np.float32).mean()
+ dc4_val = data[row_idx, CH4_IDX, frame_idx, :].astype(np.float32).mean()
+
+ bg = sras.background if (apply_bg_sub and sras.background is not None) else None
+ waveform_plot = waveform - bg if bg is not None else waveform
+
+ self.ax_wave.cla()
+ self.ax_right.cla()
+
+ if bg is not None:
+ self.ax_wave.plot(t_ns, waveform, linewidth=0.5, color="#aaaaaa",
+ label="raw", zorder=1)
+ self.ax_wave.plot(t_ns, bg, linewidth=0.5, color="#e07030",
+ linestyle="--", label="background", zorder=2)
+ self.ax_wave.plot(t_ns, waveform_plot, linewidth=0.7, color="#4488cc",
+ label="subtracted", zorder=3)
+ self.ax_wave.legend(fontsize=7, loc="upper right")
+ else:
+ self.ax_wave.plot(t_ns, waveform, linewidth=0.7, color="#4488cc")
+
+ self.ax_wave.set_xlabel("Time (ns)")
+ self.ax_wave.set_ylabel("ADC counts")
+ bg_tag = " [bg sub]" if bg is not None else ""
+ dc3_mv = adc_to_mv(dc3_val, *sras.cal(CH3_IDX))
+ dc4_mv = adc_to_mv(dc4_val, *sras.cal(CH4_IDX))
+ self.ax_wave.set_title(
+ f"CH1 RF row={row_idx} frame={frame_idx}{bg_tag}\n"
+ f"CH3={dc3_val:.1f} CH4={dc4_val:.1f} "
+ f"({dc3_mv:.2f} / {dc4_mv:.2f} mV)",
+ fontsize=8,
+ )
+
+ # FFT of the (possibly subtracted) waveform
+ power_sub = np.abs(np.fft.rfft(waveform_plot)) ** 2
+ power_sub[0] = 0.0
+ peak_mhz = f_mhz[int(np.argmax(power_sub))]
+
+ if bg is not None:
+ # Also show the unsubtracted FFT for reference
+ power_raw = np.abs(np.fft.rfft(waveform)) ** 2
+ power_raw[0] = 0.0
+ self.ax_right.plot(f_mhz, power_raw, linewidth=0.5, color="#aaaaaa",
+ label="raw FFT", zorder=1)
+
+ self.ax_right.plot(f_mhz, power_sub, linewidth=0.7, color="#4488cc",
+ label="subtracted FFT" if bg is not None else None, zorder=2)
+ self.ax_right.axvline(peak_mhz, color="tomato", linestyle="--",
+ linewidth=1.2, label=f"peak = {peak_mhz:.1f} MHz")
+ self.ax_right.set_xlabel("Frequency (MHz)")
+ self.ax_right.set_ylabel("Power (arb.)")
+ self.ax_right.set_title("FFT Power Spectrum")
+ self.ax_right.set_xlim(0, 500)
+ self.ax_right.legend(fontsize=8)
+
+ self.draw()
+
+ def show_dc_waveform(self, sras: SrasFile, angle_idx: int, ch_idx: int,
+ row_idx: int, frame_idx: int):
+ """CH3 or CH4 DC: time-domain + mean annotation."""
+ waveform = sras.data[angle_idx][row_idx, ch_idx, frame_idx, :].astype(np.float32)
+ mean_val = float(waveform.mean())
+ mean_mv = adc_to_mv(mean_val, *sras.cal(ch_idx))
+
+ self.ax_wave.cla()
+ self.ax_right.cla()
+
+ self.ax_wave.plot(sras.time_axis_ns(), waveform, linewidth=0.7, color="#4488cc")
+ self.ax_wave.axhline(mean_val, color="tomato", linestyle="--",
+ linewidth=1.2, label=f"mean = {mean_val:.2f} ADC")
+ self.ax_wave.set_xlabel("Time (ns)")
+ self.ax_wave.set_ylabel("ADC counts")
+ self.ax_wave.set_title(
+ f"{CH_NAMES[ch_idx]} DC row={row_idx} frame={frame_idx}")
+ self.ax_wave.legend(fontsize=8)
+
+ self.ax_right.text(
+ 0.5, 0.5,
+ f"DC mode\n\nmean = {mean_val:.3f} ADC\n = {mean_mv:.3f} mV",
+ ha="center", va="center",
+ transform=self.ax_right.transAxes, fontsize=11,
+ )
+ self.ax_right.set_axis_off()
+
+ self.draw()
+
+
+class AlignOverlayCanvas(FigureCanvasQTAgg):
+ """Renders the alignment wizard's multi-angle mask views and turns keyboard
+ input into translate/rotate nudge requests for whichever angle is active.
+
+ Two views of the same reprojected masks, because they answer different
+ questions. show_counts colours each pixel by *how many* angles cover it,
+ which is the at-a-glance verdict on a correlation run: a good alignment is
+ one saturated plateau, a bad one is a fringe of low-count halos.
+ show_overlay gives each angle its own colour, which is what you need while
+ nudging a specific angle by hand.
+
+ A pure input+render widget — it holds no alignment state and never
+ touches SrasFile itself; the wizard page owns all of that and decides,
+ from these signals, whether a cheap single-layer refresh or a
+ full preview-canvas rebuild is needed.
+
+ FigureCanvasQTAgg is a real QWidget, so keyPressEvent works like on any
+ other widget, but Qt only ever delivers key events to whichever widget
+ currently has focus — StrongFocus, plus grabbing focus on click and once
+ right after the dialog is shown, are both required or arrow keys
+ silently do nothing.
+
+ Rotate keys are letters (Q/E), not punctuation (comma/period or
+ brackets): Shift+letter still reports the same Qt.Key on every platform,
+ whereas Shift+comma/bracket can report a different virtual key
+ (Key_Less / Key_BraceLeft) depending on platform and keyboard layout —
+ which would silently break the "Shift = coarse step" modifier for
+ rotation specifically. Arrow keys have no such hazard.
+ """
+ nudge_translate = pyqtSignal(int, int, bool) # dir_x, dir_y in {-1,0,1}; coarse
+ nudge_rotate = pyqtSignal(int, bool) # dir in {-1,1} (CCW/CW); coarse
+
+ _TRANSLATE_KEYS = {
+ Qt.Key.Key_Left: (-1, 0),
+ Qt.Key.Key_Right: (1, 0),
+ Qt.Key.Key_Up: (0, -1),
+ Qt.Key.Key_Down: (0, 1),
+ }
+ _ROTATE_KEYS = {Qt.Key.Key_Q: 1, Qt.Key.Key_E: -1} # CCW, CW
+
+ def __init__(self, parent=None):
+ fig = Figure(figsize=(6, 6), tight_layout=True)
+ self.ax = fig.add_subplot(111)
+ super().__init__(fig)
+ self.setParent(parent)
+ self.setFocusPolicy(Qt.FocusPolicy.StrongFocus)
+ self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
+ self.mpl_connect("button_press_event", lambda _e: self.setFocus())
+
+ def show_overlay(self, rgba: np.ndarray, extent: list[float], title: str):
+ self.figure.clf()
+ self.ax = self.figure.add_subplot(111)
+ self.ax.imshow(rgba, extent=extent, origin="upper", aspect="auto")
+ self._finish(title)
+
+ def show_counts(self, counts: np.ndarray, n_angles: int,
+ extent: list[float], title: str):
+ """The mask stack coloured by how many angles cover each pixel.
+
+ A discrete colormap with integer-ticked colorbar rather than a
+ continuous one: the judgement being made is "is this a plateau at N, or
+ a fan of partial overlaps", and a region covered by one angle too few
+ has to read as its own band rather than a slightly darker shade.
+ Uncovered pixels are transparent so they cannot be mistaken for a low
+ count.
+ """
+ self.figure.clf()
+ self.ax = self.figure.add_subplot(111)
+ cmap, norm, ticks = count_colormap(n_angles)
+ im = self.ax.imshow(
+ np.asarray(counts), extent=extent, origin="upper", aspect="auto",
+ interpolation="nearest", cmap=cmap, norm=norm)
+ cb = self.figure.colorbar(im, ax=self.ax, fraction=0.046, pad=0.04,
+ ticks=ticks)
+ cb.set_label("angles overlapping")
+ self._finish(title)
+
+ def _finish(self, title: str):
+ self.ax.set_xlabel("X (mm)")
+ self.ax.set_ylabel("Y (mm)")
+ self.ax.set_title(title)
+ self.draw_idle() # coalesces rapid redraws — matters for key-repeat.
+
+ def keyPressEvent(self, event: QKeyEvent):
+ key = event.key()
+ coarse = bool(event.modifiers() & Qt.KeyboardModifier.ShiftModifier)
+ if key in self._TRANSLATE_KEYS:
+ dx, dy = self._TRANSLATE_KEYS[key]
+ self.nudge_translate.emit(dx, dy, coarse)
+ event.accept()
+ elif key in self._ROTATE_KEYS:
+ self.nudge_rotate.emit(self._ROTATE_KEYS[key], coarse)
+ event.accept()
+ else:
+ super().keyPressEvent(event)
+
diff --git a/sras_viewer/common.py b/sras_viewer/common.py
new file mode 100644
index 0000000..ad587e9
--- /dev/null
+++ b/sras_viewer/common.py
@@ -0,0 +1,178 @@
+"""Shared constants and small layout helpers for the viewer widgets."""
+
+from PyQt6.QtCore import Qt
+from PyQt6.QtWidgets import (
+ QComboBox, QDoubleSpinBox, QFormLayout, QFrame, QGroupBox, QLabel,
+ QScrollArea, QSizePolicy, QVBoxLayout, QWidget,
+)
+
+from sras_format import CH1_IDX, CH3_IDX, CH4_IDX
+
+# ---------------------------------------------------------------------------
+# Display constants
+# ---------------------------------------------------------------------------
+
+CH_LABELS = [
+ "CH1 — RF (FFT peak freq)",
+ "CH3 — Bias A (DC mean)",
+ "CH4 — Bias B (DC mean)",
+ "CH1 — Velocity (SRAS)",
+]
+
+# Combo index for the derived velocity mode (uses CH1_IDX data)
+VELOCITY_MODE_IDX = 3
+# All modes that operate on CH1 waveforms
+CH1_DERIVED_MODES = (CH1_IDX, VELOCITY_MODE_IDX)
+
+CMAPS = ["gray", "viridis", "plasma", "inferno", "hot", "jet", "RdBu_r", "seismic"]
+
+# (mode_str, status-bar unit, colorbar label) per channel index
+_CHANNEL_DISPLAY = {
+ CH1_IDX: ("RF", "Peak frequency (MHz)", "MHz"),
+ CH3_IDX: ("DC", "DC mean (mV)", "mV"),
+ CH4_IDX: ("DC", "DC mean (mV)", "mV"),
+ VELOCITY_MODE_IDX: ("Velocity", "Velocity (m/s)", "m/s"),
+}
+
+_CSS_HINT = "font-size: 11px; color: #aaa;"
+_CSS_INFO = "font-size: 11px;"
+_CSS_MUTED = "color: #888; font-size: 11px;"
+_CSS_WARN = "color: #e07000; font-size: 11px;"
+_CSS_BUSY = "color: #4a90d9; font-size: 11px;"
+
+# Side-panel column widths (the scroll areas that hold the controls).
+_LEFT_PANEL_W = 288
+_RIGHT_PANEL_W = 272
+
+# Minimum width for a spin box so its value + suffix are never clipped.
+_SPIN_MIN_W = 96
+
+
+# ---------------------------------------------------------------------------
+# Small layout helpers
+# ---------------------------------------------------------------------------
+
+class Jobs:
+ """Keys for SrasViewerWindow's background-job registry (_run_worker /
+ _job_running) and its progress dialogs — one place instead of string
+ literals scattered across window and dialogs."""
+ LOAD = "load"
+ COMPUTE = "compute"
+ DC_PRECOMPUTE = "dc_precompute"
+ BATCH = "batch"
+ # The alignment wizard's three background steps: fetching each angle's CH4
+ # image for the mask stack, registering the angles, and writing the aligned
+ # export. Separate keys because a retry of one must not be blocked by
+ # another having run, and _run_worker's busy check is per key.
+ ALIGN_MASKS = "align_masks"
+ ALIGN_CORRELATE = "align_correlate"
+ ALIGN_EXPORT = "align_export"
+
+
+def _make_dspin(lo: float, hi: float, decimals: int, *, suffix: str = "",
+ value: float | None = None, step: float | None = None) -> QDoubleSpinBox:
+ """A QDoubleSpinBox with the panel-standard construction."""
+ spin = QDoubleSpinBox()
+ spin.setRange(lo, hi)
+ spin.setDecimals(decimals)
+ if suffix:
+ spin.setSuffix(suffix)
+ if step is not None:
+ spin.setSingleStep(step)
+ if value is not None:
+ spin.setValue(value)
+ spin.setMinimumWidth(_SPIN_MIN_W)
+ return spin
+
+
+def _combo(items=(), *, min_chars: int = 10) -> QComboBox:
+ """A combo box whose size hint does not depend on its longest entry.
+
+ By default a QComboBox asks for enough width to show its widest item. These
+ hold descriptive phrases, and the side panels are fixed-width — in a scroll
+ area with the horizontal scrollbar off (`_scroll_panel`) an unconstrained
+ hint pushes the inner widget past the panel and everything on the right,
+ including the hint text, is silently clipped instead of scrolling.
+
+ *items* is a sequence of (label, data) pairs, or of plain labels.
+ """
+ combo = QComboBox()
+ combo.setSizeAdjustPolicy(
+ QComboBox.SizeAdjustPolicy.AdjustToMinimumContentsLengthWithIcon)
+ combo.setMinimumContentsLength(min_chars)
+ for item in items:
+ if isinstance(item, tuple):
+ combo.addItem(item[0], item[1])
+ else:
+ combo.addItem(item)
+ return combo
+
+
+def _axes_extent(x_axis, y_axis, dx: float, dy: float) -> list[float]:
+ """Matplotlib imshow extent with half-pixel margins, Y flipped so row 0
+ renders at the top."""
+ return [x_axis[0] - dx / 2, x_axis[-1] + dx / 2,
+ y_axis[-1] + dy / 2, y_axis[0] - dy / 2]
+
+
+def _wrap_label(text: str = "", css: str | None = None) -> QLabel:
+ """A word-wrapped QLabel that reports its *wrapped* height to the layout.
+
+ A plain word-wrapped QLabel advertises a single-line minimum height, so in a
+ fixed-width column the layout happily shrinks it and the extra lines get
+ clipped. Enabling height-for-width makes the box layout ask for the real
+ height at the column's width instead.
+ """
+ lbl = QLabel(text)
+ lbl.setWordWrap(True)
+ sp = lbl.sizePolicy()
+ sp.setVerticalPolicy(QSizePolicy.Policy.Minimum)
+ sp.setHeightForWidth(True)
+ lbl.setSizePolicy(sp)
+ if css:
+ lbl.setStyleSheet(css)
+ return lbl
+
+
+def _group(title: str) -> tuple[QGroupBox, QVBoxLayout]:
+ """A group box with consistent, non-cramped internal margins."""
+ grp = QGroupBox(title)
+ lay = QVBoxLayout(grp)
+ lay.setContentsMargins(10, 8, 10, 10)
+ lay.setSpacing(6)
+ return grp, lay
+
+
+def _form() -> QFormLayout:
+ """A label/field form layout for a narrow side panel."""
+ form = QFormLayout()
+ form.setContentsMargins(0, 0, 0, 0)
+ form.setHorizontalSpacing(8)
+ form.setVerticalSpacing(6)
+ form.setLabelAlignment(Qt.AlignmentFlag.AlignRight
+ | Qt.AlignmentFlag.AlignVCenter)
+ form.setFormAlignment(Qt.AlignmentFlag.AlignLeft | Qt.AlignmentFlag.AlignTop)
+ form.setFieldGrowthPolicy(
+ QFormLayout.FieldGrowthPolicy.AllNonFixedFieldsGrow)
+ form.setRowWrapPolicy(QFormLayout.RowWrapPolicy.DontWrapRows)
+ return form
+
+
+def _scroll_panel(inner: QWidget, width: int) -> QScrollArea:
+ """Put a side panel in a fixed-width scroll area.
+
+ Without this the panels are sized by the window: a short window squeezes the
+ controls past their minimum heights, which is what makes text overlap the
+ widget below it. Scrolling keeps every control at its natural size.
+ """
+ area = QScrollArea()
+ area.setWidget(inner)
+ area.setWidgetResizable(True)
+ area.setFrameShape(QFrame.Shape.NoFrame)
+ area.setHorizontalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAlwaysOff)
+ area.setVerticalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAsNeeded)
+ area.setFixedWidth(width)
+ area.viewport().setAutoFillBackground(False)
+ inner.setAutoFillBackground(False)
+ return area
+
diff --git a/sras_viewer/dialogs.py b/sras_viewer/dialogs.py
new file mode 100644
index 0000000..87379eb
--- /dev/null
+++ b/sras_viewer/dialogs.py
@@ -0,0 +1,415 @@
+"""FFT option dialogs.
+
+Angle alignment used to live here too, as ManualAlignmentDialog; it is now the
+alignment wizard's first page (see align_wizard.py), which needed the same
+mask-overlay editor plus the crop and export steps.
+"""
+
+from pathlib import Path
+
+from PyQt6.QtWidgets import (
+ QButtonGroup, QCheckBox, QDialog, QDialogButtonBox, QFileDialog,
+ QGroupBox, QHBoxLayout, QLabel, QLineEdit, QPushButton, QRadioButton,
+ QScrollArea, QSpinBox, QVBoxLayout, QWidget,
+)
+
+from sras_compute import PYFFTW_AVAILABLE
+from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, CH_NAMES
+
+from .common import CH_LABELS, VELOCITY_MODE_IDX, _CSS_HINT, _CSS_WARN, _group, _make_dspin, _wrap_label
+
+
+# ---------------------------------------------------------------------------
+# FFT Options dialog
+# ---------------------------------------------------------------------------
+
+class FftOptionsDialog(QDialog):
+ """Configure FFT backend and zero-padding.
+
+ Changes take effect only when the user clicks Apply. Cancel discards
+ all pending edits. The live 'frequency resolution' label updates as
+ the user adjusts the pad factor so they can see the trade-off before
+ committing.
+ """
+
+ def __init__(self, parent=None, *,
+ current_backend: str,
+ current_pad_factor: int,
+ samples_per_frame: int | None,
+ sample_rate_hz: float | None,
+ grating_um: float):
+ super().__init__(parent)
+ self.setWindowTitle("FFT Options")
+ self.setModal(True)
+ self.setMinimumWidth(380)
+
+ self._samples_per_frame = samples_per_frame
+ self._sample_rate_hz = sample_rate_hz
+ self._grating_um = grating_um
+
+ layout = QVBoxLayout(self)
+
+ # ---- Backend ---------------------------------------------------
+ grp_backend = QGroupBox("FFT Backend")
+ bl = QVBoxLayout(grp_backend)
+
+ self._btn_scipy = QRadioButton("SciPy FFT (pocketfft) (always available)")
+ self._btn_pyfftw = QRadioButton(
+ "pyFFTW (faster for large arrays)" if PYFFTW_AVAILABLE
+ else "pyFFTW (not installed — run: pip install pyfftw)")
+ self._btn_pyfftw.setEnabled(PYFFTW_AVAILABLE)
+
+ self._backend_group = QButtonGroup(self)
+ self._backend_group.addButton(self._btn_scipy, id=0)
+ self._backend_group.addButton(self._btn_pyfftw, id=1)
+
+ if current_backend == "pyfftw" and PYFFTW_AVAILABLE:
+ self._btn_pyfftw.setChecked(True)
+ else:
+ self._btn_scipy.setChecked(True)
+
+ bl.addWidget(self._btn_scipy)
+ bl.addWidget(self._btn_pyfftw)
+ layout.addWidget(grp_backend)
+
+ # ---- Zero-padding ----------------------------------------------
+ grp_zp = QGroupBox("Zero-Padding")
+ zl = QVBoxLayout(grp_zp)
+
+ pad_row = QHBoxLayout()
+ pad_row.addWidget(QLabel("Pad factor:"))
+ self._spin_pad = QSpinBox()
+ self._spin_pad.setRange(1, 256)
+ self._spin_pad.setValue(max(1, current_pad_factor))
+ self._spin_pad.setToolTip(
+ "Multiply the waveform length by this factor via zero-padding\n"
+ "before computing the FFT.\n"
+ "1 = no padding (natural length).\n"
+ "Powers of 2 (2, 4, 8 …) give the best performance."
+ )
+ self._spin_pad.valueChanged.connect(self._update_info)
+ pad_row.addWidget(self._spin_pad)
+ zl.addLayout(pad_row)
+
+ self._lbl_nfft = QLabel()
+ self._lbl_freq_res = QLabel()
+ self._lbl_vel_res = QLabel()
+ for lbl in (self._lbl_nfft, self._lbl_freq_res, self._lbl_vel_res):
+ lbl.setStyleSheet(_CSS_HINT)
+ zl.addWidget(lbl)
+
+ layout.addWidget(grp_zp)
+
+ # ---- Buttons ---------------------------------------------------
+ buttons = QDialogButtonBox()
+ buttons.addButton("Apply", QDialogButtonBox.ButtonRole.AcceptRole
+ ).clicked.connect(self.accept)
+ buttons.addButton("Cancel", QDialogButtonBox.ButtonRole.RejectRole
+ ).clicked.connect(self.reject)
+ layout.addWidget(buttons)
+
+ self._update_info()
+
+ def _update_info(self):
+ spf = self._samples_per_frame
+ sr = self._sample_rate_hz
+ pad = self._spin_pad.value()
+
+ if spf is None or sr is None:
+ self._lbl_nfft.setText("Load a file to preview FFT parameters.")
+ self._lbl_freq_res.setText("")
+ self._lbl_vel_res.setText("")
+ return
+
+ n_fft = spf * pad
+ freq_res_hz = sr / n_fft
+ freq_res_mhz = freq_res_hz / 1e6
+ # v (m/s) = freq (MHz) × grating (µm)
+ vel_res_ms = freq_res_mhz * self._grating_um
+
+ self._lbl_nfft.setText(f"FFT points: {spf} × {pad} = {n_fft:,}")
+ self._lbl_freq_res.setText(
+ f"Frequency bin: {freq_res_mhz:.4f} MHz ({freq_res_hz / 1e3:.2f} kHz)")
+ self._lbl_vel_res.setText(
+ f"Velocity bin: {vel_res_ms:.3f} m/s "
+ f"(at grating = {self._grating_um:.2f} µm)")
+
+ def get_backend(self) -> str:
+ return "pyfftw" if self._btn_pyfftw.isChecked() and PYFFTW_AVAILABLE else "scipy"
+
+ def get_pad_factor(self) -> int:
+ return max(1, self._spin_pad.value())
+
+
+# ---------------------------------------------------------------------------
+# Row-Averaged FFT Options dialog
+# ---------------------------------------------------------------------------
+
+class RowAverageFftOptionsDialog(QDialog):
+ """Configure the same-row, distance-weighted neighbor averaging applied
+ to each pixel's CH1 waveform before 'Batch Compute Row-Averaged FFT and
+ Store' re-runs the FFT peak search — a same-row SNR cleanup pass, never
+ mixing across rows/Y (see sras_compute._row_average_waveforms).
+
+ Unlike the plain FFT batch action (which stores unmasked and defers
+ masking to display time), the DC threshold here is required up front:
+ it decides which same-row neighbors are eligible to contribute to a
+ pixel's average, so it can't be deferred.
+
+ Changes take effect only when the user clicks Apply. Cancel discards
+ all pending edits.
+ """
+
+ def __init__(self, parent=None, *,
+ current_n: int,
+ current_threshold_mv: float,
+ pixel_x_mm: float | None):
+ super().__init__(parent)
+ self.setWindowTitle("Row-Averaged FFT Options")
+ self.setModal(True)
+ self.setMinimumWidth(380)
+
+ self._pixel_x_mm = pixel_x_mm
+
+ layout = QVBoxLayout(self)
+
+ # ---- Neighbor window ---------------------------------------------
+ grp_window = QGroupBox("Same-Row Neighbor Window")
+ wl = QVBoxLayout(grp_window)
+
+ n_row = QHBoxLayout()
+ n_row.addWidget(QLabel("Neighbor half-width (n):"))
+ self._spin_n = QSpinBox()
+ self._spin_n.setRange(1, 50)
+ self._spin_n.setValue(max(1, current_n))
+ self._spin_n.setToolTip(
+ "Each pixel's CH1 waveform is averaged with up to n same-row\n"
+ "neighbors on each side, distance-weighted (Gaussian) and\n"
+ "counting only neighbors that already pass the DC threshold\n"
+ "below. Never mixes across rows/Y.")
+ self._spin_n.valueChanged.connect(self._update_info)
+ n_row.addWidget(self._spin_n)
+ wl.addLayout(n_row)
+
+ self._lbl_width = QLabel()
+ self._lbl_width.setStyleSheet(_CSS_HINT)
+ wl.addWidget(self._lbl_width)
+
+ layout.addWidget(grp_window)
+
+ # ---- DC threshold ------------------------------------------------
+ grp_thr = QGroupBox("Neighbor Validity")
+ tl = QVBoxLayout(grp_thr)
+ thr_row = QHBoxLayout()
+ thr_row.addWidget(QLabel("DC threshold:"))
+ self._spin_threshold = _make_dspin(-500.0, 500.0, 3, suffix=" mV",
+ value=current_threshold_mv, step=0.025)
+ self._spin_threshold.setToolTip(
+ "A same-row neighbor only contributes to a pixel's average if\n"
+ "its own CH4 signal is at or above this threshold -- the same\n"
+ "test used for RF mask display. A pixel below threshold stays\n"
+ "masked, exactly as today; it is never rescued by its neighbors.")
+ thr_row.addWidget(self._spin_threshold)
+ tl.addLayout(thr_row)
+ layout.addWidget(grp_thr)
+
+ # ---- Buttons -----------------------------------------------------
+ buttons = QDialogButtonBox()
+ buttons.addButton("Apply", QDialogButtonBox.ButtonRole.AcceptRole
+ ).clicked.connect(self.accept)
+ buttons.addButton("Cancel", QDialogButtonBox.ButtonRole.RejectRole
+ ).clicked.connect(self.reject)
+ layout.addWidget(buttons)
+
+ self._update_info()
+
+ def _update_info(self):
+ n = self._spin_n.value()
+ if self._pixel_x_mm is None:
+ self._lbl_width.setText("Load a file to preview the window's physical width.")
+ return
+ width_um = 2 * n * self._pixel_x_mm * 1e3
+ self._lbl_width.setText(
+ f"Window: ±{n} px = {width_um:.2f} µm full width "
+ f"(pixel pitch {self._pixel_x_mm * 1e3:.3g} µm)")
+
+ def get_half_width(self) -> int:
+ return self._spin_n.value()
+
+ def get_threshold_mv(self) -> float:
+ return self._spin_threshold.value()
+
+
+# ---------------------------------------------------------------------------
+# Export Fused ROI dialog
+# ---------------------------------------------------------------------------
+
+class FusedRoiExportDialog(QDialog):
+ """Choose a value type and which angles to fuse for Export Fused ROI.
+
+ Every angle in the file is listed (a live AlignmentResult's per_angle
+ always covers every angle, and the raw-shared-grid path needs no per-
+ angle transform at all). Each checkbox is disabled — and auto-unchecked
+ — whenever *availability_fn(angle_idx, ch_idx)* is False for the
+ currently selected value type; switching the value-type radio
+ re-evaluates every checkbox live, since availability is per (angle,
+ value type) rather than just per angle — e.g. DC may be ready
+ everywhere while FFT is ready nowhere.
+ """
+
+ _VALUE_MODES = (CH1_IDX, CH3_IDX, CH4_IDX, VELOCITY_MODE_IDX)
+
+ def __init__(self, parent=None, *,
+ angles: list[tuple[int, float]],
+ availability_fn,
+ default_ch_idx: int,
+ out_dir: str,
+ stem: str,
+ grid_note: str):
+ super().__init__(parent)
+ self.setWindowTitle("Export Fused ROI")
+ self.setModal(True)
+ self.setMinimumWidth(420)
+
+ self._availability_fn = availability_fn
+ self._out_dir = out_dir
+ self._stem = stem
+ self._path_user_chosen = False
+
+ layout = QVBoxLayout(self)
+ layout.addWidget(_wrap_label(grid_note, _CSS_HINT))
+
+ # ---- Value type --------------------------------------------------
+ grp_val, vl = _group("Value to Export")
+ self._val_group = QButtonGroup(self)
+ self._val_buttons: dict[int, QRadioButton] = {}
+ for ch_idx, label in zip(self._VALUE_MODES, CH_LABELS):
+ rb = QRadioButton(label)
+ self._val_group.addButton(rb, id=ch_idx)
+ self._val_buttons[ch_idx] = rb
+ vl.addWidget(rb)
+ self._val_buttons[default_ch_idx].setChecked(True)
+ self._val_group.idClicked.connect(self._on_value_type_changed)
+ layout.addWidget(grp_val)
+
+ # ---- Angles --------------------------------------------------
+ grp_ang, al = _group("Angles to Include")
+ sel_row = QHBoxLayout()
+ btn_all = QPushButton("Select All Available")
+ btn_none = QPushButton("Select None")
+ btn_all.clicked.connect(self._on_select_all_available)
+ btn_none.clicked.connect(self._on_select_none)
+ sel_row.addWidget(btn_all)
+ sel_row.addWidget(btn_none)
+ al.addLayout(sel_row)
+
+ scroll_inner = QWidget()
+ scroll_layout = QVBoxLayout(scroll_inner)
+ self._angle_checks: dict[int, QCheckBox] = {}
+ for angle_idx, angle_deg in angles:
+ cb = QCheckBox(f"{angle_deg:.1f}° (angle {angle_idx})")
+ self._angle_checks[angle_idx] = cb
+ cb.toggled.connect(self._update_accept_enabled)
+ scroll_layout.addWidget(cb)
+ scroll = QScrollArea()
+ scroll.setWidget(scroll_inner)
+ scroll.setWidgetResizable(True)
+ scroll.setMaximumHeight(220)
+ al.addWidget(scroll)
+
+ self._lbl_none_available = _wrap_label("", _CSS_WARN)
+ al.addWidget(self._lbl_none_available)
+ layout.addWidget(grp_ang)
+
+ # ---- Output path --------------------------------------------------
+ grp_out, ol = _group("Output File")
+ path_row = QHBoxLayout()
+ self._edit_path = QLineEdit()
+ self._edit_path.setReadOnly(True)
+ path_row.addWidget(self._edit_path, 1)
+ btn_browse = QPushButton("Browse…")
+ btn_browse.clicked.connect(self._on_browse)
+ path_row.addWidget(btn_browse)
+ ol.addLayout(path_row)
+ layout.addWidget(grp_out)
+
+ # ---- Buttons -----------------------------------------------------
+ buttons = QDialogButtonBox()
+ self._btn_export = buttons.addButton(
+ "Export", QDialogButtonBox.ButtonRole.AcceptRole)
+ self._btn_export.clicked.connect(self.accept)
+ buttons.addButton("Cancel", QDialogButtonBox.ButtonRole.RejectRole
+ ).clicked.connect(self.reject)
+ layout.addWidget(buttons)
+
+ self._refresh_default_path()
+ self._apply_availability()
+
+ # ---- internals -----------------------------------------------------
+
+ def _current_ch_idx(self) -> int:
+ return self._val_group.checkedId()
+
+ def _apply_availability(self):
+ ch_idx = self._current_ch_idx()
+ n_ok = 0
+ for angle_idx, cb in self._angle_checks.items():
+ ok = self._availability_fn(angle_idx, ch_idx)
+ cb.setEnabled(ok)
+ if ok:
+ n_ok += 1
+ cb.setToolTip("")
+ else:
+ cb.setChecked(False)
+ cb.setToolTip(
+ f"No cached/stored {CH_LABELS[ch_idx]} data for this "
+ "angle yet — view it in the main window (or run "
+ "Batch Compute) first.")
+ self._lbl_none_available.setText(
+ "" if n_ok else "No angle has this value type ready yet.")
+ self._update_accept_enabled()
+
+ def _on_value_type_changed(self, _id: int):
+ self._apply_availability()
+ if not self._path_user_chosen:
+ self._refresh_default_path()
+
+ def _on_select_all_available(self):
+ for cb in self._angle_checks.values():
+ if cb.isEnabled():
+ cb.setChecked(True)
+
+ def _on_select_none(self):
+ for cb in self._angle_checks.values():
+ cb.setChecked(False)
+
+ def _refresh_default_path(self):
+ ch_idx = self._current_ch_idx()
+ name = f"{self._stem}_fused_roi_{CH_NAMES[ch_idx]}.csv"
+ self._edit_path.setText(str(Path(self._out_dir) / name))
+ self._update_accept_enabled()
+
+ def _on_browse(self):
+ path, _ = QFileDialog.getSaveFileName(
+ self, "Export Fused ROI as CSV", self._edit_path.text(),
+ "CSV files (*.csv);;All files (*)")
+ if path:
+ self._edit_path.setText(path)
+ self._path_user_chosen = True
+ self._update_accept_enabled()
+
+ def _update_accept_enabled(self):
+ any_checked = any(cb.isChecked() for cb in self._angle_checks.values())
+ self._btn_export.setEnabled(any_checked and bool(self._edit_path.text()))
+
+ # ---- getters ---------------------------------------------------------
+
+ def get_ch_idx(self) -> int:
+ return self._current_ch_idx()
+
+ def get_selected_angles(self) -> list[int]:
+ return sorted(a for a, cb in self._angle_checks.items() if cb.isChecked())
+
+ def get_output_path(self) -> str:
+ return self._edit_path.text()
+
diff --git a/sras_viewer/main_window.py b/sras_viewer/main_window.py
new file mode 100644
index 0000000..b02d6e7
--- /dev/null
+++ b/sras_viewer/main_window.py
@@ -0,0 +1,1670 @@
+"""The SrasViewerWindow main window and application entry point."""
+
+import sys
+from pathlib import Path
+
+import numpy as np
+from matplotlib.backends.backend_qtagg import NavigationToolbar2QT
+from PyQt6.QtCore import QObject, QSettings, QSignalBlocker, Qt, QThread
+from PyQt6.QtGui import QAction
+from PyQt6.QtWidgets import (
+ QApplication, QCheckBox, QComboBox, QDialog, QFileDialog, QFrame,
+ QHBoxLayout, QLabel, QMainWindow, QProgressDialog, QPushButton,
+ QSizePolicy, QSpinBox, QSplitter, QVBoxLayout, QWidget,
+)
+
+import sras_compute as compute
+from sras_compute import (
+ ManualAngleParams, apply_alignment, build_manual_alignment,
+ load_manual_alignment, save_manual_alignment, sidecar_path,
+)
+from sras_format import (
+ CH3_IDX, CH4_IDX, CH_NAMES, SrasFile, mv_to_adc,
+ _FALLBACK_YMULT_MV, _FALLBACK_YOFF_ADC,
+)
+from sras_workers import (
+ BatchCacheWorker, ComputeWorker, DcPrecomputeWorker, LoadWorker,
+)
+
+from .canvases import ImageCanvas, WaveformCanvas
+from .common import (
+ CH1_DERIVED_MODES, CH_LABELS, CMAPS, VELOCITY_MODE_IDX, _CHANNEL_DISPLAY,
+ _CSS_BUSY, _axes_extent, _CSS_HINT, _CSS_INFO, _CSS_MUTED, _CSS_WARN, _LEFT_PANEL_W,
+ _RIGHT_PANEL_W, Jobs, _combo, _form, _group, _make_dspin, _scroll_panel,
+ _wrap_label,
+)
+from .align_wizard import AlignmentWizard
+from .dialogs import (
+ FftOptionsDialog, FusedRoiExportDialog, RowAverageFftOptionsDialog,
+)
+
+# ---------------------------------------------------------------------------
+# Main window
+# ---------------------------------------------------------------------------
+
+class SrasViewerWindow(QMainWindow):
+ def __init__(self, initial_path: str | None = None):
+ super().__init__()
+ self.setWindowTitle("SRAS Scan Viewer")
+ self.resize(1560, 840)
+ self.setMinimumSize(960, 560)
+ self.setAcceptDrops(True)
+
+ self._sras: SrasFile | None = None
+ self._current_image: np.ndarray | None = None
+ self._current_angle: int = 0
+ self._current_ch: int = 0
+ self._pending_angle: int = 0
+ self._pending_ch: int = 0
+ self._pending_bg_sub: bool = True
+ self._pending_threshold: float = 50.0 # mV
+ self._pending_fft_pad_factor: int = 1
+
+ # Live background jobs, keyed by role — see _run_worker.
+ self._jobs: dict[str, tuple] = {}
+ self._progress_dlgs: dict[str, QProgressDialog] = {}
+
+ # FFT settings (configured via FFT Options dialog, persisted across
+ # sessions). IniFormat: predictable cross-platform and redirectable
+ # in tests.
+ self._settings = QSettings(QSettings.Format.IniFormat,
+ QSettings.Scope.UserScope,
+ "sras-viewer", "sras-viewer")
+ compute.set_fft_backend(str(self._settings.value("fft/backend", "scipy")))
+ try:
+ pad = int(self._settings.value("fft/pad_factor", 1))
+ except (TypeError, ValueError):
+ pad = 1
+ self._fft_pad_factor: int = max(1, min(256, pad)) # 1 = no padding
+ try:
+ row_avg_n = int(self._settings.value("fft/row_avg_n", 3))
+ except (TypeError, ValueError):
+ row_avg_n = 3
+ self._pending_row_avg_n: int = max(1, min(50, row_avg_n))
+
+ # Convert menu: batch DC/FFT compute-and-store (v6 -> v7)
+ self._batch_errors: list[str] = []
+
+ # Display-only settings (colormap, grating) never trigger a
+ # recompute — they're applied to cached data on redraw. DC images
+ # (CH3/CH4) are cheap and precomputed for every angle in the
+ # background right after load. CH1/Velocity FFT images are
+ # computed lazily (with a progress popup) the first time an
+ # angle/threshold combination is viewed — using the cached DC4
+ # image to skip the FFT entirely for masked-out pixels — and
+ # cached per (angle, threshold) so revisiting the same combination
+ # is free. bg-sub/pad are deliberately not part of the key: once an
+ # angle has any FFT image (live or from the file's own stored
+ # cache), it stays displayed regardless of those controls — see
+ # _fft_cache_key.
+ self._dc_cache: dict[tuple[int, int], np.ndarray] = {}
+ self._fft_cache: dict[tuple[int, float], np.ndarray] = {}
+ self._dc_generation: int = 0
+
+ # Angle alignment ("Fusion" menu)
+ self._alignment_result = None
+ self._aligned_cache: dict[tuple, np.ndarray] = {}
+ self._align_wizard: AlignmentWizard | None = None
+
+ self._build_ui()
+
+ if initial_path:
+ self._load_file(initial_path)
+
+ # ------------------------------------------------------------------
+ # Background job plumbing
+ # ------------------------------------------------------------------
+
+ def _run_worker(self, key: str, worker: QObject, *,
+ connect: tuple = (), quit_on: tuple = ("finished",),
+ on_done=None) -> bool:
+ """Move *worker* onto its own QThread and start it. Returns False if
+ a job under *key* is already running.
+
+ Centralises two lifetime hazards that each cost a process abort:
+
+ 1. The job is claimed in self._jobs *before* start() and before
+ anything below that can pump the Qt event loop (a
+ QProgressDialog.show() does on first display). If it weren't, a
+ re-entrant editingFinished could slip past the busy check, start a
+ second thread, and then have the first call's own assignment
+ clobber — and destroy while still running — that second QThread.
+
+ 2. thread.finished fires as the thread winds down but does not
+ guarantee the OS thread has joined. Dropping the last reference to
+ a QThread whose thread is still running logs "QThread: Destroyed
+ while thread is still running" and aborts, so wait() first.
+ """
+ if key in self._jobs:
+ return False
+
+ thread = QThread()
+ self._jobs[key] = (thread, worker, on_done) # claim before anything pumps
+ worker.moveToThread(thread)
+ thread.started.connect(worker.run)
+ for signal_name, slot in connect:
+ getattr(worker, signal_name).connect(slot)
+ for signal_name in quit_on:
+ getattr(worker, signal_name).connect(thread.quit)
+ thread.finished.connect(lambda k=key: self._on_job_finished(k))
+ thread.start()
+ return True
+
+ def _on_job_finished(self, key: str):
+ job = self._jobs.pop(key, None)
+ if job is None:
+ return
+ thread, _worker, on_done = job
+ thread.wait() # join before releasing our last reference
+ if on_done is not None:
+ on_done()
+
+ def _job_running(self, key: str) -> bool:
+ return key in self._jobs
+
+ # ------------------------------------------------------------------
+ # UI construction
+ # ------------------------------------------------------------------
+
+ def _build_ui(self):
+ central = QWidget()
+ self.setCentralWidget(central)
+ root = QHBoxLayout(central)
+ root.setContentsMargins(8, 8, 8, 8)
+ root.setSpacing(8)
+
+ root.addWidget(self._build_left_panel())
+ root.addWidget(self._build_canvases(), stretch=1)
+ root.addWidget(self._build_right_panel())
+
+ self.statusBar().showMessage("Open an .sras file to begin.")
+ self._build_menus()
+
+ def _build_left_panel(self) -> QWidget:
+ panel = QWidget()
+ panel_layout = QVBoxLayout(panel)
+ panel_layout.setContentsMargins(0, 0, 0, 0)
+ panel_layout.setSpacing(8)
+ panel_layout.addWidget(self._build_file_group())
+ panel_layout.addWidget(self._build_info_group())
+ panel_layout.addWidget(self._build_view_group())
+ panel_layout.addWidget(self._build_roi_group())
+ panel_layout.addStretch()
+ return _scroll_panel(panel, _LEFT_PANEL_W)
+
+ def _build_file_group(self) -> QWidget:
+ grp_file, fl = _group("File")
+ self.btn_open = QPushButton("Open .sras…")
+ self.btn_open.clicked.connect(self._on_open)
+ self.lbl_filename = _wrap_label("No file loaded", _CSS_MUTED)
+ fl.addWidget(self.btn_open)
+ fl.addWidget(self.lbl_filename)
+ return grp_file
+
+ def _build_info_group(self) -> QWidget:
+ grp_info, il = _group("Scan Info")
+ il.setSpacing(3)
+ self._info = {}
+ for key in ("Angles", "Rows", "Frames / row", "Samples / frame",
+ "Sample rate", "X start", "Pixel Δx", "Laser freq"):
+ lbl = _wrap_label(f"{key}: —", _CSS_INFO)
+ il.addWidget(lbl)
+ self._info[key] = lbl
+
+ # frame-count / format notes
+ self.lbl_frame_warn = _wrap_label("", _CSS_WARN)
+ il.addWidget(self.lbl_frame_warn)
+
+ # background DC-precompute progress
+ self.lbl_dc_precompute = _wrap_label("", _CSS_BUSY)
+ il.addWidget(self.lbl_dc_precompute)
+ return grp_info
+
+ def _build_view_group(self) -> QWidget:
+ grp_view, vl = _group("View Settings")
+
+ view_form = _form()
+
+ self.spin_angle = QSpinBox()
+ self.spin_angle.setRange(0, 0)
+ self.spin_angle.setEnabled(False)
+ self.spin_angle.setMinimumWidth(64)
+ # valueChanged with keyboard tracking off, not editingFinished: the
+ # latter fires only on Return or focus-out, so stepping the angle (an
+ # arrow click or Up/Down, the ordinary way to walk a scan) changed the
+ # number and left the image behind. Tracking off is what keeps
+ # valueChanged from also firing per keystroke mid-typing, which on a
+ # large scan would launch a compute for every intermediate angle.
+ self.spin_angle.setKeyboardTracking(False)
+ self.spin_angle.valueChanged.connect(self._on_view_changed)
+ self.lbl_angle_deg = QLabel("—")
+ angle_field = QWidget()
+ ar = QHBoxLayout(angle_field)
+ ar.setContentsMargins(0, 0, 0, 0)
+ ar.setSpacing(6)
+ ar.addWidget(self.spin_angle)
+ ar.addWidget(self.lbl_angle_deg)
+ ar.addStretch()
+ view_form.addRow("Angle:", angle_field)
+
+ self.combo_channel = _combo(CH_LABELS, min_chars=12)
+ self.combo_channel.setEnabled(False)
+ self.combo_channel.setSizePolicy(QSizePolicy.Policy.Expanding,
+ QSizePolicy.Policy.Fixed)
+ self.combo_channel.currentIndexChanged.connect(self._on_channel_changed)
+ view_form.addRow("Channel:", self.combo_channel)
+ vl.addLayout(view_form)
+
+ sep = QFrame()
+ sep.setFrameShape(QFrame.Shape.HLine)
+ sep.setStyleSheet("color: #555;")
+ vl.addWidget(sep)
+
+ # DC threshold (for RF / CH1 masking)
+ self.grp_threshold, tl = _group("RF Mask Threshold (CH1 only)")
+ thr_form = _form()
+ self.spin_threshold_mv = _make_dspin(-500.0, 500.0, 3, suffix=" mV",
+ value=50.0, step=0.025)
+ self.spin_threshold_mv.setEnabled(False)
+ self.spin_threshold_mv.editingFinished.connect(self._on_threshold_changed)
+ thr_form.addRow("DC threshold:", self.spin_threshold_mv)
+ tl.addLayout(thr_form)
+ self.lbl_threshold_adc = _wrap_label(
+ f"≈ {mv_to_adc(50.0):.1f} ADC counts", _CSS_MUTED)
+ tl.addWidget(self.lbl_threshold_adc)
+ vl.addWidget(self.grp_threshold)
+
+ # Background subtraction (v4+ files only)
+ self.chk_bg_sub = QCheckBox("Background subtraction (CH1 only)")
+ self.chk_bg_sub.setChecked(True)
+ self.chk_bg_sub.setEnabled(False)
+ self.chk_bg_sub.setToolTip(
+ "Subtract the stored background waveform from each CH1 frame\n"
+ "before computing the FFT (v4+ files only). Applies to angles\n"
+ "not yet computed and to future batch recomputes — it does not\n"
+ "change an image already shown or already stored in the file."
+ )
+ self.chk_bg_sub.toggled.connect(self._on_bg_sub_toggled)
+ vl.addWidget(self.chk_bg_sub)
+
+ # Aligned View (Fusion → Angle Alignment result)
+ self.chk_aligned_view = QCheckBox("Aligned View (Fusion)")
+ self.chk_aligned_view.setChecked(False)
+ self.chk_aligned_view.setEnabled(False)
+ self.chk_aligned_view.setToolTip(
+ "Show the current angle/channel resampled onto the shared,\n"
+ "rotation+translation-aligned canvas from Fusion → Angle\n"
+ "Alignment. Uncheck to see the raw per-angle scan grid."
+ )
+ self.chk_aligned_view.toggled.connect(self._on_aligned_view_toggled)
+ vl.addWidget(self.chk_aligned_view)
+
+ self.btn_export_csv = QPushButton("Export Image as CSV…")
+ self.btn_export_csv.setEnabled(False)
+ self.btn_export_csv.setToolTip(
+ "Save the current CH1 image (one scan row per CSV line).")
+ self.btn_export_csv.clicked.connect(self._on_export_csv)
+ vl.addWidget(self.btn_export_csv)
+ return grp_view
+
+ def _build_roi_group(self) -> QWidget:
+ grp_roi, rl = _group("ROI (Region of Interest)")
+
+ self.btn_draw_roi = QPushButton("Draw ROI")
+ self.btn_draw_roi.setCheckable(True)
+ self.btn_draw_roi.setEnabled(False)
+ self.btn_draw_roi.setToolTip(
+ "Arm next click+drag on the image to draw a new ROI\n"
+ "(replaces any existing one). Click again to cancel.\n"
+ "After drawing, drag inside to move, or grab corners to reshape.\n"
+ "The ROI is persistent across channels / modes / angles."
+ )
+ self.btn_draw_roi.toggled.connect(self._on_draw_roi_toggled)
+ rl.addWidget(self.btn_draw_roi)
+
+ self.btn_clear_roi = QPushButton("Clear ROI")
+ self.btn_clear_roi.setEnabled(False)
+ self.btn_clear_roi.clicked.connect(self._on_clear_roi)
+ rl.addWidget(self.btn_clear_roi)
+
+ self.btn_export_roi = QPushButton("Export ROI as CSV…")
+ self.btn_export_roi.setEnabled(False)
+ self.btn_export_roi.setToolTip(
+ "Save every pixel whose centre lies inside the ROI as CSV.\n"
+ "Columns: row, frame, x_mm, y_mm, value.\n"
+ "Corner coordinates of the quad are written in the file header."
+ )
+ self.btn_export_roi.clicked.connect(self._on_export_roi_csv)
+ rl.addWidget(self.btn_export_roi)
+
+ self.btn_export_fused_roi = QPushButton("Export Fused ROI…")
+ self.btn_export_fused_roi.setEnabled(False)
+ self.btn_export_fused_roi.clicked.connect(self._on_export_fused_roi_csv)
+ rl.addWidget(self.btn_export_fused_roi)
+
+ self.lbl_roi_center = _wrap_label("centroid: —", _CSS_HINT)
+ self.lbl_roi_size = _wrap_label("bbox: —", _CSS_HINT)
+ self.lbl_roi_npix = _wrap_label("pixels inside: —", _CSS_HINT)
+ for lbl in (self.lbl_roi_center, self.lbl_roi_size, self.lbl_roi_npix):
+ rl.addWidget(lbl)
+ return grp_roi
+
+ def _build_canvases(self) -> QWidget:
+ splitter = QSplitter(Qt.Orientation.Vertical)
+ splitter.setChildrenCollapsible(False)
+
+ img_widget = QWidget()
+ img_vl = QVBoxLayout(img_widget)
+ img_vl.setContentsMargins(0, 0, 0, 0)
+ img_vl.setSpacing(4)
+ self.image_canvas = ImageCanvas()
+ self.image_canvas.setMinimumHeight(220)
+ self.image_canvas.pixel_clicked.connect(self._on_pixel_clicked)
+ self.image_canvas.roi_changed.connect(self._update_roi_ui)
+ self.image_canvas.draw_mode_changed.connect(self._on_draw_mode_changed)
+ img_vl.addWidget(NavigationToolbar2QT(self.image_canvas, img_widget))
+ img_vl.addWidget(self.image_canvas)
+ splitter.addWidget(img_widget)
+
+ wave_widget = QWidget()
+ wave_vl = QVBoxLayout(wave_widget)
+ wave_vl.setContentsMargins(0, 0, 0, 0)
+ wave_vl.setSpacing(4)
+ self.lbl_wave_hint = QLabel(
+ "Click a pixel in the image above to inspect its waveform.")
+ self.lbl_wave_hint.setAlignment(Qt.AlignmentFlag.AlignCenter)
+ self.lbl_wave_hint.setStyleSheet(_CSS_MUTED)
+ self.wave_canvas = WaveformCanvas()
+ self.wave_canvas.setMinimumHeight(150)
+ wave_vl.addWidget(self.lbl_wave_hint)
+ wave_vl.addWidget(self.wave_canvas)
+ splitter.addWidget(wave_widget)
+
+ splitter.setStretchFactor(0, 3)
+ splitter.setStretchFactor(1, 1)
+ splitter.setSizes([580, 250])
+ return splitter
+
+ def _build_right_panel(self) -> QWidget:
+ # Velocity settings (visible only in velocity mode)
+ self.grp_velocity, vel_l = _group("Velocity Settings (CH1 only)")
+ vel_form = _form()
+ self.spin_grating_um = _make_dspin(0.1, 1000.0, 2, suffix=" µm",
+ value=25, step=0.5)
+ self.spin_grating_um.setEnabled(False)
+ self.spin_grating_um.editingFinished.connect(self._on_grating_changed)
+ vel_form.addRow("Grating size:", self.spin_grating_um)
+ vel_l.addLayout(vel_form)
+ vel_l.addWidget(_wrap_label("v (m/s) = freq (MHz) × grating (µm)",
+ "font-size: 10px; color: #888;"))
+ self.grp_velocity.setVisible(False)
+
+ grp_display, dl = _group("Display Options")
+
+ cmap_form = _form()
+ self.combo_cmap = QComboBox()
+ self.combo_cmap.addItems(CMAPS)
+ self.combo_cmap.setCurrentText("gray")
+ self.combo_cmap.setEnabled(False)
+ self.combo_cmap.setSizePolicy(QSizePolicy.Policy.Expanding,
+ QSizePolicy.Policy.Fixed)
+ self.combo_cmap.currentIndexChanged.connect(self._on_cmap_changed)
+ cmap_form.addRow("Colormap:", self.combo_cmap)
+ dl.addLayout(cmap_form)
+
+ self.chk_auto = QCheckBox("Auto-scale colormap")
+ self.chk_auto.setChecked(True)
+ self.chk_auto.toggled.connect(self._on_autoscale_toggled)
+ dl.addWidget(self.chk_auto)
+
+ range_form = _form()
+ for label, attr in (("min:", "spin_vmin"), ("max:", "spin_vmax")):
+ spin = _make_dspin(-1e9, 1e9, 4)
+ spin.setEnabled(False)
+ spin.editingFinished.connect(self._on_manual_range_changed)
+ setattr(self, attr, spin)
+ range_form.addRow(label, spin)
+ dl.addLayout(range_form)
+
+ right_panel = QWidget()
+ layout = QVBoxLayout(right_panel)
+ layout.setContentsMargins(0, 0, 0, 0)
+ layout.setSpacing(8)
+ layout.addWidget(self.grp_velocity)
+ layout.addWidget(grp_display)
+ layout.addStretch()
+ return _scroll_panel(right_panel, _RIGHT_PANEL_W)
+
+ def _build_menus(self):
+ menubar = self.menuBar()
+
+ fft_menu = menubar.addMenu("&FFT")
+ fft_act = QAction("FFT &Options…", self)
+ fft_act.setStatusTip("Configure FFT backend and zero-padding")
+ fft_act.triggered.connect(self._on_fft_options)
+ fft_menu.addAction(fft_act)
+
+ fusion_menu = menubar.addMenu("&Fusion")
+ self._wizard_act = QAction("Alignment &Wizard…", self)
+ self._wizard_act.setStatusTip(
+ "Align the angles, crop to a region of interest, and save the "
+ "aligned data as a new .sras file. Requires >1 angle.")
+ self._wizard_act.setEnabled(False)
+ self._wizard_act.triggered.connect(self._on_alignment_wizard)
+ fusion_menu.addAction(self._wizard_act)
+
+ convert_menu = menubar.addMenu("&Convert")
+ self._batch_dc_act = QAction("Batch Compute DC and &Store…", self)
+ self._batch_dc_act.setStatusTip(
+ "Select .sras files and compute+store DC images (CH3/CH4 mean) "
+ "for every angle, converting v6 files to v7 in place.")
+ self._batch_dc_act.triggered.connect(lambda: self._on_batch_compute("dc"))
+ convert_menu.addAction(self._batch_dc_act)
+
+ self._batch_fft_act = QAction("Batch Compute FFT and Sto&re…", self)
+ self._batch_fft_act.setStatusTip(
+ "Select .sras files and compute+store FFT peak-frequency images "
+ "for every angle, converting v6 files to v7 in place. Stored "
+ "images are natural-resolution (pad 1); padded views compute live.")
+ self._batch_fft_act.triggered.connect(lambda: self._on_batch_compute("fft"))
+ convert_menu.addAction(self._batch_fft_act)
+
+ self._batch_fft_rowavg_act = QAction(
+ "Batch Compute Row-A&veraged FFT and Store…", self)
+ self._batch_fft_rowavg_act.setStatusTip(
+ "Select .sras files and compute+store a same-row, distance-weighted "
+ "smoothed FFT peak-frequency image for every angle — improves SNR "
+ "on noisy regions. DC images and raw waveform data are never "
+ "touched; the raw (unsmoothed) FFT is always a recompute away. "
+ "Converts v6 files to v7 in place.")
+ self._batch_fft_rowavg_act.triggered.connect(self._on_batch_compute_row_avg)
+ convert_menu.addAction(self._batch_fft_rowavg_act)
+
+ # ------------------------------------------------------------------
+ # Drag-and-drop
+ # ------------------------------------------------------------------
+
+ def dragEnterEvent(self, event):
+ urls = event.mimeData().urls()
+ if urls and urls[0].toLocalFile().lower().endswith(".sras"):
+ event.acceptProposedAction()
+
+ def dropEvent(self, event):
+ self._load_file(event.mimeData().urls()[0].toLocalFile())
+
+ # ------------------------------------------------------------------
+ # File loading
+ # ------------------------------------------------------------------
+
+ def _on_open(self):
+ path, _ = QFileDialog.getOpenFileName(
+ self, "Open SRAS File", "", "SRAS Files (*.sras);;All Files (*)")
+ if path:
+ self._load_file(path)
+
+ def _load_file(self, path: str):
+ started = self._run_worker(
+ Jobs.LOAD, LoadWorker(path),
+ connect=(
+ ("finished", self._on_load_done),
+ ("error", lambda msg: self.statusBar().showMessage(f"Error: {msg}")),
+ ),
+ )
+ if not started:
+ return
+ self.btn_open.setEnabled(False)
+ self.statusBar().showMessage(f"Loading {Path(path).name}…")
+ self._show_progress("main", f"Loading {Path(path).name}…")
+
+ def _on_load_done(self, sras):
+ self._close_progress("main")
+ self.btn_open.setEnabled(True)
+ if sras is None:
+ return
+ self._sras = sras
+ self._current_image = None
+
+ # A manual-alignment dialog bound to the previous file must not
+ # survive a reload — its per-angle state (and the sras it was
+ # constructed against) no longer matches the new file's geometry.
+ if self._align_wizard is not None:
+ self._align_wizard.close()
+ self._align_wizard = None
+
+ # Caches (and any in-flight DC precompute) belong to the previous
+ # file's geometry — discard and start fresh. Bumping the generation
+ # counters makes any still-running worker's result get dropped when
+ # it lands.
+ self._dc_cache = {}
+ self._fft_cache = {}
+ self._dc_generation += 1
+ self.lbl_dc_precompute.setText("")
+
+ self._apply_alignment_result(None, view_checked=False)
+
+ # Silently restore a previously-saved manual alignment, if any, so
+ # the work survives closing and reopening the file.
+ sidecar = load_manual_alignment(sras)
+ if sidecar is not None:
+ try:
+ self._apply_alignment_result(
+ build_manual_alignment(
+ sras, sidecar.ref_angle_idx, sidecar.dc_threshold_mv,
+ sidecar.per_angle),
+ view_checked=True)
+ self.statusBar().showMessage(
+ f"Restored saved manual alignment from "
+ f"{sidecar_path(sras.path).name}")
+ except Exception as exc:
+ # A corrupt/foreign sidecar or a rescan that shrank n_angles
+ # below ref_angle_idx must not block opening the .sras file.
+ self.statusBar().showMessage(
+ f"Could not restore saved alignment: {exc}")
+
+ # A ROI from the previous file no longer matches the new scan's
+ # geometry, so discard it on every load.
+ self.image_canvas.clear_roi()
+
+ self.lbl_filename.setText(sras.path.name)
+
+ with QSignalBlocker(self.spin_angle):
+ self.spin_angle.setRange(0, max(0, sras.n_angles - 1))
+ self.spin_angle.setValue(0)
+
+ # DC channels are cheap and give an instant, fluid overview of a
+ # scan; CH1/Velocity require an FFT per pixel that can take minutes
+ # on a large scan, so don't default to it.
+ with QSignalBlocker(self.combo_channel):
+ self.combo_channel.setCurrentIndex(CH4_IDX)
+
+ self._update_controls_enabled(True)
+ self._on_threshold_changed() # refresh ADC label with file calibration
+ self._on_view_changed()
+ self._start_dc_precompute()
+
+ # ------------------------------------------------------------------
+ # Scan info panel
+ # ------------------------------------------------------------------
+
+ def _update_scan_info_labels(self):
+ s = self._sras
+ if s is None:
+ return
+ a = self.spin_angle.value()
+ for key, text in (
+ ("Angles", f"{s.n_angles}"),
+ ("Rows", f"{s.n_rows[a]}"),
+ ("Frames / row", f"{s.n_frames[a]}"),
+ ("Samples / frame", f"{s.samples_per_frame}"),
+ ("Sample rate", f"{s.sample_rate_hz / 1e9:.4g} GS/s"),
+ ("X start", f"{s.x_start_mm[a]:.4g} mm"),
+ ("Pixel Δx", f"{s.pixel_x_mm * 1e3:.3g} µm"),
+ ("Laser freq", f"{s.laser_freq_hz / 1e3:.4g} kHz"),
+ ):
+ self._info[key].setText(f"{key}: {text}")
+
+ notes = []
+ if s.frame_count_mismatch:
+ notes.append(f"! Header n_frames={s.n_frames_header}, "
+ f"actual={s.n_frames[a]} (scanner bug — corrected)")
+ if s.scan_aborted:
+ notes.append(f"! Scan aborted: {s.n_angles}/{s.n_angles_declared} "
+ "angles complete")
+ if s.background is not None:
+ notes.append(f"Background waveform: {len(s.background)} samples")
+ if s.version in (6, 7):
+ notes.append("v6/v7 format: rows / frames / x_start are per-angle")
+
+ n_dc = sum(1 for x in s.precomputed_dc4_mv if x is not None)
+ n_fft = sum(1 for x in s.precomputed_freq_mhz if x is not None)
+ if n_dc or n_fft:
+ bg_note = " (bg-sub)" if s.precomputed_bg_sub else " (no bg-sub)"
+ avg_note = (f", row-averaged n={s.precomputed_row_avg_n}"
+ if s.precomputed_row_avg_n else "")
+ pad_note = (f", pad {s.precomputed_pad_factor}x"
+ if s.precomputed_pad_factor > 1 else "")
+ notes.append(
+ f"Cached images: DC {n_dc}/{s.n_angles} angles, "
+ f"FFT {n_fft}/{s.n_angles} angles{bg_note if n_fft else ''}"
+ f"{avg_note if n_fft else ''}{pad_note if n_fft else ''} "
+ "— display is instant for cached angles")
+ notes += self._cache_mismatch_notes()
+ elif s.version == 7:
+ notes.append("v7 format: no cache blocks stored yet")
+ self.lbl_frame_warn.setText("\n".join(notes))
+
+ def _cache_mismatch_notes(self) -> list[str]:
+ """Informational only: whether the file's stored FFT cache was
+ computed under different bg-sub/pad settings than these controls
+ currently say. The display always shows the stored image as-is
+ regardless (see _stored_fft_image) — these controls only affect a
+ future live compute for an angle with nothing cached yet, or an
+ explicit batch recompute, never what's already on screen.
+
+ row_avg_n is compared against the file's own recorded value (a
+ self-match), so it never contributes a reason here — there's no
+ live control for it to diverge from, and the "Cached images" line
+ above already reports it.
+
+ Asks compute for the reasons rather than restating the accept rule,
+ so a new provenance field can only be added in one place.
+ """
+ s = self._sras
+ if s is None or all(x is None for x in s.precomputed_freq_mhz):
+ return []
+
+ reasons = compute.cache_mismatch_reasons(
+ s, n_fft=self._current_n_fft(),
+ apply_bg_sub=self.chk_bg_sub.isChecked(),
+ row_avg_n=s.precomputed_row_avg_n)
+ if not reasons:
+ return []
+ return ["Note: current bg-sub/pad controls differ from the stored "
+ "cache — " + "; ".join(reasons) + ". Shown as stored; use "
+ "Batch Compute to recompute with these settings."]
+
+ # ------------------------------------------------------------------
+ # Controls
+ # ------------------------------------------------------------------
+
+ def _update_controls_enabled(self, enabled: bool):
+ s = self._sras
+ has_file = enabled and s is not None
+ ch_idx = self.combo_channel.currentIndex()
+ is_ch1 = enabled and ch_idx in CH1_DERIVED_MODES
+ is_vel = enabled and ch_idx == VELOCITY_MODE_IDX
+
+ self.spin_angle.setEnabled(has_file and s.n_angles > 1)
+ self.combo_channel.setEnabled(enabled)
+ self.combo_cmap.setEnabled(enabled)
+ self.chk_auto.setEnabled(enabled)
+ manual = enabled and not self.chk_auto.isChecked()
+ self.spin_vmin.setEnabled(manual)
+ self.spin_vmax.setEnabled(manual)
+
+ # Threshold and bg-sub apply to all CH1 modes
+ self.spin_threshold_mv.setEnabled(is_ch1)
+ self.chk_bg_sub.setEnabled(has_file and s.background is not None and is_ch1)
+ self.spin_grating_um.setEnabled(is_vel)
+ self.grp_velocity.setVisible(is_vel)
+
+ self.btn_export_csv.setEnabled(is_ch1 and self._current_image is not None)
+ # ROI: always usable once a file is loaded (independent of channel)
+ self.btn_draw_roi.setEnabled(has_file)
+
+ # Batch Convert actions pick their own files, independent of
+ # whatever's currently open — only gated on no batch already running.
+ can_batch = not self._job_running(Jobs.BATCH)
+ self._batch_dc_act.setEnabled(can_batch)
+ self._batch_fft_act.setEnabled(can_batch)
+ self._batch_fft_rowavg_act.setEnabled(can_batch)
+
+ self._wizard_act.setEnabled(
+ has_file and s.n_angles > 1 and self._align_wizard is None)
+ self.chk_aligned_view.setEnabled(enabled and self._alignment_result is not None)
+ self._update_roi_ui()
+
+ def _on_channel_changed(self):
+ self._update_controls_enabled(self._sras is not None)
+ self._on_view_changed()
+
+ def _on_bg_sub_toggled(self):
+ # bg-sub no longer gates the display: it only affects a future live
+ # compute for an angle with nothing cached yet, or an explicit batch
+ # recompute — never what's already shown. Just keep the info panel's
+ # divergence note current.
+ if self._is_fft_mode():
+ self._update_scan_info_labels()
+
+ def _on_grating_changed(self):
+ # Grating is a pure post-multiply on the cached frequency image —
+ # never needs a recompute.
+ if self._sras is not None and self.combo_channel.currentIndex() == VELOCITY_MODE_IDX:
+ self._refresh_display()
+
+ def _on_threshold_changed(self):
+ mv = self.spin_threshold_mv.value()
+ cal = (self._sras.cal(CH4_IDX) if self._sras is not None
+ else (_FALLBACK_YMULT_MV, _FALLBACK_YOFF_ADC, 0.0))
+ self.lbl_threshold_adc.setText(f"≈ {mv_to_adc(mv, *cal):.1f} ADC counts")
+ # Threshold decides which pixels get an FFT at all, so changing it is
+ # a genuine cache-key change — but the recompute reuses the cached DC4
+ # image to skip masked-out pixels.
+ if self._is_fft_mode():
+ self._refresh_display()
+
+ def _on_autoscale_toggled(self, checked: bool):
+ manual = not checked
+ self.spin_vmin.setEnabled(manual and self._sras is not None)
+ self.spin_vmax.setEnabled(manual and self._sras is not None)
+ if self._sras is not None and self._current_image is not None:
+ self._redraw_image(self._current_image)
+
+ def _on_manual_range_changed(self):
+ if not self.chk_auto.isChecked() and self._current_image is not None:
+ self._redraw_image(self._current_image)
+
+ def _on_cmap_changed(self):
+ # Colormap is purely how the existing image is rendered.
+ if self._current_image is not None:
+ self._redraw_image(self._current_image)
+
+ def _on_view_changed(self):
+ if self._sras is None:
+ return
+ idx = self.spin_angle.value()
+ self.lbl_angle_deg.setText(f"({self._sras.angles_deg[idx]:.1f}°)")
+ self._update_scan_info_labels()
+ self._refresh_display()
+
+ def _on_aligned_view_toggled(self, checked: bool):
+ if self._current_image is not None:
+ self._redraw_image(self._current_image)
+
+ # ------------------------------------------------------------------
+ # CSV export
+ # ------------------------------------------------------------------
+
+ def _on_export_csv(self):
+ if self._current_image is None or self._sras is None:
+ return
+ default_name = (f"{self._sras.path.stem}_angle{self._current_angle}"
+ f"_{CH_NAMES[self._current_ch]}.csv")
+ path, _ = QFileDialog.getSaveFileName(
+ self, "Export Image as CSV",
+ str(self._sras.path.parent / default_name),
+ "CSV files (*.csv);;All files (*)")
+ if not path:
+ return
+ np.savetxt(path, self._current_image, delimiter=",", fmt="%.6g")
+ self.statusBar().showMessage(f"Exported {Path(path).name}")
+
+ def _on_export_roi_csv(self):
+ if self._current_image is None or self._sras is None:
+ return
+ roi = self.image_canvas.get_roi()
+ if roi is None:
+ self.statusBar().showMessage("No ROI — draw one first")
+ return
+ s = self._sras
+ x_axis = s.x_axis_mm(self._current_angle)
+ y_axis = s.y_positions_mm(self._current_angle)
+ mask = roi.mask_for_grid(x_axis, y_axis)
+ if not mask.any():
+ self.statusBar().showMessage("ROI does not overlap any pixel")
+ return
+ img = self._current_image
+ if img.shape != mask.shape:
+ self.statusBar().showMessage(
+ f"ROI shape {mask.shape} does not match image {img.shape}")
+ return
+
+ X, Y = np.meshgrid(np.asarray(x_axis, dtype=np.float64),
+ np.asarray(y_axis, dtype=np.float64))
+ rows_idx, frames_idx = np.where(mask)
+ n_pix = int(mask.sum())
+
+ ch_name = CH_NAMES[self._current_ch]
+ angle = self._current_angle
+ default_name = f"{s.path.stem}_angle{angle}_{ch_name}_ROI.csv"
+ path, _ = QFileDialog.getSaveFileName(
+ self, "Export ROI as CSV",
+ str(s.path.parent / default_name),
+ "CSV files (*.csv);;All files (*)")
+ if not path:
+ return
+
+ corners_str = " ".join(f"({p[0]:.6g},{p[1]:.6g})" for p in roi.corners())
+ header = (
+ f"# ROI quad corners (BL BR TR TL) mm: {corners_str}\n"
+ f"# source: {s.path.name}, channel={ch_name}, "
+ f"angle_idx={angle}, angle_deg={s.angles_deg[angle]:.4g}\n"
+ f"# n_pixels={n_pix}\n"
+ "row,frame,x_mm,y_mm,value"
+ )
+ data = np.column_stack([
+ rows_idx.astype(np.int64), frames_idx.astype(np.int64),
+ X[mask], Y[mask], img[mask].astype(np.float64),
+ ])
+ # integer columns first, floats after — use a per-column format list
+ np.savetxt(path, data, delimiter=",",
+ fmt=["%d", "%d", "%.6g", "%.6g", "%.6g"],
+ header=header, comments="")
+ self.statusBar().showMessage(
+ f"Exported ROI ({n_pix} pixels) to {Path(path).name}")
+
+ def _on_export_fused_roi_csv(self):
+ s = self._sras
+ if s is None:
+ return
+ roi = self.image_canvas.get_roi()
+ if roi is None or not self._fused_grid_ready():
+ return # button is disabled in these states; defensive no-op
+
+ angles = [(a, float(s.angles_deg[a])) for a in range(s.n_angles)]
+ grid_note = (
+ "Using the live Fusion alignment canvas."
+ if self._alignment_result is not None else
+ "Using the file's shared raw grid (no live alignment result — "
+ "this file's angles already share one grid).")
+ dlg = FusedRoiExportDialog(
+ self, angles=angles,
+ availability_fn=lambda a, c: self._cached_value_image(a, c) is not None,
+ default_ch_idx=self.combo_channel.currentIndex(),
+ out_dir=str(s.path.parent), stem=s.path.stem, grid_note=grid_note)
+ if dlg.exec() != QDialog.DialogCode.Accepted:
+ return
+ self._write_fused_roi_csv(
+ roi, dlg.get_ch_idx(), dlg.get_selected_angles(), dlg.get_output_path())
+
+ def _write_fused_roi_csv(self, roi, ch_idx: int, angle_idxs: list[int],
+ out_path: str):
+ """Mask once on the fused grid, then one value column per selected
+ angle. No row/frame index columns: once angles are fused onto one
+ shared canvas there is no single meaningful raw (row, frame) per
+ output pixel, unlike the single-angle _on_export_roi_csv above."""
+ s = self._sras
+ if s is None or not angle_idxs or not out_path:
+ return
+ x_axis, y_axis = self._fused_export_axes()
+ mask = roi.mask_for_grid(x_axis, y_axis)
+ if not mask.any():
+ self.statusBar().showMessage(
+ "ROI does not overlap any pixel on the fused grid")
+ return
+
+ columns: list[np.ndarray] = []
+ used: list[int] = []
+ skipped: list[int] = []
+ for a in angle_idxs:
+ img = self._fused_value_image(a, ch_idx)
+ if img is None or img.shape != mask.shape:
+ # Defensive only: nothing else can mutate the caches between
+ # dialog-accept and this synchronous call in a single-
+ # threaded GUI callback, so this should never trigger — skip
+ # the angle rather than abort the whole export.
+ skipped.append(a)
+ continue
+ columns.append(img[mask].astype(np.float64))
+ used.append(a)
+
+ if not columns:
+ self.statusBar().showMessage(
+ "Nothing to export — no selected angle had data")
+ return
+
+ X, Y = np.meshgrid(np.asarray(x_axis, dtype=np.float64),
+ np.asarray(y_axis, dtype=np.float64))
+ data = np.column_stack([X[mask], Y[mask], *columns])
+
+ ch_name = CH_NAMES[ch_idx]
+ corners_str = " ".join(f"({p[0]:.6g},{p[1]:.6g})" for p in roi.corners())
+ angle_list_str = ", ".join(f"{s.angles_deg[a]:.4g}°" for a in used)
+ header_cols = ["x_mm", "y_mm"] + [f"v_{s.angles_deg[a]:.4g}deg" for a in used]
+ header = (
+ f"# ROI quad corners (BL BR TR TL) mm: {corners_str}\n"
+ f"# source: {s.path.name}, value={ch_name}, "
+ f"grid={'aligned canvas' if self._alignment_result is not None else 'shared raw grid'}\n"
+ f"# angles (deg): {angle_list_str}\n"
+ f"# n_pixels={int(mask.sum())}\n"
+ + ",".join(header_cols)
+ )
+ np.savetxt(out_path, data, delimiter=",",
+ fmt=["%.6g"] * data.shape[1], header=header, comments="")
+
+ note = (f" ({len(skipped)} angle(s) skipped — no longer available)"
+ if skipped else "")
+ self.statusBar().showMessage(
+ f"Exported fused ROI ({int(mask.sum())} pixels, {len(used)} "
+ f"angle(s)) to {Path(out_path).name}{note}")
+
+ # ------------------------------------------------------------------
+ # ROI
+ # ------------------------------------------------------------------
+
+ def _on_draw_roi_toggled(self, checked: bool):
+ if checked:
+ self.image_canvas.start_drawing()
+ self.statusBar().showMessage(
+ "Click and drag on the image to draw a new rectangle.")
+ else:
+ self.image_canvas.cancel_drawing()
+
+ def _on_draw_mode_changed(self, active: bool):
+ # Keep the toggle button's visual state in sync with the canvas.
+ with QSignalBlocker(self.btn_draw_roi):
+ self.btn_draw_roi.setChecked(active)
+
+ def _on_clear_roi(self):
+ self.image_canvas.clear_roi()
+ self.statusBar().showMessage("ROI cleared")
+
+ def _update_roi_ui(self):
+ roi = self.image_canvas.get_roi()
+ if roi is None:
+ self.lbl_roi_center.setText("centroid: —")
+ self.lbl_roi_size.setText("bbox: —")
+ self.lbl_roi_npix.setText("pixels inside: —")
+ self.btn_clear_roi.setEnabled(False)
+ self.btn_export_roi.setEnabled(False)
+ self._update_fused_export_enabled()
+ return
+
+ cen = roi.centroid()
+ bbox = roi.bbox_size()
+ self.lbl_roi_center.setText(f"centroid: ({cen[0]:.3f}, {cen[1]:.3f}) mm")
+ self.lbl_roi_size.setText(f"bbox: {bbox[0]:.3f} × {bbox[1]:.3f} mm")
+
+ npix = 0
+ if self._sras is not None:
+ try:
+ # Deliberately always the raw per-angle grid, even when
+ # Aligned View is on: _on_export_roi_csv also exports on the
+ # raw grid (never synthetically-resampled pixels), so this
+ # readout must match what Export ROI actually writes.
+ mask = roi.mask_for_grid(
+ self._sras.x_axis_mm(self._current_angle),
+ self._sras.y_positions_mm(self._current_angle))
+ npix = int(mask.sum())
+ except Exception:
+ npix = 0
+ self.lbl_roi_npix.setText(f"pixels inside: {npix}")
+ self.btn_clear_roi.setEnabled(True)
+ self.btn_export_roi.setEnabled(self._current_image is not None and npix > 0)
+ self._update_fused_export_enabled()
+
+ def _update_fused_export_enabled(self):
+ s = self._sras
+ roi = self.image_canvas.get_roi() if s is not None else None
+ if s is None:
+ ready, reason = False, "Open a file first."
+ elif roi is None:
+ ready, reason = False, "Draw an ROI first."
+ elif not self._fused_grid_ready():
+ ready, reason = False, (
+ "Angles need a common grid to fuse: run Fusion → Alignment "
+ "Wizard, or open a file the wizard already exported.")
+ else:
+ ready, reason = True, (
+ "Export the ROI as one CSV with a column per selected "
+ "angle's value (choose which value and which angles in "
+ "the dialog).")
+ self.btn_export_fused_roi.setEnabled(ready)
+ self.btn_export_fused_roi.setToolTip(reason)
+
+ # ------------------------------------------------------------------
+ # Display
+ # ------------------------------------------------------------------
+
+ def _current_n_fft(self) -> int | None:
+ if self._fft_pad_factor <= 1 or self._sras is None:
+ return None
+ return self._sras.samples_per_frame * self._fft_pad_factor
+
+ def _is_fft_mode(self) -> bool:
+ """Is the selected channel an FFT-derived (CH1/Velocity) mode?"""
+ return (self._sras is not None
+ and self.combo_channel.currentIndex() in CH1_DERIVED_MODES)
+
+ def _scale_for_display(self, freq_mhz: np.ndarray, ch_idx: int) -> np.ndarray:
+ """Velocity is a pure post-multiply of the (already DC-masked)
+ cached frequency image — never worth a recompute on its own."""
+ if ch_idx == VELOCITY_MODE_IDX:
+ return freq_mhz * self.spin_grating_um.value()
+ return freq_mhz
+
+ def _fft_cache_key(self, angle_idx: int) -> tuple:
+ """Keyed by angle and DC threshold only. Once any FFT image exists
+ for an angle this session — live-computed or pulled from the file's
+ own stored cache — it stays the displayed image for that angle
+ regardless of later bg-sub/pad toggles; those only affect a future
+ live compute for an angle with nothing cached yet, or an explicit
+ batch recompute (see _stored_fft_image). Threshold stays in the key
+ because re-masking against it is free and meant to stay interactive
+ (see _on_threshold_changed)."""
+ return (angle_idx, self.spin_threshold_mv.value())
+
+ def _aligned_cache_key(self, angle_idx: int, ch_idx: int) -> tuple:
+ """Mirrors _fft_cache's key granularity so a stale aligned image is
+ never shown after bg_sub/threshold/pad/grating changes."""
+ if ch_idx in CH1_DERIVED_MODES:
+ return (*self._fft_cache_key(angle_idx), ch_idx,
+ self.spin_grating_um.value() if ch_idx == VELOCITY_MODE_IDX else None)
+ return (angle_idx, ch_idx)
+
+ def _aligned_canvas_axes(self) -> tuple[np.ndarray, np.ndarray]:
+ r = self._alignment_result
+ n_rows, n_cols = r.canvas_shape
+ return (r.canvas_origin_mm[0] + np.arange(n_cols) * r.canvas_dx_mm,
+ r.canvas_origin_mm[1] + np.arange(n_rows) * r.canvas_dy_mm)
+
+ def _get_aligned_display_image(self, raw_img: np.ndarray, angle_idx: int,
+ ch_idx: int) -> np.ndarray:
+ key = self._aligned_cache_key(angle_idx, ch_idx)
+ cached = self._aligned_cache.get(key)
+ if cached is None:
+ cached = apply_alignment(self._alignment_result, angle_idx, raw_img)
+ self._aligned_cache[key] = cached
+ return cached
+
+ # ------------------------------------------------------------------
+ # Fused ROI export (Export Fused ROI…)
+ # ------------------------------------------------------------------
+
+ def _fused_grid_ready(self) -> bool:
+ """Is there a common (x, y) grid to fuse angles onto right now?
+
+ Either a live alignment result exists (angles are reconciled onto its
+ canvas via apply_alignment), or — with no alignment run this session
+ — the open file's raw per-angle grids already coincide, which is true
+ for a .sras file the Alignment Wizard itself previously exported (see
+ scan_format.md, "Files written by the viewer's Alignment Wizard")."""
+ return (self._sras is not None
+ and (self._alignment_result is not None
+ or self._sras.angles_share_raw_grid()))
+
+ def _fused_export_axes(self) -> tuple[np.ndarray, np.ndarray]:
+ """(x_axis, y_axis) a fused ROI export masks and labels against: the
+ alignment canvas when a live result exists — the more current,
+ deliberate source of truth even if the raw grids happen to already
+ match too — else the grid every angle already shares."""
+ if self._alignment_result is not None:
+ return self._aligned_canvas_axes()
+ s = self._sras
+ return s.x_axis_mm(0), s.y_positions_mm(0)
+
+ def _fused_value_image(self, angle_idx: int, ch_idx: int) -> np.ndarray | None:
+ """One angle's image on the fused grid for (angle, channel), scaled
+ exactly like the on-screen display — so a Velocity export and the
+ on-screen Velocity image share the same scaling code. None if
+ nothing is cached/stored for this angle/channel without a compute."""
+ raw = self._cached_value_image(angle_idx, ch_idx)
+ if raw is None:
+ return None
+ img = (self._scale_for_display(raw, ch_idx)
+ if ch_idx in CH1_DERIVED_MODES else raw)
+ if self._alignment_result is not None:
+ return self._get_aligned_display_image(img, angle_idx, ch_idx)
+ return img # angles already share the raw grid — no resample needed
+
+ def _stored_fft_image(self, angle_idx: int) -> np.ndarray | None:
+ """The open file's own stored peak-frequency image for this angle,
+ masked and ready to display, or None if the file has nothing stored
+ for it.
+
+ Asks for the image at the settings it was actually computed under
+ (sras.precomputed_bg_sub / precomputed_pad_factor /
+ precomputed_row_avg_n) rather than the window's live bg-sub/pad
+ controls, so a stored image is always shown once present — those
+ controls never gate whether it's used, only what a *future* compute
+ produces. See _cache_mismatch_notes for the informational (non-
+ blocking) note when the live controls diverge from what's shown.
+ Only the DC threshold is taken live: re-masking a stored image
+ against it is free, unlike bg-sub/pad/row-averaging which are baked
+ irreversibly into the stored numbers.
+
+ allow_dc_recompute=False keeps this off the I/O path: if the mask
+ would mean reading a whole CH4 channel, this declines and the caller
+ falls through to the background worker, which reaches the same stored
+ image via compute_rf_image and pays for the mask off the GUI thread.
+ """
+ s = self._sras
+ n_fft = (s.samples_per_frame * s.precomputed_pad_factor
+ if s.precomputed_pad_factor > 1 else None)
+ return compute.cached_rf_image(
+ s, angle_idx,
+ dc_threshold_mv=self.spin_threshold_mv.value(),
+ apply_bg_sub=s.precomputed_bg_sub,
+ n_fft=n_fft,
+ row_avg_n=s.precomputed_row_avg_n,
+ dc4_mv=self._dc_cache.get((angle_idx, CH4_IDX)),
+ allow_dc_recompute=False)
+
+ def _cached_value_image(self, angle_idx: int, ch_idx: int) -> np.ndarray | None:
+ """The already-available (no compute) image for (angle, channel):
+ this window's own in-session dicts first, then the file's stored
+ v5/v7 cache blocks — the same two tiers, in the same cost order, that
+ _refresh_display and Export Fused ROI both need, so both share this
+ one lookup rather than drifting apart. None means genuinely nothing
+ is cached/stored, i.e. only a real compute could produce it — and
+ neither caller here is allowed to trigger one: _refresh_display falls
+ back to _start_compute() itself, and a fused export simply treats the
+ angle as unavailable.
+
+ For CH1/Velocity this is the unscaled, already-masked peak-frequency
+ (MHz) image; callers displaying/exporting Velocity must still run it
+ through _scale_for_display.
+ """
+ s = self._sras
+ if s is None:
+ return None
+ if ch_idx in CH1_DERIVED_MODES:
+ key = self._fft_cache_key(angle_idx)
+ raw = self._fft_cache.get(key)
+ if raw is None:
+ raw = self._stored_fft_image(angle_idx)
+ if raw is not None:
+ # Masking the stored image is cheap but not free; keep the
+ # result so revisiting this angle costs nothing at all.
+ self._fft_cache[key] = raw
+ return raw
+ # A stored DC image needs no post-processing, so the file's own
+ # parsed array is served directly — as the compute path already
+ # does, _current_image is treated as read-only by every consumer.
+ cached = self._dc_cache.get((angle_idx, ch_idx))
+ if cached is None:
+ cached = s.cached_dc_mv(angle_idx, ch_idx)
+ return cached
+
+ def _refresh_display(self):
+ """Show the image for the current angle/channel/threshold, using
+ cached data whenever possible and only falling back to a background
+ compute (with progress popup) when genuinely nothing is cached yet.
+
+ Two caches are consulted, in cost order: this window's own in-session
+ dicts, then the file's stored v5/v7 cache blocks. The second is what
+ makes a batch-computed file worth having — without it every angle
+ change queued a worker and a progress popup for an image already on
+ disk, which is exactly the cost the batch was run to avoid.
+ """
+ if self._sras is None:
+ return
+ angle_idx = self.spin_angle.value()
+ ch_idx = self.combo_channel.currentIndex()
+
+ raw = self._cached_value_image(angle_idx, ch_idx)
+ if raw is not None:
+ img = (self._scale_for_display(raw, ch_idx)
+ if ch_idx in CH1_DERIVED_MODES else raw)
+ self._show_image_now(img, angle_idx, ch_idx)
+ return
+
+ # Nothing cached for these settings — need a real compute. Changing
+ # the DC threshold changes *which* pixels get an FFT at all, so it
+ # can't be satisfied from the cache — but with the DC map already
+ # known, the recompute skips the FFT for masked-out pixels.
+ self._start_compute()
+
+ def _show_image_now(self, img: np.ndarray, angle_idx: int, ch_idx: int):
+ """Display an already-available image with no compute involved."""
+ self._current_image = img
+ self._current_angle = angle_idx
+ self._current_ch = ch_idx
+ self.btn_export_csv.setEnabled(ch_idx in CH1_DERIVED_MODES)
+ self._redraw_image(img)
+ self._update_roi_ui()
+
+ def _redraw_image(self, img: np.ndarray):
+ s = self._sras
+ angle_idx = self._current_angle
+ ch_idx = self._current_ch
+
+ aligned = (self.chk_aligned_view.isChecked()
+ and self._alignment_result is not None
+ and angle_idx in self._alignment_result.per_angle)
+ if aligned:
+ display_img = self._get_aligned_display_image(img, angle_idx, ch_idx)
+ x_axis, y_axis = self._aligned_canvas_axes()
+ else:
+ display_img = img
+ x_axis = s.x_axis_mm(angle_idx)
+ y_axis = s.y_positions_mm(angle_idx)
+
+ dx = x_axis[1] - x_axis[0] if len(x_axis) > 1 else s.pixel_x_mm
+ dy = float(y_axis[1] - y_axis[0]) if len(y_axis) > 1 else 1.0
+ extent = _axes_extent(x_axis, y_axis, dx, dy)
+
+ if self.chk_auto.isChecked():
+ vmin, vmax = float(display_img.min()), float(display_img.max())
+ for spin, val in ((self.spin_vmin, vmin), (self.spin_vmax, vmax)):
+ with QSignalBlocker(spin):
+ spin.setValue(val)
+ else:
+ vmin, vmax = self.spin_vmin.value(), self.spin_vmax.value()
+
+ angle_deg = s.angles_deg[angle_idx]
+ mode_str, unit, colorbar_label = _CHANNEL_DISPLAY[ch_idx]
+ if ch_idx == VELOCITY_MODE_IDX:
+ ch_label = f"Velocity [grating={self.spin_grating_um.value():.2f} µm]"
+ else:
+ ch_label = CH_LABELS[ch_idx]
+
+ title = f"{CH_NAMES[ch_idx]} | {mode_str} | {angle_deg:.1f}°"
+ if aligned:
+ title += " [Aligned]"
+
+ self.image_canvas.show_image(
+ display_img, extent,
+ cmap=self.combo_cmap.currentText(),
+ vmin=vmin, vmax=vmax,
+ xlabel="X (mm)", ylabel="Y (mm)",
+ title=title, colorbar_label=colorbar_label,
+ )
+ self.statusBar().showMessage(
+ f"{s.path.name} | {ch_label} @ {angle_deg:.1f}° "
+ f"| {display_img.shape[1]} × {display_img.shape[0]} px | {unit}"
+ f"{' | Aligned' if aligned else ''}"
+ )
+
+ # ------------------------------------------------------------------
+ # Background compute (only reached on a genuine cache miss)
+ # ------------------------------------------------------------------
+
+ def _start_compute(self):
+ if self._sras is None or self._job_running(Jobs.COMPUTE):
+ return # re-checked when the running compute finishes
+
+ angle_idx = self.spin_angle.value()
+ ch_idx = self.combo_channel.currentIndex()
+ is_fft = ch_idx in CH1_DERIVED_MODES
+
+ self._pending_angle = angle_idx
+ self._pending_ch = ch_idx
+ self._pending_bg_sub = self.chk_bg_sub.isChecked()
+ self._pending_threshold = self.spin_threshold_mv.value()
+ self._pending_fft_pad_factor = self._fft_pad_factor
+
+ worker = ComputeWorker(
+ self._sras, angle_idx, ch_idx,
+ apply_bg_sub=self._pending_bg_sub,
+ n_fft=self._current_n_fft(),
+ dc_threshold_mv=self._pending_threshold,
+ # Reuse the cached DC4 image (if the precompute has reached this
+ # angle) so the FFT skips masked-out pixels entirely and doesn't
+ # need to re-read the CH4 channel from disk.
+ dc4_mv=self._dc_cache.get((angle_idx, CH4_IDX)),
+ is_fft_mode=is_fft,
+ )
+ if not self._run_worker(
+ Jobs.COMPUTE, worker,
+ connect=(
+ ("finished", self._on_compute_done),
+ ("error", lambda msg: self.statusBar().showMessage(
+ f"Compute error: {msg}")),
+ ),
+ on_done=self._after_compute):
+ return
+
+ if is_fft:
+ self.statusBar().showMessage("Computing FFT…")
+ self._show_progress(
+ "main",
+ f"Computing FFT for angle {angle_idx}…\n"
+ "This can take a while on a large scan — result is cached "
+ "so revisiting this angle/mode/threshold will be instant.")
+ else:
+ self.statusBar().showMessage("Computing DC image…")
+ self._show_progress("main", f"Computing DC image for angle {angle_idx}…")
+
+ def _after_compute(self):
+ """If settings changed while the compute was running, re-dispatch
+ through the cache-aware path — the now-current combination may
+ already be cached."""
+ if (self.spin_angle.value(), self.combo_channel.currentIndex(),
+ self.chk_bg_sub.isChecked(), self.spin_threshold_mv.value(),
+ self._fft_pad_factor) != (
+ self._pending_angle, self._pending_ch, self._pending_bg_sub,
+ self._pending_threshold, self._pending_fft_pad_factor):
+ self._refresh_display()
+
+ def _on_compute_done(self, result):
+ self._close_progress("main")
+ if result is None:
+ return # cancelled mid-compute; the partial image must not cache
+ angle_idx = self._pending_angle
+ ch_idx = self._pending_ch
+
+ if ch_idx in CH1_DERIVED_MODES:
+ self._fft_cache[(angle_idx, self._pending_threshold)] = result
+ img = self._scale_for_display(result, ch_idx)
+ else:
+ img = result
+ self._dc_cache[(angle_idx, ch_idx)] = img
+
+ self._show_image_now(img, angle_idx, ch_idx)
+
+ # ------------------------------------------------------------------
+ # Background DC precompute (all angles, so switching is fluid)
+ # ------------------------------------------------------------------
+
+ def _start_dc_precompute(self):
+ if self._sras is None:
+ return
+ generation = self._dc_generation
+ n_angles = self._sras.n_angles
+
+ worker = DcPrecomputeWorker(self._sras)
+ self._run_worker(
+ Jobs.DC_PRECOMPUTE, worker,
+ connect=(
+ ("angle_done", lambda a, dc3, dc4, g=generation:
+ self._on_dc_precompute_angle_done(g, a, dc3, dc4, n_angles)),
+ ("error", lambda msg: self.statusBar().showMessage(
+ f"DC precompute error: {msg}", 5000)),
+ ),
+ quit_on=("finished", "error"),
+ )
+
+ def _on_dc_precompute_angle_done(self, generation: int, angle_idx: int,
+ dc3_mv: np.ndarray, dc4_mv: np.ndarray,
+ n_angles: int):
+ if generation != self._dc_generation:
+ return # stale result from a previously-loaded file — discard
+ self._dc_cache[(angle_idx, CH3_IDX)] = dc3_mv
+ self._dc_cache[(angle_idx, CH4_IDX)] = dc4_mv
+
+ done = sum(1 for a in range(n_angles) if (a, CH4_IDX) in self._dc_cache)
+ self.lbl_dc_precompute.setText(
+ f"Precomputing DC images: {done}/{n_angles} angles ready…"
+ if done < n_angles else "DC images ready for all angles.")
+
+ # If we just finished the angle/channel the user is currently looking
+ # at and it wasn't shown yet (they switched here before the precompute
+ # caught up and are still waiting), show it now.
+ current_ch = self.combo_channel.currentIndex()
+ if (angle_idx == self.spin_angle.value()
+ and not self._job_running(Jobs.COMPUTE)
+ and current_ch in (CH3_IDX, CH4_IDX)
+ and (self._current_angle != angle_idx or self._current_ch != current_ch)):
+ self._refresh_display()
+
+ # ------------------------------------------------------------------
+ # Pixel inspector
+ # ------------------------------------------------------------------
+
+ def _on_pixel_clicked(self, row_idx: int, frame_idx: int):
+ if self._sras is None or self._current_image is None:
+ return
+ angle_idx = self._current_angle
+ if (self.chk_aligned_view.isChecked() and self._alignment_result is not None
+ and angle_idx in self._alignment_result.per_angle):
+ # The click landed on the shared aligned canvas — invert the same
+ # canvas->raw affine used to display it back to a raw (row, frame)
+ # index before looking up the waveform.
+ t = self._alignment_result.per_angle[angle_idx]
+ raw = t.matrix @ np.array([row_idx, frame_idx], dtype=np.float64) + t.offset
+ row_idx, frame_idx = int(round(raw[0])), int(round(raw[1]))
+ n_rows_a, n_frames_a = self._sras.image_shape(angle_idx)
+ if not (0 <= row_idx < n_rows_a and 0 <= frame_idx < n_frames_a):
+ self.statusBar().showMessage(
+ "No source waveform here (padding region of the aligned canvas).")
+ return
+
+ self.lbl_wave_hint.hide()
+ if self._current_ch in CH1_DERIVED_MODES:
+ self.wave_canvas.show_rf_waveform(
+ self._sras, angle_idx, row_idx, frame_idx,
+ apply_bg_sub=self.chk_bg_sub.isChecked())
+ else:
+ self.wave_canvas.show_dc_waveform(
+ self._sras, angle_idx, self._current_ch, row_idx, frame_idx)
+
+ # ------------------------------------------------------------------
+ # Progress dialogs
+ # ------------------------------------------------------------------
+
+ def _show_progress(self, key: str, message: str, maximum: int = 0):
+ """Show (or relabel) the progress dialog under *key*. maximum=0 gives
+ an indeterminate busy indicator."""
+ dlg = self._progress_dlgs.get(key)
+ if dlg is not None:
+ dlg.setLabelText(message)
+ return
+ dlg = QProgressDialog(message, "", 0, maximum, self)
+ dlg.setWindowTitle("Please wait…")
+ dlg.setCancelButton(None)
+ dlg.setWindowModality(Qt.WindowModality.WindowModal)
+ dlg.setMinimumDuration(300) # only appears if it takes > 300 ms
+ dlg.show()
+ self._progress_dlgs[key] = dlg
+
+ def _set_progress(self, key: str, pct: int):
+ dlg = self._progress_dlgs.get(key)
+ if dlg is not None:
+ dlg.setValue(pct)
+
+ def _close_progress(self, key: str):
+ dlg = self._progress_dlgs.pop(key, None)
+ if dlg is not None:
+ dlg.close()
+
+ # ------------------------------------------------------------------
+ # Convert menu: batch DC/FFT compute-and-store (v6 -> v7)
+ # ------------------------------------------------------------------
+
+ def _on_batch_compute(self, mode: str):
+ if self._job_running(Jobs.BATCH):
+ return
+ label = "DC" if mode == "dc" else "FFT"
+ paths, _ = QFileDialog.getOpenFileNames(
+ self, f"Select .sras files to batch-compute {label}", "",
+ "SRAS files (*.sras);;All files (*)")
+ if not paths:
+ return
+
+ self._batch_errors = []
+ # Cache the FFT at the pad the viewer is actually displaying at,
+ # otherwise the batch stores images this window can never use.
+ worker = BatchCacheWorker(paths, mode, self.chk_bg_sub.isChecked(),
+ pad_factor=self._fft_pad_factor)
+ started = self._run_worker(
+ Jobs.BATCH, worker,
+ connect=(
+ ("progress", lambda pct: self._set_progress(Jobs.BATCH, pct)),
+ ("file_done", self._on_batch_file_done),
+ ("finished", lambda p=paths: self._on_batch_finished(p)),
+ ),
+ on_done=self._after_batch,
+ )
+ if not started:
+ return # a second trigger snuck in while the file dialog was open
+
+ self._batch_dc_act.setEnabled(False)
+ self._batch_fft_act.setEnabled(False)
+ self._batch_fft_rowavg_act.setEnabled(False)
+ self._show_progress(
+ Jobs.BATCH, f"Batch computing {label} for {len(paths)} file(s)…",
+ maximum=100)
+
+ def _on_batch_compute_row_avg(self):
+ if self._job_running(Jobs.BATCH):
+ return
+ dlg = RowAverageFftOptionsDialog(
+ self, current_n=self._pending_row_avg_n,
+ current_threshold_mv=self.spin_threshold_mv.value(),
+ pixel_x_mm=self._sras.pixel_x_mm if self._sras is not None else None)
+ if dlg.exec() != QDialog.DialogCode.Accepted:
+ return
+ n, threshold_mv = dlg.get_half_width(), dlg.get_threshold_mv()
+ self._pending_row_avg_n = n
+ self._settings.setValue("fft/row_avg_n", n)
+
+ paths, _ = QFileDialog.getOpenFileNames(
+ self, "Select .sras files to batch-compute Row-Averaged FFT", "",
+ "SRAS files (*.sras);;All files (*)")
+ if not paths:
+ return
+
+ self._batch_errors = []
+ worker = BatchCacheWorker(paths, "fft_rowavg", self.chk_bg_sub.isChecked(),
+ dc_threshold_mv=threshold_mv, row_avg_n=n,
+ pad_factor=self._fft_pad_factor)
+ started = self._run_worker(
+ Jobs.BATCH, worker,
+ connect=(
+ ("progress", lambda pct: self._set_progress(Jobs.BATCH, pct)),
+ ("file_done", self._on_batch_file_done),
+ ("finished", lambda p=paths: self._on_batch_finished(p)),
+ ),
+ on_done=self._after_batch,
+ )
+ if not started:
+ return # a second trigger snuck in while a dialog was open
+
+ self._batch_dc_act.setEnabled(False)
+ self._batch_fft_act.setEnabled(False)
+ self._batch_fft_rowavg_act.setEnabled(False)
+ self._show_progress(
+ Jobs.BATCH,
+ f"Batch computing row-averaged FFT (n={n}, threshold={threshold_mv:.1f} mV) "
+ f"for {len(paths)} file(s)…", maximum=100)
+
+ def _on_batch_file_done(self, path: str, err: str):
+ if err:
+ self._batch_errors.append(f"{Path(path).name} — {err}")
+ self._show_progress(Jobs.BATCH, f"Processed {Path(path).name}…")
+
+ def _on_batch_finished(self, paths: list[str]):
+ self._close_progress(Jobs.BATCH)
+
+ n_total = len(paths)
+ n_failed = len(self._batch_errors)
+ n_ok = n_total - n_failed
+ if n_failed:
+ summary = (f"Batch store: {n_ok}/{n_total} file(s) updated, "
+ f"{n_failed} failed: {'; '.join(self._batch_errors)}")
+ else:
+ summary = f"Batch store: {n_ok}/{n_total} file(s) updated."
+ self.statusBar().showMessage(summary)
+ self._batch_errors = []
+
+ # If the currently-open file was in this batch, reload it so the GUI
+ # picks up the newly-written v7 cache instead of stale state.
+ if self._sras is not None and str(self._sras.path) in paths:
+ self._load_file(str(self._sras.path))
+
+ def _after_batch(self):
+ self._batch_dc_act.setEnabled(True)
+ self._batch_fft_act.setEnabled(True)
+ self._batch_fft_rowavg_act.setEnabled(True)
+
+ # ------------------------------------------------------------------
+ # Fusion: angle alignment
+ # ------------------------------------------------------------------
+
+ def _on_alignment_wizard(self):
+ if self._sras is None or self._sras.n_angles <= 1:
+ return
+ if self._align_wizard is not None:
+ self._align_wizard.raise_()
+ self._align_wizard.activateWindow()
+ return
+
+ ref_idx = 0
+ threshold_mv = self.spin_threshold_mv.value()
+ # Only the threshold carries over from a saved alignment. The wizard's
+ # first page starts every angle pre-rotated from the stage angles and
+ # owns the parameters from there, so inheriting saved per-angle values
+ # would make "Reset to pre-rotation only" mean something different
+ # each time.
+ sidecar = load_manual_alignment(self._sras)
+ if sidecar is not None and sidecar.ref_angle_idx == ref_idx:
+ threshold_mv = sidecar.dc_threshold_mv
+
+ cached_dc4 = {a: img for (a, ch), img in self._dc_cache.items() if ch == CH4_IDX}
+ wiz = AlignmentWizard(
+ self, self._sras, ref_angle_idx=ref_idx, dc_threshold_mv=threshold_mv,
+ cached_dc4_mv=cached_dc4)
+ wiz.alignment_ready.connect(self._on_wizard_finished)
+ wiz.finished.connect(self._on_wizard_closed)
+ wiz.setAttribute(Qt.WidgetAttribute.WA_DeleteOnClose)
+ self._align_wizard = wiz
+ self._update_controls_enabled(True)
+ wiz.show()
+
+ def _on_wizard_closed(self, _result_code: int):
+ self._align_wizard = None
+ self._update_controls_enabled(self._sras is not None)
+
+ def _on_wizard_finished(self, result, out_path: str):
+ """The wizard exported a file; make the session match what was written.
+
+ Applies the *cropped* result, so Aligned View shows exactly the extent
+ that went into the file rather than the wider uncropped canvas, and
+ saves the sidecar for the input scan so reopening it lands in the same
+ place.
+ """
+ if result is None:
+ return
+ self._apply_alignment_result(result, view_checked=True)
+ note = ""
+ try:
+ save_manual_alignment(
+ self._sras, result.ref_angle_idx, result.dc_threshold_mv,
+ {a: ManualAngleParams(t.rotation_deg, t.shift_mm)
+ for a, t in result.per_angle.items()})
+ except OSError as exc:
+ note = f" (could not save the sidecar: {exc})"
+ self._update_controls_enabled(self._sras is not None)
+ nr, nc = result.canvas_shape
+ self.statusBar().showMessage(
+ f"Aligned scan written to {Path(out_path).name}; Aligned View now "
+ f"shows the exported {nc}×{nr} px region.{note}")
+ if self._current_image is not None:
+ self._refresh_display()
+
+ def _apply_alignment_result(self, result, *, view_checked: bool):
+ """Install (or clear, with result=None) the active alignment: reset
+ the aligned-image cache and set the Aligned View checkbox without
+ firing its change signal."""
+ self._alignment_result = result
+ self._aligned_cache = {}
+ with QSignalBlocker(self.chk_aligned_view):
+ self.chk_aligned_view.setChecked(view_checked)
+ self.chk_aligned_view.setEnabled(result is not None)
+
+ # ------------------------------------------------------------------
+ # FFT Options
+ # ------------------------------------------------------------------
+
+ def _on_fft_options(self):
+ dlg = FftOptionsDialog(
+ self,
+ current_backend=compute.get_fft_backend(),
+ current_pad_factor=self._fft_pad_factor,
+ samples_per_frame=self._sras.samples_per_frame if self._sras else None,
+ sample_rate_hz=self._sras.sample_rate_hz if self._sras else None,
+ grating_um=self.spin_grating_um.value(),
+ )
+ if dlg.exec() != QDialog.DialogCode.Accepted:
+ return
+ compute.set_fft_backend(dlg.get_backend())
+ self._fft_pad_factor = dlg.get_pad_factor()
+ self._settings.setValue("fft/backend", compute.get_fft_backend())
+ self._settings.setValue("fft/pad_factor", self._fft_pad_factor)
+ # Pad factor no longer gates the display: it only affects a future
+ # live compute for an angle with nothing cached yet, or an explicit
+ # batch recompute — never what's already shown. Just keep the info
+ # panel's divergence note current.
+ if self._is_fft_mode():
+ self._update_scan_info_labels()
+
+ # ------------------------------------------------------------------
+
+ def closeEvent(self, event):
+ if self._align_wizard is not None:
+ self._align_wizard.close()
+
+ # Signal every cancellable worker first, then wait. Waiting without
+ # signalling means sitting out whatever is in flight — on a large
+ # scan a single angle is ~40 s.
+ jobs = list(self._jobs.values())
+ for _thread, worker, _on_done in jobs:
+ stop = getattr(worker, "stop", None)
+ if callable(stop):
+ stop()
+ for thread, _worker, _on_done in jobs:
+ thread.quit()
+ thread.wait(5000)
+ super().closeEvent(event)
+
+
+# ---------------------------------------------------------------------------
+
+def main():
+ app = QApplication(sys.argv)
+ window = SrasViewerWindow(
+ initial_path=sys.argv[1] if len(sys.argv) > 1 else None)
+ window.show()
+ sys.exit(app.exec())
+
diff --git a/sras_viewer_requirements.txt b/sras_viewer_requirements.txt
deleted file mode 100644
index 05adf9a..0000000
--- a/sras_viewer_requirements.txt
+++ /dev/null
@@ -1,3 +0,0 @@
-PyQt6==6.10.2
-numpy==2.4.1
-matplotlib==3.10.8
diff --git a/sras_workers.py b/sras_workers.py
index 5fd0787..ff9d130 100644
--- a/sras_workers.py
+++ b/sras_workers.py
@@ -19,9 +19,8 @@ from matplotlib.figure import Figure
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_align_export import write_aligned_sras
+from sras_compute import cache_file, compute_rf_image, dc_image_mv
from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, SrasFile
# Concurrency caps. Batch conversion runs one process per file, and each of
@@ -58,6 +57,34 @@ class CancellableWorker(QObject):
return self._stop
+class _PooledWorker(CancellableWorker):
+ """Fans a per-item computation across a thread pool, emitting each result
+ from this worker's own thread as it lands (never from a pool thread).
+
+ Subclasses provide _plan() -> n_workers (stashing whatever per-run
+ context they need), _items(), _one(item) -> result, and _emit(result).
+ On stop(): queued items are dropped, in-flight ones are not waited for —
+ that is what keeps closing the window responsive on a large scan.
+ """
+ finished = pyqtSignal()
+ error = pyqtSignal(str)
+
+ def run(self):
+ try:
+ pool = ThreadPoolExecutor(max_workers=max(1, self._plan()))
+ try:
+ futures = [pool.submit(self._one, it) for it in self._items()]
+ for fut in as_completed(futures):
+ if self._stop:
+ break
+ self._emit(fut.result())
+ finally:
+ pool.shutdown(wait=not self._stop, cancel_futures=True)
+ self.finished.emit()
+ except Exception as exc:
+ self.error.emit(str(exc))
+
+
class LoadWorker(QObject):
finished = pyqtSignal(object) # SrasFile | None
error = pyqtSignal(str)
@@ -121,29 +148,29 @@ class ComputeWorker(CancellableWorker):
self.error.emit(str(exc))
-class DcPrecomputeWorker(CancellableWorker):
+class DcPrecomputeWorker(_PooledWorker):
"""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
+ self._angle_budget = 0
- def _one_angle(self, a: int) -> tuple[int, np.ndarray, np.ndarray]:
+ def _plan(self) -> int:
+ n_workers, self._angle_budget = compute.plan_angle_level(self._sras)
+ return n_workers
+
+ def _items(self):
+ return range(self._sras.n_angles)
+
+ def _one(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).
@@ -153,26 +180,8 @@ class DcPrecomputeWorker(CancellableWorker):
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))
+ def _emit(self, result):
+ self.angle_done.emit(*result)
class BatchCacheWorker(QObject):
@@ -181,9 +190,15 @@ class BatchCacheWorker(QObject):
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
+ *mode* is ``"dc"`` (CH3/CH4 mean images), ``"fft"`` (CH1 peak-frequency
images, unmasked — masking is applied at display time, same as v5's PREC
- convention).
+ convention), or ``"fft_rowavg"`` (same-row, distance-weighted CH1
+ averaging before the FFT — needs *dc_threshold_mv* and a positive
+ *row_avg_n*; see ``sras_compute.cache_file``).
+
+ Both FFT modes cache at *pad_factor*, which the caller sets from the
+ viewer's own padding — a cache stored at a pad the user is not viewing
+ at is one the display can never use.
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
@@ -195,11 +210,16 @@ class BatchCacheWorker(QObject):
file_done = pyqtSignal(str, str)
finished = pyqtSignal()
- def __init__(self, paths: list[str], mode: str, apply_bg_sub: bool):
+ def __init__(self, paths: list[str], mode: str, apply_bg_sub: bool,
+ dc_threshold_mv: float | None = None, row_avg_n: int = 0,
+ pad_factor: int = 1):
super().__init__()
self._paths = paths
self._mode = mode
self._apply_bg_sub = apply_bg_sub
+ self._dc_threshold = dc_threshold_mv
+ self._row_avg_n = row_avg_n
+ self._pad_factor = pad_factor
def _report(self, path: str, err: str, done: int, total: int):
self.file_done.emit(path, err)
@@ -224,7 +244,10 @@ class BatchCacheWorker(QObject):
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
+ compute.get_fft_backend(), per_proc_workers,
+ pad_factor=self._pad_factor,
+ dc_threshold_mv=self._dc_threshold,
+ row_avg_n=self._row_avg_n): p
for p in paths
}
for fut in as_completed(futures):
@@ -247,7 +270,11 @@ class BatchCacheWorker(QObject):
for path in paths:
try:
err = cache_file(path, self._mode, self._apply_bg_sub,
- compute.get_fft_backend(), compute._MAX_WORKERS)
+ compute.get_fft_backend(),
+ compute.default_max_workers(),
+ pad_factor=self._pad_factor,
+ dc_threshold_mv=self._dc_threshold,
+ row_avg_n=self._row_avg_n)
except Exception as exc:
err = str(exc)
done += 1
@@ -419,34 +446,9 @@ class BatchExportWorker(CancellableWorker):
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) # 0–100
- 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.
+class Ch4MaskWorker(_PooledWorker):
+ """Fetches each requested angle's CH4 (Bias B) DC image in mV, for the
+ alignment wizard's initial threshold-mask stack.
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
@@ -455,86 +457,108 @@ class Ch4MaskWorker(QObject):
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).
+ CorrelatePage._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
+ self._budget = 0
- 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))
+ def _plan(self) -> int:
+ n_workers, self._budget = compute.plan_angle_level(self._sras)
+ return n_workers
+
+ def _items(self):
+ return self._angles
+
+ def _one(self, a: int) -> tuple[int, np.ndarray]:
+ return a, dc_image_mv(self._sras, a, CH4_IDX,
+ max_workers=1, budget=self._budget)
+
+ def _emit(self, result):
+ self.angle_done.emit(*result)
-class CrossCorrelateWorker(QObject):
- """FFT phase-correlation translation for each of *angle_indices* against
- *ref_angle_idx*, for ManualAlignmentDialog's Auto Cross-Correlate button.
+class CrossCorrelateWorker(_PooledWorker):
+ """Rigid registration (rotation + translation, never scale) of each of
+ *angle_indices* against *ref_angle_idx*, for the alignment wizard's
+ Run/Re-run Correlation 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.
+ Runs on a background thread — registering a real many-angle,
+ high-resolution scan takes long enough that doing it on the GUI thread
+ would visibly freeze the dialog. Rotation is *searched*, not taken from the
+ stage's reported angle: see compute.register_angle_to_reference, which
+ seeds from that angle but scores both of its signs and refines from there.
+ dc4_mv is the dialog's own already-in-memory per-angle CH4 image — 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)
+ # angle_idx, rotation_deg, shift_x_mm, shift_y_mm, score, source
+ angle_done = pyqtSignal(int, float, float, float, float, 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):
+ dc4_mv: dict[int, np.ndarray], *, reg_kwargs: dict | None = None):
+ """*reg_kwargs* is splatted into register_angle_to_reference — every
+ registration setting the wizard exposes (sources, threshold, search
+ width, seed, signs, refine, grid sizes) travels in it, so this class
+ holds no opinion about which knobs exist and exposing another needs no
+ change here."""
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
+ self._reg_kwargs = dict(reg_kwargs or {})
- 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 _plan(self) -> int:
+ return compute.registration_workers(self._sras)
+
+ def _items(self):
+ return self._angles
+
+ def _one(self, a: int) -> tuple[int, compute.RigidFit]:
+ return a, compute.register_angle_to_reference(
+ self._sras, a, self._ref, self._dc4_mv, **self._reg_kwargs)
+
+ def _emit(self, result):
+ a, fit = result
+ self.angle_done.emit(a, fit.rotation_deg, fit.shift_mm[0],
+ fit.shift_mm[1], fit.score, fit.source)
+
+
+class AlignedExportWorker(CancellableWorker):
+ """Writes the aligned, cropped .sras on a background thread.
+
+ Unlike every other worker here this one produces a *file*, which changes
+ what cancellation has to mean: write_aligned_sras stages into a ".part"
+ sibling and removes it when should_stop() fires, so a cancelled or crashed
+ export leaves nothing behind. That matters more than it sounds — a
+ truncated .sras is not detectably broken, since the v6 parser reads a short
+ file as an aborted scan and opens it happily.
+
+ Cancellation is polled per output row chunk, the same granularity
+ CancellableWorker's docstring justifies, so closing the window never waits
+ on a multi-gigabyte write.
+ """
+ progress = pyqtSignal(int) # 0-100
+ finished = pyqtSignal(str, str) # written path ("" = none), error
+
+ def __init__(self, sras: SrasFile, result, out_path: str):
+ super().__init__()
+ self._sras = sras
+ self._result = result
+ self._out_path = out_path
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()
+ written = write_aligned_sras(
+ self._sras, self._result, self._out_path,
+ progress_cb=self.progress.emit, should_stop=self._stopped)
+ if self._stopped():
+ self.finished.emit("", "") # cancelled: no file, no error
+ else:
+ self.finished.emit(str(written), "")
except Exception as exc:
- self.error.emit(str(exc))
+ self.finished.emit("", str(exc))
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 0000000..5486080
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,16 @@
+"""Shared test setup: repo-root imports, the offscreen Qt platform, and
+hermetic QSettings (tests must not read or write the user's real viewer
+settings)."""
+
+import os
+import sys
+import tempfile
+from pathlib import Path
+
+os.environ.setdefault("QT_QPA_PLATFORM", "offscreen")
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+from PyQt6.QtCore import QSettings # noqa: E402
+
+QSettings.setPath(QSettings.Format.IniFormat, QSettings.Scope.UserScope,
+ tempfile.mkdtemp(prefix="sras_qsettings_"))
diff --git a/tests/test_align_export.py b/tests/test_align_export.py
new file mode 100644
index 0000000..0cf1360
--- /dev/null
+++ b/tests/test_align_export.py
@@ -0,0 +1,646 @@
+"""Aligned/cropped .sras export: does the written file actually hold the
+alignment the viewer showed?
+
+The export is the one place an alignment stops being a transform applied on the
+fly and becomes bytes on disk, so these tests care about two things above all:
+the file's geometry describes what was written, and the pixels in it are the
+same pixels apply_alignment would have drawn. The strongest check is the
+round-trip — register the exported file against itself and demand identity,
+which no amount of self-consistent-but-wrong index math can fake.
+
+No Qt: this exercises sras_align_export and sras_compute directly.
+"""
+
+import struct
+
+import numpy as np
+import pytest
+
+import sras_align_export as export
+import sras_compute as compute
+from sras_format import CH3_IDX, CH4_IDX, HDR_SIZE_V6, SrasFile, adc_to_mv, mv_to_adc
+import tools.make_test_sras as gen
+
+_THRESHOLD_MV = 80.0
+# Same reasoning as tests/test_alignment.py: a quarter degree is already
+# sub-pixel for this sample at the registration pitch.
+_ROT_TOL_DEG = 0.5
+_SHIFT_TOL_MM = 0.02
+
+
+def dc_mv(sras: SrasFile, angle_idx: int, ch: int = CH4_IDX) -> np.ndarray:
+ return adc_to_mv(compute.compute_dc_image(sras, angle_idx, ch), *sras.cal(ch))
+
+
+@pytest.fixture(scope="module")
+def rig(tmp_path_factory):
+ """The rotating-sample scan, its truth alignment, and its export."""
+ tmpdir = tmp_path_factory.mktemp("sras_export")
+ src_path = tmpdir / "rotating.sras"
+ meta = gen.write_rotating(src_path, n_angles=4)
+ sras = SrasFile(str(src_path))
+
+ params = {a: compute.ManualAngleParams(rot, shift)
+ for a, (rot, shift) in meta["truth"].items()}
+ result = compute.build_manual_alignment(sras, 0, _THRESHOLD_MV, params)
+
+ out_path = tmpdir / "rotating_aligned.sras"
+ export.write_aligned_sras(sras, result, out_path)
+ return type("Rig", (), dict(
+ tmpdir=tmpdir, src_path=src_path, sras=sras, meta=meta,
+ result=result, out_path=out_path, out=SrasFile(str(out_path))))
+
+
+# ---------------------------------------------------------------------------
+# Geometry and file structure
+# ---------------------------------------------------------------------------
+
+def test_output_is_v6_with_uniform_geometry(rig):
+ out, result = rig.out, rig.result
+ n_rows, n_cols = result.canvas_shape
+
+ assert out.version == 6
+ assert out.n_angles == rig.sras.n_angles
+ assert set(out.n_rows) == {n_rows}, "every angle must share the canvas rows"
+ assert set(out.n_frames) == {n_cols}, "every angle must share the canvas frames"
+ assert np.allclose(out.x_start_mm, result.canvas_origin_mm[0])
+ # x_delta must stay velocity/laser_freq or x_axis_mm() contradicts the
+ # geometry table; the canvas pitch is the reference angle's own pitch, so
+ # this is exact rather than approximate.
+ assert np.allclose(out.x_delta_mm_per_angle, rig.sras.pixel_x_mm)
+ assert out.pixel_x_mm == pytest.approx(rig.sras.pixel_x_mm)
+
+
+def test_row_table_matches_the_canvas(rig):
+ expected = (rig.result.canvas_origin_mm[1]
+ + np.arange(rig.result.canvas_shape[0]) * rig.result.canvas_dy_mm)
+ for a in range(rig.out.n_angles):
+ assert rig.out.y_positions_mm(a) == pytest.approx(expected, abs=1e-4)
+
+
+def test_angle_table_and_calibration_round_trip(rig):
+ assert rig.out.angles_deg == pytest.approx(rig.sras.angles_deg)
+ for ch in range(rig.sras.n_channels):
+ assert rig.out.cal(ch) == pytest.approx(rig.sras.cal(ch))
+ assert rig.out.samples_per_frame == rig.sras.samples_per_frame
+ assert rig.out.bytes_per_sample == rig.sras.bytes_per_sample
+ assert rig.out.n_channels == rig.sras.n_channels
+ assert rig.out.background == pytest.approx(rig.sras.background)
+
+
+def test_no_cache_tail(rig):
+ """File ends exactly at the waveform data — nothing trailing.
+
+ A stale cache tail would be indexed by the *input's* grid, so the export
+ must not carry one; asserting on the exact file size is what proves it,
+ since a v7 tail would simply be ignored by a v6 parser.
+ """
+ end = max(off + n for _, off, n in rig.out.iter_angle_blocks())
+ assert rig.out_path.stat().st_size == end
+ assert all(img is None for img in rig.out.precomputed_dc4_mv)
+
+
+def test_declared_header_size_is_v6(rig):
+ raw = rig.out_path.read_bytes()[:HDR_SIZE_V6]
+ magic, version, n_angles = struct.unpack(">4sBH", raw[:7])
+ assert (magic, version, n_angles) == (b"SRAS", 6, rig.sras.n_angles)
+
+
+# ---------------------------------------------------------------------------
+# The pixels themselves
+# ---------------------------------------------------------------------------
+
+def test_export_matches_apply_alignment(rig):
+ """The exported waveforms decode to the same DC image the viewer drew —
+ over the *whole* canvas, padding included.
+
+ Two rules have to be exactly right for this, and each fails differently:
+ * rounding must be floor(x + 0.5), not np.rint, or pixels on exact
+ half-integer boundaries pick the neighbouring source pixel;
+ * out-of-bounds must be tested on the fractional coordinate against
+ [0, n-1], not on the rounded index, or a one-pixel rim gets real data
+ where the preview shows padding.
+ Comparing every pixel rather than only the interior is what catches the
+ second one, since a rim discrepancy hides inside a `preview != 0` mask.
+ """
+ for a in range(rig.sras.n_angles):
+ preview = compute.apply_alignment(rig.result, a, dc_mv(rig.sras, a))
+ actual = dc_mv(rig.out, a)
+ assert actual.shape == preview.shape
+ # Padding matches to within half an ADC step: apply_alignment pads with
+ # literal 0.0 mV, the export with the nearest integer ADC code to 0 mV.
+ tol = abs(rig.sras.cal(CH4_IDX)[0]) / 2.0 + 1e-4
+ # Exclude the epsilon rim the export deliberately keeps and scipy drops
+ # (see test_edge_tolerance_only_affects_the_epsilon_rim).
+ sr, sc = export._src_coords(rig.result.per_angle[a],
+ np.arange(preview.shape[0]), preview.shape[1])
+ rim = export._in_bounds(sr, sc, *rig.sras.image_shape(a)) & (preview == 0.0)
+ cmp = ~rim
+ assert actual[cmp] == pytest.approx(preview[cmp], abs=tol), \
+ f"angle {a}: exported pixels differ from the aligned preview"
+ # And exactly, wherever there is real data.
+ inside = (preview != 0.0)
+ assert inside.any(), f"angle {a}: preview is entirely padding"
+ assert actual[inside] == pytest.approx(preview[inside], abs=1e-6), \
+ f"angle {a}: exported data pixels are not bit-equal to the preview"
+
+
+def test_reference_angle_is_exported_whole(rig):
+ """The reference angle must survive as a complete, exact integer crop.
+
+ It is the coordinate authority — its transform is the identity with an
+ integer offset by construction — so every one of its source pixels has to
+ appear in the export. This is what _EDGE_TOL exists for: that offset comes
+ out of the mm-space affine chain as -20 - 7e-15, and a bare `>= 0` bounds
+ test silently drops the angle's entire first row and last column.
+ """
+ ref = rig.result.ref_angle_idx
+ src_rows, src_frames = rig.sras.image_shape(ref)
+ plan = export.plan_export(rig.sras, rig.result)
+ assert plan.valid_px[ref] == src_rows * src_frames, \
+ "reference angle lost pixels to the in-bounds test"
+
+ # And the values themselves land as an exact, unrotated block.
+ src_img = dc_mv(rig.sras, ref)
+ out_img = dc_mv(rig.out, ref)
+ t = rig.result.per_angle[ref]
+ row0, col0 = (int(round(-t.offset[0])), int(round(-t.offset[1])))
+ assert np.array_equal(out_img[row0:row0 + src_rows, col0:col0 + src_frames],
+ src_img), \
+ "reference angle is not a verbatim block in the export"
+
+
+def test_edge_tolerance_only_affects_the_epsilon_rim(rig):
+ """Where the export's bounds test and scipy's disagree, the coordinate must
+ be within _EDGE_TOL of the boundary — i.e. only pixels whose scipy answer
+ was itself decided by float noise, never a real half-pixel decision."""
+ for a in range(rig.sras.n_angles):
+ t = rig.result.per_angle[a]
+ src_rows, src_frames = rig.sras.image_shape(a)
+ n_rows, n_cols = rig.result.canvas_shape
+
+ ones = np.ones((src_rows, src_frames), dtype=np.float32)
+ scipy_valid = compute.apply_alignment(rig.result, a, ones) > 0.5
+ sr, sc = export._src_coords(t, np.arange(n_rows), n_cols)
+ ours = export._in_bounds(sr, sc, src_rows, src_frames)
+
+ differ = ours != scipy_valid
+ assert not (scipy_valid & ~ours).any(), \
+ f"angle {a}: export drops pixels scipy keeps"
+ if differ.any():
+ # Every disagreement sits within the tolerance of an edge.
+ near = (np.abs(sr) <= export._EDGE_TOL)
+ near |= (np.abs(sr - (src_rows - 1)) <= export._EDGE_TOL)
+ near |= (np.abs(sc) <= export._EDGE_TOL)
+ near |= (np.abs(sc - (src_frames - 1)) <= export._EDGE_TOL)
+ assert near[differ].all(), \
+ f"angle {a}: bounds differ away from the epsilon rim"
+
+
+def test_export_matches_apply_alignment_on_ch3(rig):
+ """Channel-agnostic: the gather moves whole pixels, not per-channel images."""
+ for a in range(rig.sras.n_angles):
+ preview = compute.apply_alignment(rig.result, a,
+ dc_mv(rig.sras, a, CH3_IDX))
+ actual = dc_mv(rig.out, a, CH3_IDX)
+ inside = preview != 0.0
+ assert actual[inside] == pytest.approx(preview[inside], abs=1e-3)
+
+
+def test_padding_is_zero_mv_not_zero_adc(rig):
+ """Unreachable canvas pixels must read as ~0 mV on every channel.
+
+ Filling with literal zero ADC would decode to (0 - yoff) * ymult + yzero —
+ for this fixture's CH4 calibration that is +100 mV, well above any sensible
+ mask threshold, so the padding would masquerade as valid sample everywhere.
+ """
+ a = rig.sras.n_angles - 1
+ t = rig.result.per_angle[a]
+ n_rows, n_cols = rig.result.canvas_shape
+ src_rows, src_frames = rig.sras.image_shape(a)
+
+ sr, sc = export._src_coords(t, np.arange(n_rows), n_cols)
+ outside = ~export._in_bounds(sr, sc, src_rows, src_frames)
+ assert outside.any(), "rotated angle should leave unreachable canvas corners"
+
+ for ch in (CH3_IDX, CH4_IDX):
+ img = dc_mv(rig.out, a, ch)
+ half_step = abs(rig.sras.cal(ch)[0]) / 2.0
+ assert np.abs(img[outside]).max() <= half_step + 1e-6, \
+ f"CH{ch} padding is not within half an ADC step of 0 mV"
+
+ # And the sanity check that makes the above meaningful: zero ADC would not
+ # have passed it.
+ assert abs(adc_to_mv(0, *rig.sras.cal(CH4_IDX))) > 10.0
+
+
+def test_reregistering_the_export_is_identity(rig):
+ """The export really is aligned: registering it against its own angle 0
+ recovers no rotation and no shift.
+
+ The end-to-end check — it fails for any index error, sign flip, wrong pivot
+ or origin mistake anywhere in crop/affine/gather, in a way the
+ self-consistency tests above cannot.
+ """
+ dc4 = {a: dc_mv(rig.out, a) for a in range(rig.out.n_angles)}
+ for a in range(1, rig.out.n_angles):
+ fit = compute.register_angle_to_reference(
+ rig.out, a, 0, dc4, dc_threshold_mv=_THRESHOLD_MV,
+ seed_deg=0.0, seed_signs=(1,))
+ assert abs(fit.rotation_deg) <= _ROT_TOL_DEG, \
+ f"angle {a} still rotated by {fit.rotation_deg:.3f}° after export"
+ assert float(np.hypot(*fit.shift_mm)) <= _SHIFT_TOL_MM, \
+ f"angle {a} still shifted by {fit.shift_mm} mm after export"
+
+
+def test_export_of_int16_input(rig, tmp_path):
+ """bps=2 inputs keep their big-endian int16 dtype through the gather."""
+ src_path = tmp_path / "i16.sras"
+ gen.write(src_path, n_angles=2, samples_per_frame=16, bps=2)
+ sras = SrasFile(str(src_path))
+ result = compute.build_manual_alignment(sras, 0, 0.0, {})
+
+ out_path = tmp_path / "i16_aligned.sras"
+ export.write_aligned_sras(sras, result, out_path)
+ out = SrasFile(str(out_path))
+
+ assert out.bytes_per_sample == 2
+ assert out.data[0].dtype == np.dtype(">i2")
+ for a in range(sras.n_angles):
+ preview = compute.apply_alignment(result, a, dc_mv(sras, a))
+ actual = dc_mv(out, a)
+ inside = preview != 0.0
+ assert actual[inside] == pytest.approx(preview[inside], abs=1e-3)
+
+
+# ---------------------------------------------------------------------------
+# Cropping
+# ---------------------------------------------------------------------------
+
+def test_crop_is_a_window_of_the_full_canvas(rig):
+ """crop_alignment_result must resample exactly the sub-rectangle it names.
+
+ Asserted as bit-exact equality, not approximately: the crop composes into
+ the affine's offset by an integer number of canvas pixels, so anything but
+ an exact match means the composition is wrong.
+ """
+ n_rows, n_cols = rig.result.canvas_shape
+ row0, col0 = n_rows // 5, n_cols // 4
+ nr, nc = n_rows // 2, n_cols // 3
+ cropped = compute.crop_alignment_result(rig.result, row0, col0, nr, nc)
+
+ assert cropped.canvas_shape == (nr, nc)
+ assert cropped.canvas_origin_mm[0] == pytest.approx(
+ rig.result.canvas_origin_mm[0] + col0 * rig.result.canvas_dx_mm)
+ assert cropped.canvas_origin_mm[1] == pytest.approx(
+ rig.result.canvas_origin_mm[1] + row0 * rig.result.canvas_dy_mm)
+
+ for a in range(rig.sras.n_angles):
+ img = dc_mv(rig.sras, a)
+ full = compute.apply_alignment(rig.result, a, img)
+ assert np.array_equal(
+ compute.apply_alignment(cropped, a, img),
+ full[row0:row0 + nr, col0:col0 + nc]), \
+ f"angle {a}: cropped resample is not the same window"
+ # Rotation/shift are properties of the angle, not of the canvas.
+ assert cropped.per_angle[a].rotation_deg == rig.result.per_angle[a].rotation_deg
+ assert cropped.per_angle[a].shift_mm == rig.result.per_angle[a].shift_mm
+
+
+def test_cropped_export_round_trips(rig, tmp_path):
+ n_rows, n_cols = rig.result.canvas_shape
+ row0, col0, nr, nc = n_rows // 4, n_cols // 4, n_rows // 2, n_cols // 2
+ cropped = compute.crop_alignment_result(rig.result, row0, col0, nr, nc)
+
+ out_path = tmp_path / "cropped.sras"
+ export.write_aligned_sras(rig.sras, cropped, out_path)
+ out = SrasFile(str(out_path))
+
+ assert set(out.n_rows) == {nr} and set(out.n_frames) == {nc}
+ assert out.x_start_mm[0] == pytest.approx(cropped.canvas_origin_mm[0], abs=1e-4)
+ for a in range(rig.sras.n_angles):
+ preview = compute.apply_alignment(cropped, a, dc_mv(rig.sras, a))
+ actual = dc_mv(out, a)
+ inside = preview != 0.0
+ if inside.any():
+ assert actual[inside] == pytest.approx(preview[inside], abs=1e-3)
+
+
+def test_crop_rejects_empty_window(rig):
+ with pytest.raises(ValueError, match="empty crop"):
+ compute.crop_alignment_result(rig.result, 0, 0, 0, 10)
+ with pytest.raises(ValueError, match="empty crop"):
+ compute.crop_alignment_result(rig.result, 0, 0, 10, -1)
+
+
+# ---------------------------------------------------------------------------
+# plan_export and overlap_stats
+# ---------------------------------------------------------------------------
+
+def test_plan_export_matches_what_was_written(rig):
+ plan = export.plan_export(rig.sras, rig.result)
+ n_rows, n_cols = rig.result.canvas_shape
+ assert (plan.n_rows, plan.n_frames) == (n_rows, n_cols)
+ assert plan.n_angles == rig.sras.n_angles
+
+ data_bytes = sum(n for _, _, n in rig.out.iter_angle_blocks())
+ assert plan.total_bytes == data_bytes
+ assert plan.bytes_per_angle * plan.n_angles == plan.total_bytes
+
+ # Coverage must agree with the pixels that actually carry data. The
+ # reference angle is unrotated, so its whole footprint lands inside.
+ ref_px = np.prod(rig.sras.image_shape(0))
+ assert plan.valid_px[0] == ref_px
+ for a in range(1, rig.sras.n_angles):
+ assert 0 < plan.valid_px[a] <= n_rows * n_cols
+ assert 0.0 < plan.coverage_frac(a) < 1.0
+
+
+def test_plan_export_flags_a_crop_that_misses_an_angle(rig):
+ """A crop over a corner the rotated angles cannot reach must warn, and the
+ export must still succeed by writing that angle as padding."""
+ n_rows, n_cols = rig.result.canvas_shape
+ corner = compute.crop_alignment_result(rig.result, 0, 0,
+ max(1, n_rows // 12),
+ max(1, n_cols // 12))
+ plan = export.plan_export(rig.sras, corner)
+ empty = [a for a in range(rig.sras.n_angles) if plan.valid_px[a] == 0]
+ assert empty, "top-left canvas corner should be unreachable for some angle"
+ assert any("all padding" in w for w in plan.warnings)
+
+
+def test_overlap_stats():
+ counts = np.array([[0, 1, 2], [3, 3, 0], [0, 2, 3]])
+ stats = compute.overlap_stats(counts, 3)
+ assert stats["union_px"] == 6
+ assert stats["full_px"] == 3
+ assert stats["full_frac"] == pytest.approx(0.5)
+ assert stats["max_count"] == 3
+ assert stats["mean_count"] == pytest.approx((1 + 2 + 3 + 3 + 2 + 3) / 6)
+ assert stats["empty"] is False
+
+ empty = compute.overlap_stats(np.zeros((4, 4), dtype=int), 3)
+ assert empty["empty"] is True
+ assert empty["full_frac"] == 0.0 and empty["mean_count"] == 0.0
+
+
+def test_largest_rect_at_least():
+ # A 2x3 block of 3s with a notch that a bounding box would swallow.
+ counts = np.array([
+ [0, 0, 0, 0, 0],
+ [0, 3, 3, 3, 0],
+ [0, 3, 3, 3, 0],
+ [0, 3, 0, 3, 0],
+ ])
+ row0, col0, nr, nc = compute.largest_rect_at_least(counts, 3)
+ assert (nr * nc) == 6 and (row0, col0, nr, nc) == (1, 1, 2, 3)
+ assert (counts[row0:row0 + nr, col0:col0 + nc] >= 3).all()
+
+ # A column taller than the wide block is the better rectangle.
+ tall = np.array([[3, 3], [3, 0], [3, 0], [3, 0]])
+ r0, c0, nr2, nc2 = compute.largest_rect_at_least(tall, 3)
+ assert (r0, c0, nr2, nc2) == (0, 0, 4, 1)
+
+ assert compute.largest_rect_at_least(np.zeros((3, 3), dtype=int), 1) is None
+ # Whole-array case: no notch, so the answer is the array itself.
+ assert compute.largest_rect_at_least(np.full((3, 4), 2), 2) == (0, 0, 3, 4)
+
+
+def test_largest_rect_matches_brute_force():
+ """Randomized check against an O(n^4) reference.
+
+ The histogram sweep is short and easy to get subtly wrong — an off-by-one in
+ the stack unwind yields rectangles that are merely large, and "large but not
+ maximal" is invisible by eye on real data.
+ """
+ def brute(good):
+ n_rows, n_cols = good.shape
+ best = 0
+ for r0 in range(n_rows):
+ for r1 in range(r0 + 1, n_rows + 1):
+ run = 0
+ for g in good[r0:r1].all(axis=0):
+ run = run + 1 if g else 0
+ best = max(best, run * (r1 - r0))
+ return best
+
+ rng = np.random.default_rng(0)
+ for _ in range(200):
+ counts = rng.integers(0, 3, size=(int(rng.integers(1, 9)),
+ int(rng.integers(1, 9))))
+ got = compute.largest_rect_at_least(counts, 2)
+ expected = brute(counts >= 2)
+ if got is None:
+ assert expected == 0
+ continue
+ row0, col0, nr, nc = got
+ assert (counts[row0:row0 + nr, col0:col0 + nc] >= 2).all(), \
+ f"rectangle is not pure:\n{counts}\n{got}"
+ assert nr * nc == expected, \
+ f"not maximal ({nr * nc} < {expected}):\n{counts}\n{got}"
+
+
+def test_largest_rect_is_pure_on_the_real_fixture(rig):
+ """On real overlap counts the returned rectangle must contain only
+ full-overlap pixels — the property a bounding box would violate."""
+ n = rig.sras.n_angles
+ masks = {a: (dc_mv(rig.sras, a) >= _THRESHOLD_MV).astype(np.float32)
+ for a in range(n)}
+ counts = sum(compute.apply_alignment(rig.result, a, masks[a]) > 0.5
+ for a in range(n)).astype(int)
+ assert counts.max() == n, "fixture alignment should have a full-overlap region"
+
+ rect = compute.largest_rect_at_least(counts, n)
+ assert rect is not None
+ row0, col0, nr, nc = rect
+ assert (counts[row0:row0 + nr, col0:col0 + nc] == n).all(), \
+ "convenience crop must not include pixels some angle misses"
+
+ # And it must beat the naive bounding box, which here is impure.
+ rr, cc = np.nonzero(counts == n)
+ bbox_pure = (counts[rr.min():rr.max() + 1, cc.min():cc.max() + 1] == n).all()
+ assert not bbox_pure, "fixture no longer exercises the bounding-box hazard"
+
+
+# ---------------------------------------------------------------------------
+# Legacy inputs, validation and durability
+# ---------------------------------------------------------------------------
+
+@pytest.mark.parametrize("version", [2, 4])
+def test_legacy_input_exports_as_v6(version, tmp_path):
+ """v2-v5 inputs keep no verbatim preamble/background spans, so those
+ sections have to be re-encoded. v2 additionally has neither."""
+ src_path = tmp_path / f"legacy_v{version}.sras"
+ gen.write_legacy(src_path, version=version, n_angles=2)
+ sras = SrasFile(str(src_path))
+ result = compute.build_manual_alignment(sras, 0, 0.0, {})
+
+ out_path = tmp_path / f"legacy_v{version}_aligned.sras"
+ export.write_aligned_sras(sras, result, out_path)
+ out = SrasFile(str(out_path))
+
+ assert out.version == 6
+ assert out.n_angles == sras.n_angles
+ # A zero background rather than a zero-length one: consumers subtract it
+ # from a (spf,)-shaped row, which a length-0 array cannot broadcast against.
+ assert out.background is not None
+ assert out.background.size == sras.samples_per_frame
+ if sras.background is None:
+ assert np.all(out.background == 0)
+ assert any("no background" in w for w in
+ export.plan_export(sras, result).warnings)
+ # Calibration must survive: v2 has no preambles and falls back to the
+ # hardcoded scope constants, and the re-encoded empty preambles must land on
+ # exactly the same fallback.
+ for ch in range(sras.n_channels):
+ assert out.cal(ch) == pytest.approx(sras.cal(ch))
+ for a in range(sras.n_angles):
+ preview = compute.apply_alignment(result, a, dc_mv(sras, a))
+ actual = dc_mv(out, a)
+ inside = preview != 0.0
+ assert actual[inside] == pytest.approx(preview[inside], abs=1e-3)
+
+
+def test_too_many_rows_is_rejected_before_writing(rig, tmp_path):
+ """The geometry table stores n_rows as a u16; silently truncating would
+ write a file whose header disagrees with its own waveform block."""
+ huge = compute.crop_alignment_result(rig.result, 0, 0, 70000, 4)
+ out_path = tmp_path / "huge.sras"
+ with pytest.raises(ValueError, match="exceeds the .sras per-angle geometry"):
+ export.write_aligned_sras(rig.sras, huge, out_path)
+ assert not out_path.exists()
+ assert not out_path.with_name(out_path.name + ".part").exists()
+
+
+def test_missing_transform_is_rejected(rig, tmp_path):
+ broken = compute.crop_alignment_result(rig.result, 0, 0,
+ *rig.result.canvas_shape)
+ del broken.per_angle[1]
+ with pytest.raises(ValueError, match="no transform for angle"):
+ export.write_aligned_sras(rig.sras, broken, tmp_path / "broken.sras")
+
+
+def test_cancelled_export_leaves_nothing_behind(rig, tmp_path):
+ out_path = tmp_path / "cancelled.sras"
+ written = export.write_aligned_sras(rig.sras, rig.result, out_path,
+ should_stop=lambda: True)
+ assert written == out_path
+ assert not out_path.exists(), "cancelled export must not leave an output file"
+ assert not out_path.with_name(out_path.name + ".part").exists()
+
+
+def test_failed_write_leaves_nothing_behind(rig, tmp_path):
+ """An exception mid-write must remove the partial file: a short .sras is
+ not detectably broken — the v6 parser reads it as an aborted scan."""
+ out_path = tmp_path / "boom.sras"
+
+ def explode(_pct):
+ raise RuntimeError("boom")
+
+ with pytest.raises(RuntimeError, match="boom"):
+ export.write_aligned_sras(rig.sras, rig.result, out_path,
+ progress_cb=explode)
+ assert not out_path.exists()
+ assert not out_path.with_name(out_path.name + ".part").exists()
+
+
+def test_progress_is_monotonic_and_completes(rig, tmp_path):
+ seen: list[int] = []
+ export.write_aligned_sras(rig.sras, rig.result, tmp_path / "prog.sras",
+ progress_cb=seen.append)
+ assert seen and seen[-1] == 100
+ assert seen == sorted(seen)
+ assert all(0 <= p <= 100 for p in seen)
+
+
+def test_band_reader_path_is_byte_identical(rig, tmp_path, monkeypatch):
+ """A source block too large to hold in RAM is served from sliding bands
+ instead. That path only runs on multi-gigabyte scans, so force it with a
+ tiny budget and demand the same bytes — otherwise the one code path that
+ matters on real data is the one never tested."""
+ whole = tmp_path / "whole.sras"
+ export.write_aligned_sras(rig.sras, rig.result, whole)
+
+ monkeypatch.setattr(compute, "_TOTAL_BYTES_BUDGET", 4096)
+ banded = tmp_path / "banded.sras"
+ export.write_aligned_sras(rig.sras, rig.result, banded)
+
+ assert banded.read_bytes() == whole.read_bytes()
+
+
+def test_row_chunking_is_invariant(rig, tmp_path, monkeypatch):
+ """Output must not depend on how many rows are buffered per write."""
+ base = tmp_path / "base.sras"
+ export.write_aligned_sras(rig.sras, rig.result, base)
+
+ monkeypatch.setattr(export, "_ROW_CHUNK", 1)
+ one = tmp_path / "one.sras"
+ export.write_aligned_sras(rig.sras, rig.result, one)
+ assert one.read_bytes() == base.read_bytes()
+
+
+def test_refuses_to_overwrite_the_source(rig):
+ """The source's waveform blocks are live read-only memmaps; writing over
+ the file would corrupt the reads the gather is making from it."""
+ with pytest.raises(ValueError, match="refusing to export onto the source"):
+ export.write_aligned_sras(rig.sras, rig.result, rig.src_path)
+ assert SrasFile(str(rig.src_path)).n_angles == rig.sras.n_angles
+
+
+def test_overwrites_an_existing_file(rig, tmp_path):
+ out_path = tmp_path / "existing.sras"
+ out_path.write_bytes(b"not a scan")
+ export.write_aligned_sras(rig.sras, rig.result, out_path)
+ assert SrasFile(str(out_path)).version == 6
+
+
+# ---------------------------------------------------------------------------
+# The registration knobs the wizard exposes
+# ---------------------------------------------------------------------------
+
+def test_locked_rotation_returns_exactly_the_seed(rig):
+ """search_deg=0 + one sign + refine=False pins rotation to the seed, which
+ is what "lock rotation to the stage angle" means on the wizard's first
+ page. Only the translation may be searched."""
+ dc4 = {a: dc_mv(rig.sras, a) for a in range(rig.sras.n_angles)}
+ for a in range(1, rig.sras.n_angles):
+ nominal = compute.nominal_delta_deg(rig.sras, a, 0)
+ fit = compute.register_angle_to_reference(
+ rig.sras, a, 0, dc4, dc_threshold_mv=_THRESHOLD_MV,
+ search_deg=0.0, coarse_step_deg=2.0, seed_signs=(-1,), refine=False)
+ assert fit.rotation_deg == pytest.approx(-nominal)
+
+
+def test_seed_deg_overrides_the_stage_angle(rig):
+ """seed_deg=0.0 searches around no rotation at all, so a scan whose angles
+ are genuinely ~37° apart must fail to find them within a ±2° window —
+ proving the seed is what positions the search."""
+ dc4 = {a: dc_mv(rig.sras, a) for a in range(rig.sras.n_angles)}
+ fit = compute.register_angle_to_reference(
+ rig.sras, 1, 0, dc4, dc_threshold_mv=_THRESHOLD_MV,
+ search_deg=2.0, seed_deg=0.0, seed_signs=(1,), refine=False)
+ truth_rot = rig.meta["truth"][1][0]
+ assert abs(fit.rotation_deg) <= 2.0
+ assert abs(fit.rotation_deg - truth_rot) > 10.0
+
+
+def test_rotation_candidates_signs():
+ both = compute._rotation_candidates(10.0, 2.0, 2.0)
+ assert both == compute._rotation_candidates(10.0, 2.0, 2.0, (-1, 1)), \
+ "default must stay the both-signs sweep"
+ assert compute._rotation_candidates(10.0, 0.0, 2.0, (1,)) == [10.0]
+ assert compute._rotation_candidates(10.0, 0.0, 2.0, (-1,)) == [-10.0]
+ # A zero seed collapses the two windows; the dedupe must keep one copy.
+ assert compute._rotation_candidates(0.0, 2.0, 2.0) == [-2.0, 0.0, 2.0]
+
+
+def test_zero_mv_fill_code_is_clipped_to_dtype():
+ """mv_to_adc is unclamped, so the fill code must be clipped or the int8
+ cast wraps around to a large-magnitude value."""
+ fake = type("S", (), dict(
+ n_channels=1, samples_per_frame=2,
+ cal=lambda self, ch: (1e-6, 0.0, 5000.0)))()
+ row = export._fill_row(fake, 3, np.dtype(np.int8))
+ assert row.shape == (1, 3, 2)
+ assert row.min() == row.max() == np.iinfo(np.int8).min
+ assert mv_to_adc(0.0, 1e-6, 0.0, 5000.0) < np.iinfo(np.int8).min
diff --git a/tests/test_alignment.py b/tests/test_alignment.py
new file mode 100644
index 0000000..d9ac0e2
--- /dev/null
+++ b/tests/test_alignment.py
@@ -0,0 +1,212 @@
+"""Angle-alignment tests: does registration actually stack the scans?
+
+Builds a synthetic scan in which one sample is imaged at several *known*
+rotations and offsets (tools/make_test_sras.write_rotating) and checks that the
+alignment path recovers them, that the shared canvas is angle 0's own pixel
+grid extended, and that nothing in the result depends on any other angle's
+stage coordinates.
+
+No Qt — this exercises sras_compute directly. See tests/test_gui.py for the
+dialog and Aligned-View plumbing.
+"""
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+
+import sras_compute as compute
+from sras_format import CH4_IDX, SrasFile, adc_to_mv
+import tools.make_test_sras as gen
+
+# Registration is limited by how far a feature moves per degree: with this
+# sample's ~1 mm radius and a ~16 µm registration pitch, a quarter degree is
+# already sub-pixel, so it is the floor of what any metric can resolve here.
+_ROT_TOL_DEG = 0.5
+_SHIFT_TOL_MM = 0.02
+_STACK_IOU_MIN = 0.90
+_THRESHOLD_MV = 80.0
+
+
+def dc4_images(sras: SrasFile) -> dict[int, np.ndarray]:
+ return {a: adc_to_mv(compute.compute_dc_image(sras, a, CH4_IDX), *sras.cal(CH4_IDX))
+ for a in range(sras.n_angles)}
+
+
+def mm_transform(sras: SrasFile, result, angle_idx: int) -> np.ndarray:
+ """Recover the pure mm-space rotation from a canvas->raw affine.
+
+ matrix == D @ R^T @ A_out, where A_out and D only carry the canvas and
+ per-angle pixel pitches; undoing both must leave something orthonormal, or
+ the transform is smuggling in a scale or a shear.
+ """
+ dx_a, dy_a = compute.pixel_pitch_mm(sras, angle_idx)
+ A_out = np.array([[0.0, result.canvas_dx_mm], [result.canvas_dy_mm, 0.0]])
+ D = np.array([[0.0, 1.0 / dy_a], [1.0 / dx_a, 0.0]])
+ return np.linalg.inv(D) @ result.per_angle[angle_idx].matrix @ np.linalg.inv(A_out)
+
+
+@pytest.fixture(scope="module")
+def rig(tmp_path_factory):
+ """The rotating-sample scan plus everything computed from it once."""
+ tmpdir = tmp_path_factory.mktemp("sras_align")
+ path = tmpdir / "rotating.sras"
+ meta = gen.write_rotating(path, n_angles=5)
+ sras = SrasFile(str(path))
+ dc4 = dc4_images(sras)
+ fits = {a: compute.register_angle_to_reference(
+ sras, a, 0, dc4, dc_threshold_mv=_THRESHOLD_MV)
+ for a in range(sras.n_angles)}
+ result = compute.compute_angle_alignment(sras, 0, _THRESHOLD_MV)
+ return SimpleNamespace(path=path, sras=sras, truth=meta["truth"],
+ dc4=dc4, fits=fits, result=result)
+
+
+def test_registration_recovers_truth(rig):
+ """Per-angle rigid registration (rotation + translation, no scale)."""
+ for a, fit in rig.fits.items():
+ t_rot, t_shift = rig.truth[a]
+ rot_err = abs(fit.rotation_deg - t_rot)
+ shift_err = float(np.hypot(fit.shift_mm[0] - t_shift[0],
+ fit.shift_mm[1] - t_shift[1]))
+ assert rot_err <= _ROT_TOL_DEG, \
+ (f"angle {a}: got {fit.rotation_deg:.3f}°, truth {t_rot:.3f}° "
+ f"(err {rot_err:.3f}°)")
+ assert shift_err <= _SHIFT_TOL_MM, f"angle {a}: err {shift_err:.4f} mm"
+ assert rig.fits[0] == compute.RigidFit(0.0, (0.0, 0.0), 1.0, "reference"), \
+ "reference angle registers as exact identity"
+
+
+def test_stage_angle_sign_is_not_trusted(rig):
+ # The stage's rotational sense relative to this module's math-positive
+ # convention is not knowable from the file, and the old code hardcoded a
+ # guess. Flipping every reported angle must therefore change nothing: the
+ # search scores both signs and the images decide.
+ flipped = SrasFile(str(rig.path))
+ flipped.angles_deg = -flipped.angles_deg
+ flipped_fits = {a: compute.register_angle_to_reference(
+ flipped, a, 0, rig.dc4, dc_threshold_mv=_THRESHOLD_MV)
+ for a in range(1, flipped.n_angles)}
+ mismatches = {a: (flipped_fits[a].rotation_deg, rig.fits[a].rotation_deg)
+ for a in flipped_fits if flipped_fits[a] != rig.fits[a]}
+ assert not mismatches, \
+ f"negating every reported stage angle changed fits: {mismatches}"
+
+
+def test_stage_coordinates_are_not_consulted(rig):
+ # Move every non-reference angle's scan window somewhere else entirely.
+ # Only angle 0's coordinates may matter, so every fit must be untouched.
+ moved = SrasFile(str(rig.path))
+ for a in range(1, moved.n_angles):
+ moved.x_start_mm[a] += 13.5 * a
+ moved.y_pos_per_angle[a] = moved.y_pos_per_angle[a] - 9.25 * a
+ moved_dc4 = dc4_images(moved)
+ moved_fits = {a: compute.register_angle_to_reference(
+ moved, a, 0, moved_dc4, dc_threshold_mv=_THRESHOLD_MV)
+ for a in range(1, moved.n_angles)}
+ mismatches = {a: (round(moved_fits[a].rotation_deg, 4), rig.fits[a].rotation_deg)
+ for a in moved_fits if moved_fits[a] != rig.fits[a]}
+ assert not mismatches, \
+ f"relocating every other angle's scan window changed fits: {mismatches}"
+
+
+def test_canvas_is_reference_grid_extended(rig):
+ sras, result = rig.sras, rig.result
+ t0 = result.per_angle[0]
+ assert np.allclose(t0.matrix, np.eye(2)), \
+ f"angle 0's transform has rotation/scale/shear: {t0.matrix}"
+ assert np.allclose(t0.offset, np.round(t0.offset)), \
+ f"angle 0 does not land on whole canvas pixels: {t0.offset}"
+ assert ((result.canvas_dx_mm, result.canvas_dy_mm)
+ == compute.pixel_pitch_mm(sras, 0)), \
+ "canvas pitch is angle 0's own pitch"
+
+ n_rows, n_cols = result.canvas_shape
+ x_axis = result.canvas_origin_mm[0] + np.arange(n_cols) * result.canvas_dx_mm
+ y_axis = result.canvas_origin_mm[1] + np.arange(n_rows) * result.canvas_dy_mm
+ row0, col0 = int(round(-t0.offset[0])), int(round(-t0.offset[1]))
+ a0_rows, a0_cols = sras.image_shape(0)
+ assert np.allclose(x_axis[col0:col0 + a0_cols], sras.x_axis_mm(0)), \
+ "canvas X axis reproduces angle 0's own X coordinates"
+ assert np.allclose(y_axis[row0:row0 + a0_rows], sras.y_positions_mm(0)), \
+ "canvas Y axis reproduces angle 0's own Y coordinates"
+ assert (n_rows >= max(int(sras.n_rows[a]) for a in range(sras.n_angles))
+ and n_cols >= max(int(sras.n_frames[a]) for a in range(sras.n_angles))), \
+ f"canvas does not cover every angle's footprint: {result.canvas_shape}"
+
+
+def test_transforms_are_pure_rotations(rig):
+ """No scaling anywhere in the per-angle transforms."""
+ for a in range(rig.sras.n_angles):
+ R = mm_transform(rig.sras, rig.result, a)
+ assert (np.allclose(R @ R.T, np.eye(2), atol=1e-9)
+ and abs(abs(np.linalg.det(R)) - 1.0) < 1e-9), \
+ f"angle {a}: det={np.linalg.det(R):.6f}"
+
+
+def test_all_angles_stack(rig):
+ aligned = {a: compute.apply_alignment(rig.result, a, rig.dc4[a])
+ for a in range(rig.sras.n_angles)}
+ base = aligned[0] >= _THRESHOLD_MV
+ for a in range(1, rig.sras.n_angles):
+ other = aligned[a] >= _THRESHOLD_MV
+ iou = float((base & other).sum()) / max(1, int((base | other).sum()))
+ assert iou >= _STACK_IOU_MIN, f"angle {a}: IoU {iou:.4f}"
+
+
+def test_downsampled_preview_lands_with_full_res(rig):
+ # The wizard reprojects block-mean-downsampled masks, so the
+ # affine has to account for the factor. When it did not, every preview
+ # layer came out magnified by that factor and offset — the overlay showed a
+ # blown-up crop of each mask, which is not something you can align by eye.
+ sras, result = rig.sras, rig.result
+ pitch = (result.canvas_dx_mm, result.canvas_dy_mm)
+ a = sras.n_angles - 1
+ p = result.per_angle[a]
+ full_mask = (rig.dc4[a] >= _THRESHOLD_MV).astype(np.float32)
+ full = compute.reproject_mask(
+ sras, a, 0, full_mask, p.rotation_deg, p.shift_mm, pitch,
+ result.canvas_origin_mm, result.canvas_shape)
+ fy, fx = 4, 16
+ small = compute.reproject_mask(
+ sras, a, 0, compute.block_mean_2d(full_mask, fy, fx),
+ p.rotation_deg, p.shift_mm, (pitch[0] * fx, pitch[1] * fy),
+ result.canvas_origin_mm,
+ (result.canvas_shape[0] // fy, result.canvas_shape[1] // fx),
+ src_downsample=(fy, fx))
+
+ # Compare in mm, via each layer's own center of mass.
+ def com_mm(layer, px, py):
+ rows, cols = np.nonzero(layer > 0.5)
+ return np.array([cols.mean() * px, rows.mean() * py])
+
+ d = com_mm(small, pitch[0] * fx, pitch[1] * fy) - com_mm(full, *pitch)
+ assert (abs(d[0]) <= abs(pitch[0] * fx) and abs(d[1]) <= abs(pitch[1] * fy)), \
+ f"downsampled preview offset {d[0]:+.4f}, {d[1]:+.4f} mm"
+
+
+def test_manual_path_reproduces_geometry(rig):
+ sras, result = rig.sras, rig.result
+ params = {a: compute.ManualAngleParams(t.rotation_deg, t.shift_mm)
+ for a, t in result.per_angle.items()}
+ manual = compute.build_manual_alignment(sras, 0, _THRESHOLD_MV, params)
+ assert (manual.canvas_shape == result.canvas_shape
+ and np.allclose(manual.canvas_origin_mm, result.canvas_origin_mm)
+ and all(np.allclose(manual.per_angle[a].matrix, result.per_angle[a].matrix)
+ and np.allclose(manual.per_angle[a].offset, result.per_angle[a].offset)
+ for a in range(sras.n_angles))), \
+ "build_manual_alignment matches compute_angle_alignment for the same params"
+
+
+def test_sidecar_roundtrip(rig):
+ sras, result = rig.sras, rig.result
+ params = {a: compute.ManualAngleParams(t.rotation_deg, t.shift_mm)
+ for a, t in result.per_angle.items()}
+ compute.save_manual_alignment(sras, 0, _THRESHOLD_MV, params)
+ loaded = compute.load_manual_alignment(sras)
+ assert (loaded is not None
+ and all(np.isclose(loaded.per_angle[a].rotation_deg, params[a].rotation_deg)
+ and np.allclose(loaded.per_angle[a].shift_mm, params[a].shift_mm)
+ for a in range(sras.n_angles))), \
+ "sidecar reloads every angle's params"
+ assert compute.delete_manual_alignment(sras), "sidecar deletes cleanly"
diff --git a/tests/test_compute.py b/tests/test_compute.py
new file mode 100644
index 0000000..3c6e400
--- /dev/null
+++ b/tests/test_compute.py
@@ -0,0 +1,356 @@
+"""Behavioural tests for the compute/format layer.
+
+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.
+"""
+
+import subprocess
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+import sras_compute as compute
+from sras_compute import (
+ cache_file, compute_dc_image, compute_rf_image, dc_image_mv,
+)
+from sras_format import CH3_IDX, CH4_IDX, SrasFile, adc_to_mv
+import tools.make_test_sras as gen
+
+REPO = Path(__file__).resolve().parent.parent
+
+
+def test_cache_roundtrip(tmp_path):
+ """v6 -> v7 for DC, then FFT, asserting the first block survives the
+ second write (the carry-forward path in write_v7_cache)."""
+ path = tmp_path / "roundtrip.sras"
+ gen.write(path, n_angles=3, seed=1, samples_per_frame=64)
+
+ src = SrasFile(str(path))
+ assert 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)
+ assert err == "", err
+
+ after_dc = SrasFile(str(path))
+ assert after_dc.version == 7, f"got v{after_dc.version}"
+ assert all(x is not None for x in after_dc.precomputed_dc3_mv)
+ assert all(np.allclose(after_dc.precomputed_dc3_mv[a], expect_dc3[a], atol=1e-4)
+ for a in range(after_dc.n_angles))
+ assert all(np.allclose(after_dc.precomputed_dc4_mv[a], expect_dc4[a], atol=1e-4)
+ for a in range(after_dc.n_angles))
+ assert all(x is None for x in after_dc.precomputed_freq_mhz), "no fft block yet"
+ assert (after_dc.precomputed_dc3_mv[0].dtype == np.float32
+ and after_dc.precomputed_dc3_mv[0].dtype.byteorder in ("=", "|")), \
+ "cached images are native float32"
+ assert after_dc.precomputed_dc3_mv[0].flags.writeable
+
+ err = cache_file(str(path), "fft", True)
+ assert err == "", err
+
+ both = SrasFile(str(path))
+ assert all(x is not None for x in both.precomputed_freq_mhz)
+ assert all(np.allclose(both.precomputed_freq_mhz[a], expect_fft[a], atol=1e-3)
+ for a in range(both.n_angles))
+ assert all(np.allclose(both.precomputed_dc3_mv[a], expect_dc3[a], atol=1e-4)
+ for a in range(both.n_angles)), \
+ "DC block carried forward through the FFT write"
+ assert 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))
+ assert 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), "cached fast path == fresh compute (unmasked)"
+ assert 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), "cached fast path == fresh compute (masked)"
+
+ # Waveform data must be byte-identical to the pre-cache file.
+ orig = tmp_path / "roundtrip_orig.sras"
+ gen.write(orig, n_angles=3, seed=1, samples_per_frame=64)
+ o, n = SrasFile(str(orig)), SrasFile(str(path))
+ assert all(np.array_equal(np.asarray(o.data[a]), np.asarray(n.data[a]))
+ for a in range(o.n_angles)), \
+ "waveform data untouched by the cache write"
+
+
+def test_partial_v7_cache(tmp_path):
+ """Only some angles cached: uncached angles must compute, not read zeros.
+ This is the v5 bug the ragged normalisation fixed, checked via v7."""
+ path = tmp_path / "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))
+ assert reread.precomputed_freq_mhz[1] is None
+ assert (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)
+ assert np.any(img1 != 0) and np.allclose(img1, expected[1], atol=1e-3), \
+ "uncached angle computes rather than returning zeros"
+
+
+def test_parallel_identity(tmp_path, monkeypatch):
+ """Forcing 1 worker vs many must give identical output — catches
+ chunk-boundary and race bugs."""
+ path = tmp_path / "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))
+
+ # Shrink the budget so the outer row loop splits into many chunks, and
+ # the block size so every chunk splits into many FFT tasks — the worst
+ # case for boundary bugs.
+ monkeypatch.setattr(compute, "_TOTAL_BYTES_BUDGET", 8 * n_frames * spf * 4)
+ monkeypatch.setattr(compute, "_FFT_BLOCK", 4)
+ fft_rows = compute._plan_fft_rows(n_frames, spf, compute._TOTAL_BYTES_BUDGET)
+ assert fft_rows < n_rows, \
+ f"FFT work actually splits into multiple chunks ({fft_rows} of {n_rows})"
+ dc_rows = compute._chunk_rows_for(n_frames, spf, compute._TOTAL_BYTES_BUDGET)
+ assert dc_rows < n_rows, \
+ f"DC work actually splits into multiple chunks ({dc_rows} of {n_rows})"
+
+ monkeypatch.setattr(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)
+ rf_pad_serial = compute_rf_image(sras, 0, dc_threshold_mv=thr,
+ apply_bg_sub=True, n_fft=spf * 8)
+
+ monkeypatch.setattr(compute, "_MAX_WORKERS", 8)
+ 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)
+ rf_pad_par = compute_rf_image(sras, 0, dc_threshold_mv=thr,
+ apply_bg_sub=True, n_fft=spf * 8)
+
+ assert np.array_equal(dc_serial, dc_par), "dc image identical"
+ assert np.array_equal(rf_serial, rf_par), "rf image identical (unmasked)"
+ assert np.array_equal(rf_masked_serial, rf_masked_par), \
+ "rf image identical (masked)"
+ assert np.array_equal(rf_pad_serial, rf_pad_par), \
+ "rf image identical (masked, padded/zoom)"
+
+
+@pytest.mark.parametrize("spf,bps", [(64, 2), (37, 1)])
+def test_zoom_identity(tmp_path, monkeypatch, spf, bps):
+ """The zoom peak search must reproduce the full padded-rfft argmax
+ bit-for-bit, across pad factors, masking, bg-sub, dtype, and backend."""
+ path = tmp_path / f"zoom_{spf}.sras"
+ gen.write(path, n_angles=2, seed=6, samples_per_frame=spf, bps=bps)
+ sras = SrasFile(str(path))
+ dc4 = dc_image_mv(sras, 0, CH4_IDX)
+ thr = float(np.median(dc4))
+
+ backends = ["scipy"] + (["pyfftw"] if compute.PYFFTW_AVAILABLE else [])
+ for backend in backends:
+ monkeypatch.setattr(compute, "_fft_backend", backend)
+ for pad in (4, 8, 40):
+ n_fft = spf * pad
+ for thr_v in (None, thr):
+ for bg in (False, True):
+ ref = compute_rf_image(sras, 0, dc_threshold_mv=thr_v,
+ apply_bg_sub=bg, n_fft=n_fft,
+ exact=True)
+ zoom = compute_rf_image(sras, 0, dc_threshold_mv=thr_v,
+ apply_bg_sub=bg, n_fft=n_fft)
+ diff = int((ref != zoom).sum())
+ assert diff == 0, \
+ (f"{diff} px differ: backend={backend} pad={pad} "
+ f"thr={thr_v} bg={bg} spf={spf}")
+
+ # A threshold above every pixel masks everything: both paths must agree
+ # on an all-zero image.
+ all_masked = compute_rf_image(sras, 0, dc_threshold_mv=1e9, n_fft=spf * 8)
+ assert not all_masked.any()
+
+
+def test_zoom_identity_fuzz():
+ """Hammer _peak_bins_zoom directly with adversarial spectra: noise,
+ un-subtracted DC offsets, on-bin and off-bin tones, near-tie tone pairs,
+ and all-zero rows."""
+ import scipy.fft as scipy_fft
+
+ rng = np.random.default_rng(42)
+ for _ in range(25):
+ spf = int(rng.integers(16, 220))
+ pad = int(rng.choice([4, 5, 8, 16, 40]))
+ n_fft = spf * pad
+ n_wf = 24
+ w = rng.normal(scale=20.0, size=(n_wf, spf))
+ t = np.arange(spf)
+ # rows 0-5: pure/noisy tones (some off-bin), row 6-7: near-tie pair,
+ # row 8: big DC offset, row 9: all zeros, rest: plain noise.
+ for r in range(6):
+ f = rng.uniform(1.0, spf / 2 - 1)
+ w[r] = 60 * np.sin(2 * np.pi * f * t / spf) + w[r] * (r % 2)
+ f1, f2 = rng.uniform(2.0, spf / 2 - 2, size=2)
+ w[6] = 50 * np.sin(2 * np.pi * f1 * t / spf) \
+ + 49.9 * np.sin(2 * np.pi * f2 * t / spf)
+ w[7] = 50 * np.sin(2 * np.pi * f1 * t / spf) \
+ + 50 * np.cos(2 * np.pi * f2 * t / spf)
+ w[8] = 90 + rng.normal(scale=5.0, size=spf)
+ w[9] = 0.0
+ w = w.astype(np.float32)
+
+ S = scipy_fft.rfft(w, n=n_fft, axis=-1, workers=1)
+ P = S.real ** 2
+ P += S.imag ** 2
+ P[:, 0] = 0.0
+ ref = np.argmax(P, axis=1)
+
+ zp = compute._zoom_plan(spf, n_fft)
+ got = compute._peak_bins_zoom(w, zp)
+ bad = np.nonzero(ref != got)[0]
+ assert not len(bad), \
+ (f"spf={spf} pad={pad}: rows {bad.tolist()} picked "
+ f"{got[bad].tolist()} instead of {ref[bad].tolist()}")
+
+
+def test_nomask_equals_low_threshold(tmp_path):
+ """dc_threshold_mv=None must equal a threshold below every pixel, while
+ skipping the CH4 read."""
+ path = tmp_path / "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)
+ assert np.array_equal(none_img, low_img), \
+ f"angle {a}: None == -1e9 threshold"
+ assert len(np.unique(none_img)) > 1, \
+ f"angle {a}: image is degenerate ({len(np.unique(none_img))} unique)"
+
+
+def test_roi_mask():
+ """The bbox-restricted mask must equal a full-grid point-in-polygon test."""
+ 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)
+ assert np.array_equal(fast, slow), f"{name} ({int(slow.sum())} px inside)"
+
+ # 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)
+ assert np.array_equal(fast, slow), "descending y axis"
+
+
+def test_legacy_parse(tmp_path):
+ """v2-v4 parsing against known written data."""
+ for version in (2, 3, 4):
+ path = tmp_path / 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))
+ assert s.version == version, f"got v{s.version}"
+ assert list(s.n_rows) == [4, 4] and list(s.n_frames) == [12, 12], \
+ f"rows={list(s.n_rows)} frames={list(s.n_frames)}"
+ assert all(np.array_equal(np.asarray(s.data[a]), meta["data"][a])
+ for a in range(s.n_angles)), \
+ f"v{version} waveform data matches what was written"
+ assert (s.background is not None) == (version >= 4), \
+ f"v{version} background {'present' if version >= 4 else 'absent'}"
+ assert (isinstance(s.precomputed_freq_mhz, list)
+ and len(s.precomputed_freq_mhz) == s.n_angles), \
+ f"v{version} precomputed stores are ragged lists"
+ # DC image must equal a direct mean of the known input.
+ expect = meta["data"][0][:, CH3_IDX, :, :].astype(np.float64).mean(axis=-1)
+ assert np.allclose(compute_dc_image(s, 0, CH3_IDX), expect, atol=1e-3), \
+ f"v{version} DC image equals a direct mean"
+
+
+def test_sras_average(tmp_path):
+ """The sras_average.py CLI: frame averaging with remainder handling."""
+ src = tmp_path / "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 = tmp_path / "legacy_v4_avg.sras"
+ proc = subprocess.run(
+ [sys.executable, str(REPO / "sras_average.py"), str(src), str(dst), "--n", "4"],
+ capture_output=True, text=True, cwd=REPO)
+ assert proc.returncode == 0, (proc.stderr or proc.stdout).strip()[-200:]
+
+ avg = SrasFile(str(dst))
+ assert avg.version == 4
+ assert list(avg.n_frames) == [3, 3], f"{list(avg.n_frames)}"
+ assert (avg.n_angles == 2 and list(avg.n_rows) == [4, 4]
+ and avg.n_channels == meta["n_channels"])
+ assert np.allclose(avg.ch_ymult_mv, SrasFile(str(src)).ch_ymult_mv), \
+ "calibration preserved"
+ assert np.array_equal(avg.background, SrasFile(str(src)).background), \
+ "background preserved"
+ src_data = meta["data"]
+ expect0 = src_data[0][:, :, 0:4, :].astype(np.float32).mean(axis=2).astype(np.int16)
+ assert np.array_equal(np.asarray(avg.data[0])[:, :, 0, :], expect0), \
+ "first averaged group equals the mean of its 4 source frames"
+
+ # Remainder handling: 12 frames / 5 -> 2 full groups + 1 partial.
+ dst2 = tmp_path / "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)
+ assert list(SrasFile(str(dst2)).n_frames) == [3, 3], \
+ "partial trailing group kept by default"
+ dst3 = tmp_path / "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)
+ assert list(SrasFile(str(dst3)).n_frames) == [2, 2], \
+ "--discard-remainder drops the partial group"
+
+
+def test_unsupported_version_reported(tmp_path):
+ """cache_file must report, not raise, for a file it can't handle."""
+ bogus = tmp_path / "bogus.sras"
+ bogus.write_bytes(b"SRAS" + bytes([99]) + b"\x00" * 200)
+ err = cache_file(str(bogus), "dc", True)
+ assert err, "bad version returns an error string"
+ missing = cache_file(str(tmp_path / "does_not_exist.sras"), "dc", True)
+ assert missing, "missing file returns an error string"
diff --git a/tests/test_fused_export.py b/tests/test_fused_export.py
new file mode 100644
index 0000000..ebd8f22
--- /dev/null
+++ b/tests/test_fused_export.py
@@ -0,0 +1,75 @@
+"""SrasFile.angles_share_raw_grid(): the no-alignment-needed gating path for
+Export Fused ROI.
+
+A plain multi-angle scan gives each angle its own bounding box and stage
+x_start (scan_format.md's whole reason v6 geometry is per-angle), so it must
+read as "not shareable" without a live alignment. A file the viewer's own
+Alignment Wizard exported repeats one Per-Angle Geometry record and one Row
+Table span for every angle (scan_format.md, "Files written by the viewer's
+Alignment Wizard"), so it must read as "shareable" with no alignment needed
+at all.
+
+No Qt: this exercises sras_format/sras_compute/sras_align_export directly,
+mirroring tests/test_align_export.py.
+"""
+
+import sras_align_export as export
+import sras_compute as compute
+from sras_format import SrasFile
+import tools.make_test_sras as gen
+
+_THRESHOLD_MV = 80.0
+
+
+def test_single_angle_file_always_shares_its_grid(tmp_path):
+ path = tmp_path / "one_angle.sras"
+ gen.write(path, n_angles=1)
+ sras = SrasFile(str(path))
+ assert sras.angles_share_raw_grid()
+
+
+def test_plain_multi_angle_file_does_not_share_its_grid(tmp_path):
+ """tools.make_test_sras.write gives every angle its own geometry and
+ stage x_start (build()'s `x_start = -0.5 + 0.1 * a`), matching how real
+ v6 scans vary per angle — so this must read as "not shareable"."""
+ path = tmp_path / "plain.sras"
+ gen.write(path, n_angles=3)
+ sras = SrasFile(str(path))
+ assert not sras.angles_share_raw_grid()
+
+
+def test_wizard_exported_file_shares_its_grid(tmp_path):
+ src_path = tmp_path / "rotating.sras"
+ meta = gen.write_rotating(src_path, n_angles=4)
+ sras = SrasFile(str(src_path))
+ params = {a: compute.ManualAngleParams(rot, shift)
+ for a, (rot, shift) in meta["truth"].items()}
+ result = compute.build_manual_alignment(sras, 0, _THRESHOLD_MV, params)
+
+ out_path = tmp_path / "rotating_aligned.sras"
+ export.write_aligned_sras(sras, result, out_path)
+ out = SrasFile(str(out_path))
+
+ assert not sras.angles_share_raw_grid(), (
+ "sanity check: the *source* rotating scan must NOT already share a "
+ "grid, or this test would not actually exercise the wizard export")
+ assert out.angles_share_raw_grid()
+
+
+def test_mutated_angle_breaks_the_shared_grid(tmp_path):
+ """A loaded SrasFile's x_start_mm is a public per-angle array (as
+ tests/test_gui.py::test_alignment_geometry_is_stage_independent also
+ relies on) -- mutating one angle's start must be visible here too."""
+ src_path = tmp_path / "rotating.sras"
+ meta = gen.write_rotating(src_path, n_angles=3)
+ sras = SrasFile(str(src_path))
+ params = {a: compute.ManualAngleParams(rot, shift)
+ for a, (rot, shift) in meta["truth"].items()}
+ result = compute.build_manual_alignment(sras, 0, _THRESHOLD_MV, params)
+ out_path = tmp_path / "rotating_aligned.sras"
+ export.write_aligned_sras(sras, result, out_path)
+ out = SrasFile(str(out_path))
+ assert out.angles_share_raw_grid()
+
+ out.x_start_mm[1] += 1.0
+ assert not out.angles_share_raw_grid()
diff --git a/tests/test_gui.py b/tests/test_gui.py
new file mode 100644
index 0000000..4b184ff
--- /dev/null
+++ b/tests/test_gui.py
@@ -0,0 +1,881 @@
+"""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, the alignment wizard end to end (pre-rotation, correlation, manual
+nudging, crop, export), aligned view, ROI draw/move, and CSV export.
+
+NOTE: this module is one ordered integration sequence over a single shared
+window — the tests build on each other's state and must run in definition
+order (pytest's default within a module). Run the whole module, not single
+tests.
+"""
+
+import json
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import numpy as np
+import pytest
+from PyQt6.QtCore import QEventLoop, Qt, QTimer
+from PyQt6.QtTest import QTest
+from PyQt6.QtWidgets import QApplication, QDialog, QMessageBox
+
+import sras_compute as compute
+from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, SrasFile
+from sras_viewer import (
+ FusedRoiExportDialog, RoiQuad, SrasViewerWindow, VELOCITY_MODE_IDX,
+)
+import tools.make_test_sras as gen
+
+
+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()
+
+
+@pytest.fixture(scope="module")
+def ctx(tmp_path_factory):
+ """The shared window, test file, and cross-test state for the sequence."""
+ app = QApplication.instance() or QApplication([])
+ tmpdir = tmp_path_factory.mktemp("sras_gui")
+ path = tmpdir / "gui.sras"
+ gen.write(path, n_angles=4, seed=11, samples_per_frame=256)
+
+ win = SrasViewerWindow()
+ win.show()
+ errors: list[str] = []
+ # 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)
+
+ c = SimpleNamespace(app=app, win=win, path=path, tmpdir=tmpdir,
+ errors=errors, s=None)
+ yield c
+ if win.isVisible():
+ win.close()
+ pump(400)
+
+
+def test_load(ctx):
+ win = ctx.win
+ win._load_file(str(ctx.path))
+ assert wait_until(lambda: win._sras is not None), "file loaded"
+ ctx.s = s = win._sras
+ assert s.version == 6, f"v{s.version}"
+ assert win.combo_channel.currentIndex() == CH4_IDX, "defaults to CH4"
+ assert win._current_image is not None, "image displayed"
+ assert win.spin_angle.maximum() == s.n_angles - 1, \
+ "angle spinbox ranges over all angles"
+ assert win._info["Angles"].text() == f"Angles: {s.n_angles}", \
+ win._info["Angles"].text()
+
+
+def test_dc_precompute_all_angles(ctx):
+ win, s = ctx.win, ctx.s
+ 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)))
+ assert ok, f"every angle cached for CH3 and CH4 ({len(win._dc_cache)} entries)"
+ assert "ready for all angles" in win.lbl_dc_precompute.text(), \
+ win.lbl_dc_precompute.text()
+
+
+def test_angle_switching_from_cache(ctx):
+ win, s = ctx.win, ctx.s
+ for a in range(s.n_angles):
+ win.spin_angle.setValue(a)
+ win._on_view_changed()
+ pump(60)
+ expected = win._sras.image_shape(a)
+ assert win._current_image.shape == expected, \
+ f"angle {a} shows its own geometry {expected}, got {win._current_image.shape}"
+ assert not win._job_running("compute"), \
+ "no compute job needed for cached DC angles"
+
+
+def test_stepping_the_angle_spinbox_redraws(ctx):
+ """Clicking the angle spinbox's arrows (or pressing Up/Down in it) must
+ move the display, not just the number.
+
+ This is the ordinary way to walk a scan, and it used to do nothing: the
+ spinbox was wired on editingFinished, which QAbstractSpinBox emits only
+ on Return or focus-out — never on a step. Every other test in this module
+ called _on_view_changed() by hand and so could not have caught it.
+ """
+ win, s = ctx.win, ctx.s
+ assert s.n_angles >= 3, "need room to step in both directions"
+
+ win.spin_angle.setValue(0)
+ assert wait_until(lambda: win._current_angle == 0), "settled on angle 0"
+
+ for expected in range(1, s.n_angles):
+ win.spin_angle.stepUp()
+ assert wait_until(lambda e=expected: win._current_angle == e), \
+ f"stepping up to angle {expected} redrew the display"
+
+ win.spin_angle.stepDown()
+ assert wait_until(lambda: win._current_angle == s.n_angles - 2), \
+ "stepping down redraws too"
+
+ # Keyboard stepping goes through the same signal, so it must work as well.
+ QTest.keyClick(win.spin_angle, Qt.Key.Key_Down)
+ assert wait_until(lambda: win._current_angle == s.n_angles - 3), \
+ "Key_Down redraws"
+
+ win.spin_angle.setValue(0)
+ assert wait_until(lambda: win._current_angle == 0), "back to angle 0"
+
+
+def test_typing_an_angle_does_not_compute_intermediate_angles(ctx):
+ """Keyboard tracking must stay off: with it on, valueChanged fires per
+ keystroke, so typing "12" would dispatch a compute for angle 1 first —
+ on a real scan, a whole wasted FFT for an angle the user never asked for.
+ """
+ win, s = ctx.win, ctx.s
+ assert not win.spin_angle.keyboardTracking(), \
+ "keyboard tracking off is what makes valueChanged safe to connect"
+
+ target = s.n_angles - 1
+ assert target >= 2, "need a multi-digit-ish range to make the point"
+ win.spin_angle.setValue(0)
+ wait_until(lambda: win._current_angle == 0)
+
+ seen = []
+ win.spin_angle.valueChanged.connect(seen.append)
+ try:
+ win.spin_angle.lineEdit().selectAll()
+ QTest.keyClicks(win.spin_angle, str(target))
+ pump(60)
+ assert seen == [], f"no signal while typing, got {seen}"
+ QTest.keyClick(win.spin_angle, Qt.Key.Key_Return)
+ pump(60)
+ assert seen == [target], f"one signal on commit, got {seen}"
+ finally:
+ win.spin_angle.valueChanged.disconnect(seen.append)
+ assert wait_until(lambda: win._current_angle == target), "committed angle shown"
+
+ win.spin_angle.setValue(0)
+ assert wait_until(lambda: win._current_angle == 0), "back to angle 0"
+
+
+def test_channel_switching(ctx):
+ win = ctx.win
+ win.spin_angle.setValue(0)
+ win._on_view_changed()
+ pump(60)
+ win.combo_channel.setCurrentIndex(CH3_IDX)
+ assert wait_until(lambda: win._current_ch == CH3_IDX), "CH3 displayed"
+
+ win.combo_channel.setCurrentIndex(CH1_IDX)
+ assert wait_until(
+ lambda: win._current_ch == CH1_IDX and not win._job_running("compute")), \
+ "CH1 (FFT) computed"
+ assert len(win._fft_cache) > 0, "FFT result cached"
+ ctx.rf_img = win._current_image
+ assert len(np.unique(ctx.rf_img)) > 1, \
+ f"FFT image is degenerate ({len(np.unique(ctx.rf_img))} unique values)"
+
+
+def test_velocity_mode(ctx):
+ """Velocity mode is a pure post-multiply, no recompute."""
+ win = ctx.win
+ ctx.n_fft_before = len(win._fft_cache)
+ win.combo_channel.setCurrentIndex(VELOCITY_MODE_IDX)
+ assert wait_until(
+ lambda: win._current_ch == VELOCITY_MODE_IDX
+ and not win._job_running("compute")), "velocity displayed"
+ grating = win.spin_grating_um.value()
+ assert np.allclose(win._current_image, ctx.rf_img * grating, atol=1e-3), \
+ "velocity == freq x grating"
+ assert len(win._fft_cache) == ctx.n_fft_before, \
+ f"velocity reused the cached FFT ({ctx.n_fft_before} -> {len(win._fft_cache)})"
+ assert win.grp_velocity.isVisible(), "grating spinbox visible in velocity mode"
+
+
+def test_threshold_change_recomputes(ctx):
+ """A threshold change is a genuine cache-key change."""
+ win = ctx.win
+ 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()
+ assert wait_until(
+ lambda: not win._job_running("compute")
+ and len(win._fft_cache) > ctx.n_fft_before), "recomputed at new threshold"
+ n_zero = int((win._current_image == 0).sum())
+ assert n_zero > 0, \
+ f"masking zeroed some pixels ({n_zero} of {win._current_image.size})"
+
+
+def test_bg_sub_toggle(ctx):
+ """bg-sub no longer gates the display: it only affects a future live
+ compute for an angle with nothing cached yet, or an explicit batch
+ recompute. Toggling it on an angle that already has an FFT image must
+ leave that image on screen, untouched."""
+ win = ctx.win
+ n_before = len(win._fft_cache)
+ img_before = win._current_image
+ win.chk_bg_sub.setChecked(False)
+ pump(200)
+ assert not win._job_running("compute"), \
+ "toggling bg-sub alone must not dispatch a recompute"
+ assert len(win._fft_cache) == n_before, "no new cache entry from the toggle"
+ assert np.array_equal(win._current_image, img_before), \
+ "displayed image unchanged by the bg-sub toggle"
+ win.chk_bg_sub.setChecked(True)
+ pump(200)
+ assert not win._job_running("compute")
+ assert len(win._fft_cache) == n_before
+ assert np.array_equal(win._current_image, img_before)
+
+
+def test_roi_and_csv_export(ctx):
+ win, s = ctx.win, ctx.s
+ 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)
+ assert win.image_canvas.get_roi() is not None, "ROI registered"
+ assert ("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])
+ assert 0 < npix <= win._current_image.size, f"{npix}"
+ assert win.btn_export_roi.isEnabled(), "Export ROI enabled"
+
+ csv_path = ctx.tmpdir / "roi.csv"
+ with patch("sras_viewer.main_window.QFileDialog.getSaveFileName",
+ return_value=(str(csv_path), "")):
+ win._on_export_roi_csv()
+ assert csv_path.exists(), "ROI CSV written"
+ body = [l for l in csv_path.read_text().splitlines() if not l.startswith("#")]
+ assert len(body) == npix + 1, \
+ f"ROI CSV has {len(body)} lines for {npix} pixels (want header + one per pixel)"
+
+ img_csv = ctx.tmpdir / "img.csv"
+ with patch("sras_viewer.main_window.QFileDialog.getSaveFileName",
+ return_value=(str(img_csv), "")):
+ win._on_export_csv()
+ assert img_csv.exists(), "image CSV written"
+ arr = np.loadtxt(img_csv, delimiter=",")
+ assert (arr.shape == win._current_image.shape
+ and np.allclose(arr, win._current_image, rtol=1e-5, atol=1e-4)), \
+ "image CSV round-trips the displayed image"
+
+
+def test_roi_survives_switches(ctx):
+ win = ctx.win
+ win.spin_angle.setValue(1)
+ win._on_view_changed()
+ wait_until(lambda: not win._job_running("compute"))
+ assert win.image_canvas.get_roi() is not None, \
+ "ROI still present after angle switch"
+ win.combo_channel.setCurrentIndex(CH4_IDX)
+ wait_until(lambda: win._current_ch == CH4_IDX)
+ assert win.image_canvas.get_roi() is not None, \
+ "ROI still present after channel switch"
+
+
+def test_alignment_geometry_is_stage_independent(ctx):
+ """Local mm is anchored on each angle's array center, not its stage
+ position: that is what makes a scan's placement independent of where its
+ window happened to sit. (Registration accuracy itself is covered by
+ tests/test_alignment.py, which has a synthetic sample to register.)"""
+ win, s = ctx.win, ctx.s
+ win.spin_angle.setValue(0)
+ win._on_view_changed()
+ wait_until(lambda: not win._job_running("compute"))
+ assert win._wizard_act.isEnabled(), "alignment wizard action enabled"
+
+ n_rows, n_frames = s.image_shape(0)
+ assert np.allclose(compute._center_idx(s, 0),
+ [(n_rows - 1) / 2, (n_frames - 1) / 2]), \
+ "array center is the geometric center of the pixel grid"
+ dx0, dy0 = compute.pixel_pitch_mm(s, 0)
+ assert np.allclose(compute._local_half_extent_mm(s, 0),
+ [(n_frames - 1) / 2 * abs(dx0), (n_rows - 1) / 2 * abs(dy0)]), \
+ "local half-extent is derived from shape and pitch alone"
+ identity = {a: compute.ManualAngleParams() for a in range(s.n_angles)}
+ origin_a, shape_a = compute.canvas_for_params(s, 0, (dx0, dy0), identity)
+ moved = SrasFile(str(ctx.path))
+ for a in range(1, moved.n_angles):
+ moved.x_start_mm[a] += 7.5
+ moved.y_pos_per_angle[a] = moved.y_pos_per_angle[a] + 3.25
+ origin_b, shape_b = compute.canvas_for_params(moved, 0, (dx0, dy0), identity)
+ assert shape_a == shape_b and np.allclose(origin_a, origin_b), \
+ ("moving every non-reference angle's scan window must leave the canvas "
+ f"unchanged: {origin_a} {shape_a} vs {origin_b} {shape_b}")
+
+ # Both signs of the stage's reported angle are searched by default.
+ cands = compute._rotation_candidates(30.0, 6.0, 2.0)
+ assert min(cands) < -29.0 and max(cands) > 29.0, f"{min(cands)}..{max(cands)}"
+
+ # Whole-pixel translation must not wrap content around the edge.
+ arr = np.zeros((6, 6), dtype=np.float32)
+ arr[0, 0] = 1.0
+ assert compute._shift_into(arr, -1, -1).sum() == 0.0, \
+ "_shift_into zero-fills rather than wrapping"
+ assert compute._shift_into(arr, 2, 3)[2, 3] == 1.0, \
+ "_shift_into moves content by exactly the requested offset"
+
+
+def test_wizard_opens_prerotated(ctx):
+ """The wizard shows a mask stack before any correlation has run, built from
+ the stage angles in the file — the "pre-rotate" step. Nothing may seed from
+ the still-live automatic result; only a saved sidecar."""
+ win, s = ctx.win, ctx.s
+ win._on_alignment_wizard()
+ assert win._align_wizard is not None, "wizard opened"
+ ctx.wiz = wiz = win._align_wizard
+ ctx.p1 = p1 = wiz.page(wiz.PAGE_CORRELATE)
+
+ assert not win._job_running("align_masks"), \
+ "mask prep needed no background worker (already DC-cached)"
+ assert wait_until(lambda: p1.isComplete()), "masks ready, Next enabled"
+ assert not win._wizard_act.isEnabled(), \
+ "wizard action disabled while a wizard is open"
+
+ st = wiz.state
+ assert st.result is not None, "an AlignmentResult exists from pre-rotation alone"
+ assert st.counts is not None and st.counts.shape == st.preview_shape
+ assert 0 <= st.counts.max() <= s.n_angles
+ assert not st.fits, "no fits before a correlation has run"
+ for a in range(s.n_angles):
+ nominal = compute.nominal_delta_deg(s, a, st.ref_angle_idx)
+ expected = 0.0 if a == st.ref_angle_idx else nominal
+ assert abs(st.params[a].rotation_deg - expected) < 1e-9, \
+ f"angle {a} not pre-rotated to its stage angle"
+ assert st.params[a].shift_mm == (0.0, 0.0), \
+ "pre-rotation must not invent a translation"
+
+
+def test_wizard_reference_angle_is_locked(ctx):
+ p1, wiz = ctx.p1, ctx.wiz
+ p1.combo_active.setCurrentIndex(wiz.state.ref_angle_idx)
+ pump(30)
+ before = wiz.state.params[wiz.state.ref_angle_idx]
+ p1._on_nudge_translate(1, 0, False)
+ p1._on_nudge_rotate(1, False)
+ assert wiz.state.params[wiz.state.ref_angle_idx] == before, \
+ "reference angle untouched by nudge attempts"
+
+
+def test_wizard_nudges(ctx):
+ """Manual correction, which the wizard absorbed from the old dialog."""
+ p1, wiz, s = ctx.p1, ctx.wiz, ctx.s
+ ctx.active = active = 1 if s.n_angles > 1 else 0
+ p1.combo_active.setCurrentIndex(active)
+ pump(30)
+
+ before = wiz.state.params[active].shift_mm
+ p1._on_nudge_translate(1, 0, False)
+ fine = p1.spin_step_translate.value()
+ assert abs(wiz.state.params[active].shift_mm[0] - (before[0] + fine)) < 1e-9, \
+ "fine translate nudge moved shift_x by exactly one fine step"
+
+ before = wiz.state.params[active].shift_mm
+ p1._on_nudge_translate(0, -1, True)
+ coarse = fine * p1.spin_step_mult.value()
+ assert abs(wiz.state.params[active].shift_mm[1] - (before[1] - coarse)) < 1e-9, \
+ "coarse translate nudge uses the multiplier"
+
+ before_rot = wiz.state.params[active].rotation_deg
+ p1._on_nudge_rotate(1, False)
+ assert wiz.state.params[active].rotation_deg != before_rot
+ assert len(wiz.state.layers) == s.n_angles, \
+ "stack rebuilt for every angle after a rotation nudge"
+
+ # A nudge only reprojects the angle that moved, patching the overlap counts
+ # in place. That shortcut is only sound if it lands on exactly what a full
+ # rebuild would have produced.
+ incremental = wiz.state.counts.copy()
+ wiz.rebuild_stack()
+ assert np.array_equal(wiz.state.counts, incremental), \
+ "incremental nudge update matches a full stack rebuild"
+
+ # Real key-event wiring (keyPressEvent -> signal -> slot).
+ before = wiz.state.params[active].shift_mm
+ QTest.keyClick(p1.canvas, Qt.Key.Key_Right)
+ assert wiz.state.params[active].shift_mm[0] > before[0], \
+ "a real Right-arrow key event nudged shift_x"
+
+ # Both views render from the same reprojected layers.
+ p1.combo_view.setCurrentIndex(1)
+ pump(50)
+ p1.combo_view.setCurrentIndex(0)
+ pump(50)
+
+
+def test_wizard_correlate(ctx):
+ """Cross-correlation, for every source option, retryable."""
+ win, p1, wiz, s = ctx.win, ctx.p1, ctx.wiz, ctx.s
+ from sras_viewer.align_wizard import _CORRELATE_SOURCES
+
+ for idx, (label, _sources) in enumerate(_CORRELATE_SOURCES):
+ p1.combo_source.setCurrentIndex(idx)
+ p1.btn_correlate.click()
+ assert not p1.isComplete(), \
+ f"Next must be disabled while correlating ({label})"
+ assert wait_until(lambda: not win._job_running("align_correlate"),
+ timeout_ms=60000), f"correlation finished ({label})"
+ assert p1.isComplete(), f"Next re-enabled ({label})"
+ assert all(a in wiz.state.fits for a in range(s.n_angles)
+ if a != wiz.state.ref_angle_idx), \
+ f"every non-reference angle got a fit ({label})"
+
+ assert wiz.state.params[wiz.state.ref_angle_idx] == compute.ManualAngleParams(), \
+ "reference angle stays identity after correlation"
+ assert p1.btn_correlate.isEnabled(), "controls re-enabled when done"
+ assert p1.table.rowCount() == s.n_angles and p1.table.item(0, 0) is not None, \
+ "per-angle fit table populated"
+ assert p1.lbl_overlap.text(), "overlap summary reported"
+
+ r = wiz.state.result
+ assert len(r.per_angle) == s.n_angles, "transform for every angle"
+ assert r.per_angle[r.ref_angle_idx].shift_mm == (0.0, 0.0), \
+ "reference angle has zero shift"
+
+
+def test_wizard_retry_changes_geometry(ctx):
+ """Editing a parameter and re-running is the retry path, and it must
+ invalidate anything indexed against the old canvas."""
+ p1, wiz = ctx.p1, ctx.wiz
+ gen_before = wiz.state.geometry_generation
+ p1.spin_threshold.setValue(p1.spin_threshold.value() + 5.0)
+ p1.spin_threshold.editingFinished.emit()
+ pump(60)
+ assert wiz.state.geometry_generation > gen_before, \
+ "a threshold change rebuilt the geometry"
+
+ # Reset drops the fits and returns to pre-rotation only.
+ p1.btn_reset.click()
+ pump(60)
+ assert not wiz.state.fits, "reset cleared the fits"
+ nominal = compute.nominal_delta_deg(ctx.s, ctx.active, wiz.state.ref_angle_idx)
+ assert abs(wiz.state.params[ctx.active].rotation_deg - nominal) < 1e-9
+ assert wiz.state.params[ctx.active].shift_mm == (0.0, 0.0), \
+ "reset also drops nudged translation"
+
+ # Put a real correlation back for the pages that follow.
+ p1.btn_correlate.click()
+ assert wait_until(lambda: not ctx.win._job_running("align_correlate"),
+ timeout_ms=60000)
+
+
+def test_wizard_roi_page(ctx):
+ """The crop page: presets, and two-way sync between the drawn rectangle and
+ the numeric canvas-pixel boxes."""
+ wiz = ctx.wiz
+ wiz.next()
+ pump(150)
+ assert wiz.currentId() == wiz.PAGE_ROI, "advanced to the ROI page"
+ ctx.p2 = p2 = wiz.page(wiz.PAGE_ROI)
+ st = wiz.state
+
+ assert st.crop is not None and p2.isComplete(), \
+ "a default crop is offered on entry"
+ n_rows, n_cols = st.result.canvas_shape
+ assert st.crop[2] > 1 or n_rows == 1, \
+ f"default crop must not collapse to a single row: {st.crop}"
+
+ p2.btn_whole.click()
+ pump(50)
+ assert st.crop == (0, 0, n_rows, n_cols), "whole-canvas preset"
+
+ p2.btn_fit_union.click()
+ pump(50)
+ assert st.counts[st.crop[0]:st.crop[0] + st.crop[2],
+ st.crop[1]:st.crop[1] + st.crop[3]].sum() == st.counts.sum(), \
+ "fit-to-union must keep every covered pixel"
+
+ if p2.btn_fit_overlap.isEnabled():
+ p2.btn_fit_overlap.click()
+ pump(50)
+ row0, col0, nr, nc = st.crop
+ assert (st.counts[row0:row0 + nr, col0:col0 + nc] >= 1).all(), \
+ "full-overlap crop must not include uncovered pixels"
+
+ # Numeric -> drawn rectangle.
+ p2.btn_whole.click()
+ pump(50)
+ target = (0, 0, max(1, n_rows // 2), max(1, n_cols // 2))
+ p2.spin_rows.setValue(target[2])
+ p2.spin_cols.setValue(target[3])
+ pump(50)
+ assert st.crop == target, f"spin boxes drive the crop: {st.crop} vs {target}"
+
+ # Drawn rectangle -> numeric, round-tripping exactly.
+ x0, y0 = wiz.canvas_to_mm(target[1] - 0.5, target[0] - 0.5)
+ x1, y1 = wiz.canvas_to_mm(target[1] + target[3] - 0.5,
+ target[0] + target[2] - 0.5)
+ p2.canvas.set_roi(RoiQuad.from_bbox(min(x0, x1), min(y0, y1),
+ max(x0, x1), max(y0, y1)))
+ pump(80)
+ assert st.crop == target, \
+ f"drawn rectangle round-trips to the same crop: {st.crop} vs {target}"
+
+ # A degenerate crop blocks Next.
+ st.crop = None
+ p2.completeChanged.emit()
+ assert not p2.isComplete(), "an absent crop blocks Next"
+ p2._set_crop(*target)
+ assert p2.isComplete()
+ ctx.crop = target
+
+
+def test_wizard_crop_dropped_when_going_back(ctx):
+ """A crop is canvas-pixel indexed, so it cannot survive a re-correlation."""
+ wiz, p2 = ctx.wiz, ctx.p2
+ wiz.back()
+ pump(120)
+ assert wiz.currentId() == wiz.PAGE_CORRELATE
+ assert wiz.state.crop is None, "cleanupPage discarded the stale crop"
+ assert wiz.cropped_plan() == (None, None), \
+ "nothing derived from the dropped crop survives either"
+ wiz.next()
+ pump(150)
+ assert wiz.state.crop is not None, "a fresh default crop is offered again"
+ p2._set_crop(*ctx.crop)
+
+
+def test_wizard_export(ctx):
+ """Writing the file: Finish stays unavailable until a write succeeds."""
+ win, wiz = ctx.win, ctx.wiz
+ out = ctx.tmpdir / "wizard_aligned.sras"
+ with patch("sras_viewer.align_wizard.QMessageBox.question",
+ return_value=QMessageBox.StandardButton.Yes):
+ wiz.next()
+ pump(150)
+ assert wiz.currentId() == wiz.PAGE_SAVE, "advanced to the save page"
+ ctx.p3 = p3 = wiz.page(wiz.PAGE_SAVE)
+ cropped, _ = wiz.cropped_plan()
+ assert cropped is not None, "crop applied on leaving page 2"
+ assert cropped.canvas_shape == ctx.crop[2:], \
+ "cropped result carries the chosen shape"
+ assert not p3.isComplete(), "Finish unavailable before anything is written"
+ assert p3.lbl_summary.text(), "a summary of what will be written is shown"
+
+ with patch("sras_viewer.align_wizard.QFileDialog.getSaveFileName",
+ return_value=(str(out), "")):
+ p3.btn_browse.click()
+ assert wiz.state.out_path == str(out)
+
+ p3.btn_export.click()
+ assert wait_until(lambda: not win._job_running("align_export"),
+ timeout_ms=60000), "export finished"
+ assert wiz.state.exported_path == str(out), p3.lbl_status.text()
+ assert p3.isComplete(), "Finish available once the file exists"
+ assert out.exists()
+
+ written = SrasFile(str(out))
+ ctx.written = written
+ assert written.version == 6, "export is a v6 file"
+ assert written.n_angles == ctx.s.n_angles
+ assert all(written.image_shape(a) == ctx.crop[2:]
+ for a in range(written.n_angles)), \
+ "every angle shares the cropped grid"
+ assert not out.with_name(out.name + ".part").exists(), \
+ "no staging file left behind"
+
+
+def test_wizard_finish_applies_and_persists(ctx):
+ """Finish makes the session match the file: Aligned View shows the exported
+ extent, and the sidecar records it for the input scan."""
+ win, wiz = ctx.win, ctx.wiz
+ wiz.accept()
+ pump(250)
+ assert win._align_wizard is None, "wizard reference released"
+ assert win._wizard_act.isEnabled(), "wizard action available again"
+ assert win._alignment_result is not None
+ assert win._alignment_result.canvas_shape == ctx.crop[2:], \
+ "the *cropped* result is what the view now uses"
+ assert win.chk_aligned_view.isEnabled() and win.chk_aligned_view.isChecked()
+ pump(200)
+ assert win.image_canvas._img_shape == ctx.crop[2:], \
+ f"canvas shows the cropped extent: {win.image_canvas._img_shape}"
+
+ sidecar = compute.sidecar_path(ctx.s.path)
+ assert sidecar.exists(), "sidecar written for the input scan"
+ ctx.sidecar = sidecar
+ ctx.sidecar_raw = raw = json.loads(sidecar.read_text())
+ assert raw.get("schema_version") == compute._SIDECAR_SCHEMA_VERSION
+ assert all(raw["per_angle"][str(a)]["rotation_deg"]
+ == win._alignment_result.per_angle[a].rotation_deg
+ for a in range(ctx.s.n_angles)), \
+ "sidecar round-trips the applied rotations"
+
+ win.chk_aligned_view.setChecked(False)
+ pump(200)
+ assert win.image_canvas._img_shape == ctx.s.image_shape(0), \
+ f"unchecking returns to the raw per-angle grid: {win.image_canvas._img_shape}"
+
+
+def test_stale_schema_sidecar_ignored(ctx):
+ """An old-schema sidecar (pre-pivot/sign fix) is treated as absent."""
+ s, raw, sidecar = ctx.s, ctx.sidecar_raw, ctx.sidecar
+ stale = dict(raw)
+ stale["schema_version"] = compute._SIDECAR_SCHEMA_VERSION - 1
+ sidecar.write_text(json.dumps(stale))
+ assert compute.load_manual_alignment(s) is None, \
+ "a sidecar with an old schema_version is not loaded"
+ sidecar.write_text(json.dumps(raw)) # restore for the rest of the sequence
+
+
+def test_sidecar_restored_on_reload(ctx):
+ win, active = ctx.win, ctx.active
+ saved = json.loads(ctx.sidecar.read_text())["per_angle"][str(active)]
+ old_sras_id = id(win._sras)
+ win._load_file(str(ctx.path)) # reload the same file fresh
+ assert wait_until(
+ lambda: win._sras is not None and id(win._sras) != old_sras_id), \
+ "file reloaded"
+ ctx.s = win._sras
+ assert win._align_wizard is None, "no wizard left open across a reload"
+ assert win._alignment_result is not None, \
+ "reload restores the saved alignment automatically"
+ assert abs(win._alignment_result.per_angle[active].rotation_deg
+ - saved["rotation_deg"]) < 1e-9, \
+ "restored rotation matches what was saved"
+ assert win.chk_aligned_view.isChecked(), \
+ "Aligned View auto-checked after restoring a saved alignment"
+
+
+def test_wizard_closes_with_a_reload(ctx):
+ """An open wizard belongs to the file it was opened on."""
+ win = ctx.win
+ win._on_alignment_wizard()
+ assert win._align_wizard is not None
+ assert wait_until(
+ lambda: win._align_wizard.page(win._align_wizard.PAGE_CORRELATE).isComplete())
+ win._load_file(str(ctx.path))
+ assert wait_until(lambda: not win._job_running("load"))
+ pump(200)
+ assert win._align_wizard is None, "wizard force-closed by a reload"
+ ctx.s = win._sras
+
+
+# ---------------------------------------------------------------------------
+# Export Fused ROI
+# ---------------------------------------------------------------------------
+#
+# Two independent ways angles can end up sharing one (x, y) grid to fuse
+# onto: a live alignment result (case a, exercised on ctx.win -- the reload
+# above restored one from the sidecar), or a file that is itself a previous
+# Alignment Wizard export, whose angles already share a grid on disk with no
+# alignment result needed at all (case b, exercised on a second window
+# opened on ctx.written from test_wizard_export).
+
+def test_fused_export_gating_case_a(ctx):
+ """A live alignment result bridges the raw scan's per-angle grids --
+ angles_share_raw_grid() alone would be False here."""
+ win, s = ctx.win, ctx.s
+ assert win._alignment_result is not None, "alignment restored from sidecar"
+ assert not s.angles_share_raw_grid(), \
+ "sanity check: the raw (un-aligned) scan must not already share a grid"
+
+ x, y = win._aligned_canvas_axes()
+ roi = RoiQuad.from_bbox(float(x[1]), float(y[1]), float(x[-2]), float(y[-2]))
+ win.image_canvas.set_roi(roi)
+ pump(120)
+ assert win._fused_grid_ready()
+ assert win.btn_export_fused_roi.isEnabled()
+
+
+def test_write_fused_roi_csv_case_a_content(ctx):
+ win, s = ctx.win, ctx.s
+ roi = win.image_canvas.get_roi()
+ assert roi is not None, "ROI drawn by test_fused_export_gating_case_a"
+
+ csv_path = ctx.tmpdir / "fused_roi_case_a.csv"
+ win._write_fused_roi_csv(roi, CH4_IDX, [0, 1], str(csv_path))
+ assert csv_path.exists()
+
+ lines = csv_path.read_text().splitlines()
+ body = [l for l in lines if not l.startswith("#")]
+ header, *data_lines = body
+ assert header == (
+ f"x_mm,y_mm,v_{s.angles_deg[0]:.4g}deg,v_{s.angles_deg[1]:.4g}deg")
+
+ x, y = win._fused_export_axes()
+ mask = roi.mask_for_grid(x, y)
+ assert len(data_lines) == int(mask.sum())
+
+ data = np.array([[float(v) for v in line.split(",")] for line in data_lines])
+ expect0 = win._fused_value_image(0, CH4_IDX)[mask]
+ expect1 = win._fused_value_image(1, CH4_IDX)[mask]
+ assert np.allclose(data[:, 2], expect0, rtol=1e-5, atol=1e-4)
+ assert np.allclose(data[:, 3], expect1, rtol=1e-5, atol=1e-4)
+
+
+def test_fused_roi_dialog_availability_live_updates(ctx):
+ """Switching the value-type radio re-evaluates every angle checkbox,
+ disabling/auto-unchecking whichever ones are no longer available --
+ independent of what actually backs availability_fn, so a synthetic
+ stand-in keeps this a fast, deterministic test of the dialog itself."""
+ win, s = ctx.win, ctx.s
+ angles = [(a, float(s.angles_deg[a])) for a in range(s.n_angles)]
+ only_angle0_has_ch1 = lambda a, c: (a == 0) if c == CH1_IDX else True
+
+ dlg = FusedRoiExportDialog(
+ win, angles=angles, availability_fn=only_angle0_has_ch1,
+ default_ch_idx=CH1_IDX, out_dir=str(s.path.parent), stem=s.path.stem,
+ grid_note="test")
+ try:
+ assert dlg._angle_checks[0].isEnabled()
+ assert all(not dlg._angle_checks[a].isEnabled()
+ for a in range(1, s.n_angles))
+
+ dlg._angle_checks[0].setChecked(True)
+ dlg._val_buttons[CH4_IDX].click()
+ assert all(dlg._angle_checks[a].isEnabled() for a in range(s.n_angles)), \
+ "CH4 is available for every angle"
+ assert dlg._angle_checks[0].isChecked(), \
+ "stays checked -- still available under CH4"
+
+ if s.n_angles > 1:
+ dlg._angle_checks[1].setChecked(True)
+ dlg._val_buttons[CH1_IDX].click()
+ assert dlg._angle_checks[0].isChecked()
+ if s.n_angles > 1:
+ assert not dlg._angle_checks[1].isEnabled()
+ assert not dlg._angle_checks[1].isChecked(), \
+ "auto-unchecked: angle 1 has no data under CH1"
+ finally:
+ dlg.close()
+
+
+def test_fused_roi_dialog_select_all_none(ctx):
+ win, s = ctx.win, ctx.s
+ angles = [(a, float(s.angles_deg[a])) for a in range(s.n_angles)]
+ dlg = FusedRoiExportDialog(
+ win, angles=angles, availability_fn=lambda a, c: c == CH4_IDX,
+ default_ch_idx=CH4_IDX, out_dir=str(s.path.parent), stem=s.path.stem,
+ grid_note="test")
+ try:
+ assert not dlg._btn_export.isEnabled(), "nothing checked yet"
+ dlg._on_select_all_available()
+ assert all(cb.isChecked() for cb in dlg._angle_checks.values())
+ assert dlg._btn_export.isEnabled()
+ dlg._on_select_none()
+ assert not any(cb.isChecked() for cb in dlg._angle_checks.values())
+ assert not dlg._btn_export.isEnabled()
+ finally:
+ dlg.close()
+
+
+def test_on_export_fused_roi_csv_end_to_end(ctx):
+ win, s = ctx.win, ctx.s
+ roi = win.image_canvas.get_roi()
+ assert roi is not None, "ROI from the earlier fused-export tests is still set"
+
+ csv_path = ctx.tmpdir / "fused_roi_e2e.csv"
+ with patch("sras_viewer.main_window.FusedRoiExportDialog") as MockDlg:
+ inst = MockDlg.return_value
+ inst.exec.return_value = QDialog.DialogCode.Accepted
+ inst.get_ch_idx.return_value = CH4_IDX
+ inst.get_selected_angles.return_value = [0, 1]
+ inst.get_output_path.return_value = str(csv_path)
+ win.btn_export_fused_roi.click()
+
+ assert csv_path.exists()
+ kwargs = MockDlg.call_args.kwargs
+ assert kwargs["default_ch_idx"] == win.combo_channel.currentIndex()
+ assert kwargs["out_dir"] == str(s.path.parent)
+ assert kwargs["stem"] == s.path.stem
+
+
+def test_fused_export_no_alignment_shared_grid_path(ctx):
+ """ctx.written (from test_wizard_export) is itself a previous Alignment
+ Wizard export: opened fresh with no sidecar for its own path, so no
+ alignment result is ever restored -- but its angles already share one
+ grid on disk, so the export must work through the no-resample path."""
+ win2 = SrasViewerWindow()
+ try:
+ win2._load_file(str(ctx.written.path))
+ assert wait_until(lambda: win2._sras is not None)
+ s2 = win2._sras
+ assert win2._alignment_result is None, \
+ "no sidecar exists for this path -- nothing auto-restored"
+ assert s2.angles_share_raw_grid(), \
+ "a wizard export already shares one grid across angles"
+ assert wait_until(lambda: all((a, CH4_IDX) in win2._dc_cache
+ for a in range(s2.n_angles))), \
+ "DC precomputed for every angle"
+
+ x, y = s2.x_axis_mm(0), s2.y_positions_mm(0)
+ roi = RoiQuad.from_bbox(float(x[1]), float(y[1]), float(x[-2]), float(y[-2]))
+ win2.image_canvas.set_roi(roi)
+ pump(120)
+ assert win2._fused_grid_ready()
+ assert win2.btn_export_fused_roi.isEnabled()
+
+ aligned_cache_before = len(win2._aligned_cache)
+ angle_idxs = list(range(min(2, s2.n_angles)))
+ csv_path = ctx.tmpdir / "fused_roi_case_b.csv"
+ win2._write_fused_roi_csv(roi, CH4_IDX, angle_idxs, str(csv_path))
+ assert csv_path.exists()
+ assert len(win2._aligned_cache) == aligned_cache_before, \
+ "no-resample path must never touch apply_alignment"
+
+ header = next(l for l in csv_path.read_text().splitlines()
+ if not l.startswith("#"))
+ expected_header = "x_mm,y_mm," + ",".join(
+ f"v_{s2.angles_deg[a]:.4g}deg" for a in angle_idxs)
+ assert header == expected_header
+
+ # Break the shared grid and confirm gating flips off.
+ mutate_idx = 1 if s2.n_angles > 1 else 0
+ s2.x_start_mm[mutate_idx] += 1.0
+ assert not s2.angles_share_raw_grid()
+ win2._update_fused_export_enabled()
+ assert not win2._fused_grid_ready()
+ assert not win2.btn_export_fused_roi.isEnabled()
+ assert "Alignment Wizard" in win2.btn_export_fused_roi.toolTip()
+ finally:
+ win2.close()
+ pump(200)
+
+
+def test_pixel_inspector(ctx):
+ win = ctx.win
+ win.chk_aligned_view.setChecked(False)
+ pump(100)
+ win._on_pixel_clicked(0, 0)
+ pump(150)
+ assert win.lbl_wave_hint.isHidden(), "waveform hint hidden after a click"
+ win.combo_channel.setCurrentIndex(CH1_IDX)
+ wait_until(lambda: not win._job_running("compute"))
+ win._on_pixel_clicked(1, 1)
+ pump(150)
+ assert len(win.wave_canvas.ax_wave.lines) > 0, \
+ f"RF waveform panel rendered ({len(win.wave_canvas.ax_wave.lines)} lines)"
+
+
+def test_shutdown(ctx):
+ win = ctx.win
+ win.close()
+ pump(400)
+ assert len(win._jobs) == 0, f"all background jobs released: {list(win._jobs)}"
+
+
+def test_no_status_bar_errors(ctx):
+ unexpected = [e for e in ctx.errors if e]
+ assert not unexpected, f"status-bar errors seen: {unexpected}"
diff --git a/tests/test_row_average.py b/tests/test_row_average.py
new file mode 100644
index 0000000..448cee4
--- /dev/null
+++ b/tests/test_row_average.py
@@ -0,0 +1,278 @@
+"""Row-averaged FFT: same-row, distance-weighted CH1 waveform smoothing.
+
+Covers the properties the design depends on: the kernel is symmetric and
+n=0 is a true no-op; the masked/renormalized convolution matches an
+independent brute-force reference and gives masked neighbors exactly zero
+weight regardless of their content; background subtraction after averaging
+is algebraically identical to subtracting before; chunking/worker count
+never changes the result; and a pixel that's itself masked is never
+"rescued" by averaging.
+"""
+
+import numpy as np
+import pytest
+
+import sras_compute as compute
+from sras_compute import compute_rf_image, dc_image_mv
+from sras_format import CH4_IDX, SrasFile
+import tools.make_test_sras as gen
+
+
+def _reference_row_average(masked_waves: np.ndarray, valid: np.ndarray,
+ weights: np.ndarray) -> np.ndarray:
+ """Independent, unvectorized reference for _row_average_waveforms: for
+ each row position, sum weighted valid neighbors within the kernel's
+ radius and normalize by the actual included weight sum. Same
+ definition, computed by brute-force nested loops instead of
+ correlate1d, so it can't share a bug with the implementation."""
+ n_frames, spf = masked_waves.shape
+ n = len(weights) // 2
+ out = np.zeros_like(masked_waves)
+ for i in range(n_frames):
+ num = np.zeros(spf, dtype=np.float64)
+ den = 0.0
+ for d in range(-n, n + 1):
+ j = i + d
+ if 0 <= j < n_frames and valid[j]:
+ w = float(weights[d + n])
+ num += w * masked_waves[j].astype(np.float64)
+ den += w
+ out[i] = num / den if den > 0 else 0.0
+ return out
+
+
+# ---------------------------------------------------------------------------
+# _row_average_weights
+# ---------------------------------------------------------------------------
+
+def test_row_average_weights_shape_and_symmetry():
+ w0 = compute._row_average_weights(0)
+ assert w0.shape == (1,) and w0[0] == 1.0
+
+ for n in (1, 2, 5):
+ w = compute._row_average_weights(n)
+ assert w.shape == (2 * n + 1,)
+ assert w[n] == pytest.approx(1.0), "center tap is the peak weight"
+ assert np.allclose(w, w[::-1]), "symmetric about the center"
+ half = w[n:]
+ assert np.all(np.diff(half) < 0), "strictly decreasing away from center"
+
+
+# ---------------------------------------------------------------------------
+# _row_average_waveforms
+# ---------------------------------------------------------------------------
+
+def test_row_average_matches_hand_rolled_reference():
+ rng = np.random.default_rng(0)
+ n_frames, spf = 15, 6
+ raw = rng.integers(-50, 51, size=(n_frames, spf)).astype(np.float32)
+ valid = np.ones(n_frames, dtype=bool)
+ valid[[2, 3, 9]] = False # a run of two invalid, plus a lone invalid
+ masked = raw.copy()
+ masked[~valid] = 0.0
+
+ weights = compute._row_average_weights(3)
+ got = compute._row_average_waveforms(masked, valid, weights)
+ ref = _reference_row_average(masked, valid, weights)
+
+ assert np.allclose(got[valid], ref[valid], atol=1e-4)
+
+
+def test_row_average_edge_of_row():
+ """A window wider than the row itself must still renormalize correctly
+ at both ends -- mode='constant', cval=0.0 zero-pads both the numerator
+ and denominator, so this is not a special case, but it's the one most
+ likely to break if that padding were ever mismatched between the two."""
+ n_frames, spf = 6, 3
+ raw = np.arange(n_frames * spf, dtype=np.float32).reshape(n_frames, spf)
+ valid = np.ones(n_frames, dtype=bool)
+ weights = compute._row_average_weights(4) # window (9 taps) > n_frames (6)
+
+ got = compute._row_average_waveforms(raw, valid, weights)
+ ref = _reference_row_average(raw, valid, weights)
+ assert np.allclose(got, ref, atol=1e-4)
+
+
+def test_row_average_excludes_masked_neighbor_from_normalization():
+ """A masked neighbor must contribute zero *weight* to the normalization,
+ not participate as a legitimate zero-valued sample at full weight --
+ the two give different answers, and only the former is correct. (Note:
+ masked_waves must already be 0 at invalid positions per
+ _row_average_waveforms's contract -- that's what read_row's zero-filled
+ scratch buffer guarantees in production -- so the only way to vary
+ "what a masked position looks like" while respecting that contract is
+ whether its weight is excluded from the denominator at all.)"""
+ n_frames, spf = 9, 4
+ weights = compute._row_average_weights(2)
+
+ # A: position 4 is masked -- excluded from the weight sum entirely.
+ valid_a = np.ones(n_frames, dtype=bool)
+ valid_a[4] = False
+ masked_a = np.zeros((n_frames, spf), dtype=np.float32)
+ masked_a[valid_a] = 1.0
+ got_a = compute._row_average_waveforms(masked_a, valid_a, weights)
+
+ # B: position 4 is valid but genuinely zero-valued -- included in the
+ # weight sum, diluting neighbors' averages.
+ valid_b = np.ones(n_frames, dtype=bool)
+ masked_b = np.ones((n_frames, spf), dtype=np.float32)
+ masked_b[4] = 0.0
+ got_b = compute._row_average_waveforms(masked_b, valid_b, weights)
+
+ # Every position whose window reaches index 4 must average *higher* in
+ # A (excluded from the denominator) than in B (included as a real zero).
+ affected = [2, 3, 5, 6]
+ assert np.all(got_a[affected] > got_b[affected]), \
+ "masking must exclude a neighbor from normalization, not just zero its value"
+ # Positions outside the window (radius 2) are unaffected either way.
+ assert np.allclose(got_a[[0, 1, 7, 8]], got_b[[0, 1, 7, 8]])
+
+
+def test_background_subtracted_once_equals_subtract_then_average():
+ """Algebraic identity the implementation relies on: subtracting a fixed
+ background from the already-averaged waveform equals subtracting it
+ from every valid neighbor first, because the denominator is always the
+ *actual* included weight sum (never a fixed total)."""
+ rng = np.random.default_rng(1)
+ n_frames, spf = 11, 8
+ raw = rng.integers(-40, 41, size=(n_frames, spf)).astype(np.float32)
+ valid = np.ones(n_frames, dtype=bool)
+ valid[[1, 7]] = False
+ masked = raw.copy()
+ masked[~valid] = 0.0
+ background = rng.integers(-5, 6, size=spf).astype(np.float32)
+ weights = compute._row_average_weights(3)
+
+ # Order A (what the code does): average first, subtract background once.
+ order_a = compute._row_average_waveforms(masked, valid, weights) - background
+
+ # Order B: subtract background from every valid neighbor first (restoring
+ # the "0 at invalid positions" contract afterward), then average.
+ bg_subbed = masked - background
+ bg_subbed[~valid] = 0.0
+ order_b = compute._row_average_waveforms(bg_subbed, valid, weights)
+
+ assert np.allclose(order_a[valid], order_b[valid], atol=1e-3)
+
+
+# ---------------------------------------------------------------------------
+# compute_rf_image(row_avg_n=...) integration
+# ---------------------------------------------------------------------------
+
+def test_row_average_zero_is_identity(tmp_path):
+ """row_avg_n=0 must take the exact same code path as before this
+ feature existed (row_avg_weights stays None), not a single-tap kernel
+ that merely computes to the same answer."""
+ path = tmp_path / "zero.sras"
+ gen.write(path, n_angles=1, seed=10, samples_per_frame=64)
+ sras = SrasFile(str(path))
+ plain = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True)
+ explicit_zero = compute_rf_image(sras, 0, dc_threshold_mv=None,
+ apply_bg_sub=True, row_avg_n=0)
+ assert np.array_equal(plain, explicit_zero)
+
+
+def test_row_average_respects_own_center_mask(tmp_path):
+ """A pixel that's itself below threshold stays masked (0) after row
+ averaging -- averaging never rescues a masked pixel, matching the
+ 'valid neighbors only' design (masked pixels are excluded from other
+ pixels' averages, and are never themselves smoothed)."""
+ path = tmp_path / "center_mask.sras"
+ gen.write(path, n_angles=1, seed=13, samples_per_frame=64)
+ sras = SrasFile(str(path))
+ dc4 = dc_image_mv(sras, 0, CH4_IDX)
+ thr = float(np.percentile(dc4, 50))
+ mask = dc4 >= thr
+ assert mask.any() and not mask.all(), "threshold actually splits the image"
+
+ img = compute_rf_image(sras, 0, dc_threshold_mv=thr, apply_bg_sub=True,
+ row_avg_n=4)
+ assert np.array_equal(img == 0, ~mask), \
+ "masked pixels stay exactly 0 after row averaging; valid ones don't"
+
+
+def test_row_average_composes_with_padding(tmp_path):
+ """row_avg_n and n_fft (zero-padding) are independent knobs: using them
+ together must not raise, and must still agree with the exact (non-zoom)
+ reference path at that pad factor -- i.e. row-averaging composes with
+ the zoom peak search correctly, not just with the direct one."""
+ path = tmp_path / "padded_rowavg.sras"
+ gen.write(path, n_angles=1, seed=12, samples_per_frame=64)
+ sras = SrasFile(str(path))
+ spf = sras.samples_per_frame
+
+ raw_natural = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True)
+ avg_natural = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True,
+ row_avg_n=3)
+ avg_padded_zoom = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True,
+ row_avg_n=3, n_fft=spf * 40)
+ avg_padded_exact = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True,
+ row_avg_n=3, n_fft=spf * 40, exact=True)
+
+ assert avg_natural.shape == raw_natural.shape == avg_padded_zoom.shape
+ assert np.all(np.isfinite(avg_padded_zoom))
+ assert np.array_equal(avg_padded_zoom, avg_padded_exact), \
+ "row-averaged waveforms feed the zoom and exact FFT paths identically"
+
+
+def test_row_average_improves_snr_recovery():
+ """The actual point of the feature: averaging same-row waveforms that
+ share a true underlying tone but carry independent noise recovers that
+ tone far more reliably than any single raw (unaveraged) waveform does."""
+ rng = np.random.default_rng(42)
+ n_frames, spf = 21, 128
+ true_bin = 9
+ t = np.arange(spf)
+ tone = 15.0 * np.sin(2 * np.pi * true_bin * t / spf) # same true signal
+ # at every position
+ noise_sigma = 40.0 # much larger than the tone -- deliberately poor SNR
+ raw = (tone[None, :] + rng.normal(scale=noise_sigma, size=(n_frames, spf))
+ ).astype(np.float32)
+ valid = np.ones(n_frames, dtype=bool)
+
+ weights = compute._row_average_weights(8) # wide window: lots of averaging
+ averaged = compute._row_average_waveforms(raw, valid, weights)
+
+ raw_bins = compute._peak_bins_direct(raw, spf)
+ avg_bins = compute._peak_bins_direct(averaged, spf)
+
+ raw_hits = int(np.sum(raw_bins == true_bin))
+ avg_hits = int(np.sum(avg_bins == true_bin))
+ assert avg_hits > raw_hits, (
+ f"row averaging should recover the true bin ({true_bin}) more often "
+ f"than raw per-pixel estimates: raw {raw_hits}/{n_frames}, "
+ f"averaged {avg_hits}/{n_frames}")
+ assert avg_hits >= n_frames * 0.7, \
+ f"averaged recovery should be reliable, not just barely better: {avg_hits}/{n_frames}"
+
+
+def test_row_average_parallel_identity(tmp_path, monkeypatch):
+ """Forcing 1 worker vs many must give an identical row-averaged image --
+ catches chunk-boundary bugs (there should be none, since averaging never
+ crosses rows, but this is the empirical proof, not just inspection)."""
+ path = tmp_path / "parallel_rowavg.sras"
+ n_rows, n_frames, spf = 40, 13, 128
+ gen.write(path, n_angles=1, seed=11, samples_per_frame=spf,
+ geometry=[(n_rows, n_frames)])
+ sras = SrasFile(str(path))
+
+ monkeypatch.setattr(compute, "_TOTAL_BYTES_BUDGET", 8 * n_frames * spf * 4)
+ monkeypatch.setattr(compute, "_FFT_BLOCK", 4)
+ # row_avg_n > 0 halves the effective budget before chunk planning.
+ fft_rows = compute._plan_fft_rows(n_frames, spf, compute._TOTAL_BYTES_BUDGET // 2)
+ assert fft_rows < n_rows, \
+ f"row-averaged FFT work actually splits into multiple chunks ({fft_rows} of {n_rows})"
+
+ dc4 = dc_image_mv(sras, 0, CH4_IDX)
+ thr = float(np.median(dc4))
+
+ monkeypatch.setattr(compute, "_MAX_WORKERS", 1)
+ serial = compute_rf_image(sras, 0, dc_threshold_mv=thr, apply_bg_sub=True,
+ row_avg_n=4)
+
+ monkeypatch.setattr(compute, "_MAX_WORKERS", 8)
+ parallel = compute_rf_image(sras, 0, dc_threshold_mv=thr, apply_bg_sub=True,
+ row_avg_n=4)
+
+ assert np.array_equal(serial, parallel), \
+ "row-averaged rf image identical regardless of chunking/worker count"
diff --git a/tests/test_stored_cache.py b/tests/test_stored_cache.py
new file mode 100644
index 0000000..f802e53
--- /dev/null
+++ b/tests/test_stored_cache.py
@@ -0,0 +1,595 @@
+"""Does a batch-computed FFT cache actually spare the viewer the FFT?
+
+Storing a peak-frequency image per angle in the file is only worth doing if
+displaying it is then free. The regression this module pins down is the
+viewer's *dispatch* decision: it used to find the stored image only inside
+ComputeWorker, so after a batch every angle change still queued a background
+job behind a "Computing FFT…" popup for an image already on disk.
+
+Both layers are covered — cached_rf_image's accept/reject rules, and the
+window never reaching _start_compute for a batch-cached angle.
+"""
+
+import struct
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import numpy as np
+import pytest
+from PyQt6.QtCore import QEventLoop, QTimer
+from PyQt6.QtWidgets import QApplication, QDialog
+
+import sras_compute as compute
+import sras_format as fmt
+from sras_compute import cache_file, cached_rf_image, compute_rf_image, dc_image_mv
+from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, SrasFile
+from sras_viewer import SrasViewerWindow, VELOCITY_MODE_IDX
+import tools.make_test_sras as gen
+
+_THRESHOLD_MV = 50.0 # the viewer's own default
+
+
+def pump(ms: int = 200):
+ 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()
+
+
+@pytest.fixture(scope="module")
+def rig(tmp_path_factory):
+ """A v6 file, the FFT images a from-scratch compute gives for it, and the
+ same file after Batch Compute FFT has written them into its v7 cache."""
+ path = tmp_path_factory.mktemp("stored_cache") / "cached.sras"
+ gen.write(path, n_angles=4, seed=7, samples_per_frame=256)
+
+ src = SrasFile(str(path))
+ fresh = {a: compute_rf_image(src, a, dc_threshold_mv=_THRESHOLD_MV,
+ apply_bg_sub=True)
+ for a in range(src.n_angles)}
+ assert src.background is not None, "the fixture file must have a background"
+
+ err = cache_file(str(path), "fft", True)
+ assert err == "", err
+ cached = SrasFile(str(path))
+ assert all(x is not None for x in cached.precomputed_freq_mhz)
+ assert all(x is None for x in cached.precomputed_dc4_mv), \
+ "FFT-only batch: the mask has to come from the viewer, not the file"
+
+ return SimpleNamespace(path=path, fresh=fresh, sras=cached,
+ n_angles=cached.n_angles)
+
+
+@pytest.fixture(scope="module")
+def dc_rig(tmp_path_factory):
+ """A file that has been through Batch Compute DC and Store."""
+ path = tmp_path_factory.mktemp("stored_dc") / "dc_cached.sras"
+ gen.write(path, n_angles=4, seed=9, samples_per_frame=128)
+ err = cache_file(str(path), "dc", True)
+ assert err == "", err
+
+ sras = SrasFile(str(path))
+ assert all(x is not None for x in sras.precomputed_dc3_mv)
+ assert all(x is not None for x in sras.precomputed_dc4_mv)
+ assert len(np.unique(sras.precomputed_dc4_mv[0])) > 1, \
+ "a degenerate DC image would make the comparisons below vacuous"
+ return SimpleNamespace(path=path, sras=sras, n_angles=sras.n_angles)
+
+
+@pytest.fixture
+def no_fft(monkeypatch):
+ """Make any real FFT work loud: returns a list that stays empty unless a
+ peak search actually runs."""
+ calls = []
+ for name in ("_peak_bins_direct", "_peak_bins_zoom"):
+ original = getattr(compute, name)
+
+ def spy(*args, _f=original, **kwargs):
+ calls.append(_f.__name__)
+ return _f(*args, **kwargs)
+
+ monkeypatch.setattr(compute, name, spy)
+ return calls
+
+
+# ---------------------------------------------------------------------------
+# cached_rf_image: when may the stored image stand in for a compute?
+# ---------------------------------------------------------------------------
+
+def test_stored_image_matches_a_fresh_compute(rig, no_fft):
+ for a in range(rig.n_angles):
+ dc4 = dc_image_mv(SrasFile(str(rig.path)), a, CH4_IDX)
+ img = cached_rf_image(rig.sras, a, dc_threshold_mv=_THRESHOLD_MV,
+ apply_bg_sub=True, dc4_mv=dc4)
+ assert img is not None, f"angle {a} is cached in the file"
+ assert np.allclose(img, rig.fresh[a], atol=1e-3), \
+ f"angle {a} differs from a from-scratch compute"
+ assert not no_fft, f"the stored image was used, no FFT ran: {no_fft}"
+
+
+def test_unmasked_when_no_threshold(rig):
+ img = cached_rf_image(rig.sras, 0, dc_threshold_mv=None, apply_bg_sub=True)
+ assert img is not None and np.array_equal(img, rig.sras.precomputed_freq_mhz[0])
+ img[:] = -1.0
+ assert not np.any(rig.sras.precomputed_freq_mhz[0] == -1.0), \
+ "callers get a copy, never the file's own array"
+
+
+def test_settings_the_stored_image_cannot_serve(rig):
+ """A stored image carries one bg-sub state and one padding, so anything
+ else must fall through to a real compute rather than lie."""
+ spf = rig.sras.samples_per_frame
+ assert cached_rf_image(rig.sras, 0, _THRESHOLD_MV, apply_bg_sub=True,
+ n_fft=spf * 4) is None, \
+ "cache was written at pad 1, a pad-4 view resolves different peaks"
+ assert cached_rf_image(rig.sras, 0, _THRESHOLD_MV, apply_bg_sub=False) is None, \
+ "cache was written with bg-sub on"
+
+ uncached = SrasFile(str(rig.path))
+ uncached.precomputed_freq_mhz[1] = None
+ assert cached_rf_image(uncached, 1, _THRESHOLD_MV, apply_bg_sub=True) is None
+
+
+@pytest.mark.parametrize("pad", [2, 10])
+def test_cache_is_stored_at_the_configured_pad(tmp_path, pad):
+ """Batching at a pad factor must produce a cache that view can read back —
+ a pad-1-only cache is one the padded viewer can never use."""
+ path = tmp_path / f"pad{pad}.sras"
+ gen.write(path, n_angles=2, seed=13, samples_per_frame=256)
+ spf = SrasFile(str(path)).samples_per_frame
+
+ assert cache_file(str(path), "fft", True, "scipy", 0, pad) == ""
+ sras = SrasFile(str(path))
+ assert sras.precomputed_pad_factor == pad, "pad factor survives the round trip"
+
+ assert cached_rf_image(sras, 0, None, apply_bg_sub=True,
+ n_fft=spf * pad) is not None, f"usable at pad {pad}"
+ assert cached_rf_image(sras, 0, None, apply_bg_sub=True) is None, \
+ "not usable unpadded"
+ assert cached_rf_image(sras, 0, None, apply_bg_sub=True,
+ n_fft=spf * (pad + 1)) is None, "not usable at another pad"
+
+ # The stored numbers must be the padded ones, not pad-1 relabelled.
+ fresh = SrasFile(str(path))
+ fresh.precomputed_freq_mhz = [None] * fresh.n_angles
+ for a in range(sras.n_angles):
+ assert np.allclose(
+ compute_rf_image(sras, a, dc_threshold_mv=None, apply_bg_sub=True,
+ n_fft=spf * pad),
+ compute_rf_image(fresh, a, dc_threshold_mv=None, apply_bg_sub=True,
+ n_fft=spf * pad), atol=1e-3), \
+ f"angle {a}: stored image is the pad-{pad} answer"
+
+
+def test_cach_v1_reads_as_natural_resolution(tmp_path):
+ """Files cached before the pad factor existed must keep working: a v1 tail
+ has no pad field and is pad 1 by construction."""
+ path = tmp_path / "v1.sras"
+ gen.write(path, n_angles=2, seed=14, samples_per_frame=256)
+ assert cache_file(str(path), "fft", True) == ""
+
+ # Rewrite the tail as a genuine CACH v1 block (old header, no pad field).
+ v2 = SrasFile(str(path))
+ freq, entries = v2.precomputed_freq_mhz, list(range(v2.n_angles))
+ payload = struct.pack(fmt.CACH_HDR_FMT, fmt.CACH_MAGIC, 1, fmt.CACH_FLAG_FFT)
+ payload += struct.pack(fmt.SFFT_HDR_FMT_V1, fmt.SFFT_MAGIC,
+ fmt.SFFT_FLAG_BG_SUB, len(entries))
+ for a in entries:
+ payload += struct.pack(">H", a) + freq[a].astype(">f4").tobytes()
+ head = path.read_bytes()[:v2._cache_tail_offset()]
+ path.write_bytes(head + payload)
+
+ v1 = SrasFile(str(path))
+ assert v1.precomputed_pad_factor == 1
+ assert v1.precomputed_bg_sub is True
+ assert all(np.array_equal(v1.precomputed_freq_mhz[a], freq[a]) for a in entries), \
+ "v1 images read back unchanged"
+ assert cached_rf_image(v1, 0, None, apply_bg_sub=True) is not None
+ assert cached_rf_image(v1, 0, None, apply_bg_sub=True,
+ n_fft=v1.samples_per_frame * 10) is None
+
+
+def test_mask_read_can_be_refused(rig, no_fft):
+ """With no DC4 in hand, building the mask means reading a whole channel —
+ the GUI thread asks for None instead."""
+ assert cached_rf_image(rig.sras, 0, _THRESHOLD_MV, apply_bg_sub=True,
+ allow_dc_recompute=False) is None
+ dc4 = dc_image_mv(SrasFile(str(rig.path)), 0, CH4_IDX)
+ img = cached_rf_image(rig.sras, 0, _THRESHOLD_MV, apply_bg_sub=True,
+ dc4_mv=dc4, allow_dc_recompute=False)
+ assert img is not None and np.allclose(img, rig.fresh[0], atol=1e-3)
+ assert not no_fft, f"no FFT on either branch: {no_fft}"
+
+
+# ---------------------------------------------------------------------------
+# The viewer: no compute job at all for a batch-cached angle
+# ---------------------------------------------------------------------------
+
+def test_viewer_shows_stored_angles_without_computing(rig, no_fft, monkeypatch):
+ app = QApplication.instance() or QApplication([]) # noqa: F841
+ win = SrasViewerWindow()
+ win.show()
+
+ dispatched = []
+ original_start = type(win)._start_compute
+ monkeypatch.setattr(type(win), "_start_compute",
+ lambda self: (dispatched.append(self.spin_angle.value()),
+ original_start(self))[1])
+ try:
+ win._load_file(str(rig.path))
+ assert wait_until(lambda: win._sras is not None), "file loaded"
+ # The file carries no DC block, so the mask comes from the window's own
+ # background precompute — the state a user is in by the time they click.
+ assert wait_until(lambda: all((a, CH4_IDX) in win._dc_cache
+ for a in range(rig.n_angles))), \
+ "DC precompute finished"
+
+ dispatched.clear()
+ win.combo_channel.setCurrentIndex(CH1_IDX)
+ assert wait_until(lambda: win._current_ch == CH1_IDX), "CH1 displayed"
+
+ for a in list(range(rig.n_angles)) + [1, 0]:
+ win.spin_angle.setValue(a)
+ win._on_view_changed()
+ pump(60)
+ assert win._current_angle == a, f"angle {a} displayed"
+ assert np.allclose(win._current_image, rig.fresh[a], atol=1e-3), \
+ f"angle {a} shows the stored image"
+
+ assert dispatched == [], \
+ f"stored angles need no compute job, dispatched for {dispatched}"
+ assert not no_fft, f"no FFT ran for any stored angle: {no_fft}"
+
+ # Velocity is still a post-multiply of the same stored image.
+ win.combo_channel.setCurrentIndex(VELOCITY_MODE_IDX)
+ pump(120)
+ assert np.allclose(win._current_image,
+ rig.fresh[0] * win.spin_grating_um.value(), atol=1e-3)
+ assert dispatched == [] and not no_fft
+
+ # ...and a live control that no longer matches the stored image's own
+ # provenance must NOT force a recompute either — the stored image is
+ # shown as-is; only an explicit batch recompute changes what's shown.
+ win.combo_channel.setCurrentIndex(CH1_IDX)
+ pump(60)
+ win.chk_bg_sub.setChecked(False)
+ pump(200)
+ assert not win._job_running("compute") and not no_fft, \
+ "bg-sub off still shows the stored image, no real FFT"
+ assert np.allclose(win._current_image, rig.fresh[0], atol=1e-3)
+ finally:
+ win.close()
+ pump(300)
+
+
+def test_batch_caches_at_the_viewers_pad_factor(tmp_path, no_fft, monkeypatch):
+ """The bug a pad-10 user hits: Batch Compute FFT used to store pad-1
+ images regardless, so the padded view recomputed every angle forever.
+ """
+ path = tmp_path / "padded_gui.sras"
+ gen.write(path, n_angles=3, seed=15, samples_per_frame=256)
+
+ app = QApplication.instance() or QApplication([]) # noqa: F841
+ win = SrasViewerWindow()
+ win.show()
+
+ dispatched = []
+ original_start = type(win)._start_compute
+ monkeypatch.setattr(type(win), "_start_compute",
+ lambda self: (dispatched.append(self.spin_angle.value()),
+ original_start(self))[1])
+ try:
+ win._fft_pad_factor = 10
+ win._load_file(str(path))
+ assert wait_until(lambda: win._sras is not None), "file loaded"
+ assert wait_until(lambda: all((a, CH4_IDX) in win._dc_cache
+ for a in range(win._sras.n_angles))), \
+ "DC precompute finished"
+
+ # Convert -> Batch Compute FFT, on the open file, through the real slot.
+ with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames",
+ return_value=([str(path)], "")):
+ win._on_batch_compute("fft")
+ assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran"
+ assert wait_until(lambda: win._sras is not None
+ and win._sras.version == 7), "file reloaded as v7"
+ pump(200)
+ assert win._sras.precomputed_pad_factor == 10, \
+ f"cached at the viewer's pad, got {win._sras.precomputed_pad_factor}"
+
+ expected = {a: compute.cached_rf_image(
+ win._sras, a, dc_threshold_mv=win.spin_threshold_mv.value(),
+ apply_bg_sub=win.chk_bg_sub.isChecked(),
+ n_fft=win._current_n_fft(),
+ dc4_mv=win._dc_cache.get((a, CH4_IDX)))
+ for a in range(win._sras.n_angles)}
+ assert all(v is not None for v in expected.values()), "cache is readable at pad 10"
+
+ no_fft.clear()
+ dispatched.clear()
+ win.combo_channel.setCurrentIndex(CH1_IDX)
+ assert wait_until(lambda: win._current_ch == CH1_IDX), "CH1 displayed"
+ for a in range(win._sras.n_angles):
+ win.spin_angle.setValue(a)
+ pump(60)
+ assert np.allclose(win._current_image, expected[a], atol=1e-3), \
+ f"angle {a} served from the pad-10 cache"
+ assert dispatched == [] and not no_fft, \
+ f"no recompute at pad 10 (jobs={dispatched}, fft={no_fft})"
+ assert "unusable" not in win.lbl_frame_warn.text()
+
+ # Change the live pad control so it no longer matches the stored
+ # image's own provenance — the info panel has to say so rather than
+ # leave it a mystery, but the stored pad-10 image keeps displaying;
+ # only an explicit batch recompute would ever produce a pad-4 one.
+ win._fft_pad_factor = 4
+ win._update_scan_info_labels()
+ assert "differ from the stored cache" in win.lbl_frame_warn.text(), \
+ win.lbl_frame_warn.text()
+ assert "pad 10x" in win.lbl_frame_warn.text()
+ last_angle = win._current_angle
+ win._refresh_display()
+ assert wait_until(lambda: not win._job_running("compute")), "settled"
+ assert not no_fft, "no real FFT ran — the pad-10 cache still served the view"
+ assert np.allclose(win._current_image, expected[last_angle], atol=1e-3), \
+ "pad-10 cache still shown after the live pad control diverged"
+ finally:
+ win.close()
+ pump(300)
+
+
+def test_viewer_shows_stored_dc_without_computing(dc_rig, monkeypatch):
+ """Same for the DC half, with the background precompute silenced so the
+ file's stored block is the only thing that can be carrying the display."""
+ app = QApplication.instance() or QApplication([]) # noqa: F841
+ win = SrasViewerWindow()
+ win.show()
+
+ dispatched = []
+ original_start = type(win)._start_compute
+ monkeypatch.setattr(type(win), "_start_compute",
+ lambda self: (dispatched.append(self.spin_angle.value()),
+ original_start(self))[1])
+ monkeypatch.setattr(type(win), "_start_dc_precompute", lambda self: None)
+ try:
+ win._load_file(str(dc_rig.path))
+ assert wait_until(lambda: win._sras is not None), "file loaded"
+ pump(120)
+
+ for ch in (CH3_IDX, CH4_IDX):
+ win.combo_channel.setCurrentIndex(ch)
+ for a in range(dc_rig.n_angles):
+ win.spin_angle.setValue(a)
+ pump(60)
+ assert (win._current_angle, win._current_ch) == (a, ch), \
+ f"angle {a} on channel {ch} displayed"
+ assert np.array_equal(win._current_image,
+ dc_rig.sras.cached_dc_mv(a, ch)), \
+ f"angle {a} channel {ch} shows the file's stored DC image"
+
+ assert dispatched == [], \
+ f"stored DC angles need no compute job, dispatched for {dispatched}"
+ finally:
+ win.close()
+ pump(300)
+
+
+# ---------------------------------------------------------------------------
+# Row-averaged FFT: cache_file("fft_rowavg", ...) and its on-disk provenance
+# ---------------------------------------------------------------------------
+
+def test_row_average_flag_and_n_round_trip(tmp_path):
+ path = tmp_path / "rowavg_roundtrip.sras"
+ gen.write(path, n_angles=2, seed=20, samples_per_frame=128)
+ err = cache_file(str(path), "fft_rowavg", True, dc_threshold_mv=-1e9, row_avg_n=5)
+ assert err == "", err
+
+ sras = SrasFile(str(path))
+ assert sras.version == 7
+ assert sras.precomputed_row_avg_n == 5
+ assert all(x is not None for x in sras.precomputed_freq_mhz)
+
+
+def test_fft_rowavg_mode_requires_positive_n_and_threshold(tmp_path):
+ path = tmp_path / "rowavg_bad_args.sras"
+ gen.write(path, n_angles=1, seed=27, samples_per_frame=64)
+ assert cache_file(str(path), "fft_rowavg", True, dc_threshold_mv=0.0, row_avg_n=0)
+ assert cache_file(str(path), "fft_rowavg", True, dc_threshold_mv=None, row_avg_n=5)
+
+
+def test_raw_and_row_averaged_caches_never_cross_served(tmp_path):
+ """The central regression this feature must never allow: a raw request
+ served a row-averaged image (or vice versa), or a request at one window
+ size served a cache stored at a different one."""
+ path = tmp_path / "cross_serve.sras"
+ gen.write(path, n_angles=1, seed=21, samples_per_frame=128)
+ err = cache_file(str(path), "fft_rowavg", True, dc_threshold_mv=-1e9, row_avg_n=5)
+ assert err == "", err
+
+ sras = SrasFile(str(path))
+ assert cached_rf_image(sras, 0, None, apply_bg_sub=True, row_avg_n=0) is None, \
+ "a raw request must not be served a row-averaged cache"
+ assert cached_rf_image(sras, 0, None, apply_bg_sub=True, row_avg_n=3) is None, \
+ "a request at the wrong window size must not be served either"
+ served = cached_rf_image(sras, 0, None, apply_bg_sub=True, row_avg_n=5)
+ assert served is not None
+ assert np.array_equal(served, sras.precomputed_freq_mhz[0])
+
+
+def test_write_v7_cache_row_avg_n_carries_forward(tmp_path):
+ """A later DC-only write must leave a previously-written row-averaged
+ FFT block -- including its row_avg_n -- byte-for-byte unchanged."""
+ path = tmp_path / "carry_forward.sras"
+ gen.write(path, n_angles=2, seed=22, samples_per_frame=64)
+ err = cache_file(str(path), "fft_rowavg", True, dc_threshold_mv=-1e9, row_avg_n=7)
+ assert err == "", err
+
+ before = SrasFile(str(path))
+ assert before.precomputed_row_avg_n == 7
+ freq_before = [x.copy() for x in before.precomputed_freq_mhz]
+
+ err = cache_file(str(path), "dc", True)
+ assert err == "", err
+
+ after = SrasFile(str(path))
+ assert after.precomputed_row_avg_n == 7, "row_avg_n survives a DC-only write"
+ assert all(np.array_equal(after.precomputed_freq_mhz[a], freq_before[a])
+ for a in range(after.n_angles)), \
+ "the row-averaged FFT block itself is untouched by a DC-only write"
+
+
+def test_row_average_never_touches_dc_images(tmp_path):
+ path = tmp_path / "dc_untouched.sras"
+ gen.write(path, n_angles=2, seed=23, samples_per_frame=64)
+
+ src = SrasFile(str(path))
+ 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)]
+
+ assert cache_file(str(path), "dc", True) == ""
+ assert cache_file(str(path), "fft_rowavg", True,
+ dc_threshold_mv=-1e9, row_avg_n=6) == ""
+
+ after = SrasFile(str(path))
+ assert all(np.allclose(after.precomputed_dc3_mv[a], expect_dc3[a], atol=1e-4)
+ for a in range(after.n_angles))
+ assert all(np.allclose(after.precomputed_dc4_mv[a], expect_dc4[a], atol=1e-4)
+ for a in range(after.n_angles))
+
+
+def test_row_average_never_modifies_raw_waveform_data(tmp_path):
+ path = tmp_path / "waveform_untouched.sras"
+ gen.write(path, n_angles=2, seed=24, samples_per_frame=64)
+ orig = tmp_path / "waveform_untouched_orig.sras"
+ gen.write(orig, n_angles=2, seed=24, samples_per_frame=64)
+
+ assert cache_file(str(path), "fft_rowavg", True,
+ dc_threshold_mv=-1e9, row_avg_n=5) == ""
+
+ o, n = SrasFile(str(orig)), SrasFile(str(path))
+ assert all(np.array_equal(np.asarray(o.data[a]), np.asarray(n.data[a]))
+ for a in range(o.n_angles)), \
+ "waveform data untouched by a row-averaged cache write"
+
+
+def test_cach_v1_backward_compat_defaults_row_avg_n_zero(tmp_path):
+ """A v1 CACH tail predates row-averaged FFT caching entirely (no
+ row_avg_n byte at all) -- readers must still parse it in full, treating
+ it as row_avg_n=0. This is what protects an existing real-world v7
+ file's already-stored FFT cache from silently becoming unusable after
+ this change ships."""
+ path = tmp_path / "v1_rowavg.sras"
+ gen.write(path, n_angles=2, seed=25, samples_per_frame=128)
+ assert cache_file(str(path), "fft", True) == ""
+
+ v2 = SrasFile(str(path))
+ freq, entries = v2.precomputed_freq_mhz, list(range(v2.n_angles))
+ payload = struct.pack(fmt.CACH_HDR_FMT, fmt.CACH_MAGIC, 1, fmt.CACH_FLAG_FFT)
+ payload += struct.pack(fmt.SFFT_HDR_FMT_V1, fmt.SFFT_MAGIC,
+ fmt.SFFT_FLAG_BG_SUB, len(entries))
+ for a in entries:
+ payload += struct.pack(">H", a) + freq[a].astype(">f4").tobytes()
+ head = path.read_bytes()[:v2._cache_tail_offset()]
+ path.write_bytes(head + payload)
+
+ v1 = SrasFile(str(path))
+ assert v1.precomputed_row_avg_n == 0
+ assert all(np.array_equal(v1.precomputed_freq_mhz[a], freq[a]) for a in entries), \
+ "v1 images read back unchanged"
+ assert cached_rf_image(v1, 0, None, apply_bg_sub=True, row_avg_n=0) is not None
+ assert cached_rf_image(v1, 0, None, apply_bg_sub=True, row_avg_n=5) is None, \
+ "a v1 tail (predating this feature) can never satisfy a row-averaged request"
+
+
+def test_viewer_batch_row_average_dispatch(tmp_path, monkeypatch, no_fft):
+ """Driving the new 'Batch Compute Row-Averaged FFT and Store' action
+ end-to-end through the real menu handler: dialog values reach the
+ worker, the worker reaches cache_file, and the written file is
+ self-describing afterward. Also the exact scenario the row-averaged-FFT
+ recompute bug reported: before the fix, the display always asked for
+ row_avg_n=0 regardless of what the file actually had stored, so viewing
+ an angle after this batch action saw a phantom mismatch and launched a
+ full raw recompute on every view switch. This asserts that no longer
+ happens -- the stored row-averaged image is shown directly, with no
+ dispatched compute job and no real FFT."""
+ path = tmp_path / "rowavg_gui.sras"
+ gen.write(path, n_angles=2, seed=26, samples_per_frame=128)
+
+ class _StubDialog:
+ def __init__(self, *a, **k):
+ pass
+
+ def exec(self):
+ return QDialog.DialogCode.Accepted
+
+ def get_half_width(self):
+ return 6
+
+ def get_threshold_mv(self):
+ return -1e9 # mask nothing, keep the comparison simple
+
+ monkeypatch.setattr("sras_viewer.main_window.RowAverageFftOptionsDialog", _StubDialog)
+
+ app = QApplication.instance() or QApplication([]) # noqa: F841
+ win = SrasViewerWindow()
+ win.show()
+
+ dispatched = []
+ original_start = type(win)._start_compute
+ monkeypatch.setattr(type(win), "_start_compute",
+ lambda self: (dispatched.append(self.spin_angle.value()),
+ original_start(self))[1])
+ try:
+ win._load_file(str(path))
+ assert wait_until(lambda: win._sras is not None), "file loaded"
+ assert wait_until(lambda: all((a, CH4_IDX) in win._dc_cache
+ for a in range(win._sras.n_angles))), \
+ "DC precompute finished"
+
+ with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames",
+ return_value=([str(path)], "")):
+ win._on_batch_compute_row_avg()
+ assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran"
+ assert wait_until(lambda: win._sras is not None
+ and win._sras.version == 7), "file reloaded as v7"
+ pump(200)
+
+ assert win._sras.precomputed_row_avg_n == 6
+ assert "row-averaged n=6" in win.lbl_frame_warn.text(), win.lbl_frame_warn.text()
+
+ expected = compute.cached_rf_image(win._sras, 0, dc_threshold_mv=None,
+ apply_bg_sub=win.chk_bg_sub.isChecked(),
+ row_avg_n=6)
+ assert expected is not None, "the batch write left a readable row-averaged cache"
+
+ # The regression this batch action used to leave unfixed: viewing an
+ # angle afterward must show the stored row-averaged image directly,
+ # never fall through to a real (raw) recompute.
+ no_fft.clear()
+ dispatched.clear()
+ win.combo_channel.setCurrentIndex(CH1_IDX)
+ assert wait_until(lambda: win._current_ch == CH1_IDX), "CH1 displayed"
+ for a in range(win._sras.n_angles):
+ win.spin_angle.setValue(a)
+ pump(60)
+ assert np.allclose(win._current_image,
+ compute.cached_rf_image(
+ win._sras, a,
+ dc_threshold_mv=win.spin_threshold_mv.value(),
+ apply_bg_sub=win.chk_bg_sub.isChecked(),
+ row_avg_n=6), atol=1e-3), \
+ f"angle {a} shows the stored row-averaged image"
+ assert dispatched == [] and not no_fft, \
+ f"no recompute for row-averaged angles (jobs={dispatched}, fft={no_fft})"
+ finally:
+ win.close()
+ pump(300)
diff --git a/tools/bench_fft.py b/tools/bench_fft.py
new file mode 100644
index 0000000..68a1cc3
--- /dev/null
+++ b/tools/bench_fft.py
@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+"""Benchmark the FFT peak-search path: exact vs zoom, serial vs pooled.
+
+Reports wall time, waveforms/s, CPU utilization (utime+stime over wall, in
+cores), and verifies every variant against the exact reference image.
+
+Usage:
+ python tools/bench_fft.py # synthetic, pads 1/8/40
+ python tools/bench_fft.py --pads 40 --spf 2500 --rows 8 --frames 1024
+ python tools/bench_fft.py --real /path/big.sras --real-rows 32 --pads 40
+"""
+
+import argparse
+import resource
+import sys
+import tempfile
+import time
+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 compute_rf_image, set_fft_backend # noqa: E402
+from sras_format import SrasFile # noqa: E402
+import tools.make_test_sras as gen # noqa: E402
+from tools.check_equivalence import row_slice # noqa: E402
+
+
+def _timed(fn):
+ r0 = resource.getrusage(resource.RUSAGE_SELF)
+ t0 = time.perf_counter()
+ out = fn()
+ wall = time.perf_counter() - t0
+ r1 = resource.getrusage(resource.RUSAGE_SELF)
+ cpu = (r1.ru_utime - r0.ru_utime) + (r1.ru_stime - r0.ru_stime)
+ return out, wall, cpu / max(wall, 1e-9)
+
+
+def bench(sras, pads, backends):
+ n_wf = sum(int(sras.n_rows[a]) * int(sras.n_frames[a])
+ for a in range(sras.n_angles))
+ spf = sras.samples_per_frame
+ print(f"{n_wf} waveforms x {spf} samples, {sras.n_angles} angle(s)")
+ print(f"{'pad':>4} {'backend':>8} {'variant':>16} {'wall':>9} "
+ f"{'wf/s':>10} {'util':>6} match")
+
+ for pad in pads:
+ n_fft = spf * pad if pad > 1 else None
+ for backend in backends:
+ set_fft_backend(backend)
+
+ def run(**kw):
+ imgs = [compute_rf_image(sras, a, dc_threshold_mv=None,
+ apply_bg_sub=True, n_fft=n_fft, **kw)
+ for a in range(sras.n_angles)]
+ return np.concatenate([i.ravel() for i in imgs])
+
+ ref, wall, util = _timed(lambda: run(exact=True))
+ rows = [("exact(serial)", ref, wall, util, True)]
+ for label, kw in (("zoom(serial)", dict(max_workers=1)),
+ ("zoom(pool)", {})):
+ img, wall, util = _timed(lambda: run(**kw))
+ rows.append((label, img, wall, util, bool(np.array_equal(img, ref))))
+ for label, img, wall, util, ok in rows:
+ print(f"{pad:>4} {backend:>8} {label:>16} {wall:>8.2f}s "
+ f"{n_wf / wall:>10.0f} {util:>5.1f}x "
+ f"{'OK' if ok else 'MISMATCH'}")
+
+
+def main():
+ p = argparse.ArgumentParser(description=__doc__)
+ p.add_argument("--pads", default="1,8,40",
+ help="comma-separated pad factors (default 1,8,40)")
+ p.add_argument("--spf", type=int, default=2500)
+ p.add_argument("--rows", type=int, default=8)
+ p.add_argument("--frames", type=int, default=1024)
+ p.add_argument("--backends", default=None,
+ help="comma-separated (default: scipy,pyfftw if available)")
+ p.add_argument("--real", help="path to a real .sras file")
+ p.add_argument("--real-rows", type=int, default=32,
+ help="rows of angle 0 to use from the real file")
+ args = p.parse_args()
+
+ pads = [int(x) for x in args.pads.split(",")]
+ if args.backends:
+ backends = args.backends.split(",")
+ else:
+ backends = ["scipy"] + (["pyfftw"] if compute.PYFFTW_AVAILABLE else [])
+
+ if args.real:
+ sras = row_slice(SrasFile(args.real), 0, args.real_rows)
+ sras.data = [sras.data[0]]
+ sras.n_angles = 1
+ bench(sras, pads, backends)
+ else:
+ with tempfile.TemporaryDirectory(prefix="sras_bench_") as tmp:
+ path = Path(tmp) / "bench.sras"
+ gen.write(path, n_angles=1, seed=0, samples_per_frame=args.spf,
+ geometry=[(args.rows, args.frames)])
+ bench(SrasFile(str(path)), pads, backends)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/tools/check_equivalence.py b/tools/check_equivalence.py
index fa11968..9488040 100644
--- a/tools/check_equivalence.py
+++ b/tools/check_equivalence.py
@@ -1,13 +1,9 @@
#!/usr/bin/env python3
-"""Golden-output equivalence harness for the sras-viewer refactor.
+"""Golden-output equivalence harness for compute-path refactors.
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.
+for each. Run it before a refactor to capture a baseline, then again after
+and diff the two reports — every line must match.
Hashes canonicalise to native little-endian float64 before hashing, so a
deliberate dtype/byte-order change that preserves values does not show up as
@@ -29,23 +25,11 @@ 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
+from sras_format import SrasFile, CH3_IDX, CH4_IDX, adc_to_mv # noqa: E402
+from sras_compute import ( # noqa: E402
+ apply_alignment, compute_angle_alignment, compute_dc_image, compute_rf_image,
+)
+import tools.make_test_sras as gen # noqa: E402
def h(arr) -> str:
@@ -63,8 +47,8 @@ def row_slice(sras: SrasFile, angle_idx: int, n_rows: int) -> SrasFile:
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]
+ 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
@@ -109,7 +93,7 @@ def check_file(path: Path, lines: list[str], tag: str,
for bg in (False, True):
if bg and s.background is None:
continue
- for pad in (1, 2):
+ for pad in (1, 2, 4, 8, 40):
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,
@@ -135,7 +119,7 @@ def check_alignment(path: Path, lines: list[str], tag: str):
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)
+ res = compute_angle_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}")
@@ -165,7 +149,7 @@ def main():
help="directory for generated synthetic files")
args = p.parse_args()
- lines = [f"# layout: {_LAYOUT}", f"# numpy: {np.__version__}"]
+ lines = [f"# numpy: {np.__version__}"]
scratch = Path(args.scratch)
synth = scratch / "equiv_synth.sras"
@@ -179,6 +163,19 @@ def main():
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)
+ # A legacy v4 file exercises the uniform-geometry legacy layout through
+ # the same DC/FFT battery.
+ synth_v4 = scratch / "equiv_synth_v4.sras"
+ gen.write_legacy(synth_v4, version=4, n_angles=2, n_rows=6,
+ n_frames=14, samples_per_frame=48, seed=5)
+ check_file(synth_v4, lines, "v4", angles=[0, 1], n_rows=None)
+
+ # A big-endian int16 v6 file (real acquisitions are >i2; the other
+ # synthetics are int8).
+ synth_i16 = scratch / "equiv_synth_i16.sras"
+ gen.write(synth_i16, n_angles=2, seed=9, samples_per_frame=64, bps=2)
+ check_file(synth_i16, lines, "int16", angles=[0, 1], n_rows=None)
+
if args.real:
real = Path(args.real)
if real.exists():
diff --git a/tools/make_test_sras.py b/tools/make_test_sras.py
index a9bf3e2..f27e76f 100644
--- a/tools/make_test_sras.py
+++ b/tools/make_test_sras.py
@@ -18,9 +18,12 @@ import numpy as np
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
-# Packed straight from sras_format.py's own constants (not a hand-copy) so
-# this generator can never silently drift from what the real parser expects.
-from sras_format import GEO_FMT_V6, HDR_FMT as HDR_FMT_LEGACY, HDR_FMT_V6 # noqa: E402
+# Single source of truth for the byte layout: the reader's own constants.
+# The byte *assembly* below stays independent, so a writer bug can't be
+# masked by a matching reader bug.
+from sras_format import HDR_FMT as HDR_FMT_LEGACY # noqa: E402
+from sras_format import GEO_FMT_V6, HDR_FMT_V6 # noqa: E402
+from sras_compute import _rotation_matrix as _rot # noqa: E402
# Per-angle (n_rows, n_frames) — deliberately different per angle so ragged
# geometry handling is actually exercised.
@@ -46,12 +49,12 @@ def _preamble(ymult_v: float, yoff_adc: float, yzero_v: float) -> bytes:
def build(n_angles: int, seed: int, samples_per_frame: int,
- geometry: list[tuple[int, int]] | None = None) -> tuple[bytes, dict]:
+ geometry: list[tuple[int, int]] | None = None,
+ bps: int = 1) -> 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.
@@ -104,7 +107,9 @@ def build(n_angles: int, seed: int, samples_per_frame: int,
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()
+ # bps=2 stores the same values big-endian int16, exercising the
+ # reader's >i2 memmap path.
+ out += (block.astype(">i2") if bps == 2 else block).tobytes()
meta = {
"n_angles": n_angles,
@@ -123,12 +128,139 @@ def build(n_angles: int, seed: int, samples_per_frame: int,
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)
+ geometry: list[tuple[int, int]] | None = None,
+ bps: int = 1) -> dict:
+ payload, meta = build(n_angles, seed, samples_per_frame, geometry, bps=bps)
path.write_bytes(payload)
return meta
+# ---------------------------------------------------------------------------
+# Rotating-sample scan: one shape, imaged at several known rotations
+# ---------------------------------------------------------------------------
+#
+# The scan the angle-alignment path actually has to solve: every angle images
+# the *same* sample at a different known rotation and offset, and a correct
+# alignment stacks them all back into one shape. Two properties are
+# deliberately hostile:
+#
+# * every angle gets a different window size and a different, meaningless
+# stage x_start / y0 — alignment must ignore per-angle stage coordinates
+# entirely, so any code that reads them will visibly fail here;
+# * the pixel grid is strongly anisotropic (5 µm along x, 50 µm along y),
+# like the real instrument, so any registration that rotates raw indices
+# instead of millimetres shears the image and cannot converge.
+
+_ROT_DX_MM = 0.005 # x pitch, from velocity/laser_freq below
+_ROT_DY_MM = 0.05 # row spacing
+_ROT_BG_MV = 4.0
+_ROT_FG_MV = 160.0
+
+
+# How far the sample sits from the rotation axis. Non-zero on purpose: on the
+# real instrument every angle's scan window is centred on the rotation axis
+# while the sample is not, so each scan sees the sample somewhere else along a
+# circle. That offset is exactly what a wrong rotation pivot turns into a ring
+# of scans instead of a stack, so a centred test sample would hide the bug.
+_ROT_SAMPLE_OFFSET_MM = (0.55, 0.40)
+
+
+def _sample_shape_mv(u: np.ndarray, v: np.ndarray) -> np.ndarray:
+ """An asymmetric test sample in its own mm frame, chirally distinct at
+ every rotation (no 180° ambiguity) and with structure at several radii so
+ rotation is well determined."""
+ u = u - _ROT_SAMPLE_OFFSET_MM[0]
+ v = v - _ROT_SAMPLE_OFFSET_MM[1]
+ img = np.full(u.shape, _ROT_BG_MV, dtype=np.float32)
+ img[((u / 0.85) ** 2 + (v / 0.40) ** 2) <= 1.0] = _ROT_FG_MV # bar
+ img[(np.abs(u - 0.55) <= 0.22) & (np.abs(v - 0.62) <= 0.22)] = _ROT_FG_MV # nub
+ img[((u + 0.75) ** 2 + (v + 0.30) ** 2) <= 0.20 ** 2] = _ROT_FG_MV # dot
+ return img
+
+
+def write_rotating(path: Path, n_angles: int = 5, samples_per_frame: int = 4,
+ seed: int = 0) -> dict:
+ """Write a v6 file whose CH4 DC image is one sample seen at n_angles known
+ rotations, and return the ground truth each angle should register to.
+
+ ``truth[a] = (rotation_deg, (shift_x_mm, shift_y_mm))`` is the rigid map
+ from angle *a*'s local mm (origin at its own array center) to angle 0's —
+ exactly what ``register_angle_to_reference`` is supposed to recover.
+ """
+ rng = np.random.default_rng(seed)
+ n_ch, bps = 3, 1
+ cal = [(1.5625e-3, -87.04, 0.0), (2.0e-3, -60.0, 1.0e-3), (2.5e-3, -40.0, -2.0e-3)]
+ ymult_mv, yoff, yzero_mv = cal[2][0] * 1000, cal[2][1], cal[2][2] * 1000
+
+ stage_angles, geom, x_starts, y_starts, thetas, offsets = [], [], [], [], [], []
+ for a in range(n_angles):
+ stage = -37.0 * a # what the rotation stage reports
+ stage_angles.append(stage)
+ # The true image rotation is the negative of the stage's reported
+ # angle: the stage's positive sense is the opposite of math-positive
+ # (x toward y) in scan mm. Nothing may depend on knowing that — the
+ # registration search tries both signs.
+ thetas.append(-stage)
+ offsets.append((0.0, 0.0) if a == 0
+ else (float(rng.uniform(-0.3, 0.3)), float(rng.uniform(-0.3, 0.3))))
+ # A different window per angle, all centred on the same array center —
+ # the real instrument grows each angle's axis-aligned bounding box to
+ # cover the rotated ROI. Sized so the off-axis sample stays inside every
+ # window at every angle, keeping the expected result unambiguous.
+ geom.append((88 + 8 * a, 780 + 60 * a))
+ # Meaningless per-angle stage positions: correct alignment never reads
+ # them, so scattering them proves it.
+ x_starts.append(float(20.0 + rng.uniform(-6.0, 6.0)))
+ y_starts.append(float(30.0 + rng.uniform(-6.0, 6.0)))
+
+ out = bytearray()
+ out += struct.pack(
+ HDR_FMT_V6, b"SRAS", 6, n_angles,
+ x_starts[0], y_starts[0], 1.0, 1.0, _ROT_DY_MM,
+ _VELOCITY_MM_S, _VELOCITY_MM_S / _ROT_DX_MM, # velocity/freq -> 5 µm pitch
+ samples_per_frame, _SAMPLE_RATE_HZ, bps, n_ch,
+ )
+ out += np.array(stage_angles, dtype=">f4").tobytes()
+ for a, (n_rows, n_frames) in enumerate(geom):
+ out += struct.pack(GEO_FMT_V6, x_starts[a], 1.0, n_frames, n_rows)
+ for a, (n_rows, _) in enumerate(geom):
+ out += (y_starts[a] + np.arange(n_rows) * _ROT_DY_MM).astype(">f4").tobytes()
+ for ymult_v, yoff_a, yzero_v in cal:
+ p = _preamble(ymult_v, yoff_a, 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()
+
+ truth, dc4_images = {}, []
+ for a, (n_rows, n_frames) in enumerate(geom):
+ # Local mm of every pixel, measured from this angle's own array center.
+ lx = (np.arange(n_frames) - (n_frames - 1) / 2.0) * _ROT_DX_MM
+ ly = (np.arange(n_rows) - (n_rows - 1) / 2.0) * _ROT_DY_MM
+ gx, gy = np.meshgrid(lx, ly)
+ # local = R(theta) @ sample + offset, so sample = R(theta)^T @ (local - offset)
+ rel = np.stack([gx - offsets[a][0], gy - offsets[a][1]], axis=-1)
+ s = rel @ _rot(thetas[a]) # == rel @ R^T.T == R^T @ rel
+ dc4 = _sample_shape_mv(s[..., 0], s[..., 1])
+ dc4_images.append(dc4)
+
+ inv = _rot(-thetas[a])
+ truth[a] = (-thetas[a],
+ tuple(float(v) for v in -(inv @ np.array(offsets[a]))))
+
+ adc4 = np.clip(np.round((dc4 - yzero_mv) / ymult_mv + yoff), -128, 127).astype(np.int8)
+ block = np.zeros((n_rows, n_ch, n_frames, samples_per_frame), dtype=np.int8)
+ block[:, 2] = adc4[:, :, None] # CH4 carries the sample
+ block[:, 1] = 10 # CH3 flat
+ block[:, 0] = rng.integers(-40, 41, size=(n_rows, n_frames, samples_per_frame),
+ dtype=np.int8) # CH1 noise
+ out += block.tobytes()
+
+ path.write_bytes(bytes(out))
+ return {"n_angles": n_angles, "geometry": geom, "stage_angles_deg": stage_angles,
+ "truth": truth, "dc4_mv": dc4_images, "x_starts": x_starts,
+ "y_starts": y_starts, "dx_mm": _ROT_DX_MM, "dy_mm": _ROT_DY_MM}
+
+
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:
diff --git a/tools/test_gui.py b/tools/test_gui.py
deleted file mode 100644
index 482e11a..0000000
--- a/tools/test_gui.py
+++ /dev/null
@@ -1,485 +0,0 @@
-#!/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())
diff --git a/tools/test_refactor.py b/tools/test_refactor.py
deleted file mode 100644
index cbcb93e..0000000
--- a/tools/test_refactor.py
+++ /dev/null
@@ -1,365 +0,0 @@
-#!/usr/bin/env python3
-"""Behavioural tests for the sras-viewer refactor.
-
-Covers what the golden-hash harness can't: the v6->v7 cache round-trip
-(including block carry-forward), parallel-vs-serial identity, the no-mask
-fast path, and the ROI bounding-box mask optimisation.
-
-Usage: python tools/test_refactor.py [--scratch DIR]
-"""
-
-import argparse
-import shutil
-import sys
-import tempfile
-from pathlib import Path
-
-import numpy as np
-
-sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
-
-import sras_compute as compute # noqa: E402
-from sras_compute import ( # noqa: E402
- cache_file, compute_dc_image, compute_rf_image, dc_image_mv,
-)
-from sras_format import CH3_IDX, CH4_IDX, SrasFile, adc_to_mv # noqa: E402
-import tools.make_test_sras as gen # noqa: E402
-
-_failures: list[str] = []
-
-
-def check(name: str, ok: bool, detail: str = ""):
- print(f" {'PASS' if ok else 'FAIL'} {name}" + (f" — {detail}" if detail else ""))
- if not ok:
- _failures.append(name)
-
-
-def test_cache_roundtrip(scratch: Path):
- """v6 -> v7 for DC, then FFT, asserting the first block survives the
- second write (the carry-forward path in write_v7_cache)."""
- print("\ncache round-trip (v6 -> v7, both blocks)")
- path = scratch / "roundtrip.sras"
- gen.write(path, n_angles=3, seed=1, samples_per_frame=64)
-
- src = SrasFile(str(path))
- check("source is v6", src.version == 6, f"got v{src.version}")
- expect_dc3 = [dc_image_mv(src, a, CH3_IDX) for a in range(src.n_angles)]
- expect_dc4 = [dc_image_mv(src, a, CH4_IDX) for a in range(src.n_angles)]
- expect_fft = [compute_rf_image(src, a, dc_threshold_mv=None, apply_bg_sub=True)
- for a in range(src.n_angles)]
-
- err = cache_file(str(path), "dc", True)
- check("dc cache_file succeeded", err == "", err)
-
- after_dc = SrasFile(str(path))
- check("version flipped to 7", after_dc.version == 7, f"got v{after_dc.version}")
- check("dc3 stored for every angle",
- all(x is not None for x in after_dc.precomputed_dc3_mv))
- check("dc3 values round-trip",
- all(np.allclose(after_dc.precomputed_dc3_mv[a], expect_dc3[a], atol=1e-4)
- for a in range(after_dc.n_angles)))
- check("dc4 values round-trip",
- all(np.allclose(after_dc.precomputed_dc4_mv[a], expect_dc4[a], atol=1e-4)
- for a in range(after_dc.n_angles)))
- check("no fft block yet",
- all(x is None for x in after_dc.precomputed_freq_mhz))
- check("cached images are native float32",
- after_dc.precomputed_dc3_mv[0].dtype == np.float32
- and after_dc.precomputed_dc3_mv[0].dtype.byteorder in ("=", "|"),
- str(after_dc.precomputed_dc3_mv[0].dtype.byteorder))
- check("cached images are writable",
- after_dc.precomputed_dc3_mv[0].flags.writeable)
-
- err = cache_file(str(path), "fft", True)
- check("fft cache_file succeeded", err == "", err)
-
- both = SrasFile(str(path))
- check("fft stored for every angle",
- all(x is not None for x in both.precomputed_freq_mhz))
- check("fft values round-trip",
- all(np.allclose(both.precomputed_freq_mhz[a], expect_fft[a], atol=1e-3)
- for a in range(both.n_angles)))
- check("DC block carried forward through the FFT write",
- all(np.allclose(both.precomputed_dc3_mv[a], expect_dc3[a], atol=1e-4)
- for a in range(both.n_angles)))
- check("bg_sub flag persisted", both.precomputed_bg_sub is True)
-
- # The fast path must reproduce a fresh compute, and masking must still
- # apply on top of a cached (unmasked) image.
- fresh = SrasFile(str(path))
- fresh.precomputed_freq_mhz = [None] * fresh.n_angles
- dc4 = dc_image_mv(both, 0, CH4_IDX)
- thr = float(np.median(dc4))
- check("cached fast path == fresh compute (unmasked)",
- np.allclose(compute_rf_image(both, 0, dc_threshold_mv=None, apply_bg_sub=True),
- compute_rf_image(fresh, 0, dc_threshold_mv=None, apply_bg_sub=True),
- atol=1e-3))
- check("cached fast path == fresh compute (masked)",
- np.allclose(compute_rf_image(both, 0, dc_threshold_mv=thr, apply_bg_sub=True),
- compute_rf_image(fresh, 0, dc_threshold_mv=thr, apply_bg_sub=True),
- atol=1e-3))
-
- # Waveform data must be byte-identical to the pre-cache file.
- orig = scratch / "roundtrip_orig.sras"
- gen.write(orig, n_angles=3, seed=1, samples_per_frame=64)
- o, n = SrasFile(str(orig)), SrasFile(str(path))
- check("waveform data untouched by the cache write",
- all(np.array_equal(np.asarray(o.data[a]), np.asarray(n.data[a]))
- for a in range(o.n_angles)))
-
-
-def test_partial_v7_cache(scratch: Path):
- """Only some angles cached: uncached angles must compute, not read zeros.
- This is the v5 bug the ragged normalisation fixed, checked via v7."""
- print("\npartial cache (only some angles stored)")
- path = scratch / "partial.sras"
- gen.write(path, n_angles=3, seed=2, samples_per_frame=64)
-
- src = SrasFile(str(path))
- expected = [compute_rf_image(src, a, dc_threshold_mv=None, apply_bg_sub=True)
- for a in range(src.n_angles)]
- partial = [expected[0], None, expected[2]] # angle 1 deliberately absent
- src.write_v7_cache(new_freq_mhz=partial, new_bg_sub=True)
-
- reread = SrasFile(str(path))
- check("angle 1 is not cached", reread.precomputed_freq_mhz[1] is None)
- check("angles 0 and 2 are cached",
- reread.precomputed_freq_mhz[0] is not None
- and reread.precomputed_freq_mhz[2] is not None)
- img1 = compute_rf_image(reread, 1, dc_threshold_mv=None, apply_bg_sub=True)
- check("uncached angle computes rather than returning zeros",
- np.any(img1 != 0) and np.allclose(img1, expected[1], atol=1e-3))
-
-
-def test_parallel_identity(scratch: Path):
- """Forcing 1 worker vs many must give identical output — catches
- chunk-boundary and race bugs."""
- print("\nparallel vs serial identity")
- path = scratch / "parallel.sras"
- # Many rows, so the row loop actually splits into several chunks.
- n_rows, n_frames, spf = 48, 9, 256
- gen.write(path, n_angles=1, seed=3, samples_per_frame=spf,
- geometry=[(n_rows, n_frames)])
- sras = SrasFile(str(path))
-
- saved_budget, saved_workers = compute._TOTAL_BYTES_BUDGET, compute._MAX_WORKERS
- try:
- # Shrink the budget so chunk_rows collapses to 1 and every row is
- # its own chunk — the worst case for boundary bugs.
- compute._TOTAL_BYTES_BUDGET = 8 * n_frames * spf * 4
-
- compute._MAX_WORKERS = 1
- dc_serial = compute_dc_image(sras, 0, CH4_IDX)
- rf_serial = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True)
- dc4 = adc_to_mv(dc_serial, *sras.cal(CH4_IDX))
- thr = float(np.median(dc4))
- rf_masked_serial = compute_rf_image(sras, 0, dc_threshold_mv=thr,
- apply_bg_sub=True)
-
- chunk_rows, n_workers = compute._plan_chunks(
- n_rows, n_frames, spf, live_multiplier=compute._FFT_LIVE_MULTIPLIER)
- check("serial plan uses 1 worker", n_workers == 1, f"chunk_rows={chunk_rows}")
- check("work actually splits into multiple chunks", chunk_rows < n_rows,
- f"chunk_rows={chunk_rows} of {n_rows} rows")
-
- compute._MAX_WORKERS = 8
- chunk_rows, n_workers = compute._plan_chunks(
- n_rows, n_frames, spf, live_multiplier=compute._FFT_LIVE_MULTIPLIER)
- check("parallel plan uses >1 worker", n_workers > 1,
- f"chunk_rows={chunk_rows} workers={n_workers}")
-
- dc_par = compute_dc_image(sras, 0, CH4_IDX)
- rf_par = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True)
- rf_masked_par = compute_rf_image(sras, 0, dc_threshold_mv=thr, apply_bg_sub=True)
-
- check("dc image identical", np.array_equal(dc_serial, dc_par))
- check("rf image identical (unmasked)", np.array_equal(rf_serial, rf_par))
- check("rf image identical (masked)",
- np.array_equal(rf_masked_serial, rf_masked_par))
- finally:
- compute._TOTAL_BYTES_BUDGET, compute._MAX_WORKERS = saved_budget, saved_workers
-
-
-def test_nomask_equals_low_threshold(scratch: Path):
- """dc_threshold_mv=None must equal a threshold below every pixel, while
- skipping the CH4 read."""
- print("\nno-mask path")
- path = scratch / "nomask.sras"
- gen.write(path, n_angles=2, seed=4, samples_per_frame=128)
- sras = SrasFile(str(path))
- for a in range(sras.n_angles):
- none_img = compute_rf_image(sras, a, dc_threshold_mv=None, apply_bg_sub=True)
- low_img = compute_rf_image(sras, a, dc_threshold_mv=-1e9, apply_bg_sub=True)
- check(f"angle {a}: None == -1e9 threshold",
- np.array_equal(none_img, low_img))
- check(f"angle {a}: image is non-degenerate",
- len(np.unique(none_img)) > 1, f"{len(np.unique(none_img))} unique")
-
-
-def test_roi_mask():
- """The bbox-restricted mask must equal a full-grid point-in-polygon test."""
- print("\nROI mask (bbox fast path vs full grid)")
- from matplotlib.path import Path as MplPath
- from sras_viewer import RoiQuad
-
- rng = np.random.default_rng(0)
- x = np.linspace(-2.0, 3.0, 137)
- y = np.linspace(1.0, 4.0, 91)
-
- cases = {
- "axis-aligned rect": np.array([[0.0, 1.5], [1.0, 1.5], [1.0, 3.0], [0.0, 3.0]]),
- "skewed quad": np.array([[-0.5, 1.2], [1.7, 1.9], [1.2, 3.4], [-1.0, 2.6]]),
- "entirely outside": np.array([[8.0, 8.0], [9.0, 8.0], [9.0, 9.0], [8.0, 9.0]]),
- "covers whole grid": np.array([[-9.0, -9.0], [9.0, -9.0], [9.0, 9.0], [-9.0, 9.0]]),
- "straddles left edge": np.array([[-4.0, 2.0], [0.5, 2.0], [0.5, 3.0], [-4.0, 3.0]]),
- }
- for _ in range(5):
- cases[f"random {_}"] = rng.uniform([-2.5, 0.5], [3.5, 4.5], size=(4, 2))
-
- for name, pts in cases.items():
- roi = RoiQuad(pts)
- fast = roi.mask_for_grid(x, y)
- X, Y = np.meshgrid(x.astype(np.float64), y.astype(np.float64))
- slow = MplPath(pts).contains_points(
- np.column_stack([X.ravel(), Y.ravel()])).reshape(X.shape)
- check(f"{name} ({int(slow.sum())} px inside)", np.array_equal(fast, slow))
-
- # Descending y axis (images are stored top-down in some scans).
- roi = RoiQuad(cases["skewed quad"])
- y_desc = y[::-1]
- fast = roi.mask_for_grid(x, y_desc)
- X, Y = np.meshgrid(x.astype(np.float64), y_desc.astype(np.float64))
- slow = MplPath(cases["skewed quad"]).contains_points(
- np.column_stack([X.ravel(), Y.ravel()])).reshape(X.shape)
- check("descending y axis", np.array_equal(fast, slow))
-
-
-def test_legacy_parse_and_average(scratch: Path):
- """v2-v4 parsing plus the sras_average.py rewrite (which now streams via
- SrasFile rather than slurping the whole file)."""
- import subprocess
- print("\nlegacy formats (v2-v4) and sras_average")
- repo = Path(__file__).resolve().parent.parent
-
- for version in (2, 3, 4):
- path = scratch / f"legacy_v{version}.sras"
- meta = gen.write_legacy(path, version=version, n_angles=2, n_rows=4,
- n_frames=12, samples_per_frame=32, seed=version)
- s = SrasFile(str(path))
- check(f"v{version} parses", s.version == version, f"got v{s.version}")
- check(f"v{version} geometry uniform across angles",
- list(s.n_rows) == [4, 4] and list(s.n_frames) == [12, 12],
- f"rows={list(s.n_rows)} frames={list(s.n_frames)}")
- check(f"v{version} waveform data matches what was written",
- all(np.array_equal(np.asarray(s.data[a]), meta["data"][a])
- for a in range(s.n_angles)))
- check(f"v{version} background {'present' if version >= 4 else 'absent'}",
- (s.background is not None) == (version >= 4))
- check(f"v{version} precomputed stores are ragged lists",
- isinstance(s.precomputed_freq_mhz, list)
- and len(s.precomputed_freq_mhz) == s.n_angles)
- # DC image must equal a direct mean of the known input.
- expect = meta["data"][0][:, CH3_IDX, :, :].astype(np.float64).mean(axis=-1)
- check(f"v{version} DC image equals a direct mean",
- np.allclose(compute_dc_image(s, 0, CH3_IDX), expect, atol=1e-3))
-
- src = scratch / "legacy_v4.sras"
- meta = gen.write_legacy(src, version=4, n_angles=2, n_rows=4, n_frames=12,
- samples_per_frame=32, seed=4)
- dst = scratch / "legacy_v4_avg.sras"
- if dst.exists():
- dst.unlink()
- proc = subprocess.run(
- [sys.executable, str(repo / "sras_average.py"), str(src), str(dst), "--n", "4"],
- capture_output=True, text=True, cwd=repo)
- check("sras_average ran", proc.returncode == 0,
- (proc.stderr or proc.stdout).strip()[-200:])
-
- if dst.exists():
- avg = SrasFile(str(dst))
- check("averaged file parses", avg.version == 4)
- check("frame count divided by 4",
- list(avg.n_frames) == [3, 3], f"{list(avg.n_frames)}")
- check("angles/rows/channels unchanged",
- avg.n_angles == 2 and list(avg.n_rows) == [4, 4]
- and avg.n_channels == meta["n_channels"])
- check("calibration preserved",
- np.allclose(avg.ch_ymult_mv, SrasFile(str(src)).ch_ymult_mv))
- check("background preserved",
- np.array_equal(avg.background, SrasFile(str(src)).background))
- src_data = meta["data"]
- # int16 (not float32) before .mean(): matches average_rows' own
- # float64-accumulator behavior for integer input, so this doesn't
- # drift from what average_rows actually guarantees.
- expect0 = src_data[0][:, :, 0:4, :].astype(np.int16).mean(axis=2).astype(np.int16)
- check("first averaged group equals the mean of its 4 source frames",
- np.array_equal(np.asarray(avg.data[0])[:, :, 0, :], expect0))
-
- # Remainder handling: 12 frames / 5 -> 2 full groups + 1 partial.
- dst2 = scratch / "legacy_v4_avg5.sras"
- proc2 = subprocess.run([sys.executable, str(repo / "sras_average.py"),
- str(src), str(dst2), "--n", "5"],
- capture_output=True, text=True, cwd=repo)
- check("sras_average ran (--n 5)", proc2.returncode == 0,
- (proc2.stderr or proc2.stdout).strip()[-200:])
- if dst2.exists():
- check("partial trailing group kept by default",
- list(SrasFile(str(dst2)).n_frames) == [3, 3],
- f"{list(SrasFile(str(dst2)).n_frames)}")
- dst3 = scratch / "legacy_v4_avg5d.sras"
- proc3 = subprocess.run([sys.executable, str(repo / "sras_average.py"),
- str(src), str(dst3), "--n", "5", "--discard-remainder"],
- capture_output=True, text=True, cwd=repo)
- check("sras_average ran (--n 5 --discard-remainder)", proc3.returncode == 0,
- (proc3.stderr or proc3.stdout).strip()[-200:])
- if dst3.exists():
- check("--discard-remainder drops the partial group",
- list(SrasFile(str(dst3)).n_frames) == [2, 2],
- f"{list(SrasFile(str(dst3)).n_frames)}")
-
-
-def test_unsupported_version_reported(scratch: Path):
- """cache_file must report, not raise, for a file it can't handle."""
- print("\nerror reporting")
- bogus = scratch / "bogus.sras"
- bogus.write_bytes(b"SRAS" + bytes([99]) + b"\x00" * 200)
- err = cache_file(str(bogus), "dc", True)
- check("bad version returns an error string", bool(err), err)
- missing = cache_file(str(scratch / "does_not_exist.sras"), "dc", True)
- check("missing file returns an error string", bool(missing), missing)
-
-
-def main():
- p = argparse.ArgumentParser(description=__doc__)
- p.add_argument("--scratch")
- args = p.parse_args()
-
- tmp = None
- if args.scratch:
- scratch = Path(args.scratch)
- scratch.mkdir(parents=True, exist_ok=True)
- else:
- tmp = tempfile.mkdtemp(prefix="sras_test_")
- scratch = Path(tmp)
-
- try:
- test_cache_roundtrip(scratch)
- test_partial_v7_cache(scratch)
- test_parallel_identity(scratch)
- test_nomask_equals_low_threshold(scratch)
- test_roi_mask()
- test_legacy_parse_and_average(scratch)
- test_unsupported_version_reported(scratch)
- finally:
- if tmp:
- shutil.rmtree(tmp, ignore_errors=True)
-
- print()
- if _failures:
- print(f"{len(_failures)} FAILURE(S): " + ", ".join(_failures))
- sys.exit(1)
- print("All checks passed.")
-
-
-if __name__ == "__main__":
- main()