From 1dd46e3ed49221a2d65d3d75cf6760833ee7c167 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Wed, 9 Sep 2026 10:45:13 -0700 Subject: [PATCH 1/7] Generalize normalized views into declarative recipes (issue #40) Adapts VCSCArrayNormalized/VCSRArrayNormalized to fishidaho's proposed y = (g(x*a*b) - c) * s form (#40), with five built-in recipes (raw, cp10k_log1p, parafac2, scanpy, pearson). .normalized(view, recalculate=True) lets a caller switch recipes without recomputing statistics once each has been built; VCSCAnnData.normalized() additionally records a/b/c/s into obs/varm/uns (named vsparse_a/b/c/s) and tracks staleness across indexing, so recalculate=False can opt into carried-over (population-stale) stats instead of paying for a fresh O(nnz) pass. Matmul/matmat kernels are generalized to the same g/s parameterization so they keep working as a view for every recipe. Adds recipe-correctness tests against an independent numpy reference and benchmarks comparing the view path to materialize-then-multiply for every recipe (both timing and peak allocation), gated in baselines.json. Co-Authored-By: Claude Sonnet 5 --- benchmarks/baselines.json | 84 +++++- benchmarks/cases.py | 75 ++++++ benchmarks/run.py | 4 +- src/vsparse/__init__.py | 3 + src/vsparse/_anndata_class.py | 99 ++++++- src/vsparse/_base.py | 109 ++++---- src/vsparse/_construct.py | 4 +- src/vsparse/_io.py | 4 +- src/vsparse/_ivcsc.py | 4 +- src/vsparse/_norm_common.py | 455 ++++++++++++++++++++++++++------- src/vsparse/_ops.py | 6 +- src/vsparse/_rapid_load.py | 4 +- src/vsparse/_vcs_matmul.py | 95 +++++-- src/vsparse/_vcs_norm.py | 26 +- tests/test_vcs_norm_recipes.py | 254 ++++++++++++++++++ 15 files changed, 1031 insertions(+), 195 deletions(-) create mode 100644 tests/test_vcs_norm_recipes.py diff --git a/benchmarks/baselines.json b/benchmarks/baselines.json index 3d1117d..aed662d 100644 --- a/benchmarks/baselines.json +++ b/benchmarks/baselines.json @@ -5,7 +5,9 @@ "indices_bytes_per_nonzero": 1.1, "vs_scipy_ratio": 1.1, "peak_alloc_mb": 2.0, - "time_ratio_vs_scipy": 4.0 + "time_ratio_vs_scipy": 4.0, + "peak_alloc_mb_view": 2.0, + "time_ratio_view_over_materialize": 4.0 }, "cases": { "layout_bytes_per_nonzero": { @@ -33,6 +35,86 @@ }, "minor_selection_peak_mb": { "peak_alloc_mb": 124.5457 + }, + "normalized_cp10k_log1p_matmat_vs_materialize": { + "time_ratio_view_over_materialize": 0.1444, + "peak_alloc_mb_view": 2.7262 + }, + "normalized_cp10k_log1p_matvec_vs_materialize": { + "time_ratio_view_over_materialize": 0.0864, + "peak_alloc_mb_view": 0.3858 + }, + "normalized_cp10k_log1p_rmatmat_vs_materialize": { + "time_ratio_view_over_materialize": 0.1639, + "peak_alloc_mb_view": 3.6521 + }, + "normalized_cp10k_log1p_rmatvec_vs_materialize": { + "time_ratio_view_over_materialize": 0.094, + "peak_alloc_mb_view": 0.1321 + }, + "normalized_parafac2_matmat_vs_materialize": { + "time_ratio_view_over_materialize": 0.1286, + "peak_alloc_mb_view": 2.7262 + }, + "normalized_parafac2_matvec_vs_materialize": { + "time_ratio_view_over_materialize": 0.1264, + "peak_alloc_mb_view": 0.3858 + }, + "normalized_parafac2_rmatmat_vs_materialize": { + "time_ratio_view_over_materialize": 0.1686, + "peak_alloc_mb_view": 3.6521 + }, + "normalized_parafac2_rmatvec_vs_materialize": { + "time_ratio_view_over_materialize": 0.0781, + "peak_alloc_mb_view": 0.1321 + }, + "normalized_pearson_matmat_vs_materialize": { + "time_ratio_view_over_materialize": 0.052, + "peak_alloc_mb_view": 2.7262 + }, + "normalized_pearson_matvec_vs_materialize": { + "time_ratio_view_over_materialize": 0.0364, + "peak_alloc_mb_view": 0.3858 + }, + "normalized_pearson_rmatmat_vs_materialize": { + "time_ratio_view_over_materialize": 0.1295, + "peak_alloc_mb_view": 3.6521 + }, + "normalized_pearson_rmatvec_vs_materialize": { + "time_ratio_view_over_materialize": 0.0397, + "peak_alloc_mb_view": 0.1321 + }, + "normalized_raw_matmat_vs_materialize": { + "time_ratio_view_over_materialize": 0.0518, + "peak_alloc_mb_view": 2.7262 + }, + "normalized_raw_matvec_vs_materialize": { + "time_ratio_view_over_materialize": 0.0427, + "peak_alloc_mb_view": 0.3858 + }, + "normalized_raw_rmatmat_vs_materialize": { + "time_ratio_view_over_materialize": 0.1011, + "peak_alloc_mb_view": 3.6521 + }, + "normalized_raw_rmatvec_vs_materialize": { + "time_ratio_view_over_materialize": 0.0257, + "peak_alloc_mb_view": 0.1321 + }, + "normalized_scanpy_matmat_vs_materialize": { + "time_ratio_view_over_materialize": 0.1513, + "peak_alloc_mb_view": 2.7262 + }, + "normalized_scanpy_matvec_vs_materialize": { + "time_ratio_view_over_materialize": 0.1004, + "peak_alloc_mb_view": 0.3858 + }, + "normalized_scanpy_rmatmat_vs_materialize": { + "time_ratio_view_over_materialize": 0.1887, + "peak_alloc_mb_view": 3.6521 + }, + "normalized_scanpy_rmatvec_vs_materialize": { + "time_ratio_view_over_materialize": 0.1171, + "peak_alloc_mb_view": 0.1321 } } } diff --git a/benchmarks/cases.py b/benchmarks/cases.py index eb3be8b..1d69c62 100644 --- a/benchmarks/cases.py +++ b/benchmarks/cases.py @@ -5,6 +5,7 @@ import numpy as np from benchmarks.harness import ( + best_time, integer_counts_csr, peak_alloc_mb, ratio_vs_scipy, @@ -139,6 +140,80 @@ def matmat_vs_scipy() -> dict[str, float]: return {"time_ratio_vs_scipy": ratio_vs_scipy(lambda: v @ B, lambda: mat @ B)} +# -- normalized views (issue #40 recipes): view-op vs materialize-then-op --- +# +# For every recipe, the view-based matmul/matvec should cost less, both in +# time and in peak allocation, than fully materializing the (dense, +# implicit-zero-filling) normalized matrix and multiplying that -- the whole +# point of a *view*. ``time_ratio_view_over_materialize`` < 1 and +# ``peak_alloc_mb_view`` < ``peak_alloc_mb_materialize`` are the expectation +# for every case below. + + +def _normalized_bench(recipe: str, *, vector: bool) -> Callable[[], dict[str, float]]: + def bench() -> dict[str, float]: + from vsparse import VCSRArray + + mat = integer_counts_csr(20_000, 2_000, density=0.05) + v = VCSRArray.from_scipy(mat) + nv = v.normalized(recipe) + rng = np.random.default_rng(0) + B = rng.normal(size=mat.shape[1]) if vector else rng.normal(size=(mat.shape[1], 8)) + + def via_view() -> np.ndarray: + return nv @ B + + def via_materialize() -> np.ndarray: + return nv.toarray() @ B + + return { + "time_ratio_view_over_materialize": best_time(via_view) / best_time(via_materialize), + "peak_alloc_mb_view": peak_alloc_mb(via_view), + "peak_alloc_mb_materialize": peak_alloc_mb(via_materialize), + } + + bench.__name__ = f"normalized_{recipe}_{'matvec' if vector else 'matmat'}_vs_materialize" + return bench + + +def _normalized_rbench(recipe: str, *, vector: bool) -> Callable[[], dict[str, float]]: + def bench() -> dict[str, float]: + from vsparse import VCSCArray + + mat = integer_counts_csr(20_000, 2_000, density=0.05) + v = VCSCArray.from_scipy(mat) + nv = v.normalized(recipe) + rng = np.random.default_rng(0) + B = rng.normal(size=mat.shape[0]) if vector else rng.normal(size=(8, mat.shape[0])) + + def via_view() -> np.ndarray: + return B @ nv + + def via_materialize() -> np.ndarray: + return B @ nv.toarray() + + return { + "time_ratio_view_over_materialize": best_time(via_view) / best_time(via_materialize), + "peak_alloc_mb_view": peak_alloc_mb(via_view), + "peak_alloc_mb_materialize": peak_alloc_mb(via_materialize), + } + + bench.__name__ = f"normalized_{recipe}_{'rmatvec' if vector else 'rmatmat'}_vs_materialize" + return bench + + +def _register_normalized_benchmarks() -> None: + from vsparse import RECIPES + + for recipe in sorted(RECIPES): + for vector in (False, True): + fast(_normalized_bench(recipe, vector=vector)) + fast(_normalized_rbench(recipe, vector=vector)) + + +_register_normalized_benchmarks() + + # -- larger, for the scheduled job ------------------------------------------- diff --git a/benchmarks/run.py b/benchmarks/run.py index e1143e6..50cda8d 100644 --- a/benchmarks/run.py +++ b/benchmarks/run.py @@ -42,9 +42,7 @@ def _compare(results: dict[str, dict[str, float]], baselines: dict) -> list[str] if ceiling is None: continue # recorded for context, not gated if value > ceiling: - failures.append( - f"{case}.{metric}: {value:.4g} exceeds ceiling {ceiling:.4g}" - ) + failures.append(f"{case}.{metric}: {value:.4g} exceeds ceiling {ceiling:.4g}") return failures diff --git a/src/vsparse/__init__.py b/src/vsparse/__init__.py index 4baeb5a..dce69d8 100644 --- a/src/vsparse/__init__.py +++ b/src/vsparse/__init__.py @@ -9,10 +9,13 @@ from vsparse._anndata import from_anndata, to_layer from vsparse._anndata_class import VCSCAnnData from vsparse._base import VCSCArray, VCSRArray +from vsparse._norm_common import RECIPES, Recipe from vsparse._rapid_load import load_and_normalize, load_packed from vsparse._vcs_norm import VCSCArrayNormalized, VCSRArrayNormalized __all__ = [ + "RECIPES", + "Recipe", "VCSCAnnData", "VCSCArray", "VCSCArrayNormalized", diff --git a/src/vsparse/_anndata_class.py b/src/vsparse/_anndata_class.py index 58e0ca2..1ff2be0 100644 --- a/src/vsparse/_anndata_class.py +++ b/src/vsparse/_anndata_class.py @@ -12,6 +12,8 @@ from vsparse import _compression, _io from vsparse._base import VCSCArray, VCSRArray, _VCSBase +from vsparse._norm_common import DEFAULT_RECIPE, resolve_recipe +from vsparse._vcs_norm import VCSCArrayNormalized, VCSRArrayNormalized if TYPE_CHECKING: from collections.abc import Mapping @@ -26,6 +28,13 @@ _FIELD_KEYS = (*_DF_KEYS, *_MAPPING_KEYS) _STORE_FORMATS = ("vcsc", "ivcsc") +# Names statistics are recorded under in .obs/.varm/.uns -- see .normalized(). +_VSPARSE_UNS_KEY = "vsparse" +_VSPARSE_OBS_A = "vsparse_a" +_VSPARSE_VARM_B = "vsparse_b" +_VSPARSE_VARM_C = "vsparse_c" +_VSPARSE_VARM_S = "vsparse_s" + def _as_slice_index(idx: Any, n: int) -> Any: """Turn a bare int index into a length-1 slice, matching anndata's own convention.""" @@ -108,6 +117,7 @@ def __init__( ) self._vcs_X: _AnyVCS | None = None self._vcs_raw_X: _AnyVCS | None = None + self._vcs_norm_cache: dict[str, Any] = {} shape = kwargs.pop("shape", None) if shape is None and X is None and "obs" not in kwargs: shape = (0, 0) @@ -131,7 +141,12 @@ def X(self, value: Any) -> None: value = vcls.from_scipy(value) else: _check_vcs_type(value, "X") - if value is not None and hasattr(self, "_obs") and hasattr(self, "_var") and value.shape != self.shape: + if ( + value is not None + and hasattr(self, "_obs") + and hasattr(self, "_var") + and value.shape != self.shape + ): raise ValueError(f"X shape {value.shape} does not match adata shape {self.shape}") self._vcs_X = value @@ -169,12 +184,25 @@ def __getitem__(self, index: Any) -> VCSCAnnData: # ty: ignore[invalid-method-o obs = cast(pd.DataFrame, self.obs).iloc[oidx].copy() var = cast(pd.DataFrame, self.var).iloc[vidx].copy() + # obs["vsparse_a"]/varm["vsparse_b"/"c"/"s"] get windowed along with + # obs/varm below, same as any other per-cell/per-gene column -- each + # kept cell/gene keeps its own precomputed factor. But population-level + # aspects of those statistics (e.g. the median depth or per-gene mean + # baked into them) were derived from the *pre-subset* population, so + # they no longer reflect this narrower one -- mark them stale rather + # than silently pretending they were computed fresh. Call + # .normalized(recalculate=True) to recompute for this subset, or + # recalculate=False to keep using these (stale) values. + uns = self.uns + if _VSPARSE_UNS_KEY in uns: + uns = {**uns, _VSPARSE_UNS_KEY: {**uns[_VSPARSE_UNS_KEY], "stale": True}} + return VCSCAnnData( X=_subset_2d(self._vcs_X, oidx, vidx), raw_X=_subset_2d(self._vcs_raw_X, oidx, vidx), obs=obs, var=var, - uns=self.uns, + uns=uns, obsm={k: _subset_1d(v, oidx) for k, v in self.obsm.items() if k is not None}, varm={k: _subset_1d(v, vidx) for k, v in self.varm.items() if k is not None}, obsp={k: _subset_2d(v, oidx, oidx) for k, v in self.obsp.items() if k is not None}, @@ -182,6 +210,73 @@ def __getitem__(self, index: Any) -> VCSCAnnData: # ty: ignore[invalid-method-o layers={k: _subset_2d(v, oidx, vidx) for k, v in self.layers.items() if k is not None}, ) + # -- normalization ---------------------------------------------------------- + + def normalized(self, view: str = DEFAULT_RECIPE, *, recalculate: bool = True) -> Any: + """A normalized view of ``X`` -- see :meth:`vsparse._base._VCSBase.normalized`. + + Also records the recipe's statistics -- per-cell ``a`` in + ``obs["vsparse_a"]``, per-gene ``b``/``c``/``s`` in + ``varm["vsparse_b"]``/``["vsparse_c"]``/``["vsparse_s"]``, and the + active recipe name plus a ``stale`` flag in ``uns["vsparse"]`` -- so + they persist across :meth:`write_h5ad`/:meth:`write_zarr` and can be + reused instead of recomputed. + + ``recalculate=False`` reuses whatever is already recorded for + ``view``: first an in-memory view already built this session (however + it was built), else what's stored in ``obs``/``varm``/``uns`` (built + fresh, but skipping the ``O(nnz)`` statistics passes) if it matches + ``view`` and the current shape. ``uns["vsparse"]["stale"]`` is set to + ``True`` after indexing (see :meth:`__getitem__`), since population + statistics like a depth median or a per-gene mean no longer reflect + the subset -- ``recalculate=False`` reuses them anyway; the default + ``recalculate=True`` always recomputes fresh (and clears ``stale``). + """ + if self._vcs_X is None: + raise ValueError("normalized() requires X to be set") + recipe = resolve_recipe(view) + cache = self._vcs_norm_cache + if not recalculate: + cached = cache.get(recipe.name) + if cached is not None: + return cached + stored = self.uns.get(_VSPARSE_UNS_KEY) + if ( + stored is not None + and stored.get("recipe") == recipe.name + and _VSPARSE_OBS_A in self.obs + and len(self.obs[_VSPARSE_OBS_A]) == self.n_obs + and _VSPARSE_VARM_B in self.varm + and _VSPARSE_VARM_C in self.varm + and _VSPARSE_VARM_S in self.varm + and len(self.varm[_VSPARSE_VARM_B]) == self.n_vars + ): + nview_cls = ( + VCSCArrayNormalized + if isinstance(self._vcs_X, VCSCArray) + else VCSRArrayNormalized + ) + nview = nview_cls.from_stats( + self._vcs_X, + recipe, + a=np.asarray(self.obs[_VSPARSE_OBS_A], dtype=np.float64), + b=np.asarray(self.varm[_VSPARSE_VARM_B], dtype=np.float64).reshape(-1), + c=np.asarray(self.varm[_VSPARSE_VARM_C], dtype=np.float64).reshape(-1), + s=np.asarray(self.varm[_VSPARSE_VARM_S], dtype=np.float64).reshape(-1), + stale=bool(stored.get("stale", False)), + ) + cache[recipe.name] = nview + return nview + + nview = self._vcs_X.normalized(recipe.name, recalculate=True) + cache[recipe.name] = nview + self.obs[_VSPARSE_OBS_A] = nview.a + self.varm[_VSPARSE_VARM_B] = nview.b + self.varm[_VSPARSE_VARM_C] = nview.c + self.varm[_VSPARSE_VARM_S] = nview.s + self.uns[_VSPARSE_UNS_KEY] = {"recipe": recipe.name, "stale": False} + return nview + # -- conversion ------------------------------------------------------------- @classmethod diff --git a/src/vsparse/_base.py b/src/vsparse/_base.py index 055b003..64c60ed 100644 --- a/src/vsparse/_base.py +++ b/src/vsparse/_base.py @@ -21,6 +21,7 @@ from vsparse._indexutils import is_full_slice as _is_full_slice from vsparse._indexutils import normalize_major_idx as _normalize_major_idx from vsparse._indexutils import smallest_index_dtype as _smallest_index_dtype +from vsparse._norm_common import DEFAULT_RECIPE as _DEFAULT_RECIPE __all__ = ["VCSCArray", "VCSRArray"] @@ -39,7 +40,7 @@ class _VCSBase: # __rmatmul__/__rmul__ instead of trying to broadcast us as an ndarray. __array_ufunc__ = None - __slots__ = ("indices", "major_ptr", "shape", "value_ptr", "values") + __slots__ = ("_norm_cache", "indices", "major_ptr", "shape", "value_ptr", "values") def __init__( self, @@ -56,9 +57,7 @@ def __init__( indices = np.asarray(indices) if major_ptr.shape[0] != n_major + 1: - raise ValueError( - f"major_ptr has length {major_ptr.shape[0]}, expected {n_major + 1}" - ) + raise ValueError(f"major_ptr has length {major_ptr.shape[0]}, expected {n_major + 1}") if value_ptr.shape[0] != values.shape[0] + 1: raise ValueError("value_ptr must have length len(values) + 1") if major_ptr[-1] != values.shape[0]: @@ -79,6 +78,7 @@ def __init__( self.values = values self.value_ptr = value_ptr self.indices = indices + self._norm_cache: dict[str, Any] = {} # -- axis bookkeeping ------------------------------------------------ @@ -112,8 +112,7 @@ def n_unique(self) -> int: def __repr__(self) -> str: # pragma: no cover - cosmetic cls = type(self).__name__ return ( - f"<{cls} shape={self.shape} dtype={self.dtype} " - f"nnz={self.nnz} n_unique={self.n_unique}>" + f"<{cls} shape={self.shape} dtype={self.dtype} nnz={self.nnz} n_unique={self.n_unique}>" ) def copy(self): @@ -200,12 +199,38 @@ def _transpose_major(self) -> _VCSBase: ) return other_cls(self.shape, major_ptr, values, value_ptr, indices) - def normalized(self) -> Any: - """A read-depth-normalized, log-transformed, mean-centered *view* -- see :mod:`vsparse._vcs_norm`.""" + def normalized(self, view: str = _DEFAULT_RECIPE, *, recalculate: bool = True) -> Any: + """A normalized *view* of this array -- see :mod:`vsparse._vcs_norm`/:mod:`vsparse._norm_common`. + + Parameters + ---------- + view + Which normalization recipe to apply -- one of + :data:`vsparse._norm_common.RECIPES` (``"raw"``, ``"cp10k_log1p"``, + ``"parafac2"`` (the default), ``"scanpy"``, ``"pearson"``). + recalculate + If ``True`` (the default), (re)compute the recipe's statistics + fresh from this array. If ``False``, reuse a previously computed + view for ``view`` on this same array, if one exists (from an + earlier ``.normalized(view, ...)`` call, with any value of + ``recalculate``) -- this is how switching between recipes avoids + recomputing each one every time. If no such view has been + computed yet, one is still computed (there is nothing to reuse). + """ + from vsparse._norm_common import resolve_recipe from vsparse._vcs_norm import VCSCArrayNormalized, VCSRArrayNormalized + recipe = resolve_recipe(view) + cache = self._norm_cache + if not recalculate: + cached = cache.get(recipe.name) + if cached is not None: + return cached + cls = VCSCArrayNormalized if self._format == "csc" else VCSRArrayNormalized - return cls(self) + result = cls(self, recipe) + cache[recipe.name] = result + return result def log1p(self) -> _VCSBase: """Elementwise ``log1p``. Structural zeros stay zero implicitly.""" @@ -229,7 +254,10 @@ def _major_sums(self) -> np.ndarray: def _minor_sums(self) -> np.ndarray: """A parallel scatter-add over every nonzero to get per-minor-index totals.""" return _ops.minor_sums( - self.values, self.value_ptr, self.indices, self.n_minor, + self.values, + self.value_ptr, + self.indices, + self.n_minor, _ops.accumulator_threads(self.n_minor), ) @@ -255,7 +283,9 @@ def _major_nnz(self) -> np.ndarray: def _minor_nnz(self) -> np.ndarray: """A parallel scatter over every nonzero to get per-minor-index counts.""" return _ops.minor_counts( - self.value_ptr, self.indices, self.n_minor, + self.value_ptr, + self.indices, + self.n_minor, _ops.accumulator_threads(self.n_minor), ) @@ -355,9 +385,7 @@ def __add__(self, other): if np.isscalar(other): if other == 0: return self.copy() - raise NotImplementedError( - "adding a nonzero scalar to a sparse array is not supported" - ) + raise NotImplementedError("adding a nonzero scalar to a sparse array is not supported") return self._elementwise(other, "__add__") __radd__ = __add__ @@ -465,15 +493,11 @@ def __matmul__(self, other): other_arr = np.asarray(other) if other_arr.ndim == 1: if other_arr.shape[0] != self.shape[1]: - raise ValueError( - f"shapes {self.shape} and {other_arr.shape} not aligned" - ) + raise ValueError(f"shapes {self.shape} and {other_arr.shape} not aligned") return self._dot_right(other_arr) if other_arr.ndim == 2: if other_arr.shape[0] != self.shape[1]: - raise ValueError( - f"shapes {self.shape} and {other_arr.shape} not aligned" - ) + raise ValueError(f"shapes {self.shape} and {other_arr.shape} not aligned") return self._dot_right_mat(other_arr) return NotImplemented @@ -481,15 +505,11 @@ def __rmatmul__(self, other): other_arr = np.asarray(other) if other_arr.ndim == 1: if other_arr.shape[0] != self.shape[0]: - raise ValueError( - f"shapes {other_arr.shape} and {self.shape} not aligned" - ) + raise ValueError(f"shapes {other_arr.shape} and {self.shape} not aligned") return self._dot_left(other_arr) if other_arr.ndim == 2: if other_arr.shape[1] != self.shape[0]: - raise ValueError( - f"shapes {other_arr.shape} and {self.shape} not aligned" - ) + raise ValueError(f"shapes {other_arr.shape} and {self.shape} not aligned") return self._dot_left_mat(other_arr) return NotImplemented @@ -503,9 +523,11 @@ def _select_major(self, key: Any) -> _VCSBase: new_major_ptr = np.zeros(idx.shape[0] + 1, dtype=np.int64) np.cumsum(counts, out=new_major_ptr[1:]) - value_slots = np.concatenate( - [np.arange(s, e) for s, e in zip(starts, ends, strict=True)] - ) if idx.shape[0] else np.empty(0, dtype=np.int64) + value_slots = ( + np.concatenate([np.arange(s, e) for s, e in zip(starts, ends, strict=True)]) + if idx.shape[0] + else np.empty(0, dtype=np.int64) + ) new_values = self.values[value_slots] v_starts = self.value_ptr[value_slots] @@ -513,14 +535,14 @@ def _select_major(self, key: Any) -> _VCSBase: idx_counts = v_ends - v_starts new_value_ptr = np.zeros(value_slots.shape[0] + 1, dtype=np.int64) np.cumsum(idx_counts, out=new_value_ptr[1:]) - new_indices = np.concatenate( - [self.indices[s:e] for s, e in zip(v_starts, v_ends, strict=True)] - ) if value_slots.shape[0] else np.empty(0, dtype=self.indices.dtype) + new_indices = ( + np.concatenate([self.indices[s:e] for s, e in zip(v_starts, v_ends, strict=True)]) + if value_slots.shape[0] + else np.empty(0, dtype=self.indices.dtype) + ) n_minor = self.n_minor - new_shape = ( - (n_minor, idx.shape[0]) if self._format == "csc" else (idx.shape[0], n_minor) - ) + new_shape = (n_minor, idx.shape[0]) if self._format == "csc" else (idx.shape[0], n_minor) return type(self)(new_shape, new_major_ptr, new_values, new_value_ptr, new_indices) def _major_range(self, start: int, stop: int) -> _VCSBase: @@ -532,9 +554,7 @@ def _major_range(self, start: int, stop: int) -> _VCSBase: u0, u1 = int(self.major_ptr[start]), int(self.major_ptr[stop]) k0, k1 = int(self.value_ptr[u0]), int(self.value_ptr[u1]) n_sel = stop - start - new_shape = ( - (self.n_minor, n_sel) if self._format == "csc" else (n_sel, self.n_minor) - ) + new_shape = (self.n_minor, n_sel) if self._format == "csc" else (n_sel, self.n_minor) return type(self)( new_shape, self.major_ptr[start : stop + 1] - u0, @@ -569,8 +589,13 @@ def _select_minor(self, key: Any) -> _VCSBase: # Keep the parent's index dtype, as every other structural op does. new_indices = np.empty(int(new_value_ptr[-1]), dtype=self.indices.dtype) _ops.minor_select_fill( - self.value_ptr, self.indices, offsets, positions, - kept_slots, new_value_ptr, new_indices, + self.value_ptr, + self.indices, + offsets, + positions, + kept_slots, + new_value_ptr, + new_indices, ) # Slots keep their original order, so each major slice owns a @@ -598,9 +623,7 @@ def __getitem__(self, key): if isinstance(row_key, int | np.integer) and isinstance(col_key, int | np.integer): return self.to_scipy()[row_key, col_key] - major_key, minor_key = ( - (col_key, row_key) if self._format == "csc" else (row_key, col_key) - ) + major_key, minor_key = (col_key, row_key) if self._format == "csc" else (row_key, col_key) if _is_full_slice(major_key) and _is_full_slice(minor_key): return self.copy() diff --git a/src/vsparse/_construct.py b/src/vsparse/_construct.py index 958bdd6..3ca22bc 100644 --- a/src/vsparse/_construct.py +++ b/src/vsparse/_construct.py @@ -195,7 +195,9 @@ def transpose_major( values_out = sorted_value[group_starts] value_ptr_out = np.concatenate([group_starts, [nnz]]).astype(np.int64) - major_ptr_out = np.searchsorted(sorted_major[group_starts], np.arange(n_minor + 1)).astype(np.int64) + major_ptr_out = np.searchsorted(sorted_major[group_starts], np.arange(n_minor + 1)).astype( + np.int64 + ) # The output's minor axis is the input's major axis, so that's the bound. indices_out = sorted_minor.astype(smallest_index_dtype(n_major), copy=False) diff --git a/src/vsparse/_io.py b/src/vsparse/_io.py index 002f8c8..d7220c9 100644 --- a/src/vsparse/_io.py +++ b/src/vsparse/_io.py @@ -60,9 +60,7 @@ def _make_read_vcs(cls: type[_VCSBase]): def _read(elem: GroupStorageType, *, _reader: Reader) -> _VCSBase: shape_vals = [int(s) for s in np.asarray(elem.attrs["shape"]).tolist()] shape = (shape_vals[0], shape_vals[1]) - arrays = { - name: cast(np.ndarray, _reader.read_elem(elem[name])) for name in _ARRAY_KEYS - } + arrays = {name: cast(np.ndarray, _reader.read_elem(elem[name])) for name in _ARRAY_KEYS} return cls( shape, arrays["major_ptr"], diff --git a/src/vsparse/_ivcsc.py b/src/vsparse/_ivcsc.py index 1831fce..c779d94 100644 --- a/src/vsparse/_ivcsc.py +++ b/src/vsparse/_ivcsc.py @@ -236,7 +236,9 @@ def _group_chunk_boundaries( return chunk_group, chunk_byte -def _unpack_parallel(value_ptr: np.ndarray, buf: np.ndarray, out: np.ndarray, n_chunks: int) -> None: +def _unpack_parallel( + value_ptr: np.ndarray, buf: np.ndarray, out: np.ndarray, n_chunks: int +) -> None: chunk_group, chunk_byte = _group_chunk_boundaries(value_ptr, buf, n_chunks) _decode_chunks(value_ptr, buf, out, chunk_group, chunk_byte) diff --git a/src/vsparse/_norm_common.py b/src/vsparse/_norm_common.py index e02eac3..625ee6a 100644 --- a/src/vsparse/_norm_common.py +++ b/src/vsparse/_norm_common.py @@ -1,46 +1,47 @@ """Shared statistics/materialization logic for normalized VCS views. :mod:`vsparse._vcs_norm` (:class:`~vsparse.VCSCArrayNormalized`/:class:`~vsparse. -VCSRArrayNormalized`) wraps a raw VCS array and behaves like the -read-depth-normalized, log-transformed, mean-centered matrix that -:func:`vsparse._rapid_load.load_and_normalize` builds -- without ever -materializing it. Centering makes every implicit structural zero a nonzero -value, so a real materialization is an ``n_rows * n_cols`` dense array; the -whole point of a "view" here is to avoid paying for that until (and unless) -the caller actually asks for it. - -Statistics are computed once, at construction, over the whole wrapped -array. Indexing a view windows that matrix and keeps those statistics; -:meth:`~NormalizedViewBase.select` renormalizes a subset on its own terms. - -What's precomputed, once, at construction: - -- ``row_scale``: per-row (cell) total raw counts, scaled to a median of 1 -- - the read-depth normalization factor. -- ``gene_scale``: per-column (gene) sum of read-depth-scaled raw counts -- - the per-gene normalization factor used inside the log transform. -- ``col_mean``: the mean, over *all* rows (including implicit zeros, whose - transformed value is exactly ``log10(1) == 0``), of the transformed value - in that column -- the centering offset. - -Given those three (all ``O(n_rows)``/``O(n_cols)``-sized, not ``O(nnz)``), -any entry's final value is ``log10(1 + 1000 * raw / row_scale / gene_scale) -- col_mean``, computable independently per entry. Computing the statistics -themselves still requires touching every nonzero (twice: once for -``gene_scale``, once more for ``col_mean``, since the transform needs -``gene_scale`` first) -- exactly like :mod:`vsparse._rapid_load`'s reference -implementation -- so that part is done with parallel numba kernels below, -specialized per storage format: +VCSRArrayNormalized`) wraps a raw VCS array and behaves like a +read-depth-normalized, transformed matrix -- without ever materializing it. +The transform follows the declarative recipe proposed in +https://github.com/meyer-lab/vsparse/issues/40#issuecomment-5546322485:: + + y[i, j] = (g(x[i, j] * a[i] * b[j]) - c[j]) * s[j] + +where ``a`` is a per-cell scale (depth normalization), ``b`` a per-gene +scale, ``g`` a monotone scalar function with ``g(0) == 0``, ``c`` a per-gene +center, and ``s`` a per-gene post-scale. :data:`RECIPES` collects several +common instantiations of this shape (``"raw"``, ``"cp10k_log1p"``, +``"parafac2"``, ``"scanpy"``, ``"pearson"``). + +Three properties of this shape keep it representable as a view rather than a +materialized ``n_rows * n_cols`` array: + +1. **Separability** -- ``a`` depends only on the cell index, ``b``/``c``/``s`` + only on the gene index, so the whole state of a normalization is + ``O(n_cells + n_genes)``, never ``O(n_cells * n_genes)``. +2. **Entrywise independence** -- ``y[i, j]`` depends only on ``x[i, j]``, so + the transform fuses into any kernel already visiting a nonzero. +3. **Rank-1 dense part** -- because ``g(0) = 0``, every structural (implicit) + zero maps to exactly ``-c[j] * s[j]``, so the full dense matrix is a rank-1 + "baseline" plus a sparse "delta" that is zero off the stored nonzeros. + +Statistics (``a``, ``b``, ``c``, ``s``) are computed once, at construction, +over the whole wrapped array, and cached on the view -- see +:meth:`~NormalizedViewBase.select`/:meth:`_VCSBase.normalized` for how a +subset or an alternate recipe is (re)computed. Computing them still requires +touching every nonzero (once for ``b``, once more for ``c``/``s``, since +those need ``b`` first), done with parallel numba kernels below, specialized +per storage format: - major=columns (VCSC): both passes collapse into one, fully parallel over columns with no cross-thread writes -- each column's own elements - carry everything needed to compute both its ``gene_scale`` and its - ``col_mean`` (:func:`_column_stats_major_is_col`). + carry everything needed to compute both its ``b`` and its ``c``/``s`` + (:func:`_column_stats_major_is_col`, :func:`_column_gstats_major_is_col`). - major=rows (VCSR): each pass is a scatter-add across columns from - many rows, so it's parallelized like :mod:`vsparse._rapid_load`'s - ``_scaled_col_sums`` -- row-chunked with thread-local partial column - arrays, reduced by summing across threads (:func:`_scaled_col_sums_vcs`, - :func:`_transformed_col_sums_vcs`). + many rows, so it's parallelized row-chunked with thread-local partial + column arrays, reduced by summing across threads + (:func:`_scaled_col_sums_vcs`, :func:`_gstats_col_sums_vcs`). These kernels only need ``major_ptr``/``values``/``value_ptr``/``indices`` arrays, from the plain (:mod:`vsparse._base`) array types. The matmul kernels @@ -51,38 +52,127 @@ from __future__ import annotations +from dataclasses import dataclass from typing import Any import numba import numpy as np -__all__ = ["NormalizedViewBase"] +__all__ = ["DEFAULT_RECIPE", "RECIPES", "NormalizedViewBase", "Recipe", "resolve_recipe"] + + +# -- recipes ------------------------------------------------------------------ + +G_IDENTITY = 0 +G_LOG1P = 1 +G_LOG1P_1000X = 2 +G_SQRT = 3 + + +@dataclass(frozen=True, slots=True) +class Recipe: + """One instantiation of ``y = (g(x * a * b) - c) * s``.""" + + name: str + #: Target used by ``a``: ``None`` -> ``a = 1`` for every cell (no depth + #: normalization); ``"median"`` -> ``a = median(depth) / depth``; + #: a float -> ``a = target / depth``. + depth_target: float | str | None + #: Whether ``b`` is ``1 / sum_i(x[i, j] * a[i])`` (``True``) or ``1`` (``False``). + gene_scale: bool + g_code: int + center: bool + post_scale: bool + + +RECIPES: dict[str, Recipe] = { + "raw": Recipe("raw", None, False, G_IDENTITY, False, False), + "cp10k_log1p": Recipe("cp10k_log1p", 1e4, False, G_LOG1P, False, False), + "parafac2": Recipe("parafac2", "median", True, G_LOG1P_1000X, True, False), + "scanpy": Recipe("scanpy", 1e4, False, G_LOG1P, True, True), + "pearson": Recipe("pearson", 1.0, True, G_SQRT, True, True), +} +DEFAULT_RECIPE = "parafac2" + + +def resolve_recipe(view: str | Recipe) -> Recipe: + if isinstance(view, Recipe): + return view + try: + return RECIPES[view] + except KeyError: + raise ValueError( + f"unknown normalization view {view!r}; choose from {sorted(RECIPES)}" + ) from None + + +# -- elementwise transform ---------------------------------------------------- + + +@numba.njit(cache=True, inline="always") +def _g(x: float, g_code: int) -> float: + if g_code == G_LOG1P: + return np.log1p(x) + if g_code == G_LOG1P_1000X: + return np.log10(1.0 + 1000.0 * x) + if g_code == G_SQRT: + return np.sqrt(x) if x > 0.0 else 0.0 + return x + + +def _g_np(x: np.ndarray, g_code: int) -> np.ndarray: + """Vectorized (non-numba) counterpart of :func:`_g`, for the ``__getitem__`` path.""" + if g_code == G_LOG1P: + return np.log1p(x) + if g_code == G_LOG1P_1000X: + return np.log10(1.0 + 1000.0 * x) + if g_code == G_SQRT: + return np.sqrt(np.clip(x, 0.0, None)) + return x # -- statistics: major=columns -- fused, no cross-thread writes ------------- @numba.njit(cache=True, parallel=True) -def _column_stats_major_is_col(major_ptr, values, value_ptr, indices, row_scale, n_rows): +def _column_stats_major_is_col(major_ptr, values, value_ptr, indices, row_scale): + """``gsum[j] = sum_i values[i, j] / row_scale[i]`` -- the raw material for ``b``.""" n_major = major_ptr.shape[0] - 1 - gene_scale = np.zeros(n_major, dtype=np.float64) - col_mean = np.zeros(n_major, dtype=np.float64) + gsum = np.zeros(n_major, dtype=np.float64) for j in numba.prange(n_major): # ty: ignore[not-iterable] gs = 0.0 for u in range(major_ptr[j], major_ptr[j + 1]): v = values[u] for k in range(value_ptr[u], value_ptr[u + 1]): gs += v / row_scale[indices[k]] - gene_scale[j] = gs - if gs > 0.0 and n_rows > 0: - cs = 0.0 - for u in range(major_ptr[j], major_ptr[j + 1]): - v = values[u] - for k in range(value_ptr[u], value_ptr[u + 1]): - scaled = v / row_scale[indices[k]] / gs - cs += np.log10(1.0 + 1000.0 * scaled) - col_mean[j] = cs / n_rows - return gene_scale, col_mean + gsum[j] = gs + return gsum + + +@numba.njit(cache=True, parallel=True) +def _column_gstats_major_is_col( + major_ptr, values, value_ptr, indices, row_scale, gene_scale, g_code +): + """Per-column sum/sum-of-squares of ``g(x / row_scale / gene_scale)`` over stored entries.""" + n_major = major_ptr.shape[0] - 1 + col_sum = np.zeros(n_major, dtype=np.float64) + col_sumsq = np.zeros(n_major, dtype=np.float64) + for j in numba.prange(n_major): # ty: ignore[not-iterable] + gs = gene_scale[j] + if gs <= 0.0: + continue + s0 = 0.0 + s1 = 0.0 + for u in range(major_ptr[j], major_ptr[j + 1]): + v = values[u] + for k in range(value_ptr[u], value_ptr[u + 1]): + scaled = v / row_scale[indices[k]] / gs + gy = _g(scaled, g_code) + s0 += gy + s1 += gy * gy + col_sum[j] = s0 + col_sumsq[j] = s1 + return col_sum, col_sumsq # -- statistics: major=rows -- scatter-add passes ---------------------------- @@ -107,16 +197,18 @@ def _scaled_col_sums_vcs(major_ptr, values, value_ptr, indices, row_scale, n_col @numba.njit(cache=True, parallel=True) -def _transformed_col_sums_vcs( - major_ptr, values, value_ptr, indices, row_scale, gene_scale, n_cols, nthreads +def _gstats_col_sums_vcs( + major_ptr, values, value_ptr, indices, row_scale, gene_scale, g_code, n_cols, nthreads ): n_major = major_ptr.shape[0] - 1 chunk = (n_major + nthreads - 1) // nthreads - partial = np.zeros((nthreads, n_cols), dtype=np.float64) + partial_sum = np.zeros((nthreads, n_cols), dtype=np.float64) + partial_sumsq = np.zeros((nthreads, n_cols), dtype=np.float64) for t in numba.prange(nthreads): # ty: ignore[not-iterable] start = t * chunk end = min(n_major, start + chunk) - local = partial[t] + local_sum = partial_sum[t] + local_sumsq = partial_sumsq[t] for i in range(start, end): rs = row_scale[i] for u in range(major_ptr[i], major_ptr[i + 1]): @@ -126,37 +218,62 @@ def _transformed_col_sums_vcs( gs = gene_scale[c] if gs > 0.0: scaled = v / rs / gs - local[c] += np.log10(1.0 + 1000.0 * scaled) - return partial.sum(axis=0) + gy = _g(scaled, g_code) + local_sum[c] += gy + local_sumsq[c] += gy * gy + return partial_sum.sum(axis=0), partial_sumsq.sum(axis=0) # -- full materialization ----------------------------------------------------- # -# Both kernels start from an ``out`` already filled with ``-col_mean`` -# (the value every implicit structural zero takes) and only overwrite the -# entries that are actually stored -- each parallelized over the major axis, -# which owns disjoint rows (major=rows) or columns (major=columns) of -# ``out``, so there's no cross-thread write. +# Both kernels start from an ``out`` already filled with ``-col_offset`` +# (``col_post_scale * col_mean`` -- the value every implicit structural zero +# takes) and only overwrite the entries that are actually stored -- each +# parallelized over the major axis, which owns disjoint rows (major=rows) or +# columns (major=columns) of ``out``, so there's no cross-thread write. @numba.njit(cache=True, parallel=True) -def _fill_normalized_major_is_col(major_ptr, values, value_ptr, indices, row_scale, gene_scale, col_mean, out): +def _fill_normalized_major_is_col( + major_ptr, + values, + value_ptr, + indices, + row_scale, + gene_scale, + col_mean, + col_post_scale, + g_code, + out, +): n_major = major_ptr.shape[0] - 1 for j in numba.prange(n_major): # ty: ignore[not-iterable] gs = gene_scale[j] - if gs == 0.0: + if gs <= 0.0: continue cm = col_mean[j] + s = col_post_scale[j] for u in range(major_ptr[j], major_ptr[j + 1]): v = values[u] for k in range(value_ptr[u], value_ptr[u + 1]): r = indices[k] scaled = v / row_scale[r] / gs - out[r, j] = np.log10(1.0 + 1000.0 * scaled) - cm + out[r, j] = (_g(scaled, g_code) - cm) * s @numba.njit(cache=True, parallel=True) -def _fill_normalized_major_is_row(major_ptr, values, value_ptr, indices, row_scale, gene_scale, col_mean, out): +def _fill_normalized_major_is_row( + major_ptr, + values, + value_ptr, + indices, + row_scale, + gene_scale, + col_mean, + col_post_scale, + g_code, + out, +): n_major = major_ptr.shape[0] - 1 for i in numba.prange(n_major): # ty: ignore[not-iterable] rs = row_scale[i] @@ -165,10 +282,10 @@ def _fill_normalized_major_is_row(major_ptr, values, value_ptr, indices, row_sca for k in range(value_ptr[u], value_ptr[u + 1]): c = indices[k] gs = gene_scale[c] - if gs == 0.0: + if gs <= 0.0: continue scaled = v / rs / gs - out[i, c] = np.log10(1.0 + 1000.0 * scaled) - col_mean[c] + out[i, c] = (_g(scaled, g_code) - col_mean[c]) * col_post_scale[c] def _prep_key(key: Any) -> Any: @@ -178,53 +295,175 @@ def _prep_key(key: Any) -> Any: return key +def _compute_row_scale(arr: Any, recipe: Recipe) -> np.ndarray: + n_rows, _n_cols = arr.shape + if recipe.depth_target is None: + return np.ones(n_rows, dtype=np.float64) + + row_totals = np.asarray(arr.sum(axis=1), dtype=np.float64) + if recipe.depth_target == "median": + target = float(np.median(row_totals)) if row_totals.shape[0] else 0.0 + else: + target = float(recipe.depth_target) + if target <= 0.0: + return np.ones(n_rows, dtype=np.float64) + row_scale = row_totals / target + row_scale[row_scale == 0.0] = 1.0 # rows with no counts: avoid div-by-zero (unused otherwise) + return row_scale + + class NormalizedViewBase: """Shared implementation for the normalized VCSC/VCSR views. Subclasses fix ``_format`` (``"csc"``/``"csr"``) and supply ``__matmul__``/``__rmatmul__`` wired to :mod:`vsparse._vcs_matmul`. + + Internally, ``row_scale``/``gene_scale`` hold ``1 / a``/``1 / b`` (the + reciprocals of the recipe's ``a``/``b``) -- that's the form the numba + kernels above divide by. ``a``/``b`` (as named in the recipe) are exposed + via the :attr:`a`/:attr:`b` properties. """ _format: str __array_ufunc__ = None - __slots__ = ("_arr", "col_mean", "gene_scale", "row_scale") + __slots__ = ("_arr", "col_mean", "col_post_scale", "gene_scale", "recipe", "row_scale", "stale") - def __init__(self, arr: Any) -> None: + def __init__( + self, arr: Any, recipe: str | Recipe = DEFAULT_RECIPE, *, stale: bool = False + ) -> None: if arr._format != self._format: raise ValueError( f"{type(self).__name__} wraps a {self._format!r}-format array, " f"got {type(arr).__name__}" ) self._arr = arr + self.recipe = resolve_recipe(recipe) + self.stale = stale n_rows, n_cols = arr.shape - row_totals = np.asarray(arr.sum(axis=1), dtype=np.float64) - median = float(np.median(row_totals)) if row_totals.shape[0] else 0.0 - if median > 0.0: - row_scale = row_totals / median - row_scale[row_scale == 0.0] = 1.0 - else: - row_scale = np.ones(n_rows, dtype=np.float64) + row_scale = _compute_row_scale(arr, self.recipe) self.row_scale = row_scale indices = arr.indices # decode once; shared by both statistics passes below - if self._format == "csc": - gene_scale, col_mean = _column_stats_major_is_col( - arr.major_ptr, arr.values, arr.value_ptr, indices, row_scale, n_rows - ) + if self.recipe.gene_scale: + if self._format == "csc": + gsum = _column_stats_major_is_col( + arr.major_ptr, arr.values, arr.value_ptr, indices, row_scale + ) + else: + nthreads = numba.get_num_threads() + gsum = _scaled_col_sums_vcs( + arr.major_ptr, arr.values, arr.value_ptr, indices, row_scale, n_cols, nthreads + ) + gene_scale = gsum else: - nthreads = numba.get_num_threads() - gene_scale = _scaled_col_sums_vcs( - arr.major_ptr, arr.values, arr.value_ptr, indices, row_scale, n_cols, nthreads - ) - col_sum = _transformed_col_sums_vcs( - arr.major_ptr, arr.values, arr.value_ptr, indices, row_scale, gene_scale, n_cols, nthreads - ) - col_mean = col_sum / n_rows if n_rows > 0 else np.zeros(n_cols, dtype=np.float64) + gene_scale = np.ones(n_cols, dtype=np.float64) self.gene_scale = gene_scale + + if self.recipe.center or self.recipe.post_scale: + if self._format == "csc": + col_sum, col_sumsq = _column_gstats_major_is_col( + arr.major_ptr, + arr.values, + arr.value_ptr, + indices, + row_scale, + gene_scale, + self.recipe.g_code, + ) + else: + nthreads = numba.get_num_threads() + col_sum, col_sumsq = _gstats_col_sums_vcs( + arr.major_ptr, + arr.values, + arr.value_ptr, + indices, + row_scale, + gene_scale, + self.recipe.g_code, + n_cols, + nthreads, + ) + mean = col_sum / n_rows if n_rows > 0 else np.zeros(n_cols, dtype=np.float64) + variance = np.clip( + col_sumsq / n_rows - mean**2 if n_rows > 0 else np.zeros(n_cols), 0.0, None + ) + std = np.sqrt(variance) + col_mean = mean if self.recipe.center else np.zeros(n_cols, dtype=np.float64) + if self.recipe.post_scale: + with np.errstate(divide="ignore", invalid="ignore"): + col_post_scale = np.where(std > 0.0, 1.0 / std, 1.0) + else: + col_post_scale = np.ones(n_cols, dtype=np.float64) + else: + col_mean = np.zeros(n_cols, dtype=np.float64) + col_post_scale = np.ones(n_cols, dtype=np.float64) self.col_mean = col_mean + self.col_post_scale = col_post_scale + + @classmethod + def from_stats( + cls, + arr: Any, + recipe: str | Recipe, + a: np.ndarray, + b: np.ndarray, + c: np.ndarray, + s: np.ndarray, + *, + stale: bool = True, + ) -> NormalizedViewBase: + """Build directly from precomputed ``a``/``b``/``c``/``s``, skipping the ``O(nnz)`` passes. + + Used to carry a previously-computed normalization (e.g. from + ``.obs``/``.varm``) onto a (possibly different) array without + recalculating -- see ``recalculate=False`` on + :meth:`vsparse._base._VCSBase.normalized`. Defaults to marking the + result :attr:`stale`, since the caller is asserting these statistics + rather than deriving them from ``arr`` itself. + """ + if arr._format != cls._format: + raise ValueError( + f"{cls.__name__} wraps a {cls._format!r}-format array, got {type(arr).__name__}" + ) + self = object.__new__(cls) + self._arr = arr + self.recipe = resolve_recipe(recipe) + self.stale = stale + a = np.asarray(a, dtype=np.float64) + b = np.asarray(b, dtype=np.float64) + self.row_scale = np.where(a > 0.0, 1.0 / a, 0.0) + self.gene_scale = np.where(b > 0.0, 1.0 / b, 0.0) + self.col_mean = np.asarray(c, dtype=np.float64) + self.col_post_scale = np.asarray(s, dtype=np.float64) + # Subclasses with the ``_dual_arr`` slot (VCS views) need it initialized + # too, since ``__init__`` (which normally does) is bypassed here. + self._dual_arr = None + return self + + # -- recipe-facing statistics (a/b/c/s, as named in the issue) ------------- + + @property + def a(self) -> np.ndarray: + """Per-cell scale -- ``1 / row_scale``.""" + return np.where(self.row_scale > 0.0, 1.0 / self.row_scale, 0.0) + + @property + def b(self) -> np.ndarray: + """Per-gene scale -- ``1 / gene_scale``.""" + return np.where(self.gene_scale > 0.0, 1.0 / self.gene_scale, 0.0) + + @property + def c(self) -> np.ndarray: + """Per-gene center.""" + return self.col_mean + + @property + def s(self) -> np.ndarray: + """Per-gene post-scale.""" + return self.col_post_scale @property def shape(self) -> tuple[int, int]: @@ -235,24 +474,42 @@ def dtype(self) -> np.dtype: return np.dtype(np.float64) def __repr__(self) -> str: # pragma: no cover - cosmetic - return f"<{type(self).__name__} shape={self.shape} dtype={self.dtype}>" + stale = " stale=True" if self.stale else "" + return f"<{type(self).__name__} shape={self.shape} dtype={self.dtype} recipe={self.recipe.name!r}{stale}>" # -- materialization ------------------------------------------------------ def toarray(self) -> np.ndarray: - """The full normalized, log-transformed, mean-centered matrix, densely.""" + """The full normalized matrix, densely.""" n_rows, n_cols = self.shape - out = np.broadcast_to(-self.col_mean, (n_rows, n_cols)).copy() + baseline = -self.col_mean * self.col_post_scale + out = np.broadcast_to(baseline, (n_rows, n_cols)).copy() arr = self._arr if self._format == "csc": _fill_normalized_major_is_col( - arr.major_ptr, arr.values, arr.value_ptr, arr.indices, - self.row_scale, self.gene_scale, self.col_mean, out, + arr.major_ptr, + arr.values, + arr.value_ptr, + arr.indices, + self.row_scale, + self.gene_scale, + self.col_mean, + self.col_post_scale, + self.recipe.g_code, + out, ) else: _fill_normalized_major_is_row( - arr.major_ptr, arr.values, arr.value_ptr, arr.indices, - self.row_scale, self.gene_scale, self.col_mean, out, + arr.major_ptr, + arr.values, + arr.value_ptr, + arr.indices, + self.row_scale, + self.gene_scale, + self.col_mean, + self.col_post_scale, + self.recipe.g_code, + out, ) return out @@ -270,7 +527,7 @@ def select(self, rows: Any = slice(None), cols: Any = slice(None)) -> Any: sub: Any = self._arr[_prep_key(rows), _prep_key(cols)] if not isinstance(sub, type(self._arr)): sub = type(self._arr).from_scipy(sub) - return type(self)(sub) + return type(self)(sub, self.recipe) # -- on-the-fly elementwise access ------------------------------------------ @@ -294,17 +551,17 @@ def __getitem__(self, key: Any) -> np.ndarray: rs = np.asarray(self.row_scale)[row_key].reshape(-1, 1) gs = np.asarray(self.gene_scale)[col_key].reshape(1, -1) cm = np.asarray(self.col_mean)[col_key].reshape(1, -1) + s = np.asarray(self.col_post_scale)[col_key].reshape(1, -1) with np.errstate(divide="ignore", invalid="ignore"): scaled = np.where(gs > 0.0, dense_raw / rs / gs, 0.0) - return np.log10(1.0 + 1000.0 * scaled) - cm + return (_g_np(scaled, self.recipe.g_code) - cm) * s # -- explicitly-unsupported operations ------------------------------------ def _unsupported(self, op: str) -> Any: raise RuntimeError( - f"{op} is not supported on {type(self).__name__}. " - "Call .toarray() first if you need it." + f"{op} is not supported on {type(self).__name__}. Call .toarray() first if you need it." ) def __add__(self, other: Any) -> Any: diff --git a/src/vsparse/_ops.py b/src/vsparse/_ops.py index e81fd70..e9d3792 100644 --- a/src/vsparse/_ops.py +++ b/src/vsparse/_ops.py @@ -167,8 +167,6 @@ def minor_counts(value_ptr, indices, n_minor, nthreads): return partial.sum(axis=0) - - def major_matvec(major_ptr, values, value_ptr, indices, x, n_major, n_minor): return _major_matvec(major_ptr, values, value_ptr, indices, np.asarray(x), n_major, n_minor) @@ -223,7 +221,9 @@ def minor_select_counts(value_ptr, indices, fanout): return counts -def minor_select_fill(value_ptr, indices, offsets, positions, kept_slots, new_value_ptr, out_indices): +def minor_select_fill( + value_ptr, indices, offsets, positions, kept_slots, new_value_ptr, out_indices +): """Write the remapped minor indices for the surviving slots, in place.""" _minor_select_fill( value_ptr, indices, offsets, positions, kept_slots, new_value_ptr, out_indices diff --git a/src/vsparse/_rapid_load.py b/src/vsparse/_rapid_load.py index 3bc2864..b5e12c0 100644 --- a/src/vsparse/_rapid_load.py +++ b/src/vsparse/_rapid_load.py @@ -518,9 +518,7 @@ def load_and_normalize( del indices, data, row_indptr normalized = _normalize_and_transform(new_indptr, out_indices, out_data, n_kept_genes) - X = csr_array( - (normalized, out_indices, new_indptr), shape=(kept_rows.shape[0], n_kept_genes) - ) + X = csr_array((normalized, out_indices, new_indptr), shape=(kept_rows.shape[0], n_kept_genes)) if "obs" in kwargs and "var" in kwargs: adata = ad.AnnData(**kwargs) # ty: ignore[invalid-argument-type] diff --git a/src/vsparse/_vcs_matmul.py b/src/vsparse/_vcs_matmul.py index bf253e9..a6f16b7 100644 --- a/src/vsparse/_vcs_matmul.py +++ b/src/vsparse/_vcs_matmul.py @@ -1,13 +1,16 @@ """Parallel numba kernels for ``VCSCArrayNormalized``/``VCSRArrayNormalized`` @ dense. -For row-scale ``r``, gene-scale ``g``, and centering ``c = col_mean``, +For row-scale ``r = 1/a``, gene-scale ``gs = 1/b``, transform ``g``, center +``c = col_mean``, and post-scale ``s = col_post_scale`` (see +:mod:`vsparse._norm_common` for the recipes these come from), - A_norm = broadcast_rows(-c) + Delta, Delta[i, j] = log10(1 + 1000 * raw[i, j] / r[i] / g[j]) + A_norm = broadcast_rows(-s * c) + Delta, Delta[i, j] = s[j] * g(raw[i, j] / r[i] / gs[j]) -with ``Delta`` exactly zero off the structural nonzeros, so -``A_norm @ B = ones(n_rows) (x) (-c @ B) + Delta @ B`` and ``B @ A_norm = -(B.sum(axis=1)) (x) (-c) + B @ Delta``, with only the genuinely ``O(nnz)`` -``Delta @ B`` / ``B @ Delta`` needing a kernel. +with ``Delta`` exactly zero off the structural nonzeros (since every recipe's +``g`` has ``g(0) == 0``), so ``A_norm @ B = ones(n_rows) (x) (-(s * c) @ B) + +Delta @ B`` and ``B @ A_norm = (B.sum(axis=1)) (x) (-(s * c)) + B @ Delta``, +with only the genuinely ``O(nnz)`` ``Delta @ B`` / ``B @ Delta`` needing a +kernel. Both ``self @ B`` and ``B @ self`` need a kernel that's *major-aligned*: parallelizing safely over the major axis requires the major axis to be @@ -23,8 +26,9 @@ chunk of major slices at a time so the extra memory is one chunk's worth rather than a second copy of the array. -``row_scale``/``gene_scale``/``col_mean`` are per-row/per-column statistics, -so a chunk just takes the slice of them its own axis covers. +``row_scale``/``gene_scale``/``col_mean``/``col_post_scale`` are per-row/ +per-column statistics, so a chunk just takes the slice of them its own axis +covers. """ from __future__ import annotations @@ -34,6 +38,8 @@ import numba import numpy as np +from vsparse._norm_common import _g + if TYPE_CHECKING: from vsparse._vcs_norm import _VCSNormalizedBase @@ -44,7 +50,9 @@ @numba.njit(cache=True, parallel=True) -def _vcsr_matmul_delta(major_ptr, values, value_ptr, indices, row_scale, gene_scale, B, out): +def _vcsr_matmul_delta( + major_ptr, values, value_ptr, indices, row_scale, gene_scale, col_post_scale, g_code, B, out +): n_major = major_ptr.shape[0] - 1 k = B.shape[1] for i in numba.prange(n_major): # ty: ignore[not-iterable] @@ -55,7 +63,7 @@ def _vcsr_matmul_delta(major_ptr, values, value_ptr, indices, row_scale, gene_sc col = indices[kk] gs = gene_scale[col] if gs > 0.0: - delta = np.log10(1.0 + 1000.0 * (v / rs / gs)) + delta = col_post_scale[col] * _g(v / rs / gs, g_code) for c in range(k): out[i, c] += delta * B[col, c] @@ -75,19 +83,22 @@ def _vcsr_matmul_delta(major_ptr, values, value_ptr, indices, row_scale, gene_sc @numba.njit(cache=True, parallel=True) -def _vcsc_rmatmul_delta(major_ptr, values, value_ptr, indices, row_scale, gene_scale, Bt, out_t): +def _vcsc_rmatmul_delta( + major_ptr, values, value_ptr, indices, row_scale, gene_scale, col_post_scale, g_code, Bt, out_t +): n_major = major_ptr.shape[0] - 1 p = Bt.shape[1] for j in numba.prange(n_major): # ty: ignore[not-iterable] gs = gene_scale[j] if gs == 0.0: continue + s = col_post_scale[j] acc = out_t[j] for u in range(major_ptr[j], major_ptr[j + 1]): v = values[u] for kk in range(value_ptr[u], value_ptr[u + 1]): row = indices[kk] - delta = np.log10(1.0 + 1000.0 * (v / row_scale[row] / gs)) + delta = s * _g(v / row_scale[row] / gs, g_code) brow = Bt[row] for c in range(p): acc[c] += delta * brow[c] @@ -169,12 +180,22 @@ def normalized_at_dense(nview: _VCSNormalizedBase, other: Any) -> np.ndarray: B, squeeze = _prep_dense(other, n_cols) out = np.zeros((arr.shape[0], B.shape[1]), dtype=np.float64) + g_code = nview.recipe.g_code + # self @ B is major-aligned for VCSR; a VCSC array needs regrouping. src, bounds = _aligned_source(nview, "csr") if src is not None: _vcsr_matmul_delta( - src.major_ptr, src.values, src.value_ptr, src.indices, - nview.row_scale, nview.gene_scale, B, out, + src.major_ptr, + src.values, + src.value_ptr, + src.indices, + nview.row_scale, + nview.gene_scale, + nview.col_post_scale, + g_code, + B, + out, ) else: # Column chunk at a time: Delta @ B == sum over chunks of @@ -183,12 +204,20 @@ def normalized_at_dense(nview: _VCSNormalizedBase, other: Any) -> np.ndarray: for start, stop in bounds: chunk = arr._major_range(start, stop)._transpose_major() _vcsr_matmul_delta( - chunk.major_ptr, chunk.values, chunk.value_ptr, chunk.indices, - nview.row_scale, nview.gene_scale[start:stop], - np.ascontiguousarray(B[start:stop]), out, + chunk.major_ptr, + chunk.values, + chunk.value_ptr, + chunk.indices, + nview.row_scale, + nview.gene_scale[start:stop], + nview.col_post_scale[start:stop], + g_code, + np.ascontiguousarray(B[start:stop]), + out, ) - baseline = (-nview.col_mean) @ B # (k,): every row's implicit-zero contribution + offset = nview.col_mean * nview.col_post_scale + baseline = (-offset) @ B # (k,): every row's implicit-zero contribution out += baseline[None, :] return out[:, 0] if squeeze else out @@ -207,24 +236,42 @@ def dense_at_normalized(nview: _VCSNormalizedBase, other: Any) -> np.ndarray: Bt = np.ascontiguousarray(B2.T) # (n_rows, p) -- see the note above the kernel out_t = np.zeros((arr.shape[1], p), dtype=np.float64) + g_code = nview.recipe.g_code + # B @ self is major-aligned for VCSC; a VCSR array needs regrouping. src, bounds = _aligned_source(nview, "csc") if src is not None: _vcsc_rmatmul_delta( - src.major_ptr, src.values, src.value_ptr, src.indices, - nview.row_scale, nview.gene_scale, Bt, out_t, + src.major_ptr, + src.values, + src.value_ptr, + src.indices, + nview.row_scale, + nview.gene_scale, + nview.col_post_scale, + g_code, + Bt, + out_t, ) else: # Row chunk at a time, accumulating into the same output. for start, stop in bounds: chunk = arr._major_range(start, stop)._transpose_major() _vcsc_rmatmul_delta( - chunk.major_ptr, chunk.values, chunk.value_ptr, chunk.indices, - nview.row_scale[start:stop], nview.gene_scale, - np.ascontiguousarray(Bt[start:stop]), out_t, + chunk.major_ptr, + chunk.values, + chunk.value_ptr, + chunk.indices, + nview.row_scale[start:stop], + nview.gene_scale, + nview.col_post_scale, + g_code, + np.ascontiguousarray(Bt[start:stop]), + out_t, ) out = np.ascontiguousarray(out_t.T) - baseline = B2.sum(axis=1)[:, None] * (-nview.col_mean)[None, :] # (m, n_cols) + offset = nview.col_mean * nview.col_post_scale + baseline = B2.sum(axis=1)[:, None] * (-offset)[None, :] # (m, n_cols) out += baseline return out[0, :] if squeeze else out diff --git a/src/vsparse/_vcs_norm.py b/src/vsparse/_vcs_norm.py index 01e12a4..183b831 100644 --- a/src/vsparse/_vcs_norm.py +++ b/src/vsparse/_vcs_norm.py @@ -1,20 +1,21 @@ -"""Normalized, mean-centered *views* of :class:`~vsparse.VCSCArray`/:class:`~vsparse.VCSRArray`. +"""Normalized *views* of :class:`~vsparse.VCSCArray`/:class:`~vsparse.VCSRArray`. :class:`VCSCArrayNormalized`/:class:`VCSRArrayNormalized` wrap a plain -(unpacked) VCSC/VCSR array and behave like the read-depth-normalized, -log-transformed, mean-centered matrix that :func:`vsparse._rapid_load. -load_and_normalize` builds -- without ever materializing it. See -:mod:`vsparse._norm_common` for the shared statistics/materialization logic -(:class:`~vsparse._norm_common.NormalizedViewBase`) and :mod:`vsparse._vcs_matmul` -for the matmul kernels these use (a direct per-nonzero walk over the -already-decoded ``indices`` array). +(unpacked) VCSC/VCSR array and behave like a normalized matrix -- read-depth +normalized, transformed, optionally centered/scaled, per one of +:data:`vsparse._norm_common.RECIPES` (the default, ``"parafac2"``, matches +what :func:`vsparse._rapid_load.load_and_normalize` builds) -- without ever +materializing it. See :mod:`vsparse._norm_common` for the shared recipe/ +statistics/materialization logic (:class:`~vsparse._norm_common.NormalizedViewBase`) +and :mod:`vsparse._vcs_matmul` for the matmul kernels these use (a direct +per-nonzero walk over the already-decoded ``indices`` array). """ from __future__ import annotations from typing import TYPE_CHECKING, Any -from vsparse._norm_common import NormalizedViewBase +from vsparse._norm_common import DEFAULT_RECIPE, NormalizedViewBase, Recipe if TYPE_CHECKING: from vsparse._base import _VCSBase @@ -27,13 +28,14 @@ class _VCSNormalizedBase(NormalizedViewBase): __slots__ = ("_dual_arr",) - def __init__(self, arr: _VCSBase) -> None: - super().__init__(arr) + def __init__( + self, arr: _VCSBase, recipe: str | Recipe = DEFAULT_RECIPE, *, stale: bool = False + ) -> None: + super().__init__(arr, recipe, stale=stale) # Opposite-format copy of `arr`, cached by vsparse._vcs_matmul when # regrouping the whole array fits one chunk's budget. self._dual_arr: _VCSBase | None = None - def __matmul__(self, other: Any) -> Any: """``self @ other`` for a dense ``other`` -- see :mod:`vsparse._vcs_matmul`.""" from vsparse._vcs_matmul import normalized_at_dense diff --git a/tests/test_vcs_norm_recipes.py b/tests/test_vcs_norm_recipes.py new file mode 100644 index 0000000..d9ef897 --- /dev/null +++ b/tests/test_vcs_norm_recipes.py @@ -0,0 +1,254 @@ +"""Tests for normalization recipes (issue #40): multiple views, caching, and staleness.""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest +import scipy.sparse as sp + +from vsparse import ( + RECIPES, + VCSCAnnData, + VCSCArray, + VCSCArrayNormalized, + VCSRArray, + VCSRArrayNormalized, +) + + +@pytest.fixture(params=[VCSCArray, VCSRArray]) +def vcls(request): + return request.param + + +def _scipy_for(vcls, dense): + return sp.csc_array(dense) if vcls is VCSCArray else sp.csr_array(dense) + + +def _norm_cls(vcls): + return VCSCArrayNormalized if vcls is VCSCArray else VCSRArrayNormalized + + +def _reference(dense: np.ndarray, recipe: str) -> np.ndarray: + """A plain-numpy version of ``y = (g(x * a * b) - c) * s`` for each recipe.""" + depth = dense.sum(axis=1) + with np.errstate(divide="ignore", invalid="ignore"): + if recipe == "raw": + a = np.ones_like(depth) + elif recipe in ("cp10k_log1p", "scanpy"): + a = np.where(depth > 0, 1e4 / depth, 1.0) + elif recipe == "parafac2": + median = np.median(depth) + a = np.where(depth > 0, median / depth, 1.0) if median > 0 else np.ones_like(depth) + elif recipe == "pearson": + a = np.where(depth > 0, 1.0 / depth, 1.0) + else: + raise ValueError(recipe) + + scaled = dense * a[:, None] + if recipe in ("parafac2", "pearson"): + gsum = scaled.sum(axis=0) + with np.errstate(divide="ignore", invalid="ignore"): + b = np.where(gsum > 0, 1.0 / gsum, 0.0) + else: + b = np.ones(dense.shape[1]) + + x = scaled * b[None, :] + if recipe in ("cp10k_log1p", "scanpy"): + g = np.log1p(x) + elif recipe == "parafac2": + g = np.log10(1.0 + 1000.0 * x) + elif recipe == "pearson": + g = np.sqrt(np.clip(x, 0.0, None)) + else: + g = x + + if recipe in ("parafac2", "scanpy", "pearson"): + c = g.mean(axis=0) + else: + c = np.zeros(dense.shape[1]) + + if recipe in ("scanpy", "pearson"): + std = g.std(axis=0) + with np.errstate(divide="ignore", invalid="ignore"): + s = np.where(std > 0, 1.0 / std, 1.0) + else: + s = np.ones(dense.shape[1]) + + return (g - c) * s + + +# -- per-recipe numerical correctness ----------------------------------------- + + +@pytest.mark.parametrize("recipe", sorted(RECIPES)) +def test_toarray_matches_reference_for_every_recipe(dense, vcls, recipe): + if dense.sum() == 0: + pytest.skip("all-zero matrix: median row total is 0") + v = vcls.from_scipy(_scipy_for(vcls, dense)) + nv = v.normalized(recipe) + assert isinstance(nv, _norm_cls(vcls)) + assert nv.recipe.name == recipe + np.testing.assert_allclose(nv.toarray(), _reference(dense, recipe), atol=1e-6) + + +@pytest.mark.parametrize("recipe", sorted(RECIPES)) +def test_matmul_matches_reference_for_every_recipe(dense, vcls, recipe): + if dense.sum() == 0: + pytest.skip("all-zero matrix: median row total is 0") + v = vcls.from_scipy(_scipy_for(vcls, dense)) + nv = v.normalized(recipe) + ref = _reference(dense, recipe) + + rng = np.random.default_rng(11) + B = rng.normal(size=(dense.shape[1], 3)) + np.testing.assert_allclose(nv @ B, ref @ B, atol=1e-5) + + Bl = rng.normal(size=(3, dense.shape[0])) + np.testing.assert_allclose(Bl @ nv, Bl @ ref, atol=1e-5) + + +def test_unknown_view_raises(vcls, dense): + v = vcls.from_scipy(_scipy_for(vcls, dense)) + with pytest.raises(ValueError, match="unknown normalization view"): + v.normalized("not-a-recipe") + + +# -- a/b/c/s exposed as named in the issue ------------------------------------ + + +def test_abcs_properties_are_named_per_issue_40(vcls, dense): + if dense.sum() == 0: + pytest.skip("all-zero matrix: median row total is 0") + v = vcls.from_scipy(_scipy_for(vcls, dense)) + nv = v.normalized("scanpy") + assert nv.a.shape == (dense.shape[0],) + assert nv.b.shape == (dense.shape[1],) + assert nv.c.shape == (dense.shape[1],) + assert nv.s.shape == (dense.shape[1],) + # scanpy: b == 1 everywhere (no per-gene scale) + np.testing.assert_allclose(nv.b, 1.0) + + +# -- .normalized(view, recalculate=...) caching ------------------------------- + + +def test_recalculate_false_reuses_cached_view(vcls, dense): + if dense.sum() == 0: + pytest.skip("all-zero matrix: median row total is 0") + v = vcls.from_scipy(_scipy_for(vcls, dense)) + nv1 = v.normalized("parafac2") + nv2 = v.normalized("parafac2", recalculate=False) + assert nv1 is nv2 + + +def test_recalculate_true_always_recomputes(vcls, dense): + if dense.sum() == 0: + pytest.skip("all-zero matrix: median row total is 0") + v = vcls.from_scipy(_scipy_for(vcls, dense)) + nv1 = v.normalized("parafac2") + nv2 = v.normalized("parafac2", recalculate=True) + assert nv1 is not nv2 + np.testing.assert_allclose(nv1.toarray(), nv2.toarray()) + + +def test_switching_views_without_recalculate_uses_independent_caches(vcls, dense): + """Each recipe gets its own cache slot; switching doesn't disturb the others.""" + if dense.sum() == 0: + pytest.skip("all-zero matrix: median row total is 0") + v = vcls.from_scipy(_scipy_for(vcls, dense)) + parafac2 = v.normalized("parafac2") + raw = v.normalized("raw", recalculate=False) # not cached yet -> computed once + assert not np.allclose(parafac2.toarray(), raw.toarray()) + + # Switching back doesn't recompute -- same object, same values as before. + back = v.normalized("parafac2", recalculate=False) + assert back is parafac2 + back_raw = v.normalized("raw", recalculate=False) + assert back_raw is raw + + +def test_indexing_the_raw_array_starts_with_an_empty_cache(vcls, dense): + """A freshly indexed array is a new object with nothing to reuse.""" + if dense.sum() == 0 or dense.shape[0] < 2: + pytest.skip("shape too small or all-zero") + v = vcls.from_scipy(_scipy_for(vcls, dense)) + v.normalized("parafac2") + sub = v[0:1, :] + assert sub._norm_cache == {} + + +# -- select() carries the recipe forward -------------------------------------- + + +def test_select_keeps_the_same_recipe(vcls, dense): + if dense.sum() == 0 or dense.shape[0] < 2: + pytest.skip("shape too small or all-zero") + v = vcls.from_scipy(_scipy_for(vcls, dense)) + nv = v.normalized("scanpy") + sub = nv.select(np.arange(min(2, dense.shape[0]))) + assert sub.recipe.name == "scanpy" + np.testing.assert_allclose( + sub.toarray(), _reference(dense[: sub.shape[0]], "scanpy"), atol=1e-6 + ) + + +# -- VCSCAnnData: obs/varm/uns storage and staleness -------------------------- + + +def _small_adata(rng: np.random.Generator, n_obs=24, n_vars=10) -> VCSCAnnData: + dense = rng.poisson(1.5, size=(n_obs, n_vars)).astype(np.float64) + obs = pd.DataFrame(index=[f"c{i}" for i in range(n_obs)]) + var = pd.DataFrame(index=[f"g{i}" for i in range(n_vars)]) + return VCSCAnnData(X=VCSCArray.from_scipy(sp.csc_array(dense)), obs=obs, var=var) + + +def test_anndata_normalized_records_obs_varm_uns(): + rng = np.random.default_rng(0) + adata = _small_adata(rng) + nv = adata.normalized("parafac2") + + np.testing.assert_allclose(adata.obs["vsparse_a"].to_numpy(), nv.a) + np.testing.assert_allclose(np.asarray(adata.varm["vsparse_b"]).reshape(-1), nv.b) + np.testing.assert_allclose(np.asarray(adata.varm["vsparse_c"]).reshape(-1), nv.c) + np.testing.assert_allclose(np.asarray(adata.varm["vsparse_s"]).reshape(-1), nv.s) + assert adata.uns["vsparse"] == {"recipe": "parafac2", "stale": False} + + +def test_anndata_normalized_recalculate_false_hits_in_memory_cache(): + rng = np.random.default_rng(1) + adata = _small_adata(rng) + nv1 = adata.normalized("parafac2") + nv2 = adata.normalized("parafac2", recalculate=False) + assert nv1 is nv2 + + +def test_anndata_indexing_marks_normalization_stale(): + rng = np.random.default_rng(2) + adata = _small_adata(rng) + nv = adata.normalized("parafac2") + + sub = adata[0:8, :] + assert sub.uns["vsparse"] == {"recipe": "parafac2", "stale": True} + # Parent is untouched by the child's bookkeeping. + assert adata.uns["vsparse"]["stale"] is False + + # recalculate=False: reuse the carried-over (stale) per-cell/per-gene + # values without recomputing -- exactly the parent's own values, windowed. + reused = sub.normalized("parafac2", recalculate=False) + assert reused.stale is True + np.testing.assert_allclose(reused.a, nv.a[0:8]) + np.testing.assert_allclose(reused.b, nv.b) + + # recalculate=True: fresh statistics for just this subset, no longer stale. + fresh = sub.normalized("parafac2", recalculate=True) + assert fresh.stale is False + assert sub.uns["vsparse"]["stale"] is False + assert not np.allclose(fresh.b, nv.b) # different population -> different gene scale + + +def test_anndata_normalized_requires_x(): + adata = VCSCAnnData(obs=pd.DataFrame(index=["a"]), var=pd.DataFrame(index=["g"])) + with pytest.raises(ValueError, match="requires X"): + adata.normalized() From e7a53e5c2cb47fe075105f8a005dee5d22e61577 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Wed, 9 Sep 2026 10:52:05 -0700 Subject: [PATCH 2/7] Set a uniform 0.3 ceiling for time_ratio_view_over_materialize Replaces the per-case recorded-with-margin ceilings with a flat cutoff across every normalized-view benchmark, rather than one keyed to each case's specific measured ratio. Co-Authored-By: Claude Sonnet 5 --- benchmarks/baselines.json | 40 +++++++++++++++++++-------------------- 1 file changed, 20 insertions(+), 20 deletions(-) diff --git a/benchmarks/baselines.json b/benchmarks/baselines.json index aed662d..c94072b 100644 --- a/benchmarks/baselines.json +++ b/benchmarks/baselines.json @@ -37,83 +37,83 @@ "peak_alloc_mb": 124.5457 }, "normalized_cp10k_log1p_matmat_vs_materialize": { - "time_ratio_view_over_materialize": 0.1444, + "time_ratio_view_over_materialize": 0.3, "peak_alloc_mb_view": 2.7262 }, "normalized_cp10k_log1p_matvec_vs_materialize": { - "time_ratio_view_over_materialize": 0.0864, + "time_ratio_view_over_materialize": 0.3, "peak_alloc_mb_view": 0.3858 }, "normalized_cp10k_log1p_rmatmat_vs_materialize": { - "time_ratio_view_over_materialize": 0.1639, + "time_ratio_view_over_materialize": 0.3, "peak_alloc_mb_view": 3.6521 }, "normalized_cp10k_log1p_rmatvec_vs_materialize": { - "time_ratio_view_over_materialize": 0.094, + "time_ratio_view_over_materialize": 0.3, "peak_alloc_mb_view": 0.1321 }, "normalized_parafac2_matmat_vs_materialize": { - "time_ratio_view_over_materialize": 0.1286, + "time_ratio_view_over_materialize": 0.3, "peak_alloc_mb_view": 2.7262 }, "normalized_parafac2_matvec_vs_materialize": { - "time_ratio_view_over_materialize": 0.1264, + "time_ratio_view_over_materialize": 0.3, "peak_alloc_mb_view": 0.3858 }, "normalized_parafac2_rmatmat_vs_materialize": { - "time_ratio_view_over_materialize": 0.1686, + "time_ratio_view_over_materialize": 0.3, "peak_alloc_mb_view": 3.6521 }, "normalized_parafac2_rmatvec_vs_materialize": { - "time_ratio_view_over_materialize": 0.0781, + "time_ratio_view_over_materialize": 0.3, "peak_alloc_mb_view": 0.1321 }, "normalized_pearson_matmat_vs_materialize": { - "time_ratio_view_over_materialize": 0.052, + "time_ratio_view_over_materialize": 0.3, "peak_alloc_mb_view": 2.7262 }, "normalized_pearson_matvec_vs_materialize": { - "time_ratio_view_over_materialize": 0.0364, + "time_ratio_view_over_materialize": 0.3, "peak_alloc_mb_view": 0.3858 }, "normalized_pearson_rmatmat_vs_materialize": { - "time_ratio_view_over_materialize": 0.1295, + "time_ratio_view_over_materialize": 0.3, "peak_alloc_mb_view": 3.6521 }, "normalized_pearson_rmatvec_vs_materialize": { - "time_ratio_view_over_materialize": 0.0397, + "time_ratio_view_over_materialize": 0.3, "peak_alloc_mb_view": 0.1321 }, "normalized_raw_matmat_vs_materialize": { - "time_ratio_view_over_materialize": 0.0518, + "time_ratio_view_over_materialize": 0.3, "peak_alloc_mb_view": 2.7262 }, "normalized_raw_matvec_vs_materialize": { - "time_ratio_view_over_materialize": 0.0427, + "time_ratio_view_over_materialize": 0.3, "peak_alloc_mb_view": 0.3858 }, "normalized_raw_rmatmat_vs_materialize": { - "time_ratio_view_over_materialize": 0.1011, + "time_ratio_view_over_materialize": 0.3, "peak_alloc_mb_view": 3.6521 }, "normalized_raw_rmatvec_vs_materialize": { - "time_ratio_view_over_materialize": 0.0257, + "time_ratio_view_over_materialize": 0.3, "peak_alloc_mb_view": 0.1321 }, "normalized_scanpy_matmat_vs_materialize": { - "time_ratio_view_over_materialize": 0.1513, + "time_ratio_view_over_materialize": 0.3, "peak_alloc_mb_view": 2.7262 }, "normalized_scanpy_matvec_vs_materialize": { - "time_ratio_view_over_materialize": 0.1004, + "time_ratio_view_over_materialize": 0.3, "peak_alloc_mb_view": 0.3858 }, "normalized_scanpy_rmatmat_vs_materialize": { - "time_ratio_view_over_materialize": 0.1887, + "time_ratio_view_over_materialize": 0.3, "peak_alloc_mb_view": 3.6521 }, "normalized_scanpy_rmatvec_vs_materialize": { - "time_ratio_view_over_materialize": 0.1171, + "time_ratio_view_over_materialize": 0.3, "peak_alloc_mb_view": 0.1321 } } From 52209ddbc98fd11abcfd84a76ee33e8be15032fe Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Wed, 9 Sep 2026 10:58:46 -0700 Subject: [PATCH 3/7] Re-fuse VCSC gene_scale/center-scale into one pass per column For VCSC, a column's gsum (b) depends only on that column's own nonzeros, so it's already final by the time the g-transform pass over the same nonzeros needs it -- no need to wait for other columns. The refactor to per-recipe statistics had split this into two separate numba dispatches (_column_stats_major_is_col + _column_gstats_major_is_col), doubling parallel-region launch overhead for parafac2/pearson without doing any extra elementwise work. Fuses them back into one kernel, parameterized by which parts a given recipe actually needs (need_b/need_gstats), matching the pre-refactor VCSC fusion. VCSR keeps two separate passes, since gene_scale there isn't final until every row has scattered into it. Co-Authored-By: Claude Sonnet 5 --- src/vsparse/_norm_common.py | 118 +++++++++++++++++++----------------- 1 file changed, 63 insertions(+), 55 deletions(-) diff --git a/src/vsparse/_norm_common.py b/src/vsparse/_norm_common.py index 625ee6a..43a3255 100644 --- a/src/vsparse/_norm_common.py +++ b/src/vsparse/_norm_common.py @@ -34,10 +34,10 @@ those need ``b`` first), done with parallel numba kernels below, specialized per storage format: -- major=columns (VCSC): both passes collapse into one, fully parallel - over columns with no cross-thread writes -- each column's own elements - carry everything needed to compute both its ``b`` and its ``c``/``s`` - (:func:`_column_stats_major_is_col`, :func:`_column_gstats_major_is_col`). +- major=columns (VCSC): both passes collapse into one fused kernel, fully + parallel over columns with no cross-thread writes -- each column's own + elements carry everything needed to compute both its ``b`` and its + ``c``/``s`` (:func:`_column_stats_major_is_col`). - major=rows (VCSR): each pass is a scatter-add across columns from many rows, so it's parallelized row-chunked with thread-local partial column arrays, reduced by summing across threads @@ -135,44 +135,43 @@ def _g_np(x: np.ndarray, g_code: int) -> np.ndarray: @numba.njit(cache=True, parallel=True) -def _column_stats_major_is_col(major_ptr, values, value_ptr, indices, row_scale): - """``gsum[j] = sum_i values[i, j] / row_scale[i]`` -- the raw material for ``b``.""" - n_major = major_ptr.shape[0] - 1 - gsum = np.zeros(n_major, dtype=np.float64) - for j in numba.prange(n_major): # ty: ignore[not-iterable] - gs = 0.0 - for u in range(major_ptr[j], major_ptr[j + 1]): - v = values[u] - for k in range(value_ptr[u], value_ptr[u + 1]): - gs += v / row_scale[indices[k]] - gsum[j] = gs - return gsum - - -@numba.njit(cache=True, parallel=True) -def _column_gstats_major_is_col( - major_ptr, values, value_ptr, indices, row_scale, gene_scale, g_code +def _column_stats_major_is_col( + major_ptr, values, value_ptr, indices, row_scale, need_b, need_gstats, g_code ): - """Per-column sum/sum-of-squares of ``g(x / row_scale / gene_scale)`` over stored entries.""" + """Per-column ``gsum`` (raw material for ``b``) and sum/sum-of-squares of ``g(scaled)``. + + Fused into one pass per column (rather than two separate dispatches): + unlike the VCSR scatter passes below, a VCSC column's ``gsum`` depends + only on that column's own nonzeros, so it's already final by the time + the second (``g``-transform) loop over the same nonzeros needs it -- + no need to wait for every other column to finish first. + """ n_major = major_ptr.shape[0] - 1 + gsum = np.ones(n_major, dtype=np.float64) col_sum = np.zeros(n_major, dtype=np.float64) col_sumsq = np.zeros(n_major, dtype=np.float64) for j in numba.prange(n_major): # ty: ignore[not-iterable] - gs = gene_scale[j] - if gs <= 0.0: - continue - s0 = 0.0 - s1 = 0.0 - for u in range(major_ptr[j], major_ptr[j + 1]): - v = values[u] - for k in range(value_ptr[u], value_ptr[u + 1]): - scaled = v / row_scale[indices[k]] / gs - gy = _g(scaled, g_code) - s0 += gy - s1 += gy * gy - col_sum[j] = s0 - col_sumsq[j] = s1 - return col_sum, col_sumsq + gs = 1.0 + if need_b: + gs = 0.0 + for u in range(major_ptr[j], major_ptr[j + 1]): + v = values[u] + for k in range(value_ptr[u], value_ptr[u + 1]): + gs += v / row_scale[indices[k]] + gsum[j] = gs + if need_gstats and gs > 0.0: + s0 = 0.0 + s1 = 0.0 + for u in range(major_ptr[j], major_ptr[j + 1]): + v = values[u] + for k in range(value_ptr[u], value_ptr[u + 1]): + scaled = v / row_scale[indices[k]] / gs + gy = _g(scaled, g_code) + s0 += gy + s1 += gy * gy + col_sum[j] = s0 + col_sumsq[j] = s1 + return gsum, col_sum, col_sumsq # -- statistics: major=rows -- scatter-add passes ---------------------------- @@ -347,33 +346,37 @@ def __init__( self.row_scale = row_scale indices = arr.indices # decode once; shared by both statistics passes below - if self.recipe.gene_scale: - if self._format == "csc": - gsum = _column_stats_major_is_col( - arr.major_ptr, arr.values, arr.value_ptr, indices, row_scale - ) - else: - nthreads = numba.get_num_threads() - gsum = _scaled_col_sums_vcs( - arr.major_ptr, arr.values, arr.value_ptr, indices, row_scale, n_cols, nthreads - ) - gene_scale = gsum - else: - gene_scale = np.ones(n_cols, dtype=np.float64) - self.gene_scale = gene_scale + need_b = self.recipe.gene_scale + need_gstats = self.recipe.center or self.recipe.post_scale - if self.recipe.center or self.recipe.post_scale: - if self._format == "csc": - col_sum, col_sumsq = _column_gstats_major_is_col( + if self._format == "csc": + # One fused pass per column for both -- see _column_stats_major_is_col. + if need_b or need_gstats: + gene_scale, col_sum, col_sumsq = _column_stats_major_is_col( arr.major_ptr, arr.values, arr.value_ptr, indices, row_scale, - gene_scale, + need_b, + need_gstats, self.recipe.g_code, ) else: + gene_scale = np.ones(n_cols, dtype=np.float64) + col_sum = col_sumsq = np.zeros(n_cols, dtype=np.float64) + else: + # VCSR can't fuse these: gene_scale[c] isn't final until every row + # has been scattered into it, so the g-transform pass has to wait + # for the whole first pass to finish -- two genuinely separate passes. + if need_b: + nthreads = numba.get_num_threads() + gene_scale = _scaled_col_sums_vcs( + arr.major_ptr, arr.values, arr.value_ptr, indices, row_scale, n_cols, nthreads + ) + else: + gene_scale = np.ones(n_cols, dtype=np.float64) + if need_gstats: nthreads = numba.get_num_threads() col_sum, col_sumsq = _gstats_col_sums_vcs( arr.major_ptr, @@ -386,6 +389,11 @@ def __init__( n_cols, nthreads, ) + else: + col_sum = col_sumsq = np.zeros(n_cols, dtype=np.float64) + self.gene_scale = gene_scale + + if need_gstats: mean = col_sum / n_rows if n_rows > 0 else np.zeros(n_cols, dtype=np.float64) variance = np.clip( col_sumsq / n_rows - mean**2 if n_rows > 0 else np.zeros(n_cols), 0.0, None From a4e99f8bbe6de65e369634cd866f3bb708b9bf2c Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Wed, 9 Sep 2026 11:02:07 -0700 Subject: [PATCH 4/7] Fix ty type-check failure: _dual_arr unresolved on NormalizedViewBase from_stats() set self._dual_arr directly, but that attribute is only declared on the _VCSNormalizedBase subclass (in _vcs_norm.py), not on the NormalizedViewBase base class where from_stats() lives -- ty can't verify it exists on Self there. Replaces the direct attribute set with an _init_extra() hook (no-op on the base class, overridden by _VCSNormalizedBase to initialize _dual_arr), called from both __init__ and from_stats(). Co-Authored-By: Claude Sonnet 5 --- src/vsparse/_norm_common.py | 12 +++++++++--- src/vsparse/_vcs_norm.py | 7 ++++++- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/src/vsparse/_norm_common.py b/src/vsparse/_norm_common.py index 43a3255..4385d69 100644 --- a/src/vsparse/_norm_common.py +++ b/src/vsparse/_norm_common.py @@ -446,11 +446,17 @@ def from_stats( self.gene_scale = np.where(b > 0.0, 1.0 / b, 0.0) self.col_mean = np.asarray(c, dtype=np.float64) self.col_post_scale = np.asarray(s, dtype=np.float64) - # Subclasses with the ``_dual_arr`` slot (VCS views) need it initialized - # too, since ``__init__`` (which normally does) is bypassed here. - self._dual_arr = None + self._init_extra() return self + def _init_extra(self) -> None: + """Hook for subclasses with extra per-instance state (e.g. ``_dual_arr``). + + ``__init__`` normally initializes that state itself; :meth:`from_stats` + builds an instance via ``object.__new__`` instead, bypassing it, so it + calls this explicitly. A no-op here; overridden where needed. + """ + # -- recipe-facing statistics (a/b/c/s, as named in the issue) ------------- @property diff --git a/src/vsparse/_vcs_norm.py b/src/vsparse/_vcs_norm.py index 183b831..bd9fdbf 100644 --- a/src/vsparse/_vcs_norm.py +++ b/src/vsparse/_vcs_norm.py @@ -28,13 +28,18 @@ class _VCSNormalizedBase(NormalizedViewBase): __slots__ = ("_dual_arr",) + _dual_arr: _VCSBase | None + def __init__( self, arr: _VCSBase, recipe: str | Recipe = DEFAULT_RECIPE, *, stale: bool = False ) -> None: super().__init__(arr, recipe, stale=stale) + self._init_extra() + + def _init_extra(self) -> None: # Opposite-format copy of `arr`, cached by vsparse._vcs_matmul when # regrouping the whole array fits one chunk's budget. - self._dual_arr: _VCSBase | None = None + self._dual_arr = None def __matmul__(self, other: Any) -> Any: """``self @ other`` for a dense ``other`` -- see :mod:`vsparse._vcs_matmul`.""" From 61db4359ee5c177e59592ec3a48d0933ffefc6aa Mon Sep 17 00:00:00 2001 From: fishidaho Date: Thu, 10 Sep 2026 14:30:19 -0700 Subject: [PATCH 5/7] Fix VCSCAnnData.normalized() rejecting caller-built Recipe objects `VCSCAnnData.normalized()` resolved its `view` argument to a `Recipe`, then handed `recipe.name` back down to `_VCSBase.normalized()`, which re-resolved that name through `RECIPES`. Any `Recipe` the caller built themselves is not in `RECIPES`, so the round trip raised `ValueError: unknown normalization view`, even though the same object works at the array level. Pass the resolved `Recipe` straight down instead, and widen both signatures to `str | Recipe` to say so. The `_norm_cache`/`_vcs_norm_cache` keys move from `recipe.name` to the resolved `Recipe` (a frozen, slotted dataclass, so hashable). Two distinct caller-built recipes are now free to share a `name` without silently handing back each other's cached view. Also validate `g_code` in `resolve_recipe()`: `_g()` falls through to the identity branch for anything it doesn't recognize, so a typo'd code would have silently produced an untransformed matrix rather than an error. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_019ZwuJMcTRUacxqARwjMngE --- src/vsparse/_anndata_class.py | 14 +++---- src/vsparse/_base.py | 14 ++++--- src/vsparse/_norm_common.py | 10 +++++ tests/test_vcs_norm_recipes.py | 67 ++++++++++++++++++++++++++++++++++ 4 files changed, 92 insertions(+), 13 deletions(-) diff --git a/src/vsparse/_anndata_class.py b/src/vsparse/_anndata_class.py index 1ff2be0..ee8ab2b 100644 --- a/src/vsparse/_anndata_class.py +++ b/src/vsparse/_anndata_class.py @@ -12,7 +12,7 @@ from vsparse import _compression, _io from vsparse._base import VCSCArray, VCSRArray, _VCSBase -from vsparse._norm_common import DEFAULT_RECIPE, resolve_recipe +from vsparse._norm_common import DEFAULT_RECIPE, Recipe, resolve_recipe from vsparse._vcs_norm import VCSCArrayNormalized, VCSRArrayNormalized if TYPE_CHECKING: @@ -117,7 +117,7 @@ def __init__( ) self._vcs_X: _AnyVCS | None = None self._vcs_raw_X: _AnyVCS | None = None - self._vcs_norm_cache: dict[str, Any] = {} + self._vcs_norm_cache: dict[Any, Any] = {} shape = kwargs.pop("shape", None) if shape is None and X is None and "obs" not in kwargs: shape = (0, 0) @@ -212,7 +212,7 @@ def __getitem__(self, index: Any) -> VCSCAnnData: # ty: ignore[invalid-method-o # -- normalization ---------------------------------------------------------- - def normalized(self, view: str = DEFAULT_RECIPE, *, recalculate: bool = True) -> Any: + def normalized(self, view: str | Recipe = DEFAULT_RECIPE, *, recalculate: bool = True) -> Any: """A normalized view of ``X`` -- see :meth:`vsparse._base._VCSBase.normalized`. Also records the recipe's statistics -- per-cell ``a`` in @@ -237,7 +237,7 @@ def normalized(self, view: str = DEFAULT_RECIPE, *, recalculate: bool = True) -> recipe = resolve_recipe(view) cache = self._vcs_norm_cache if not recalculate: - cached = cache.get(recipe.name) + cached = cache.get(recipe) if cached is not None: return cached stored = self.uns.get(_VSPARSE_UNS_KEY) @@ -265,11 +265,11 @@ def normalized(self, view: str = DEFAULT_RECIPE, *, recalculate: bool = True) -> s=np.asarray(self.varm[_VSPARSE_VARM_S], dtype=np.float64).reshape(-1), stale=bool(stored.get("stale", False)), ) - cache[recipe.name] = nview + cache[recipe] = nview return nview - nview = self._vcs_X.normalized(recipe.name, recalculate=True) - cache[recipe.name] = nview + nview = self._vcs_X.normalized(recipe, recalculate=True) + cache[recipe] = nview self.obs[_VSPARSE_OBS_A] = nview.a self.varm[_VSPARSE_VARM_B] = nview.b self.varm[_VSPARSE_VARM_C] = nview.c diff --git a/src/vsparse/_base.py b/src/vsparse/_base.py index 64c60ed..0f810f0 100644 --- a/src/vsparse/_base.py +++ b/src/vsparse/_base.py @@ -22,6 +22,7 @@ from vsparse._indexutils import normalize_major_idx as _normalize_major_idx from vsparse._indexutils import smallest_index_dtype as _smallest_index_dtype from vsparse._norm_common import DEFAULT_RECIPE as _DEFAULT_RECIPE +from vsparse._norm_common import Recipe as _Recipe __all__ = ["VCSCArray", "VCSRArray"] @@ -78,7 +79,7 @@ def __init__( self.values = values self.value_ptr = value_ptr self.indices = indices - self._norm_cache: dict[str, Any] = {} + self._norm_cache: dict[Any, Any] = {} # -- axis bookkeeping ------------------------------------------------ @@ -199,15 +200,16 @@ def _transpose_major(self) -> _VCSBase: ) return other_cls(self.shape, major_ptr, values, value_ptr, indices) - def normalized(self, view: str = _DEFAULT_RECIPE, *, recalculate: bool = True) -> Any: + def normalized(self, view: str | _Recipe = _DEFAULT_RECIPE, *, recalculate: bool = True) -> Any: """A normalized *view* of this array -- see :mod:`vsparse._vcs_norm`/:mod:`vsparse._norm_common`. Parameters ---------- view - Which normalization recipe to apply -- one of + Which normalization recipe to apply: the name of one of :data:`vsparse._norm_common.RECIPES` (``"raw"``, ``"cp10k_log1p"``, - ``"parafac2"`` (the default), ``"scanpy"``, ``"pearson"``). + ``"parafac2"`` (the default), ``"scanpy"``, ``"pearson"``), or a + :class:`~vsparse._norm_common.Recipe` built by the caller. recalculate If ``True`` (the default), (re)compute the recipe's statistics fresh from this array. If ``False``, reuse a previously computed @@ -223,13 +225,13 @@ def normalized(self, view: str = _DEFAULT_RECIPE, *, recalculate: bool = True) - recipe = resolve_recipe(view) cache = self._norm_cache if not recalculate: - cached = cache.get(recipe.name) + cached = cache.get(recipe) if cached is not None: return cached cls = VCSCArrayNormalized if self._format == "csc" else VCSRArrayNormalized result = cls(self, recipe) - cache[recipe.name] = result + cache[recipe] = result return result def log1p(self) -> _VCSBase: diff --git a/src/vsparse/_norm_common.py b/src/vsparse/_norm_common.py index 4385d69..57e3503 100644 --- a/src/vsparse/_norm_common.py +++ b/src/vsparse/_norm_common.py @@ -95,8 +95,18 @@ class Recipe: DEFAULT_RECIPE = "parafac2" +#: Every ``g`` :func:`_g` knows how to apply. A ``Recipe`` carrying anything +#: else would silently fall through to the identity branch, so reject it here. +_G_CODES = frozenset({G_IDENTITY, G_LOG1P, G_LOG1P_1000X, G_SQRT}) + + def resolve_recipe(view: str | Recipe) -> Recipe: if isinstance(view, Recipe): + if view.g_code not in _G_CODES: + raise ValueError( + f"recipe {view.name!r} has unknown g_code {view.g_code!r}; " + f"choose from {sorted(_G_CODES)}" + ) return view try: return RECIPES[view] diff --git a/tests/test_vcs_norm_recipes.py b/tests/test_vcs_norm_recipes.py index d9ef897..c1a3c80 100644 --- a/tests/test_vcs_norm_recipes.py +++ b/tests/test_vcs_norm_recipes.py @@ -9,6 +9,7 @@ from vsparse import ( RECIPES, + Recipe, VCSCAnnData, VCSCArray, VCSCArrayNormalized, @@ -252,3 +253,69 @@ def test_anndata_normalized_requires_x(): adata = VCSCAnnData(obs=pd.DataFrame(index=["a"]), var=pd.DataFrame(index=["g"])) with pytest.raises(ValueError, match="requires X"): adata.normalized() + + +# -- caller-built Recipe objects ---------------------------------------------- + + +def _custom_recipe() -> Recipe: + """A recipe that is *not* in RECIPES: cp10k + log1p, centered and variance-scaled.""" + return Recipe("custom_cp10k_scaled", 1e4, False, RECIPES["scanpy"].g_code, True, True) + + +def test_custom_recipe_works_on_the_array(vcls, dense): + if dense.sum() == 0: + pytest.skip("all-zero matrix: median row total is 0") + v = vcls.from_scipy(_scipy_for(vcls, dense)) + nv = v.normalized(_custom_recipe()) + assert nv.recipe.name == "custom_cp10k_scaled" + # Same (a, b, g, c, s) as "scanpy", so it must agree with that reference. + np.testing.assert_allclose(nv.toarray(), _reference(dense, "scanpy"), atol=1e-6) + + +def test_custom_recipe_works_on_the_anndata(): + """Regression: VCSCAnnData.normalized() used to hand ``recipe.name`` back to + the array, which re-resolved it through RECIPES and raised for anything the + caller built themselves.""" + rng = np.random.default_rng(0) + adata = _small_adata(rng) + recipe = _custom_recipe() + nv = adata.normalized(recipe) + + assert nv.recipe is recipe + assert adata.uns["vsparse"]["recipe"] == "custom_cp10k_scaled" + np.testing.assert_allclose(adata.obs["vsparse_a"].to_numpy(), nv.a) + np.testing.assert_allclose( + nv.toarray(), _reference(np.asarray(adata.X.toarray()), "scanpy"), atol=1e-6 + ) + + +def test_custom_recipe_round_trips_through_recalculate_false(): + rng = np.random.default_rng(0) + adata = _small_adata(rng) + recipe = _custom_recipe() + ref = adata.normalized(recipe).toarray() + adata._vcs_norm_cache.clear() # force the obs/varm/uns path, not the in-memory one + again = adata.normalized(recipe, recalculate=False) + np.testing.assert_allclose(again.toarray(), ref) + + +def test_two_custom_recipes_sharing_a_name_do_not_collide(vcls, dense): + """The cache keys on the Recipe itself, so a shared ``name`` is not a shared slot.""" + if dense.sum() == 0: + pytest.skip("all-zero matrix: median row total is 0") + v = vcls.from_scipy(_scipy_for(vcls, dense)) + centered = Recipe("dup", 1e4, False, RECIPES["scanpy"].g_code, True, False) + plain = Recipe("dup", 1e4, False, RECIPES["scanpy"].g_code, False, False) + nv_centered = v.normalized(centered) + nv_plain = v.normalized(plain, recalculate=False) + assert nv_plain is not nv_centered + np.testing.assert_allclose(nv_plain.c, 0.0) + assert v.normalized(centered, recalculate=False) is nv_centered + + +def test_recipe_with_an_unknown_g_code_is_rejected(vcls, dense): + """``_g`` falls through to the identity for an unrecognized code -- fail loudly instead.""" + v = vcls.from_scipy(_scipy_for(vcls, dense)) + with pytest.raises(ValueError, match="unknown g_code"): + v.normalized(Recipe("bogus", None, False, 99, False, False)) From f1b1fa35dc4bbe13ab1364cf2dd60adf00c8ff2d Mon Sep 17 00:00:00 2001 From: fishidaho Date: Thu, 10 Sep 2026 14:01:48 -0700 Subject: [PATCH 6/7] Bound the normalization cache and stop it pinning O(nnz) duals `_norm_cache`/`_vcs_norm_cache` were plain dicts holding every computed view strongly, with no eviction. Two problems followed. A view can carry a `_dual_arr` -- a whole opposite-format copy of the array, cached by the matmul kernels whenever regrouping fits one chunk's budget. Because the cache held views strongly, dropping every reference to a view did not release its dual: the array kept it alive for its own lifetime. Touching all five built-in recipes on an 8k x 1.5k array and discarding each view left 5 duals resident, 8.4 MB against 2.3 MB of index+value bytes for the array itself. And nothing capped the number of entries, so a caller cycling through recipes accumulated one entry per distinct recipe forever. Replace both with `_NormCache`: an LRU bounded at `NORM_CACHE_MAXSIZE` (4) that holds each view *weakly* and its statistics strongly. The statistics are `O(n_rows + n_cols)` -- 0.55 MB per entry on a 60k x 3k array, measured equal to the predicted `8 * (n_rows + 3 * n_cols)` -- so retaining those is cheap, while nothing `O(nnz)` survives the caller dropping a view. `recalculate=False` keeps its contract: a hit on a live view returns that exact object, so identity is stable while the caller holds it; a hit on a collected one rebuilds from the retained statistics, still skipping the `O(nnz)` passes (1334x faster than recomputing, measured). That rebuild goes through a new `_from_internal_stats()` rather than `from_stats()`, because `a`/`b` are exposed as reciprocals of the internal `row_scale`/`gene_scale` and `1 / (1 / x)` is not exactly involutive in float64. Restored views are now bit-identical to a fresh recompute, which is asserted rather than approximated in the tests. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_019ZwuJMcTRUacxqARwjMngE --- src/vsparse/_anndata_class.py | 10 +-- src/vsparse/_base.py | 10 ++- src/vsparse/_norm_common.py | 145 +++++++++++++++++++++++++++++++-- tests/test_vcs_norm_recipes.py | 72 +++++++++++++++- 4 files changed, 219 insertions(+), 18 deletions(-) diff --git a/src/vsparse/_anndata_class.py b/src/vsparse/_anndata_class.py index ee8ab2b..f514b83 100644 --- a/src/vsparse/_anndata_class.py +++ b/src/vsparse/_anndata_class.py @@ -12,7 +12,7 @@ from vsparse import _compression, _io from vsparse._base import VCSCArray, VCSRArray, _VCSBase -from vsparse._norm_common import DEFAULT_RECIPE, Recipe, resolve_recipe +from vsparse._norm_common import DEFAULT_RECIPE, Recipe, _NormCache, resolve_recipe from vsparse._vcs_norm import VCSCArrayNormalized, VCSRArrayNormalized if TYPE_CHECKING: @@ -117,7 +117,7 @@ def __init__( ) self._vcs_X: _AnyVCS | None = None self._vcs_raw_X: _AnyVCS | None = None - self._vcs_norm_cache: dict[Any, Any] = {} + self._vcs_norm_cache = _NormCache() shape = kwargs.pop("shape", None) if shape is None and X is None and "obs" not in kwargs: shape = (0, 0) @@ -237,7 +237,7 @@ def normalized(self, view: str | Recipe = DEFAULT_RECIPE, *, recalculate: bool = recipe = resolve_recipe(view) cache = self._vcs_norm_cache if not recalculate: - cached = cache.get(recipe) + cached = cache.get(recipe, self._vcs_X) if cached is not None: return cached stored = self.uns.get(_VSPARSE_UNS_KEY) @@ -265,11 +265,11 @@ def normalized(self, view: str | Recipe = DEFAULT_RECIPE, *, recalculate: bool = s=np.asarray(self.varm[_VSPARSE_VARM_S], dtype=np.float64).reshape(-1), stale=bool(stored.get("stale", False)), ) - cache[recipe] = nview + cache.put(recipe, nview) return nview nview = self._vcs_X.normalized(recipe, recalculate=True) - cache[recipe] = nview + cache.put(recipe, nview) self.obs[_VSPARSE_OBS_A] = nview.a self.varm[_VSPARSE_VARM_B] = nview.b self.varm[_VSPARSE_VARM_C] = nview.c diff --git a/src/vsparse/_base.py b/src/vsparse/_base.py index 0f810f0..001266c 100644 --- a/src/vsparse/_base.py +++ b/src/vsparse/_base.py @@ -23,6 +23,7 @@ from vsparse._indexutils import smallest_index_dtype as _smallest_index_dtype from vsparse._norm_common import DEFAULT_RECIPE as _DEFAULT_RECIPE from vsparse._norm_common import Recipe as _Recipe +from vsparse._norm_common import _NormCache __all__ = ["VCSCArray", "VCSRArray"] @@ -79,7 +80,7 @@ def __init__( self.values = values self.value_ptr = value_ptr self.indices = indices - self._norm_cache: dict[Any, Any] = {} + self._norm_cache = _NormCache() # -- axis bookkeeping ------------------------------------------------ @@ -218,6 +219,9 @@ def normalized(self, view: str | _Recipe = _DEFAULT_RECIPE, *, recalculate: bool ``recalculate``) -- this is how switching between recipes avoids recomputing each one every time. If no such view has been computed yet, one is still computed (there is nothing to reuse). + The cache holds at most + :data:`~vsparse._norm_common.NORM_CACHE_MAXSIZE` recipes, and holds + each view weakly -- see :class:`~vsparse._norm_common._NormCache`. """ from vsparse._norm_common import resolve_recipe from vsparse._vcs_norm import VCSCArrayNormalized, VCSRArrayNormalized @@ -225,13 +229,13 @@ def normalized(self, view: str | _Recipe = _DEFAULT_RECIPE, *, recalculate: bool recipe = resolve_recipe(view) cache = self._norm_cache if not recalculate: - cached = cache.get(recipe) + cached = cache.get(recipe, self) if cached is not None: return cached cls = VCSCArrayNormalized if self._format == "csc" else VCSRArrayNormalized result = cls(self, recipe) - cache[recipe] = result + cache.put(recipe, result) return result def log1p(self) -> _VCSBase: diff --git a/src/vsparse/_norm_common.py b/src/vsparse/_norm_common.py index 57e3503..ae961fc 100644 --- a/src/vsparse/_norm_common.py +++ b/src/vsparse/_norm_common.py @@ -52,13 +52,22 @@ from __future__ import annotations +import weakref +from collections import OrderedDict from dataclasses import dataclass from typing import Any import numba import numpy as np -__all__ = ["DEFAULT_RECIPE", "RECIPES", "NormalizedViewBase", "Recipe", "resolve_recipe"] +__all__ = [ + "DEFAULT_RECIPE", + "NORM_CACHE_MAXSIZE", + "RECIPES", + "NormalizedViewBase", + "Recipe", + "resolve_recipe", +] # -- recipes ------------------------------------------------------------------ @@ -116,6 +125,86 @@ def resolve_recipe(view: str | Recipe) -> Recipe: ) from None +# -- cache -------------------------------------------------------------------- + +#: How many distinct recipes one array keeps statistics for. The cache exists so +#: that switching back and forth between a handful of recipes doesn't repeat the +#: ``O(nnz)`` passes; past that many, the least recently used entry is dropped. +#: Each entry costs ``8 * (n_rows + 3 * n_cols)`` bytes, dominated by +#: ``row_scale`` -- ~800 MB per entry at 100M cells, so this is deliberately small. +NORM_CACHE_MAXSIZE = 4 + + +class _NormCache: + """Bounded LRU of computed normalizations for one array, keyed by :class:`Recipe`. + + Holds each view's *statistics* strongly and the view itself only weakly. + That split matters: a view can carry an ``O(nnz)`` ``_dual_arr`` (a whole + opposite-format copy of the array, cached by the matmul kernels), and a + cache that kept views alive would pin one of those per recipe for as long + as the array lived, even after the caller had dropped every reference. + The statistics are ``O(n_rows + n_cols)``, so retaining those is cheap. + + A hit on a still-live view hands back that exact object, so + ``recalculate=False`` is identity-stable for as long as the caller holds + it. A hit on a view that has since been collected rebuilds one from the + retained statistics -- which is what the cache is actually for, since that + skips the ``O(nnz)`` passes. + """ + + __slots__ = ("_entries", "maxsize") + + def __init__(self, maxsize: int = NORM_CACHE_MAXSIZE) -> None: + # recipe -> (weakref to the view, view class, row_scale, gene_scale, + # col_mean, col_post_scale, stale). Insertion-ordered, so the + # first key is the least recently used. + self._entries: OrderedDict[Recipe, tuple[Any, ...]] = OrderedDict() + self.maxsize = maxsize + + def __len__(self) -> int: + return len(self._entries) + + def __contains__(self, recipe: Recipe) -> bool: + return recipe in self._entries + + def clear(self) -> None: + self._entries.clear() + + def get(self, recipe: Recipe, arr: Any) -> Any | None: + """The cached view for ``recipe`` over ``arr``, or ``None`` if there is none.""" + entry = self._entries.get(recipe) + if entry is None: + return None + self._entries.move_to_end(recipe) + ref, cls, row_scale, gene_scale, col_mean, col_post_scale, stale = entry + view = ref() + if view is not None and view._arr is arr: + return view + return cls._from_internal_stats( + arr, + recipe, + row_scale=row_scale, + gene_scale=gene_scale, + col_mean=col_mean, + col_post_scale=col_post_scale, + stale=stale, + ) + + def put(self, recipe: Recipe, view: Any) -> None: + self._entries[recipe] = ( + weakref.ref(view), + type(view), + view.row_scale, + view.gene_scale, + view.col_mean, + view.col_post_scale, + view.stale, + ) + self._entries.move_to_end(recipe) + while len(self._entries) > self.maxsize: + self._entries.popitem(last=False) + + # -- elementwise transform ---------------------------------------------------- @@ -337,7 +426,16 @@ class NormalizedViewBase: __array_ufunc__ = None - __slots__ = ("_arr", "col_mean", "col_post_scale", "gene_scale", "recipe", "row_scale", "stale") + __slots__ = ( + "__weakref__", # so _NormCache can hold a view without pinning it + "_arr", + "col_mean", + "col_post_scale", + "gene_scale", + "recipe", + "row_scale", + "stale", + ) def __init__( self, arr: Any, recipe: str | Recipe = DEFAULT_RECIPE, *, stale: bool = False @@ -446,16 +544,45 @@ def from_stats( raise ValueError( f"{cls.__name__} wraps a {cls._format!r}-format array, got {type(arr).__name__}" ) + a = np.asarray(a, dtype=np.float64) + b = np.asarray(b, dtype=np.float64) + return cls._from_internal_stats( + arr, + resolve_recipe(recipe), + row_scale=np.where(a > 0.0, 1.0 / a, 0.0), + gene_scale=np.where(b > 0.0, 1.0 / b, 0.0), + col_mean=np.asarray(c, dtype=np.float64), + col_post_scale=np.asarray(s, dtype=np.float64), + stale=stale, + ) + + @classmethod + def _from_internal_stats( + cls, + arr: Any, + recipe: Recipe, + *, + row_scale: np.ndarray, + gene_scale: np.ndarray, + col_mean: np.ndarray, + col_post_scale: np.ndarray, + stale: bool, + ) -> NormalizedViewBase: + """Attach already-computed *internal* statistics to a fresh view. + + The reciprocals :attr:`a`/:attr:`b` expose are not exactly involutive in + float64, so anything restoring a view it built earlier (see + :class:`_NormCache`) has to carry these arrays rather than round-trip + through ``a``/``b``. + """ self = object.__new__(cls) self._arr = arr - self.recipe = resolve_recipe(recipe) + self.recipe = recipe self.stale = stale - a = np.asarray(a, dtype=np.float64) - b = np.asarray(b, dtype=np.float64) - self.row_scale = np.where(a > 0.0, 1.0 / a, 0.0) - self.gene_scale = np.where(b > 0.0, 1.0 / b, 0.0) - self.col_mean = np.asarray(c, dtype=np.float64) - self.col_post_scale = np.asarray(s, dtype=np.float64) + self.row_scale = row_scale + self.gene_scale = gene_scale + self.col_mean = col_mean + self.col_post_scale = col_post_scale self._init_extra() return self diff --git a/tests/test_vcs_norm_recipes.py b/tests/test_vcs_norm_recipes.py index c1a3c80..7b9fc40 100644 --- a/tests/test_vcs_norm_recipes.py +++ b/tests/test_vcs_norm_recipes.py @@ -2,6 +2,9 @@ from __future__ import annotations +import gc +import weakref + import numpy as np import pandas as pd import pytest @@ -16,6 +19,7 @@ VCSRArray, VCSRArrayNormalized, ) +from vsparse._norm_common import NORM_CACHE_MAXSIZE @pytest.fixture(params=[VCSCArray, VCSRArray]) @@ -177,7 +181,7 @@ def test_indexing_the_raw_array_starts_with_an_empty_cache(vcls, dense): v = vcls.from_scipy(_scipy_for(vcls, dense)) v.normalized("parafac2") sub = v[0:1, :] - assert sub._norm_cache == {} + assert len(sub._norm_cache) == 0 # -- select() carries the recipe forward -------------------------------------- @@ -319,3 +323,69 @@ def test_recipe_with_an_unknown_g_code_is_rejected(vcls, dense): v = vcls.from_scipy(_scipy_for(vcls, dense)) with pytest.raises(ValueError, match="unknown g_code"): v.normalized(Recipe("bogus", None, False, 99, False, False)) + + +# -- cache eviction ----------------------------------------------------------- + + +def test_cache_does_not_pin_a_dropped_view(vcls, dense): + """The cache holds views weakly, so it can't keep an O(nnz) _dual_arr alive.""" + if dense.sum() == 0: + pytest.skip("all-zero matrix: median row total is 0") + v = vcls.from_scipy(_scipy_for(vcls, dense)) + ref = weakref.ref(v.normalized("parafac2")) + gc.collect() + assert ref() is None, "cached view outlived the caller's last reference" + assert len(v._norm_cache) == 1, "the statistics themselves should still be cached" + + +def test_cache_rebuilds_an_evicted_view_from_retained_statistics(vcls, dense): + """A collected view is rebuilt exactly, without redoing the O(nnz) passes.""" + if dense.sum() == 0: + pytest.skip("all-zero matrix: median row total is 0") + v = vcls.from_scipy(_scipy_for(vcls, dense)) + first = v.normalized("scanpy") + expected = first.toarray() + row_scale, gene_scale = first.row_scale, first.gene_scale + del first + gc.collect() + + again = v.normalized("scanpy", recalculate=False) + # Bit-identical, not merely close: the internal arrays are carried over + # rather than round-tripped through the a/b reciprocals. + np.testing.assert_array_equal(again.row_scale, row_scale) + np.testing.assert_array_equal(again.gene_scale, gene_scale) + np.testing.assert_array_equal(again.toarray(), expected) + + +def test_cache_is_bounded_and_evicts_least_recently_used(vcls, dense): + if dense.sum() == 0: + pytest.skip("all-zero matrix: median row total is 0") + v = vcls.from_scipy(_scipy_for(vcls, dense)) + names = sorted(RECIPES) + assert len(names) > NORM_CACHE_MAXSIZE, "test needs more recipes than the cache holds" + + held = [v.normalized(n) for n in names] + assert len(v._norm_cache) == NORM_CACHE_MAXSIZE + # The first recipe touched is the one dropped. + assert RECIPES[names[0]] not in v._norm_cache + assert RECIPES[names[-1]] in v._norm_cache + assert len(held) == len(names) + + +def test_cache_hit_is_identity_stable_while_the_caller_holds_the_view(vcls, dense): + if dense.sum() == 0: + pytest.skip("all-zero matrix: median row total is 0") + v = vcls.from_scipy(_scipy_for(vcls, dense)) + nv = v.normalized("parafac2") + assert v.normalized("parafac2", recalculate=False) is nv + + +def test_anndata_cache_does_not_pin_a_dropped_view(): + rng = np.random.default_rng(0) + adata = _small_adata(rng) + ref = weakref.ref(adata.normalized("scanpy")) + gc.collect() + assert ref() is None + # Still reusable -- from obs/varm/uns if not from the retained statistics. + assert adata.normalized("scanpy", recalculate=False) is not None From f6d04e3bef5fbb7e113d872f2e4dccf736411b1b Mon Sep 17 00:00:00 2001 From: fishidaho Date: Thu, 10 Sep 2026 14:21:31 -0700 Subject: [PATCH 7/7] Gate normalized-view CPU cost against a sparse baseline, not a dense one Every normalized-view benchmark compared `nv @ B` against `nv.toarray() @ B` -- a dense materialization of the implicit-zero-filled matrix, which nobody would actually perform. It is where the "20-100x faster, 150-200x less peak memory" figures come from, and it flatters the view for a reason that has nothing to do with the kernels: `time_ratio_view_over_materialize` clears its 0.3 ceiling no matter what the per-nonzero cost is. Add the comparison the prototype benchmarks used instead. `_sparse_delta()` builds `Delta` as a scipy CSR -- exactly what a `to_csr()` on the view would return -- and the baseline multiplies that and adds the same rank-1 correction the view adds. Both sides then do identical arithmetic, asserted to rtol=1e-9 inside the case, so the ratio isolates the kernel. CPU is the metric that matters here, not wall: our kernels are `parallel=True` against a single-threaded scipy, so wall time hides a per-nonzero cost a shared workstation still pays. Measured with numba pinned to one thread, on 60k x 2k @ 5% (6M nonzeros): raw (identity) 1.53x scipy pearson (sqrt) 1.89x scanpy (log1p) 5.12x cp10k_log1p (log1p) 5.10x parafac2 (log10(1+1000x)) 5.87x So the transform is essentially the whole cost, consistent with the 2.13x the prototype measured at 336M nonzeros. `normalized_cpu_vs_sparse_1t` gates these; the parallel wall/CPU ratios are recorded for context only, since they depend on the runner's core count. Three things the measurement needed to be reproducible: - `benchmarks/run.py` gives any case whose name ends in `_1t` a subprocess with `NUMBA_NUM_THREADS=1`. It has to be set before numba is imported -- capping the pool at runtime leaves the idle workers spinning and `process_time()` counts them, which swung the ratio 2x between runs. - 6M nonzeros rather than 2M, and `repeat=15`. The identity-transform kernels land within ~1.5x of scipy, so shorter runs left jitter a large share of the ratio. - `gc.collect()` between recipes. Each builds its own ~150 MB delta, and without it the recipes measured last ran against a progressively more fragmented heap -- scanpy alone swung between 2.8x and 6.1x. Ceilings are set from the observed maximum over five runs x 1.5 rather than from one `--record` sample: the ratios have a low tail, and recording off one of those puts the ceiling under the steady-state value. Verified clean over four consecutive runs. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_019ZwuJMcTRUacxqARwjMngE --- benchmarks/baselines.json | 44 ++++++++++- benchmarks/cases.py | 137 +++++++++++++++++++++++++++++++++ benchmarks/harness.py | 18 +++++ benchmarks/run.py | 8 ++ tests/test_vcs_norm_recipes.py | 4 +- 5 files changed, 209 insertions(+), 2 deletions(-) diff --git a/benchmarks/baselines.json b/benchmarks/baselines.json index c94072b..251c274 100644 --- a/benchmarks/baselines.json +++ b/benchmarks/baselines.json @@ -7,7 +7,12 @@ "peak_alloc_mb": 2.0, "time_ratio_vs_scipy": 4.0, "peak_alloc_mb_view": 2.0, - "time_ratio_view_over_materialize": 4.0 + "time_ratio_view_over_materialize": 4.0, + "cpu_ratio_1t_cp10k_log1p": 1.5, + "cpu_ratio_1t_parafac2": 1.5, + "cpu_ratio_1t_pearson": 1.5, + "cpu_ratio_1t_raw": 1.5, + "cpu_ratio_1t_scanpy": 1.5 }, "cases": { "layout_bytes_per_nonzero": { @@ -115,6 +120,43 @@ "normalized_scanpy_rmatvec_vs_materialize": { "time_ratio_view_over_materialize": 0.3, "peak_alloc_mb_view": 0.1321 + }, + "normalized_cp10k_log1p_matmat_vs_sparse": { + "peak_alloc_mb_view": 2.7262 + }, + "normalized_cp10k_log1p_matvec_vs_sparse": { + "peak_alloc_mb_view": 0.3858 + }, + "normalized_parafac2_matmat_vs_sparse": { + "peak_alloc_mb_view": 2.7262 + }, + "normalized_parafac2_matvec_vs_sparse": { + "peak_alloc_mb_view": 0.3858 + }, + "normalized_pearson_matmat_vs_sparse": { + "peak_alloc_mb_view": 2.7262 + }, + "normalized_pearson_matvec_vs_sparse": { + "peak_alloc_mb_view": 0.3858 + }, + "normalized_raw_matmat_vs_sparse": { + "peak_alloc_mb_view": 2.7262 + }, + "normalized_raw_matvec_vs_sparse": { + "peak_alloc_mb_view": 0.3858 + }, + "normalized_scanpy_matmat_vs_sparse": { + "peak_alloc_mb_view": 2.7262 + }, + "normalized_scanpy_matvec_vs_sparse": { + "peak_alloc_mb_view": 0.3858 + }, + "normalized_cpu_vs_sparse_1t": { + "cpu_ratio_1t_cp10k_log1p": 7.815, + "cpu_ratio_1t_parafac2": 8.91, + "cpu_ratio_1t_pearson": 2.85, + "cpu_ratio_1t_raw": 2.445, + "cpu_ratio_1t_scanpy": 7.845 } } } diff --git a/benchmarks/cases.py b/benchmarks/cases.py index 1d69c62..f813ab0 100644 --- a/benchmarks/cases.py +++ b/benchmarks/cases.py @@ -1,10 +1,12 @@ from __future__ import annotations +import gc from collections.abc import Callable import numpy as np from benchmarks.harness import ( + best_cpu_time, best_time, integer_counts_csr, peak_alloc_mb, @@ -214,6 +216,141 @@ def _register_normalized_benchmarks() -> None: _register_normalized_benchmarks() +# -- normalized views vs the *sparse* baseline ------------------------------- +# +# The cases above compare against `nv.toarray() @ B`, a dense materialization +# nobody would actually perform -- it makes the view look good for a reason +# that has nothing to do with the kernels. The honest baseline is the one the +# prototype benchmarks used: build the sparse delta once, multiply that with +# scipy, and add the same rank-1 correction the view adds. Both sides then do +# identical arithmetic and the ratio isolates the kernel. +# +# CPU, not just wall, is the point here. Our kernels are `parallel=True` and +# scipy's are single-threaded, so wall time hides a per-nonzero cost that a +# shared workstation still pays -- and every recipe but "raw" applies a +# transcendental to every stored nonzero. + + +def _sparse_delta(nv, mat): + """``Delta`` as scipy CSR: what ``to_csr()`` on the view would return. + + ``Delta[i, j] = s[j] * g(x[i, j] / row_scale[i] / gene_scale[j])`` on the + stored nonzeros and exactly zero off them, so + ``A_norm = Delta + 1 (x) (-c * s)``. + """ + import scipy.sparse as sp + + from vsparse._norm_common import _g_np + + coo = mat.tocoo() + gs = nv.gene_scale[coo.col] + with np.errstate(divide="ignore", invalid="ignore"): + scaled = np.where(gs > 0.0, coo.data / nv.row_scale[coo.row] / gs, 0.0) + data = nv.col_post_scale[coo.col] * _g_np(scaled, nv.recipe.g_code) + return sp.csr_array((data, (coo.row, coo.col)), shape=mat.shape) + + +def _normalized_vs_sparse(recipe: str, *, vector: bool) -> Callable[[], dict[str, float]]: + def bench() -> dict[str, float]: + from vsparse import VCSRArray + + mat = integer_counts_csr(20_000, 2_000, density=0.05) + v = VCSRArray.from_scipy(mat) + nv = v.normalized(recipe) + rng = np.random.default_rng(0) + B = rng.normal(size=mat.shape[1]) if vector else rng.normal(size=(mat.shape[1], 8)) + + delta = _sparse_delta(nv, mat) + offset = -(nv.col_mean * nv.col_post_scale) + + def via_view() -> np.ndarray: + return nv @ B + + def via_sparse() -> np.ndarray: + return delta @ B + (offset @ B) + + np.testing.assert_allclose(via_view(), via_sparse(), rtol=1e-9, atol=1e-9) + + # Context, and core-count dependent: the kernel is `parallel=True` + # against a single-threaded scipy. The gated, machine-portable CPU + # comparison lives in `normalized_cpu_vs_sparse_1t` below, which is + # also the only case here that pays for CPU timing -- doing it in + # every case doubled the suite's runtime for a number nothing gates. + return { + "wall_ratio_view_over_sparse": best_time(via_view) / best_time(via_sparse), + "peak_alloc_mb_view": peak_alloc_mb(via_view), + "peak_alloc_mb_sparse_delta": peak_alloc_mb(lambda: _sparse_delta(nv, mat)), + } + + bench.__name__ = f"normalized_{recipe}_{'matvec' if vector else 'matmat'}_vs_sparse" + return bench + + +def _register_sparse_baseline_benchmarks() -> None: + from vsparse import RECIPES + + for recipe in sorted(RECIPES): + for vector in (False, True): + fast(_normalized_vs_sparse(recipe, vector=vector)) + + +_register_sparse_baseline_benchmarks() + + +@fast +def normalized_cpu_vs_sparse_1t() -> dict[str, float]: + """Per-nonzero CPU cost of each recipe's kernel, against the same math in scipy. + + Runs with numba pinned to one thread -- `benchmarks.run` gives any case + whose name ends in `_1t` a `NUMBA_NUM_THREADS=1` subprocess, which has to + happen before numba is imported. Capping the pool from inside the process + is not enough: the idle workers still spin, and `process_time()` counts + every thread, which made the ratio swing by 2x between runs. + + Pinned, both sides are single-threaded and the ratio isolates what a + nonzero costs us over scipy -- almost entirely `g`. It is stable to well + under a percent between runs and does not depend on the runner's core + count, which is what makes it gateable. + """ + from vsparse import RECIPES, VCSRArray + + # Deliberately larger than the cases above (6M nonzeros, not 2M). Both + # sides here land within a small multiple of each other, so the ratio only + # settles once each measurement is long enough to swamp scheduling jitter. + mat = integer_counts_csr(60_000, 2_000, density=0.05) + v = VCSRArray.from_scipy(mat) + rng = np.random.default_rng(0) + B = rng.normal(size=(mat.shape[1], 8)) + + out: dict[str, float] = {} + for recipe in sorted(RECIPES): + # Each recipe builds its own 6M-nonzero delta (~150 MB). Left to the + # allocator, the recipes measured last ran against a progressively more + # fragmented heap and their ratios swung by 2x; dropping the previous + # one first keeps every recipe on the same footing. + gc.collect() + nv = v.normalized(recipe) + delta = _sparse_delta(nv, mat) + offset = -(nv.col_mean * nv.col_post_scale) + + def via_view(nv=nv): + return nv @ B + + def via_sparse(delta=delta, offset=offset): + return delta @ B + (offset @ B) + + np.testing.assert_allclose(via_view(), via_sparse(), rtol=1e-9, atol=1e-9) + # repeat=15: the identity-transform kernels land within ~1.5x of scipy, + # so at the default repeat count run-to-run jitter was a large share of + # the ratio and the gate was not reproducible across processes. + out[f"cpu_ratio_1t_{recipe}"] = best_cpu_time(via_view, repeat=15) / best_cpu_time( + via_sparse, repeat=15 + ) + del nv, delta, via_view, via_sparse + + return out + + # -- larger, for the scheduled job ------------------------------------------- diff --git a/benchmarks/harness.py b/benchmarks/harness.py index 7a9627d..ada58dc 100644 --- a/benchmarks/harness.py +++ b/benchmarks/harness.py @@ -39,6 +39,24 @@ def best_time(fn: Callable[[], Any], repeat: int = 7) -> float: return best +def best_cpu_time(fn: Callable[[], Any], repeat: int = 7) -> float: + """Best *CPU* time over ``repeat`` runs, in seconds. Warms up first. + + Wall time alone flatters every numba kernel here: they are ``parallel=True`` + while scipy's sparse matmul is single-threaded, so a kernel can be several + times faster on the clock while burning an order of magnitude more CPU. + On a shared workstation -- the machine this project is aimed at -- that + difference is what the user actually pays. + """ + fn() + best = float("inf") + for _ in range(repeat): + start = time.process_time() + fn() + best = min(best, time.process_time() - start) + return best + + def ratio_vs_scipy(ours: Callable[[], Any], theirs: Callable[[], Any], repeat: int = 7) -> float: """``our time / scipy's time`` for the same work.""" return best_time(ours, repeat) / best_time(theirs, repeat) diff --git a/benchmarks/run.py b/benchmarks/run.py index 50cda8d..bb40711 100644 --- a/benchmarks/run.py +++ b/benchmarks/run.py @@ -12,6 +12,7 @@ import argparse import json +import os import subprocess import sys from pathlib import Path @@ -20,11 +21,18 @@ def _run_one_in_subprocess(name: str) -> dict[str, float]: + env = os.environ.copy() + if name.endswith("_1t"): + # Has to be set before numba is imported, which is why these cases are + # named rather than configured: capping the pool at runtime leaves the + # idle workers spinning, and process_time() counts them. + env["NUMBA_NUM_THREADS"] = "1" proc = subprocess.run( [sys.executable, "-m", "benchmarks.run", "--emit", name], capture_output=True, text=True, cwd=Path(__file__).resolve().parent.parent, + env=env, check=False, ) if proc.returncode != 0: diff --git a/tests/test_vcs_norm_recipes.py b/tests/test_vcs_norm_recipes.py index 7b9fc40..f91b01e 100644 --- a/tests/test_vcs_norm_recipes.py +++ b/tests/test_vcs_norm_recipes.py @@ -285,12 +285,14 @@ def test_custom_recipe_works_on_the_anndata(): adata = _small_adata(rng) recipe = _custom_recipe() nv = adata.normalized(recipe) + raw = adata.X + assert raw is not None assert nv.recipe is recipe assert adata.uns["vsparse"]["recipe"] == "custom_cp10k_scaled" np.testing.assert_allclose(adata.obs["vsparse_a"].to_numpy(), nv.a) np.testing.assert_allclose( - nv.toarray(), _reference(np.asarray(adata.X.toarray()), "scanpy"), atol=1e-6 + nv.toarray(), _reference(np.asarray(raw.toarray()), "scanpy"), atol=1e-6 )