diff --git a/src/vsparse/_norm_common.py b/src/vsparse/_norm_common.py index 842611e..e02eac3 100644 --- a/src/vsparse/_norm_common.py +++ b/src/vsparse/_norm_common.py @@ -9,6 +9,10 @@ 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 -- @@ -252,16 +256,29 @@ def toarray(self) -> np.ndarray: ) return out + # -- 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. + + Returns a view, not a dense array, so it still composes with + ``@``/:meth:`toarray`. + """ + # 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)] + if not isinstance(sub, type(self._arr)): + sub = type(self._arr).from_scipy(sub) + return type(self)(sub) + # -- on-the-fly elementwise access ------------------------------------------ def __getitem__(self, key: Any) -> np.ndarray: - """Compute just the requested sub-block, on the fly, from the raw data. + """``toarray()[key]``, computed on the fly without materializing the full matrix. - Only the raw counts for the requested rows/columns are ever - decompressed (via the underlying array's own indexing, which stays - compact for a major-axis-only slice); the transform/centering - formula is then applied to that small block directly, using the - precomputed per-row/per-column statistics -- never the full matrix. + Applies this view's statistics, so it is not the same as normalizing + the selected sub-matrix; use :meth:`select` for that. """ if isinstance(key, tuple): if len(key) != 2: diff --git a/tests/test_norm_selection.py b/tests/test_norm_selection.py new file mode 100644 index 0000000..bd3dbb3 --- /dev/null +++ b/tests/test_norm_selection.py @@ -0,0 +1,144 @@ +from __future__ import annotations + +import numpy as np +import pytest +import scipy.sparse as sp + +from vsparse import 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) -> np.ndarray: + """Read-depth normalize, log-transform, and mean-center a dense matrix directly.""" + 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 _mixed_population(seed: int = 0) -> tuple[np.ndarray, np.ndarray]: + """Two cell types differing in both sequencing depth and marker genes.""" + rng = np.random.default_rng(seed) + n_cells, n_genes = 240, 80 + dense = np.empty((n_cells, n_genes)) + dense[: n_cells // 2] = rng.poisson(0.6, size=(n_cells // 2, n_genes)) + dense[n_cells // 2 :] = rng.poisson(4.0, size=(n_cells // 2, n_genes)) + dense[n_cells // 2 :, :16] *= 6 # markers expressed only in the second type + + mask = np.zeros(n_cells, dtype=bool) + mask[n_cells // 2 :] = True + return dense, mask + + +def _relative_frobenius(got: np.ndarray, want: np.ndarray) -> float: + return float(np.linalg.norm(got - want) / np.linalg.norm(want)) + + +# -- select(): statistics recomputed for the selection ----------------------- + + +def test_select_matches_normalizing_the_selection_directly(vcls): + """select() gives the same answer as normalizing those cells on their own.""" + dense, mask = _mixed_population() + nv = vcls.from_scipy(_scipy_for(vcls, dense)).normalized() + + got = nv.select(mask).toarray() + want = _reference(dense[mask]) + + assert _relative_frobenius(got, want) < 1e-12 + np.testing.assert_allclose(got, want, atol=1e-10) + + +def test_select_equals_selecting_on_the_raw_array_first(vcls): + """select() matches selecting on the raw array and normalizing after.""" + dense, mask = _mixed_population() + arr = vcls.from_scipy(_scipy_for(vcls, dense)) + + raw_sub = arr[mask, :] + assert isinstance(raw_sub, vcls) + + np.testing.assert_allclose( + arr.normalized().select(mask).toarray(), + raw_sub.normalized().toarray(), + atol=1e-12, + ) + + +def test_select_returns_a_view_that_still_composes(vcls): + """The result is a view, so it still composes with @ and toarray().""" + dense, mask = _mixed_population() + nv = vcls.from_scipy(_scipy_for(vcls, dense)).normalized() + + sub = nv.select(mask) + assert isinstance(sub, _norm_cls(vcls)) + assert sub.shape == (int(mask.sum()), dense.shape[1]) + + rng = np.random.default_rng(3) + B = rng.normal(size=(dense.shape[1], 4)) + np.testing.assert_allclose(sub @ B, _reference(dense[mask]) @ B, atol=1e-8) + + +def test_select_columns_and_both_axes(vcls): + """Two index arrays select a sub-block, not a pointwise diagonal.""" + dense, _ = _mixed_population() + nv = vcls.from_scipy(_scipy_for(vcls, dense)).normalized() + rows = np.arange(0, dense.shape[0], 7) + cols = np.arange(0, dense.shape[1], 5) + assert rows.shape != cols.shape + + np.testing.assert_allclose( + nv.select(cols=cols).toarray(), _reference(dense[:, cols]), atol=1e-10 + ) + np.testing.assert_allclose( + nv.select(rows, cols).toarray(), _reference(dense[np.ix_(rows, cols)]), atol=1e-10 + ) + + +def test_select_everything_is_the_whole_view(vcls, dense): + if dense.sum() == 0: + pytest.skip("all-zero matrix: median row total is 0") + nv = vcls.from_scipy(_scipy_for(vcls, dense)).normalized() + np.testing.assert_allclose(nv.select().toarray(), nv.toarray(), atol=1e-12) + + +# -- __getitem__: a window that keeps the parent's statistics ---------------- + + +def test_getitem_is_a_window_into_the_parent_matrix(vcls): + """Indexing gives the same values as the fully materialized view.""" + dense, mask = _mixed_population() + nv = vcls.from_scipy(_scipy_for(vcls, dense)).normalized() + + full = nv.toarray() + np.testing.assert_allclose(nv[mask, :], full[mask, :], atol=1e-10) + np.testing.assert_allclose(nv[0:4, 0:4], full[0:4, 0:4], atol=1e-10) + + +def test_getitem_and_select_disagree_substantially_on_a_real_selection(vcls): + """Windowing and renormalizing diverge by tens of percent on a real selection.""" + dense, mask = _mixed_population() + nv = vcls.from_scipy(_scipy_for(vcls, dense)).normalized() + + want = _reference(dense[mask]) + windowed = np.asarray(nv[mask, :]) + renormalized = nv.select(mask).toarray() + + assert _relative_frobenius(renormalized, want) < 1e-12 + assert _relative_frobenius(windowed, want) > 0.1