Fix/bicv alignment and diagnostics - #550
Merged
Merged
Conversation
Contributor
Author
fishidaho
force-pushed
the
fix/bicv-alignment-and-diagnostics
branch
from
September 11, 2026 21:33
e842613 to
06ea4fa
Compare
aarmey
approved these changes
Sep 11, 2026
fishidaho
added this pull request to stack #553
September 12, 2026 00:58
…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
force-pushed
the
fix/bicv-alignment-and-diagnostics
branch
from
September 12, 2026 01:17
8435d14 to
9c9108b
Compare
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.
The goal of this PR is to fix #546 and close #547 as well as #548.
#546 —
_bicv_trialmisalignedZagainst the rows it is regressed onThe held-out gene regression built its design matrix by concatenating per-condition blocks in condition order:
but regressed it against
X_train_test_genes, which stays in the training cells' natural row order.P_train[i]lists conditioni's cells in their within-condition order, so concatenating yieldsZsorted by condition. The two coincide only when conditions happen to be stored contiguously._train_cell_loadingsscatters 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.
#548 —
bicvreturned aggregated R2X onlyBiCV rows now also carry:
Train Block R2XNTrainGenes,NTestGenes,NTrainCells,NTestCellsSeedThese are columns rather than a new
Metricvalue, soplot_bicv_r2xis 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.mddocuments the new columns.#547 — three dense blocks materialized per trial
_bicv_trialdensified three blocks of the expression matrix on every trial of every rank:X[train_cells][:, test_genes]X[test_cells][:, train_genes]X[test_cells][:, test_genes]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:
C_testnow solves the normal equations rather thanlstsqon the tall design, which is never formed.lstsqon 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], whichrun_parafac2needs as an AnnData. Removing it would need a gene mask inparafac2.compress_dataset. We had proposed this earlier but agreed that putting cross-validation logic inparafac2was out of scope for that repository.