From b65c32168c347590032aeb2466f8d7ae16373ad1 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Sun, 13 Sep 2026 12:24:31 -0700 Subject: [PATCH 1/3] Keep BiCV's held-out cell blocks lazy for duck-typed backends _bicv_trial evaluated a rank's fit by row-restricting X.X with a boolean cell mask (X_mat[train_cell_mask]/X_mat[test_cell_mask]) and handing the result to rmatmul/calc_W, deliberately avoiding materializing the raw data ("reaches the raw data through products... rather than materialising a block of it"). That holds for a plain ndarray or scipy-sparse X_mat, but not for a vsparse normalized view (e.g. BAL-Pf2's lazy-normalized-view AnnData): bracket indexing such a view is documented to always eagerly build a dense ndarray for the selection. At BAL-Pf2's real scale (1.3M cells), a single train/test split (~50% of all cells) would materialize on the order of tens of GB, once or twice per BiCV trial, across every rank/repeat in a sweep -- never surfaced before because an unrelated vsparse memory issue always killed these runs earlier in the pipeline. Adds _restrict_rows(X_mat, mask), which uses vsparse's new select(recalculate=False) (meyer-lab/vsparse#50) to stay a genuinely lazy view for a duck-typed backend, falling back to ordinary indexing for plain dense/sparse X_mat (unaffected, still cheap for those). Also adds a third branch to _test_block_moments for the same duck-typed case: selects the cell subset lazily, then streams it in bounded row chunks rather than ever materializing the whole subset as one dense block. Temporarily pins vsparse to its (as yet unmerged/unreleased) select-without-recalculate branch in pyproject.toml -- drop once meyer-lab/vsparse#50 merges and releases, reverting to the plain PyPI version constraint. Co-Authored-By: Claude Sonnet 5 --- pyproject.toml | 8 ++ scrise/rank_selection.py | 60 +++++++++++-- scrise/tests/test_rank_selection.py | 133 ++++++++++++++++++++++++++++ uv.lock | 10 +-- 4 files changed, 197 insertions(+), 14 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 0852ae0c..00d17ecc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,14 @@ fbuild = "analysis.figures.common:genFigure" requires = ["uv_build>=0.12.0,<0.13"] build-backend = "uv_build" +# TEMPORARY: pin vsparse to an unreleased PR branch for select(recalculate=False) +# (meyer-lab/vsparse#50), needed by rank_selection's BiCV evaluation to stay a +# lazy view over a large test-cell block instead of eagerly materializing it. +# Drop this once #50 merges and a release picks it up, reverting to the plain +# PyPI version constraint above. +[tool.uv.sources] +vsparse = { git = "https://github.com/meyer-lab/vsparse", branch = "select-without-recalculate" } + [tool.uv.build-backend] module-name = ["scrise"] module-root = "" 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..cc8cd0ea 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", git = "https://github.com/meyer-lab/vsparse?branch=select-without-recalculate" }, ] provides-extras = ["gpu"] @@ -2471,8 +2471,8 @@ wheels = [ [[package]] name = "vsparse" -version = "0.2.0" -source = { registry = "https://pypi.org/simple" } +version = "0.3.0" +source = { git = "https://github.com/meyer-lab/vsparse?branch=select-without-recalculate#87a9a13c9d0cc84f59608b5595289ba5f9147f78" } dependencies = [ { name = "anndata" }, { name = "hdf5plugin" }, @@ -2480,10 +2480,6 @@ 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" } -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" }, -] [[package]] name = "watchdog" From 2b5e886683d13d4fa54c54bf38a682550d9a9863 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Sun, 13 Sep 2026 12:26:03 -0700 Subject: [PATCH 2/3] Point the temporary vsparse pin at the combined testing branch Avoids a conflicting-git-ref resolution error for downstream consumers (e.g. BAL-Pf2) that pin vsparse's combined bal-pf2-gpu-testing branch (carrying both #49's matmul kernel and #50's select(recalculate=False)) rather than #50's branch alone. Co-Authored-By: Claude Sonnet 5 --- pyproject.toml | 15 +++++++++------ uv.lock | 4 ++-- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 00d17ecc..e675e0b2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,13 +39,16 @@ fbuild = "analysis.figures.common:genFigure" requires = ["uv_build>=0.12.0,<0.13"] build-backend = "uv_build" -# TEMPORARY: pin vsparse to an unreleased PR branch for select(recalculate=False) -# (meyer-lab/vsparse#50), needed by rank_selection's BiCV evaluation to stay a -# lazy view over a large test-cell block instead of eagerly materializing it. -# Drop this once #50 merges and a release picks it up, reverting to the plain -# PyPI version constraint above. +# TEMPORARY: pin vsparse to an unreleased combined testing branch (carrying +# both meyer-lab/vsparse#49's matmul kernel and meyer-lab/vsparse#50's +# select(recalculate=False), which rank_selection's BiCV evaluation needs to +# stay a lazy view over a large test-cell block instead of eagerly +# materializing it) rather than #50's own branch directly, so downstream +# consumers pinning the same combined branch (e.g. BAL-Pf2) don't hit a +# conflicting-git-ref resolution error. Drop this once #49/#50 merge and a +# release picks them up, reverting to the plain PyPI version constraint above. [tool.uv.sources] -vsparse = { git = "https://github.com/meyer-lab/vsparse", branch = "select-without-recalculate" } +vsparse = { git = "https://github.com/meyer-lab/vsparse", branch = "bal-pf2-gpu-testing" } [tool.uv.build-backend] module-name = ["scrise"] diff --git a/uv.lock b/uv.lock index cc8cd0ea..4b126c9e 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", git = "https://github.com/meyer-lab/vsparse?branch=select-without-recalculate" }, + { name = "vsparse", git = "https://github.com/meyer-lab/vsparse?branch=bal-pf2-gpu-testing" }, ] provides-extras = ["gpu"] @@ -2472,7 +2472,7 @@ wheels = [ [[package]] name = "vsparse" version = "0.3.0" -source = { git = "https://github.com/meyer-lab/vsparse?branch=select-without-recalculate#87a9a13c9d0cc84f59608b5595289ba5f9147f78" } +source = { git = "https://github.com/meyer-lab/vsparse?branch=bal-pf2-gpu-testing#0beb1af8dfd155a2a21fa61feb824b1c11be65a3" } dependencies = [ { name = "anndata" }, { name = "hdf5plugin" }, From 499d5bfbfbd27cbb8bec5be93f8a0f4f0c810f2b Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Mon, 14 Sep 2026 10:25:10 -0700 Subject: [PATCH 3/3] Switch vsparse dependency back to PyPI, pinned to >=0.4.0 The matmul kernel and lazy select(recalculate=False) this branch needed have been released, so drop the temporary git pin to the combined testing branch in favor of the PyPI release. Co-Authored-By: Claude Sonnet 5 --- pyproject.toml | 13 +------------ uv.lock | 10 +++++++--- 2 files changed, 8 insertions(+), 15 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index e675e0b2..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", @@ -39,17 +39,6 @@ fbuild = "analysis.figures.common:genFigure" requires = ["uv_build>=0.12.0,<0.13"] build-backend = "uv_build" -# TEMPORARY: pin vsparse to an unreleased combined testing branch (carrying -# both meyer-lab/vsparse#49's matmul kernel and meyer-lab/vsparse#50's -# select(recalculate=False), which rank_selection's BiCV evaluation needs to -# stay a lazy view over a large test-cell block instead of eagerly -# materializing it) rather than #50's own branch directly, so downstream -# consumers pinning the same combined branch (e.g. BAL-Pf2) don't hit a -# conflicting-git-ref resolution error. Drop this once #49/#50 merge and a -# release picks them up, reverting to the plain PyPI version constraint above. -[tool.uv.sources] -vsparse = { git = "https://github.com/meyer-lab/vsparse", branch = "bal-pf2-gpu-testing" } - [tool.uv.build-backend] module-name = ["scrise"] module-root = "" diff --git a/uv.lock b/uv.lock index 4b126c9e..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", git = "https://github.com/meyer-lab/vsparse?branch=bal-pf2-gpu-testing" }, + { name = "vsparse", specifier = ">=0.4.0" }, ] provides-extras = ["gpu"] @@ -2471,8 +2471,8 @@ wheels = [ [[package]] name = "vsparse" -version = "0.3.0" -source = { git = "https://github.com/meyer-lab/vsparse?branch=bal-pf2-gpu-testing#0beb1af8dfd155a2a21fa61feb824b1c11be65a3" } +version = "0.4.0" +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anndata" }, { name = "hdf5plugin" }, @@ -2480,6 +2480,10 @@ dependencies = [ { name = "numpy" }, { name = "scipy" }, ] +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/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]] name = "watchdog"