diff --git a/spyde/actions/find_vectors_action.py b/spyde/actions/find_vectors_action.py index 6901c792..61f75d27 100644 --- a/spyde/actions/find_vectors_action.py +++ b/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 diff --git a/spyde/actions/lifecycle.py b/spyde/actions/lifecycle.py index e9ca8e89..2b424fde 100644 --- a/spyde/actions/lifecycle.py +++ b/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: `` 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. diff --git a/spyde/signals/diffraction_vectors.py b/spyde/signals/diffraction_vectors.py index 639fe2f5..e3aab612 100644 --- a/spyde/signals/diffraction_vectors.py +++ b/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,15 +33,18 @@ 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 @@ -45,8 +52,8 @@ def _build_nav_offsets( 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,19 +64,9 @@ 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: @@ -77,54 +74,12 @@ def _build_nav_offsets( 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 diff --git a/spyde/signals/ragged_store.py b/spyde/signals/ragged_store.py new file mode 100644 index 00000000..77e7ff8d --- /dev/null +++ b/spyde/signals/ragged_store.py @@ -0,0 +1,670 @@ +""" +RaggedStore — ragged per-navigation-position column storage over a scan grid. + +The shared base for SpyDE's genuinely ragged result families (diffraction +vectors; particles), extracted from ``SpyDEDiffractionVectors``. It owns the +NAV-SPACE storage model only — signal-space (detector axes, kernel radii, +rendering) stays with the owning subclass. + +Storage model (the whole contract) +---------------------------------- +offsets : (n_positions + 1,) int64 CSR row pointers over the C-ordered flat + nav grid; ``offsets[p]:offsets[p+1]`` is position p's rows. +columns : one flat 1-D array per schema column, all length n_rows, dtype per + schema. MAY be strided views into a single packed (n_rows, n_cols) array + (``from_packed``) — the packed layout is then the subclass's public ABI; + the base never copies or reorders a packed backing. +full_nav_shape : all nav dims, outermost first, rank >= 1. +nav_axes : axis records duck-typing .scale/.offset/.size/.units/.name; scan + calibration only. + +Because the grid is rectangular, every OUTER level of a multi-level CSR index +is a zero-copy strided view of the leaf level (``leaf[::prod(inner dims)]``) — +so the base stores ONE offsets array and derives the rest +(:func:`derive_levels`). + +Lifecycle +--------- +STRUCTURE-FROZEN after ``finalize()``: row count/order, offsets and the column +set never change; cell VALUES may be overwritten in place by the owning +subclass (cache invalidation is the subclass's problem, as today). + +Threading +--------- +``append_batch()`` is the only mutating call (mutex-guarded); every read +accessor is valid only after ``finalize()`` and thereafter lock-free. An +unfinalized store carries a ``threading.Lock`` and is therefore not picklable; +a finalized one drops it and pickles like a plain container. +""" +from __future__ import annotations + +import threading +from dataclasses import dataclass +from typing import Callable, ClassVar, Dict, List, Optional, Sequence, Tuple + +import numpy as np + +ColumnDef = Tuple[str, str] # (name, numpy dtype string), e.g. ("kx", "f4") + + +@dataclass +class _AxisLite: + """Minimal axis record so results 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 = "" + + +# ── grid arithmetic ────────────────────────────────────────────────────────── + +def _grid_strides(full_nav_shape: Sequence[int]) -> np.ndarray: + """int64 C-order strides of the nav grid, outermost first (innermost = 1).""" + n = len(full_nav_shape) + strides = np.ones(n, dtype=np.int64) + for i in range(n - 2, -1, -1): + strides[i] = strides[i + 1] * int(full_nav_shape[i + 1]) + return strides + + +def flat_leaf_index(index_arrays: Sequence[np.ndarray], + full_nav_shape: Sequence[int]) -> np.ndarray: + """Per-row flat leaf position from the integer index columns. + + ``index_arrays`` holds one 1-D array per nav dim, outermost first. Values + are cast to int64 TRANSIENTLY (stored columns are never cast — the f4 + packed ABI stays f4); negative values clamp to 0, which is what lets a + subclass sentinel (the vectors' ``time = -1`` for "no axis") classify as + position 0 rather than corrupt the bincount. + """ + n_rows = int(len(index_arrays[0])) if index_arrays else 0 + strides = _grid_strides(full_nav_shape) + flat = np.zeros(n_rows, dtype=np.int64) + for vals, stride in zip(index_arrays, strides): + v = np.asarray(vals).astype(np.int64) + v = np.where(v < 0, np.int64(0), v) + flat += v * int(stride) + return flat + + +def build_leaf_offsets(index_arrays: Sequence[np.ndarray], + full_nav_shape: Sequence[int]) -> np.ndarray: + """(n_positions + 1,) int64 CSR row pointers by bincount — O(n_rows), no + sort, no copy. Rows must already be sorted by flat position for the + resulting slices to be meaningful (counts are order-independent).""" + n_positions = int(np.prod(np.asarray(full_nav_shape, dtype=np.int64))) + flat = flat_leaf_index(index_arrays, full_nav_shape) + counts = np.bincount(flat, minlength=n_positions).astype(np.int64) + if counts.shape[0] != n_positions: + raise ValueError( + f"index columns address position {int(flat.max())} outside the " + f"{tuple(full_nav_shape)} nav grid") + leaf = np.zeros(n_positions + 1, dtype=np.int64) + np.cumsum(counts, out=leaf[1:]) + return leaf + + +def derive_levels(leaf: np.ndarray, + full_nav_shape: Sequence[int]) -> List[np.ndarray]: + """Multi-level CSR offsets, outermost first, derived from the leaf level. + + The rectangular-grid stride identity: level k is ``leaf[::prod(inner + dims)]`` — a ZERO-COPY strided view, element-for-element identical to the + historically materialised outer levels. The innermost entry is the leaf + array itself (identity, not a view), preserving the 4-D + ``offsets is nav_offsets[-1]`` alias downstream.""" + strides = _grid_strides(full_nav_shape) + return [leaf[:: int(s)] for s in strides[:-1]] + [leaf] + + +class RaggedStore: + """Ragged per-navigation-position column store over a rectangular scan grid. + + See the module docstring for the storage model and lifecycle contract. + + Construction paths: + + * ``RaggedStore(columns, offsets, full_nav_shape, ...)`` — FINALIZED from a + dict of 1-D column arrays + valid offsets (zero-copy; arrays adopted). + * ``from_packed(packed, full_nav_shape, index_columns=...)`` — FINALIZED + from an already-sorted packed (n_rows, n_cols) single-dtype array. + * ``streaming(...)`` then ``append_batch()`` × N then ``finalize()`` — THE + one seam for chunked computes; batches may arrive in any chunk order. + + Subclasses declare ``columns_schema`` (ORDER load-bearing — for a packed + backing it IS the column order) and may bump ``format_version`` under the + append-only-columns rule. Exactly two persistence extension points exist: + ``_save_extra`` / ``_load_extra``. Subclasses with their own constructor + ABI (dataclass fields) keep their constructors and wire the base state over + the same memory (see ``SpyDEDiffractionVectors.__post_init__``). + """ + + # ORDER load-bearing; append-only + bump format_version on change. + columns_schema: ClassVar[Tuple[ColumnDef, ...]] = () + format_version: ClassVar[int] = 1 + + def __init__(self, columns: Dict[str, np.ndarray], offsets: np.ndarray, + full_nav_shape: Sequence[int], *, + index_columns: Sequence[str] = (), + nav_axes: Sequence[object] = (), + params: Optional[dict] = None, + provenance: Optional[dict] = None): + """Construct FINALIZED from dict-of-1D-arrays + valid offsets.""" + names = self._schema_names() + if not names: + raise TypeError( + f"{type(self).__name__} declares no columns_schema") + full_nav_shape = tuple(int(s) for s in full_nav_shape) + if not full_nav_shape: + raise ValueError("full_nav_shape must have rank >= 1") + offsets = np.asarray(offsets, dtype=np.int64) + n_positions = int(np.prod(np.asarray(full_nav_shape, dtype=np.int64))) + if offsets.shape != (n_positions + 1,): + raise ValueError( + f"offsets shape {offsets.shape} != ({n_positions + 1},) for " + f"nav grid {full_nav_shape}") + n_rows = int(offsets[-1]) + cols: Dict[str, np.ndarray] = {} + given = dict(columns) + if set(given) != set(names): + raise ValueError( + f"columns {sorted(given)} do not match schema {list(names)}") + for name in names: # schema order, zero-copy + arr = np.asarray(given[name]) + if arr.ndim != 1 or arr.shape[0] != n_rows: + raise ValueError( + f"column {name!r} must be 1-D of length {n_rows}, " + f"got shape {arr.shape}") + cols[name] = arr + self.full_nav_shape = full_nav_shape + self.nav_axes = list(nav_axes or []) + self.params = dict(params) if params else {} + self.provenance = provenance + self.offsets = offsets + self._wire(columns=cols, packed=None, offsets=offsets, levels=None, + index_columns=self._check_index_columns(index_columns)) + + # ── schema helpers ─────────────────────────────────────────────────────── + + @classmethod + def _schema_names(cls) -> Tuple[str, ...]: + return tuple(n for n, _ in cls.columns_schema) + + @classmethod + def _check_index_columns(cls, index_columns: Sequence[str]) -> Tuple[str, ...]: + idx = tuple(index_columns) + names = cls._schema_names() + for c in idx: + if c not in names: + raise ValueError(f"index column {c!r} not in schema {list(names)}") + return idx + + def _wire(self, *, columns: Optional[Dict[str, np.ndarray]], + packed: Optional[np.ndarray], offsets: np.ndarray, + levels: Optional[List[np.ndarray]], + index_columns: Sequence[str]) -> None: + """Install the FINALIZED read state. ``levels=None`` derives the outer + levels as zero-copy strided views of *offsets* (the leaf).""" + self._columns = columns + self._packed = packed + self._offsets = offsets + self._levels = (levels if levels is not None + else derive_levels(offsets, self.full_nav_shape)) + self._index_columns = tuple(index_columns) + self._finalized = True + self._staged = None + self._staged_packed = None + self._staged_rows = 0 + self._append_lock = None # locks don't pickle; frozen stores need none + + def _require_finalized(self) -> None: + if not getattr(self, "_finalized", False): + raise RuntimeError( + f"{type(self).__name__} is not finalized — call finalize() " + "before reading") + + # ── constructors ───────────────────────────────────────────────────────── + + @classmethod + def from_packed(cls, packed: np.ndarray, full_nav_shape: Sequence[int], *, + index_columns: Sequence[str] = (), + nav_axes: Sequence[object] = (), + params: Optional[dict] = None, + provenance: Optional[dict] = None) -> "RaggedStore": + """FINALIZED from a packed (n_rows, n_cols) single-dtype array already + sorted outermost-nav-first. Columns are zero-copy strided views; + offsets by bincount over the index columns — O(n_rows), no sort, no + copy (the packed array is ADOPTED, never reordered). + + Only valid for classes whose usable state is pure base state — a + subclass with its own constructor ABI (dataclass fields) must build + through its own constructors instead. + """ + names = cls._schema_names() + if not names: + raise TypeError(f"{cls.__name__} declares no columns_schema") + packed = np.asarray(packed) + if packed.ndim != 2 or packed.shape[1] != len(names): + raise ValueError( + f"packed must be (n_rows, {len(names)}); got {packed.shape}") + inst = cls.__new__(cls) + inst.full_nav_shape = tuple(int(s) for s in full_nav_shape) + if not inst.full_nav_shape: + raise ValueError("full_nav_shape must have rank >= 1") + inst.nav_axes = list(nav_axes or []) + inst.params = dict(params) if params else {} + inst.provenance = provenance + idx = cls._check_index_columns(index_columns) + if len(idx) != len(inst.full_nav_shape): + raise ValueError( + f"index_columns {idx} must name one column per nav dim " + f"{inst.full_nav_shape}") + index_arrays = [packed[:, names.index(c)] for c in idx] + leaf = build_leaf_offsets(index_arrays, inst.full_nav_shape) + inst.offsets = leaf + inst._wire(columns=None, packed=packed, offsets=leaf, levels=None, + index_columns=idx) + return inst + + # ── streaming fill: THE one seam for chunked computes ──────────────────── + + @classmethod + def streaming(cls, full_nav_shape: Sequence[int], *, + index_columns: Sequence[str], + nav_axes: Sequence[object] = (), + params: Optional[dict] = None, + provenance: Optional[dict] = None) -> "RaggedStore": + """UNFINALIZED store accepting ``append_batch()``; reads raise until + ``finalize()``.""" + names = cls._schema_names() + if not names: + raise TypeError(f"{cls.__name__} declares no columns_schema") + inst = cls.__new__(cls) + inst.full_nav_shape = tuple(int(s) for s in full_nav_shape) + if not inst.full_nav_shape: + raise ValueError("full_nav_shape must have rank >= 1") + idx = cls._check_index_columns(index_columns) + if len(idx) != len(inst.full_nav_shape): + raise ValueError( + f"index_columns {idx} must name one column per nav dim " + f"{inst.full_nav_shape}") + inst.nav_axes = list(nav_axes or []) + inst.params = dict(params) if params else {} + inst.provenance = provenance + inst.offsets = None + inst._columns = None + inst._packed = None + inst._offsets = None + inst._levels = None + inst._index_columns = idx + inst._finalized = False + inst._staged = [] # dict-of-columns batches + inst._staged_packed = [] # packed (m, n_cols) blocks + inst._staged_rows = 0 + inst._append_lock = threading.Lock() + return inst + + def append_batch(self, batch) -> int: + """Stage rows; returns the total number of rows staged so far. + + *batch* is either an ``(m, n_cols)`` packed block (single-dtype + schema) or a dict of equal-length 1-D arrays covering EVERY schema + column (the integer index columns included). Batch kinds cannot be + mixed within one store. Thread-safe; callable from dask done-callbacks + in ANY chunk order. O(1) list append — no offsets maintenance during + the fill (deliberate: leaf offsets can't be maintained incrementally + under out-of-order chunks; batch-then-build matches the historical + one-shot build). + """ + lock = getattr(self, "_append_lock", None) + if getattr(self, "_finalized", False) or lock is None: + raise RuntimeError( + "append_batch() after finalize() — the structure is frozen") + names = self._schema_names() + with lock: + if self._finalized: # finalized while we waited for the lock + raise RuntimeError( + "append_batch() after finalize() — the structure is frozen") + if isinstance(batch, dict): + if self._staged_packed: + raise TypeError("cannot mix dict and packed batches") + if set(batch) != set(names): + raise ValueError( + f"batch columns {sorted(batch)} do not match schema " + f"{list(names)}") + cols = {n: np.asarray(batch[n]) for n in names} + lengths = {c.shape[0] for c in cols.values()} + if any(c.ndim != 1 for c in cols.values()) or len(lengths) > 1: + raise ValueError("batch columns must be equal-length 1-D arrays") + m = lengths.pop() if lengths else 0 + self._staged.append(cols) + else: + if self._staged: + raise TypeError("cannot mix dict and packed batches") + block = np.asarray(batch) + if block.ndim != 2 or block.shape[1] != len(names): + raise ValueError( + f"packed batch must be (m, {len(names)}); got {block.shape}") + m = int(block.shape[0]) + self._staged_packed.append(block) + self._staged_rows += int(m) + return self._staged_rows + + def finalize(self) -> "RaggedStore": + """Concatenate, STABLE-sort by flat nav position, build offsets, freeze. + + Idempotent. The O(n) already-sorted check keeps a sorted input + byte-identical — and, for a SINGLE sorted packed batch, object- + identical (the block is adopted, not copied), which is what keeps the + ``from_arrays``/orchestrate path exactly as it was. Row order within a + position = batch arrival order then in-batch order (stable) — what + keeps overlay spot order and the report embed's CSR-contiguity + assumption deterministic. + """ + lock = getattr(self, "_append_lock", None) + if getattr(self, "_finalized", False) or lock is None: + return self + with lock: + if self._finalized: # lost the race to another finalizer + return self + names = self._schema_names() + n_positions = int(np.prod( + np.asarray(self.full_nav_shape, dtype=np.int64))) + if self._staged_packed: + blocks = self._staged_packed + packed = blocks[0] if len(blocks) == 1 else np.concatenate(blocks) + index_arrays = [packed[:, names.index(c)] + for c in self._index_columns] + flat = flat_leaf_index(index_arrays, self.full_nav_shape) + if flat.size and np.any(np.diff(flat) < 0): + order = np.argsort(flat, kind="stable") + packed = packed[order] + flat = flat[order] + leaf = self._leaf_from_sorted(flat, n_positions) + columns = None + else: + cols: Dict[str, np.ndarray] = {} + for name, dt in self.columns_schema: + parts = [b[name] for b in self._staged] + if len(parts) == 1: + cols[name] = parts[0] + elif parts: + cols[name] = np.concatenate(parts) + else: + cols[name] = np.zeros(0, dtype=np.dtype(dt)) + index_arrays = [cols[c] for c in self._index_columns] + flat = flat_leaf_index(index_arrays, self.full_nav_shape) + if flat.size and np.any(np.diff(flat) < 0): + order = np.argsort(flat, kind="stable") + cols = {n: c[order] for n, c in cols.items()} + flat = flat[order] + leaf = self._leaf_from_sorted(flat, n_positions) + columns = cols + packed = None + self.offsets = leaf + self._wire(columns=columns, packed=packed, offsets=leaf, + levels=None, index_columns=self._index_columns) + return self + + @staticmethod + def _leaf_from_sorted(flat: np.ndarray, n_positions: int) -> np.ndarray: + counts = np.bincount(flat, minlength=n_positions).astype(np.int64) + if counts.shape[0] != n_positions: + raise ValueError( + f"index columns address position {int(flat.max())} outside " + f"the {n_positions}-position nav grid") + leaf = np.zeros(n_positions + 1, dtype=np.int64) + np.cumsum(counts, out=leaf[1:]) + return leaf + + # ── index arithmetic (shared with subclasses; hot path — no guards) ────── + + 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 _row_range(self, nav_indices: tuple) -> Tuple[int, int]: + """Half-open row range for one FULL nav index. O(n_dims) arithmetic.""" + p = self._flat_pos(nav_indices) + lev = self._levels[-1] + return int(lev[p]), int(lev[p + 1]) + + def _prefix_row_range(self, nav_indices: tuple) -> Tuple[int, int]: + """Half-open row range for a FULL or PREFIX nav index (partial indexing + reads the outer CSR level — O(1) via the stride identity). Exactly the + historical ``slice_at`` arithmetic, degenerate cases included.""" + n = len(nav_indices) + if n == len(self.full_nav_shape): + return self._row_range(nav_indices) + 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 + lev = self._levels[n - 1] + return int(lev[pos]), int(lev[pos + 1]) + + def _slice_flat(self, nav_indices: tuple) -> np.ndarray: + """Packed rows for one FULL nav index (packed backing only).""" + s, e = self._row_range(nav_indices) + return self._packed[s:e] + + # ── access ─────────────────────────────────────────────────────────────── + + def offset_levels(self) -> List[np.ndarray]: + """Multi-level CSR offsets, outermost first, leaf last. For base-built + stores the outer levels are zero-copy strided views of the leaf; a + subclass wired over a stored list returns that list's entries verbatim.""" + self._require_finalized() + return list(self._levels) + + def column(self, name: str) -> np.ndarray: + """The full flat column — zero-copy (a strided view for a packed + backing, the adopted array otherwise).""" + self._require_finalized() + if self._columns is not None: + return self._columns[name] + names = self._schema_names() + try: + i = names.index(name) + except ValueError: + raise KeyError(name) from None + return self._packed[:, i] + + def at(self, *nav_index: int) -> Dict[str, np.ndarray]: + """Zero-copy column views for ONE full nav index, in schema order. + + Subclasses with a pinned packed return (the vectors' ``(N, 6)`` slice) + override this; the base contract is the dict form.""" + self._require_finalized() + if len(nav_index) != len(self.full_nav_shape): + raise ValueError( + f"at() needs {len(self.full_nav_shape)} indices, " + f"got {len(nav_index)}") + s, e = self._row_range(nav_index) + return {name: self.column(name)[s:e] for name in self._schema_names()} + + def slice_at(self, *nav_prefix: int) -> Dict[str, np.ndarray]: + """All rows under an outer-index prefix, O(1) via the stride identity. + Returns the dict-of-column-views form (packed subclasses override).""" + self._require_finalized() + s, e = self._prefix_row_range(nav_prefix) + return {name: self.column(name)[s:e] for name in self._schema_names()} + + def flatten(self) -> np.ndarray: + """The full packed (n_rows, n_cols) buffer (packed backing only).""" + self._require_finalized() + if self._packed is None: + raise TypeError("flatten() requires a packed backing") + return self._packed + + def counts(self) -> np.ndarray: + """(n_positions,) int64 rows per leaf position — np.diff over the leaf + offsets. O(positions), never O(rows).""" + self._require_finalized() + return np.diff(self._levels[-1]) + + def count_map(self) -> np.ndarray: + """(full_nav_shape) int64 rows per position. O(positions) via + ``np.diff(offsets)`` — NOT implemented as ``map(None, 'count')`` + because count_map sits on the progressive-fill hot path; their + equivalence is pinned by test instead.""" + return self.counts().reshape(self.full_nav_shape) + + def map(self, column: Optional[str], reducer="mean", *, + fill: float = np.nan) -> np.ndarray: + """(full_nav_shape) per-position reduction of a column. + + ``reducer`` is one of ``'sum'|'mean'|'max'|'min'|'median'|'std'| + 'count'`` or a callable ``(rows_i,) -> scalar``. ``'count'`` ignores + *column* (``map(None, 'count') == count_map()`` — the pinned law) and + returns int64; ``'sum'`` returns float64 with empty positions 0; every + other reducer returns float64 with empty positions set to *fill* + (NaN — a frame with no rows has no mean; zero would plot a fake + event). 'median'/'std'/callable loop over non-empty positions in + Python — O(positions) calls; the array reducers are vectorised. + """ + self._require_finalized() + counts = self.counts() + n_positions = counts.shape[0] + if reducer == "count": + return counts.reshape(self.full_nav_shape) + col = np.asarray(self.column(column)) + leaf = np.asarray(self._levels[-1]) + nonempty = counts > 0 + if reducer in ("sum", "mean"): + pos_ids = np.repeat(np.arange(n_positions, dtype=np.int64), counts) + sums = np.bincount(pos_ids, weights=col.astype(np.float64), + minlength=n_positions) + if reducer == "sum": + return sums.reshape(self.full_nav_shape) + out = np.full(n_positions, fill, dtype=np.float64) + out[nonempty] = sums[nonempty] / counts[nonempty] + return out.reshape(self.full_nav_shape) + if reducer in ("max", "min"): + out = np.full(n_positions, fill, dtype=np.float64) + if nonempty.any(): + # reduceat over the STARTS of non-empty positions only: empty + # positions contribute no rows, so consecutive non-empty starts + # delimit exactly one position's rows (reduceat's empty-segment + # misbehaviour never comes into play). + starts = leaf[:-1][nonempty] + ufunc = np.maximum if reducer == "max" else np.minimum + out[nonempty] = ufunc.reduceat(col.astype(np.float64), starts) + return out.reshape(self.full_nav_shape) + if reducer in ("median", "std") or callable(reducer): + fn = {"median": np.median, "std": np.std}.get(reducer, reducer) + out = np.full(n_positions, fill, dtype=np.float64) + for p in np.nonzero(nonempty)[0]: + out[p] = fn(col[leaf[p]:leaf[p + 1]]) + return out.reshape(self.full_nav_shape) + raise ValueError(f"unknown reducer {reducer!r}") + + # ── persistence (versioned from day one; append-only columns) ──────────── + + def _save_extra(self) -> Tuple[Dict[str, np.ndarray], dict]: + """Subclass extension point: extra arrays + JSON-safe extra meta to + persist alongside the base payload.""" + return {}, {} + + @classmethod + def _load_extra(cls, arrays: Dict[str, np.ndarray], meta: dict) -> dict: + """Subclass extension point: attributes (name -> value) to set on the + loaded instance, built from what ``_save_extra`` persisted.""" + return {} + + def save(self, path: str) -> None: + """Save to a compressed ``.npz`` — versioned, self-describing (schema + + index columns + nav grid/axes in ``meta_json``), columns stored one + array each so the append-only rule can pad older files on load.""" + self._require_finalized() + import json + arrays: Dict[str, np.ndarray] = { + "offsets": np.ascontiguousarray(self._levels[-1]), + } + for name in self._schema_names(): + arrays[f"col_{name}"] = np.ascontiguousarray(self.column(name)) + extra_arrays, extra_meta = self._save_extra() + for k, v in (extra_arrays or {}).items(): + arrays[f"extra_{k}"] = np.asarray(v) + axes_meta = [ + dict(scale=float(ax.scale), offset=float(ax.offset), + size=int(ax.size), + units=str(getattr(ax, "units", "") or ""), + name=str(getattr(ax, "name", "") or "")) + for ax in (self.nav_axes or []) + ] + meta = { + "format_version": int(self.format_version), + "class": type(self).__name__, + "columns": [[n, d] for n, d in self.columns_schema], + "index_columns": list(self._index_columns), + "full_nav_shape": [int(s) for s in self.full_nav_shape], + "nav_axes": axes_meta, + "params": self.params or {}, + "provenance": self.provenance, + "extra": extra_meta or {}, + } + np.savez_compressed( + path, + meta_json=np.frombuffer( + json.dumps(meta, default=str).encode("utf-8"), dtype=np.uint8), + **arrays, + ) + + @classmethod + def load(cls, path: str) -> "RaggedStore": + """Load a store saved with :meth:`save`. + + An older file whose columns are a PREFIX of ``cls.columns_schema`` is + padded with zero-filled arrays for the missing columns (the proven + append-only rule); a file whose ``format_version`` exceeds the class's + is rejected. nav_axes come back as :class:`_AxisLite` records. + """ + import json + with np.load(path) as z: + meta = json.loads(bytes(z["meta_json"]).decode("utf-8")) + file_version = int(meta.get("format_version", 0)) + if file_version > cls.format_version: + raise ValueError( + f"{path}: format_version {file_version} is newer than " + f"{cls.__name__}.format_version {cls.format_version}") + file_names = [n for n, _ in + (tuple(c) for c in meta.get("columns", []))] + names = list(cls._schema_names()) + if file_names != names[:len(file_names)]: + raise ValueError( + f"{path}: saved columns {file_names} are not a prefix of " + f"the {cls.__name__} schema {names} (append-only rule)") + offsets = np.asarray(z["offsets"], dtype=np.int64) + n_rows = int(offsets[-1]) if offsets.size else 0 + cols: Dict[str, np.ndarray] = {} + for name, dt in cls.columns_schema: + key = f"col_{name}" + if key in z.files: + cols[name] = np.asarray(z[key]) + else: + cols[name] = np.zeros(n_rows, dtype=np.dtype(dt)) + extra_arrays = {k[len("extra_"):]: np.asarray(z[k]) + for k in z.files if k.startswith("extra_")} + nav_axes = [_AxisLite(**a) for a in meta.get("nav_axes", [])] + inst = cls( + cols, offsets, tuple(int(s) for s in meta["full_nav_shape"]), + index_columns=tuple(meta.get("index_columns", ()) or ()), + nav_axes=nav_axes, + params=meta.get("params") or {}, + provenance=meta.get("provenance"), + ) + for k, v in (cls._load_extra(extra_arrays, + meta.get("extra", {}) or {}) or {}).items(): + setattr(inst, k, v) + return inst diff --git a/spyde/tests/migrated/test_ragged_store.py b/spyde/tests/migrated/test_ragged_store.py new file mode 100644 index 00000000..c8eb8eb0 --- /dev/null +++ b/spyde/tests/migrated/test_ragged_store.py @@ -0,0 +1,388 @@ +""" +RaggedStore — the shared ragged per-nav-position column store. + +Pins the base contract SpyDEDiffractionVectors (and later the particle store) +is wired over: streaming out-of-order fill == from_packed of the sorted input, +derived outer offset levels are zero-copy strided views matching the historical +``_build_nav_offsets`` output, the ``count_map() == map(None, 'count')`` law, +the versioned save/load format (append-only column padding included), rank-1 +and rank-3 grids, and the structure-frozen lifecycle. +""" +from __future__ import annotations + +import numpy as np +import pytest + +from spyde.signals.ragged_store import ( + RaggedStore, _AxisLite, build_leaf_offsets, derive_levels, +) + + +class _EventStore(RaggedStore): + """Rank-2 store with a heterogeneous (dict-of-columns) schema.""" + columns_schema = (("iy", "i8"), ("ix", "i8"), ("val", "f4")) + + +class _SeriesStore(RaggedStore): + """Rank-1 store (a time series of ragged events).""" + columns_schema = (("t", "i8"), ("val", "f4")) + + +class _StackStore(RaggedStore): + """Rank-3 packed store mirroring the 5-D vectors layout (single dtype).""" + columns_schema = (("t", "f4"), ("iy", "f4"), ("ix", "f4"), ("val", "f4")) + + +def _random_packed(rng, full_nav_shape, max_per_pos=5): + """A SORTED packed (N, 4) f4 buffer for _StackStore with random ragged + counts (some positions empty).""" + rows = [] + grid = np.ndindex(*full_nav_shape) + for pos in grid: + n = int(rng.integers(0, max_per_pos + 1)) + for _ in range(n): + rows.append([*pos, rng.uniform(0, 100)]) + if not rows: + return np.zeros((0, len(_StackStore.columns_schema)), dtype=np.float32) + return np.asarray(rows, dtype=np.float32) + + +class TestStreamingEqualsFromPacked: + def test_out_of_order_packed_batches(self): + rng = np.random.default_rng(7) + shape = (3, 4, 5) + packed = _random_packed(rng, shape) + # Shuffle rows, then split into ragged batches — the out-of-order + # chunk-arrival case. + shuffled = packed[rng.permutation(len(packed))] + splits = np.array_split(shuffled, 4) + + st = _StackStore.streaming(shape, index_columns=("t", "iy", "ix")) + total = 0 + for b in splits: + total = st.append_batch(b) + assert total == len(packed) + st.finalize() + + # Reference: from_packed of the stable-sorted concatenation (same + # arrival order finalize saw). + flat = (shuffled[:, 0].astype(np.int64) * 20 + + shuffled[:, 1].astype(np.int64) * 5 + + shuffled[:, 2].astype(np.int64)) + ref_sorted = shuffled[np.argsort(flat, kind="stable")] + ref = _StackStore.from_packed(ref_sorted, shape, + index_columns=("t", "iy", "ix")) + + np.testing.assert_array_equal(st.offset_levels()[-1], + ref.offset_levels()[-1]) + for name in ("t", "iy", "ix", "val"): + np.testing.assert_array_equal(st.column(name), ref.column(name)) + np.testing.assert_array_equal(st.count_map(), ref.count_map()) + + def test_dict_batches_match_packed(self): + rng = np.random.default_rng(11) + shape = (4, 3) + iy = rng.integers(0, 4, 40) + ix = rng.integers(0, 3, 40) + val = rng.uniform(0, 1, 40).astype(np.float32) + + st = _EventStore.streaming(shape, index_columns=("iy", "ix")) + st.append_batch({"iy": iy[:25], "ix": ix[:25], "val": val[:25]}) + st.append_batch({"iy": iy[25:], "ix": ix[25:], "val": val[25:]}) + st.finalize() + + order = np.argsort(iy * 3 + ix, kind="stable") + np.testing.assert_array_equal(st.column("iy"), iy[order]) + np.testing.assert_array_equal(st.column("ix"), ix[order]) + np.testing.assert_array_equal(st.column("val"), val[order]) + counts = np.bincount(iy * 3 + ix, minlength=12).reshape(4, 3) + np.testing.assert_array_equal(st.count_map(), counts) + + def test_single_sorted_packed_batch_is_adopted_not_copied(self): + """The already-sorted fast path keeps the block object-identical — + what keeps the from_arrays/orchestrate path byte-for-byte as it was.""" + rng = np.random.default_rng(3) + shape = (2, 3, 4) + packed = _random_packed(rng, shape) + st = _StackStore.streaming(shape, index_columns=("t", "iy", "ix")) + st.append_batch(packed) + st.finalize() + assert st.flatten() is packed + + def test_mixing_batch_kinds_raises(self): + st = _EventStore.streaming((2, 2), index_columns=("iy", "ix")) + st.append_batch({"iy": [0], "ix": [1], "val": [2.0]}) + with pytest.raises(TypeError): + st.append_batch(np.zeros((1, 3))) + + def test_zero_batches_finalizes_empty(self): + st = _EventStore.streaming((2, 3), index_columns=("iy", "ix")) + st.finalize() + assert st.count_map().shape == (2, 3) + assert st.count_map().sum() == 0 + assert st.column("val").shape == (0,) + assert all(v.shape == (0,) for v in st.at(1, 2).values()) + + +class TestDerivedLevels: + def test_levels_match_build_nav_offsets_and_are_views(self): + """The rectangular-grid stride identity: every outer level equals the + historically materialised ``_build_nav_offsets`` output, but as a + ZERO-COPY strided view of the leaf.""" + from spyde.signals.diffraction_vectors import ( + _build_nav_offsets, N_COLS, COL_NAV_X, COL_NAV_Y, COL_TIME, + ) + rng = np.random.default_rng(5) + shape = (3, 4, 5) + packed = _random_packed(rng, shape) # (t, iy, ix, val) + + # The same rows in the vectors (N, 6) layout for the historical builder. + flat6 = np.zeros((len(packed), N_COLS), dtype=np.float32) + flat6[:, COL_TIME] = packed[:, 0] + flat6[:, COL_NAV_Y] = packed[:, 1] + flat6[:, COL_NAV_X] = packed[:, 2] + legacy = _build_nav_offsets(flat6, shape) + + st = _StackStore.from_packed(packed, shape, + index_columns=("t", "iy", "ix")) + levels = st.offset_levels() + assert len(levels) == len(legacy) == 3 + for got, want in zip(levels, legacy): + np.testing.assert_array_equal(got, want) + leaf = levels[-1] + for outer in levels[:-1]: + assert np.shares_memory(outer, leaf) + + def test_rank1_single_level(self): + st = _SeriesStore.streaming((4,), index_columns=("t",)) + st.append_batch({"t": [2, 0, 2, 3], "val": [1.0, 2.0, 3.0, 4.0]}) + st.finalize() + levels = st.offset_levels() + assert len(levels) == 1 + np.testing.assert_array_equal(levels[0], [0, 1, 1, 3, 4]) + + def test_columns_are_zero_copy_views_of_packed(self): + rng = np.random.default_rng(2) + shape = (2, 2, 2) + packed = _random_packed(rng, shape) + st = _StackStore.from_packed(packed, shape, + index_columns=("t", "iy", "ix")) + assert np.shares_memory(st.column("val"), packed) + if len(packed): + row = st.at(0, 0, 0) if st.counts()[0] else st.at(*np.unravel_index( + int(np.argmax(st.counts())), shape)) + assert all(np.shares_memory(v, packed) for v in row.values() + if v.size) + + +class TestCountMapMapLaw: + def _store(self, seed=9, shape=(4, 6)): + rng = np.random.default_rng(seed) + n = 60 + iy = rng.integers(0, shape[0], n) + # Leave column 0 empty so the law covers empty positions too. + ix = rng.integers(1, shape[1], n) + val = rng.uniform(-5, 5, n).astype(np.float32) + st = _EventStore.streaming(shape, index_columns=("iy", "ix")) + st.append_batch({"iy": iy, "ix": ix, "val": val}) + return st.finalize(), iy, ix, val + + def test_count_map_equals_map_count(self): + st, *_ = self._store() + cm = st.count_map() + mapped = st.map(None, "count") + assert cm.dtype == mapped.dtype == np.int64 + np.testing.assert_array_equal(cm, mapped) + assert (cm[:, 0] == 0).all() # the law holds over empties + + def test_map_reducers_match_manual_loop(self): + st, iy, ix, val = self._store() + shape = st.full_nav_shape + for reducer, fn in [("sum", np.sum), ("mean", np.mean), + ("max", np.max), ("min", np.min), + ("median", np.median), ("std", np.std)]: + got = st.map("val", reducer) + for y in range(shape[0]): + for x in range(shape[1]): + rows = val[(iy == y) & (ix == x)] + if not len(rows): + want = 0.0 if reducer == "sum" else np.nan + else: + want = fn(rows.astype(np.float64)) + np.testing.assert_allclose( + got[y, x], want, atol=1e-6, err_msg=f"{reducer} ({y},{x})") + + def test_map_callable_and_fill(self): + st, *_ = self._store() + got = st.map("val", lambda rows: float(len(rows)), fill=-1.0) + counts = st.count_map() + np.testing.assert_array_equal(got[counts > 0], + counts[counts > 0].astype(np.float64)) + assert (got[counts == 0] == -1.0).all() + + +class TestSaveLoad: + def _store(self): + st = _EventStore.streaming( + (2, 3), index_columns=("iy", "ix"), + nav_axes=[_AxisLite(scale=2.0, offset=1.0, size=2, units="nm", name="y"), + _AxisLite(scale=3.0, offset=-1.0, size=3, units="nm", name="x")], + params={"threshold": 0.5}, + provenance={"action": "test", "spyde_version": "0"}) + st.append_batch({"iy": [1, 0, 1], "ix": [2, 0, 2], "val": [7.0, 8.0, 9.0]}) + return st.finalize() + + def test_round_trip(self, tmp_path): + st = self._store() + path = str(tmp_path / "store.npz") + st.save(path) + back = _EventStore.load(path) + assert back.full_nav_shape == (2, 3) + np.testing.assert_array_equal(back.offset_levels()[-1], + st.offset_levels()[-1]) + for name in ("iy", "ix", "val"): + np.testing.assert_array_equal(back.column(name), st.column(name)) + assert back.params == {"threshold": 0.5} + assert back.provenance["action"] == "test" + assert len(back.nav_axes) == 2 + assert back.nav_axes[0].scale == 2.0 and back.nav_axes[1].name == "x" + np.testing.assert_array_equal(back.count_map(), st.count_map()) + + def test_older_file_prefix_columns_are_padded(self, tmp_path): + """Append-only rule: a file written before a column existed loads with + that column zero-filled at the schema dtype.""" + + class _V1(RaggedStore): + columns_schema = (("iy", "i8"), ("ix", "i8"), ("val", "f4")) + + class _V2(RaggedStore): + columns_schema = (("iy", "i8"), ("ix", "i8"), ("val", "f4"), + ("weight", "f8")) + format_version = 2 + + old = _V1.streaming((2, 2), index_columns=("iy", "ix")) + old.append_batch({"iy": [0, 1], "ix": [1, 0], "val": [3.0, 4.0]}) + old.finalize() + path = str(tmp_path / "old.npz") + old.save(path) + + new = _V2.load(path) + np.testing.assert_array_equal(new.column("val"), old.column("val")) + pad = new.column("weight") + assert pad.dtype == np.float64 and pad.shape == (2,) + assert (pad == 0).all() + + def test_newer_format_version_is_rejected(self, tmp_path): + class _New(RaggedStore): + columns_schema = (("iy", "i8"), ("ix", "i8"), ("val", "f4")) + format_version = 99 + + st = _New.streaming((1, 1), index_columns=("iy", "ix")) + st.append_batch({"iy": [0], "ix": [0], "val": [1.0]}) + st.finalize() + path = str(tmp_path / "new.npz") + st.save(path) + with pytest.raises(ValueError, match="format_version"): + _EventStore.load(path) + + def test_non_prefix_columns_are_rejected(self, tmp_path): + class _Other(RaggedStore): + columns_schema = (("a", "f4"), ("b", "f4")) + + st = self._store() + path = str(tmp_path / "store.npz") + st.save(path) + with pytest.raises(ValueError, match="prefix"): + _Other.load(path) + + +class TestRanks: + def test_rank1_round_trip_access(self): + st = _SeriesStore.streaming((5,), index_columns=("t",)) + st.append_batch({"t": [4, 1, 1], "val": [9.0, 5.0, 6.0]}) + st.finalize() + assert st.count_map().shape == (5,) + np.testing.assert_array_equal(st.count_map(), [0, 2, 0, 0, 1]) + np.testing.assert_array_equal(st.at(1)["val"], [5.0, 6.0]) + np.testing.assert_array_equal(st.at(4)["val"], [9.0]) + assert st.at(0)["val"].shape == (0,) + + def test_rank3_full_and_prefix_slicing(self): + rng = np.random.default_rng(13) + shape = (2, 3, 4) + packed = _random_packed(rng, shape) + st = _StackStore.from_packed(packed, shape, + index_columns=("t", "iy", "ix")) + # Full index: every position's rows match a mask over the raw buffer. + for t in range(2): + for iy in range(3): + for ix in range(4): + mask = ((packed[:, 0] == t) & (packed[:, 1] == iy) + & (packed[:, 2] == ix)) + np.testing.assert_array_equal( + st.at(t, iy, ix)["val"], packed[mask, 3]) + # Prefix: slice_at(t) and slice_at(t, iy) are O(1) outer-level reads. + for t in range(2): + np.testing.assert_array_equal( + st.slice_at(t)["val"], packed[packed[:, 0] == t, 3]) + for iy in range(3): + mask = (packed[:, 0] == t) & (packed[:, 1] == iy) + np.testing.assert_array_equal( + st.slice_at(t, iy)["val"], packed[mask, 3]) + + +class TestStructureFrozen: + def _store(self): + st = _EventStore.streaming((2, 2), index_columns=("iy", "ix")) + st.append_batch({"iy": [0, 1], "ix": [0, 1], "val": [1.0, 2.0]}) + return st.finalize() + + def test_append_after_finalize_raises(self): + st = self._store() + with pytest.raises(RuntimeError, match="frozen"): + st.append_batch({"iy": [0], "ix": [0], "val": [3.0]}) + + def test_reads_before_finalize_raise(self): + st = _EventStore.streaming((2, 2), index_columns=("iy", "ix")) + for read in (st.count_map, st.offset_levels, + lambda: st.column("val"), lambda: st.at(0, 0), + lambda: st.map("val")): + with pytest.raises(RuntimeError, match="finalized"): + read() + + def test_finalize_is_idempotent(self): + st = self._store() + leaf = st.offset_levels()[-1] + assert st.finalize() is st + assert st.offset_levels()[-1] is leaf + + def test_cell_values_stay_writable_in_place(self): + """Structure-frozen, not value-frozen: the owning subclass may + overwrite cell VALUES through the column views.""" + st = self._store() + st.column("val")[0] = 42.0 + np.testing.assert_array_equal(st.at(0, 0)["val"], [42.0]) + + def test_packed_values_writable_through_backing(self): + packed = np.array([[0, 0, 1.0], [1, 1, 2.0]], dtype=np.float32) + + class _P(RaggedStore): + columns_schema = (("iy", "f4"), ("ix", "f4"), ("val", "f4")) + + st = _P.from_packed(packed, (2, 2), index_columns=("iy", "ix")) + packed[0, 2] = 5.0 + np.testing.assert_array_equal(st.at(0, 0)["val"], [5.0]) + + +class TestBuilders: + def test_build_leaf_offsets_out_of_grid_raises(self): + with pytest.raises(ValueError, match="outside"): + build_leaf_offsets([np.array([5]), np.array([0])], (2, 2)) + + def test_derive_levels_shapes(self): + leaf = np.arange(0, 25, dtype=np.int64) # (3*2*4)+1 = 25 entries + levels = derive_levels(leaf, (3, 2, 4)) + assert [len(v) for v in levels] == [4, 7, 25] + assert levels[-1] is leaf + np.testing.assert_array_equal(levels[0], leaf[::8]) + np.testing.assert_array_equal(levels[1], leaf[::4])