diff --git a/pyproject.toml b/pyproject.toml index 0852ae0c..f7089d4c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ dependencies = [ "anndata>=0.13", "pacmap>=0.9", "tqdm>=4.66.1", - "vsparse>=0.2.0", + "vsparse>=0.4.0", "hdf5plugin>=7.0.0", "matplotlib>=3.8", "seaborn>=0.13.2", diff --git a/scrise/rank_selection.py b/scrise/rank_selection.py index feb9a4cf..a53b5a2a 100644 --- a/scrise/rank_selection.py +++ b/scrise/rank_selection.py @@ -105,6 +105,31 @@ def _holdout_scale( # temporaries to a few hundred MB regardless of how large the matrix is. _MOMENT_CHUNK_NNZ = 50_000_000 +# Row-chunk size (in bytes of the dense block materialized per chunk) when +# streaming a duck-typed backend (e.g. a vsparse normalized view) that has no +# CSR internals to stream directly, but does support a lazy, stats-preserving +# `select()`. Keeps a chunk's dense materialization bounded regardless of how +# large the selected block is. +_MOMENT_CHUNK_BUDGET_BYTES = 64 << 20 + + +def _restrict_rows(X_mat: Any, mask: np.ndarray) -> Any: + """Row-restrict ``X_mat`` by a boolean mask, staying lazy where possible. + + For a duck-typed backend that supports ``select()`` (e.g. a vsparse + normalized view), keeps the view's *existing* statistics fixed rather + than renormalizing the selected rows on their own -- correct here, since + the caller is evaluating a model fit against a slice, not treating that + slice as its own dataset -- and, unlike bracket indexing on such a view, + never eagerly materializes the (potentially huge) selection as a dense + array. Plain dense/scipy-sparse ``X_mat`` falls back to ordinary + indexing, which is already cheap for those. + """ + select = getattr(X_mat, "select", None) + if select is not None: + return select(mask, recalculate=False) + return X_mat[mask] + def _test_block_moments( X_mat: Any, cell_mask: np.ndarray, gene_idx: np.ndarray @@ -116,8 +141,8 @@ def _test_block_moments( sums = np.zeros(n_genes) squares = np.zeros(n_genes) - if not sps.issparse(X_mat): - block = np.asarray(X_mat)[cell_mask] + if isinstance(X_mat, np.ndarray): + block = X_mat[cell_mask] # Both moments accumulate in float64. A float32 block summed in its own # dtype carries ~1e-4 relative error over a few hundred thousand rows, # which would reach the reported R2X through `ss_tot`. @@ -126,6 +151,25 @@ def _test_block_moments( np.sum(block.astype(np.float64) ** 2, axis=0)[gene_idx], ) + if not sps.issparse(X_mat): + # A duck-typed backend (e.g. a vsparse normalized view): select the + # cell subset lazily (keeping the view's existing statistics fixed, + # rather than renormalizing this subset on its own -- see + # `_restrict_rows`) and stream it in row chunks, rather than ever + # materializing the whole (potentially huge) subset as one dense + # block. + sub = X_mat.select(cell_mask, recalculate=False) + n_sub_rows = sub.shape[0] + chunk_rows = max(1, _MOMENT_CHUNK_BUDGET_BYTES // (n_genes * 8)) + start = 0 + while start < n_sub_rows: + stop = min(n_sub_rows, start + chunk_rows) + block = np.asarray(sub[start:stop], dtype=np.float64) + sums += block.sum(axis=0) + squares += np.sum(block**2, axis=0) + start = stop + return sums[gene_idx], squares[gene_idx] + mat = X_mat.tocsr() if X_mat.format != "csr" else X_mat n_rows = mat.shape[0] start = 0 @@ -229,7 +273,8 @@ def _bicv_trial( cond_train = cond_idx[train_cell_mask] Z = _cell_loadings(P_train, B, A, cond_train, n_cond) ZtY = np.asarray( - rmatmul(np.ascontiguousarray(Z.T), X_mat[train_cell_mask]), dtype=np.float64 + rmatmul(np.ascontiguousarray(Z.T), _restrict_rows(X_mat, train_cell_mask)), + dtype=np.float64, )[:, test_gene_idx] ZtY -= np.outer(Z.sum(axis=0), means_test_genes) # `lstsq` on the rank x rank system keeps the minimum-norm @@ -240,7 +285,8 @@ def _bicv_trial( C_full = np.zeros((n_genes, C.shape[1])) C_full[train_gene_mask] = C cond_test = cond_idx[test_cell_mask] - W_test = calc_W(X_mat[test_cell_mask], means, C_full) + X_test = _restrict_rows(X_mat, test_cell_mask) + W_test = calc_W(X_test, means, C_full) cond_slices_test = condition_slices(cond_test, n_cond) P_test, _ = project_data(W_test, [A, B, C], cond_slices_test) @@ -248,9 +294,9 @@ def _bicv_trial( # held-out loadings need rescaling for the held-out slice's size. L = _cell_loadings(P_test, B, A, cond_test, n_cond) L *= _holdout_scale(cond_train, cond_test, n_cond) - LtY = np.asarray( - rmatmul(np.ascontiguousarray(L.T), X_mat[test_cell_mask]), dtype=np.float64 - )[:, test_gene_idx] + LtY = np.asarray(rmatmul(np.ascontiguousarray(L.T), X_test), dtype=np.float64)[ + :, test_gene_idx + ] LtY -= np.outer(L.sum(axis=0), means_test_genes) col_sums, col_squares = _test_block_moments(X_mat, test_cell_mask, test_gene_idx) diff --git a/scrise/tests/test_rank_selection.py b/scrise/tests/test_rank_selection.py index 3155ea59..16da0e25 100644 --- a/scrise/tests/test_rank_selection.py +++ b/scrise/tests/test_rank_selection.py @@ -507,6 +507,139 @@ def test_block_moments_stream_across_more_than_one_row_block(monkeypatch): np.testing.assert_allclose(squares, np.sum(block**2, axis=0), rtol=1e-12) +# -- duck-typed backend: a vsparse normalized view --------------------------- +# +# `X.X` for a real vsparse-backed dataset (e.g. BAL-Pf2's lazy-normalized-view +# AnnData) is neither a plain ndarray nor scipy-sparse -- it's a +# VCSRArrayNormalized/VCSCArrayNormalized. Bracket indexing it with a large +# boolean mask (the old `X_mat[cell_mask]`) eagerly materializes the whole +# selection as dense, which is fine at test scale but would be tens of GB at +# BAL-Pf2's real scale. `_restrict_rows`/`_test_block_moments`'s duck-typed +# branch exist to avoid that; these tests exercise them directly. + + +def _normalized_view(rng, shape=(120, 30), vsparse_cls=None): + import vsparse + + vsparse_cls = vsparse_cls or vsparse.VCSRArray + dense = rng.random(shape) + dense[dense < 0.6] = 0.0 + mat = ( + sps.csr_array(dense) + if vsparse_cls is vsparse.VCSRArray + else sps.csc_array(dense) + ) + return dense, vsparse_cls.from_scipy(mat).normalized() + + +@pytest.mark.parametrize("fmt", ["csr", "csc"]) +def test_restrict_rows_matches_bracket_indexing_for_a_normalized_view(fmt): + import vsparse + + rng = np.random.default_rng(2) + vsparse_cls = vsparse.VCSRArray if fmt == "csr" else vsparse.VCSCArray + _dense, nv = _normalized_view(rng, vsparse_cls=vsparse_cls) + mask = rng.random(nv.shape[0]) < 0.6 + + from ..rank_selection import _restrict_rows + + restricted = _restrict_rows(nv, mask) + assert isinstance(restricted, type(nv)) # stayed a lazy view, not densified + np.testing.assert_allclose( + np.asarray(restricted.toarray()), np.asarray(nv[mask, :]) + ) + + +def test_restrict_rows_is_a_no_op_passthrough_for_plain_arrays(): + from ..rank_selection import _restrict_rows + + rng = np.random.default_rng(3) + dense = rng.random((20, 5)) + mask = rng.random(20) < 0.5 + + np.testing.assert_allclose(_restrict_rows(dense, mask), dense[mask]) + np.testing.assert_allclose( + _restrict_rows(sps.csr_array(dense), mask).toarray(), dense[mask] + ) + + +def test_block_moments_match_a_normalized_view_reference(): + """`_test_block_moments`'s duck-typed branch matches a dense reference.""" + import vsparse + + rng = np.random.default_rng(4) + dense, nv = _normalized_view(rng, vsparse_cls=vsparse.VCSRArray) + mask = rng.random(dense.shape[0]) < 0.7 + gene_idx = np.sort(rng.choice(dense.shape[1], size=11, replace=False)) + + sums, squares = _test_block_moments(nv, mask, gene_idx) + block = np.asarray(nv[mask, :])[:, gene_idx] + np.testing.assert_allclose(sums, block.sum(axis=0), rtol=1e-10) + np.testing.assert_allclose(squares, np.sum(block**2, axis=0), rtol=1e-10) + + +def test_block_moments_stream_a_normalized_view_across_more_than_one_chunk(monkeypatch): + """Force several row chunks so the normalized-view streaming path is exercised.""" + import vsparse + + import scrise.rank_selection as rs + + monkeypatch.setattr(rs, "_MOMENT_CHUNK_BUDGET_BYTES", 64) # forces 1-row chunks + rng = np.random.default_rng(5) + dense, nv = _normalized_view(rng, shape=(50, 12), vsparse_cls=vsparse.VCSRArray) + mask = rng.random(dense.shape[0]) < 0.8 + gene_idx = np.arange(12) + + sums, squares = rs._test_block_moments(nv, mask, gene_idx) + block = np.asarray(nv[mask, :]) + np.testing.assert_allclose(sums, block.sum(axis=0), rtol=1e-10) + np.testing.assert_allclose(squares, np.sum(block**2, axis=0), rtol=1e-10) + + +def test_block_moments_streaming_peak_is_bounded_by_the_chunk_budget_not_the_block(): + """The row-chunk streaming loop's own peak tracks the chunk size, not the + selected block's total size. + + ``select(recalculate=False)`` on a boolean mask does its own one-time + ``O(nnz)`` structural rebuild (the same cost bracket indexing or any + other full-array boolean selection pays in vsparse today) -- that part + isn't what this streaming code is trying to bound, so it's built once, + outside the measurement, isolating the chunk loop itself. + """ + import tracemalloc + + import vsparse + + rng = np.random.default_rng(6) + dense, nv = _normalized_view(rng, shape=(4000, 400), vsparse_cls=vsparse.VCSRArray) + mask = np.ones(dense.shape[0], dtype=bool) # select everything + sub = nv.select(mask, recalculate=False) + n_rows = sub.shape[0] + + def _stream(chunk_rows: int) -> int: + tracemalloc.start() + try: + tracemalloc.reset_peak() + start = 0 + while start < n_rows: + stop = min(n_rows, start + chunk_rows) + block = np.asarray(sub[start:stop], dtype=np.float64) + block.sum(axis=0) + start = stop + return tracemalloc.get_traced_memory()[1] + finally: + tracemalloc.stop() + + _stream(50) # warm up any lazy imports/JIT before measuring + small_chunk_peak = _stream(50) + large_chunk_peak = _stream(n_rows) # the whole block in one "chunk" + + # A peak that scaled with the full block regardless of chunk size would + # make these roughly equal; bounded streaming keeps the small-chunk peak + # well under the whole-block one. + assert small_chunk_peak < large_chunk_peak / 4 + + # Held-out slice scaling diff --git a/uv.lock b/uv.lock index bdfe0ae6..dd9b6e6f 100644 --- a/uv.lock +++ b/uv.lock @@ -2177,7 +2177,7 @@ requires-dist = [ { name = "tensorly", specifier = ">=0.9.0" }, { name = "tlviz", specifier = ">=0.1.1" }, { name = "tqdm", specifier = ">=4.66.1" }, - { name = "vsparse", specifier = ">=0.2.0" }, + { name = "vsparse", specifier = ">=0.4.0" }, ] provides-extras = ["gpu"] @@ -2471,7 +2471,7 @@ wheels = [ [[package]] name = "vsparse" -version = "0.2.0" +version = "0.4.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anndata" }, @@ -2480,9 +2480,9 @@ dependencies = [ { name = "numpy" }, { name = "scipy" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/3d/a5/78db8f12cb1dfe2e0a2f7d637efef572ffb75578d01e34f2c1d2a4327747/vsparse-0.2.0.tar.gz", hash = "sha256:aba6221ec50ce3cbe6bc149d69216202c66b2361c9b7a7ba09e3f3bc21a52cff", size = 37836, upload-time = "2026-09-04T12:50:17.51Z" } +sdist = { url = "https://files.pythonhosted.org/packages/90/36/c4e280f14af925b71c223777389bfa3ad472f4540375880bfe9b50ac6225/vsparse-0.4.0.tar.gz", hash = "sha256:72faa01087fcdd3fd72467638303c53838ff226685b0126cf5c060fc54626742", size = 47657, upload-time = "2026-09-14T17:12:39.176Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/28/14/f2a51a5aa3a3d131429ed7d238d9a13d32dd10493b3019e69e9fd4399420/vsparse-0.2.0-py3-none-any.whl", hash = "sha256:567035afa9a51c5edfc3103a9f98f95568f4ba5e2a9477f73031b703a376d277", size = 45084, upload-time = "2026-09-04T12:50:16.137Z" }, + { url = "https://files.pythonhosted.org/packages/ff/4e/3a167ca192bcb10623cbf403e76798d600e249c7c26f0bd28f24d0ed15b6/vsparse-0.4.0-py3-none-any.whl", hash = "sha256:a294b647dd8fac06746a8d8ffd1a778a31d60a6387d50f5915a406a37d768f2c", size = 54961, upload-time = "2026-09-14T17:12:38.011Z" }, ] [[package]]