Files
sras-viewer/sras_workers.py
T
Thomas Ales [M S E] a8b962e317 Add min peak frequency floor (CACH v4) and Batch Export View as Images
Two strands of in-progress work, committed together because they overlap
in sras_workers.py and main_window.py.

Min peak frequency floor:
- CACH tail bumped to version 4, adding u32 min_freq_khz provenance in
  fixed-point kHz (a float32 20.1 reads back as 20.10000038 and would
  report a spurious mismatch forever). v1-v3 tails read as no floor.
- Stored FFT caches are accepted when the reader's floor is at or above
  the stored one, since a higher floor is re-applicable by masking.
- Floor plumbed through compute_rf_image, BatchCacheWorker and the viewer.

Batch Export View as Images:
- New sras_render.py holds draw_view_image, shared by the Qt canvas and
  the headless exporter so a PNG cannot drift from what the GUI shows.
  Deliberately Qt-free so it is importable in a pool subprocess.
- BatchExportImagesWorker renders the current view settings across many
  files, process-pooled with an inline fallback, reporting per-file
  output names so the caller can flag same-stem collisions.
- _axes_extent extracted into sras_format for both render paths.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-08-12 10:17:43 -05:00

541 lines
22 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/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
from sras_render import export_view_image
# 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,
min_freq_mhz: float = 0.0):
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
self._min_freq_mhz = min_freq_mhz
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,
min_freq_mhz=self._min_freq_mhz)
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. *min_freq_mhz* travels the same
way: the viewer's live min-peak-freq floor, recorded in the store as
provenance so a reader knows which bins the peak search considered.
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, min_freq_mhz: float = 0.0):
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
self._min_freq_mhz = min_freq_mhz
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,
per_proc_workers,
pad_factor=self._pad_factor,
dc_threshold_mv=self._dc_threshold,
row_avg_n=self._row_avg_n,
min_freq_mhz=self._min_freq_mhz): 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.default_max_workers(),
pad_factor=self._pad_factor,
dc_threshold_mv=self._dc_threshold,
row_avg_n=self._row_avg_n,
min_freq_mhz=self._min_freq_mhz)
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 BatchExportImagesWorker(QObject):
"""Renders the view settings captured by the caller at trigger time
(angle/channel/threshold/etc.) to one PNG per file in *paths*, via
sras_render.export_view_image.
Same process-pool-with-inline-fallback strategy as BatchCacheWorker
above (same _BATCH_MAX_PROCS / _BATCH_POOL_MIN_BYTES thresholds): an
FFT-derived (CH1/Velocity) view is exactly the same expensive per-file
compute Batch Compute FFT already parallelizes this way. Never mutates
*paths* — each file is opened read-only — so unlike BatchCacheWorker
there is no version gate.
Emits progress(int) (0-100 by files completed), file_done(str, str, str)
(path, error message or "", output filename this file targeted — set
even on most failures so the caller can flag same-stem collisions across
the batch without any cross-process bookkeeping), and finished().
"""
progress = pyqtSignal(int)
file_done = pyqtSignal(str, str, str)
finished = pyqtSignal()
def __init__(self, paths: list[str], **render_kwargs):
"""*render_kwargs* is exactly export_view_image's keyword-only
settings (out_dir, angle_idx, ch_idx, is_fft_mode, is_velocity,
dc_threshold_mv, apply_bg_sub, pad_factor, min_freq_mhz, grating_um,
cmap, auto_scale, vmin, vmax, highlight_masked, mode_str,
colorbar_label, mask_color) — bundled rather than repeated as
positional params across __init__/_run_pooled/_run_inline."""
super().__init__()
self._paths = paths
self._kw = render_kwargs
def _report(self, path: str, err: str, out_name: str, done: int, total: int):
self.file_done.emit(path, err, out_name)
self.progress.emit(int(done / max(1, total) * 100))
def _run_pooled(self, paths: list[str], n_procs: int) -> list[str]:
"""Same contract as BatchCacheWorker._run_pooled: returns the paths
that never got a real answer because the pool itself died, so the
caller can retry them in-process."""
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(export_view_image, p,
max_workers=per_proc_workers, **self._kw): p
for p in paths
}
for fut in as_completed(futures):
path = futures[fut]
try:
err, out_name = fut.result()
except BrokenProcessPool:
unresolved.append(path)
continue
except Exception as exc:
err, out_name = str(exc), ""
done += 1
self._report(path, err, out_name, done, len(paths))
return unresolved
def _run_inline(self, paths: list[str], done: int, total: int):
for path in paths:
try:
err, out_name = export_view_image(
path, max_workers=compute.default_max_workers(), **self._kw)
except Exception as exc:
err, out_name = str(exc), ""
done += 1
self._report(path, err, out_name, 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 export_view_image
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], *, reg_kwargs: dict | None = None):
"""*reg_kwargs* is splatted into register_angle_to_reference — every
registration setting the wizard exposes (sources, threshold, search
width, seed, signs, refine, grid sizes) travels in it, so this class
holds no opinion about which knobs exist and exposing another needs no
change here."""
super().__init__()
self._sras = sras
self._ref = ref_angle_idx
self._angles = angle_indices
self._dc4_mv = dc4_mv
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, **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))