Skip to content

Fix/bicv alignment and diagnostics - #550

Merged
aarmey merged 5 commits into
mainfrom
fix/bicv-alignment-and-diagnostics
Sep 12, 2026
Merged

Fix/bicv alignment and diagnostics#550
aarmey merged 5 commits into
mainfrom
fix/bicv-alignment-and-diagnostics

Conversation

@fishidaho

Copy link
Copy Markdown
Contributor

The goal of this PR is to fix #546 and close #547 as well as #548.


#546_bicv_trial misaligned Z against the rows it is regressed on

The held-out gene regression built its design matrix by concatenating per-condition blocks in condition order:

Z = np.concatenate(
    [(P_train[i] @ B) * A[i] for i in range(n_cond) if np.any(cond_train == i)],
    axis=0,
)

but regressed it against X_train_test_genes, which stays in the training cells' natural row order. P_train[i] lists condition i's cells in their within-condition order, so concatenating yields Z sorted by condition. The two coincide only when conditions happen to be stored contiguously.

_train_cell_loadings scatters each condition's block back to the positions its cells actually occupy. Same asymptotic cost, and identical output when conditions are contiguous.

Might be worth retrying the hESC trial now to see if the R2X truly does recover when there is meaningful structure left.


#548bicv returned aggregated R2X only

BiCV rows now also carry:

column why
Train Block R2X In-sample R2X on the block the model was fit to
NTrainGenes, NTestGenes, NTrainCells, NTestCells The realised block sizes
Seed The seed determining that trial's splits and initialisation

These are columns rather than a new Metric value, so plot_bicv_r2x is unchanged. test_bicv_shape_and_range's column assertion is relaxed from an exact set to a subset, since the return shape is intentionally wider. docs/tutorial.md documents the new columns.


#547 — three dense blocks materialized per trial

_bicv_trial densified three blocks of the expression matrix on every trial of every rank:

block purpose
X[train_cells][:, test_genes] regress the held-out gene loadings
X[test_cells][:, train_genes] project the held-out cells
X[test_cells][:, test_genes] score the reconstruction

Each is O(cells × genes) where the matrix itself is O(nnz). Measured on a 12,000 × 4,000 dataset at 10% density the blocks alone are more than twice the data they come from. At cohort scale (1.2M × 34k, 20%/20% holdout) that is roughly 115 GB of dense allocation per trial against a 44 GB matrix.

All three are now reached through products against the full matrix, with the restriction applied by zeroing the dense operand so the result is the submatrix's without the copy:

Z^T (X[train, test] − 1 μ^T)  =  (Z_full^T X)[:, test] − (Z^T 1) μ^T

C_test now solves the normal equations rather than lstsq on the tall design, which is never formed. lstsq on the rank × rank system keeps minimum-norm behaviour when the fit is rank deficient, at some cost in conditioning.

The remaining per-trial copy is X[train_cells][:, train_genes], which run_parafac2 needs as an AnnData. Removing it would need a gene mask in parafac2.compress_dataset. We had proposed this earlier but agreed that putting cross-validation logic in parafac2 was out of scope for that repository.

@fishidaho
fishidaho requested a review from aarmey September 11, 2026 18:32
@fishidaho

Copy link
Copy Markdown
Contributor Author

@aarmey, I am not sure what experiment you had initially ran with #546 with the hESC data, but it might be worth rerunning it with the same normalization scheme you had done prior to see if the rank is no longer artificially low.

@fishidaho
fishidaho force-pushed the fix/bicv-alignment-and-diagnostics branch from e842613 to 06ea4fa Compare September 11, 2026 21:33
@fishidaho
fishidaho added this pull request to stack #553 September 12, 2026 00:58
fishidaho and others added 5 commits September 11, 2026 18:17
…inst

Fixes #546.

`_bicv_trial` built the held-out gene regression's design matrix by
concatenating per-condition blocks in condition order:

    Z = np.concatenate(
        [(P_train[i] @ B) * A[i] for i in range(n_cond) if np.any(cond_train == i)],
        axis=0,
    )

but regressed it against `X_train_test_genes`, which stays in the training
cells' natural row order. `P_train[i]` lists condition `i`'s cells in their
within-condition order, so concatenating produces `Z` sorted by condition. The
two orderings coincide only when conditions happen to be stored as contiguous
blocks. On pooled data -- a Perturb-seq screen, or anything not physically
sorted by sample before sequencing -- conditions are interleaved, and row `k` of
`Z` then describes a different cell from row `k` of the expression it is paired
with. Every held-out R2X the function reports is computed from that
least-squares fit.

`_train_cell_loadings` scatters each condition's block back to the positions its
cells actually occupy. Same asymptotic cost, and identical output in the
contiguous case.

Three of the four new tests fail against the previous implementation and pass
against this one; the fourth asserts the no-op on contiguous input, and passes
either way by construction. That split is deliberate -- without the
"differ when interleaved" test, an implementation that quietly reverted to
concatenation would still satisfy the rest.

The missing-condition case is also covered: when a fold leaves some condition
with no training cells, concatenation produced the right number of rows in the
wrong order, which is the same defect in a form that is even harder to notice.

Why every existing test missed this: `_make_test_data` builds conditions as
contiguous blocks, which is precisely the case where the bug is invisible. The
end-to-end test added here stores the same cells interleaved.

Reported by @aarmey, with the diagnosis and the suggested fix.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Closes #548.

`bicv` returned `Rank / Repeat / Metric / R2X`, which is enough to plot a curve
but not enough to audit the rank it implies. BiCV rows now also carry:

    Train Block R2X   in-sample R2X on the block the model was fit to
    NTrainGenes, NTestGenes, NTrainCells, NTestCells
    Seed              the seed determining that trial's splits and init

`Train Block R2X` is the one that changes what the output can be used for. The
useful reading of a BiCV curve is the held-out score turning over while the
in-sample score keeps climbing, and the existing "Fit R2X" metric cannot supply
that comparison: it is a separate fit, on the full dataset rather than on the
training block, so it answers a different question. The value was already being
computed and discarded.

`_bicv_trial` now takes an explicit `seed` and builds its own generator instead
of drawing from a shared one, which is what makes `Seed` meaningful -- a single
trial can be replayed from the value in the table. A test asserts that replay
reproduces the recorded scores exactly.

The block sizes are reported rather than assumed because the cell split is
stratified per condition, so the realised fraction is only approximately the
requested one, and the spread across repeats has to be read against the block
each repeat actually used.

These columns are NaN on "Fit R2X" rows, which come from an unsplit fit.

The additions are columns rather than a new `Metric` value, so
`plot_bicv_r2x` -- which colours by `Metric` -- is unchanged. The
`test_bicv_shape_and_range` column assertion is relaxed from an exact set to a
subset, since the return shape is intentionally wider.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Closes #547.

`_bicv_trial` densified three blocks of the expression matrix on every trial of
every rank:

    X[train_cells][:, test_genes]     to regress the held-out gene loadings
    X[test_cells][:, train_genes]     to project the held-out cells
    X[test_cells][:, test_genes]      to score the reconstruction

Each is O(cells x genes) where the matrix itself is O(nnz). On a 12,000 x 4,000
dataset at 10% density they measure 37.2, 37.0 and 9.3 MB against a 38 MB
sparse matrix -- so the blocks alone are more than twice the data they are drawn
from. At cohort scale (1.2M x 34k, 20%/20% holdout) that is roughly 115 GB of
dense allocation per trial against a 44 GB matrix, which is simply not runnable.

All three are now reached through products against the *full* matrix, with the
cell or gene restriction applied by zeroing the dense operand. A masked-out gene
contributes zero to every product it appears in, so the result is the
submatrix's, without the copy:

    Z^T (X[train, test] - 1 mu^T)  =  (Z_full^T X)[:, test] - (Z^T 1) mu^T

The same identity supplies the cross term in the score. A gene factor that is
zero on the held-out genes makes `calc_W` ignore them -- including in its
`means @ C` centering -- so projecting the held-out cells needs no slice either.
Writing the reconstruction as `L @ C_test^T` over all test cells lets the score
decompose into three small pieces: the block's column moments, one sparse
product, and two rank x rank Grams.

`_test_block_moments` is the only place raw data is newly summarised. It streams
in row blocks, because the per-nonzero column lookup it needs is otherwise as
large as the matrix -- the same trap the norms elsewhere fall into.

`C_test` now solves the normal equations rather than `lstsq` on the tall design,
which is never formed. `lstsq` on the rank x rank system keeps the minimum-norm
behaviour when the fit is rank deficient, at some cost in conditioning.

**This is an optimisation, so the numbers must not move.** Four tests fit a
model and score it both ways -- dense and streamed -- on dense and sparse input,
two seeds, and on deliberately interleaved conditions, asserting the held-out
R2X agrees to 1e-6 relative. Two more pin `_test_block_moments` against a
materialised block, one of them with the chunk size forced small so the
multi-block accumulation path is actually exercised.

`_train_cell_loadings` is renamed `_cell_loadings`: it now assembles the
per-cell loadings for the test cells as well, which is what lets the score be
written as a single product.

The remaining per-trial copy is `X[train_cells][:, train_genes]`, which
`run_parafac2` needs as an AnnData. Removing it would need a gene mask in
`parafac2.compress_dataset`; that was proposed and declined as putting
cross-validation logic in the wrong library, so it stays.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Three review findings on this branch, none changing any score --
`_bicv_trial` still reproduces an independent dense reference to 7.6e-08
on both contiguous and interleaved row orders, bit-for-bit what it
returned before these edits.

`_test_block_moments` accumulated the dense branch's column sums in the
input dtype while widening the squares to float64 in the same return
statement. On a float32 block the sums drifted ~5e-04 relative, and the
two branches disagreed by that much on identical data -- which reaches
the reported R2X through `ss_tot`, itself only O(1e-2). Both moments now
accumulate in float64, with a test over 200k float32 rows that pins the
two branches to each other; the previous test could not catch this,
using a float64 array of 120 rows.

The cell restriction now slices the matrix rows instead of zeroing a
full-height dense operand, dropping two (n_obs x rank) allocations and
making each product proportional to the rows actually wanted rather than
to the whole dataset -- `calc_W` in particular was computing W for every
cell and discarding ~80%. The operands stay float64 and are now
explicitly C-contiguous. That deliberately departs from the dtype
convention in `calc_W`/`parafac_update`, so there is a comment saying
why: a float32 operand measures ~2x faster but costs ~3e-2 relative
error on the product, which this score cannot absorb.

Subscripting `X.X` needs the `assert`/`cast` that `extract_dataset_info`
uses, since its union admits `None` and `ZappyArray`; `ty` is clean over
`scrise/` again.

Finally, the diagnostics table moved below the paragraph introducing the
returned frame, which it previously preceded while that paragraph still
listed only four columns, and it gained the `Seed` row the docstring
already documented.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_016VJRB1RMUAb1id11XZ6KWS
@aarmey
aarmey force-pushed the fix/bicv-alignment-and-diagnostics branch from 8435d14 to 9c9108b Compare September 12, 2026 01:17
@aarmey
aarmey merged commit efe0a7d into main Sep 12, 2026
6 checks passed
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.

bicv materializes four AnnData copies per trial bicv(): _bicv_trial misaligns Z against X_train_test_genes when conditions aren't stored contiguously

2 participants