Thomas Ales
6d0e30b9ce
The alignment machinery never modified a scan: it produced an
AlignmentResult and every consumer resampled on the fly. That is right
for a viewer but means the aligned stack cannot leave the process, so no
other tool can read it and reopening the scan redoes the registration.
sras_align_export.write_aligned_sras bakes an alignment into a new v6
file: each angle is gathered onto the shared canvas by nearest neighbour,
so every output angle shares one grid and the file opens already aligned.
Registering the export against itself returns identity, which is the test
that pins the whole index chain.
Three details that are easy to get wrong and are now covered:
* Rounding must be floor(x + 0.5), not np.rint. scipy's order=0 rounds
halves away from zero, and the canvas is snapped to the reference's
pixel grid, so exact halves are common rather than hypothetical.
* Out-of-bounds must be tested on the fractional coordinate against
[0, n-1], not on the rounded index, or a one-pixel rim gets real data
everywhere the Aligned View shows padding.
* ...but with a tolerance, because the mm-space affine chain lands an
exactly-integer transform a few times 1e-13 off. A bare >= 0 drops
the *reference* angle's entire first row and last column.
Padding is the per-channel ADC code nearest 0 mV, not zero: zero ADC
decodes to ~+100 mV on real calibration and would masquerade as sample.
Source rows are served from sliding in-RAM bands. A rotated angle maps
one output row to a diagonal across the source, so indexing a memmap in
output order refaults nearly the whole angle per row — terabytes of
paging for a gigabyte of data.
Also in compute: crop_alignment_result (a crop is pure index translation,
so it folds into the affine's offset rather than becoming a second
transform), overlap_stats, largest_rect_at_least for a crop that stays
inside the overlap region, and seed/sign/refine knobs on
register_angle_to_reference whose defaults leave existing behaviour and
tests byte-identical.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
1896 lines
86 KiB
Python
1896 lines
86 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, 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))
|
|
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 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 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); no custom zero-padding is
|
|
requested (n_fft is None) — the store is always natural-resolution;
|
|
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.
|
|
|
|
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 n_fft is not None
|
|
or sras.precomputed_bg_sub != (apply_bg_sub and sras.background is not None)
|
|
or sras.precomputed_row_avg_n != row_avg_n):
|
|
return None
|
|
freq_img = cached_freq.copy()
|
|
if dc_threshold_mv is not None:
|
|
# DC4 mask, in priority order: already-cached DC block, caller-
|
|
# supplied image, or a fresh (cheap — no FFT) recompute.
|
|
dc4_img = sras.cached_dc_mv(angle_idx, CH4_IDX)
|
|
if dc4_img is None:
|
|
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[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)
|
|
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_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
|
|
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,
|
|
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, natural-resolution, 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.
|
|
|
|
The stored FFT cache is always natural-resolution (pad 1): the v7 SFFT
|
|
block records no pad factor, and padded views compute live fast enough
|
|
(see _peak_bins_zoom) that caching them is not worth a format change.
|
|
"""
|
|
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')"
|
|
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)
|
|
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)
|
|
for a in range(n)]
|
|
sras.write_v7_cache(new_freq_mhz=freq, new_bg_sub=effective_bg)
|
|
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, 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)
|
|
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],
|
|
*, 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 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 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.
|
|
|
|
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
|