Skip to content

Compress a BiCV trial's held-out split once, reused across every rank - #557

Merged
aarmey merged 7 commits into
mainfrom
bicv-compress-once-per-trial
Sep 14, 2026
Merged

aarmey merged 7 commits into
mainfrom
bicv-compress-once-per-trial

Conversation

@aarmey

@aarmey aarmey commented Sep 13, 2026

Copy link
Copy Markdown
Member

Summary

Stacked on #556 (this branch is based on it and should merge after it).

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 over the (train-block) data.

This was unnecessary: CompressedData is explicitly designed to support
fitting any rank <= L_g/L_c ("enabling fast rank sweeps without
touching the raw data again" — its own docstring), and parafac2_nd
already accepts a precomputed CompressedData directly. A BiCV trial's
train/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, ...): when compression_kwarg is given,
    compresses X_in exactly once via compress_dataset, sized for
    max(ranks) (via compress_dataset's own rank parameter, which sets
    the compression dimensions for the largest rank you intend to fit), 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 internal compression shortcut, since there
    is no CompressedData object to hoist out of 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 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 per
    repeat, evaluated at every rank
    , instead of one independent split per
    (rank, repeat) pair. Same for the in-sample fit: the full dataset is
    compressed 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/
  • New tests: compress_dataset is called exactly 1 + n_repeats
    times for an n-rank sweep (not n * (1 + n_repeats)), each call
    sized for the largest requested rank; the compression_kwarg-less
    path is unaffected (still one run_parafac2 call per (rank, repeat) pair, as before).
  • All existing correctness tests pass unchanged, including the
    dense-reference comparison tests (test_streamed_scoring_matches_*)
    and the seed-reproducibility test, confirming the single-rank,
    no-compression_kwarg code path is bit-for-bit unaffected.

Related

aarmey and others added 4 commits September 13, 2026 12:24
_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
Base automatically changed from vsparse-lazy-bicv-blocks to main September 14, 2026 17:26
aarmey and others added 3 commits September 14, 2026 10:28
…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>
@aarmey
aarmey merged commit e6dfb94 into main Sep 14, 2026
5 checks passed
@aarmey
aarmey deleted the bicv-compress-once-per-trial branch September 14, 2026 17:36
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