Thomas Ales [M S E]
105b9514a5
# Conflicts: # sras_compute.py # sras_viewer.py # sras_workers.py # tools/make_test_sras.py # tools/test_refactor.py
2002 lines
91 KiB
Python
2002 lines
91 KiB
Python
#!/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 atexit
|
|
import json
|
|
import os
|
|
import threading
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
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, MAX_PAD_FACTOR, SrasFile,
|
|
adc_to_mv)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# FFT backend
|
|
# ---------------------------------------------------------------------------
|
|
|
|
try:
|
|
import pyfftw
|
|
PYFFTW_AVAILABLE = True
|
|
except ImportError:
|
|
PYFFTW_AVAILABLE = False
|
|
|
|
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. "numpy" is accepted as a legacy alias for "scipy"."""
|
|
global _fft_backend
|
|
if name == "numpy":
|
|
name = "scipy"
|
|
_fft_backend = name if (name == "pyfftw" and PYFFTW_AVAILABLE) else "scipy"
|
|
|
|
|
|
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)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 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
|
|
# ---------------------------------------------------------------------------
|
|
#
|
|
# 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)
|
|
|
|
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,
|
|
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_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,
|
|
budget: int | None = None) -> tuple[int, int]:
|
|
"""(chunk_rows, n_workers) such that all concurrent chunks together fit
|
|
the budget: n_workers x live_multiplier x chunk_rows x n_frames x
|
|
samples x 4 bytes <= budget.
|
|
|
|
The worker count is chosen *first* and the chunk sized to it. Sizing the
|
|
chunk first is the trap: _chunk_rows_for spends whatever budget it is
|
|
given, so a single chunk would always consume the lot and leave room for
|
|
exactly one worker — no concurrency, on precisely the large scans that
|
|
need it most.
|
|
|
|
A caller that is itself running several of these concurrently 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.
|
|
"""
|
|
total = _TOTAL_BYTES_BUDGET if budget is None else max(1, budget)
|
|
cap = max_workers if max_workers is not None else _MAX_WORKERS
|
|
bytes_per_row = max(1, live_multiplier * n_frames * samples_per_frame * 4)
|
|
|
|
# A chunk is at least one row, so that alone caps how many can be live.
|
|
n_workers = int(max(1, min(cap, n_rows, total // bytes_per_row)))
|
|
chunk_rows = _chunk_rows_for(n_frames, samples_per_frame,
|
|
max(1, total // (n_workers * live_multiplier)))
|
|
# With chunk_rows known, more workers than chunks buys nothing.
|
|
n_workers = int(max(1, min(n_workers, -(-n_rows // chunk_rows))))
|
|
return chunk_rows, n_workers
|
|
|
|
|
|
def _map_row_chunks(n_rows: int, chunk_rows: int, n_workers: int, fn,
|
|
should_stop=None):
|
|
"""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.
|
|
|
|
*should_stop* is polled per chunk so a cancelled job abandons work within
|
|
one chunk rather than one whole angle — on a large scan an angle is
|
|
~40 s, which is far too long to block application shutdown.
|
|
"""
|
|
bounds = [(r0, min(r0 + chunk_rows, n_rows))
|
|
for r0 in range(0, n_rows, chunk_rows)]
|
|
|
|
def run(b):
|
|
if should_stop is not None and should_stop():
|
|
return
|
|
fn(*b)
|
|
|
|
if n_workers <= 1 or len(bounds) == 1:
|
|
for b in bounds:
|
|
run(b)
|
|
return
|
|
with ThreadPoolExecutor(max_workers=min(n_workers, len(bounds))) as pool:
|
|
list(pool.map(run, bounds))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Image computation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def plan_angle_level(sras: SrasFile,
|
|
live_multiplier: int = 1) -> tuple[int, int]:
|
|
"""(n_angle_workers, per_angle_budget) for a caller that parallelises
|
|
over angles instead of over rows.
|
|
|
|
Callers that parallelise over angles must not also let each angle
|
|
parallelise over rows: the two levels multiply, in threads and in memory.
|
|
Each angle is therefore given max_workers=1 and the returned budget share,
|
|
which together keep total live buffers within _TOTAL_BYTES_BUDGET.
|
|
"""
|
|
biggest = max(range(sras.n_angles), key=lambda a: int(sras.n_frames[a]))
|
|
_, n_workers = _plan_chunks(int(sras.n_rows[biggest]),
|
|
int(sras.n_frames[biggest]),
|
|
sras.samples_per_frame,
|
|
live_multiplier=live_multiplier)
|
|
n_workers = max(1, min(n_workers, sras.n_angles))
|
|
return n_workers, max(1, _TOTAL_BYTES_BUDGET // n_workers)
|
|
|
|
|
|
def compute_dc_image(sras: SrasFile, angle_idx: int, ch_idx: int,
|
|
max_workers: int | None = None,
|
|
budget: int | None = None,
|
|
should_stop=None) -> np.ndarray:
|
|
"""Mean of each waveform → (n_rows, n_frames) float32, in ADC counts.
|
|
|
|
Pass max_workers=1 when the caller is already parallelising over angles.
|
|
If *should_stop* ever returns True the result is incomplete — the caller
|
|
is expected to be abandoning it.
|
|
"""
|
|
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,
|
|
max_workers=max_workers, budget=budget)
|
|
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, should_stop=should_stop)
|
|
return img
|
|
|
|
|
|
def dc_image_mv(sras: SrasFile, angle_idx: int, ch_idx: int,
|
|
max_workers: int | None = None, budget: int | None = None,
|
|
should_stop=None) -> 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, max_workers=max_workers,
|
|
budget=budget, should_stop=should_stop),
|
|
*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,
|
|
n_fft: int | None = None,
|
|
dc4_mv: np.ndarray | None = None,
|
|
max_workers: int | None = None,
|
|
budget: int | None = None,
|
|
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
|
|
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.
|
|
|
|
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) 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 = 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 --------------------------------------------------
|
|
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)
|
|
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:
|
|
# already-cached DC block (skips reading CH4 raw waveforms entirely),
|
|
# then the caller-supplied image, else recompute per chunk below.
|
|
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
|
|
|
|
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:
|
|
if dc4_full is not None:
|
|
dc4_chunk = dc4_full[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
|
|
|
|
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)
|
|
|
|
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,)
|
|
|
|
if pool is None:
|
|
for i in range(r1 - r0):
|
|
read_row(i)
|
|
else:
|
|
list(pool.map(read_row, range(r1 - r0)))
|
|
|
|
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
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Batch cache (Convert menu) — process-pool entry point
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def cache_file(path: str, mode: str, apply_bg_sub: bool,
|
|
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", "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
|
|
|
|
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
|
|
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 = _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)
|
|
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. 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, n_fft=n_fft)
|
|
for a in range(n)]
|
|
# 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)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Angle alignment (Fusion menu)
|
|
#
|
|
# 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) 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
|
|
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]
|
|
|
|
|
|
@dataclass
|
|
class ManualAngleParams:
|
|
"""One angle's manual-alignment state, independent of any canvas.
|
|
|
|
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)
|
|
|
|
|
|
# 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). 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 _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 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:
|
|
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 nominal_delta_deg(sras: SrasFile, angle_idx: int, ref_idx: int) -> float:
|
|
"""The rotation-stage's own reported angle change between two angles.
|
|
|
|
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])
|
|
|
|
|
|
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 _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.
|
|
|
|
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.
|
|
|
|
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.
|
|
"""
|
|
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 - shift) - off) + np.asarray(src_center_idx, dtype=np.float64)
|
|
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 image) 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_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 // fy) * fy, (w // fx) * fx
|
|
if h2 == 0 or w2 == 0:
|
|
return img
|
|
trimmed = img[:h2, :w2]
|
|
return trimmed.reshape(h2 // fy, fy, w2 // fx, fx).mean(axis=(1, 3))
|
|
|
|
|
|
# ---- 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 _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 _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)))
|
|
|
|
|
|
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 RigidFit(0.0, (0.0, 0.0), 1.0, "reference")
|
|
|
|
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")
|
|
|
|
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)
|
|
|
|
# ---- 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)
|
|
|
|
if best[3] == "none":
|
|
return RigidFit(0.0, (0.0, 0.0), -1.0, "none")
|
|
|
|
# ---- 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)
|
|
|
|
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 = 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:
|
|
"""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,
|
|
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
|
|
# Parallel over angles, serial within each — see plan_angle_level.
|
|
n_workers, angle_budget = plan_angle_level(sras)
|
|
|
|
# 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: 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),
|
|
*sras.cal(CH4_IDX))
|
|
tick(0, 25)
|
|
return dc4
|
|
|
|
dc4_mv = dict(enumerate(_parallel_map(dc4_for, range(n), n_workers)))
|
|
|
|
# ---- 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()
|
|
|
|
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
|
|
|
|
fits = dict(enumerate(_parallel_map(
|
|
fit_for, range(n), _registration_workers(sras, fine_dim))))
|
|
|
|
# ---- 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 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 reproject_mask(sras: SrasFile, angle_idx: int, ref_angle_idx: int,
|
|
mask: np.ndarray, rotation_deg: float,
|
|
shift_mm: tuple[float, float],
|
|
canvas_pitch_mm: tuple[float, float],
|
|
canvas_origin_mm: tuple[float, float],
|
|
canvas_shape: tuple[int, int],
|
|
*, 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 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)
|
|
|
|
|
|
def build_manual_alignment(sras: SrasFile, ref_angle_idx: int,
|
|
dc_threshold_mv: float,
|
|
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 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.
|
|
|
|
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.
|
|
"""
|
|
params = {a: per_angle_params.get(a, ManualAngleParams())
|
|
for a in range(sras.n_angles)}
|
|
params[ref_angle_idx] = ManualAngleParams()
|
|
return _result_from_params(sras, ref_angle_idx, dc_threshold_mv, params)
|
|
|
|
|
|
# ---- Sidecar persistence (<name>.sras.align.json) -------------------------
|
|
# 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:
|
|
ref_angle_idx: int
|
|
dc_threshold_mv: float
|
|
per_angle: dict[int, ManualAngleParams]
|
|
|
|
|
|
def sidecar_path(sras_path) -> Path:
|
|
"""<name>.sras.align.json next to the scan file. A thin, independently
|
|
testable function since save/load/delete and the GUI's status messages
|
|
all need the identical path."""
|
|
p = Path(sras_path)
|
|
return p.with_name(p.name + ".align.json")
|
|
|
|
|
|
# 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. 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,
|
|
"ref_angle_idx": ref_angle_idx,
|
|
"dc_threshold_mv": dc_threshold_mv,
|
|
"per_angle": {
|
|
str(a): {"rotation_deg": p.rotation_deg, "shift_mm": list(p.shift_mm)}
|
|
for a, p in per_angle.items()
|
|
},
|
|
}
|
|
# Write to a temp file in the same directory then rename over the target
|
|
# — os.replace() is atomic on both POSIX and Windows, so a crash mid-write
|
|
# leaves either the old sidecar or the new one, never a truncated/corrupt
|
|
# file that load_manual_alignment would silently treat as "no alignment".
|
|
tmp_path = path.with_name(path.name + f".tmp{os.getpid()}")
|
|
tmp_path.write_text(json.dumps(payload, indent=2))
|
|
os.replace(tmp_path, path)
|
|
return path
|
|
|
|
|
|
def load_manual_alignment(sras: SrasFile) -> ManualAlignmentSidecar | None:
|
|
"""Read <sras.path>'s sidecar JSON if present, else None — never raises:
|
|
a missing, corrupt, foreign, or version-mismatched JSON file must not
|
|
block opening the .sras file itself (see SrasViewerWindow._on_load_done).
|
|
Per-angle entries for an angle index no longer present in *sras* (e.g.
|
|
re-scanned with fewer angles) are silently dropped.
|
|
"""
|
|
path = sidecar_path(sras.path)
|
|
if not path.exists():
|
|
return None
|
|
try:
|
|
raw = json.loads(path.read_text())
|
|
if raw.get("schema_version") != _SIDECAR_SCHEMA_VERSION:
|
|
return None
|
|
per_angle = {
|
|
int(a): ManualAngleParams(
|
|
float(v["rotation_deg"]),
|
|
(float(v["shift_mm"][0]), float(v["shift_mm"][1])))
|
|
for a, v in raw.get("per_angle", {}).items()
|
|
if int(a) < sras.n_angles
|
|
}
|
|
return ManualAlignmentSidecar(
|
|
ref_angle_idx=int(raw.get("ref_angle_idx", 0)),
|
|
dc_threshold_mv=float(raw.get("dc_threshold_mv", 0.0)),
|
|
per_angle=per_angle)
|
|
except (OSError, ValueError, KeyError, TypeError, IndexError,
|
|
json.JSONDecodeError):
|
|
return None
|
|
|
|
|
|
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
|
|
(the wizard's Clear path) surfaces them rather than silently
|
|
pretending the destructive action succeeded."""
|
|
path = sidecar_path(sras.path)
|
|
try:
|
|
path.unlink()
|
|
return True
|
|
except FileNotFoundError:
|
|
return False
|