Add select(recalculate=False): a lazy, stats-preserving row/col selection - #50
Merged
Conversation
…tion select() previously always recomputed statistics fresh for the selected sub-array. A caller that instead wants to evaluate a fixed (e.g. train-derived) normalization against a different slice -- such as bi-cross-validation scoring a held-out test block against train-derived statistics -- had no lazy option: bracket indexing (__getitem__) keeps the parent's statistics but is documented to always eagerly materialize the selection as a dense ndarray, which is fine for a small window but not for a selection covering a large fraction of a huge array. select(rows, cols, recalculate=False) reuses this view's existing a/b/c/s (row_scale sliced by rows -- exact, since it's a per-cell quantity; the per-gene stats sliced by cols, unchanged by which rows are selected) via the existing from_stats() classmethod, and returns a real view rather than an array -- so it composes with `@`/toarray() and stays lazy regardless of selection size, unlike __getitem__. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This was referenced Sep 13, 2026
aarmey
added a commit
to meyer-lab/RISE
that referenced
this pull request
Sep 14, 2026
* 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 <noreply@anthropic.com>
* 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 <noreply@anthropic.com>
* 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 <noreply@anthropic.com>
---------
Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
aarmey
added a commit
to meyer-lab/RISE
that referenced
this pull request
Sep 14, 2026
…#557) * 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 <noreply@anthropic.com> * 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 <noreply@anthropic.com> * Compress a BiCV trial's held-out split once, reused across every rank bicv() previously called run_parafac2 once per (rank, repeat) pair, and run_parafac2 unconditionally recompressed its input via compress_dataset whenever compression_kwarg was given -- so a 100-rank x 3-repeat sweep ran ~300 full CANDELINC compressions of freshly-split raw data, each an O(nnz) pass, even though CompressedData is explicitly designed to support fitting any rank <= its own L_g/L_c ("enabling fast rank sweeps without touching the raw data again" -- parafac2.compress.CompressedData) and parafac2_nd already accepts a precomputed CompressedData directly. Each BiCV trial's train/test split is independent of rank (only the downstream fit is), and the in-sample fit's full dataset doesn't change across ranks either -- so both were needlessly recompressing the same data at every rank. Adds _fit_at_ranks(X_in, ranks, ...), which -- when compression_kwarg is given -- compresses X_in exactly once via compress_dataset, sized for max(ranks) via compress_dataset's own rank parameter, then calls parafac2_nd(compressed, rank=r) directly for every r in ranks, reusing that one compressed representation. Without compression_kwarg, behavior is unchanged: each rank still goes through run_parafac2's own per-call compression shortcut, since there's no CompressedData object to hoist out in that path. _bicv_trial now takes `ranks: Sequence[int]` instead of a single `rank` and returns one result dict per rank, computing the once-per-trial split and its downstream references (test-block moments, restricted train/test views) exactly once and reusing them for every rank's held-out scoring, rather than recomputing per rank. bicv()'s main loop restructures to match: one held-out split per repeat, evaluated at every rank, instead of one independent split per (rank, repeat) pair. New tests confirm compress_dataset is called exactly (1 + n_repeats) times for an n-rank sweep (not n * (1 + n_repeats)), each sized for the largest requested rank, and that the compression_kwarg-less path is unaffected. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014MHUrbvYnPr1naj92oYTab * Point the vsparse pin at the new to-scipy-sparse-dtype branch vsparse#49/#50 (the fixes this pin originally existed for) have merged to vsparse's main. Repointing to meyer-lab/vsparse@to-scipy-sparse-dtype (main plus vsparse#51's still-open dtype argument for to_scipy_sparse, which RISE doesn't itself need) rather than main directly, so downstream consumers pinning the same branch for #51 (e.g. BAL-Pf2) don't hit a conflicting-git-ref resolution error against RISE's own vsparse pin. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014MHUrbvYnPr1naj92oYTab * Pin parafac2 to at least the latest PyPI release (1.7.0) Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Remove the temporary vsparse git pin now that 0.4.0 is on PyPI Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
select()previously always recomputed statistics fresh for the selectedsub-array. A caller that instead wants to evaluate a fixed (e.g.
train-derived) normalization against a different slice had no lazy option
for that: bracket indexing (
__getitem__) keeps the parent's statistics,but is documented to always eagerly materialize the selection as a dense
ndarray(toarray()[key]) -- fine for a small window, but not for aselection covering a large fraction of a huge array.
This surfaced while getting BAL-Pf2's BiCV rank-selection sweep working at
its real scale (1.3M cells):
scrise.rank_selection._bicv_trialscores amodel fit by indexing a normalized view with a boolean cell mask covering
~50% of all cells (
X_mat[test_cell_mask]), which -- via__getitem__--eagerly builds a dense array on the order of tens of GB, for something that
should stay a lazy sparse view all the way through to the eventual
dense-times-sparse product.
Changes
select(rows, cols, *, recalculate=True)gains therecalculatekeyword.recalculate=Falsereuses this view's existinga/b/c/s(row_scalesliced by
rows-- exact, since it's a per-cell quantity; the per-genestatistics sliced by
cols, unaffected by which rows are selected) via theexisting
from_stats()classmethod, rather than recomputing from theselected sub-array's own data. The default (
recalculate=True) isunchanged.
The result is a real view (not an array), so it composes with
@/toarray()and stays lazy regardless of selection size, unlike__getitem__-- the caller decides when (or whether) to ever materializeit, and even then the O(nnz) machinery is bounded rather than allocating a
dense block sized to the selection up front.
Test plan
uv run pytest-- 1246 passed, 49 skipped, 3 pre-existing failuresunrelated to this change (identical on
mainwithout this diff)uv run ruff check ./uv run ruff format --check .uv run ty check src/tests/test_norm_selection.py: matches__getitem__'svalues, returns a genuinely lazy view that still composes with
@,composes with a column selection too, and the default
(
recalculate=True) is unaffected by the new keyword.