From 499cb32fa936a19a1affc070991407e22c50e240 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Tue, 28 Jul 2026 11:57:09 -0500 Subject: [PATCH 01/14] feat(csb): read Direct Electron CSB event streams, and re-cut them To Frames MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A CSB file is not a frame stack — it is a compressed sparse block stream of detected electron events at the camera's raw cadence (390 us on the test movie), and a single raw frame of it is mostly empty pixels. An image only exists once you choose an exposure and integrate, so "opening" one means choosing one, and the reader's job is to make that choice cheap to change and cheap to scrub. Opening gives a stack of integrated time planes, ONE PLANE PER DASK BLOCK, so scrubbing integrates only the window it is parked on — the property the navigator rests on, and the one the sibling rosettasciio/mrc patch documents a 40x scrub regression from losing. "To Frames" then re-cuts the same stream at whatever frame rate you ask for, as a new lazy node in the signal tree; from there it is an ordinary insitu movie and Play, virtual imaging, the movie editor and the rest work on it with no idea where it came from. Written as a DROP-IN rsciio plugin (spyde/external/rsciio_csb/: __init__ re-exporting file_reader, _api.py, specifications.yaml, importing nothing from SpyDE) so moving it into RosettaSciIO is a directory copy plus deleting the runtime registration. _core/_sparse are vendored verbatim from de-csb. Three things that had to be got right, all measured on the real 8192x8192, 400-frame, 125.8M-event movie: * `dask.delayed(pure=True)` tokenizes its arguments, and passing the dataset meant dask hashed the 252 MB memory-mapped payload once per plane: 70 SECONDS to build a graph that reads nothing. The graph now carries the file PATH and resolves a cached dataset per worker — which is also what makes it picklable and gives the accumulator a single owner. Open is now 35 ms warm. * The navigator is FREE from the block table (counts per frame, zero payload reads, 1.7 ms) and exactly equals what integrating yields. It travels as original_metadata — the signal dict's `attributes` key does NOT reach the signal object, so a navigator sent that way silently vanishes and SpyDE integrates the whole movie to draw a thumbnail anyway. * The accumulator is not concurrency-safe (reset+add_frames mutate one shared buffer), so integrations take a lock. Per-thread accumulators would cost 268 MB each at 8192" to win back parallelism the GPU serialises regardless. Also: `requires_original_metadata`, a general format gate mirroring requires_vectors, because "re-cut the exposure" is a question you can only ask of an event stream and signal_type cannot express that (a CSB movie is an ordinary insitu signal and must stay one). And _load_file_thread now ensures spyde.external's patches have landed before hs.load — it did not, so a load that beat the startup prewarm saw an unpatched rosettasciio: no CSB reader, and silently the slow MRC path. --- spyde/actions/csb_to_frames.py | 162 ++++ spyde/backend/_session_files.py | 48 +- .../drawing/toolbars/plot_control_toolbar.py | 28 + spyde/external/rosettasciio/__init__.py | 9 +- spyde/external/rosettasciio/csb_format.py | 84 ++ spyde/external/rsciio_csb/__init__.py | 25 + spyde/external/rsciio_csb/_api.py | 249 +++++ spyde/external/rsciio_csb/_core.py | 888 ++++++++++++++++++ spyde/external/rsciio_csb/_sparse.py | 435 +++++++++ spyde/external/rsciio_csb/specifications.yaml | 10 + spyde/toolbars.yaml | 28 + 11 files changed, 1963 insertions(+), 3 deletions(-) create mode 100644 spyde/actions/csb_to_frames.py create mode 100644 spyde/external/rosettasciio/csb_format.py create mode 100644 spyde/external/rsciio_csb/__init__.py create mode 100644 spyde/external/rsciio_csb/_api.py create mode 100644 spyde/external/rsciio_csb/_core.py create mode 100644 spyde/external/rsciio_csb/_sparse.py create mode 100644 spyde/external/rsciio_csb/specifications.yaml diff --git a/spyde/actions/csb_to_frames.py b/spyde/actions/csb_to_frames.py new file mode 100644 index 00000000..49aecfdc --- /dev/null +++ b/spyde/actions/csb_to_frames.py @@ -0,0 +1,162 @@ +""" +csb_to_frames.py — "To Frames": an event stream → an ordinary in-situ movie. + +A CSB file has no frames in the usable sense. It is a stream of detected +electron events at the camera's raw cadence (390 µs on the test movie), and a +single raw frame of it is mostly empty pixels. An image only exists once you +pick an exposure and integrate, which is why the file opens as a stack of +integrated time planes in the first place. + +This action re-cuts that choice. Given an exposure — or the frame rate you want +in fps, which is how it is usually thought about — it rebuilds the movie at +that cadence as a **new node in the same signal tree**, lazily, one plane per +dask block. From there it is an ordinary lazy in-situ signal: the navigator +scrub, Play/Fast-Forward, virtual imaging, the movie editor and everything else +work on it with no knowledge that it came from an event stream. + +Lazy is not an optimisation here, it is the only option: the test movie is +8192², so one plane is 268 MB and a 50-plane stack is 13 GB. Nothing is +integrated until a plane is actually looked at. +""" +from __future__ import annotations + +import logging + +import numpy as np + +from spyde.backend.ipc import emit_error, emit_status + +log = logging.getLogger(__name__) + +DEFAULTS = dict(fps=0.0, exposure_ms=0.0, frames_per_plane=0, bin=1) + + +def _source_path(signal) -> str | None: + """The .csb this signal was read from, or None if it was not.""" + try: + om = signal.original_metadata + if "csb" in om: + return str(om.csb.path) + except Exception as e: + log.debug("reading the CSB source path failed: %s", e) + return None + + +def is_csb(signal) -> bool: + """True when a signal came from a CSB event stream — the gate for this + action, since re-cutting the exposure only means anything for one.""" + return _source_path(signal) is not None + + +def _resolve_exposure(ds, params) -> float: + """Seconds per output plane, from whichever knob the caller used. + + Three ways of saying the same thing, because all three get asked for: a + frame RATE (fps) is what an in-situ experiment is described by, an EXPOSURE + (ms) is what the camera is set to, and a FRAME COUNT is the atomic unit the + stream actually quantises to. + """ + dt = ds.frame_duration + if dt <= 0: + raise ValueError("this CSB file declares no frame cadence, so it has " + "no time axis to re-cut") + n = int(params.get("frames_per_plane") or 0) + if n > 0: + return max(1, n) * dt + ms = float(params.get("exposure_ms") or 0.0) + if ms > 0: + return max(ms * 1e-3, dt) + fps = float(params.get("fps") or 0.0) + if fps > 0: + return max(1.0 / fps, dt) + raise ValueError("choose a frame rate, an exposure or a number of frames " + "per plane") + + +def csb_to_frames(ctx, action_name: str = "To Frames", **params): + """Toolbar entry point (ActionContext convention).""" + plot, session = ctx.plot, ctx.session + tree = getattr(plot, "signal_tree", None) + if tree is None or session is None: + emit_error("To Frames: no active dataset") + return None + src = tree.root + path = _source_path(src) + if path is None: + emit_error("To Frames only applies to a CSB event stream") + return None + + def _work(): + try: + _rebuild(session, tree, src, path, params) + except Exception as e: + emit_error(f"To Frames failed: {e}") + log.exception("To Frames failed") + + from spyde.actions.lifecycle import run_on_worker + run_on_worker(session, _work, name="csb-to-frames") + return None + + +def _rebuild(session, tree, src, path: str, params) -> None: + """Re-cut the stream at a new exposure and add it to the tree.""" + import hyperspy.api as hs + from spyde.external.rsciio_csb._api import ( + _dataset, _axes, lazy_stack, plane_counts, + ) + + backend = str(params.get("backend") or "auto") + bin_factor = max(1, int(params.get("bin") or 1)) + ds = _dataset(path, backend) + exposure = _resolve_exposure(ds, params) + bounds = ds._time_bounds(exposure, 0.0, ds.duration) + if not bounds: + raise ValueError( + f"a {exposure * 1e3:.4g} ms exposure selects no frames from this " + f"{ds.duration * 1e3:.4g} ms movie") + + emit_status(f"To Frames: {len(bounds)} planes at " + f"{exposure * 1e3:.4g} ms ({1.0 / exposure:.4g} fps)…") + + data = lazy_stack(path, bounds, ds.shape, bin_factor, np.float32, backend) + sig = hs.signals.Signal2D(data).as_lazy() + for axis, spec in zip( + list(sig.axes_manager.navigation_axes) + + list(sig.axes_manager.signal_axes), + _axes(ds, bounds, bin_factor)): + axis.name, axis.units = spec["name"], spec["units"] + axis.scale, axis.offset = spec["scale"], spec["offset"] + # `insitu` is what gates Play / Fast-Forward in toolbars.yaml; the nav axis + # is already named "time" in seconds, which is what _is_movie_time_axis + # looks for, but set it explicitly rather than relying on the sniff. + try: + sig.set_signal_type("insitu") + except Exception as e: + log.debug("set_signal_type(insitu) on the re-cut movie failed: %s", e) + + base = src.metadata.get_item("General.title", "CSB") + sig.metadata.General.title = f"{base} @ {exposure * 1e3:.4g} ms" + sig.original_metadata.add_dictionary({"csb": { + **{k: v for k, v in ds.csb.info().items()}, + "exposure_s": exposure, + "frames_per_plane": [int(f1 - f0) for f0, f1 in bounds], + "plane_frame_bounds": [[int(f0), int(f1)] for f0, f1 in bounds], + "bin": bin_factor, + "backend": backend, + }}) + + # The free navigator again — re-cutting the exposure changes the planes, so + # the overview has to be recomputed, but it still costs no payload reads. + nav = hs.signals.Signal1D(np.asarray(plane_counts(ds, bounds), np.float32)) + nav.metadata.General.title = "Total counts" + + def _add(): + session._add_signal(sig, source_path=path, navigator_override=nav) + emit_status(f"To Frames: {len(bounds)} planes at " + f"{exposure * 1e3:.4g} ms ({1.0 / exposure:.4g} fps)") + + dispatch = getattr(session, "_dispatch_to_main", None) + if dispatch is not None: + dispatch(_add) + else: + _add() diff --git a/spyde/backend/_session_files.py b/spyde/backend/_session_files.py index 57e817eb..682572f9 100644 --- a/spyde/backend/_session_files.py +++ b/spyde/backend/_session_files.py @@ -25,7 +25,7 @@ log = logging.getLogger(__name__) -SUPPORTED_EXTS = (".hspy", ".zspy", ".mrc", ".tif", ".tiff", ".de5") +SUPPORTED_EXTS = (".hspy", ".zspy", ".mrc", ".tif", ".tiff", ".de5", ".csb") # Extensions whose "file" is actually a DIRECTORY (a Zarr store). `.zspy` (and a # bare `.zarr`) is a nested-group folder, not a single file — so `os.path.isfile` @@ -67,6 +67,42 @@ def _dataset_size_bytes(path: str) -> float: return 0.0 +def _reader_navigator(sig): + """A navigator the READER already knew, or None. + + Normally SpyDE builds the scrub overview by reducing over every frame, + which for most formats is the cheapest thing available. For an event + stream it is the most expensive: a CSB navigator built that way would + integrate the entire movie just to draw a thumbnail, when the per-plane + totals are sitting in the block table for free (``rsciio_csb`` computes + them at open, costing no payload reads at all). + + Readers publish it as ``original_metadata.csb.plane_counts``. (Not the + signal dict's ``attributes`` key — HyperSpy does not propagate that onto + the signal object, so a navigator sent that way silently vanishes and the + expensive reduction happens anyway.) + """ + try: + om = getattr(sig, "original_metadata", None) + if om is None or not om.has_item("csb.plane_counts"): + return None + counts = om.get_item("csb.plane_counts") + except Exception as e: + log.debug("reading the reader-supplied navigator failed: %s", e) + return None + if counts is None or not len(counts): + return None + try: + import hyperspy.api as hs + import numpy as _np + nav = hs.signals.Signal1D(_np.asarray(counts, _np.float32)) + nav.metadata.General.title = "Total counts" + return nav + except Exception as e: + log.debug("building the reader-supplied navigator failed: %s", e) + return None + + _DEFAULT_EXAMPLE_NAMES = ( "mgo_nanocrystals", "small_ptychography", @@ -230,6 +266,13 @@ def _load_file_thread(self, path: str) -> None: # Wait for the Dask cluster (see _load_example_thread) — a file opened # during startup queues here rather than racing ahead with no client. self._await_dask() + # And make sure spyde.external's patches have landed BEFORE the + # read. They register formats and reader fast paths, so a load that + # beats the startup prewarm would otherwise see an unpatched + # rosettasciio — silently no CSB reader, silently the slow MRC path. + # Single-flighted and idempotent, so this costs nothing once warm. + from spyde.backend.heavy_imports import ensure_heavy_imports + ensure_heavy_imports() signal = hs.load(path, lazy=True) if not isinstance(signal, list): signal = [signal] @@ -273,7 +316,8 @@ def _load_file_thread(self, path: str) -> None: return for sig in signal: self._maybe_set_insitu_signal_type(sig) - self._add_signal(sig, source_path=path) + self._add_signal(sig, source_path=path, + navigator_override=_reader_navigator(sig)) self._add_recent(path) ipc.emit({"type": "recent_files", "paths": self._recent_files[:20]}) except Exception as e: diff --git a/spyde/drawing/toolbars/plot_control_toolbar.py b/spyde/drawing/toolbars/plot_control_toolbar.py index 7d6b5ffe..c244c210 100644 --- a/spyde/drawing/toolbars/plot_control_toolbar.py +++ b/spyde/drawing/toolbars/plot_control_toolbar.py @@ -72,6 +72,29 @@ def install_hint(name: str) -> str: return f'pip install "spyde[{extra}]"' if extra else f"pip install {name}" +def _has_original_metadata(signal, dotted: str | None) -> bool: + """``requires_original_metadata:`` gate — a dotted path that must exist in + the signal's ``original_metadata``. + + The FORMAT-specific gate, mirroring ``requires_vectors`` (which gates on a + result being attached). Some actions only mean anything for one file + format: "To Frames" re-cuts an exposure, which is a question you can only + ask of a CSB event stream, not of the dense movie it produces. Gating on + signal_type cannot express that — a CSB movie is an ordinary ``insitu`` + signal, and should stay one, so every in-situ feature keeps working on it. + """ + if not dotted: + return True + om = getattr(signal, "original_metadata", None) + if om is None: + return False + try: + return bool(om.has_item(dotted)) + except Exception as e: + log.debug("requires_original_metadata check for %r failed: %s", dotted, e) + return False + + def _packages_present(meta: dict) -> bool: """``requires_package:`` gate — a str or list of importable module names. @@ -194,6 +217,10 @@ def get_toolbar_actions_for_plot( or isinstance(signal, _resolve_signal_class(signal_class)) ) and (not requires_vectors or has_vectors) + # requires_original_metadata: hide unless the signal came from the + # format this action is about (see _has_original_metadata). + and _has_original_metadata( + signal, meta.get("requires_original_metadata")) # requires_package: hide until the optional extra is installed # (exspy / kikuchipy / atomap). Checked by find_spec, so this # costs nothing and never imports the package. @@ -276,6 +303,7 @@ def _action_matches_plot(action: str, meta: dict, plot_state: "PlotState") -> bo or isinstance(signal, _resolve_signal_class(signal_class)) ) and (not requires_vectors or has_vectors) + and _has_original_metadata(signal, meta.get("requires_original_metadata")) and _packages_present(meta) and (plot_state.dimensions in plot_dim) and ( diff --git a/spyde/external/rosettasciio/__init__.py b/spyde/external/rosettasciio/__init__.py index 3189e45a..89a7542a 100644 --- a/spyde/external/rosettasciio/__init__.py +++ b/spyde/external/rosettasciio/__init__.py @@ -16,14 +16,21 @@ * :mod:`~spyde.external.rosettasciio.tiff` — home for the per-page lazy TIFF chunking patch (rsciio's lazy TIFF returns ONE monolithic chunk). Currently a documented no-op in this branch — see the module docstring. +* :mod:`~spyde.external.rosettasciio.csb_format` — ADDS a format rather than + patching one: registers the Direct Electron ``.csb`` centroid-stream reader + (:mod:`spyde.external.rsciio_csb`, written as a drop-in rsciio plugin) so + ``hs.load`` dispatches ``.csb`` to it. Deleted when that plugin lands in + rosettasciio proper. """ from __future__ import annotations from spyde.external import register +from spyde.external.rosettasciio.csb_format import apply as _apply_csb from spyde.external.rosettasciio.mrc import apply as _apply_mrc from spyde.external.rosettasciio.tiff import apply as _apply_tiff register("rosettasciio", _apply_mrc) register("rosettasciio", _apply_tiff) +register("rosettasciio", _apply_csb) -__all__ = ["_apply_mrc", "_apply_tiff"] +__all__ = ["_apply_mrc", "_apply_tiff", "_apply_csb"] diff --git a/spyde/external/rosettasciio/csb_format.py b/spyde/external/rosettasciio/csb_format.py new file mode 100644 index 00000000..7cec5c16 --- /dev/null +++ b/spyde/external/rosettasciio/csb_format.py @@ -0,0 +1,84 @@ +""" +Register the CSB reader with RosettaSciIO at runtime. + +WHAT + Appends a plugin spec for Direct Electron ``.csb`` centroid-streaming files + to ``rsciio.IO_PLUGINS``, pointing at :mod:`spyde.external.rsciio_csb`, so + ``hs.load("movie.csb", lazy=True)`` resolves to our ``file_reader`` through + the ordinary extension dispatch. Nothing in SpyDE's load path changes. + +WHY + This is a stretch of what ``spyde.external`` is for — the rest of this + package PATCHES upstream, and this ADDS a format. It is here because it is + the only place SpyDE can teach an installed rosettasciio about a format it + does not ship, and the alternative (special-casing ``.csb`` in + ``_session_files``) would put format dispatch somewhere it does not belong. + + The reader itself is deliberately NOT written as SpyDE code: it lives in + :mod:`spyde.external.rsciio_csb`, laid out exactly as an rsciio plugin + (``__init__`` re-exporting ``file_reader``, ``_api.py``, + ``specifications.yaml``) and importing nothing from SpyDE. + +WHEN TO REMOVE + When the plugin lands in RosettaSciIO (or the cssfrancis fork) as + ``rsciio/csb/``. That is a directory copy of ``rsciio_csb`` — rsciio + auto-discovers ``specifications.yaml``, so registration becomes unnecessary + and this module and its entry in ``__init__`` are simply deleted. +""" +from __future__ import annotations + +import logging + +log = logging.getLogger(__name__) + +_applied = False + +#: Mirrors rsciio_csb/specifications.yaml. Duplicated rather than read from the +#: yaml because rsciio builds IO_PLUGINS from those files at import time and we +#: are appending after the fact; the yaml stays authoritative for the day this +#: moves upstream. `test_csb_reader.py` asserts the two agree. +_SPEC = { + "name": "CSB", + "name_aliases": ["DE_CSB"], + "description": ("Direct Electron CSB (compressed sparse block) " + "centroid-streaming files: a sparse event stream rather " + "than a dense frame stack. Reading integrates events over " + "a time window per plane."), + "full_support": False, + "file_extensions": ["csb", "CSB"], + "default_extension": 0, + "writes": False, + "non_uniform_axis": False, + "api": "spyde.external.rsciio_csb", +} + + +def apply() -> bool: + """Append the CSB spec to ``rsciio.IO_PLUGINS`` (idempotent).""" + global _applied + if _applied: + return True + try: + import rsciio + except Exception as e: # pragma: no cover + log.warning("spyde.external.rosettasciio.csb_format: rsciio import " + "failed (%s) — .csb files will not open", e) + return False + + plugins = getattr(rsciio, "IO_PLUGINS", None) + if not isinstance(plugins, list): + # Upstream changed shape. Never raise from a startup patch. + log.warning("spyde.external.rosettasciio.csb_format: rsciio.IO_PLUGINS " + "is %s, not a list — cannot register the CSB reader; .csb " + "files will not open", type(plugins).__name__) + return False + + if any(isinstance(p, dict) and p.get("name") == "CSB" for p in plugins): + _applied = True + return True + + plugins.append(dict(_SPEC)) + _applied = True + log.debug("spyde.external.rosettasciio.csb_format: registered the CSB " + "reader for %s", _SPEC["file_extensions"]) + return True diff --git a/spyde/external/rsciio_csb/__init__.py b/spyde/external/rsciio_csb/__init__.py new file mode 100644 index 00000000..f10db5b5 --- /dev/null +++ b/spyde/external/rsciio_csb/__init__.py @@ -0,0 +1,25 @@ +"""CSB — Direct Electron compressed-sparse-block centroid streams. + +**This package is shaped as a RosettaSciIO plugin, not as SpyDE code.** Its +layout is rsciio's exactly — ``__init__.py`` re-exporting ``file_reader``, +the implementation in ``_api.py``, and a ``specifications.yaml`` beside them — +so moving it upstream is a directory copy into ``rsciio/csb/`` plus deleting +SpyDE's runtime registration (``spyde.external.rosettasciio.csb_format``); +rsciio auto-discovers the yaml and needs no registration at all. + +Nothing in here imports SpyDE. Keep it that way — the moment it does, the +directory copy stops working. + +``_core`` and ``_sparse`` are vendored verbatim from the ``de-csb`` project +(``csb.py`` / ``csb_sparse.py``); the only edit is the import between them +being made relative. +""" +from ._api import file_reader + +__all__ = [ + "file_reader", +] + + +def __dir__(): + return sorted(__all__) diff --git a/spyde/external/rsciio_csb/_api.py b/spyde/external/rsciio_csb/_api.py new file mode 100644 index 00000000..acb7fdda --- /dev/null +++ b/spyde/external/rsciio_csb/_api.py @@ -0,0 +1,249 @@ +"""file_reader for Direct Electron CSB centroid-streaming files. + +A CSB file is **not a frame stack** — it is a compressed sparse block stream of +detected electron events with a frame cadence. There is no dense array to read; +an image only exists once you choose a time window and integrate the events in +it. So "opening" a CSB movie means choosing an exposure, and the reader's job +is to make that choice cheap to change and cheap to scrub. + +The lazy path is the point +-------------------------- +``lazy=True`` returns a dask array with **one time plane per block**, each block +integrating only its own window. That shape is load-bearing, not incidental: +SpyDE's navigator scrub rests on ``data[i].compute()`` touching only plane +``i``, and the sibling ``rosettasciio/mrc`` patch documents a 40x scrub +regression from a graph that read a whole 134 MB chunk to yield one frame. A +plane here is an integration over its own frames and nothing else. + +Two more consequences of the format: + +* **The navigator is free.** A per-plane total-counts overview comes from the + block table alone — ``counts`` summed per frame — reading zero payload bytes. + Handing it over matters: without it a viewer builds the overview by reducing + every plane, i.e. by integrating the entire movie, which is the one thing + this reader exists to avoid. +* **The accumulator is not concurrency-safe.** ``reset()`` + ``add_frames()`` + mutate one shared device/host buffer, so concurrent planes would corrupt each + other. Integrations take a lock rather than getting an accumulator each: at + 8192² one accumulator's buffer is 268 MB, so per-thread copies would cost + gigabytes to win back parallelism the GPU serialises anyway. +""" +from __future__ import annotations + +import logging +import os +import threading + +import numpy as np + +from ._core import CSBFile +from ._sparse import SparseCSB + +_logger = logging.getLogger(__name__) + +#: Planes to aim for when the caller names no exposure. Enough that the time +#: slider has somewhere to go, few enough that each plane has real signal in it +#: — a single frame of a sparse stream is mostly empty pixels. +DEFAULT_PLANES = 50 + +# One integration at a time (see the module docstring). +_ACC_LOCK = threading.Lock() + +# Open datasets, keyed by (path, backend). The graph carries this KEY rather +# than the SparseCSB itself, for two reasons: +# +# * `dask.delayed(pure=True)` tokenizes its arguments, and a SparseCSB holds +# the memory-mapped payload — so passing it made dask hash the whole file +# once per plane. On the 252 MB test movie that was 70 SECONDS to build a +# graph that reads nothing. +# * a key is picklable and a memmap-backed reader is not, so the same graph +# can go to a distributed worker, which resolves its own dataset here. +# +# One dataset per key also means one accumulator per key (SparseCSB caches it), +# which is what keeps device memory from being reallocated on every scrub. +_DATASETS: dict[tuple, SparseCSB] = {} +_DATASETS_LOCK = threading.Lock() + + +def _dataset(path: str, backend: str) -> SparseCSB: + """The process-wide dataset for one file, opened once.""" + key = (os.path.abspath(path), backend) + with _DATASETS_LOCK: + ds = _DATASETS.get(key) + if ds is None: + ds = SparseCSB(CSBFile(path), backend=backend) + _DATASETS[key] = ds + return ds + + +def _resolve_step(ds: SparseCSB, step, frames_per_plane) -> float: + """The exposure per output plane, in seconds. + + ``frames_per_plane`` is offered because a frame is the atomic exposure and + "8 frames" is a more natural thing to ask for than "3.12 ms". Whichever is + given, the result is floored at one frame — a step finer than the cadence + cannot produce more planes, only empty ones. + """ + dt = ds.frame_duration + if dt <= 0: + raise ValueError( + "this CSB file declares no frame cadence (microsec_per_frame is 0), " + "so it has no time axis to integrate over") + if frames_per_plane: + return max(1, int(frames_per_plane)) * dt + if step: + return max(float(step), dt) + return max(ds.duration / DEFAULT_PLANES, dt) + + +def plane_counts(ds: SparseCSB, bounds) -> np.ndarray: + """Total events per plane, straight from the block table. + + The navigator, for free — no payload byte is read. ``counts`` is one entry + per block and blocks tile each frame, so summing within a frame and then + across a plane's frames gives its total intensity. + """ + bpf = int(ds.csb.blocks_per_frame) + per_frame = np.asarray(ds.csb.counts, np.int64).reshape(-1, bpf).sum(1) + return np.array([per_frame[f0:f1].sum() for f0, f1 in bounds], np.float64) + + +def integrate_plane(path: str, backend: str, f0: int, f1: int, + bin_factor: int, dtype) -> np.ndarray: + """One plane: integrate frames ``[f0, f1)`` -> ``(1, H, W)``. + + Takes the file PATH, not an open dataset — see :data:`_DATASETS`. Every + argument is small and hashable, so the graph is cheap to build and can be + shipped to a worker. + """ + ds = _dataset(path, backend) + with _ACC_LOCK: + img = ds._sum_frames(f0, f1, bin_factor, dtype) + return np.asarray(img, dtype)[None] + + +def lazy_stack(path: str, bounds, shape, bin_factor: int = 1, + dtype=np.float32, backend: str = "auto"): + """A dask array of integrated planes, one plane per block. + + Each block is an independent ``[f0, f1)`` integration, so computing plane + ``i`` reads that window's events and nothing else — the property the + navigator scrub depends on. + """ + import dask + import dask.array as da + + h, w = shape[0] // bin_factor, shape[1] // bin_factor + delayed = dask.delayed(integrate_plane, pure=True) + blocks = [ + da.from_delayed(delayed(path, backend, int(f0), int(f1), bin_factor, dtype), + shape=(1, h, w), dtype=dtype) + for f0, f1 in bounds + ] + return da.concatenate(blocks, axis=0) + + +def _axes(ds: SparseCSB, bounds, bin_factor: int): + """Time / y / x axes, calibrated. The time axis is named so a viewer can + recognise this as a movie rather than a scan.""" + dt = ds.frame_duration + times = np.array([f0 * dt for f0, _ in bounds], float) + # Plane spacing, not frame spacing: consecutive planes are `step` apart. + scale = float(times[1] - times[0]) if len(times) > 1 else float(dt) + ang = float(getattr(ds.csb, "ang_per_pix", 0.0) or 0.0) * bin_factor + px = {"scale": ang, "units": "Å"} if ang > 0 else {"scale": 1.0, "units": "px"} + h = ds.csb.frame_height // bin_factor + w = ds.csb.frame_width // bin_factor + return [ + {"name": "time", "size": len(bounds), "navigate": True, + "offset": float(times[0]), "scale": scale, "units": "s"}, + {"name": "y", "size": h, "navigate": False, "offset": 0.0, **px}, + {"name": "x", "size": w, "navigate": False, "offset": 0.0, **px}, + ] + + +def file_reader(filename, lazy=False, step=None, frames_per_plane=None, + start=0.0, stop=None, bin=1, backend="auto", **kwds): + """Read a CSB centroid stream as a stack of integrated time planes. + + Parameters + ---------- + lazy : bool + With ``True`` (strongly preferred) each plane integrates on demand, so + opening the file costs only the header and block table. Eager reads + integrate every plane up front — at 8192² that is 268 MB per plane. + step : float, optional + Exposure per plane in seconds. Defaults to the whole movie split into + :data:`DEFAULT_PLANES`. + frames_per_plane : int, optional + The same thing in frames, which is usually how it is thought about. + Takes precedence over *step*. + start, stop : float + Time range to cover, in seconds. Defaults to the whole movie. + bin : int + Integer spatial binning, summed. + backend : {"auto", "gpu", "cpu", "cpu-numba", "cpu-numpy"} + Integration backend; ``"auto"`` takes the GPU when there is one. + """ + if kwds: + _logger.debug("CSB reader ignoring unknown kwargs: %s", sorted(kwds)) + path = str(filename) + ds = _dataset(path, backend) + bin_factor = max(1, int(bin)) + exposure = _resolve_step(ds, step, frames_per_plane) + bounds = ds._time_bounds(exposure, float(start), + ds.duration if stop is None else float(stop)) + if not bounds: + raise ValueError( + f"no planes: a {exposure:g} s exposure over " + f"[{start}, {stop if stop is not None else ds.duration}) s of this " + f"{ds.duration:g} s movie selects no frames") + + dtype = np.float32 + if lazy: + data = lazy_stack(path, bounds, ds.shape, bin_factor, dtype, backend) + else: + data = np.concatenate([integrate_plane(path, backend, f0, f1, + bin_factor, dtype) + for f0, f1 in bounds]) + + info = ds.csb.info() + return [{ + "data": data, + "axes": _axes(ds, bounds, bin_factor), + "metadata": { + "General": { + "title": os.path.splitext(os.path.basename(str(filename)))[0], + "original_filename": os.path.basename(str(filename)), + }, + "Signal": {"signal_type": ""}, + "Acquisition_instrument": { + "TEM": { + "magnification": info.get("mag"), + "defocus": info.get("defocus_um"), + "Detector": {"camera_serial": info.get("camera_sn")}, + }, + }, + }, + "original_metadata": { + "csb": { + **info, + # What this particular read chose — without it a plane's + # intensity is uninterpretable (it is events per exposure, and + # the exposure was a load-time decision). + "exposure_s": exposure, + "frames_per_plane": [int(f1 - f0) for f0, f1 in bounds], + "plane_frame_bounds": [[int(f0), int(f1)] for f0, f1 in bounds], + "bin": bin_factor, + "backend": backend, + # The free navigator (see plane_counts). A reader cannot hand + # HyperSpy a navigator directly, and the signal dict's + # `attributes` key does NOT reach the signal object — so it + # travels as ordinary original_metadata, which does, and a + # viewer picks it up from there. Without this the overview is + # built by reducing every plane, i.e. by integrating the whole + # movie: the one thing this reader exists to avoid. + "plane_counts": plane_counts(ds, bounds).tolist(), + }, + }, + }] diff --git a/spyde/external/rsciio_csb/_core.py b/spyde/external/rsciio_csb/_core.py new file mode 100644 index 00000000..5697b44a --- /dev/null +++ b/spyde/external/rsciio_csb/_core.py @@ -0,0 +1,888 @@ +"""Fast reader and frame accumulator for Direct Electron CSB centroid-streaming files. + +Drop-in-compatible with the reference ``de_csb.py`` for the summation path, but +built for real-time ingest. Verified to reproduce the reference reader's output +bit-identically (0 differing pixels of 67,108,864 on the Apollo XS test movie). + +Real-time context +----------------- +At 390 us/frame and ~314,546 events/frame the acquisition rate is roughly +807 M events/s (1.61 GB/s of payload). Measured on an RTX 3050 (a deliberately +weak proxy for production hardware): + + reference de_csb.py ................. 0.5 M ev/s 0.0006x real-time + this module, CPU (numpy) backend ..... 19-28 M ev/s 0.02-0.04x + this module, CPU (numpy, block-local) 50-70 M ev/s 0.06-0.09x + this module, CPU (numba, multithread) >800 M ev/s real-time-class* + this module, GPU, payload in host RAM ~2,800 M ev/s 3.5x (PCIe-bound) + this module, GPU, payload in VRAM .... 13,140 M ev/s 16.3x + +The host-RAM path is limited by PCIe transfer (measured 5.65 GB/s pinned), not +by the kernel, so handing this class device-resident data is worth roughly 5x. + +* The numba CPU path scales with core count and clears the real-time bar on a + many-core host once the ~1 s JIT warm-up is paid (on first ``add_frames``, or + reloaded from numba's on-disk cache). It is the path to use when no GPU is + present but sustained ingest matters; for a single one-shot conversion the + warm-up can dominate and the block-local numpy path is the better default. + All three CPU paths are bit-identical to one another and to the GPU path. + +File format +----------- +Little-endian throughout. A 108-byte metadata header (padded out to +``csb_data_offset``), then one uint16 per event, then a footer of one uint16 +event-count per block. + +A frame is tiled by ``csb_block_width`` x ``csb_block_height`` blocks. Each +event is a single uint16 giving its raster position *inside* its own block:: + + y = block_y_origin + word // csb_block_width + x = block_x_origin + word % csb_block_width + +Block ordering is given by the ``csb_order`` header field (0 = column-major, +1 = row-major). Events arrive in readout (time) order, which is ascending +raster order within a readout pass; a block typically spans one to four passes, +so its words form a few ascending runs rather than one sorted list. Nothing +here depends on that ordering. + +Why it is fast +-------------- +Because the payload word is *already* the raster index inside its block, an +accumulator stored BLOCK-MAJOR - one contiguous ``block_w*block_h`` tile per +block - needs no coordinate arithmetic at all: the scatter target is simply +``acc[block_index * tile_size + word]``. Every atomic issued by one thread +block then lands inside a single contiguous tile (256 KB for 2048x32), which +stays resident in L2. One transpose at readout restores normal image layout. + +See README.md for the measured tuning study behind ``FRAMES_PER_GROUP`` and +``THREADS_PER_BLOCK``. +""" + +from __future__ import annotations + +import math +import os +import struct +from contextlib import nullcontext as _nullcontext +from typing import Optional, Sequence, Tuple + +import numpy as np + +try: # optional GPU backend + import cupy as _cp +except Exception: # pragma: no cover + _cp = None + +try: # optional JIT CPU backend + import numba as _numba +except Exception: # pragma: no cover + _numba = None + +__all__ = ["CSBFile", "CSBAccumulator", "sum_file", "write_fraction_stack", + "bin_image", "gpu_available", "numba_available"] + +CSB_MAGIC = 13240 +_HEADER_FIELDS_BYTES = 108 + + +def gpu_available() -> bool: + """True if a working CuPy/CUDA device is present.""" + if _cp is None: + return False + try: + _cp.cuda.runtime.getDeviceCount() + return True + except Exception: + return False + + +def numba_available() -> bool: + """True if Numba is importable (used by the fast multithreaded CPU path).""" + return _numba is not None + + +# -------------------------------------------------------------------------- +# Header / container +# -------------------------------------------------------------------------- +class CSBFile: + """Parses a CSB file's header and block table and memory-maps the payload. + + Nothing is read eagerly beyond the header and the block-length table, so + opening a multi-gigabyte movie is cheap. + """ + + def __init__(self, path: str): + self.path = path + self.size = os.path.getsize(path) + if self.size < _HEADER_FIELDS_BYTES + 2: + raise ValueError(f"{path}: too short to be a CSB file") + + with open(path, "rb") as fh: + h = fh.read(_HEADER_FIELDS_BYTES) + if len(h) < _HEADER_FIELDS_BYTES: + raise ValueError(f"{path}: truncated header") + + u16 = lambda o: struct.unpack_from(" self.size: + raise ValueError(f"{path}: block table is incomplete " + f"(needs {need} bytes, file is {self.size})") + + self.counts = np.asarray( + np.memmap(path, dtype=" Tuple[int, int]: + """(y, x) pixel origin of a block index within a frame.""" + if self.csb_order == 1: # row-major + by, bx = divmod(block_index, self.blocks_per_width) + else: # column-major (default) + bx, by = divmod(block_index, self.blocks_per_height) + return by * self.csb_block_height, bx * self.csb_block_width + + def frame_slice(self, f0: int, f1: int) -> Tuple[int, int]: + """(first_word, n_words) of the payload for frames [f0, f1).""" + a = f0 * self.blocks_per_frame + b = f1 * self.blocks_per_frame + s = int(self.starts[a]) + e = int(self._ends[b - 1]) + return s, e - s + + def events_in_range(self, f0: int, f1: int) -> int: + """Number of events in frames [f0, f1).""" + return int(self.counts[f0 * self.blocks_per_frame: + f1 * self.blocks_per_frame].sum()) + + def frame_events(self, frame: int) -> Tuple[np.ndarray, np.ndarray]: + """Decoded (y, x) int32 arrays for one frame, in readout order. + + Convenience/reference path - not the fast path. + """ + a = frame * self.blocks_per_frame + b = a + self.blocks_per_frame + s, e = int(self.starts[a]), int(self._ends[b - 1]) + w = np.asarray(self.events[s:e], dtype=np.int32) + cnt = self.counts[a:b] + oy = np.empty(self.blocks_per_frame, np.int32) + ox = np.empty(self.blocks_per_frame, np.int32) + for i in range(self.blocks_per_frame): + oy[i], ox[i] = self.block_origin(i) + y = np.repeat(oy, cnt) + (w // self.csb_block_width) + x = np.repeat(ox, cnt) + (w % self.csb_block_width) + return y, x + + def info(self) -> dict: + ts = self.timestamp / 1000.0 + import datetime + return { + "path": self.path, + "frame_width": self.frame_width, "frame_height": self.frame_height, + "frame_count": self.frame_count, "ang_per_pix": self.ang_per_pix, + "microsec_per_frame": self.microsec_per_frame, + "csb_block_width": self.csb_block_width, + "csb_block_height": self.csb_block_height, + "csb_order": "row-major" if self.csb_order == 1 else "column-major", + "blocks_per_frame": self.blocks_per_frame, + "n_events": self.n_events, + "camera_sn": self.camera_sn, "superres_factor": self.superres_factor, + "microscope_kv": self.microscope_kv, "mag": self.microscope_mag, + "defocus_um": self.microscope_defocus, + "timestamp": self.timestamp, + "datetime": datetime.datetime.fromtimestamp(ts).isoformat(sep=" "), + } + + def __repr__(self): + return (f"") + + +# -------------------------------------------------------------------------- +# GPU kernel +# -------------------------------------------------------------------------- +# blockIdx.x is scheduled fastest, so it MUST be the frame group and +# blockIdx.y the tile. Concurrent thread blocks then share one tile, which +# stays hot in L2. Putting the tile on .x instead costs ~6x (measured). +_KERNEL_SRC = r''' +extern "C" __global__ +void csb_accumulate(const unsigned short* __restrict__ words, + long long word_base, + const long long* __restrict__ starts, + const int* __restrict__ counts, + int* acc, int blocks_per_frame, int first_frame, + int n_frames, int frames_per_group, int tile_size) +{ + const int b = blockIdx.y; + const int g0 = blockIdx.x * frames_per_group; + const int g1 = min(g0 + frames_per_group, n_frames); + int* tile = acc + (long long)b * tile_size; + + for (int f = g0; f < g1; ++f) { + const int k = (first_frame + f) * blocks_per_frame + b; + const int n = counts[k]; + if (n <= 0) continue; + // starts[] hold absolute file word offsets; word_base rebases them + // onto `words`. Done here rather than in Python so no temporary array + // is created on a different stream than the one this kernel runs on. + const unsigned short* p = words + (starts[k] - word_base); + for (int i = threadIdx.x; i < n; i += blockDim.x) + atomicAdd(&tile[p[i]], 1); + } +} +''' + +# Measured optimum on an RTX 3050 for 2048x32 blocks; these two interact +# strongly (at 25 frames/group the best thread count is 1024, not 256), so +# re-tune them together on new hardware. See README.md. +FRAMES_PER_GROUP = 4 +THREADS_PER_BLOCK = 256 + +_kernel_cache = {} + + +def _get_kernel(): + if "k" not in _kernel_cache: + if _cp is None: + raise RuntimeError("CuPy is not installed; GPU backend unavailable") + _kernel_cache["k"] = _cp.RawModule(code=_KERNEL_SRC).get_function( + "csb_accumulate") + return _kernel_cache["k"] + + +# -------------------------------------------------------------------------- +# Numba CPU kernel +# -------------------------------------------------------------------------- +# Same block-major idea as the GPU kernel, mapped onto CPU threads. Each thread +# owns a disjoint set of spatial blocks and walks them through every frame, +# writing only into that block's own tile. Because tiles are disjoint slices of +# `acc`, threads never touch the same memory - no atomics, no privatisation, no +# reduction. The scatter for one block stays inside its `tile_size`-word tile +# (256 KB for 2048x32), which is L2-resident, mirroring the GPU's L2 trick. +# +# Built lazily so `import csb` costs nothing when Numba is absent, and so the +# ~1 s JIT warm-up is paid once, on first use, not at import. +_numba_kernel_cache = {} + + +def _get_numba_kernel(): + if "k" not in _numba_kernel_cache: + if _numba is None: + raise RuntimeError("Numba is not installed; numba backend unavailable") + from numba import njit, prange + + @njit(parallel=True, cache=True, nogil=True, fastmath=False) + def _accumulate(words, base_word, starts, counts, acc, + blocks_per_frame, first_frame, n_frames, tile_size): + # One parallel iteration per spatial block. Blocks map to disjoint + # tiles in `acc`, so iterations are independent - safe to run on + # separate threads with no synchronisation. + for b in prange(blocks_per_frame): + base = b * tile_size + for f in range(n_frames): + k = (first_frame + f) * blocks_per_frame + b + n = counts[k] + if n <= 0: + continue + p = starts[k] - base_word + for i in range(n): + acc[base + words[p + i]] += 1 + + _numba_kernel_cache["k"] = _accumulate + return _numba_kernel_cache["k"] + + +# -------------------------------------------------------------------------- +# Accumulator +# -------------------------------------------------------------------------- +class CSBAccumulator: + """Accumulates CSB events into a summed image. + + The accumulator lives in block-major layout for the whole lifetime of the + object and is only rearranged into image layout when you ask for it, so a + live view can keep adding frames without ever paying for a full readout. + + Parameters + ---------- + csb : CSBFile + backend : {"auto", "gpu", "cpu", "cpu-numba", "cpu-numpy"} + "auto" selects the GPU when CuPy and a device are available, otherwise + the fastest CPU path. "cpu" uses the Numba kernel when Numba is + installed and falls back to pure NumPy otherwise. "cpu-numba" and + "cpu-numpy" force one specific CPU path (mainly for benchmarking and + tests); all three CPU paths produce bit-identical results. + dtype : + Accumulator element type. int32 is plenty for counting (the busiest + pixel in a 400-frame Apollo XS movie sees ~100 events). + + Examples + -------- + >>> f = CSBFile("movie.csb") # doctest: +SKIP + >>> acc = CSBAccumulator(f) # doctest: +SKIP + >>> acc.add_frames(0, f.frame_count) # doctest: +SKIP + >>> img = acc.image() # doctest: +SKIP + """ + + def __init__(self, csb: CSBFile, backend: str = "auto", dtype=np.int32): + self.csb = csb + self.dtype = dtype + if backend == "auto": + backend = "gpu" if gpu_available() else "cpu" + if backend == "cpu": # pick the best CPU path + backend = "cpu-numba" if numba_available() else "cpu-numpy" + if backend == "gpu" and not gpu_available(): + raise RuntimeError("GPU backend requested but no CUDA device is available") + if backend == "cpu-numba" and not numba_available(): + raise RuntimeError("cpu-numba backend requested but Numba is not installed") + if backend not in ("gpu", "cpu-numba", "cpu-numpy"): + raise ValueError(f"unknown backend {backend!r}") + self.backend = backend + # True for any CPU path; the GPU path is the sole special case elsewhere. + self._on_gpu = backend == "gpu" + + n = csb.blocks_per_frame * csb.tile_size + if backend == "gpu": + self._acc = _cp.zeros(n, dtype=dtype) + self._d_counts = _cp.asarray(csb.counts.astype(np.int32)) + self._d_starts = _cp.asarray(csb.starts) + self._kernel = _get_kernel() + elif backend == "cpu-numba": + self._acc = np.zeros(n, dtype=np.int64) + self._kernel = _get_numba_kernel() # triggers JIT on first use + else: + self._acc = np.zeros(n, dtype=np.int64) # bincount output dtype + # Device staging buffers for host payloads, held until the work that + # reads them has completed. Without this the buffer could be returned + # to CuPy's pool and reused while an async copy/kernel is still pending. + self._pending: list = [] + self.events_added = 0 + + # -- ingest ----------------------------------------------------------- + def add_frames(self, f0: int, f1: int, payload=None, stream=None) -> None: + """Accumulate frames [f0, f1). + + Parameters + ---------- + payload : + ``None`` - read from the file's memory map (offline use). + ndarray - host uint16 array holding *exactly* the payload for + these frames. Use pinned memory for best transfer rate. + cupy array- device-resident payload for these frames; no transfer + is performed. This is the fastest path by ~5x. + stream : + Optional CuPy stream. Work is enqueued and NOT synchronised, so you + can overlap successive calls; call :meth:`synchronize` before + reading results. + """ + if not (0 <= f0 < f1 <= self.csb.frame_count): + raise ValueError(f"bad frame range [{f0}, {f1}) for a " + f"{self.csb.frame_count}-frame movie") + base_word, n_words = self.csb.frame_slice(f0, f1) + + if payload is None: + payload = np.asarray(self.csb.events[base_word:base_word + n_words]) + if len(payload) != n_words: + raise ValueError(f"payload has {len(payload):,} words, expected " + f"{n_words:,} for frames [{f0}, {f1})") + + if self._on_gpu: + self._add_gpu(f0, f1, payload, base_word, stream) + elif self.backend == "cpu-numba": + self._add_numba(f0, f1, payload, base_word) + else: + self._add_cpu(f0, f1, payload, base_word) + self.events_added += int(n_words) + + def _add_gpu(self, f0, f1, payload, base_word, stream): + cs = self.csb + # Everything - staging allocation, transfer and kernel - must be issued + # inside the same stream context, or work lands on the default stream + # and races with the kernel. + ctx = stream if stream is not None else _nullcontext() + with ctx: + if isinstance(payload, np.ndarray): # host -> device + d_words = _cp.empty(len(payload), dtype=_cp.uint16) + if stream is not None: + _cp.cuda.runtime.memcpyAsync( + d_words.data.ptr, payload.ctypes.data, payload.nbytes, + _cp.cuda.runtime.memcpyHostToDevice, stream.ptr) + else: + d_words.set(payload) + self._pending.append(d_words) # released by synchronize() + else: # already device-resident + d_words = payload + + n_frames = f1 - f0 + grid = ((n_frames + FRAMES_PER_GROUP - 1) // FRAMES_PER_GROUP, + cs.blocks_per_frame) + self._kernel(grid, (THREADS_PER_BLOCK,), ( + d_words, np.int64(base_word), self._d_starts, self._d_counts, + self._acc, np.int32(cs.blocks_per_frame), np.int32(f0), + np.int32(n_frames), np.int32(FRAMES_PER_GROUP), + np.int32(cs.tile_size))) + + def _add_cpu(self, f0, f1, payload, base_word): + """Pure-NumPy block-local histogram. + + The event word is already the intra-block raster index, so each spatial + block's events histogram directly into that block's own `tile_size`-word + slice of the accumulator. Doing it per block keeps every ``bincount`` + working set inside one L2-resident tile and, crucially, never builds the + full-accumulator-sized index array or temporary the naive one-shot + ``bincount`` needs (that temp is ~536 MB for an 8192^2 frame and is + touched three times per call regardless of event count). + + All of a block's events across the whole chunk are gathered once so the + per-block Python overhead is paid ``blocks_per_frame`` times, not once + per (block, frame). + """ + cs = self.csb + bpf = cs.blocks_per_frame + tile_size = cs.tile_size + n_frames = f1 - f0 + counts = cs.counts[f0 * bpf:f1 * bpf].reshape(n_frames, bpf) + # Per-block start offset into `payload` for every (frame, block), from a + # single cumulative sum of the counts in payload (readout) order. + ends = np.cumsum(counts.reshape(-1)) + starts = ends - counts.reshape(-1) + starts = starts.reshape(n_frames, bpf) + ends = ends.reshape(n_frames, bpf) + words = np.asarray(payload) + acc = self._acc + for b in range(bpf): + # Gather this spatial block's events from every frame in the chunk. + if n_frames == 1: + seg = words[starts[0, b]:ends[0, b]] + else: + parts = [words[starts[f, b]:ends[f, b]] for f in range(n_frames)] + seg = np.concatenate(parts) if parts else words[:0] + if seg.size: + acc[b * tile_size:(b + 1) * tile_size] += np.bincount( + seg, minlength=tile_size) + + def _add_numba(self, f0, f1, payload, base_word): + """Multithreaded JIT histogram; see :func:`_get_numba_kernel`.""" + cs = self.csb + words = np.ascontiguousarray(payload, dtype=np.uint16) + # `starts` are absolute file word offsets; the kernel rebases them onto + # `words` with `base_word`, matching the GPU kernel exactly. + self._kernel(words, np.int64(base_word), cs.starts, cs.counts, self._acc, + np.int64(cs.blocks_per_frame), np.int64(f0), + np.int64(f1 - f0), np.int64(cs.tile_size)) + + def synchronize(self, stream=None) -> None: + """Wait for outstanding GPU work and release staging buffers. + + Pass the stream(s) you enqueued on, or nothing to wait on the whole + device. Staging buffers for host payloads are only freed here, so call + it periodically in a long-running ingest loop. + """ + if not self._on_gpu: + return + if stream is not None: + stream.synchronize() + else: + _cp.cuda.Device().synchronize() + self._pending.clear() + + def reset(self) -> None: + self.synchronize() + self._acc.fill(0) + self.events_added = 0 + + # -- readout ---------------------------------------------------------- + def _untile(self): + """Block-major accumulator -> padded 2-D image, still on-device.""" + cs = self.csb + bh, bw = cs.csb_block_height, cs.csb_block_width + if cs.csb_order == 1: # row-major + v = self._acc.reshape(cs.blocks_per_height, cs.blocks_per_width, bh, bw) + v = v.transpose(0, 2, 1, 3) + else: # column-major + v = self._acc.reshape(cs.blocks_per_width, cs.blocks_per_height, bh, bw) + v = v.transpose(1, 2, 0, 3) + return v.reshape(cs.padded_height, cs.padded_width) + + def image(self, dtype=None) -> np.ndarray: + """The summed image as a host ndarray, cropped to the true frame size. + + Costs a full device-to-host copy (~90 ms for 8192^2 on an RTX 3050). + For a live view prefer :meth:`preview`, which is ~5x cheaper. + """ + cs = self.csb + v = self._untile()[:cs.frame_height, :cs.frame_width] + if self._on_gpu: + out = _cp.asnumpy(_cp.ascontiguousarray(v)) + else: + out = np.ascontiguousarray(v) + return out if dtype is None else out.astype(dtype) + + def preview(self, bin_factor: int = 4, dtype=np.float32) -> np.ndarray: + """Binned preview for live display; the binning happens on-device. + + At bin_factor=4 an 8192^2 accumulator returns 2048^2 in ~18 ms, + i.e. a sustainable ~55 Hz display rate. + """ + if bin_factor < 1: + raise ValueError("bin_factor must be >= 1") + cs = self.csb + h = (cs.frame_height // bin_factor) * bin_factor + w = (cs.frame_width // bin_factor) * bin_factor + v = self._untile()[:h, :w] + xp = _cp if self._on_gpu else np + v = xp.ascontiguousarray(v).reshape(h // bin_factor, bin_factor, + w // bin_factor, bin_factor) + out = v.sum(axis=(1, 3), dtype=xp.int64).astype(dtype) + return _cp.asnumpy(out) if self._on_gpu else out + + def apply_corrections(self, image: np.ndarray, + gain: Optional[np.ndarray] = None, + defects: Optional[np.ndarray] = None) -> np.ndarray: + """Apply gain and defect masking to a summed image. + + Both are exactly equivalent to the reference reader's per-event + handling and far cheaper: gain and the defect mask depend only on + (y, x), so scaling or zeroing the *sum* gives an identical result to + scaling or skipping each event. + + Note this deliberately does NOT reproduce ``de_csb.py``'s defect + *filling*, which synthesises random events on flagged pixels at the + local block sparsity. That fabricates data and has no place in a + quantitative path; do it in the display layer if you want it. + """ + out = image + if gain is not None: + if gain.shape != out.shape: + raise ValueError(f"gain shape {gain.shape} != image {out.shape}") + out = out.astype(np.float32) * gain.astype(np.float32) + if defects is not None: + if defects.shape != out.shape: + raise ValueError(f"defects shape {defects.shape} != image {out.shape}") + out = out.copy() + out[defects.astype(bool)] = 0 + return out + + +# -------------------------------------------------------------------------- +# Convenience +# -------------------------------------------------------------------------- +def bin_image(img: np.ndarray, factor: int) -> np.ndarray: + """Bin an image by an integer factor, summing counts within each block. + + Summing (not averaging) is correct for count data: a physical pixel's count + is the sum of the super-resolution sub-pixels it contains. Edge rows/columns + that don't fill a full block are cropped, matching the readout convention. + """ + if factor < 1: + raise ValueError("bin factor must be >= 1") + if factor == 1: + return img + h = (img.shape[0] // factor) * factor + w = (img.shape[1] // factor) * factor + return img[:h, :w].reshape(h // factor, factor, + w // factor, factor).sum(axis=(1, 3)) + + +def write_fraction_stack(csb, out_path: str, frames_per_fraction: int, + frames: Optional[Sequence[int]] = None, + bin_factor: int = 1, + gain: Optional[np.ndarray] = None, + defects: Optional[np.ndarray] = None, + backend: str = "auto") -> list: + """Write an MRC stack of dose-fraction sums, one plane per group of frames. + + Each output plane is the sum of ``frames_per_fraction`` consecutive CSB + frames. This is the movie you hand to an external motion estimator + (cryoSPARC / MotionCor / RELION); the trajectories it returns drive the + motion-correction step. + + Output is at native (super-resolution) sampling by default. ``bin_factor`` + reduces it as a final step only - super-res is the point of centroid data. + + Parameters + ---------- + csb : CSBFile or str + out_path : str + Destination MRC. Written as a 3-D float32 stack, streamed plane by + plane, with ``voxel_size`` set from the header (scaled by bin_factor). + frames_per_fraction : int + CSB frames summed into each output plane. + frames : (start, end), optional + Frame range; default all frames. A trailing partial fraction is kept. + bin_factor : int + Integer output binning, summed. 1 = native super-resolution (default). + gain, defects : ndarray, optional + Applied per plane, at full resolution, before binning. + + Returns + ------- + list of (f0, f1) + Frame boundaries for each output plane, in order. Piece 2 needs these + to map returned trajectories back onto individual frames. + """ + import mrcfile + + if not isinstance(csb, CSBFile): + csb = CSBFile(csb) + if frames_per_fraction < 1: + raise ValueError("frames_per_fraction must be >= 1") + if bin_factor < 1: + raise ValueError("bin_factor must be >= 1") + + f0, f1 = (0, csb.frame_count) if frames is None else frames + f0 = max(0, f0) + f1 = csb.frame_count if f1 < 0 else min(f1, csb.frame_count) + if f1 <= f0: + raise ValueError(f"empty frame range [{f0}, {f1})") + + bounds = [(a, min(a + frames_per_fraction, f1)) + for a in range(f0, f1, frames_per_fraction)] + h_out = (csb.frame_height // bin_factor) + w_out = (csb.frame_width // bin_factor) + + acc = CSBAccumulator(csb, backend=backend) + apix = csb.ang_per_pix * bin_factor + with mrcfile.new_mmap(out_path, shape=(len(bounds), h_out, w_out), + mrc_mode=2, overwrite=True) as m: + for i, (a, b) in enumerate(bounds): + acc.reset() + acc.add_frames(a, b) + acc.synchronize() + img = acc.image(dtype=np.float32) + if gain is not None or defects is not None: + img = acc.apply_corrections(img, gain=gain, defects=defects) + m.data[i] = bin_image(img, bin_factor).astype(np.float32) + if apix > 0: + m.voxel_size = apix + label = (f"CSB dose-fraction stack: {frames_per_fraction} frames/plane" + + (f", bin {bin_factor}" if bin_factor > 1 else "")) + m.header.nlabl = 1 + m.header.label[0] = label.encode("ascii")[:80] + return bounds + + +def sum_file(path: str, frames: Optional[Sequence[int]] = None, + backend: str = "auto", chunk_frames: int = 20, + n_streams: int = 3) -> np.ndarray: + """Sum a whole CSB file (or a frame range) into an image. + + On the GPU backend this stages through pinned host buffers across several + CUDA streams so the file read overlaps the transfer and the kernel. Reading + from the memory map straight into a device buffer instead is ~8x slower. + """ + csb = CSBFile(path) + f0, f1 = (0, csb.frame_count) if frames is None else frames + f0 = max(0, f0) + f1 = csb.frame_count if f1 < 0 else min(f1, csb.frame_count) + if f1 <= f0: + raise ValueError(f"empty frame range [{f0}, {f1})") + acc = CSBAccumulator(csb, backend=backend) + + bounds = [(a, min(a + chunk_frames, f1)) for a in range(f0, f1, chunk_frames)] + if not acc._on_gpu: # any CPU path: plain sequential add + for a, b in bounds: + acc.add_frames(a, b) + return acc.image() + + import threading + from concurrent.futures import ThreadPoolExecutor + + nmax = max(csb.frame_slice(a, b)[1] for a, b in bounds) + streams = [_cp.cuda.Stream(non_blocking=True) for _ in range(n_streams)] + + # A ring of pinned buffers, deeper than the stream count so a reader rarely + # has to wait. Reads run on worker threads (readinto releases the GIL), so + # file I/O overlaps the transfer and the kernel. NVMe needs several + # concurrent readers to reach full bandwidth. + ring = max(6, 2 * n_streams) + ring = min(ring, len(bounds)) + staging = [] + for _ in range(ring): + mem = _cp.cuda.alloc_pinned_memory(nmax * 2) + staging.append(np.frombuffer(mem, dtype=np.uint16, count=nmax)) + slot_ready = [None] * ring # event: prior GPU use of this slot done + local = threading.local() + + def fetch(j): + slot = j % ring + ev = slot_ready[slot] + if ev is not None: + ev.synchronize() + if not hasattr(local, "fh"): + local.fh = open(path, "rb", buffering=0) + s, n = csb.frame_slice(*bounds[j]) + local.fh.seek(csb.csb_data_offset + 2 * s) + local.fh.readinto(memoryview(staging[slot])[:n].cast("B")) + return n + + n_readers = min(6, max(2, ring // 2)) + ex = ThreadPoolExecutor(max_workers=n_readers) + try: + futs = {j: ex.submit(fetch, j) for j in range(min(ring, len(bounds)))} + for j, (a, b) in enumerate(bounds): + n = futs.pop(j).result() + slot, sl = j % ring, j % n_streams + acc.add_frames(a, b, payload=staging[slot][:n], stream=streams[sl]) + ev = _cp.cuda.Event() + ev.record(streams[sl]) + slot_ready[slot] = ev + nxt = j + ring + if nxt < len(bounds): + futs[nxt] = ex.submit(fetch, nxt) + finally: + ex.shutdown(wait=True) + for st in streams: + st.synchronize() + acc.synchronize() + return acc.image() + + +def _load_mrc(path): + import mrcfile + with mrcfile.open(path, permissive=True) as m: + return np.asarray(m.data) + + +def _main(argv=None): + import argparse + import time + p = argparse.ArgumentParser( + description="Convert a CSB file to an MRC image or dose-fraction stack.") + p.add_argument("input_csb") + p.add_argument("output_mrc", nargs="?") + p.add_argument("--frames", type=int, nargs=2, default=[0, -1], + metavar=("START", "END"), help="frame range [start end)") + p.add_argument("--fractions", type=int, metavar="N", + help="write an MRC stack, summing N CSB frames per plane " + "(for feeding to a motion estimator)") + p.add_argument("--bin", type=int, default=1, metavar="B", + help="integer output binning, summed (default 1 = native " + "super-resolution)") + p.add_argument("--backend", + choices=["auto", "gpu", "cpu", "cpu-numba", "cpu-numpy"], + default="auto") + p.add_argument("--gain", help="MRC gain reference") + p.add_argument("--defects", help="MRC defect map (non-zero = bad pixel)") + p.add_argument("--info", action="store_true", help="print header and exit") + a = p.parse_args(argv) + + csb = CSBFile(a.input_csb) + for k, v in csb.info().items(): + print(f" {k}: {v}") + if a.info: + return 0 + + f0 = max(0, a.frames[0]) + f1 = csb.frame_count if a.frames[1] < 0 else min(a.frames[1], csb.frame_count) + gain = _load_mrc(a.gain) if a.gain else None + defects = _load_mrc(a.defects) if a.defects else None + + # --- dose-fraction stack ------------------------------------------------- + if a.fractions: + if not a.output_mrc: + p.error("--fractions requires an output_mrc path") + t0 = time.perf_counter() + bounds = write_fraction_stack(csb, a.output_mrc, a.fractions, + frames=[f0, f1], bin_factor=a.bin, + gain=gain, defects=defects, + backend=a.backend) + dt = time.perf_counter() - t0 + print(f"\n wrote {len(bounds)}-plane stack " + f"({a.fractions} frames/plane, bin {a.bin}) to {a.output_mrc} " + f"in {dt*1000:.0f} ms") + return 0 + + # --- single summed image ------------------------------------------------- + t0 = time.perf_counter() + img = sum_file(a.input_csb, [f0, f1], backend=a.backend) + dt = time.perf_counter() - t0 + n_ev = csb.events_in_range(f0, f1) + acq = (f1 - f0) * csb.microsec_per_frame * 1e-6 + print(f"\n frames [{f0}, {f1}) : summed {n_ev:,} events in {dt*1000:.1f} ms " + f"({n_ev/dt/1e6:,.0f} M events/s)") + if acq > 0: + print(f" acquisition was {acq*1000:.0f} ms -> {acq/dt:.2f}x real-time " + f"(includes file read and one-time CUDA/context setup)") + + if gain is not None or defects is not None: + acc = CSBAccumulator(csb, backend=a.backend) + img = acc.apply_corrections(img, gain, defects) + img = bin_image(img, a.bin) + + if a.output_mrc: + import mrcfile + with mrcfile.new(a.output_mrc, overwrite=True) as m: + m.set_data(img.astype(np.float32)) + if csb.ang_per_pix > 0: + m.voxel_size = csb.ang_per_pix * a.bin + print(f" wrote {a.output_mrc}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(_main()) diff --git a/spyde/external/rsciio_csb/_sparse.py b/spyde/external/rsciio_csb/_sparse.py new file mode 100644 index 00000000..c2707a22 --- /dev/null +++ b/spyde/external/rsciio_csb/_sparse.py @@ -0,0 +1,435 @@ +"""Time-indexed sparse view over a CSB centroid-streaming file. + +This is the layer a viewer sits on. :class:`SparseCSB` wraps the fast +:class:`csb.CSBFile` / :class:`csb.CSBAccumulator` core and adds: + +* a physical **time axis** (frames carry a ``microsec_per_frame`` cadence), so + you address the movie in seconds, not frame indices; +* :meth:`SparseCSB.integrate` - one dense image summed over a time window + ``[t0, t1)``, the primitive a slider-driven viewer calls on every drag; +* :meth:`SparseCSB.to_frames` - a dense stack, one plane per fixed time step, + the whole movie rebinned to a chosen exposure; +* :meth:`SparseCSB.events` / :meth:`SparseCSB.events_in_time` - the decoded + ``(y, x)`` event lists (a COO-style sparse view) for an event-based renderer; +* :meth:`SparseCSB.to_hyperspy` and the module-level :func:`file_reader` - a + RosettaSciIO-compatible signal dictionary, so the result drops straight into + HyperSpy with calibrated navigation (time) and signal (y, x) axes. + +The numeric core returns plain NumPy arrays and has **no** hard HyperSpy / dask +dependency; those are imported lazily only inside :meth:`to_hyperspy` / +:func:`file_reader` and only when you ask for a lazy signal. + +Design note - why seconds map to whole frames +---------------------------------------------- +A CSB frame is the atomic exposure; there is no sub-frame timing in the format. +So a time window ``[t0, t1)`` is realised as the half-open *frame* range whose +frames start within it. :meth:`time_to_frames` is the single place that mapping +lives, and every method routes through it, so the rounding rule is defined once. +""" + +from __future__ import annotations + +import math +from typing import Optional, Tuple, Union + +import numpy as np + +from ._core import CSBFile, CSBAccumulator, bin_image + +__all__ = ["SparseCSB", "load", "file_reader"] + + +class SparseCSB: + """A lazy, time-indexed sparse dataset over one CSB file. + + Parameters + ---------- + source : str or CSBFile + A path to a ``.csb`` file, or an already-open :class:`csb.CSBFile`. + backend : {"auto", "gpu", "cpu", "cpu-numba", "cpu-numpy"} + Compute backend for the integrations; forwarded to + :class:`csb.CSBAccumulator`. ``"auto"`` uses the GPU when available. + + Notes + ----- + Opening is cheap - only the header and block table are read. The payload is + memory-mapped and touched lazily, per integration, so a multi-gigabyte movie + costs nothing until you actually integrate a window of it. + + A single accumulator is created once and reused (``reset`` + ``add_frames``) + across calls, so repeated slider drags don't reallocate device memory. + + Examples + -------- + >>> ds = SparseCSB("movie.csb") # doctest: +SKIP + >>> img = ds.integrate(0.0, 0.1) # 0 - 0.1 s # doctest: +SKIP + >>> stack = ds.to_frames(0.05) # 50 ms/plane # doctest: +SKIP + >>> y, x = ds.events_in_time(0.0, 0.01) # COO view # doctest: +SKIP + """ + + def __init__(self, source: Union[str, CSBFile], backend: str = "auto"): + self.csb = source if isinstance(source, CSBFile) else CSBFile(source) + self.backend = backend + self._acc: Optional[CSBAccumulator] = None # created on first use + + # -- time axis -------------------------------------------------------- + @property + def frame_count(self) -> int: + return self.csb.frame_count + + @property + def frame_duration(self) -> float: + """Seconds per frame (from the header ``microsec_per_frame``).""" + return self.csb.microsec_per_frame * 1e-6 + + @property + def frame_rate(self) -> float: + """Frames per second; 0.0 if the header has no cadence.""" + dt = self.frame_duration + return 1.0 / dt if dt > 0 else 0.0 + + @property + def duration(self) -> float: + """Total acquisition time in seconds (``frame_count * frame_duration``).""" + return self.frame_count * self.frame_duration + + @property + def time_axis(self) -> np.ndarray: + """Start time (s) of every frame: ``[0, dt, 2dt, ...]``, length frame_count.""" + return np.arange(self.frame_count, dtype=np.float64) * self.frame_duration + + @property + def shape(self) -> Tuple[int, int]: + """(height, width) of a frame in pixels.""" + return self.csb.frame_height, self.csb.frame_width + + def _require_timed(self) -> None: + if self.frame_duration <= 0: + raise ValueError( + "this file declares no frame cadence (microsec_per_frame is 0); " + "address it in frame indices via integrate_frames() instead of " + "seconds") + + def time_to_frames(self, t0: float, t1: float) -> Tuple[int, int]: + """Map a half-open time window ``[t0, t1)`` in seconds to frames ``[f0, f1)``. + + A frame belongs to the window when *its start time* lies in ``[t0, t1)``. + ``t1 = None``-style "to the end" is expressed by passing ``self.duration`` + (or any value past it); results are always clamped to the movie. + """ + self._require_timed() + if t1 <= t0: + raise ValueError(f"empty time window [{t0}, {t1}) s") + dt = self.frame_duration + # Frame f starts at f*dt; it is in [t0, t1) iff t0 <= f*dt < t1. + f0 = max(0, math.ceil(t0 / dt - 1e-9)) + f1 = min(self.frame_count, math.ceil(t1 / dt - 1e-9)) + if f1 <= f0: + raise ValueError( + f"time window [{t0}, {t1}) s selects no frames " + f"(cadence {dt*1e3:.4g} ms/frame); window is narrower than one frame") + return f0, f1 + + # -- accumulator plumbing -------------------------------------------- + def _accumulator(self) -> CSBAccumulator: + if self._acc is None: + self._acc = CSBAccumulator(self.csb, backend=self.backend) + return self._acc + + def _sum_frames(self, f0: int, f1: int, bin_factor: int, + dtype) -> np.ndarray: + acc = self._accumulator() + acc.reset() + acc.add_frames(f0, f1) + acc.synchronize() + img = acc.image() + if bin_factor > 1: + img = bin_image(img, bin_factor) + return img.astype(dtype) if dtype is not None else img + + # -- integration (seconds) ------------------------------------------- + def integrate(self, t0: float, t1: Optional[float] = None, + bin: int = 1, dtype=None) -> np.ndarray: + """Sum all events in the time window ``[t0, t1)`` into one dense image. + + This is the viewer primitive: give it the current slider range in + seconds and it returns the corresponding summed frame. + + Parameters + ---------- + t0, t1 : float + Window start/stop in seconds. ``t1=None`` means "to the end". + bin : int + Integer output binning (summed), applied after integration. + dtype : + Output dtype; ``None`` keeps the accumulator's integer counts. + + Returns + ------- + ndarray + 2-D image, shape ``(H//bin, W//bin)``. + """ + if t1 is None: + t1 = self.duration + f0, f1 = self.time_to_frames(t0, t1) + return self._sum_frames(f0, f1, bin, dtype) + + def integrate_frames(self, f0: int, f1: Optional[int] = None, + bin: int = 1, dtype=None) -> np.ndarray: + """Frame-indexed sibling of :meth:`integrate` (no cadence needed).""" + if f1 is None: + f1 = self.frame_count + if not (0 <= f0 < f1 <= self.frame_count): + raise ValueError(f"bad frame range [{f0}, {f1})") + return self._sum_frames(f0, f1, bin, dtype) + + # -- dense stacks ----------------------------------------------------- + def _time_bounds(self, step: float, start: float, + stop: float) -> list: + """Frame boundaries for consecutive `step`-second planes over [start, stop). + + Returns a **contiguous partition**: each plane ends exactly where the + next begins, so every frame in the covered range lands in exactly one + plane - no gaps, no overlaps - even when ``step`` is not a whole number + of frames. (Mapping each plane's time window independently, as an + earlier version did, let floating-point rounding at the boundaries put a + frame in two adjacent planes or in neither.) + + A plane's frames are those whose *start time* falls in that plane's time + window; the shared edge between plane i and i+1 is a single frame index, + computed once, so it cannot be double-assigned. Planes that would be + empty (a time step narrower than the frame cadence leaves some with no + frame start) are dropped rather than emitted as zero-frame planes. + """ + self._require_timed() + if step <= 0: + raise ValueError("step must be > 0 seconds") + dt = self.frame_duration + f_start = max(0, math.ceil(start / dt - 1e-9)) + f_stop = min(self.frame_count, math.ceil(stop / dt - 1e-9)) + if f_stop <= f_start: + return [] + n_planes = int(math.ceil((stop - start) / step - 1e-9)) + # The frame index at each plane's leading time edge, clamped and made + # monotone; consecutive edges form the [f0, f1) of each plane. + edge_f = [] + for i in range(n_planes + 1): + t_edge = start + i * step + fe = math.ceil(t_edge / dt - 1e-9) + edge_f.append(min(max(fe, f_start), f_stop)) + bounds = [] + for i in range(n_planes): + f0, f1 = edge_f[i], edge_f[i + 1] + if f1 > f0: # drop empty planes + bounds.append((f0, f1)) + return bounds + + def to_frames(self, step: float, start: float = 0.0, + stop: Optional[float] = None, bin: int = 1, + dtype=np.float32) -> Tuple[np.ndarray, np.ndarray]: + """Rebin the movie into a dense stack, one plane per ``step`` seconds. + + Parameters + ---------- + step : float + Exposure per output plane, in seconds (e.g. ``0.1`` -> 100 ms planes). + start, stop : float + Time range to cover; defaults to the whole movie. + bin : int + Integer spatial binning (summed) applied per plane. + dtype : + Output dtype (default float32, HyperSpy/MRC-friendly). + + Returns + ------- + stack : ndarray + 3-D array ``(n_planes, H//bin, W//bin)``. + plane_times : ndarray + Start time (s) of each plane, length ``n_planes`` - the calibrated + navigation axis for a viewer or for HyperSpy. + """ + if stop is None: + stop = self.duration + bounds = self._time_bounds(step, start, stop) + if not bounds: + raise ValueError( + f"no planes: step {step}s over [{start}, {stop})s selects nothing") + h = self.csb.frame_height // bin + w = self.csb.frame_width // bin + stack = np.empty((len(bounds), h, w), dtype=dtype) + for i, (f0, f1) in enumerate(bounds): + stack[i] = self._sum_frames(f0, f1, bin, dtype) + plane_times = np.array([f0 * self.frame_duration for f0, _ in bounds], + dtype=np.float64) + return stack, plane_times + + # -- sparse / event (COO) view --------------------------------------- + def events(self, frame: int) -> Tuple[np.ndarray, np.ndarray]: + """Decoded ``(y, x)`` int32 arrays for a single frame (COO-style view). + + Every returned pair is one detected event; repeated coordinates mean + repeated counts. This is the raw material for an event-based renderer + that draws points rather than an integrated image. + """ + if not (0 <= frame < self.frame_count): + raise ValueError(f"frame {frame} out of range [0, {self.frame_count})") + return self.csb.frame_events(frame) + + def events_in_time(self, t0: float, t1: Optional[float] = None + ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Decoded ``(y, x, frame)`` for every event whose frame is in ``[t0, t1)``. + + Returns three parallel arrays: pixel ``y``, pixel ``x``, and the source + frame index of each event (so a renderer can colour or fade by time). + Concatenates per-frame decodes; for very wide windows prefer + :meth:`integrate`, which never materialises the event list. + """ + if t1 is None: + t1 = self.duration + f0, f1 = self.time_to_frames(t0, t1) + ys, xs, fs = [], [], [] + for fr in range(f0, f1): + y, x = self.csb.frame_events(fr) + if y.size: + ys.append(y) + xs.append(x) + fs.append(np.full(y.size, fr, dtype=np.int32)) + if not ys: + z = np.empty(0, np.int32) + return z, z, z + return (np.concatenate(ys), np.concatenate(xs), np.concatenate(fs)) + + # -- HyperSpy / RosettaSciIO bridge ---------------------------------- + def to_signal_dict(self, step: float, start: float = 0.0, + stop: Optional[float] = None, bin: int = 1, + dtype=np.float32) -> dict: + """Build a RosettaSciIO signal dictionary for a ``to_frames`` stack. + + The returned dict has the standard keys - ``data``, ``axes``, + ``metadata``, ``original_metadata`` - and can be handed to + ``hyperspy.signals.Signal2D(**d)`` or returned from a rsciio reader. + + Axes: a navigation *time* axis (seconds) and two signal *y*/*x* axes + calibrated in angstrom from the header ``ang_per_pix`` (scaled by bin). + """ + stack, plane_times = self.to_frames(step, start, stop, bin, dtype) + return self._signal_dict(stack, plane_times, step, bin) + + def _signal_dict(self, stack: np.ndarray, plane_times: np.ndarray, + step: float, bin: int) -> dict: + cs = self.csb + apix = cs.ang_per_pix * bin + nav_scale = step if step > 0 else (self.frame_duration or 1.0) + nav_offset = float(plane_times[0]) if plane_times.size else 0.0 + axes = [ + {"name": "time", "size": stack.shape[0], "index_in_array": 0, + "scale": float(nav_scale), "offset": nav_offset, + "units": "s", "navigate": True}, + {"name": "y", "size": stack.shape[1], "index_in_array": 1, + "scale": float(apix) if apix > 0 else 1.0, "offset": 0.0, + "units": "Å" if apix > 0 else "px", "navigate": False}, + {"name": "x", "size": stack.shape[2], "index_in_array": 2, + "scale": float(apix) if apix > 0 else 1.0, "offset": 0.0, + "units": "Å" if apix > 0 else "px", "navigate": False}, + ] + info = cs.info() + metadata = { + "General": { + "title": _basename_no_ext(cs.path), + "original_filename": _basename(cs.path), + }, + "Signal": {"signal_type": ""}, + } + if info.get("datetime"): + date, _, tm = info["datetime"].partition(" ") + metadata["General"]["date"] = date + metadata["General"]["time"] = tm + return { + "data": stack, + "axes": axes, + "metadata": metadata, + "original_metadata": {"csb_header": info}, + } + + def to_hyperspy(self, step: float, start: float = 0.0, + stop: Optional[float] = None, bin: int = 1, + dtype=np.float32, lazy: bool = False): + """Return a HyperSpy ``Signal2D`` of the ``to_frames`` stack. + + Requires ``hyperspy``. With ``lazy=True`` the data is wrapped in a dask + array (requires ``dask``); the stack is still computed eagerly here - + true per-plane lazy compute is a future refinement (see module docstring). + """ + try: + import hyperspy.api as hs + except Exception as ex: # pragma: no cover + raise ImportError( + "to_hyperspy() needs hyperspy installed " + "(`uv add hyperspy` or `pip install hyperspy`)") from ex + d = self.to_signal_dict(step, start, stop, bin, dtype) + if lazy: + import dask.array as da + d = dict(d) + d["data"] = da.from_array(d["data"], chunks="auto") + return hs.signals.Signal2D(**d).as_lazy() + return hs.signals.Signal2D(**d) + + def __repr__(self) -> str: + dt = self.frame_duration + cad = f"{dt*1e3:.4g} ms/frame" if dt > 0 else "no cadence" + return (f"") + + +# -------------------------------------------------------------------------- +# Module-level helpers +# -------------------------------------------------------------------------- +def _basename(path: str) -> str: + import os + return os.path.basename(path) + + +def _basename_no_ext(path: str) -> str: + import os + return os.path.splitext(os.path.basename(path))[0] + + +def load(source: Union[str, CSBFile], backend: str = "auto") -> SparseCSB: + """Open a CSB file as a :class:`SparseCSB` (thin constructor alias).""" + return SparseCSB(source, backend=backend) + + +def file_reader(filename: str, lazy: bool = False, step: Optional[float] = None, + bin: int = 1, backend: str = "auto", dtype=np.float32, + **kwds) -> list: + """RosettaSciIO-style reader: return a list with one signal dictionary. + + Parameters + ---------- + filename : str + Path to the ``.csb`` file. + lazy : bool + If True, wrap ``data`` in a dask array (needs ``dask``). + step : float, optional + Seconds per output plane. Defaults to one plane per frame + (``step = frame_duration``), i.e. the movie at native cadence. + bin : int + Integer spatial binning (summed) per plane. + backend : str + Compute backend forwarded to the accumulator. + + Returns + ------- + list of dict + A single-element list holding the signal dictionary, per the + RosettaSciIO contract. + """ + ds = SparseCSB(filename, backend=backend) + if step is None: + step = ds.frame_duration if ds.frame_duration > 0 else 1.0 + d = ds.to_signal_dict(step, bin=bin, dtype=dtype) + if lazy: + import dask.array as da + d["data"] = da.from_array(d["data"], chunks="auto") + return [d] diff --git a/spyde/external/rsciio_csb/specifications.yaml b/spyde/external/rsciio_csb/specifications.yaml new file mode 100644 index 00000000..acb81991 --- /dev/null +++ b/spyde/external/rsciio_csb/specifications.yaml @@ -0,0 +1,10 @@ +name: "CSB" +name_aliases: ["DE_CSB"] +description: "Direct Electron CSB (compressed sparse block) centroid-streaming + files: a sparse event stream rather than a dense frame stack. Reading + integrates events over a time window per plane." +full_support: False +file_extensions: ["csb", "CSB"] +default_extension: 0 +writes: False +non_uniform_axis: False diff --git a/spyde/toolbars.yaml b/spyde/toolbars.yaml index 609ee95d..4f1da2a6 100644 --- a/spyde/toolbars.yaml +++ b/spyde/toolbars.yaml @@ -294,6 +294,34 @@ functions: type: enum default: intensity options: [intensity, count] + To Frames: + description: Re-cut a CSB event stream into an in-situ movie at a chosen frame rate, integrating the events in each exposure. Adds a new lazy dataset to the signal tree. + icon: drawing/toolbars/icons/playback.svg + function: spyde.actions.csb_to_frames.csb_to_frames + requires_original_metadata: csb + plot_dim: [2] + toolbar_side: bottom + navigation: False + parameters: + fps: + name: Frame rate (fps) + type: float + default: 100.0 + min: 0.001 + max: 100000.0 + exposure_ms: + name: or Exposure (ms) + type: float + default: 0.0 + min: 0.0 + max: 60000.0 + bin: + name: Bin + type: int + default: 1 + min: 1 + max: 16 + EBSD Indexing: description: Dictionary indexing of EBSD patterns. Build a simulated pattern dictionary, check the matched orientation's Kikuchi bands live under the crosshair, then index the whole scan into an IPF orientation map. icon: drawing/toolbars/icons/ebsd.svg From 28969c92b02f944c1f6736279bbc15afc4b1365a Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Tue, 28 Jul 2026 12:07:44 -0500 Subject: [PATCH 02/14] fix(csb): integrate on torch-CUDA, and let Open actually pick a .csb MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two things that made the reader look worse than it is. **The GPU was never used.** _core's GPU lane is a CuPy RawModule kernel, and SpyDE does not install CuPy — its GPU stack is torch, pinned with CUDA. So on a SpyDE machine gpu_available() was False for want of a package nobody has, and every integration fell to the CPU. _torch.py adds the missing lane without touching the vendored _core/_sparse (which stay diffable against upstream de-csb): a CSB event word IS its intra-block raster index, so the whole accumulate is a scatter-add of ones at block_id * tile_size + word — no custom kernel, just index_add_. Measured on the real 8192², 125.8M-event movie against numba, bit-identical both ways: one plane (what the slider integrates) 487 ms -> 156 ms 3.1x whole movie (125.8M events) 544 ms -> 624 ms 0.9x The per-plane win is the one that matters; the whole-movie case is transfer-bound, not arithmetic-bound, which is also why the words go across PCIe as int32 and are widened to the int64 index_add_ wants on the device (sending int64 cost another 50 ms/plane). numba stays available for whole-file sums via backend="cpu-numba". **Open could not select a .csb.** SUPPORTED_EXTS had it but the Electron file dialogs carry their own extension filters, in two places, and neither listed it — so the format was registered, loadable by path, and unreachable from the menu. Both updated, plus a dedicated "DE CSB (.csb)" filter row. --- electron/src/main/index.ts | 5 +- spyde/external/rsciio_csb/_api.py | 35 ++++++- spyde/external/rsciio_csb/_torch.py | 138 ++++++++++++++++++++++++++++ 3 files changed, 174 insertions(+), 4 deletions(-) create mode 100644 spyde/external/rsciio_csb/_torch.py diff --git a/electron/src/main/index.ts b/electron/src/main/index.ts index 81025e2c..ff3d26e5 100644 --- a/electron/src/main/index.ts +++ b/electron/src/main/index.ts @@ -476,10 +476,11 @@ function buildMenu(): void { const result = await dialog.showOpenDialog(win!, { properties: ['openFile', 'multiSelections'], filters: [ - { name: 'EM Data', extensions: ['hspy', 'zspy', 'mrc', 'tif', 'tiff', 'de5'] }, + { name: 'EM Data', extensions: ['hspy', 'zspy', 'mrc', 'tif', 'tiff', 'de5', 'csb'] }, { name: 'HyperSpy', extensions: ['hspy', 'zspy'] }, { name: 'MRC', extensions: ['mrc'] }, { name: 'TIFF', extensions: ['tif', 'tiff'] }, + { name: 'DE CSB (.csb)', extensions: ['csb'] }, ], }) if (!result.canceled) { @@ -607,7 +608,7 @@ ipcMain.handle('spyde:open-file', async () => { const result = await dialog.showOpenDialog(win!, { properties: ['openFile', 'multiSelections'], filters: [ - { name: 'EM Data', extensions: ['hspy', 'zspy', 'mrc', 'tif', 'tiff', 'de5'] }, + { name: 'EM Data', extensions: ['hspy', 'zspy', 'mrc', 'tif', 'tiff', 'de5', 'csb'] }, ], }) if (!result.canceled) { diff --git a/spyde/external/rsciio_csb/_api.py b/spyde/external/rsciio_csb/_api.py index acb7fdda..031509db 100644 --- a/spyde/external/rsciio_csb/_api.py +++ b/spyde/external/rsciio_csb/_api.py @@ -65,13 +65,44 @@ _DATASETS_LOCK = threading.Lock() +class _TorchSparseCSB(SparseCSB): + """A SparseCSB whose integrations run on torch-CUDA. + + ``_core``'s GPU lane needs CuPy, which SpyDE does not install — its GPU + stack is torch. Subclassing rather than editing ``_sparse`` keeps both + vendored modules diffable against upstream de-csb. + """ + + def _accumulator(self): + if self._acc is None: + from ._torch import TorchAccumulator + self._acc = TorchAccumulator(self.csb) + return self._acc + + def _dataset(path: str, backend: str) -> SparseCSB: - """The process-wide dataset for one file, opened once.""" + """The process-wide dataset for one file, opened once. + + ``backend="auto"`` prefers torch-CUDA when there is a device, because that + is the GPU SpyDE actually ships; it falls back to ``_core``'s own auto + (CuPy if somehow present, else numba, else numpy). + """ key = (os.path.abspath(path), backend) with _DATASETS_LOCK: ds = _DATASETS.get(key) if ds is None: - ds = SparseCSB(CSBFile(path), backend=backend) + csb = CSBFile(path) + use_torch = backend == "torch" + if backend == "auto": + from ._core import gpu_available + from ._torch import torch_gpu_available + use_torch = not gpu_available() and torch_gpu_available() + if use_torch: + ds = _TorchSparseCSB(csb, backend=backend) + _logger.debug("CSB %s: integrating on torch-CUDA", + os.path.basename(path)) + else: + ds = SparseCSB(csb, backend=backend) _DATASETS[key] = ds return ds diff --git a/spyde/external/rsciio_csb/_torch.py b/spyde/external/rsciio_csb/_torch.py new file mode 100644 index 00000000..5f479f0a --- /dev/null +++ b/spyde/external/rsciio_csb/_torch.py @@ -0,0 +1,138 @@ +"""A torch-CUDA accumulator for CSB, alongside the vendored cupy one. + +``_core``'s GPU path is a CuPy ``RawModule`` kernel. SpyDE's GPU stack is +**torch** — it is what the fitting engine, the vector-orientation fit and the +EBSD indexing all run on, and it is what the installer pins with CUDA — so on a +SpyDE machine ``gpu_available()`` is False for want of a package nobody +installed, and a 125 M-event movie integrates on the CPU at ~0.7 s per plane +when the GPU could do it in milliseconds. + +This module adds the missing lane without touching ``_core``/``_sparse``, which +are vendored verbatim from de-csb and want to stay diffable against upstream. + +The kernel, in one line +----------------------- +A CSB event word IS its intra-block raster index, and block ``b``'s events +belong in ``acc[b*tile_size : (b+1)*tile_size]``. So the whole accumulate is a +scatter-add of ones at ``block_id * tile_size + word`` — no custom kernel +needed, just ``index_add_``, which is exactly the shape CLAUDE.md's GPU notes +say to reach for instead of a per-item loop. + +Memory is the one thing to be careful about. ``torch.bincount`` over the full +accumulator would allocate a fresh int64 the size of the frame (536 MB at +8192²) on every call; ``index_add_`` into a persistent int32 accumulator +allocates only ``ones`` at the EVENT count instead (~10 MB for one plane). +""" +from __future__ import annotations + +import logging + +import numpy as np + +log = logging.getLogger(__name__) + + +def torch_gpu_available() -> bool: + """True when torch can see a CUDA device.""" + try: + import torch + return bool(torch.cuda.is_available()) + except Exception as e: + log.debug("torch CUDA unavailable: %s", e) + return False + + +class TorchAccumulator: + """Duck-compatible with :class:`.._core.CSBAccumulator` for the summation + path (``reset`` / ``add_frames`` / ``synchronize`` / ``image``).""" + + def __init__(self, csb, device=None, dtype=None): + import torch + + self.csb = csb + self.device = torch.device(device or "cuda") + self.dtype = dtype or torch.int32 + n = int(csb.blocks_per_frame) * int(csb.tile_size) + self._acc = torch.zeros(n, dtype=self.dtype, device=self.device) + # Block index within a frame for each (frame, block) segment, in the + # payload's readout order. Uploaded once and reused every call. + self._seg_block = torch.arange(int(csb.blocks_per_frame), + dtype=torch.int64, device=self.device) + self._counts = torch.as_tensor(np.asarray(csb.counts, np.int64), + device=self.device) + + # -- accumulate -------------------------------------------------------- + def reset(self) -> None: + self._acc.zero_() + + def add_frames(self, f0: int, f1: int, payload=None, stream=None) -> None: + """Histogram the events of frames ``[f0, f1)`` into the accumulator.""" + import torch + + cs = self.csb + bpf = int(cs.blocks_per_frame) + n_seg = (int(f1) - int(f0)) * bpf + if n_seg <= 0: + return + seg0 = int(f0) * bpf + counts = self._counts[seg0:seg0 + n_seg] + n_ev = int(counts.sum().item()) + if n_ev == 0: + return + + # The payload words for this window are contiguous: `starts` are + # absolute file offsets and the frames are consecutive. `cs.events` is + # the memory-mapped None: + import torch + if self.device.type == "cuda": + torch.cuda.synchronize(self.device) + + # -- read out ---------------------------------------------------------- + def _untile(self): + """Block-major accumulator -> padded 2-D image, still on device. + + Mirrors ``_core.CSBAccumulator._untile`` exactly; the two must agree or + the same file reads differently depending on which backend answered. + """ + cs = self.csb + bh, bw = cs.csb_block_height, cs.csb_block_width + if cs.csb_order == 1: # row-major + v = self._acc.reshape(cs.blocks_per_height, cs.blocks_per_width, bh, bw) + v = v.permute(0, 2, 1, 3) + else: # column-major + v = self._acc.reshape(cs.blocks_per_width, cs.blocks_per_height, bh, bw) + v = v.permute(1, 2, 0, 3) + return v.reshape(cs.padded_height, cs.padded_width) + + def image(self, dtype=None) -> np.ndarray: + cs = self.csb + v = self._untile()[:cs.frame_height, :cs.frame_width] + out = v.contiguous().cpu().numpy() + return out if dtype is None else out.astype(dtype) + + def preview(self, bin_factor: int = 4, dtype=np.float32) -> np.ndarray: + if bin_factor < 1: + raise ValueError("bin_factor must be >= 1") + cs = self.csb + h = (cs.frame_height // bin_factor) * bin_factor + w = (cs.frame_width // bin_factor) * bin_factor + v = self._untile()[:h, :w].contiguous() + v = v.reshape(h // bin_factor, bin_factor, w // bin_factor, bin_factor) + return v.sum(dim=(1, 3), dtype=__import__("torch").int64).cpu().numpy( + ).astype(dtype) From e1677925dc07d088eb94dc34871e76add443d888 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Tue, 28 Jul 2026 12:22:04 -0500 Subject: [PATCH 03/14] perf(csb): pin the readback; verify the whole path in the app MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Drove the real 8192², 125.8M-event movie through the app (csb_movie.spec.ts): it opens as an insitu movie with the free block-table navigator, the time slider integrates a fresh window per position, and To Frames adds a second dataset to the tree ("16 planes at 10 ms (100 fps)"). Profiling a plane says the events were never the problem: memmap read + host cast + H2D + repeat_interleave + idx + index_add_ ~10 ms image() — the D2H of the 8192² result ~212 ms 2.5M events accumulate at ~250 M ev/s; shipping 268 MB back over PCIe is 95% of the wall clock. So the copy now goes through a reusable pinned buffer (212 -> 168 ms; a pageable D2H bounces through a driver staging buffer), and the docstring points at preview() for anything that only needs looking at — at bin 4 that is 1/16th of the bytes and 36 ms/plane against 179. Still bit-identical to numba, and 3.6x on the per-plane window the slider uses. The remaining full-res cost is the device-side untile copy plus the transfer itself; a 268 MB plane has a hard floor here and the way past it is to stop moving whole planes (a TileBackend over the accumulator), not micro-tuning. --- electron/tests/csb_movie.spec.ts | 100 ++++++++++++++++++++++++++++ spyde/external/rsciio_csb/_torch.py | 28 +++++++- 2 files changed, 126 insertions(+), 2 deletions(-) create mode 100644 electron/tests/csb_movie.spec.ts diff --git a/electron/tests/csb_movie.spec.ts b/electron/tests/csb_movie.spec.ts new file mode 100644 index 00000000..243009fa --- /dev/null +++ b/electron/tests/csb_movie.spec.ts @@ -0,0 +1,100 @@ +/** + * csb_movie.spec.ts — a real Direct Electron CSB event stream, in the app. + * + * Opens the 8192², 400-frame, 125.8M-event test movie, checks it arrives as a + * scrubbable in-situ movie, drags the time slider (which integrates a fresh + * window per position — the whole point of the format) and runs To Frames. + * + * Skips itself when the file is absent, so this is not a CI gate; it is the + * "look at the pixels" check CLAUDE.md requires for anything UI. + */ +import { existsSync } from 'fs' +import { test, expect } from '@playwright/test' +const { + launchApp, backendAction, waitForSubwindowCount, countColorPixels, sigWindow, + navWindow, +} = require('./_harness.cjs') + +const CSB = 'C:\\Users\\CarterFrancis\\Downloads' + + '\\directelectron_csb-data-for-testing-de_csb-py_2026-07-22_1739' + + '\\20240604_00001_ces_movie_234.csb' +const SHOTS = 'csb_shots' + +let ctx: Awaited> + +test.skip(!existsSync(CSB), 'CSB test movie not present on this machine') +test.setTimeout(600_000) + +test.beforeAll(async () => { + ctx = await launchApp({ dask: true, env: { SPYDE_LOG_LEVEL: 'INFO' } }) + await backendAction(ctx.page, 'open_file', { path: CSB }) + // 8192² planes and a first-touch numba/CUDA warm-up — give it room. + await waitForSubwindowCount(ctx.page, 2, 300_000) +}) + +test.afterAll(async () => { + await ctx?.app?.close() +}) + +test('a CSB stream opens as a scrubbable in-situ movie', async () => { + const { page } = ctx + await page.screenshot({ path: `${SHOTS}/01-opened.png` }) + + // The reader names the nav axis "time" in seconds, which is what tags it + // insitu — and that is what turns on Play / Fast-Forward. + await expect(navWindow(page)).toBeVisible({ timeout: 60_000 }) + await expect(sigWindow(page)).toBeVisible() + await expect.poll(() => countColorPixels(page, 'bright'), { + timeout: 240_000, message: 'the integrated plane never painted', + }).toBeGreaterThan(0) + await sigWindow(page).screenshot({ path: `${SHOTS}/02-plane.png` }) + ctx.assertNoJsErrors() +}) + +test('dragging the time slider integrates a new window', async () => { + const { page } = ctx + const nav = navWindow(page) + const box = await nav.boundingBox() + if (!box) throw new Error('no navigator window') + + await page.screenshot({ path: `${SHOTS}/03-before-scrub.png` }) + // Drag along the 1-D time navigator: each position is a different exposure + // window, so the displayed image must change. + const y = box.y + box.height * 0.55 + await page.mouse.move(box.x + box.width * 0.25, y) + await page.mouse.down() + await page.mouse.move(box.x + box.width * 0.75, y, { steps: 12 }) + await page.mouse.up() + + await expect.poll(() => countColorPixels(page, 'bright'), { + timeout: 240_000, message: 'nothing painted after the scrub', + }).toBeGreaterThan(0) + await page.screenshot({ path: `${SHOTS}/04-after-scrub.png` }) + await sigWindow(page).screenshot({ path: `${SHOTS}/05-scrubbed-plane.png` }) + ctx.assertNoJsErrors() +}) + +test('To Frames adds an in-situ dataset to the signal tree', async () => { + const { page } = ctx + const sig = sigWindow(page) + await sig.getByTestId('subwindow-title').click() + await sig.getByTestId('subwindow-titlebar').hover() + + const button = sig.getByTestId('action-btn-To Frames') + await expect(button, 'the To Frames action never appeared — the ' + + 'requires_original_metadata gate did not match').toBeVisible({ timeout: 30_000 }) + await page.screenshot({ path: `${SHOTS}/06-toolbar.png` }) + + const before = await page.getByTestId('subwindow').count() + await button.click() + await page.screenshot({ path: `${SHOTS}/07-to-frames-params.png` }) + // The param popout's Run. + const run = page.getByRole('button', { name: /^Run$/ }).first() + if (await run.isVisible().catch(() => false)) await run.click() + + await expect.poll(() => page.getByTestId('subwindow').count(), { + timeout: 300_000, message: 'To Frames never opened a new dataset', + }).toBeGreaterThan(before) + await page.screenshot({ path: `${SHOTS}/08-to-frames-result.png` }) + ctx.assertNoJsErrors() +}) diff --git a/spyde/external/rsciio_csb/_torch.py b/spyde/external/rsciio_csb/_torch.py index 5f479f0a..8207388e 100644 --- a/spyde/external/rsciio_csb/_torch.py +++ b/spyde/external/rsciio_csb/_torch.py @@ -121,11 +121,35 @@ def _untile(self): return v.reshape(cs.padded_height, cs.padded_width) def image(self, dtype=None) -> np.ndarray: + """The summed image as a host ndarray. + + THIS is where a plane's time goes, not the events: accumulating 2.5 M + events takes ~10 ms, while shipping the 8192² result back over PCIe + takes ~200 ms. So the copy is done once, into pinned memory — a + pageable D2H has to bounce through a driver staging buffer, which on + this box costs roughly 2x. Use :meth:`preview` for anything that only + needs to be looked at; at bin 4 it moves 1/16th of the bytes. + """ + import torch + cs = self.csb - v = self._untile()[:cs.frame_height, :cs.frame_width] - out = v.contiguous().cpu().numpy() + v = self._untile()[:cs.frame_height, :cs.frame_width].contiguous() + host = self._pinned_like(v) + host.copy_(v, non_blocking=False) + out = host.numpy() return out if dtype is None else out.astype(dtype) + def _pinned_like(self, v): + """A reusable pinned host buffer matching *v* — allocating (and + page-locking) 268 MB per plane would cost more than the copy saves.""" + import torch + + buf = getattr(self, "_pin", None) + if buf is None or buf.shape != v.shape or buf.dtype != v.dtype: + buf = torch.empty(v.shape, dtype=v.dtype, pin_memory=True) + self._pin = buf + return buf + def preview(self, bin_factor: int = 4, dtype=np.float32) -> np.ndarray: if bin_factor < 1: raise ValueError("bin_factor must be >= 1") From 67a86b32d73abd9cadf8e5c3d72876bf440c4b16 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Tue, 28 Jul 2026 12:49:28 -0500 Subject: [PATCH 04/14] =?UTF-8?q?perf(csb):=20bin=20on=20the=20device=20an?= =?UTF-8?q?d=20cache=20planes=20=E2=80=94=20scrubbing=20354=20->=2037=20ms?= =?UTF-8?q?/step?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit de-csb's README already answered this, under "Readout is a separate cost — don't count it against real-time": pulling the full 8192² image back caps you at ~11 Hz while the 4x-binned preview runs at ~84 Hz, so display from the binned one and call image() only when something is measured. My own profile agreed independently — accumulating 2.5M events is ~10 ms, the readback ~200. So a CuPy kernel is the wrong lever. The repo's table has it at 2,230 M ev/s against this torch path's ~250 M, which sounds decisive until you notice it is 9x on the 10 ms, i.e. ~5% of a plane. Left alone; `auto` already prefers _core's CuPy kernel if anyone installs it. What actually moved the needle, on the real 8192² movie: bin on the DEVICE _sparse's _sum_frames pulls the full image and bins on the host, so asking for bin=4 cost MORE than bin=1 — the whole 268 MB transfer plus a host bin. _TorchSparseCSB overrides it to use preview(), which bins before the copy: 1/16th of the bytes. auto bin default keeps a plane's longest edge at 2048, so an 8192² movie scrubs 2048² planes. bin=1 still available. plane cache 512 MB LRU keyed by the same tuple the graph carries. A drag walks back and forth over the same planes and each repeat was paying the readback again. first visit 330 -> 57 ms revisit 177 -> 22 ms 24-step drag 354 -> 37 ms/step (~27 fps) Bit-identical to numba at bin 1, 2 and 4, counts conserved. Binning is a SUM, and on data this sparse (~1.9 counts/px for the ENTIRE movie) it is also what makes a plane visible at all — the bin=1 view is near-black speckle, the binned one shows the particles. --- spyde/actions/csb_to_frames.py | 6 +- spyde/external/rsciio_csb/_api.py | 121 ++++++++++++++++++++++++++++-- 2 files changed, 117 insertions(+), 10 deletions(-) diff --git a/spyde/actions/csb_to_frames.py b/spyde/actions/csb_to_frames.py index 49aecfdc..b9340709 100644 --- a/spyde/actions/csb_to_frames.py +++ b/spyde/actions/csb_to_frames.py @@ -28,7 +28,7 @@ log = logging.getLogger(__name__) -DEFAULTS = dict(fps=0.0, exposure_ms=0.0, frames_per_plane=0, bin=1) +DEFAULTS = dict(fps=0.0, exposure_ms=0.0, frames_per_plane=0, bin=0) def _source_path(signal) -> str | None: @@ -102,12 +102,12 @@ def _rebuild(session, tree, src, path: str, params) -> None: """Re-cut the stream at a new exposure and add it to the tree.""" import hyperspy.api as hs from spyde.external.rsciio_csb._api import ( - _dataset, _axes, lazy_stack, plane_counts, + _dataset, _axes, _resolve_bin, lazy_stack, plane_counts, ) backend = str(params.get("backend") or "auto") - bin_factor = max(1, int(params.get("bin") or 1)) ds = _dataset(path, backend) + bin_factor = _resolve_bin(ds, params.get("bin")) exposure = _resolve_exposure(ds, params) bounds = ds._time_bounds(exposure, 0.0, ds.duration) if not bounds: diff --git a/spyde/external/rsciio_csb/_api.py b/spyde/external/rsciio_csb/_api.py index 031509db..a5b68f61 100644 --- a/spyde/external/rsciio_csb/_api.py +++ b/spyde/external/rsciio_csb/_api.py @@ -33,6 +33,7 @@ import logging import os import threading +from collections import OrderedDict import numpy as np @@ -46,6 +47,18 @@ #: — a single frame of a sparse stream is mostly empty pixels. DEFAULT_PLANES = 50 +#: Longest plane edge to integrate at by default. A CSB frame can be 8192², and +#: at float32 that is a 268 MB readback per plane — which IS the cost of a +#: scrub (accumulating the events is ~10 ms of it). de-csb's own README says +#: the same thing under "Readout is a separate cost": the full image caps you +#: at ~11 Hz while a 4x-binned one runs at ~84 Hz, so display from the binned +#: one and pull the full image only when something is actually measured. +#: +#: Binning is a SUM, so no counts are lost — and on data this sparse (the test +#: movie is ~1.9 counts/px for the entire 125.8M-event movie) binning is what +#: makes a plane visible at all. Pass ``bin=1`` for the full grid. +DEFAULT_MAX_EDGE = 2048 + # One integration at a time (see the module docstring). _ACC_LOCK = threading.Lock() @@ -79,6 +92,23 @@ def _accumulator(self): self._acc = TorchAccumulator(self.csb) return self._acc + def _sum_frames(self, f0, f1, bin_factor, dtype): + """Integrate, binning ON THE DEVICE before the readback. + + ``_sparse``'s version pulls the full image and bins on the host, so + asking for bin=4 costs MORE than bin=1 — it pays the full 268 MB + transfer and then bins it. Binning first is the entire reason + ``preview()`` exists: it moves 1/16th of the bytes, which is where a + scrub's time actually goes. + """ + acc = self._accumulator() + acc.reset() + acc.add_frames(int(f0), int(f1)) + acc.synchronize() + if bin_factor and int(bin_factor) > 1: + return acc.preview(int(bin_factor), dtype) + return acc.image(dtype) + def _dataset(path: str, backend: str) -> SparseCSB: """The process-wide dataset for one file, opened once. @@ -127,6 +157,23 @@ def _resolve_step(ds: SparseCSB, step, frames_per_plane) -> float: return max(ds.duration / DEFAULT_PLANES, dt) +def _resolve_bin(ds: SparseCSB, requested) -> int: + """Spatial binning per plane. ``None``/0 picks one for the frame size. + + The default keeps a plane's longest edge at or under + :data:`DEFAULT_MAX_EDGE` by the next power of two, so an 8192² movie + integrates 2048² planes (16 MB, ~36 ms) instead of 8192² ones (268 MB, + ~180 ms). An explicit ``bin`` — including ``bin=1`` — is always honoured. + """ + if requested: + return max(1, int(requested)) + edge = max(ds.shape) + factor = 1 + while edge // factor > DEFAULT_MAX_EDGE and factor < 16: + factor *= 2 + return factor + + def plane_counts(ds: SparseCSB, bounds) -> np.ndarray: """Total events per plane, straight from the block table. @@ -139,18 +186,75 @@ def plane_counts(ds: SparseCSB, bounds) -> np.ndarray: return np.array([per_frame[f0:f1].sum() for f0, f1 in bounds], np.float64) +#: Bytes of integrated planes to keep. Scrubbing revisits planes constantly — +#: a drag goes back and forth over the same handful — and re-integrating one is +#: a fresh readback, which is the whole cost (see integrate_plane). Sized to +#: hold a decent run of binned planes (16 MB each at 2048²) or a couple of +#: full-resolution ones. +PLANE_CACHE_BYTES = 512 << 20 + +_plane_cache: "OrderedDict[tuple, np.ndarray]" = OrderedDict() +_plane_cache_bytes = 0 +_PLANE_CACHE_LOCK = threading.Lock() + + +def _cache_get(key): + with _PLANE_CACHE_LOCK: + hit = _plane_cache.get(key) + if hit is not None: + _plane_cache.move_to_end(key) # LRU + return hit + + +def _cache_put(key, arr: np.ndarray) -> None: + global _plane_cache_bytes + nbytes = int(arr.nbytes) + if nbytes > PLANE_CACHE_BYTES: + return # one plane over budget + with _PLANE_CACHE_LOCK: + if key in _plane_cache: + return + _plane_cache[key] = arr + _plane_cache_bytes += nbytes + while _plane_cache_bytes > PLANE_CACHE_BYTES and len(_plane_cache) > 1: + _k, old = _plane_cache.popitem(last=False) + _plane_cache_bytes -= int(old.nbytes) + + +def clear_plane_cache() -> None: + """Drop every cached plane (a different exposure invalidates nothing — + the key carries the window — so this is only for tests and teardown).""" + global _plane_cache_bytes + with _PLANE_CACHE_LOCK: + _plane_cache.clear() + _plane_cache_bytes = 0 + + def integrate_plane(path: str, backend: str, f0: int, f1: int, bin_factor: int, dtype) -> np.ndarray: """One plane: integrate frames ``[f0, f1)`` -> ``(1, H, W)``. Takes the file PATH, not an open dataset — see :data:`_DATASETS`. Every - argument is small and hashable, so the graph is cheap to build and can be - shipped to a worker. + argument is small and hashable, so the graph is cheap to build, it can be + shipped to a worker, and it doubles as the cache key. + + Cached because a scrub is not a linear pass: dragging the slider walks back + and forth over the same few planes, and each repeat would otherwise pay the + full readback again. Accumulating the events is ~10 ms; pulling the result + back is the rest. """ + key = (os.path.abspath(path), backend, int(f0), int(f1), int(bin_factor), + np.dtype(dtype).str) + hit = _cache_get(key) + if hit is not None: + return hit + ds = _dataset(path, backend) with _ACC_LOCK: img = ds._sum_frames(f0, f1, bin_factor, dtype) - return np.asarray(img, dtype)[None] + out = np.asarray(img, dtype)[None] + _cache_put(key, out) + return out def lazy_stack(path: str, bounds, shape, bin_factor: int = 1, @@ -194,7 +298,7 @@ def _axes(ds: SparseCSB, bounds, bin_factor: int): def file_reader(filename, lazy=False, step=None, frames_per_plane=None, - start=0.0, stop=None, bin=1, backend="auto", **kwds): + start=0.0, stop=None, bin=None, backend="auto", **kwds): """Read a CSB centroid stream as a stack of integrated time planes. Parameters @@ -211,8 +315,11 @@ def file_reader(filename, lazy=False, step=None, frames_per_plane=None, Takes precedence over *step*. start, stop : float Time range to cover, in seconds. Defaults to the whole movie. - bin : int - Integer spatial binning, summed. + bin : int, optional + Integer spatial binning, summed. Defaults to whatever keeps a plane's + longest edge at :data:`DEFAULT_MAX_EDGE` — see :func:`_resolve_bin`, + and note that binning is what makes a scrub interactive AND what makes + sparse counted data visible. Pass ``1`` for the full pixel grid. backend : {"auto", "gpu", "cpu", "cpu-numba", "cpu-numpy"} Integration backend; ``"auto"`` takes the GPU when there is one. """ @@ -220,7 +327,7 @@ def file_reader(filename, lazy=False, step=None, frames_per_plane=None, _logger.debug("CSB reader ignoring unknown kwargs: %s", sorted(kwds)) path = str(filename) ds = _dataset(path, backend) - bin_factor = max(1, int(bin)) + bin_factor = _resolve_bin(ds, bin) exposure = _resolve_step(ds, step, frames_per_plane) bounds = ds._time_bounds(exposure, float(start), ds.duration if stop is None else float(stop)) From 4dbaa52eb923e7762ba365e462fc20769a327eed Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Tue, 28 Jul 2026 13:20:08 -0500 Subject: [PATCH 05/14] =?UTF-8?q?fix(csb):=20calibrate=20the=20supplied=20?= =?UTF-8?q?navigator=20=E2=80=94=20every=20scrub=20position=20hit=20plane?= =?UTF-8?q?=200?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The navigator I hand to _add_signal was a bare Signal1D of the block-table counts, with no axis calibration. But a 1-D selector turns the widget position into an index with (x - offset) / scale read from THE NAVIGATOR PLOT'S OWN signal axis (selector1d._signal_axis) — so the widget's x in seconds was being divided by a scale of 1.0, and every position on a 0.156 s movie resolved to index 0. Inspected directly: SIGNAL nav axis : time, 50, scale 0.00312, units s NAVIGATOR sig ax: , 50, scale 1.0, The navigator now carries the source signal's navigation calibration, in both places one is built (the reader path and To Frames). The spec that was supposed to cover this was worthless three times over, which is worth recording because each failure looked green: * counting bright pixels page-wide never changes with the plane, and the navigator's own canvas keeps it non-zero regardless; * reading canvas pixels with getImageData returns all zeros — the image is on a WebGPU canvas (the trap the vectors-embed work hit). It "failed" without the fix for the wrong reason: it read 0:0 either way; * page.mouse.click on the navigator moves the CURSOR, not the widget — the readout showed x:0.0297 while the selector marker sat at 0. It now hashes a screenshot of the signal window (composited truth) and moves via test_nav_drag, and it does prove the per-plane read/paint path tracks. HONEST LIMIT: test_nav_drag sets indices directly, so it does NOT exercise the selector's axis conversion and would not have caught this bug. Driving the VLine widget from a spec has no hook today (insitu_playback works around the same gap); the calibration is verified by inspection, not by test. --- electron/tests/csb_movie.spec.ts | 62 ++++++++++++++++++++++++++------ spyde/actions/csb_to_frames.py | 7 ++-- spyde/backend/_session_files.py | 29 ++++++++++++--- 3 files changed, 80 insertions(+), 18 deletions(-) diff --git a/electron/tests/csb_movie.spec.ts b/electron/tests/csb_movie.spec.ts index 243009fa..75408410 100644 --- a/electron/tests/csb_movie.spec.ts +++ b/electron/tests/csb_movie.spec.ts @@ -9,6 +9,7 @@ * "look at the pixels" check CLAUDE.md requires for anything UI. */ import { existsSync } from 'fs' +import { createHash } from 'crypto' import { test, expect } from '@playwright/test' const { launchApp, backendAction, waitForSubwindowCount, countColorPixels, sigWindow, @@ -51,26 +52,65 @@ test('a CSB stream opens as a scrubbable in-situ movie', async () => { ctx.assertNoJsErrors() }) -test('dragging the time slider integrates a new window', async () => { +/** A cheap signature of what the signal window is actually showing. + * + * Counting bright pixels page-wide is NOT enough — it stays non-zero no + * matter which plane is displayed (and the navigator's own canvas keeps it + * non-zero regardless), which is exactly how "every scrub position showed + * plane 0" passed a green test. This samples the displayed pixels so two + * different planes cannot produce the same number. */ +async function planeSignature(page): Promise { + // A hash of the SIGNAL window as composited. + // + // Two cheaper things were tried and both lied. Counting bright pixels + // page-wide never changes with the plane. Reading canvas pixels via + // getImageData returns all zeros here, because the image is drawn on a + // WebGPU canvas — the same trap the vectors-embed work hit, where buffer + // asserts "pass" against an overlay canvas that holds nothing. A screenshot + // is the composited truth, which is what "the image doesn't move" means. + try { + const buf = await sigWindow(page).screenshot() + return createHash('sha1').update(buf).digest('hex') + } catch { + return '' + } +} + +test('dragging the time slider shows a DIFFERENT integrated plane', async () => { const { page } = ctx const nav = navWindow(page) const box = await nav.boundingBox() if (!box) throw new Error('no navigator window') await page.screenshot({ path: `${SHOTS}/03-before-scrub.png` }) - // Drag along the 1-D time navigator: each position is a different exposure - // window, so the displayed image must change. - const y = box.y + box.height * 0.55 - await page.mouse.move(box.x + box.width * 0.25, y) - await page.mouse.down() - await page.mouse.move(box.x + box.width * 0.75, y, { steps: 12 }) - await page.mouse.up() + const before = await planeSignature(page) + expect(before, 'no plane was on screen to begin with').not.toBe('') - await expect.poll(() => countColorPixels(page, 'bright'), { - timeout: 240_000, message: 'nothing painted after the scrub', - }).toBeGreaterThan(0) + // Move the navigator with test_nav_drag, not the mouse. Clicking the plot + // moves the CURSOR, not the VLine widget — the readout showed x:0.0297 while + // the selector marker sat at 0, so a mouse-driven check "passed" on an + // unrelated repaint. This drives the same path a real drag ends in. + const goTo = async (plane: number) => { + await backendAction(page, 'test_nav_drag', { targets: [[plane, 0]] }) + } + await goTo(30) + await expect.poll(() => planeSignature(page), { + timeout: 60_000, + message: 'the displayed plane did not change when the slider moved — every ' + + 'position is resolving to the same index', + }).not.toBe(before) + + const moved = await planeSignature(page) await page.screenshot({ path: `${SHOTS}/04-after-scrub.png` }) await sigWindow(page).screenshot({ path: `${SHOTS}/05-scrubbed-plane.png` }) + + // And a third position differs from the second, so it is tracking rather + // than flipping between two cached frames. + await goTo(12) + await expect.poll(() => planeSignature(page), { + timeout: 60_000, message: "the plane stopped tracking after one move", + }).not.toBe(moved) + await sigWindow(page).screenshot({ path: `${SHOTS}/05b-third-position.png` }) ctx.assertNoJsErrors() }) diff --git a/spyde/actions/csb_to_frames.py b/spyde/actions/csb_to_frames.py index b9340709..bdde399f 100644 --- a/spyde/actions/csb_to_frames.py +++ b/spyde/actions/csb_to_frames.py @@ -147,8 +147,11 @@ def _rebuild(session, tree, src, path: str, params) -> None: # The free navigator again — re-cutting the exposure changes the planes, so # the overview has to be recomputed, but it still costs no payload reads. - nav = hs.signals.Signal1D(np.asarray(plane_counts(ds, bounds), np.float32)) - nav.metadata.General.title = "Total counts" + # Calibrated from the NEW signal: a 1-D selector derives its index from the + # navigator's own axis, so an uncalibrated one pins every scrub position to + # plane 0. + from spyde.backend._session_files import calibrated_nav_signal + nav = calibrated_nav_signal(plane_counts(ds, bounds), sig) def _add(): session._add_signal(sig, source_path=path, navigator_override=nav) diff --git a/spyde/backend/_session_files.py b/spyde/backend/_session_files.py index 682572f9..7fdd1e6b 100644 --- a/spyde/backend/_session_files.py +++ b/spyde/backend/_session_files.py @@ -93,16 +93,35 @@ def _reader_navigator(sig): if counts is None or not len(counts): return None try: - import hyperspy.api as hs - import numpy as _np - nav = hs.signals.Signal1D(_np.asarray(counts, _np.float32)) - nav.metadata.General.title = "Total counts" - return nav + return calibrated_nav_signal(counts, sig) except Exception as e: log.debug("building the reader-supplied navigator failed: %s", e) return None +def calibrated_nav_signal(counts, sig): + """A 1-D navigator carrying *sig*'s NAVIGATION calibration. + + The calibration is not cosmetic. A 1-D selector turns the widget's position + into an index with ``(x - offset) / scale`` read from **the navigator + plot's own signal axis** (``selector1d._signal_axis``) — so an uncalibrated + navigator means dividing a position in seconds by a scale of 1, and every + position on a 0.156 s movie resolves to index 0. The image then never + changes as you scrub, which looks like a broken reader rather than a + missing axis. + """ + import hyperspy.api as hs + import numpy as _np + + nav = hs.signals.Signal1D(_np.asarray(counts, _np.float32)) + src = sig.axes_manager.navigation_axes[0] + dst = nav.axes_manager.signal_axes[0] + dst.name, dst.units = src.name, src.units + dst.scale, dst.offset = float(src.scale), float(src.offset) + nav.metadata.General.title = "Total counts" + return nav + + _DEFAULT_EXAMPLE_NAMES = ( "mgo_nanocrystals", "small_ptychography", From 98a43d0e27a14ca8070b697cdd98cddcfbde0cb5 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Tue, 28 Jul 2026 13:43:00 -0500 Subject: [PATCH 06/14] feat(csb): make the raw camera frame reachable from To Frames MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Integrating by default is right — a raw CSB frame is 390 us and ~0.005 counts/px, so on its own it is very nearly an empty image. But that is exactly why it has to be reachable: "the default is a sum" is only a safe default if you can still see what was summed. `frames_per_plane` was already the knob (it takes precedence over fps and exposure in _resolve_exposure) but it was a hs.load kwarg only, so from the UI there was no way to ask for one. Now a To Frames parameter: 1 gives one plane per raw frame — 400 planes of 390 us here — and 0 keeps the fps/exposure fields in charge. Verified against the real movie: default 50 planes, 3.12 ms, 8 raw frames each 0.603 counts/px (bin 4) raw 400 planes, 0.39 ms, 1 raw frame each 0.075 counts/px The status line now reports the raw-frame count alongside the exposure — "16 planes at 10 ms (100 fps), 32 raw frames each", or "1 raw frame each — no integration", because the time alone does not tell you how much of the stream went into a plane. Also exposed bin as "0 = auto" rather than hiding the auto-binning behind a default of 1 that no longer meant anything. --- spyde/actions/csb_to_frames.py | 11 +++++++++-- spyde/toolbars.yaml | 17 ++++++++++++++--- 2 files changed, 23 insertions(+), 5 deletions(-) diff --git a/spyde/actions/csb_to_frames.py b/spyde/actions/csb_to_frames.py index bdde399f..b8a0a58f 100644 --- a/spyde/actions/csb_to_frames.py +++ b/spyde/actions/csb_to_frames.py @@ -115,8 +115,15 @@ def _rebuild(session, tree, src, path: str, params) -> None: f"a {exposure * 1e3:.4g} ms exposure selects no frames from this " f"{ds.duration * 1e3:.4g} ms movie") + # Report the exposure in RAW FRAMES as well as in time. "16 planes at + # 10 ms" does not tell you how much of the stream each one summed, and at + # 1 frame/plane you are looking at the camera's actual output rather than + # any integration of it — which is a different thing to be looking at. + per_plane = bounds[0][1] - bounds[0][0] + raw = ("1 raw frame each — no integration" if per_plane == 1 + else f"{per_plane} raw frames each") emit_status(f"To Frames: {len(bounds)} planes at " - f"{exposure * 1e3:.4g} ms ({1.0 / exposure:.4g} fps)…") + f"{exposure * 1e3:.4g} ms ({1.0 / exposure:.4g} fps), {raw}…") data = lazy_stack(path, bounds, ds.shape, bin_factor, np.float32, backend) sig = hs.signals.Signal2D(data).as_lazy() @@ -156,7 +163,7 @@ def _rebuild(session, tree, src, path: str, params) -> None: def _add(): session._add_signal(sig, source_path=path, navigator_override=nav) emit_status(f"To Frames: {len(bounds)} planes at " - f"{exposure * 1e3:.4g} ms ({1.0 / exposure:.4g} fps)") + f"{exposure * 1e3:.4g} ms ({1.0 / exposure:.4g} fps), {raw}") dispatch = getattr(session, "_dispatch_to_main", None) if dispatch is not None: diff --git a/spyde/toolbars.yaml b/spyde/toolbars.yaml index 4f1da2a6..2794dfe0 100644 --- a/spyde/toolbars.yaml +++ b/spyde/toolbars.yaml @@ -315,11 +315,22 @@ functions: default: 0.0 min: 0.0 max: 60000.0 + # The raw camera frame is the atomic exposure — 390 us on the test movie, + # ~0.005 counts/px, which is why integrating is the default. But it has to + # be REACHABLE: 1 here gives one plane per raw frame, which is the only + # way to see what the camera actually recorded rather than a sum of it. + # Takes precedence over both fields above (see _resolve_exposure). + frames_per_plane: + name: or Raw frames per plane + type: int + default: 0 + min: 0 + max: 100000 bin: - name: Bin + name: Bin (0 = auto) type: int - default: 1 - min: 1 + default: 0 + min: 0 max: 16 EBSD Indexing: From fd7b1cd47fbdf230ec6f033a8d8f0bed3ecc1ad1 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Tue, 28 Jul 2026 14:02:01 -0500 Subject: [PATCH 07/14] =?UTF-8?q?feat(dock):=20Sum=20frames=20=E2=80=94=20?= =?UTF-8?q?give=20a=20POINT=20navigator=20a=20width,=20everywhere?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A point selector reads one navigation position. This lets it read n and sum them, keeping the single pointer. Distinct from Integrate, which hands you a draggable span with two edges to place: this is for when the exposure is a property of the acquisition rather than a region you are choosing — summing n frames of an in-situ movie for signal, or picking the exposure of a sparse event stream, without turning the slider into a range. Generic, not CSB-specific, per review: any 1-D (movie/time) navigator gets it. A 2-D navigator does not, because "n frames" has no unambiguous direction there and Integrate already covers it — the backend simply omits sum_frames and the dock hides the control. Downstream needs no new path: emitting n index rows is exactly what a region selector emits, so the existing region-integration machinery (array_cache, region_sum's threaded bands and incremental +-1) handles it unchanged. The ladder is powers of two, labelled with the rate each produces when the axis is time ("8 frames — 40 fps"). Frame counts rather than round frame rates because a rate cannot generally divide the cadence — asking for 60 fps on a 2564 fps camera silently rounds to 43, so the number picked is not the number you get. Two things that would have been silent bugs: * the live selector is IntegratingSelector1D, a COMPOSITE. Its __getattr__ delegates reads only, so `composite.sum_frames = 8` bound on the wrapper while the inner line selector kept 1 — the attribute read back as 8 and nothing summed. It is an explicit property now. * the window SLIDES at the ends rather than being clipped. A clipped window folds onto the edge index and sums the same frame several times, quietly making the first and last positions brighter than the rest. --- .../src/components/PlotControlDock.tsx | 44 +++++++++++++++ .../src/renderer/src/kernel/SpyDEContext.tsx | 14 +++++ electron/tests/csb_movie.spec.ts | 25 +++++++++ spyde/backend/_session_actions.py | 3 ++ spyde/backend/session.py | 33 ++++++++++++ spyde/drawing/plots/multiplot_manager.py | 27 +++++++++- spyde/drawing/selectors/selector1d.py | 54 ++++++++++++++++++- 7 files changed, 197 insertions(+), 3 deletions(-) diff --git a/electron/src/renderer/src/components/PlotControlDock.tsx b/electron/src/renderer/src/components/PlotControlDock.tsx index 798bf134..06ccdd22 100644 --- a/electron/src/renderer/src/components/PlotControlDock.tsx +++ b/electron/src/renderer/src/components/PlotControlDock.tsx @@ -647,13 +647,57 @@ export function PlotControlDock() { ))} + {/* Sum N frames under a POINT selector. Distinct from Integrate: + that gives you a draggable span with two edges to place; this + keeps the single pointer and widens what it reads — n frames of + a movie summed for signal, or the exposure of a sparse event + stream, without turning the slider into a range. Only offered + for a 1-D navigator (the backend omits sumFrames otherwise). */} + {navSelectors.filter((s) => s.sumFrames != null && s.mode === 'crosshair') + .map((s) => ( +
+ Sum frames + sendAction('set_selector_sum', + { frames: Number(v), selector_id: s.selectorId }, s.windowId)} + /> +
+ ))} )} ) } +/** The Sum-frames ladder: powers of two up to the navigation length. + * + * Powers of two rather than round frame rates because a rate cannot generally + * divide the acquisition cadence — asking for "60 fps" on a 2564 fps camera + * silently rounds to 43 frames, so the number you picked is not the number + * you got. A frame count is always exact, and the rate it produces is shown + * beside it when the axis is time. */ +function sumFrameOptions(navSize: number, navScale: number) { + const opts: { value: string; label: string }[] = [] + for (let n = 1; n <= Math.max(1, navSize); n *= 2) { + const rate = navScale > 0 ? 1 / (n * navScale) : 0 + const hz = rate >= 1000 ? `${(rate / 1000).toFixed(1)} kfps` + : rate >= 1 ? `${rate.toFixed(0)} fps` + : rate > 0 ? `${rate.toFixed(2)} fps` : '' + const frames = n === 1 ? '1 frame' : `${n} frames` + opts.push({ value: String(n), label: hz ? `${frames} — ${hz}` : frames }) + if (n >= 64) break + } + return opts +} + const styles: Record = { + sumRow: { + display: 'flex', alignItems: 'center', gap: 8, marginTop: 6, + }, + sumLabel: { fontSize: 11.5, color: '#a6adc8', flex: '0 0 auto' }, dock: { width: 300, flexShrink: 0, height: '100%', diff --git a/electron/src/renderer/src/kernel/SpyDEContext.tsx b/electron/src/renderer/src/kernel/SpyDEContext.tsx index c78061b1..53ce3ecf 100644 --- a/electron/src/renderer/src/kernel/SpyDEContext.tsx +++ b/electron/src/renderer/src/kernel/SpyDEContext.tsx @@ -139,6 +139,15 @@ export interface SelectorInfo { selectorId?: number /** Widget colour — the dock row's dot. */ color?: string + /** How many navigation positions a POINT selector sums (1 = plain + * crosshair). Only sent for a 1-D (movie/time) navigator; its absence is + * what hides the control on a 2-D one, where "n frames" has no direction. */ + sumFrames?: number + /** Length of that navigation axis, so the dock can cap the ladder. */ + navSize?: number + /** Seconds per navigation position (0 when the axis isn't time), so a + * summed window can be labelled with the rate it works out to. */ + navScale?: number } /** The named navigators a navigator window offers (its top chip strip). */ export interface NavigatorOptions { names: string[]; current?: string | null } @@ -1111,6 +1120,11 @@ export function SpyDEProvider({ children }: { children: React.ReactNode }) { // creation-time title/colour on a mode-only re-emit. ...(msg.title != null ? { title: msg.title } : {}), ...(msg.color != null ? { color: msg.color } : {}), + ...(msg.sum_frames != null + ? { sumFrames: Number(msg.sum_frames) } : {}), + ...(msg.nav_size != null ? { navSize: Number(msg.nav_size) } : {}), + ...(msg.nav_scale != null + ? { navScale: Number(msg.nav_scale) } : {}), }, }) break diff --git a/electron/tests/csb_movie.spec.ts b/electron/tests/csb_movie.spec.ts index 75408410..cf4c6290 100644 --- a/electron/tests/csb_movie.spec.ts +++ b/electron/tests/csb_movie.spec.ts @@ -114,6 +114,31 @@ test('dragging the time slider shows a DIFFERENT integrated plane', async () => ctx.assertNoJsErrors() }) +test('Sum frames widens the point selector without moving it', async () => { + const { page } = ctx + await backendAction(page, 'test_nav_drag', { targets: [[25, 0]] }) + const one = await planeSignature(page) + + const dd = page.getByTestId('selector-sum-frames') + await expect(dd, 'no Sum-frames control on a 1-D navigator').toBeVisible({ + timeout: 30_000, + }) + await page.screenshot({ path: `${SHOTS}/09-sum-frames-control.png` }) + + // Themed dropdown: click the trigger, then the option (selectOption is dead + // on it — see Dropdown.tsx). + await dd.click() + await page.getByTestId('selector-sum-frames-opt-8').click() + + await expect.poll(() => planeSignature(page), { + timeout: 60_000, + message: 'summing 8 frames did not change the image', + }).not.toBe(one) + await page.screenshot({ path: `${SHOTS}/10-summed-8.png` }) + await sigWindow(page).screenshot({ path: `${SHOTS}/11-summed-plane.png` }) + ctx.assertNoJsErrors() +}) + test('To Frames adds an in-situ dataset to the signal tree', async () => { const { page } = ctx const sig = sigWindow(page) diff --git a/spyde/backend/_session_actions.py b/spyde/backend/_session_actions.py index 4f2bcba1..ad677366 100644 --- a/spyde/backend/_session_actions.py +++ b/spyde/backend/_session_actions.py @@ -164,6 +164,9 @@ def dispatch_action(self, msg: dict) -> None: # dialog) on the active signal, so the E2E workflow can be driven # headlessly / in Playwright. payload={"phase":"si"|"ag"} (default si). self._run_test_orientation(plot, payload) + elif action == "set_selector_sum": + self.set_selector_sum(window_id, int(payload.get("frames", 1)), + payload.get("selector_id")) elif action == "set_selector_mode": self.set_selector_mode(window_id, bool(payload.get("integrate")), payload.get("selector_id")) diff --git a/spyde/backend/session.py b/spyde/backend/session.py index 2d642d87..25717991 100644 --- a/spyde/backend/session.py +++ b/spyde/backend/session.py @@ -465,6 +465,39 @@ def set_selector_mode(self, window_id: int, integrate: bool, except Exception as e: log.warning("set_selector_mode failed: %s", e) + def set_selector_sum(self, window_id: int, frames: int, + selector_id: int | None = None) -> None: + """Set how many navigation positions a POINT selector sums. + + 1 is a plain crosshair. Higher keeps the single pointer and widens what + it reads — n frames of an in-situ movie summed for signal, or the + exposure of a sparse event stream chosen without turning the slider + into a draggable range (which is what Integrate is for). + """ + sel = None + if selector_id is not None: + window_id, sel = getattr(self, "_nav_selectors_by_id", {}).get( + selector_id, (window_id, None)) + if sel is None: + sel = getattr(self, "_nav_selectors", {}).get(window_id) + if sel is None or not hasattr(sel, "sum_frames"): + return + try: + sel.sum_frames = max(1, int(frames)) + # Re-read at the new width immediately; the pointer has not moved, + # so nothing else would trigger it. + if hasattr(sel, "update_data"): + sel.update_data(force=True) + emit({ + "type": "selector_info", + "window_id": window_id, + "selector_id": id(sel), + "color": getattr(sel, "color", None), + "sum_frames": int(sel.sum_frames), + }) + except Exception as e: + log.warning("set_selector_sum failed: %s", e) + def _select_signal_node(self, plot, signal_id) -> None: """Switch to the signal-tree node with the given id (the id(node.signal) emitted in the signal_tree message). The pick can come from ANY of the diff --git a/spyde/drawing/plots/multiplot_manager.py b/spyde/drawing/plots/multiplot_manager.py index 09dd8538..3fa1b9dd 100644 --- a/spyde/drawing/plots/multiplot_manager.py +++ b/spyde/drawing/plots/multiplot_manager.py @@ -231,14 +231,37 @@ def add_navigation_selector_and_signal_plot( idx = len(self.navigation_selectors[plot_window]) self.session.register_nav_selector(plot_window.window_id, selector) from spyde.backend.ipc import emit - emit({ + # `sum_frames` / `nav_size` are only sent for a selector that can + # widen a POINT into a summed window — a 1-D (movie/time) navigator. + # Their absence is what tells the dock not to offer the control on a + # 2-D navigator, where "n frames" has no unambiguous direction and + # Integrate is the right tool anyway. + info = { "type": "selector_info", "window_id": plot_window.window_id, "selector_id": id(selector), "color": color, "mode": mode, "title": "Navigator" if idx <= 1 else f"Navigator {idx}", - }) + } + if hasattr(selector, "sum_frames"): + info["sum_frames"] = int(getattr(selector, "sum_frames", 1) or 1) + try: + ax = (selector.current_plot.plot_state.current_signal + .axes_manager.signal_axes[0]) + info["nav_size"] = int(ax.size) + # Seconds per navigation position, so the dock can label a + # summed window with the effective rate it produces. Only + # meaningful on a time axis; 0 tells it to label in frames + # alone. + units = str(getattr(ax, "units", "") or "").lower() + info["nav_scale"] = (float(ax.scale) + if units in ("s", "sec", "second", + "seconds") else 0.0) + except Exception: + info["nav_size"] = 0 + info["nav_scale"] = 0.0 + emit(info) except Exception as e: logger.debug("emitting navigator selector_info failed: %s", e) diff --git a/spyde/drawing/selectors/selector1d.py b/spyde/drawing/selectors/selector1d.py index 424169d4..c821c648 100644 --- a/spyde/drawing/selectors/selector1d.py +++ b/spyde/drawing/selectors/selector1d.py @@ -71,13 +71,49 @@ def _get_plot1d(self): def _on_pointer_up(self, event): self.update_data() + #: How many consecutive navigation positions this point sums. 1 is a plain + #: crosshair; >1 keeps the single pointer but reads a window around it. + #: + #: This is NOT the same thing as Integrate. Integrate hands you a draggable + #: span with two edges to place; this keeps one pointer and gives it a + #: width, which is what you want when the exposure is a property of the + #: acquisition rather than a region you are choosing — summing n frames of + #: an in-situ movie for signal, or picking the exposure of a sparse event + #: stream (1 raw frame vs 8) without turning the slider into a range. + #: + #: Downstream this needs no special path at all: emitting n index rows is + #: exactly what a region selector emits, so the existing region-integration + #: machinery (array_cache, region_sum's threaded bands and incremental +-1) + #: handles it unchanged. + sum_frames: int = 1 + def _get_selected_indices(self) -> np.ndarray: if self._widget is None: return np.array([[0]]) scale, offset = _signal_axis(self) pos = float(self._widget.x) index = int(round((pos - offset) / scale)) - return np.array([[index]]) + n = max(1, int(getattr(self, "sum_frames", 1) or 1)) + if n == 1: + return np.array([[index]]) + # Centre the window on the pointer, then SLIDE it to stay in range + # rather than letting the clip in _get_selected_indices_and_clip fold + # it — a clipped window would sum the edge frame several times and + # quietly report a brighter first/last position than the rest. + size = self._nav_size() + n = min(n, size) if size else n + start = index - n // 2 + if size: + start = max(0, min(start, size - n)) + return np.arange(start, start + n, dtype=int).reshape(-1, 1) + + def _nav_size(self) -> int: + """Length of the navigation axis this selector indexes, or 0.""" + try: + return int(self.current_plot.plot_state.current_signal + .axes_manager.signal_axes[0].size) + except Exception: + return 0 def translate_pixels(self, shift_x: int) -> None: if self._widget is not None: @@ -381,6 +417,22 @@ def __init__( except Exception as e: logger.debug("pushing 1-D panel overlay state failed: %s", e) + @property + def sum_frames(self) -> int: + """How many positions the POINT sub-selector sums. + + An explicit property, not left to ``__getattr__``: that only delegates + READS of undefined attributes, so a plain ``composite.sum_frames = 8`` + would bind on the composite and the inner line selector — the one that + actually computes the indices — would keep its default of 1. Silently: + the attribute reads back as 8 and nothing sums. + """ + return int(getattr(self._inf_line_selector, "sum_frames", 1) or 1) + + @sum_frames.setter + def sum_frames(self, n) -> None: + self._inf_line_selector.sum_frames = max(1, int(n or 1)) + def __getattr__(self, name): """Delegate undefined attributes to the active sub-selector.""" if name in ("selector", "_inf_line_selector", "_linear_region_selector"): From 4f813fbddea26ac671993c6fbd061afaff1ee8dd Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Tue, 28 Jul 2026 21:48:27 -0500 Subject: [PATCH 08/14] feat(dock): move the frame-width picker onto the Point button MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The picker now hangs off the Point button as a bare caret, and the button carries the width it is reading as a subtle badge — "✛ Point 8f" — so you can see what the pointer is summing without opening anything. Dropdown gained a `triggerText` override for this: the menu keeps the full "8 frames — 40 fps" labels while the trigger shows nothing, because repeating the selection next to a control that already displays it is just noise. Also fixes a swallowed exception in set_selector_sum. It called update_data(force=True), which that method does not accept; the raise happened BEFORE the confirming emit, so the width applied to the selector and the dock never heard about it. The pointer really was summing 8 frames while the button still said 1 — and the test passed anyway, because it only checked that the image changed. It now asserts the button reports the width back. NOT flipping the CSB load default to raw cadence (asked for, conditional on performance — it fails that condition). Summing n positions means n accumulator passes and n readbacks, where pre-integrating them is one of each, and the readback is where the time goes: one pre-integrated 8-frame plane (today) 37.6 ms eight raw planes summed, cold 340.8 ms 9.1x sliding one step, 7 of 8 cached 108.6 ms 2.9x Identical results, so this is purely cost. Raw cadence stays reachable through To Frames -> Raw frames per plane = 1, which pays it only when asked for. --- .../src/renderer/src/components/Dropdown.tsx | 15 +++++- .../src/components/PlotControlDock.tsx | 54 +++++++++++-------- electron/tests/csb_movie.spec.ts | 7 +++ spyde/backend/session.py | 7 ++- 4 files changed, 56 insertions(+), 27 deletions(-) diff --git a/electron/src/renderer/src/components/Dropdown.tsx b/electron/src/renderer/src/components/Dropdown.tsx index 526d72a2..18792552 100644 --- a/electron/src/renderer/src/components/Dropdown.tsx +++ b/electron/src/renderer/src/components/Dropdown.tsx @@ -17,12 +17,19 @@ */ import React from 'react' -export function Dropdown({ value, options, onChange, testid, width }: { +export function Dropdown({ + value, options, onChange, testid, width, triggerText, +}: { value: T options: readonly { value: T; label: string }[] onChange: (v: T) => void testid: string width?: number | string + /** Override what the TRIGGER shows, while the menu keeps the full labels. + * For a dropdown attached to another control, where the selection is + * already displayed beside it and repeating the whole label would just be + * noise — pass '' for a bare caret. */ + triggerText?: string }) { const [open, setOpen] = React.useState(false) // Auto drop-UP when the menu would clip the bottom of the window (e.g. the @@ -63,7 +70,11 @@ export function Dropdown({ value, options, onChange, testid, w style={{ ...S.trigger, ...(open ? S.triggerOpen : {}) }} onClick={toggle} > - {current?.label ?? String(value)} + {triggerText !== '' && ( + + {triggerText ?? current?.label ?? String(value)} + + )} ▾ {open && ( diff --git a/electron/src/renderer/src/components/PlotControlDock.tsx b/electron/src/renderer/src/components/PlotControlDock.tsx index 06ccdd22..5d3f79a8 100644 --- a/electron/src/renderer/src/components/PlotControlDock.tsx +++ b/electron/src/renderer/src/components/PlotControlDock.tsx @@ -631,12 +631,34 @@ export function PlotControlDock() { /> + {/* …and the picker hangs off it, a bare caret since the count is + already on the button. */} + {s.sumFrames != null && s.mode === 'crosshair' && ( + sendAction('set_selector_sum', + { frames: Number(v), selector_id: s.selectorId }, s.windowId)} + /> + )} , 57 x waiting for element to be visible" is exactly that shape, and it is not a timing problem however much it reads like one. * a {open && (
= { textAlign: 'left', }, triggerOpen: { borderColor: '#45475a', background: '#181825' }, + triggerBare: { + background: 'transparent', border: 'none', padding: '0 5px 0 2px', + width: 'auto', + }, triggerLabel: { overflow: 'hidden', textOverflow: 'ellipsis', whiteSpace: 'nowrap' }, caret: { fontSize: 9, color: '#6c7086', flex: '0 0 auto' }, // The panel is a copy of MenuBar's `styles.dropdown` / `styles.item`. diff --git a/electron/src/renderer/src/components/PlotControlDock.tsx b/electron/src/renderer/src/components/PlotControlDock.tsx index 5d3f79a8..493bac4a 100644 --- a/electron/src/renderer/src/components/PlotControlDock.tsx +++ b/electron/src/renderer/src/components/PlotControlDock.tsx @@ -629,36 +629,49 @@ export function PlotControlDock() { style={{ ...styles.selectorDot, background: s.color ?? '#00e676' }} title={s.title ?? 'Navigator'} /> - + {s.sumFrames != null && s.mode === 'crosshair' && ( + sendAction('set_selector_sum', + { frames: Number(v), selector_id: s.selectorId }, s.windowId)} + /> )} - - {/* …and the picker hangs off it, a bare caret since the count is - already on the button. */} - {s.sumFrames != null && s.mode === 'crosshair' && ( - sendAction('set_selector_sum', - { frames: Number(v), selector_id: s.selectorId }, s.windowId)} - /> - )} +