Keep BiCV's held-out cell blocks lazy for duck-typed backends

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",

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,17 +285,18 @@ 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)
 
     # Score the held-out block. `A` carries the training slice's energy, so the
     # 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)

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