Compress a BiCV trial's held-out split once, reused across every rank - #557
Merged
Merged
Conversation
_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>
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>
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
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
…r-trial # Conflicts: # scrise/rank_selection.py # uv.lock
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
Stacked on #556 (this branch is based on it and should merge after it).
bicv()previously calledrun_parafac2once per(rank, repeat)pair,and
run_parafac2unconditionally recompressed its input viacompress_datasetwhenevercompression_kwargwas given — so a100-rank x 3-repeat sweep ran ~300 full CANDELINC compressions of
freshly-split raw data, each an
O(nnz)pass over the (train-block) data.This was unnecessary:
CompressedDatais explicitly designed to supportfitting any rank
<= L_g/L_c("enabling fast rank sweeps withouttouching the raw data again" — its own docstring), and
parafac2_ndalready accepts a precomputed
CompressedDatadirectly. A BiCV trial'strain/test split is independent of rank (only the downstream PARAFAC2 fit
depends on rank), and the in-sample fit's full dataset doesn't change
across ranks either — so both paths were recompressing the same data at
every rank for no reason.
Changes
_fit_at_ranks(X_in, ranks, ...): whencompression_kwargis given,compresses
X_inexactly once viacompress_dataset, sized formax(ranks)(viacompress_dataset's ownrankparameter, which setsthe compression dimensions for the largest rank you intend to fit), then
calls
parafac2_nd(compressed, rank=r)directly for everyrinranks, reusing that one compressed representation. Withoutcompression_kwarg, behavior is unchanged: each rank still goes throughrun_parafac2's own per-call internal compression shortcut, since thereis no
CompressedDataobject to hoist out of that path._bicv_trialnow takesranks: Sequence[int]instead of a singlerank, and returns one result dict per rank — computing theonce-per-trial split and its rank-independent 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 perrepeat, evaluated at every rank, instead of one independent split per
(rank, repeat)pair. Same for the in-sample fit: the full dataset iscompressed once (sized for the largest rank) and reused across ranks.
Test plan
uv run pytest scrise/tests/— 133 passed, 1 skipped (pre-existing,unrelated)
uv run ruff check ./uv run ruff format --check .uv run ty check scrise/compress_datasetis called exactly1 + n_repeatstimes for an
n-rank sweep (notn * (1 + n_repeats)), each callsized for the largest requested rank; the
compression_kwarg-lesspath is unaffected (still one
run_parafac2call per(rank, repeat)pair, as before).dense-reference comparison tests (
test_streamed_scoring_matches_*)and the seed-reproducibility test, confirming the single-rank,
no-
compression_kwargcode path is bit-for-bit unaffected.Related
this builds on)