tka/sras-viewer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Split into four modules, deduplicate, and parallelize the compute paths

Browse Source 
Structure
  sras_format.py   parsing/writing, calibration, axes  (numpy + struct)
  sras_compute.py  DC/FFT images, alignment, batch cache  (+ scipy)
  sras_workers.py  Qt background workers
  sras_viewer.py   ROI, canvases, dialog, main window

format+compute import in 0.59s with no Qt or matplotlib (vs 2.96s for the
full app), which is what makes a spawn-based process pool worth using.

Deduplication
  - SrasFile.cal() and compute.dc_image_mv() replace the ADC->mV
    calibration incantation that appeared at ten call sites.
  - _read_preambles/_read_background/_set_calibration are shared by the
    legacy and v6 parsers instead of duplicated.
  - v5 PREC images are normalised into the same ragged per-angle list
    v6/v7 uses, removing four isinstance() branches and shortening
    compute_rf_image's fast path.
  - _run_worker replaces five copies of the QThread setup and five
    near-identical teardown methods; the two lifetime hazards they
    guarded against are now documented once, authoritatively.
  - _on_channel_changed defers to _update_controls_enabled rather than
    re-deriving the same six enable rules.
  - _corners_in_ref_frame, _CHANNEL_DISPLAY dict, unified progress-dialog
    helper, dead _on_roi_angle_edited stub removed.
  - sras_average.py builds on SrasFile instead of carrying a second copy
    of the header format, and streams per angle instead of loading the
    whole file (was a ~3x file-size peak).
  Executable lines: 2390 -> 2254, despite adding all of the below.

Parallelism
  - compute_rf_image/compute_dc_image map row chunks over a thread pool.
    Worker count is derived from the memory budget rather than the core
    count: on a large scan chunk_rows is already floored at 1 row (~75 MB
    of float32 at 7507x2500), so only the worker count can bound peak RAM.
  - DcPrecomputeWorker computes angles on a pool, emitting each result
    from its own QThread as it lands.
  - BatchCacheWorker runs one process per file via cache_file(); only
    paths and scalars cross the boundary. Per-process thread counts are
    divided so the two levels don't oversubscribe. Drops the old pass 1,
    which fully parsed every file just to weight a progress bar.
  - dc_threshold_mv=None means "no mask" and skips the CH4 read entirely
    — the batch FFT job previously read all of CH4 to compare against a
    threshold of -1e9.
  - Alignment mask and correlation stages map over angles; fft2/ifft2 use
    workers=-1.

Fixes found on the way
  - A partially-cached v5 file showed an all-zero image for uncached
    angles: the fast-path check was file-wide, not per-angle.
  - v7 cached images were read-only big-endian views; now native float32.

Verified: tools/check_equivalence.py produces byte-identical hashes for
167 outputs (DC/FFT across channels, angles, bg-sub, pad factors and
thresholds, plus the full alignment result) against ed0eba4, on both
synthetic files and a real 496 GB 17-angle v6 scan.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Thomas Ales
2026-07-30 23:38:08 -05:00
parent 95c273dce5
commit 55a4c5e42f
5 changed files with 2218 additions and 2275 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sras_compute.py
+577
View File
@@ -0,0 +1,577 @@
#!/usr/bin/env python3
"""Image computation and angle alignment for .sras scans.
Depends only on numpy/scipy (+ optional pyfftw) and sras_format, so a
multiprocessing child can import it without loading Qt or matplotlib —
which matters because Python 3.14 on macOS spawns rather than forks.
"""
import os
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
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
# ---------------------------------------------------------------------------
# FFT backend
# ---------------------------------------------------------------------------
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()
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."""
global _fft_backend
_fft_backend = name if (name != "pyfftw" or PYFFTW_AVAILABLE) else "numpy"
def get_fft_backend() -> str:
return _fft_backend
def _do_rfft(x: np.ndarray, n: int | None = None, axis: int = -1,
workers: int = 1) -> np.ndarray:
"""Dispatch rfft to the selected backend with optional multithreading."""
if _fft_backend == "pyfftw" and PYFFTW_AVAILABLE:
return pyfftw.interfaces.numpy_fft.rfft(x, n=n, axis=axis, threads=workers)
return scipy_fft.rfft(x, n=n, axis=axis, workers=workers)
# ---------------------------------------------------------------------------
# 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.
_TOTAL_BYTES_BUDGET = int(os.environ.get("SRAS_MEM_BUDGET_MB", 1536)) * 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 _chunk_rows_for(n_frames: int, samples_per_frame: int,
budget: int = _TOTAL_BYTES_BUDGET) -> int:
bytes_per_row = max(1, n_frames * samples_per_frame * 4) # float32
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) -> tuple[int, int]:
"""(chunk_rows, n_workers) sized so live_multiplier x chunk_rows x
n_frames x samples x 4 bytes x n_workers stays within the budget."""
per_chunk_budget = max(1, _TOTAL_BYTES_BUDGET // live_multiplier)
chunk_rows = _chunk_rows_for(n_frames, samples_per_frame, per_chunk_budget)
chunk_bytes = max(1, live_multiplier * chunk_rows * n_frames * samples_per_frame * 4)
cap = max_workers if max_workers is not None else _MAX_WORKERS
n_workers = int(max(1, min(cap,
_TOTAL_BYTES_BUDGET // chunk_bytes,
-(-n_rows // chunk_rows))))
return chunk_rows, n_workers
def _map_row_chunks(n_rows: int, chunk_rows: int, n_workers: int, fn):
"""Apply fn(r0, r1) over row chunks, in parallel when it pays.
Chunks write to disjoint output slices, so no locking is needed. numpy
ufuncs, scipy's pocketfft, and memmap page faults all release the GIL, so
threads give real parallelism here — and on a very large file they also
keep many more page-fault requests in flight, which is what the I/O path
wants.
"""
bounds = [(r0, min(r0 + chunk_rows, n_rows))
for r0 in range(0, n_rows, chunk_rows)]
if n_workers <= 1 or len(bounds) == 1:
for r0, r1 in bounds:
fn(r0, r1)
return
with ThreadPoolExecutor(max_workers=min(n_workers, len(bounds))) as pool:
list(pool.map(lambda b: fn(*b), bounds))
# ---------------------------------------------------------------------------
# Image computation
# ---------------------------------------------------------------------------
def compute_dc_image(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray:
"""Mean of each waveform → (n_rows, n_frames) float32, in ADC counts."""
n_rows, n_frames = sras.image_shape(angle_idx)
data = sras.data[angle_idx]
chunk_rows, n_workers = _plan_chunks(n_rows, n_frames, sras.samples_per_frame)
img = np.empty((n_rows, n_frames), dtype=np.float32)
def chunk(r0: int, r1: int):
img[r0:r1] = data[r0:r1, ch_idx, :, :].astype(np.float32).mean(axis=-1)
_map_row_chunks(n_rows, chunk_rows, n_workers, chunk)
return img
def dc_image_mv(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray:
"""DC image for (angle, channel) in mV, preferring a stored v5/v7 cache."""
cached = sras.cached_dc_mv(angle_idx, ch_idx)
if cached is not None:
return cached
return adc_to_mv(compute_dc_image(sras, angle_idx, ch_idx), *sras.cal(ch_idx))
def compute_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) -> 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
never run for them, since that's the expensive part. The masked-out
CH1 samples are also never *read*: the boolean mask is applied to the
raw memmap slice before any dtype conversion, so numpy only pages in
the bytes for pixels that pass the threshold (an untouched memmap page
is never read from disk).
*dc_threshold_mv* of None means "mask nothing", and skips reading and
averaging CH4 entirely — worth a third of the I/O when caching a whole
file, where the mask is applied later at display time.
If the DC4 image for this angle is already known (e.g. from the
DC-channel precompute cache), pass it as *dc4_mv* (mV, shape
(n_rows, n_frames)) to reuse it instead of re-reading CH4 here.
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.
"""
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
# ---- Chunked FFT path --------------------------------------------------
freq_axis = sras.freq_axis_mhz(n_fft)
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)
background = sras.background if (apply_bg_sub and sras.background is not None) else None
cal4 = sras.cal(CH4_IDX)
def chunk(r0: int, r1: int):
if dc_threshold_mv is None:
valid = None
else:
if dc4_mv is not None:
dc4_chunk = dc4_mv[r0:r1]
else:
dc4_chunk = adc_to_mv(
data[r0:r1, CH4_IDX, :, :].astype(np.float32).mean(axis=-1), *cal4)
valid = dc4_chunk >= dc_threshold_mv # True = run the FFT
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)
if background is not None:
waves -= 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]
else:
img[r0:r1] = freq_axis[peak_bins]
_map_row_chunks(n_rows, chunk_rows, n_workers, chunk)
return img
# ---------------------------------------------------------------------------
# Batch cache (Convert menu) — process-pool entry point
# ---------------------------------------------------------------------------
def cache_file(path: str, mode: str, apply_bg_sub: bool,
fft_backend: str = "numpy", max_workers: 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.
"""
global _MAX_WORKERS
try:
set_fft_backend(fft_backend)
if max_workers:
_MAX_WORKERS = max_workers
sras = SrasFile(path)
if sras.version not in (6, 7):
return (f"unsupported version {sras.version} — only v6/v7 files "
"can be batch-cached")
n = sras.n_angles
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)]
sras.write_v7_cache(new_dc3_mv=dc3, new_dc4_mv=dc4)
else:
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.
freq = [compute_rf_image(sras, a, dc_threshold_mv=None,
apply_bg_sub=effective_bg)
for a in range(n)]
sras.write_v7_cache(new_freq_mhz=freq, new_bg_sub=effective_bg)
return ""
except Exception as exc:
return str(exc)
# ---------------------------------------------------------------------------
# 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.
# ---------------------------------------------------------------------------
@dataclass
class AngleTransform:
rotation_deg: float
shift_mm: tuple[float, float] # (dx_mm, dy_mm) found by phase correlation
matrix: np.ndarray # (2,2): canvas (row,col) -> this angle's raw (row,col)
offset: np.ndarray # (2,)
@dataclass
class AlignmentResult:
ref_angle_idx: int
dc_threshold_mv: float
canvas_shape: tuple[int, int] # (n_rows, n_cols)
canvas_dx_mm: float
canvas_dy_mm: float
canvas_origin_mm: tuple[float, float] # mm at canvas pixel index (0, 0)
per_angle: dict[int, AngleTransform]
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)."""
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 _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 _rotation_matrix(theta_deg: float) -> np.ndarray:
t = np.radians(theta_deg)
c, s = np.cos(t), np.sin(t)
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:
return float(sras.angles_deg[angle_idx] - sras.angles_deg[ref_idx])
def _corners_in_ref_frame(sras: SrasFile, angle_idx: int, ref_idx: int,
shift_mm=(0.0, 0.0)) -> np.ndarray:
"""Angle *angle_idx*'s bbox corners, rotated about its own centroid into
the reference frame and translated by *shift_mm*. Shape (4, 2)."""
R = _rotation_matrix(_theta_deg(sras, angle_idx, ref_idx))
c_a = np.array(_bbox_center_mm(sras, angle_idx))
c_ref = np.array(_bbox_center_mm(sras, 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]
) -> tuple[np.ndarray, np.ndarray]:
"""matrix, offset s.t. raw_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
where theta = angles_deg[angle_idx] - angles_deg[ref_idx], c_ref/c_a are
each angle's own raw-bbox mm centroid (the rotation pivot — this keeps
rotated content centered, minimizing required canvas padding), and
A_out/D are the index<->mm scaling matrices for the canvas pitch and
this angle's own native pitch respectively.
"""
Rinv = _rotation_matrix(_theta_deg(sras, angle_idx, ref_idx)).T
cx_a, cy_a = _bbox_center_mm(sras, angle_idx)
cx_ref, cy_ref = _bbox_center_mm(sras, 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])
matrix = D @ Rinv @ A_out
offset = D @ Rinv @ (b_out - c_ref_v - shift) + D @ (c_a_v - origin_a)
return matrix, offset
def apply_alignment(result: AlignmentResult, angle_idx: int, img: np.ndarray,
order: int = 0) -> np.ndarray:
"""Resample any already-computed 2D image for `angle_idx` (same shape as
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
for any channel's image of that angle."""
t = result.per_angle[angle_idx]
return scipy_ndimage.affine_transform(
img.astype(np.float32, copy=False), t.matrix, offset=t.offset,
output_shape=result.canvas_shape, order=order,
mode="constant", cval=0.0)
def _block_mean_downsample(img: np.ndarray, factor: int) -> np.ndarray:
if factor <= 1:
return img
h, w = img.shape
h2, w2 = (h // factor) * factor, (w // factor) * factor
trimmed = img[:h2, :w2]
return trimmed.reshape(h2 // factor, factor, w2 // factor, factor).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)
def _working_canvas_for_pair(sras: SrasFile, ref_idx: int, a_idx: int,
dx: float, dy: 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 center) 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)])
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 _parallel_map(fn, items, n_workers: int) -> list:
"""fn over items, in order, threaded when it pays."""
items = list(items)
if n_workers <= 1 or len(items) <= 1:
return [fn(x) for x in items]
with ThreadPoolExecutor(max_workers=min(n_workers, len(items))) as pool:
return list(pool.map(fn, items))
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."""
n = sras.n_angles # already the *complete*-angle count for aborted v6 scans
dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx)
n_workers = max(1, min(_MAX_WORKERS, n))
# progress_cb fires from pool threads, so the counter behind it must be
# atomic. list.append is, and len() of a list is a consistent read.
_ticks: list[int] = []
def tick(base: int, span: int):
if progress_cb:
_ticks.append(1)
progress_cb(base + int(min(len(_ticks), n) / n * span))
# ---- Step 1: binarized CH4 mask per angle, native per-angle grid -----
def mask_for(a: int) -> np.ndarray:
dc4 = adc_to_mv(compute_dc_image(sras, a, CH4_IDX), *sras.cal(CH4_IDX))
m = (dc4 >= dc_threshold_mv).astype(np.float32)
tick(0, 25)
return m
masks = dict(enumerate(_parallel_map(mask_for, range(n), n_workers)))
# ---- Step 2: coarse correlation stage (downsample first, then rotate)
# Downsampling before affine_transform (not after) is what keeps this
# tractable for a v6 scan with thousands of rows/frames per angle.
max_dim = max(max(m.shape) for m in masks.values())
factor = max(1, int(np.ceil(max_dim / 1024)))
dx_c, dy_c = dx_ref * factor, dy_ref * factor
masks_small = {a: _block_mean_downsample(m, factor) for a, m in masks.items()}
_ticks.clear()
def shift_for(a: int) -> tuple[float, float]:
if a == ref_angle_idx:
tick(25, 50)
return (0.0, 0.0)
work_origin, work_shape = _working_canvas_for_pair(
sras, ref_angle_idx, a, dx_c, dy_c, margin_frac=0.3)
# Coarse canvas->raw affine at native pitch, then rescale by /factor
# so it maps coarse-canvas idx -> coarse (downsampled) raw idx —
# exact for block-mean downsampling (up to the trimmed remainder).
m_a, o_a = _build_affine_canvas_to_raw(
sras, a, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin)
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)
rotated_a = scipy_ndimage.affine_transform(
masks_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(
masks_small[ref_angle_idx], m_ref / factor, offset=o_ref / factor,
output_shape=work_shape, order=0, mode="constant", cval=0.0)
dr, dc = _phase_correlate_shift(embedded_ref, rotated_a)
tick(25, 50)
return (dc * dx_c, dr * dy_c)
shifts_mm = dict(enumerate(_parallel_map(shift_for, range(n), n_workers)))
# ---- Step 3: union bounding box over all angles (rotation+shift applied)
corners = np.vstack([_corners_in_ref_frame(sras, a, ref_angle_idx, 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)
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