feat(signals): extract RaggedStore from SpyDEDiffractionVectors

spyde/actions/find_vectors_action.py

@@ -456,7 +456,8 @@ def _finalize(tree, vecs) -> None:
         tree.root._clear_cache_dask_data()
     except Exception as e:
         log.debug("clearing stale cached dask array failed: %s", e)
-    tree.diffraction_vectors = vecs
+    from spyde.actions.lifecycle import attach_container
+    attach_container(tree, vecs, name="diffraction_vectors")
 
     # Paint the count map onto the SPATIAL (2-D) navigator plot. For a 5-D stack
     # the navigator is multi-level; _first_nav_plot may return the OUTER (1-D

spyde/actions/lifecycle.py

@@ -90,6 +90,38 @@ def is_current(owner, key: str, gen: int) -> bool:
 
 # ── the find-vectors attach gap ───────────────────────────────────────────────
 
+def attach_container(tree, store, *, name: str):
+    """Attach a ragged result container to *tree* under attribute *name* —
+    THE seam between a batch compute finalizing and the tree carrying its
+    result (``tree.diffraction_vectors`` today; ``tree.particles`` next).
+
+    setattr + provenance stamp: a container carrying no ``provenance`` record
+    of its own inherits the tree's commit provenance (the dict
+    ``commit._stamp_provenance`` stores as ``tree._commit_provenance``), so a
+    saved container self-describes the way a committed tree does. Readers are
+    unchanged: :func:`resolve_vectors`, the ``requires_vectors`` toolbar gate
+    and the Save hook all read the attribute this helper sets.
+
+    Later-PR design (recorded here, deliberately NOT implemented in this PR):
+    toolbar YAML grows a generic ``requires_container: <name>`` key with
+    ``requires_vectors`` kept as its alias, and the ``commit.*`` entry points
+    (``open_result_tree`` / ``commit_result_tree``) grow a ``container=``
+    kwarg routed through this helper — one attach/gate/commit seam for every
+    ragged family instead of each action hand-rolling the setattr.
+
+    Returns *store*.
+    """
+    setattr(tree, name, store)
+    try:
+        if getattr(store, "provenance", None) is None:
+            prov = getattr(tree, "_commit_provenance", None)
+            if prov:
+                store.provenance = dict(prov)
+    except Exception as e:
+        log.debug("stamping container provenance failed: %s", e)
+    return store
+
+
 def resolve_vectors(session, plot):
     """Resolve ``(tree, diffraction_vectors)`` for an action.
 

spyde/signals/diffraction_vectors.py

@@ -2,7 +2,11 @@
 
 import numpy as np
 from dataclasses import dataclass, field
-from typing import List, Optional
+from typing import ClassVar, List, Optional
+
+# _AxisLite lives in ragged_store now; re-exported here for its existing
+# importers (orientation_map, dense_diffraction_vectors, tests, …).
+from spyde.signals.ragged_store import RaggedStore, _AxisLite  # noqa: F401
 
 # Column indices — never use bare integer literals outside this module
 COL_NAV_X     = 0
@@ -29,24 +33,27 @@ def _build_nav_offsets(
 
     The flat_buffer must be sorted outermost-nav-dim first (e.g. t → iy → ix).
 
+    Thin wrapper over the RaggedStore builders: picks the vectors' index
+    columns by rank (the ``time = -1`` 4-D sentinel clamps to 0, as always),
+    builds the leaf level by bincount, and derives the outer levels as
+    ZERO-COPY strided views of it (element-for-element identical to the
+    historically materialised arrays — the grid is rectangular, so
+    ``level[k] == leaf[::prod(inner dims)]``).
+
     Layout
     ------
-    Returns one offsets array per navigation dimension, outermost first.
-    The outermost N-1 levels use **uniform strides** (the grid is always
-    rectangular) so they can be stored as simple `(dim_size + 1,)` arrays
-    where `offsets[k]` = k * (product of inner dim sizes).
-
-    Only the innermost level (x_offsets, over all leaf positions) stores
-    actual variable-length counts — vectors per (t, iy, ix) position.
+    Returns one offsets array per navigation dimension, outermost first; only
+    the innermost level (over all leaf positions) stores variable-length
+    counts.
 
     Lookup for nav indices (i0, i1, …, iN) where N = len(full_nav_shape):
         flat_pos = i0 * stride[0] + i1 * stride[1] + … + iN
         s = nav_offsets[-1][flat_pos]
         e = nav_offsets[-1][flat_pos + 1]
         return flat_buffer[s:e]
 
-    The outer arrays (nav_offsets[0..N-2]) are included for the partial-index
-    API (slice_at(t) → all vectors at time t).  They store vector offsets:
+    The outer arrays (nav_offsets[0..N-2]) serve the partial-index API
+    (slice_at(t) → all vectors at time t, O(1)):
         nav_offsets[k][i] = sum of vectors in all leaf positions with outer
                             index < i along dimension k.
 
@@ -57,74 +64,22 @@ def _build_nav_offsets(
                        y_vec_offsets (n_t*nav_y+1,),
                        x_vec_offsets (n_t*nav_y*nav_x+1,)]
     """
-    n_dims = len(full_nav_shape)
-    n_patterns = int(np.prod(full_nav_shape))
+    from spyde.signals.ragged_store import build_leaf_offsets, derive_levels
 
-    if len(flat_buffer) == 0:
-        nav_offsets = []
-        product = 1
-        for dim_size in full_nav_shape:
-            product *= dim_size
-            nav_offsets.append(np.zeros(product + 1, dtype=np.int64))
-        return nav_offsets
-
-    # ── Build the innermost (leaf) offsets: one entry per (i0,…,iN) position ──
-    # Map each vector to its flat leaf index using stored coordinate columns.
+    n_dims = len(full_nav_shape)
     if n_dims == 2:
         col_seq = [COL_NAV_Y, COL_NAV_X]
     elif n_dims == 3:
         col_seq = [COL_TIME, COL_NAV_Y, COL_NAV_X]
     else:
         raise NotImplementedError("nav_offsets for >3 nav dims requires explicit outer columns")
 
-    strides = np.ones(n_dims, dtype=np.int64)
-    for i in range(n_dims - 2, -1, -1):
-        strides[i] = strides[i + 1] * full_nav_shape[i + 1]
-
-    flat_leaf = np.zeros(len(flat_buffer), dtype=np.int64)
-    for dim, col in enumerate(col_seq):
-        vals = flat_buffer[:, col].astype(np.int64)
-        if col == COL_TIME:
-            vals = np.where(vals < 0, np.int64(0), vals)
-        flat_leaf += vals * strides[dim]
-
-    leaf_counts = np.bincount(flat_leaf, minlength=n_patterns).astype(np.int64)
-    innermost = np.zeros(n_patterns + 1, dtype=np.int64)
-    np.cumsum(leaf_counts, out=innermost[1:])
-
-    # ── Build outer levels by summing leaf_counts over inner-dimension groups ──
-    # outer_level[k] stores the cumulative vector counts at dimension k,
-    # collapsing all inner dimensions.  This lets slice_at(t) return the
-    # exact flat_buffer slice for time step t in O(1).
-    #
-    # Example: full_nav_shape=(3, 4, 4), leaf_counts shape (48,):
-    #   y_level: sum groups of nav_x=4  → shape (12,)  [n_t * nav_y]
-    #   t_level: sum groups of nav_y=4  → shape (3,)   [n_t]
-    level_offsets = [innermost]
-    group_counts = leaf_counts.copy()
-
-    for dim in range(n_dims - 1, 0, -1):
-        group_size = full_nav_shape[dim]
-        n_outer = len(group_counts) // group_size
-        outer_counts = group_counts.reshape(n_outer, group_size).sum(axis=1)
-        outer_off = np.zeros(n_outer + 1, dtype=np.int64)
-        np.cumsum(outer_counts, out=outer_off[1:])
-        level_offsets.append(outer_off)
-        group_counts = outer_counts
-
-    level_offsets.reverse()  # now outermost-first
-    return level_offsets
-
-
-@dataclass
-class _AxisLite:
-    """Minimal axis record so vectors loaded from disk can be rendered
-    without HyperSpy axes objects (duck-types .scale/.offset/.size/.units/.name)."""
-    scale: float = 1.0
-    offset: float = 0.0
-    size: int = 0
-    units: str = ""
-    name: str = ""
+    if len(flat_buffer) == 0:
+        index_arrays = [np.zeros(0, dtype=np.int64)] * n_dims
+    else:
+        index_arrays = [flat_buffer[:, c] for c in col_seq]
+    leaf = build_leaf_offsets(index_arrays, full_nav_shape)
+    return derive_levels(leaf, full_nav_shape)
 
 
 def _render_disks_block(
@@ -195,7 +150,7 @@ def _render_disks_block(
 
 
 @dataclass
-class SpyDEDiffractionVectors:
+class SpyDEDiffractionVectors(RaggedStore):
     """
     Flat-buffer CSR storage for diffraction vectors across a scan.
 
@@ -242,6 +197,11 @@ class SpyDEDiffractionVectors:
         Falls back to numpy when CUDA is unavailable.
     """
 
+    # The packed-buffer column order IS the ABI (COL_* indexed straight into
+    # numpy slices and GPU tensors) — agreement with COLUMN_NAMES/COL_* is
+    # asserted at import, right below the class.
+    columns_schema: ClassVar[tuple] = tuple((n, "f4") for n in COLUMN_NAMES)
+
     flat_buffer:    np.ndarray          # (N_total, 6) float32
     nav_offsets:    List[np.ndarray]    # outermost-first CSR levels
     nav_shape:      tuple               # (nav_y, nav_x)
@@ -266,27 +226,38 @@ class SpyDEDiffractionVectors:
     _kdtree:      Optional[object]     = field(default=None, repr=False)
     _gpu_buffer:  Optional[object]     = field(default=None, repr=False)  # torch.Tensor on CUDA
 
-    # ── Internal helpers ──────────────────────────────────────────────────────
-
-    def _flat_pos(self, nav_indices: tuple) -> int:
-        """Convert nav indices to a flat leaf position using grid strides."""
-        pos = 0
-        stride = 1
-        for idx, dim_size in zip(reversed(nav_indices), reversed(self.full_nav_shape)):
-            pos += int(idx) * stride
-            stride *= dim_size
-        return pos
-
-    def _slice_flat(self, nav_indices: tuple) -> np.ndarray:
-        """
-        Return the flat_buffer slice for the given nav_indices.
-        Uses the innermost nav_offsets[-1] via arithmetic flat position.
-        O(n_dims) — pure arithmetic, no pointer chasing.
+    def __post_init__(self):
+        """Wire the RaggedStore base state over the SAME memory the dataclass
+        fields hold — pure aliasing, no copies, no rebuilt arrays.
+
+        The base's leaf offsets alias ``nav_offsets[-1]`` and its packed
+        backing IS ``flat_buffer``; the level list is the stored
+        ``nav_offsets`` list itself (kept verbatim — tests construct with
+        hand-built lists, including a degenerate outer level, and those stay
+        the source of truth for this instance). The subclass field ``offsets``
+        keeps its legacy 4-D-only semantics (None for 5-D) and shadows the
+        base's public alias; base internals never read it. Identity aliasing
+        is what keeps pickling size and the 4-D
+        ``offsets is nav_offsets[-1]`` relationship unchanged.
         """
-        flat_pos = self._flat_pos(nav_indices)
-        s = int(self.nav_offsets[-1][flat_pos])
-        e = int(self.nav_offsets[-1][flat_pos + 1])
-        return self.flat_buffer[s:e]
+        self._packed = self.flat_buffer
+        self._columns = None
+        self._levels = self.nav_offsets
+        self._offsets = self.nav_offsets[-1] if len(self.nav_offsets or []) else None
+        self._index_columns = (
+            ("time", "nav_y", "nav_x") if len(self.full_nav_shape) == 3
+            else ("nav_y", "nav_x"))
+        self._finalized = True
+        self._staged = None
+        self._staged_packed = None
+        self._staged_rows = 0
+        self._append_lock = None
+
+    # ── Internal helpers ──────────────────────────────────────────────────────
+    # _flat_pos and _slice_flat are inherited from RaggedStore: the base's
+    # packed backing IS flat_buffer and its levels ARE the stored nav_offsets
+    # list (see __post_init__), so the inherited bodies read the same memory
+    # the historical methods did.
 
     def _frame_slice(self, t: int) -> np.ndarray:
         """
@@ -316,23 +287,7 @@ def slice_at(self, *nav_indices: int) -> np.ndarray:
             vecs.slice_at(t)             — all vectors at time step t (O(1))
             vecs.slice_at(t, iy)         — all vectors at time t, row iy (O(1))
         """
-        n = len(nav_indices)
-        n_dims = len(self.full_nav_shape)
-        if n == n_dims:
-            return self._slice_flat(nav_indices)
-        # Partial index: use the outer-level vector offsets
-        # nav_offsets[n-1] has cumulative vector counts at level n-1
-        # (e.g., nav_offsets[0] = time-level offsets)
-        # Compute flat position at this level
-        level_shape = self.full_nav_shape[:n]
-        pos = 0
-        stride = 1
-        for idx, dim_size in zip(reversed(nav_indices), reversed(level_shape)):
-            pos += int(idx) * stride
-            stride *= dim_size
-        level_idx = n - 1  # which nav_offsets level to use
-        s = int(self.nav_offsets[level_idx][pos])
-        e = int(self.nav_offsets[level_idx][pos + 1])
+        s, e = self._prefix_row_range(nav_indices)
         return self.flat_buffer[s:e]
 
     def at(self, iy: int, ix: int) -> np.ndarray:
@@ -409,12 +364,8 @@ def count_map_series(self) -> np.ndarray:
         stack / time dimension; this is the natural navigator for a 5-D stack
         (scrub the stack axis → that slice's spatial counts)."""
         nav_y, nav_x = self.nav_shape
-        x_off = self.nav_offsets[-1]
-        if self.n_time == 0:
-            counts = np.diff(x_off).reshape(1, nav_y, nav_x)
-        else:
-            counts = np.diff(x_off).reshape(self.n_time, nav_y, nav_x)
-        return counts.astype(np.int32)
+        return self.counts().reshape(
+            max(1, self.n_time), nav_y, nav_x).astype(np.int32)
 
     def count_map_at_t(self, t: int) -> np.ndarray:
         """(nav_y, nav_x) int32 — vector count at time step t.
@@ -429,9 +380,8 @@ def count_map_at_t(self, t: int) -> np.ndarray:
         t = int(np.clip(t, 0, self.n_time - 1))
         return self.count_map_series()[t]
 
-    def flatten(self) -> np.ndarray:
-        """Return the full (N_total, 6) flat buffer."""
-        return self.flat_buffer
+    # flatten() is inherited from RaggedStore (the packed backing IS
+    # flat_buffer, so it returns the same (N_total, 6) object it always did).
 
     # ── Virtual imaging ───────────────────────────────────────────────────────
 
@@ -1007,13 +957,32 @@ def from_arrays(
         Primary constructor.  flat_buffer must already be sorted outermost-first.
         nav_offsets are built automatically.
         nav_shape is always the last two dims of full_nav_shape.
+
+        Internally this is the RaggedStore streaming seam — one staged batch,
+        then ``finalize()`` — so the ordinary construction path transitively
+        exercises the exact machinery chunked computes stream through.
+        finalize()'s already-sorted fast path ADOPTS the buffer (no copy, no
+        reorder), keeping this byte- and identity-compatible with the
+        historical build; an unsorted buffer (previously a silent contract
+        violation yielding broken CSR slices) now gets stable-sorted.
         """
-        nav_offsets = _build_nav_offsets(flat_buffer, full_nav_shape)
+        full_nav_shape = tuple(int(s) for s in full_nav_shape)
+        if len(full_nav_shape) == 2:
+            index_columns = ("nav_y", "nav_x")
+        elif len(full_nav_shape) == 3:
+            index_columns = ("time", "nav_y", "nav_x")
+        else:
+            raise NotImplementedError(
+                "nav_offsets for >3 nav dims requires explicit outer columns")
+        staged = cls.streaming(full_nav_shape, index_columns=index_columns)
+        staged.append_batch(np.asarray(flat_buffer))
+        staged.finalize()
+        nav_offsets = staged.offset_levels()
         nav_shape = full_nav_shape[-2:]
         # Legacy offsets: only meaningful for 4D
         offsets = nav_offsets[-1] if len(full_nav_shape) == 2 else None
         return cls(
-            flat_buffer=flat_buffer,
+            flat_buffer=staged.flatten(),
             nav_offsets=nav_offsets,
             nav_shape=nav_shape,
             full_nav_shape=full_nav_shape,
@@ -1050,3 +1019,13 @@ def from_ragged(
                 flat_buffer[s:e, COL_KX:COL_KY + 1] = arr
 
         return cls.from_arrays(flat_buffer, nav_shape, **kwargs)
+
+
+# The packed layout is the pinned ABI of orchestrate, strain_mapping,
+# vector_orientation(_gpu), vectors_embed and live_frames — COL_* index
+# straight into numpy slices and GPU tensors, so a schema/constant divergence
+# would produce plausible garbage, not an error. Freeze the agreement at import.
+assert tuple(n for n, _ in SpyDEDiffractionVectors.columns_schema) == COLUMN_NAMES
+assert (COL_NAV_X, COL_NAV_Y, COL_KX, COL_KY,
+        COL_TIME, COL_INTENSITY) == tuple(range(N_COLS))
+assert len(COLUMN_NAMES) == N_COLS