Thomas Ales
1652 lines
74 KiB
Python
1652 lines
74 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, PAD_FACTOR_MAX, 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 _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))
|
|
|
|
|
|
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,
|
|
allow_dc_recompute: bool = True) -> np.ndarray | None:
|
|
"""The stored peak-frequency image for one angle (v5 PREC or v7 CACH),
|
|
masked and ready to display — or None if no stored image can serve these
|
|
settings and a real FFT is needed.
|
|
|
|
Stored images are unmasked, so they serve any threshold — but they carry
|
|
one bg-sub state and one zero-padding, and a view asking for either of the
|
|
others is asking for different numbers, not a different rendering of the
|
|
same ones. Padding is the sharp edge: a pad-10 view resolves peaks on a
|
|
10x finer bin grid, so handing back a pad-1 image would quietly undo the
|
|
setting. v5 PREC and CACH v1 tails record no pad and are natural
|
|
resolution by construction.
|
|
|
|
Cheap enough to call on the GUI thread — a copy plus a comparison —
|
|
*except* when the CH4 mask has to be read from disk to build it. Pass
|
|
*allow_dc_recompute* False to return None in that case instead, and leave
|
|
the whole-channel read to a worker.
|
|
"""
|
|
cached_freq = sras.precomputed_freq_mhz[angle_idx]
|
|
# n_fft None means "no padding", i.e. exactly samples_per_frame points.
|
|
want_n_fft = n_fft if n_fft is not None else sras.samples_per_frame
|
|
if (cached_freq is None
|
|
or want_n_fft != sras.precomputed_pad_factor * sras.samples_per_frame
|
|
or sras.precomputed_bg_sub != (apply_bg_sub and sras.background is not None)):
|
|
return None
|
|
|
|
if dc_threshold_mv is None:
|
|
return cached_freq.copy()
|
|
|
|
# DC4 mask, in priority order: stored 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_img is None:
|
|
if not allow_dc_recompute:
|
|
return None
|
|
dc4_img = adc_to_mv(compute_dc_image(sras, angle_idx, CH4_IDX),
|
|
*sras.cal(CH4_IDX))
|
|
|
|
freq_img = cached_freq.copy()
|
|
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) -> 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.
|
|
|
|
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. That test is
|
|
cached_rf_image, which the viewer also applies before it ever dispatches
|
|
a compute at all.
|
|
"""
|
|
# ---- Fast path: precomputed image (v5 PREC or v7 CACH) ----------------
|
|
cached = cached_rf_image(sras, angle_idx, dc_threshold_mv,
|
|
apply_bg_sub=apply_bg_sub, n_fft=n_fft,
|
|
dc4_mv=dc4_mv)
|
|
if cached is not None:
|
|
return cached
|
|
|
|
# ---- Chunked FFT path --------------------------------------------------
|
|
n_rows, n_frames = sras.image_shape(angle_idx)
|
|
data = sras.data[angle_idx]
|
|
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)
|
|
|
|
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)
|
|
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_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
|
|
|
|
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
|
|
# 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 + i, CH1_IDX]
|
|
dst = waves[offs[i]:offs[i + 1]]
|
|
dst[:] = raw[valid[i]] if 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) -> 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, worker cap and pad factor are passed explicitly because
|
|
module globals do not survive a spawn.
|
|
|
|
*pad_factor* is stored alongside the images: a cache is only usable by a
|
|
view asking for the same padding, so caching at the viewer's own setting
|
|
is the difference between a batch that pays off and one that can never be
|
|
read back (see cached_rf_image).
|
|
"""
|
|
global _MAX_WORKERS
|
|
try:
|
|
if mode not in ("dc", "fft"):
|
|
return f"unknown cache mode {mode!r} (expected 'dc' or 'fft')"
|
|
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_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)
|
|
else:
|
|
effective_bg = apply_bg_sub and sras.background is not None
|
|
pad = max(1, min(PAD_FACTOR_MAX, int(pad_factor)))
|
|
n_fft = None if pad <= 1 else sras.samples_per_frame * pad
|
|
# 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)]
|
|
sras.write_v7_cache(new_freq_mhz=freq, new_bg_sub=effective_bg,
|
|
new_pad_factor=pad)
|
|
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) -> list[float]:
|
|
"""Coarse rotation candidates: a window around *both* signs of the stage's
|
|
reported angle change. Scoring both 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."""
|
|
out: list[float] = []
|
|
for center in (-nominal_deg, 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) -> 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 both
|
|
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.
|
|
"""
|
|
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)
|
|
thetas = _rotation_candidates(nominal, search_deg, coarse_step_deg)
|
|
|
|
# ---- 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)
|
|
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],
|
|
*, margin_frac: float = 0.0, snap: bool = True
|
|
) -> 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).
|
|
|
|
snap=True aligns the canvas grid 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". It requires pitch_mm to be the reference's own
|
|
pitch; the manual-alignment preview passes a coarser pitch and snap=False.
|
|
|
|
margin_frac pads the box on every side: 0 for a final canvas, nonzero for
|
|
ManualAlignmentDialog's preview canvas, which needs headroom so an ordinary
|
|
translation nudge never has to trigger a full canvas resize (an extreme
|
|
nudge can still push content past this padding; accepted, and cheap to
|
|
recover from by re-opening the dialog).
|
|
"""
|
|
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)
|
|
if margin_frac:
|
|
pad_x, pad_y = (x_max - x_min) * margin_frac, (y_max - y_min) * margin_frac
|
|
x_min, x_max = x_min - pad_x, x_max + pad_x
|
|
y_min, y_max = y_min - pad_y, y_max + pad_y
|
|
|
|
center = ref_center_mm(sras, ref_angle_idx)
|
|
if snap:
|
|
# 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)
|
|
else:
|
|
origin_ref = np.array([x_min, y_min if dy > 0 else y_max])
|
|
shape = (int(np.ceil((y_max - y_min) / abs(dy))) + 1,
|
|
int(np.ceil((x_max - x_min) / dx)) + 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, snap=True)
|
|
|
|
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 _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.
|
|
|
|
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
|
|
# 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 ManualAlignmentDialog's own downsampled preview calls that per keystroke
|
|
# — see that class's docstring for how it limits each nudge to reprojecting only
|
|
# the actively-edited angle.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
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 ManualAlignmentDialog's
|
|
live preview repeatedly calls (once per keystroke, for only the
|
|
actively-nudged angle). *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()
|
|
},
|
|
}
|
|
path.write_text(json.dumps(payload, indent=2))
|
|
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
|
|
(ManualAlignmentDialog._on_clear) surfaces them rather than silently
|
|
pretending the destructive action succeeded."""
|
|
path = sidecar_path(sras.path)
|
|
try:
|
|
path.unlink()
|
|
return True
|
|
except FileNotFoundError:
|
|
return False
|