From e6ac786718bd225d30623aa83394a3afd09c4705 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Sat, 12 Sep 2026 19:05:12 -0700 Subject: [PATCH] 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)