From e6ac786718bd225d30623aa83394a3afd09c4705 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Sat, 12 Sep 2026 19:05:12 -0700 Subject: [PATCH 1/6] Add norm_sq/slice_norms/to_scipy_sparse to normalized views parafac2 touches its input matrix through matmul/rmatmul (already supported via __matmul__/__rmatmul__) plus a squared-Frobenius-norm reduction it currently only knows how to compute for a plain np.ndarray or scipy.sparse array. Add that as norm_sq()/slice_norms() on NormalizedViewBase, computed with new numba kernels that reuse the same O(nnz) traversal as the existing matmul/toarray kernels, so a normalized view can satisfy parafac2's duck-typed backend contract without materializing anything. Also add to_scipy_sparse() (the uncentered, scaled sparse term with the same sparsity pattern as the raw array) and a means property (the per-gene correction to subtract from it), for code that only knows how to move a plain NumPy/SciPy array onto a device -- e.g. so a normalized view can be materialized into a real (and, for typical single-cell data, tiny) sparse array before running through parafac2's existing CuPy/MLX GPU path, rather than needing new GPU-native kernels of its own. Co-Authored-By: Claude Sonnet 5 --- src/vsparse/_norm_common.py | 319 ++++++++++++++++++++++++++++++ tests/test_property_norm_stats.py | 88 +++++++++ tests/test_vcs_norm.py | 46 +++++ 3 files changed, 453 insertions(+) create mode 100644 tests/test_property_norm_stats.py diff --git a/src/vsparse/_norm_common.py b/src/vsparse/_norm_common.py index ae961fc..9b6d406 100644 --- a/src/vsparse/_norm_common.py +++ b/src/vsparse/_norm_common.py @@ -386,6 +386,205 @@ def _fill_normalized_major_is_row( out[i, c] = (_g(scaled, g_code) - col_mean[c]) * col_post_scale[c] +# -- squared-norm statistics (whole array + per-condition slices) ----------- +# +# Both reductions use the same ``||A_norm||_F^2 = sum(Delta^2) - 2 sum(Delta * +# offset[col]) + n_rows * sum(offset^2)`` expansion the *sparse-plus-external- +# means* code path in ``parafac2.utils.calc_norm_sq``/``calc_slice_norms`` +# uses, just carried out against ``Delta`` (this view's own uncentered, +# scaled sparse term -- see :mod:`vsparse._vcs_matmul`) instead of a plain +# scipy ``data``/``means`` pair, since ``offset = col_post_scale * col_mean`` +# is exactly the external ``means`` that convention expects. +# +# For VCSR (major=rows), a whole major slice belongs to exactly one row, so +# ``sum(Delta^2)``/``sum(Delta * offset)`` are pure scalar (norm_sq) or +# per-condition (slice_norms) reductions with no cross-thread write conflict: +# numba recognizes plain ``+=`` accumulation in a ``prange`` loop as a +# reduction. For VCSC (major=cols), a column's nonzeros span many rows/ +# conditions, so ``slice_norms`` needs the same thread-chunked scatter as +# :func:`_gstats_col_sums_vcs`; ``norm_sq`` only ever needs a scalar, so it +# stays reduction-only even there. + + +@numba.njit(cache=True, parallel=True) +def _norm_sq_terms_major_is_row( + major_ptr, values, value_ptr, indices, row_scale, gene_scale, col_mean, col_post_scale, g_code +): + n_major = major_ptr.shape[0] - 1 + total_sq = 0.0 + total_cross = 0.0 + for i in numba.prange(n_major): # ty: ignore[not-iterable] + rs = row_scale[i] + for u in range(major_ptr[i], major_ptr[i + 1]): + v = values[u] + for k in range(value_ptr[u], value_ptr[u + 1]): + col = indices[k] + gs = gene_scale[col] + if gs <= 0.0: + continue + s = col_post_scale[col] + delta = s * _g(v / rs / gs, g_code) + total_sq += delta * delta + total_cross += delta * col_mean[col] * s + return total_sq, total_cross + + +@numba.njit(cache=True, parallel=True) +def _norm_sq_terms_major_is_col( + major_ptr, values, value_ptr, indices, row_scale, gene_scale, col_mean, col_post_scale, g_code +): + n_major = major_ptr.shape[0] - 1 + total_sq = 0.0 + total_cross = 0.0 + 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] + offset = col_mean[j] * s + col_sq = 0.0 + col_sum = 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]): + row = indices[k] + delta = s * _g(v / row_scale[row] / gs, g_code) + col_sq += delta * delta + col_sum += delta + total_sq += col_sq + total_cross += col_sum * offset + return total_sq, total_cross + + +@numba.njit(cache=True, parallel=True) +def _slice_norm_terms_major_is_row( + major_ptr, + values, + value_ptr, + indices, + row_scale, + gene_scale, + col_mean, + col_post_scale, + g_code, + condition_idxs, + n_cond, + nthreads, +): + n_major = major_ptr.shape[0] - 1 + chunk = (n_major + nthreads - 1) // nthreads + partial_sq = np.zeros((nthreads, n_cond), dtype=np.float64) + partial_cross = np.zeros((nthreads, n_cond), dtype=np.float64) + for t in numba.prange(nthreads): # ty: ignore[not-iterable] + start = t * chunk + end = min(n_major, start + chunk) + loc_sq = partial_sq[t] + loc_cross = partial_cross[t] + for i in range(start, end): + rs = row_scale[i] + cond = condition_idxs[i] + row_sq = 0.0 + row_cross = 0.0 + for u in range(major_ptr[i], major_ptr[i + 1]): + v = values[u] + for k in range(value_ptr[u], value_ptr[u + 1]): + col = indices[k] + gs = gene_scale[col] + if gs <= 0.0: + continue + s = col_post_scale[col] + delta = s * _g(v / rs / gs, g_code) + row_sq += delta * delta + row_cross += delta * col_mean[col] * s + loc_sq[cond] += row_sq + loc_cross[cond] += row_cross + return partial_sq.sum(axis=0), partial_cross.sum(axis=0) + + +@numba.njit(cache=True, parallel=True) +def _slice_norm_terms_major_is_col( + major_ptr, + values, + value_ptr, + indices, + row_scale, + gene_scale, + col_mean, + col_post_scale, + g_code, + condition_idxs, + n_cond, + nthreads, +): + n_major = major_ptr.shape[0] - 1 + chunk = (n_major + nthreads - 1) // nthreads + partial_sq = np.zeros((nthreads, n_cond), dtype=np.float64) + partial_cross = np.zeros((nthreads, n_cond), dtype=np.float64) + for t in numba.prange(nthreads): # ty: ignore[not-iterable] + start = t * chunk + end = min(n_major, start + chunk) + loc_sq = partial_sq[t] + loc_cross = partial_cross[t] + for j in range(start, end): + gs = gene_scale[j] + if gs <= 0.0: + continue + s = col_post_scale[j] + cm = col_mean[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]): + row = indices[k] + cond = condition_idxs[row] + delta = s * _g(v / row_scale[row] / gs, g_code) + loc_sq[cond] += delta * delta + loc_cross[cond] += delta * cm * s + return partial_sq.sum(axis=0), partial_cross.sum(axis=0) + + +# -- uncentered sparse materialization --------------------------------------- +# +# Unlike ``toarray``'s dense fill, these write only the structural nonzeros +# (``Delta``, see :mod:`vsparse._vcs_matmul`) into a flat ``nnz``-length +# buffer, in exactly the position ``indices`` already gives each one -- so +# ``(data, indices, value_ptr[major_ptr])`` is directly a valid scipy CSR/CSC +# triple with the same sparsity pattern as the underlying raw array. The +# per-gene mean correction (:attr:`NormalizedViewBase.means`) is left for the +# caller to subtract externally, matching the convention +# ``parafac2.utils.calc_norm_sq``/``calc_W`` already use for a sparse ``X`` +# plus a separate ``means`` vector. + + +@numba.njit(cache=True, parallel=True) +def _materialize_delta_major_is_row( + major_ptr, values, value_ptr, indices, row_scale, gene_scale, col_post_scale, g_code, data +): + n_major = major_ptr.shape[0] - 1 + for i in numba.prange(n_major): # ty: ignore[not-iterable] + rs = row_scale[i] + for u in range(major_ptr[i], major_ptr[i + 1]): + v = values[u] + for k in range(value_ptr[u], value_ptr[u + 1]): + col = indices[k] + gs = gene_scale[col] + data[k] = col_post_scale[col] * _g(v / rs / gs, g_code) if gs > 0.0 else 0.0 + + +@numba.njit(cache=True, parallel=True) +def _materialize_delta_major_is_col( + major_ptr, values, value_ptr, indices, row_scale, gene_scale, col_post_scale, g_code, data +): + n_major = major_ptr.shape[0] - 1 + for j in numba.prange(n_major): # ty: ignore[not-iterable] + gs = gene_scale[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]): + row = indices[k] + data[k] = s * _g(v / row_scale[row] / gs, g_code) if gs > 0.0 else 0.0 + + def _prep_key(key: Any) -> Any: """Turn a bare int into a length-1 list, so fancy indexing never drops that axis.""" if isinstance(key, int | np.integer): @@ -616,6 +815,17 @@ def s(self) -> np.ndarray: """Per-gene post-scale.""" return self.col_post_scale + @property + def means(self) -> np.ndarray: + """Per-gene mean-correction vector, ``col_post_scale * col_mean``. + + This is the ``means`` a caller following the ``parafac2``-style + convention (a sparse/uncentered matrix plus a separate per-column + ``means`` vector) should subtract externally: ``self.toarray() == + self.to_scipy_sparse().toarray() - self.means``. + """ + return self.col_mean * self.col_post_scale + @property def shape(self) -> tuple[int, int]: return self._arr.shape @@ -664,6 +874,115 @@ def toarray(self) -> np.ndarray: ) return out + def to_scipy_sparse(self) -> Any: + """The uncentered, scaled sparse ``Delta`` term, as a real scipy sparse array. + + Same sparsity pattern as the underlying raw array (a ``csr_array`` + for a VCSR-backed view, ``csc_array`` for VCSC), with :attr:`means` + left to subtract externally -- see :attr:`means`. Useful for handing + this view to code (such as ``parafac2``'s CuPy/MLX GPU backends) + that only knows how to move a plain NumPy/SciPy array onto a device, + rather than this view's own ``__matmul__``/``__rmatmul__``. + """ + import scipy.sparse as sp + + arr = self._arr + data = np.empty(arr.nnz, dtype=np.float64) + if self._format == "csc": + _materialize_delta_major_is_col( + arr.major_ptr, + arr.values, + arr.value_ptr, + arr.indices, + self.row_scale, + self.gene_scale, + self.col_post_scale, + self.recipe.g_code, + data, + ) + ctor = sp.csc_array + else: + _materialize_delta_major_is_row( + arr.major_ptr, + arr.values, + arr.value_ptr, + arr.indices, + self.row_scale, + self.gene_scale, + self.col_post_scale, + self.recipe.g_code, + data, + ) + ctor = sp.csr_array + + indptr = arr.value_ptr[arr.major_ptr] + return ctor((data, arr.indices, indptr), shape=self.shape) + + def norm_sq(self) -> float: + """Squared Frobenius norm of the full normalized matrix, in ``O(nnz + n_cols)``.""" + arr = self._arr + args = ( + 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, + ) + if self._format == "csc": + total_sq, total_cross = _norm_sq_terms_major_is_col(*args) + else: + total_sq, total_cross = _norm_sq_terms_major_is_row(*args) + + n_rows = self.shape[0] + offset_sq_sum = float(np.sum(self.means**2)) + return float(total_sq - 2.0 * total_cross + n_rows * offset_sq_sum) + + def slice_norms(self, condition_idxs: Any, n_cond: int) -> np.ndarray: + """Per-condition Frobenius norm of the normalized matrix's rows. + + Parameters + ---------- + condition_idxs : array-like of int + Condition index (in ``[0, n_cond)``) for each row. + n_cond : int + The total number of conditions. + + Returns + ------- + np.ndarray + Length-``n_cond`` array of each condition's rows' Frobenius norm. + """ + idxs = np.asarray(condition_idxs, dtype=np.int64) + arr = self._arr + nthreads = numba.get_num_threads() + args = ( + 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, + idxs, + n_cond, + nthreads, + ) + if self._format == "csc": + total_sq, total_cross = _slice_norm_terms_major_is_col(*args) + else: + total_sq, total_cross = _slice_norm_terms_major_is_row(*args) + + counts = np.bincount(idxs, minlength=n_cond).astype(np.float64) + offset_sq_sum = float(np.sum(self.means**2)) + sq = total_sq - 2.0 * total_cross + counts * offset_sq_sum + return np.sqrt(np.clip(sq, 0.0, None)) + # -- selection --------------------------------------------------------------- def select(self, rows: Any = slice(None), cols: Any = slice(None)) -> Any: diff --git a/tests/test_property_norm_stats.py b/tests/test_property_norm_stats.py new file mode 100644 index 0000000..ed93f76 --- /dev/null +++ b/tests/test_property_norm_stats.py @@ -0,0 +1,88 @@ +"""Property-based tests for the norm_sq/slice_norms/to_scipy_sparse trio. + +These give ``VCSCArrayNormalized``/``VCSRArrayNormalized`` the pieces a +duck-typed ``parafac2`` backend needs beyond ``__matmul__``/``__rmatmul__`` +(see https://github.com/meyer-lab/parafac2's ``parafac2.utils`` module +docstring): a squared-Frobenius-norm reduction, a per-condition-group version +of the same, and a way to materialize the underlying sparse term as a real +scipy array plus its external mean correction. Each is checked directly +against a dense/numpy reference built from :meth:`toarray`. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import scipy.sparse as sp +from _hypothesis_strategies import dense_matrices, slow_first_call +from hypothesis import assume, given +from hypothesis import strategies as st + +from vsparse import RECIPES, VCSCArray, VCSRArray + + +def _scipy_for(vcls, dense): + return sp.csc_array(dense) if vcls is VCSCArray else sp.csr_array(dense) + + +@st.composite +def condition_groups(draw, *, n_rows: int): + """A random condition index (in ``[0, n_cond)``) for each of ``n_rows`` rows.""" + if n_rows == 0: + return np.zeros(0, dtype=np.int64), 1 + n_cond = draw(st.integers(1, max(1, n_rows))) + idxs = draw(st.lists(st.integers(0, n_cond - 1), min_size=n_rows, max_size=n_rows)) + return np.asarray(idxs, dtype=np.int64), n_cond + + +@pytest.mark.parametrize("vcls", [VCSCArray, VCSRArray]) +@pytest.mark.parametrize("recipe", sorted(RECIPES)) +@slow_first_call +@given(dense=dense_matrices()) +def test_norm_sq_matches_dense_reference(vcls, recipe, dense): + assume(dense.sum() > 0) + nv = vcls.from_scipy(_scipy_for(vcls, dense)).normalized(recipe) + expected = float(np.sum(nv.toarray() ** 2)) + assert nv.norm_sq() >= -1e-6 + np.testing.assert_allclose(nv.norm_sq(), expected, rtol=1e-6, atol=1e-6) + + +@pytest.mark.parametrize("vcls", [VCSCArray, VCSRArray]) +@slow_first_call +@given(dense=dense_matrices(), data=st.data()) +def test_slice_norms_matches_dense_reference(vcls, dense, data): + assume(dense.sum() > 0) + idxs, n_cond = data.draw(condition_groups(n_rows=dense.shape[0])) + + nv = vcls.from_scipy(_scipy_for(vcls, dense)).normalized() + ref = nv.toarray() + expected = np.array([np.linalg.norm(ref[idxs == i]) for i in range(n_cond)]) + + np.testing.assert_allclose(nv.slice_norms(idxs, n_cond), expected, rtol=1e-5, atol=1e-4) + + +@pytest.mark.parametrize("vcls", [VCSCArray, VCSRArray]) +@pytest.mark.parametrize("recipe", sorted(RECIPES)) +@slow_first_call +@given(dense=dense_matrices()) +def test_to_scipy_sparse_plus_means_reconstructs_toarray(vcls, recipe, dense): + assume(dense.sum() > 0) + nv = vcls.from_scipy(_scipy_for(vcls, dense)).normalized(recipe) + + sparse = nv.to_scipy_sparse() + assert sparse.dtype == np.float64 + assert sparse.shape == dense.shape + assert isinstance(sparse, sp.csc_array if vcls is VCSCArray else sp.csr_array) + + reconstructed = sparse.toarray() - nv.means + np.testing.assert_allclose(reconstructed, nv.toarray(), atol=1e-6) + + +@pytest.mark.parametrize("vcls", [VCSCArray, VCSRArray]) +@slow_first_call +@given(dense=dense_matrices()) +def test_to_scipy_sparse_has_the_same_sparsity_pattern_as_the_raw_array(vcls, dense): + assume(dense.sum() > 0) + raw = _scipy_for(vcls, dense) + nv = vcls.from_scipy(raw).normalized() + assert nv.to_scipy_sparse().nnz == raw.nnz diff --git a/tests/test_vcs_norm.py b/tests/test_vcs_norm.py index 5da4f6c..09f60d1 100644 --- a/tests/test_vcs_norm.py +++ b/tests/test_vcs_norm.py @@ -150,3 +150,49 @@ def test_transpose_major_roundtrip(dense, vcls): assert dual.shape == v.shape assert dual._format != v._format np.testing.assert_allclose(dual.toarray(), dense) + + +# -- norm_sq / slice_norms / to_scipy_sparse: numeric correctness is +# property-tested in test_property_norm_stats.py ----------------------------- + + +def test_norm_sq_all_zero_matrix_is_zero(vcls): + dense = np.zeros((5, 4)) + v = vcls.from_scipy(_scipy_for(vcls, dense)) + nv = v.normalized() + assert nv.norm_sq() == 0.0 + + +def test_slice_norms_all_zero_matrix_is_all_zero(vcls): + dense = np.zeros((5, 4)) + v = vcls.from_scipy(_scipy_for(vcls, dense)) + nv = v.normalized() + np.testing.assert_allclose(nv.slice_norms(np.array([0, 0, 1, 1, 1]), 2), 0.0) + + +def test_to_scipy_sparse_all_zero_matrix_is_empty(vcls): + dense = np.zeros((5, 4)) + v = vcls.from_scipy(_scipy_for(vcls, dense)) + nv = v.normalized() + sparse = nv.to_scipy_sparse() + assert sparse.shape == (5, 4) + assert sparse.nnz == 0 + np.testing.assert_allclose(nv.means, 0.0) + + +def test_to_scipy_sparse_matches_format(dense, vcls): + 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() + sparse = nv.to_scipy_sparse() + assert isinstance(sparse, sp.csc_array if vcls is VCSCArray else sp.csr_array) + assert sparse.dtype == np.float64 + + +def test_means_property_matches_c_times_s(dense, vcls): + 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() + np.testing.assert_allclose(nv.means, nv.c * nv.s) From ab73fba19bba77c2cbcf064e094462d295175966 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Sat, 12 Sep 2026 19:50:22 -0700 Subject: [PATCH 2/6] Fix VCSCAnnData.copy() silently dropping X/raw_X VCSCAnnData stores X/raw_X in private _vcs_X/_vcs_raw_X attributes (anndata's own X validation rejects a VCSCArray/VCSRArray directly), but never overrode copy(), so the inherited anndata.AnnData.copy() copies the standard (unused, always-None) _X attribute instead: X silently comes back None, and the returned object is downgraded to a plain AnnData rather than preserving this class (or a subclass, such as one overriding the X property for a lazy-normalized view). Any caller relying on `adata[mask].copy()` -- e.g. parafac2's BiCV train/test splitting -- hits this immediately. Add an explicit override that copies every field, including a real VCSCArray/VCSRArray copy of X/raw_X, and returns type(self). Co-Authored-By: Claude Sonnet 5 --- src/vsparse/_anndata_class.py | 36 +++++++++++++++++++ tests/test_anndata_class.py | 68 +++++++++++++++++++++++++++++++++++ 2 files changed, 104 insertions(+) diff --git a/src/vsparse/_anndata_class.py b/src/vsparse/_anndata_class.py index f514b83..6e79d6a 100644 --- a/src/vsparse/_anndata_class.py +++ b/src/vsparse/_anndata_class.py @@ -2,6 +2,7 @@ from __future__ import annotations +import copy as _copy from types import MappingProxyType from typing import TYPE_CHECKING, Any, cast @@ -69,6 +70,16 @@ def _subset_2d(v: Any, oidx: Any, vidx: Any) -> Any: return np.asarray(v)[oidx][:, vidx] +def _copy_value(v: Any) -> Any: + """A deep-enough copy of one obs/var/obsm/varm/obsp/varp/layers value. + + Every value anndata can hold there -- a DataFrame, ndarray, scipy sparse + array, or a :class:`~vsparse.VCSCArray`/:class:`~vsparse.VCSRArray` -- + implements its own ``.copy()``. + """ + return None if v is None else v.copy() + + def _check_vcs_type(value: Any, name: str) -> None: if value is not None and not isinstance(value, _VCS_TYPES): raise TypeError( @@ -210,6 +221,31 @@ 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}, ) + def copy(self) -> VCSCAnnData: # ty: ignore[invalid-method-override] + """A deep copy, preserving the VCSC/VCSR-backed ``X``/``raw_X``. + + The inherited :meth:`anndata.AnnData.copy` only knows how to copy the + standard private ``_X`` attribute, which this class never sets (see + the class docstring: ``X``/``raw_X`` live in ``_vcs_X``/``_vcs_raw_X`` + instead) -- so it silently drops them (``X`` comes back ``None``) and + returns a plain ``AnnData`` rather than this class. This copies every + field explicitly instead, including a real ``VCSCArray``/ + ``VCSRArray`` copy of ``X``/``raw_X``, and returns ``type(self)`` so a + subclass (e.g. one overriding the ``X`` property) round-trips too. + """ + return type(self)( + X=_copy_value(self._vcs_X), + raw_X=_copy_value(self._vcs_raw_X), + obs=cast(pd.DataFrame, self.obs).copy(), + var=cast(pd.DataFrame, self.var).copy(), + uns=_copy.deepcopy(dict(self.uns)), + obsm={k: _copy_value(v) for k, v in self.obsm.items() if k is not None}, + varm={k: _copy_value(v) for k, v in self.varm.items() if k is not None}, + obsp={k: _copy_value(v) for k, v in self.obsp.items() if k is not None}, + varp={k: _copy_value(v) for k, v in self.varp.items() if k is not None}, + layers={k: _copy_value(v) for k, v in self.layers.items() if k is not None}, + ) + # -- normalization ---------------------------------------------------------- def normalized(self, view: str | Recipe = DEFAULT_RECIPE, *, recalculate: bool = True) -> Any: diff --git a/tests/test_anndata_class.py b/tests/test_anndata_class.py index 1c4c06d..7cbd717 100644 --- a/tests/test_anndata_class.py +++ b/tests/test_anndata_class.py @@ -251,3 +251,71 @@ def test_getitem_single_int_row(base_adata, dense): assert sub.shape == (1, dense.shape[1]) assert isinstance(sub.X, VCSCArray) np.testing.assert_allclose(sub.X.toarray(), dense[0:1]) + + +# -- copy() -- the inherited anndata.AnnData.copy() only knows how to copy +# the standard private _X attribute, which this class never sets (X/raw_X +# live in _vcs_X/_vcs_raw_X instead), so it silently drops them. + + +def test_copy_preserves_x_and_raw_x(base_adata, dense): + va = VCSCAnnData.from_anndata(base_adata) + c = va.copy() + assert type(c) is VCSCAnnData + assert isinstance(c.X, VCSCArray) + assert isinstance(c.raw_X, VCSCArray) + np.testing.assert_allclose(c.X.toarray(), dense) + np.testing.assert_allclose(c.raw_X.toarray(), dense) + + +def test_copy_is_independent_of_the_original(base_adata, dense): + if dense.shape[0] == 0 or dense.shape[1] == 0: + pytest.skip("shape too small") + va = VCSCAnnData.from_anndata(base_adata, include_raw=False) + c = va.copy() + assert c.X is not va.X + assert c.obs is not va.obs + + c.obs["grp"] = "mutated" + assert list(va.obs["grp"]) != list(c.obs["grp"]) + + +def test_copy_preserves_obs_var_uns(base_adata, dense): + base_adata.uns["note"] = {"k": "v"} + va = VCSCAnnData.from_anndata(base_adata, include_raw=False) + c = va.copy() + assert list(c.obs["grp"]) == list(va.obs["grp"]) + assert list(c.var["gene"]) == list(va.var["gene"]) + assert c.uns["note"] == {"k": "v"} + + +def test_copy_preserves_a_normalized_x_subclass(base_adata, dense): + """A subclass overriding the `X` getter (as BAL-Pf2's own + lazy-normalized-view AnnData does) must round-trip through copy().""" + if dense.sum() == 0: + pytest.skip("all-zero matrix: median row total is 0") + + class _NormalizedView(VCSCAnnData): + @property + def X(self): + return self.normalized("parafac2") + + nv = _NormalizedView.from_anndata(base_adata, include_raw=False) + c = nv.copy() + assert type(c) is _NormalizedView + c_x, nv_x = c.X, nv.X + assert c_x is not None + assert nv_x is not None + np.testing.assert_allclose(c_x.toarray(), nv_x.toarray()) + + +def test_copy_of_a_slice_round_trips_x(base_adata, dense): + """The exact pattern a caller like `parafac2`'s BiCV split uses: + `adata[obs_mask][:, var_mask].copy()`.""" + if dense.shape[0] < 2 or dense.shape[1] < 2: + pytest.skip("shape too small") + va = VCSCAnnData.from_anndata(base_adata, include_raw=False) + sub = va[0:2, 0:2] + c = sub.copy() + assert c.X is not None + np.testing.assert_allclose(c.X.toarray(), dense[0:2, 0:2]) From b356d9db54561b415787d7aec8692ff202507ff6 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Sat, 12 Sep 2026 20:00:16 -0700 Subject: [PATCH 3/6] Fix VCSCAnnData.to_memory() dropping X/raw_X the same way copy() did The inherited anndata.AnnData.to_memory() has the same root cause as copy() (previous commit): it only knows about the standard, unused _X attribute, not this class's _vcs_X/_vcs_raw_X, and reconstructs a plain AnnData that silently loses X. This class never actually supports a lazily backed X/raw_X, so to_memory() now just delegates to copy(). scrise's BiCV rank selection calls to_memory() on its input before the train/test split loop, so this was hit immediately after fixing copy(). Co-Authored-By: Claude Sonnet 5 --- src/vsparse/_anndata_class.py | 18 ++++++++++++++++++ tests/test_anndata_class.py | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+) diff --git a/src/vsparse/_anndata_class.py b/src/vsparse/_anndata_class.py index 6e79d6a..c7cb95c 100644 --- a/src/vsparse/_anndata_class.py +++ b/src/vsparse/_anndata_class.py @@ -246,6 +246,24 @@ def copy(self) -> VCSCAnnData: # ty: ignore[invalid-method-override] layers={k: _copy_value(v) for k, v in self.layers.items() if k is not None}, ) + def to_memory(self, *, copy: bool = False) -> VCSCAnnData: + """Return this object with its data loaded into memory. + + The inherited :meth:`anndata.AnnData.to_memory` has the same problem + as the inherited ``copy()`` (see above): it iterates the object's + *standard* attributes, which never includes this class's ``X``/ + ``raw_X`` (held in ``_vcs_X``/``_vcs_raw_X`` instead), and reconstructs + a plain ``AnnData`` -- so ``X`` silently comes back ``None``. + + This class never actually supports a lazily backed ``X``/``raw_X`` + (they're always eagerly-held ``VCSCArray``/``VCSRArray`` instances), + so there is never anything to load -- this always returns a full + :meth:`copy` instead, regardless of ``copy`` (unlike plain + ``AnnData``, where ``copy=False`` can skip copying arrays already in + memory; here everything already is, so the distinction doesn't apply). + """ + return self.copy() + # -- normalization ---------------------------------------------------------- def normalized(self, view: str | Recipe = DEFAULT_RECIPE, *, recalculate: bool = True) -> Any: diff --git a/tests/test_anndata_class.py b/tests/test_anndata_class.py index 7cbd717..e4e56ba 100644 --- a/tests/test_anndata_class.py +++ b/tests/test_anndata_class.py @@ -319,3 +319,35 @@ def test_copy_of_a_slice_round_trips_x(base_adata, dense): c = sub.copy() assert c.X is not None np.testing.assert_allclose(c.X.toarray(), dense[0:2, 0:2]) + + +# -- to_memory() -- same underlying problem as copy(): the inherited +# anndata.AnnData.to_memory() doesn't know about _vcs_X/_vcs_raw_X either. + + +def test_to_memory_preserves_x_and_raw_x(base_adata, dense): + va = VCSCAnnData.from_anndata(base_adata) + m = va.to_memory() + assert type(m) is VCSCAnnData + assert isinstance(m.X, VCSCArray) + assert isinstance(m.raw_X, VCSCArray) + np.testing.assert_allclose(m.X.toarray(), dense) + np.testing.assert_allclose(m.raw_X.toarray(), dense) + + +def test_to_memory_preserves_a_normalized_x_subclass(base_adata, dense): + if dense.sum() == 0: + pytest.skip("all-zero matrix: median row total is 0") + + class _NormalizedView(VCSCAnnData): + @property + def X(self): + return self.normalized("parafac2") + + nv = _NormalizedView.from_anndata(base_adata, include_raw=False) + m = nv.to_memory() + assert type(m) is _NormalizedView + m_x, nv_x = m.X, nv.X + assert m_x is not None + assert nv_x is not None + np.testing.assert_allclose(m_x.toarray(), nv_x.toarray()) From 2bb7fe616c7d1b176d91e9d6b11d272fde08b9e2 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Sat, 12 Sep 2026 20:27:26 -0700 Subject: [PATCH 4/6] Fix VCSCAnnData.__getitem__ dropping a subclass's X override __getitem__ already returned a new, eagerly-copied object (rather than a lazy view) since anndata's own view machinery can't handle _vcs_X, but it hardcoded the returned type to VCSCAnnData instead of type(self) -- so a subclass overriding X (e.g. one that always hands back a value normalized fresh from _vcs_X, as with copy()/to_memory() in the previous two commits) silently reverted to the raw, un-normalized array after any slice. scrise's BiCV rank selection slices its input for every train/test split, so this surfaced immediately after fixing copy()/to_memory(): the sliced object's X was the raw VCSRArray, which doesn't implement the norm_sq() a normalized view needs. Co-Authored-By: Claude Sonnet 5 --- src/vsparse/_anndata_class.py | 7 ++++++- tests/test_anndata_class.py | 23 +++++++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/src/vsparse/_anndata_class.py b/src/vsparse/_anndata_class.py index c7cb95c..ee284b0 100644 --- a/src/vsparse/_anndata_class.py +++ b/src/vsparse/_anndata_class.py @@ -187,6 +187,11 @@ def __getitem__(self, index: Any) -> VCSCAnnData: # ty: ignore[invalid-method-o this class doesn't use (see the class docstring), so it can't be reused here. ``X``/``raw_X`` stay VCSC/VCSR-backed either way, via that array type's own indexing. + + Returns ``type(self)``, not a bare ``VCSCAnnData``, so a subclass + overriding ``X`` (e.g. one that always hands back a normalized view + computed fresh from ``_vcs_X``, as with :meth:`copy`/:meth:`to_memory`) + keeps that behavior after slicing. """ oidx, vidx = self._normalize_indices(index) oidx = _as_slice_index(oidx, self.n_obs) @@ -208,7 +213,7 @@ def __getitem__(self, index: Any) -> VCSCAnnData: # ty: ignore[invalid-method-o if _VSPARSE_UNS_KEY in uns: uns = {**uns, _VSPARSE_UNS_KEY: {**uns[_VSPARSE_UNS_KEY], "stale": True}} - return VCSCAnnData( + return type(self)( X=_subset_2d(self._vcs_X, oidx, vidx), raw_X=_subset_2d(self._vcs_raw_X, oidx, vidx), obs=obs, diff --git a/tests/test_anndata_class.py b/tests/test_anndata_class.py index e4e56ba..141ef17 100644 --- a/tests/test_anndata_class.py +++ b/tests/test_anndata_class.py @@ -253,6 +253,29 @@ def test_getitem_single_int_row(base_adata, dense): np.testing.assert_allclose(sub.X.toarray(), dense[0:1]) +def test_getitem_preserves_a_subclass_overriding_x(base_adata, dense): + """A subclass overriding the `X` getter (as BAL-Pf2's own + lazy-normalized-view AnnData does) must keep that behavior after + slicing, not silently fall back to the raw (un-normalized) array.""" + if dense.shape[0] < 2 or dense.sum() == 0: + pytest.skip("shape too small or all-zero matrix") + + class _NormalizedView(VCSCAnnData): + @property + def X(self): + return self.normalized("parafac2") + + nv = _NormalizedView.from_anndata(base_adata, include_raw=False) + sub = nv[0:2, :] + assert type(sub) is _NormalizedView + sub_x = sub.X + assert sub_x is not None + assert not isinstance(sub_x, VCSCArray) # the normalized view, not the raw array + + expected = VCSCArray.from_scipy(sp.csc_array(dense[0:2])).normalized("parafac2") + np.testing.assert_allclose(sub_x.toarray(), expected.toarray()) + + # -- copy() -- the inherited anndata.AnnData.copy() only knows how to copy # the standard private _X attribute, which this class never sets (X/raw_X # live in _vcs_X/_vcs_raw_X instead), so it silently drops them. From eaf7e300be875dbb7bb3815a979a79ca58c0e371 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Sun, 13 Sep 2026 11:42:45 -0700 Subject: [PATCH 5/6] Replace chunked-transpose matmul fallback with a native minor-axis kernel The misaligned direction of a normalized view's matmul (self@B for VCSC, B@self for VCSR) previously regrouped the array into the other VCS format a chunk of major slices at a time (_chunk_bounds/_transpose_major), caching a full opposite-format copy (_dual_arr) when the whole array fit one chunk's budget. At scale (~2.3B nonzeros, 1.3M x 12.2K), this needed ~1,150 chunks per pass, repeated across every power iteration and every rank/trial in a BiCV sweep -- tens of thousands of variably-sized alloc/free cycles that fragmented the allocator and drove RSS to ~140 GB on a shared machine, even though no single chunk's live memory was large. Replace it with a direct minor-axis kernel that walks the array's own storage as-is: each thread gets a private, full-output-sized accumulator and scatters into it while owning a disjoint range of major slices, summed across threads at the end -- the same pattern _ops.py already uses for minor_sums/minor_counts/minor_extrema. accumulator_threads caps the thread count to a fixed byte budget, so the accumulator block is always nthreads * output_dim * width * 8 bytes: bounded, independent of nnz, and allocated once per call instead of thousands of times. For a huge output axis (e.g. millions of cells) that caps down to a single thread, trading parallelism for a hard memory bound rather than the previous unbounded chunk churn. Removes _chunk_bounds/_aligned_source/_build_dual/_dual_arr and the chunked-transpose test suite; adds tests/test_minor_axis_matmul.py covering correctness against a dense reference, thread-count invariance, that no second copy of the array is ever built, and that peak memory stays bounded by the accumulator budget rather than nnz. Co-Authored-By: Claude Sonnet 5 --- src/vsparse/_norm_common.py | 12 +- src/vsparse/_vcs_matmul.py | 263 +++++++++++++++++++------------- src/vsparse/_vcs_norm.py | 22 +-- tests/test_chunked_transpose.py | 147 ------------------ tests/test_minor_axis_matmul.py | 125 +++++++++++++++ tests/test_vcs_norm.py | 19 ++- tests/test_vcs_norm_recipes.py | 2 +- 7 files changed, 300 insertions(+), 290 deletions(-) delete mode 100644 tests/test_chunked_transpose.py create mode 100644 tests/test_minor_axis_matmul.py diff --git a/src/vsparse/_norm_common.py b/src/vsparse/_norm_common.py index 9b6d406..1bb4835 100644 --- a/src/vsparse/_norm_common.py +++ b/src/vsparse/_norm_common.py @@ -138,12 +138,10 @@ def resolve_recipe(view: str | Recipe) -> Recipe: 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. + Holds each view's *statistics* strongly and the view itself only weakly, + so a cache hit never pins a view alive for longer than the caller's own + references would. 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 @@ -786,7 +784,7 @@ def _from_internal_stats( return self def _init_extra(self) -> None: - """Hook for subclasses with extra per-instance state (e.g. ``_dual_arr``). + """Hook for subclasses with extra per-instance state to initialize. ``__init__`` normally initializes that state itself; :meth:`from_stats` builds an instance via ``object.__new__`` instead, bypassing it, so it diff --git a/src/vsparse/_vcs_matmul.py b/src/vsparse/_vcs_matmul.py index a6f16b7..f00043b 100644 --- a/src/vsparse/_vcs_matmul.py +++ b/src/vsparse/_vcs_matmul.py @@ -18,17 +18,34 @@ ``self @ B``, :class:`~vsparse.VCSCArray` for ``B @ self`` -- see :func:`_vcsr_matmul_delta`/:func:`_vcsc_rmatmul_delta`). The other direction on a given array (:class:`~vsparse.VCSCArray` for ``self @ B``, -:class:`~vsparse.VCSRArray` for ``B @ self``) doesn't have that alignment -- -rather than run a scatter kernel there (thread-local output-shaped -accumulators, reduced across threads: cache-unfriendly, and memory-hungry -enough for wide ``B`` to need throttling), the storage is regrouped into the -other VCS format via :meth:`~vsparse._base._VCSBase._transpose_major`, a -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``/``col_post_scale`` are per-row/ -per-column statistics, so a chunk just takes the slice of them its own axis -covers. +:class:`~vsparse.VCSRArray` for ``B @ self``) doesn't have that alignment: +parallelizing naively over major slices there would have different threads +scatter-add into the *same* output row/column, a data race. + +That misaligned direction is handled the same way :mod:`vsparse._ops` +already handles the minor-axis reductions (``minor_sums``/``minor_counts``/ +``minor_extrema``): each thread gets its own private, full-output-sized +accumulator and walks a disjoint contiguous range of major slices, scattering +into that private copy; the per-thread copies are summed once at the end +(:func:`_vcsc_matmul_delta_minor`/:func:`_vcsr_rmatmul_delta_minor`). No +second copy of the sparse array's structure is ever built, and no chunk of it +is ever regrouped into the other VCS format -- the kernel walks the array's +own ``indices`` exactly as stored, so the cost of a call is ``O(nnz * width)`` +work plus one fixed-size accumulator allocation, never a sequence of +allocate/free cycles that scales with array size. + +The one thing this can't avoid is that the accumulator block itself costs +``nthreads * out_major_dim * width * 8`` bytes. +:func:`~vsparse._ops.accumulator_threads` (shared with the reduction +kernels) caps ``nthreads`` to keep that block under a fixed budget -- for a +huge output axis (e.g. millions of cells) that caps down to a single thread, +trading parallelism for a hard memory bound. That is the right trade at +that scale: a small, constant footprint and a correct answer, rather than +either an unbounded thread count or a chunked-transpose fallback whose total +allocation churn (thousands of variably-sized chunk buffers, over many +power iterations and many ranks/trials in a BiCV sweep) is what previously +fragmented the allocator and drove RSS up to ~140 GB on a shared machine +(see the project issue this replaced). """ from __future__ import annotations @@ -39,6 +56,7 @@ import numpy as np from vsparse._norm_common import _g +from vsparse._ops import accumulator_threads if TYPE_CHECKING: from vsparse._vcs_norm import _VCSNormalizedBase @@ -104,60 +122,97 @@ def _vcsc_rmatmul_delta( acc[c] += delta * brow[c] -# ``Delta @ B`` splits over the contracted axis and ``B @ Delta`` over rows, -# so a contiguous range of major slices can be regrouped on its own, -# accumulated into the shared output, and dropped. +# -- misaligned direction: per-thread private accumulators, no regrouping ---- -_CHUNK_BUDGET_BYTES = 128 << 20 # 128 MiB of transient regrouping per chunk - -# transpose_major sorts globally over the chunk's nonzeros, so its peak is -# several nnz-sized temporaries. Deliberately generous, since underestimating -# means the budget doesn't hold. -_TRANSPOSE_BYTES_PER_NNZ = 64 - - -def _chunk_bounds(arr, budget_bytes: int) -> list[tuple[int, int]]: - """Contiguous ``[start, stop)`` major-slice ranges, each within the byte budget.""" - n_major = arr.n_major - if n_major == 0: - return [] - max_nnz = max(1, budget_bytes // _TRANSPOSE_BYTES_PER_NNZ) - if arr.nnz <= max_nnz: - return [(0, n_major)] - - # nnz of major slices [0, j) -- value_ptr indexed by the group boundary. - cumulative = arr.value_ptr[arr.major_ptr] - bounds = [] - start = 0 - while start < n_major: - # Furthest stop whose chunk stays under budget; always advance by >= 1. - stop = int(np.searchsorted(cumulative, cumulative[start] + max_nnz, side="right")) - 1 - stop = min(max(stop, start + 1), n_major) - bounds.append((start, stop)) - start = stop - return bounds +@numba.njit(cache=True, parallel=True) +def _vcsc_matmul_delta_minor( + major_ptr, + values, + value_ptr, + indices, + row_scale, + gene_scale, + col_post_scale, + g_code, + B, + nthreads, + n_rows, +): + """``Delta @ B`` for a VCSC array (major=columns), walked directly (no VCSR regroup). -def _aligned_source(nview: _VCSNormalizedBase, needed_format: str): - """``(array, None)`` to run one major-aligned pass, or ``(None, chunk bounds)``.""" - arr = nview._arr - if arr._format == needed_format: - return arr, None - if nview._dual_arr is not None: - return nview._dual_arr, None - bounds = _chunk_bounds(arr, _CHUNK_BUDGET_BYTES) - if len(bounds) <= 1: - # Regrouping the whole array already fits the per-call budget, so - # keeping it costs no extra peak memory and saves every later call. - return _build_dual(nview), None - return None, bounds + ``B`` is ``(n_cols, k)``; returns ``(n_rows, k)``. Parallel-safe because + each thread scatters into its own private ``(n_rows, k)`` accumulator + while owning a disjoint, contiguous range of columns. + """ + n_major = major_ptr.shape[0] - 1 # n_cols + k = B.shape[1] + chunk = (n_major + nthreads - 1) // nthreads + partial = np.zeros((nthreads, n_rows, k), dtype=np.float64) + for t in numba.prange(nthreads): # ty: ignore[not-iterable] + start = t * chunk + stop = min(n_major, start + chunk) + local = partial[t] + for j in range(start, stop): + gs = gene_scale[j] + if gs == 0.0: + continue + s = col_post_scale[j] + brow = B[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 = s * _g(v / row_scale[row] / gs, g_code) + acc = local[row] + for c in range(k): + acc[c] += delta * brow[c] + return partial.sum(axis=0) -def _build_dual(nview: _VCSNormalizedBase): - """Build and cache the opposite-format copy of ``nview``'s array, once.""" - if nview._dual_arr is None: - nview._dual_arr = nview._arr._transpose_major() - return nview._dual_arr +@numba.njit(cache=True, parallel=True) +def _vcsr_rmatmul_delta_minor( + major_ptr, + values, + value_ptr, + indices, + row_scale, + gene_scale, + col_post_scale, + g_code, + Bt, + nthreads, + n_cols, +): + """``B @ Delta`` for a VCSR array (major=rows), walked directly (no VCSC regroup). + + ``Bt`` is ``(n_rows, p)`` (``B`` transposed); returns ``(n_cols, p)``. + Parallel-safe because each thread scatters into its own private + ``(n_cols, p)`` accumulator while owning a disjoint, contiguous range of + rows. + """ + n_major = major_ptr.shape[0] - 1 # n_rows + p = Bt.shape[1] + chunk = (n_major + nthreads - 1) // nthreads + partial = np.zeros((nthreads, n_cols, p), dtype=np.float64) + for t in numba.prange(nthreads): # ty: ignore[not-iterable] + start = t * chunk + stop = min(n_major, start + chunk) + local = partial[t] + for i in range(start, stop): + rs = row_scale[i] + brow = Bt[i] + for u in range(major_ptr[i], major_ptr[i + 1]): + v = values[u] + for kk in range(value_ptr[u], value_ptr[u + 1]): + col = indices[kk] + gs = gene_scale[col] + if gs > 0.0: + delta = col_post_scale[col] * _g(v / rs / gs, g_code) + acc = local[col] + for c in range(p): + acc[c] += delta * brow[c] + return partial.sum(axis=0) # -- public entry points: dense correction + sparse delta ------------------- @@ -182,14 +237,13 @@ def normalized_at_dense(nview: _VCSNormalizedBase, other: Any) -> np.ndarray: 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: + if arr._format == "csr": + # self @ B is major-aligned for VCSR. _vcsr_matmul_delta( - src.major_ptr, - src.values, - src.value_ptr, - src.indices, + arr.major_ptr, + arr.values, + arr.value_ptr, + arr.indices, nview.row_scale, nview.gene_scale, nview.col_post_scale, @@ -198,23 +252,21 @@ def normalized_at_dense(nview: _VCSNormalizedBase, other: Any) -> np.ndarray: out, ) else: - # Column chunk at a time: Delta @ B == sum over chunks of - # Delta[:, chunk] @ B[chunk, :], so each chunk's contribution - # accumulates into the same output and is then discarded. - 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], - nview.col_post_scale[start:stop], - g_code, - np.ascontiguousarray(B[start:stop]), - out, - ) + # arr is VCSC: self @ B is the misaligned direction for this format. + nthreads = accumulator_threads(arr.shape[0], bytes_per_element=8 * B.shape[1]) + out += _vcsc_matmul_delta_minor( + arr.major_ptr, + arr.values, + arr.value_ptr, + arr.indices, + nview.row_scale, + nview.gene_scale, + nview.col_post_scale, + g_code, + B, + nthreads, + arr.shape[0], + ) offset = nview.col_mean * nview.col_post_scale baseline = (-offset) @ B # (k,): every row's implicit-zero contribution @@ -238,14 +290,13 @@ def dense_at_normalized(nview: _VCSNormalizedBase, other: Any) -> np.ndarray: 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: + if arr._format == "csc": + # B @ self is major-aligned for VCSC. _vcsc_rmatmul_delta( - src.major_ptr, - src.values, - src.value_ptr, - src.indices, + arr.major_ptr, + arr.values, + arr.value_ptr, + arr.indices, nview.row_scale, nview.gene_scale, nview.col_post_scale, @@ -254,21 +305,21 @@ def dense_at_normalized(nview: _VCSNormalizedBase, other: Any) -> np.ndarray: 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, - nview.col_post_scale, - g_code, - np.ascontiguousarray(Bt[start:stop]), - out_t, - ) + # arr is VCSR: B @ self is the misaligned direction for this format. + nthreads = accumulator_threads(arr.shape[1], bytes_per_element=8 * p) + out_t += _vcsr_rmatmul_delta_minor( + arr.major_ptr, + arr.values, + arr.value_ptr, + arr.indices, + nview.row_scale, + nview.gene_scale, + nview.col_post_scale, + g_code, + Bt, + nthreads, + arr.shape[1], + ) out = np.ascontiguousarray(out_t.T) offset = nview.col_mean * nview.col_post_scale diff --git a/src/vsparse/_vcs_norm.py b/src/vsparse/_vcs_norm.py index bd9fdbf..eaab3cd 100644 --- a/src/vsparse/_vcs_norm.py +++ b/src/vsparse/_vcs_norm.py @@ -13,12 +13,9 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any +from typing import Any -from vsparse._norm_common import DEFAULT_RECIPE, NormalizedViewBase, Recipe - -if TYPE_CHECKING: - from vsparse._base import _VCSBase +from vsparse._norm_common import NormalizedViewBase __all__ = ["VCSCArrayNormalized", "VCSRArrayNormalized"] @@ -26,20 +23,7 @@ class _VCSNormalizedBase(NormalizedViewBase): """Shared implementation for :class:`VCSCArrayNormalized`/:class:`VCSRArrayNormalized`.""" - __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 = None + __slots__ = () def __matmul__(self, other: Any) -> Any: """``self @ other`` for a dense ``other`` -- see :mod:`vsparse._vcs_matmul`.""" diff --git a/tests/test_chunked_transpose.py b/tests/test_chunked_transpose.py deleted file mode 100644 index 09d7279..0000000 --- a/tests/test_chunked_transpose.py +++ /dev/null @@ -1,147 +0,0 @@ -from __future__ import annotations - -import tracemalloc -from itertools import pairwise - -import numpy as np -import pytest -import scipy.sparse as sp - -from vsparse import VCSCArray, VCSRArray -from vsparse._vcs_matmul import _TRANSPOSE_BYTES_PER_NNZ, _chunk_bounds - - -@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 _reference(dense: np.ndarray) -> np.ndarray: - row_totals = dense.sum(axis=1) - row_scale = row_totals / np.median(row_totals) - row_scale[row_scale == 0.0] = 1.0 - scaled = dense / row_scale[:, None] - gene_scale = scaled.sum(axis=0) - with np.errstate(divide="ignore", invalid="ignore"): - normalized = np.where(gene_scale > 0, scaled / gene_scale[None, :], 0.0) - transformed = np.log10(1.0 + 1000.0 * normalized) - return transformed - transformed.mean(axis=0, keepdims=True) - - -@pytest.mark.parametrize("budget", [1, 200, 1 << 30]) -def test_chunks_tile_the_major_axis(dense, vcls, budget): - """Chunks must cover every major slice exactly once, at any budget.""" - v = vcls.from_scipy(_scipy_for(vcls, dense)) - bounds = _chunk_bounds(v, budget) - - assert bounds[0][0] == 0 - assert bounds[-1][1] == v.n_major - assert all(start < stop for start, stop in bounds) - for (_, prev_stop), (start, _) in pairwise(bounds): - assert start == prev_stop - - -def test_chunks_respect_the_budget(vcls, rng): - """A chunk may only exceed the budget when it is a single, indivisible slice.""" - dense = rng.integers(0, 4, size=(60, 50)).astype(np.float64) - v = vcls.from_scipy(_scipy_for(vcls, dense)) - budget = 40 * _TRANSPOSE_BYTES_PER_NNZ - - cumulative = v.value_ptr[v.major_ptr] - for start, stop in _chunk_bounds(v, budget): - chunk_nnz = int(cumulative[stop] - cumulative[start]) - assert chunk_nnz <= 40 or stop - start == 1 - - -def test_empty_array_has_no_chunks(vcls): - shape = (0, 4) if vcls is VCSRArray else (4, 0) - v = vcls.from_scipy(_scipy_for(vcls, np.zeros(shape))) - assert _chunk_bounds(v, 1 << 20) == [] - - -def test_major_range_is_a_zero_copy_view(vcls, rng): - """Chunking is only affordable because the buffers are shared, not gathered.""" - dense = rng.integers(0, 4, size=(30, 20)).astype(np.float64) - v = vcls.from_scipy(_scipy_for(vcls, dense)) - - chunk = v._major_range(1, 4) - assert np.shares_memory(chunk.values, v.values) - assert chunk.nnz == 0 or np.shares_memory(chunk.indices, v.indices) - - -def test_major_range_matches_ordinary_slicing(vcls, rng): - """A chunk must hold the same sub-array a caller would get by slicing.""" - dense = rng.integers(0, 4, size=(30, 20)).astype(np.float64) - v = vcls.from_scipy(_scipy_for(vcls, dense)) - start, stop = 2, 7 - - native = v[:, start:stop] if vcls is VCSCArray else v[start:stop, :] - np.testing.assert_allclose(v._major_range(start, stop).toarray(), native.toarray()) - - -@pytest.mark.parametrize("budget", [1, 200, 4000, 1 << 30]) -def test_matmul_is_chunk_size_invariant(monkeypatch, vcls, rng, budget): - """Same answer whether it runs in one chunk or one major slice at a time.""" - import vsparse._vcs_matmul as vcs_matmul - - dense = rng.integers(0, 5, size=(40, 24)).astype(np.float64) - v = vcls.from_scipy(_scipy_for(vcls, dense)) - ref = _reference(dense) - B = rng.normal(size=(dense.shape[1], 3)) - Bl = rng.normal(size=(2, dense.shape[0])) - - monkeypatch.setattr(vcs_matmul, "_CHUNK_BUDGET_BYTES", budget) - nv = v.normalized() - - np.testing.assert_allclose(nv @ B, ref @ B, atol=1e-7) - np.testing.assert_allclose(Bl @ nv, Bl @ ref, atol=1e-7) - np.testing.assert_allclose(nv @ B[:, 0], ref @ B[:, 0], atol=1e-7) - np.testing.assert_allclose(Bl[0] @ nv, Bl[0] @ ref, atol=1e-7) - - -def test_multi_chunk_matmul_caches_nothing(monkeypatch, vcls, rng): - """Past one chunk no full dual is built, which is the memory guarantee.""" - import vsparse._vcs_matmul as vcs_matmul - - dense = rng.integers(1, 5, size=(40, 24)).astype(np.float64) - v = vcls.from_scipy(_scipy_for(vcls, dense)) - monkeypatch.setattr(vcs_matmul, "_CHUNK_BUDGET_BYTES", 200) - assert len(_chunk_bounds(v, 200)) > 1 - - nv = v.normalized() - nv @ rng.normal(size=(dense.shape[1], 2)) - rng.normal(size=(2, dense.shape[0])) @ nv - - assert nv._dual_arr is None - - -def test_misaligned_matmul_peak_is_bounded_by_the_chunk_budget(monkeypatch, rng): - """Peak memory has to track the budget, not the size of the array.""" - import vsparse._vcs_matmul as vcs_matmul - - dense = rng.integers(1, 5, size=(1200, 400)).astype(np.float64) - v = VCSCArray.from_scipy(sp.csc_array(dense)) - nnz_bytes = v.nnz * v.indices.dtype.itemsize - B = rng.normal(size=(dense.shape[1], 2)) - - v.normalized() @ B # warm up the JIT before measuring - - monkeypatch.setattr(vcs_matmul, "_CHUNK_BUDGET_BYTES", 64 * _TRANSPOSE_BYTES_PER_NNZ) - nv = v.normalized() - - tracemalloc.start() - try: - before = tracemalloc.get_traced_memory()[0] - tracemalloc.reset_peak() - out = nv @ B - peak = tracemalloc.get_traced_memory()[1] - finally: - tracemalloc.stop() - - assert nv._dual_arr is None - assert peak - before < nnz_bytes - np.testing.assert_allclose(out, _reference(dense) @ B, atol=1e-7) diff --git a/tests/test_minor_axis_matmul.py b/tests/test_minor_axis_matmul.py new file mode 100644 index 0000000..b558cb4 --- /dev/null +++ b/tests/test_minor_axis_matmul.py @@ -0,0 +1,125 @@ +"""Misaligned-direction normalized-view matmul: native minor-axis kernel. + +Replaces the old chunked-transpose/dual-copy approach (see +``vsparse._vcs_matmul``): the misaligned direction now walks the array's own +storage directly with per-thread private accumulators, never building a +second copy of the sparse structure or regrouping any part of it into the +other VCS format. +""" + +from __future__ import annotations + +import tracemalloc + +import numpy as np +import pytest +import scipy.sparse as sp + +from vsparse import VCSCArray, VCSRArray +from vsparse._ops import accumulator_threads + + +@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 _reference(dense: np.ndarray) -> np.ndarray: + row_totals = dense.sum(axis=1) + row_scale = row_totals / np.median(row_totals) + row_scale[row_scale == 0.0] = 1.0 + scaled = dense / row_scale[:, None] + gene_scale = scaled.sum(axis=0) + with np.errstate(divide="ignore", invalid="ignore"): + normalized = np.where(gene_scale > 0, scaled / gene_scale[None, :], 0.0) + transformed = np.log10(1.0 + 1000.0 * normalized) + return transformed - transformed.mean(axis=0, keepdims=True) + + +def test_misaligned_matmul_matches_reference(dense, vcls): + """Both matmul directions match a dense reference, whichever direction is misaligned.""" + 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() + ref = _reference(dense) + + rng = np.random.default_rng(7) + B = rng.normal(size=(dense.shape[1], 3)) + Bl = rng.normal(size=(2, dense.shape[0])) + + np.testing.assert_allclose(nv @ B, ref @ B, atol=1e-7) + np.testing.assert_allclose(Bl @ nv, Bl @ ref, atol=1e-7) + np.testing.assert_allclose(nv @ B[:, 0], ref @ B[:, 0], atol=1e-7) + np.testing.assert_allclose(Bl[0] @ nv, Bl[0] @ ref, atol=1e-7) + + +def test_matmul_result_is_accumulator_thread_count_invariant(monkeypatch, vcls, rng): + """Same answer regardless of how many private accumulators the kernel splits work across.""" + import vsparse._vcs_matmul as vcs_matmul + + dense = rng.integers(0, 5, size=(40, 24)).astype(np.float64) + v = vcls.from_scipy(_scipy_for(vcls, dense)) + ref = _reference(dense) + B = rng.normal(size=(dense.shape[1], 3)) + Bl = rng.normal(size=(2, dense.shape[0])) + nv = v.normalized() + + for forced_threads in (1, 2, 8): + monkeypatch.setattr( + vcs_matmul, "accumulator_threads", lambda *a, _n=forced_threads, **k: _n + ) + np.testing.assert_allclose(nv @ B, ref @ B, atol=1e-7) + np.testing.assert_allclose(Bl @ nv, Bl @ ref, atol=1e-7) + + +def test_no_second_copy_of_the_array_is_built(vcls, rng): + """The normalized view carries nothing beyond the original array and O(n_rows+n_cols) stats.""" + dense = rng.integers(1, 5, size=(30, 20)).astype(np.float64) + v = vcls.from_scipy(_scipy_for(vcls, dense)) + nv = v.normalized() + + B = rng.normal(size=(dense.shape[1], 2)) + Bl = rng.normal(size=(2, dense.shape[0])) + nv @ B + Bl @ nv + + assert not hasattr(nv, "_dual_arr") + # The wrapped array's own buffers are untouched/unreplaced by any matmul call. + assert nv._arr is v + + +def test_misaligned_matmul_peak_is_bounded_by_the_accumulator_budget(rng): + """Peak memory tracks the (fixed) accumulator budget, not the size of the array.""" + dense = rng.integers(1, 5, size=(1200, 400)).astype(np.float64) + v = VCSCArray.from_scipy(sp.csc_array(dense)) + nnz_bytes = v.nnz * v.indices.dtype.itemsize + B = rng.normal(size=(dense.shape[1], 2)) + + v.normalized() @ B # warm up the JIT before measuring + + nv = v.normalized() + tracemalloc.start() + try: + before = tracemalloc.get_traced_memory()[0] + tracemalloc.reset_peak() + out = nv @ B + peak = tracemalloc.get_traced_memory()[1] + finally: + tracemalloc.stop() + + # Accumulator block is nthreads * n_rows * width * 8 bytes -- independent + # of nnz, so it stays far below a per-nonzero cost for this shape. + assert peak - before < nnz_bytes + np.testing.assert_allclose(out, _reference(dense) @ B, atol=1e-7) + + +def test_accumulator_threads_degrades_to_one_for_a_huge_output_axis(): + """A huge output axis (e.g. millions of cells) must not blow the accumulator budget.""" + huge_axis = 2_000_000 + wide_b = 200 # e.g. rank + oversampling in a randomized SVD + assert accumulator_threads(huge_axis, bytes_per_element=8 * wide_b) == 1 diff --git a/tests/test_vcs_norm.py b/tests/test_vcs_norm.py index 09f60d1..16858c3 100644 --- a/tests/test_vcs_norm.py +++ b/tests/test_vcs_norm.py @@ -121,13 +121,17 @@ def test_all_zero_matrix_does_not_crash(vcls): np.testing.assert_allclose(nv.toarray(), np.zeros((5, 4))) -def test_small_array_dual_is_cached_lazily(dense, vcls): - """An array whose whole regrouping fits one chunk is transposed once and kept.""" +def test_matmul_never_builds_a_second_copy_of_the_array(dense, vcls): + """Neither matmul direction builds an opposite-format copy of the array. + + See ``tests/test_minor_axis_matmul.py`` for the misaligned direction's own + correctness/memory-bound coverage; this just confirms the view stays thin. + """ 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() - assert nv._dual_arr is None # not built at construction + assert nv._arr is v rng = np.random.default_rng(12) B = rng.normal(size=(dense.shape[1], 3)) @@ -135,13 +139,8 @@ def test_small_array_dual_is_cached_lazily(dense, vcls): nv @ B # major-aligned for VCSR self@B; misaligned for VCSC Bl @ nv # major-aligned for VCSC B@self; misaligned for VCSR - dual_after_matmul = nv._dual_arr - assert dual_after_matmul is not None - assert dual_after_matmul._format != v._format - - nv @ B - Bl @ nv - assert nv._dual_arr is dual_after_matmul # reused, not rebuilt + assert nv._arr is v # still the original array, not replaced by a dual + assert not hasattr(nv, "_dual_arr") def test_transpose_major_roundtrip(dense, vcls): diff --git a/tests/test_vcs_norm_recipes.py b/tests/test_vcs_norm_recipes.py index 4d5a7e0..e9a4a3e 100644 --- a/tests/test_vcs_norm_recipes.py +++ b/tests/test_vcs_norm_recipes.py @@ -296,7 +296,7 @@ def test_recipe_with_an_unknown_g_code_is_rejected(vcls, dense): 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.""" + """The cache holds views weakly, so a dropped view is actually collected.""" if dense.sum() == 0: pytest.skip("all-zero matrix: median row total is 0") v = vcls.from_scipy(_scipy_for(vcls, dense)) From 0beb1af8dfd155a2a21fa61feb824b1c11be65a3 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Sun, 13 Sep 2026 12:11:37 -0700 Subject: [PATCH 6/6] Add select(recalculate=False): a lazy, stats-preserving row/col selection select() previously always recomputed statistics fresh for the selected sub-array. A caller that instead wants to evaluate a fixed (e.g. train-derived) normalization against a different slice -- such as bi-cross-validation scoring a held-out test block against train-derived statistics -- had no lazy option: bracket indexing (__getitem__) keeps the parent's statistics but is documented to always eagerly materialize the selection as a dense ndarray, which is fine for a small window but not for a selection covering a large fraction of a huge array. select(rows, cols, recalculate=False) reuses this view's existing a/b/c/s (row_scale sliced by rows -- exact, since it's a per-cell quantity; the per-gene stats sliced by cols, unchanged by which rows are selected) via the existing from_stats() classmethod, and returns a real view rather than an array -- so it composes with `@`/toarray() and stays lazy regardless of selection size, unlike __getitem__. Co-Authored-By: Claude Sonnet 5 --- src/vsparse/_norm_common.py | 46 ++++++++++++++++++++++++++----- tests/test_norm_selection.py | 52 ++++++++++++++++++++++++++++++++++++ 2 files changed, 91 insertions(+), 7 deletions(-) diff --git a/src/vsparse/_norm_common.py b/src/vsparse/_norm_common.py index 1bb4835..5a10f9d 100644 --- a/src/vsparse/_norm_common.py +++ b/src/vsparse/_norm_common.py @@ -983,18 +983,50 @@ def slice_norms(self, condition_idxs: Any, n_cond: int) -> np.ndarray: # -- selection --------------------------------------------------------------- - def select(self, rows: Any = slice(None), cols: Any = slice(None)) -> Any: - """A normalized view of the selected sub-array, with statistics recomputed for it. + def select( + self, rows: Any = slice(None), cols: Any = slice(None), *, recalculate: bool = True + ) -> Any: + """A normalized view of the selected sub-array. Returns a view, not a dense array, so it still composes with - ``@``/:meth:`toarray`. + ``@``/:meth:`toarray` without ever materializing the selection. + + Parameters + ---------- + recalculate + If ``True`` (the default), statistics are recomputed fresh from + the selected sub-array -- a column selection then re-derives + read depth from only the selected columns, which is rarely what + a caller wants (select genes first and normalize after if it + matters). If ``False``, this view's *existing* statistics are + reused instead: ``row_scale``/``a`` sliced by ``rows`` (a + per-cell quantity, so subsetting is exact, not an + approximation), and the per-gene statistics (``gene_scale``/ + ``col_mean``/``col_post_scale``, i.e. ``b``/``c``/``s``) sliced + by ``cols`` unchanged. This is the right choice for evaluating a + model fit on a held-out slice against statistics computed on + the whole (or a different) array -- e.g. bi-cross-validation + scoring a test block against train-derived statistics -- and, + unlike bracket indexing (:meth:`__getitem__`), stays a lazy view + no matter how large the selection is: nothing gets materialized + until (and unless) the caller calls :meth:`toarray` or `@`s it, + and even then the O(nnz) work is bounded, not an eager dense + allocation sized to the selection. """ - # A column selection re-derives read depth from only the selected - # columns, which is rarely what a caller wants; select genes first - # and normalize after if it matters. - sub: Any = self._arr[_prep_key(rows), _prep_key(cols)] + row_key, col_key = _prep_key(rows), _prep_key(cols) + sub: Any = self._arr[row_key, col_key] if not isinstance(sub, type(self._arr)): sub = type(self._arr).from_scipy(sub) + if not recalculate: + return type(self).from_stats( + sub, + self.recipe, + np.asarray(self.a)[row_key], + np.asarray(self.b)[col_key], + np.asarray(self.c)[col_key], + np.asarray(self.s)[col_key], + stale=self.stale, + ) return type(self)(sub, self.recipe) # -- on-the-fly elementwise access ------------------------------------------ diff --git a/tests/test_norm_selection.py b/tests/test_norm_selection.py index bc4921f..48c9cd1 100644 --- a/tests/test_norm_selection.py +++ b/tests/test_norm_selection.py @@ -102,6 +102,58 @@ def test_select_returns_a_view_that_still_composes(vcls): np.testing.assert_allclose(sub @ B, _reference(dense[mask]) @ B, atol=1e-8) +# -- select(recalculate=False): a window that stays a lazy view ------------- + + +def test_select_no_recalculate_matches_getitem_values(vcls): + """select(recalculate=False) keeps the parent's statistics, like __getitem__.""" + dense, mask = _mixed_population() + nv = vcls.from_scipy(_scipy_for(vcls, dense)).normalized() + + want = np.asarray(nv[mask, :]) + got = nv.select(mask, recalculate=False).toarray() + + np.testing.assert_allclose(got, want, atol=1e-10) + + +def test_select_no_recalculate_returns_a_lazy_view(vcls): + """Unlike __getitem__, this never eagerly materializes the selection.""" + dense, mask = _mixed_population() + nv = vcls.from_scipy(_scipy_for(vcls, dense)).normalized() + + sub = nv.select(mask, recalculate=False) + assert isinstance(sub, _norm_cls(vcls)) + assert sub.shape == (int(mask.sum()), dense.shape[1]) + + rng = np.random.default_rng(5) + B = rng.normal(size=(dense.shape[1], 4)) + Bl = rng.normal(size=(3, int(mask.sum()))) + want = np.asarray(nv[mask, :]) + np.testing.assert_allclose(sub @ B, want @ B, atol=1e-8) + np.testing.assert_allclose(Bl @ sub, Bl @ want, atol=1e-8) + + +def test_select_no_recalculate_composes_with_column_selection_too(vcls): + """A column selection just slices the existing per-gene stats, not recomputed.""" + dense, mask = _mixed_population() + nv = vcls.from_scipy(_scipy_for(vcls, dense)).normalized() + cols = np.array([0, 3, 7, 15, 40]) + + want = np.asarray(nv[mask, :])[:, cols] + got = nv.select(mask, cols, recalculate=False).toarray() + np.testing.assert_allclose(got, want, atol=1e-10) + + +def test_select_no_recalculate_default_recalculate_still_recomputes(vcls): + """Bare select(mask) is unaffected by the new keyword-only argument.""" + dense, mask = _mixed_population() + nv = vcls.from_scipy(_scipy_for(vcls, dense)).normalized() + + np.testing.assert_allclose( + nv.select(mask).toarray(), nv.select(mask, recalculate=True).toarray(), atol=1e-12 + ) + + # -- __getitem__: a window that keeps the parent's statistics ----------------