Thomas Ales [M S E]
48d560fb00
Two strands, committed together because they touch the same four files and answer the same question: an exported PNG must be the image the user was looking at when they triggered the export. Dropped view settings (bug fix): - BatchExportWorker forwarded the DC threshold, bg-sub, pad and row- averaging to compute_rf_image but never min_freq_mhz, so every RF and Velocity map was exported with no min peak frequency floor. The pixels the floor exists to reject came back at their pre-floor peaks. - That also decided *which source* the image came from: requesting floor 0 makes a stored v7 cache look like a match (cache_mismatch_reasons only rejects a higher stored floor), so a batch-computed file exported the cached peak-frequency block verbatim - numbers from an earlier Batch Compute run, while the screen showed that same cache re-masked against the live floor. The symptom was a newly loaded file exporting a stale-looking velocity map, since the floor is a display control that survives a load while the window's own FFT cache does not. - export_view_image had the same class of defect with a different parameter: it dropped row_avg_n, so a file whose stored cache is row- averaged displayed from the store but exported as a full raw recompute - a different, and much slower, image. It now reads sras.precomputed_row_avg_n per file, mirroring _stored_fft_image and _start_compute, and derived per file for the same reason n_fft is. Export at the live canvas size: - Both exporters now take the canvas's current figure size instead of a hard-coded 7x5. draw_view_image uses aspect="auto", so the figure box is what sets the map's proportions - a view sized wide on screen was being squeezed into a different shape on disk. Read at trigger time, in inches, so a resize mid-batch cannot change images later in the same run and a HiDPI display exports like a standard one. - sanitize_figsize clamps a degenerate size (collapsed splitter pane, minimized window): a wrong-looking aspect ratio must never be the reason a batch loses an image. Tests: BatchExportWorker had no coverage at all. Four new tests cover the floor being applied, a stored cache never being served unfloored, the menu-to-worker wiring, and the row-averaged cache case; each was verified to fail against the unfixed code. 166 passed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
704 lines
29 KiB
Python
704 lines
29 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
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
from matplotlib.backends.backend_agg import FigureCanvasAgg
|
||
from matplotlib.figure import Figure
|
||
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 CH1_IDX, CH3_IDX, CH4_IDX, SrasFile
|
||
from sras_render import export_view_image, sanitize_figsize
|
||
|
||
# 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,
|
||
row_avg_n: int = 0,
|
||
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._row_avg_n = row_avg_n
|
||
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,
|
||
row_avg_n=self._row_avg_n,
|
||
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()
|
||
|
||
|
||
@dataclass
|
||
class ExportChannel:
|
||
"""One row of a batch image export: which raw channel to read, whether
|
||
to apply the Velocity post-multiply, and the fixed display range/labels
|
||
to render it with.
|
||
|
||
Kept free of sras_viewer's display constants (CH_LABELS, VELOCITY_MODE_IDX,
|
||
etc.) so BatchExportWorker has no dependency on the GUI module — the
|
||
caller resolves labels/units/filename tags once, up front.
|
||
"""
|
||
ch_idx: int # CH1_IDX, CH3_IDX, or CH4_IDX -- which raw data to read
|
||
is_velocity: bool # True only for the derived Velocity map (post-multiply of CH1 freq)
|
||
vmin: float
|
||
vmax: float
|
||
label: str # e.g. "CH3 -- Bias A (DC mean)"
|
||
unit: str # colorbar units, e.g. "mV"
|
||
tag: str # filename tag: "CH1", "CH3", "CH4", "VEL"
|
||
|
||
|
||
def _render_map_png(img: np.ndarray, extent: list[float], *, cmap: str,
|
||
vmin: float, vmax: float, title: str, colorbar_label: str,
|
||
out_path: str, figsize: tuple[float, float] | None = None):
|
||
"""Render one map image to *out_path* with a fixed vmin/vmax, using a
|
||
headless Agg canvas so this never touches the GUI thread's interactive
|
||
matplotlib backend. Layout mirrors ImageCanvas.show_image.
|
||
|
||
*figsize* is the live canvas's size in inches, so the PNG comes out at
|
||
the shape the view was being read at instead of a fixed 7x5 -- with
|
||
aspect="auto" below, the figure box is what sets the map's proportions.
|
||
None falls back to sras_render.DEFAULT_FIGSIZE."""
|
||
fig = Figure(figsize=sanitize_figsize(figsize), tight_layout=True)
|
||
FigureCanvasAgg(fig)
|
||
ax = fig.add_subplot(111)
|
||
im = ax.imshow(img, aspect="auto", origin="upper", extent=extent,
|
||
cmap=cmap, vmin=vmin, vmax=vmax, interpolation="nearest")
|
||
cb = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
|
||
if colorbar_label:
|
||
cb.set_label(colorbar_label)
|
||
ax.set_xlabel("X (mm)")
|
||
ax.set_ylabel("Y (mm)")
|
||
ax.set_title(title)
|
||
fig.savefig(out_path, dpi=150)
|
||
|
||
|
||
class BatchExportWorker(CancellableWorker):
|
||
"""Renders and saves one PNG per (angle, selected channel) for an
|
||
already-open SrasFile, with each channel's vmin/vmax held fixed across
|
||
every angle so the colorbar is directly comparable image to image.
|
||
|
||
Takes the SrasFile object directly (like ComputeWorker/DcPrecomputeWorker)
|
||
rather than a path -- this runs against the file already open in the GUI,
|
||
not an arbitrary batch of files, so there is no need to reopen it in a
|
||
subprocess the way BatchCacheWorker does.
|
||
|
||
Every setting that decides what an FFT-derived pixel *is* has to arrive
|
||
here explicitly, because compute_rf_image defaults each one to "off" and
|
||
an omitted argument is therefore not a no-op -- it silently exports a
|
||
different image than the one on screen. *min_freq_mhz* is the one with
|
||
teeth: dropping it does not just skip a mask, it also makes this file's
|
||
stored cache (written by Batch Compute FFT at some earlier floor) look
|
||
like a match, so the export hands back the pre-floor peaks the user
|
||
raised the floor to get rid of. It mirrors
|
||
SrasViewerWindow._stored_fft_image, which takes the floor live for the
|
||
same reason -- re-masking a stored image against a *higher* floor is
|
||
free, so the displayed map always honors the spin box.
|
||
|
||
Emits progress(int) (0-100 over all angle x channel pairs), file_done(str,
|
||
str) (output path, error message or ""), and finished() -- same shape as
|
||
BatchCacheWorker.
|
||
"""
|
||
progress = pyqtSignal(int)
|
||
file_done = pyqtSignal(str, str)
|
||
finished = pyqtSignal()
|
||
|
||
def __init__(self, sras: SrasFile, channels: list[ExportChannel],
|
||
output_dir: str, prefix: str, *, cmap: str,
|
||
apply_bg_sub: bool, dc_threshold_mv: float,
|
||
n_fft: int | None, grating_um: float,
|
||
min_freq_mhz: float = 0.0,
|
||
figsize: tuple[float, float] | None = None):
|
||
super().__init__()
|
||
self._sras = sras
|
||
self._channels = channels
|
||
self._output_dir = Path(output_dir)
|
||
self._prefix = prefix
|
||
self._cmap = cmap
|
||
self._apply_bg_sub = apply_bg_sub
|
||
self._dc_threshold_mv = dc_threshold_mv
|
||
self._n_fft = n_fft
|
||
self._grating_um = grating_um
|
||
self._min_freq_mhz = min_freq_mhz
|
||
self._figsize = figsize
|
||
|
||
def _angle_extent(self, angle_idx: int) -> list[float]:
|
||
s = self._sras
|
||
x_axis = s.x_axis_mm(angle_idx)
|
||
y_axis = s.y_positions_mm(angle_idx)
|
||
dx = x_axis[1] - x_axis[0] if len(x_axis) > 1 else s.pixel_x_mm
|
||
dy = float(y_axis[1] - y_axis[0]) if len(y_axis) > 1 else 1.0
|
||
return [x_axis[0] - dx / 2, x_axis[-1] + dx / 2,
|
||
y_axis[-1] + dy / 2, y_axis[0] - dy / 2]
|
||
|
||
def run(self):
|
||
try:
|
||
s = self._sras
|
||
n_angles = s.n_angles
|
||
total = n_angles * len(self._channels)
|
||
done = 0
|
||
needs_freq = any(c.ch_idx == CH1_IDX for c in self._channels)
|
||
|
||
for angle_idx in range(n_angles):
|
||
if self._stop:
|
||
break
|
||
extent = self._angle_extent(angle_idx)
|
||
angle_deg = s.angles_deg[angle_idx]
|
||
|
||
freq_mhz = None
|
||
if needs_freq:
|
||
freq_mhz = compute_rf_image(
|
||
s, angle_idx, dc_threshold_mv=self._dc_threshold_mv,
|
||
apply_bg_sub=self._apply_bg_sub, n_fft=self._n_fft,
|
||
should_stop=self._stopped, row_avg_n=s.precomputed_row_avg_n,
|
||
min_freq_mhz=self._min_freq_mhz)
|
||
|
||
for channel in self._channels:
|
||
if self._stop:
|
||
break
|
||
out_path = (self._output_dir
|
||
/ f"{self._prefix}_angle{angle_idx:02d}_{channel.tag}.png")
|
||
try:
|
||
if channel.ch_idx == CH1_IDX:
|
||
img = (freq_mhz * self._grating_um if channel.is_velocity
|
||
else freq_mhz)
|
||
else:
|
||
img = dc_image_mv(s, angle_idx, channel.ch_idx,
|
||
should_stop=self._stopped)
|
||
title = f"{channel.tag} | {angle_deg:.1f}°"
|
||
colorbar_label = (f"{channel.label} ({channel.unit})"
|
||
if channel.unit else channel.label)
|
||
_render_map_png(
|
||
img, extent, cmap=self._cmap,
|
||
vmin=channel.vmin, vmax=channel.vmax,
|
||
title=title, colorbar_label=colorbar_label,
|
||
out_path=str(out_path), figsize=self._figsize)
|
||
self.file_done.emit(str(out_path), "")
|
||
except Exception as exc:
|
||
self.file_done.emit(str(out_path), str(exc))
|
||
done += 1
|
||
self.progress.emit(int(done / max(1, total) * 100))
|
||
|
||
self.finished.emit()
|
||
except Exception as exc:
|
||
self.file_done.emit("", str(exc))
|
||
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))
|