Skip to content

Add select(recalculate=False): a lazy, stats-preserving row/col selection - #50

Merged
aarmey merged 1 commit into
mainfrom
select-without-recalculate
Sep 13, 2026
Merged

Add select(recalculate=False): a lazy, stats-preserving row/col selection#50
aarmey merged 1 commit into
mainfrom
select-without-recalculate

Conversation

@aarmey

@aarmey aarmey commented Sep 13, 2026

Copy link
Copy Markdown
Member

Summary

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 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 a
selection 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_trial scores a
model 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 the recalculate keyword.
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
statistics sliced by cols, unaffected by which rows are selected) via the
existing from_stats() classmethod, rather than recomputing from the
selected sub-array's own data. The default (recalculate=True) is
unchanged.

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 materialize
it, 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 failures
    unrelated to this change (identical on main without this diff)
  • uv run ruff check . / uv run ruff format --check .
  • uv run ty check src/
  • New tests in tests/test_norm_selection.py: matches __getitem__'s
    values, 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.

…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>
@aarmey
aarmey merged commit 43d78cb into main Sep 13, 2026
6 checks passed
@aarmey
aarmey deleted the select-without-recalculate branch September 13, 2026 20:33
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant