Thomas Ales
8348ad313c
tests/test_stored_cache.py exercised two features that were never built, so six of its tests had been failing on main. Both are now implemented. Pad-factor caching. A padded view could not use a stored FFT cache at all: the store was pad 1 by definition and cached_rf_image rejected any n_fft outright, so a user working at a pad factor got nothing from batch-computing a file. CACH tail v3 records the pad the images were resolved at, cache_file computes at a requested pad, and the batch actions pass the viewer's own pad down — while still refusing a store resolved at a different pad, since a padded FFT interpolates between the natural bins and so resolves genuinely different peak frequencies. v1/v2 tails read as pad 1 and keep working. Stored-cache dispatch. _refresh_display only ever consulted this window's in-session dicts, so after a batch every angle change still queued a worker and a progress popup for an image already on disk — the exact cost the batch was run to avoid. It now checks the file's own DC/FFT blocks first, asking with allow_dc_recompute=False so the GUI thread never touches I/O. When the stored cache genuinely cannot serve the view, the scan info panel says why rather than leaving the silent recompute a mystery. Two bugs surfaced on the way: - The angle spinbox was wired on editingFinished, which QAbstractSpinBox emits only on Return or focus-out — never on a step. Clicking its arrows, the ordinary way to walk a scan, moved the number and left the image behind. Now valueChanged with keyboard tracking off, which fires on a step and once on commit, but not per keystroke mid-typing. Every existing GUI test called _on_view_changed() by hand and so could not have caught this; two new tests pin it and fail against the old wiring. - A plain "fft" batch over a file previously cached with fft_rowavg carried the old row_avg_n forward, labelling raw images as row-averaged. It now writes row_avg_n=0 explicitly. Verified: 111 passed (was 103 passed / 6 failed), and tools/check_equivalence.py is byte-identical to the pre-change baseline. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
433 lines
17 KiB
Python
433 lines
17 KiB
Python
#!/usr/bin/env python3
|
||
"""Background workers for the SRAS viewer.
|
||
|
||
Every worker is a plain QObject moved onto its own QThread by
|
||
SrasViewerWindow._run_worker, exposing signals only. Workers must never touch
|
||
GUI-thread-owned state (the display caches in particular) — they take
|
||
everything they need through their constructor and hand results back by signal.
|
||
"""
|
||
|
||
import os
|
||
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
|
||
from concurrent.futures.process import BrokenProcessPool
|
||
|
||
import numpy as np
|
||
from PyQt6.QtCore import QObject, pyqtSignal
|
||
|
||
import sras_compute as compute
|
||
from sras_align_export import write_aligned_sras
|
||
from sras_compute import cache_file, compute_rf_image, dc_image_mv
|
||
from sras_format import CH3_IDX, CH4_IDX, SrasFile
|
||
|
||
# Concurrency caps. Batch conversion runs one process per file, and each of
|
||
# those processes threads internally, so the two must be divided rather than
|
||
# both set to the core count. Files also commonly sit on one external drive,
|
||
# where a dozen concurrent readers is slower than a few — hence the low
|
||
# default, overridable from the environment.
|
||
_BATCH_MAX_PROCS = int(os.environ.get("SRAS_BATCH_PROCS", 0)) or min(
|
||
4, os.cpu_count() or 2)
|
||
|
||
# Spawning a pool costs roughly a second of interpreter startup (each child
|
||
# re-imports the entry module). That is noise against a multi-GB scan but
|
||
# dominates a batch of small files, where it would make the job *slower* —
|
||
# so below this total size the batch just runs in the worker thread.
|
||
_BATCH_POOL_MIN_BYTES = int(os.environ.get("SRAS_BATCH_POOL_MIN_MB", 512)) * 1024 * 1024
|
||
|
||
|
||
class CancellableWorker(QObject):
|
||
"""A worker whose compute polls stop() between row chunks.
|
||
|
||
Without this a shutdown has to wait out whatever is in flight, and on a
|
||
large scan a single angle is ~40 s — far too long to block closing the
|
||
window. Chunk-level polling bounds the wait to one chunk instead.
|
||
"""
|
||
|
||
def __init__(self):
|
||
super().__init__()
|
||
self._stop = False
|
||
|
||
def stop(self):
|
||
self._stop = True
|
||
|
||
def _stopped(self) -> bool:
|
||
return self._stop
|
||
|
||
|
||
class _PooledWorker(CancellableWorker):
|
||
"""Fans a per-item computation across a thread pool, emitting each result
|
||
from this worker's own thread as it lands (never from a pool thread).
|
||
|
||
Subclasses provide _plan() -> n_workers (stashing whatever per-run
|
||
context they need), _items(), _one(item) -> result, and _emit(result).
|
||
On stop(): queued items are dropped, in-flight ones are not waited for —
|
||
that is what keeps closing the window responsive on a large scan.
|
||
"""
|
||
finished = pyqtSignal()
|
||
error = pyqtSignal(str)
|
||
|
||
def run(self):
|
||
try:
|
||
pool = ThreadPoolExecutor(max_workers=max(1, self._plan()))
|
||
try:
|
||
futures = [pool.submit(self._one, it) for it in self._items()]
|
||
for fut in as_completed(futures):
|
||
if self._stop:
|
||
break
|
||
self._emit(fut.result())
|
||
finally:
|
||
pool.shutdown(wait=not self._stop, cancel_futures=True)
|
||
self.finished.emit()
|
||
except Exception as exc:
|
||
self.error.emit(str(exc))
|
||
|
||
|
||
class LoadWorker(QObject):
|
||
finished = pyqtSignal(object) # SrasFile | None
|
||
error = pyqtSignal(str)
|
||
|
||
def __init__(self, path: str):
|
||
super().__init__()
|
||
self._path = path
|
||
|
||
def run(self):
|
||
try:
|
||
self.finished.emit(SrasFile(self._path))
|
||
except Exception as exc:
|
||
self.error.emit(str(exc))
|
||
self.finished.emit(None)
|
||
|
||
|
||
class ComputeWorker(CancellableWorker):
|
||
"""Computes one displayable image for (angle, channel).
|
||
|
||
For CH1/Velocity (FFT-derived) channels, the FFT is only run for pixels
|
||
whose DC4 (Bias B) mean is at or above dc_threshold_mv — masked pixels are
|
||
left at 0 MHz without ever being FFT'd, since that's the expensive part of
|
||
a scan. If the DC4 image for this angle is already known, pass it in as
|
||
*dc4_mv* to skip re-reading the CH4 channel from disk entirely.
|
||
|
||
Emits a plain ``np.ndarray`` already in display units.
|
||
"""
|
||
finished = pyqtSignal(object)
|
||
error = pyqtSignal(str)
|
||
|
||
def __init__(self, sras: SrasFile, angle_idx: int, ch_idx: int,
|
||
apply_bg_sub: bool = True, n_fft: int | None = None,
|
||
dc_threshold_mv: float = 0.0,
|
||
dc4_mv: np.ndarray | None = None,
|
||
is_fft_mode: bool = False):
|
||
super().__init__()
|
||
self._sras = sras
|
||
self._angle = angle_idx
|
||
self._ch = ch_idx
|
||
self._apply_bg_sub = apply_bg_sub
|
||
self._n_fft = n_fft
|
||
self._dc_threshold = dc_threshold_mv
|
||
self._dc4_mv = dc4_mv
|
||
self._is_fft_mode = is_fft_mode
|
||
|
||
def run(self):
|
||
try:
|
||
if self._is_fft_mode:
|
||
img = compute_rf_image(
|
||
self._sras, self._angle, dc_threshold_mv=self._dc_threshold,
|
||
apply_bg_sub=self._apply_bg_sub, n_fft=self._n_fft,
|
||
dc4_mv=self._dc4_mv, should_stop=self._stopped)
|
||
else:
|
||
img = dc_image_mv(self._sras, self._angle, self._ch,
|
||
should_stop=self._stopped)
|
||
# On cancellation the image is only partly filled, so hand back
|
||
# None rather than something that would be cached as real. The
|
||
# signal still fires either way — it is what quits the thread.
|
||
self.finished.emit(None if self._stop else img)
|
||
except Exception as exc:
|
||
self.error.emit(str(exc))
|
||
|
||
|
||
class DcPrecomputeWorker(_PooledWorker):
|
||
"""Computes CH3/CH4 DC images for every angle in the background.
|
||
|
||
DC images are cheap (a per-waveform mean, no FFT) compared to the
|
||
CH1/Velocity FFT, so precomputing them for the whole file right after load
|
||
makes switching angles instant while on a DC channel, and also means the
|
||
FFT masking step (which needs a DC4 image) rarely has to wait on anything.
|
||
"""
|
||
angle_done = pyqtSignal(int, np.ndarray, np.ndarray) # angle_idx, dc3_mv, dc4_mv
|
||
|
||
def __init__(self, sras: SrasFile):
|
||
super().__init__()
|
||
self._sras = sras
|
||
self._angle_budget = 0
|
||
|
||
def _plan(self) -> int:
|
||
n_workers, self._angle_budget = compute.plan_angle_level(self._sras)
|
||
return n_workers
|
||
|
||
def _items(self):
|
||
return range(self._sras.n_angles)
|
||
|
||
def _one(self, a: int) -> tuple[int, np.ndarray, np.ndarray]:
|
||
# max_workers=1 *and* a budget share: this call is one of several
|
||
# concurrent angles, and both the thread count and the buffer size
|
||
# have to be divided (see compute.plan_angle_level).
|
||
kw = dict(max_workers=1, budget=self._angle_budget,
|
||
should_stop=self._stopped)
|
||
return (a,
|
||
dc_image_mv(self._sras, a, CH3_IDX, **kw),
|
||
dc_image_mv(self._sras, a, CH4_IDX, **kw))
|
||
|
||
def _emit(self, result):
|
||
self.angle_done.emit(*result)
|
||
|
||
|
||
class BatchCacheWorker(QObject):
|
||
"""Batch-computes and stores DC or FFT images into each of *paths*'s v7
|
||
CACH tail, in place — converting v6 sources to v7 on first use, or updating
|
||
an existing v7 file's cache blocks without disturbing whatever the other
|
||
block already holds.
|
||
|
||
*mode* is ``"dc"`` (CH3/CH4 mean images), ``"fft"`` (CH1 peak-frequency
|
||
images, unmasked — masking is applied at display time, same as v5's PREC
|
||
convention), or ``"fft_rowavg"`` (same-row, distance-weighted CH1
|
||
averaging before the FFT — needs *dc_threshold_mv* and a positive
|
||
*row_avg_n*; see ``sras_compute.cache_file``).
|
||
|
||
Both FFT modes cache at *pad_factor*, which the caller sets from the
|
||
viewer's own padding — a cache stored at a pad the user is not viewing
|
||
at is one the display can never use.
|
||
|
||
Files are processed one per subprocess: they are fully independent, each
|
||
opens its own memmap and writes only its own bytes, and only path strings
|
||
and scalars cross the process boundary. Emits ``progress(int)`` (0–100 by
|
||
files completed), ``file_done(str, str)`` (path, error message or "") so
|
||
one file's failure doesn't abort the batch, and ``finished()``.
|
||
"""
|
||
progress = pyqtSignal(int)
|
||
file_done = pyqtSignal(str, str)
|
||
finished = pyqtSignal()
|
||
|
||
def __init__(self, paths: list[str], mode: str, apply_bg_sub: bool,
|
||
dc_threshold_mv: float | None = None, row_avg_n: int = 0,
|
||
pad_factor: int = 1):
|
||
super().__init__()
|
||
self._paths = paths
|
||
self._mode = mode
|
||
self._apply_bg_sub = apply_bg_sub
|
||
self._dc_threshold = dc_threshold_mv
|
||
self._row_avg_n = row_avg_n
|
||
self._pad_factor = pad_factor
|
||
|
||
def _report(self, path: str, err: str, done: int, total: int):
|
||
self.file_done.emit(path, err)
|
||
self.progress.emit(int(done / max(1, total) * 100))
|
||
|
||
def _run_pooled(self, paths: list[str], n_procs: int) -> list[str]:
|
||
"""Process the batch across *n_procs* subprocesses. Returns the paths
|
||
that never got a real answer because the pool itself died, so the
|
||
caller can retry them in-process.
|
||
|
||
Under spawn each child re-imports the entry module, so the batch must
|
||
survive that going wrong (an unguarded __main__, a frozen build, a
|
||
sandbox that forbids subprocesses) rather than reporting every file as
|
||
failed — hence the retry list instead of a per-file error.
|
||
"""
|
||
# Each child threads internally; divide the machine rather than
|
||
# letting every process claim every core.
|
||
per_proc_workers = max(1, (os.cpu_count() or 4) // n_procs)
|
||
unresolved: list[str] = []
|
||
done = 0
|
||
|
||
with ProcessPoolExecutor(max_workers=n_procs) as executor:
|
||
futures = {
|
||
executor.submit(cache_file, p, self._mode, self._apply_bg_sub,
|
||
compute.get_fft_backend(), per_proc_workers,
|
||
pad_factor=self._pad_factor,
|
||
dc_threshold_mv=self._dc_threshold,
|
||
row_avg_n=self._row_avg_n): p
|
||
for p in paths
|
||
}
|
||
for fut in as_completed(futures):
|
||
path = futures[fut]
|
||
try:
|
||
err = fut.result()
|
||
except BrokenProcessPool:
|
||
unresolved.append(path)
|
||
continue
|
||
except Exception as exc:
|
||
err = str(exc)
|
||
done += 1
|
||
self._report(path, err, done, len(paths))
|
||
|
||
return unresolved
|
||
|
||
def _run_inline(self, paths: list[str], done: int, total: int):
|
||
"""Fallback / single-file path: compute in this thread. Still uses the
|
||
full core count internally, since nothing else is competing."""
|
||
for path in paths:
|
||
try:
|
||
err = cache_file(path, self._mode, self._apply_bg_sub,
|
||
compute.get_fft_backend(),
|
||
compute.default_max_workers(),
|
||
pad_factor=self._pad_factor,
|
||
dc_threshold_mv=self._dc_threshold,
|
||
row_avg_n=self._row_avg_n)
|
||
except Exception as exc:
|
||
err = str(exc)
|
||
done += 1
|
||
self._report(path, err, done, total)
|
||
|
||
def _worth_pooling(self, paths: list[str]) -> bool:
|
||
if len(paths) < 2:
|
||
return False
|
||
total = 0
|
||
for p in paths:
|
||
try:
|
||
total += os.path.getsize(p)
|
||
except OSError:
|
||
pass # unreadable files are reported by cache_file
|
||
return total >= _BATCH_POOL_MIN_BYTES
|
||
|
||
def run(self):
|
||
paths = self._paths
|
||
n_procs = max(1, min(_BATCH_MAX_PROCS, len(paths)))
|
||
|
||
if not self._worth_pooling(paths):
|
||
self._run_inline(paths, 0, len(paths))
|
||
self.finished.emit()
|
||
return
|
||
|
||
try:
|
||
unresolved = self._run_pooled(paths, n_procs)
|
||
except Exception:
|
||
# The pool could not be created or collapsed wholesale.
|
||
unresolved = list(paths)
|
||
|
||
if unresolved:
|
||
self._run_inline(unresolved, len(paths) - len(unresolved), len(paths))
|
||
|
||
self.finished.emit()
|
||
|
||
|
||
class Ch4MaskWorker(_PooledWorker):
|
||
"""Fetches each requested angle's CH4 (Bias B) DC image in mV, for the
|
||
alignment wizard's initial threshold-mask stack.
|
||
|
||
Reuses dc_image_mv, which prefers a stored v5/v7 cache over recomputing
|
||
from raw waveforms, so this only does real work for a file that hasn't
|
||
gone through the v7 "Convert" batch step and for angles the main
|
||
window's own DcPrecomputeWorker (which runs automatically right after
|
||
every file load) hasn't reached yet. In the common case — the user opens
|
||
Fusion -> Manual Alignment after DC precompute has already finished —
|
||
*angle_indices* is empty and this worker is never even constructed (see
|
||
CorrelatePage._start_mask_prep).
|
||
"""
|
||
angle_done = pyqtSignal(int, np.ndarray) # angle_idx, dc4_mv
|
||
|
||
def __init__(self, sras: SrasFile, angle_indices: list[int]):
|
||
super().__init__()
|
||
self._sras = sras
|
||
self._angles = angle_indices
|
||
self._budget = 0
|
||
|
||
def _plan(self) -> int:
|
||
n_workers, self._budget = compute.plan_angle_level(self._sras)
|
||
return n_workers
|
||
|
||
def _items(self):
|
||
return self._angles
|
||
|
||
def _one(self, a: int) -> tuple[int, np.ndarray]:
|
||
return a, dc_image_mv(self._sras, a, CH4_IDX,
|
||
max_workers=1, budget=self._budget)
|
||
|
||
def _emit(self, result):
|
||
self.angle_done.emit(*result)
|
||
|
||
|
||
class CrossCorrelateWorker(_PooledWorker):
|
||
"""Rigid registration (rotation + translation, never scale) of each of
|
||
*angle_indices* against *ref_angle_idx*, for the alignment wizard's
|
||
Run/Re-run Correlation button.
|
||
|
||
Runs on a background thread — registering a real many-angle,
|
||
high-resolution scan takes long enough that doing it on the GUI thread
|
||
would visibly freeze the dialog. Rotation is *searched*, not taken from the
|
||
stage's reported angle: see compute.register_angle_to_reference, which
|
||
seeds from that angle but scores both of its signs and refines from there.
|
||
dc4_mv is the dialog's own already-in-memory per-angle CH4 image — this
|
||
worker does no fetching of its own.
|
||
"""
|
||
# angle_idx, rotation_deg, shift_x_mm, shift_y_mm, score, source
|
||
angle_done = pyqtSignal(int, float, float, float, float, str)
|
||
|
||
def __init__(self, sras: SrasFile, ref_angle_idx: int, angle_indices: list[int],
|
||
dc4_mv: dict[int, np.ndarray], *,
|
||
sources: tuple[str, ...], dc_threshold_mv: float,
|
||
search_deg: float, reg_kwargs: dict | None = None):
|
||
"""*reg_kwargs* is splatted into register_angle_to_reference on top of
|
||
the named arguments — the wizard's rotation-search controls (seed,
|
||
signs, refine, grid sizes) go through here, so exposing another knob
|
||
needs no change to this class."""
|
||
super().__init__()
|
||
self._sras = sras
|
||
self._ref = ref_angle_idx
|
||
self._angles = angle_indices
|
||
self._dc4_mv = dc4_mv
|
||
self._sources = sources
|
||
self._threshold = dc_threshold_mv
|
||
self._search_deg = search_deg
|
||
self._reg_kwargs = dict(reg_kwargs or {})
|
||
|
||
def _plan(self) -> int:
|
||
return compute.registration_workers(self._sras)
|
||
|
||
def _items(self):
|
||
return self._angles
|
||
|
||
def _one(self, a: int) -> tuple[int, compute.RigidFit]:
|
||
return a, compute.register_angle_to_reference(
|
||
self._sras, a, self._ref, self._dc4_mv,
|
||
dc_threshold_mv=self._threshold, sources=self._sources,
|
||
search_deg=self._search_deg, **self._reg_kwargs)
|
||
|
||
def _emit(self, result):
|
||
a, fit = result
|
||
self.angle_done.emit(a, fit.rotation_deg, fit.shift_mm[0],
|
||
fit.shift_mm[1], fit.score, fit.source)
|
||
|
||
|
||
class AlignedExportWorker(CancellableWorker):
|
||
"""Writes the aligned, cropped .sras on a background thread.
|
||
|
||
Unlike every other worker here this one produces a *file*, which changes
|
||
what cancellation has to mean: write_aligned_sras stages into a ".part"
|
||
sibling and removes it when should_stop() fires, so a cancelled or crashed
|
||
export leaves nothing behind. That matters more than it sounds — a
|
||
truncated .sras is not detectably broken, since the v6 parser reads a short
|
||
file as an aborted scan and opens it happily.
|
||
|
||
Cancellation is polled per output row chunk, the same granularity
|
||
CancellableWorker's docstring justifies, so closing the window never waits
|
||
on a multi-gigabyte write.
|
||
"""
|
||
progress = pyqtSignal(int) # 0-100
|
||
finished = pyqtSignal(str, str) # written path ("" = none), error
|
||
|
||
def __init__(self, sras: SrasFile, result, out_path: str):
|
||
super().__init__()
|
||
self._sras = sras
|
||
self._result = result
|
||
self._out_path = out_path
|
||
|
||
def run(self):
|
||
try:
|
||
written = write_aligned_sras(
|
||
self._sras, self._result, self._out_path,
|
||
progress_cb=self.progress.emit, should_stop=self._stopped)
|
||
if self._stopped():
|
||
self.finished.emit("", "") # cancelled: no file, no error
|
||
else:
|
||
self.finished.emit(str(written), "")
|
||
except Exception as exc:
|
||
self.finished.emit("", str(exc))
|