From 64b92a47cd44abb7fd780f8332b993c9a4e95b67 Mon Sep 17 00:00:00 2001 From: fishidaho Date: Fri, 11 Sep 2026 16:43:34 -0700 Subject: [PATCH 1/2] Correct the held-out slice scale, and hold out half of each mode PARAFAC2 writes slice k as `P_k (B diag(a_k)) C^T` with `P_k` orthonormal, so the reconstruction's norm is `||(B diag(a_k)) C^T||` regardless of how many cells the slice holds: all of a slice's energy sits in `a_k`. BiCV fit `a_k` against the training cells and then reused it unchanged for the held-out cells, and since a slice's energy grows like sqrt(n_k) that overstated the predicted block by sqrt(n_train_k / n_test_k) -- a factor of two at the old default split. The symptom was a metric that got worse as the model got better. On noiseless, exactly-rank-4 data fit at the true rank, the training block scored 0.9996 while the held-out block scored -0.034: the fit had recovered the structure exactly and the score still called it worse than predicting the mean. On the IBD cohort the curve was negative at every rank and fell monotonically as rank rose, putting its argmax at rank 1 and firing the "rank at the edge of tested ranks" warning. The error was purely one of magnitude -- a single global rescale of the prediction recovered R^2 = 0.970. `_holdout_scale` corrects each condition by sqrt(n_test_k / n_train_k), which needs only the cell counts and is exact when the two halves of a condition have comparable per-cell energy -- true here, since the split is random within each condition. Three candidates were measured (this one, re-fitting `a_k` per held-out slice from the train genes, and unconstrained least-squares cell loadings); all three fixed the bug and none scored above zero on pure noise, so the narrowest was taken. Both held-out fractions now default to one half, following Owen and Perry (https://arxiv.org/abs/0908.2062), who report that "in simulated examples we find that a method leaving out half the rows and half the columns performs well". That is also the one split at which the old score happened to be unbiased, so it is the setting under which older numbers remain comparable, and it recovers the true rank in every simulated scenario tested -- which the correction alone did not at 0.2. Measured after the change: noiseless rank-4 data at its own rank scores 0.97-0.99 at every fraction (was -0.034); the spread across fractions 0.1-0.7 falls from 3.75 to 0.02; argmax lands on the true rank for (4, noise 0.3), (6, noise 0.3) and (4, noise 1.0); pure noise still scores negative. Row-order invariance from the alignment fix underneath is undisturbed, agreeing to 8.9e-16 between contiguous and fully interleaved conditions. The dense reference in the tests derives the same factor independently from the counts, so it still pins the streaming algebra rather than the old behaviour. The tutorial's "peaks" claim is softened to describe the plateau the curve actually forms, since past the plateau neighbouring ranks differ only in the fourth decimal. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_016VJRB1RMUAb1id11XZ6KWS --- docs/tutorial.md | 4 +- scrise/rank_selection.py | 54 ++++++++++-- scrise/tests/test_rank_selection.py | 127 +++++++++++++++++++++++++++- 3 files changed, 174 insertions(+), 11 deletions(-) diff --git a/docs/tutorial.md b/docs/tutorial.md index 041fdefe..4b79d7ce 100644 --- a/docs/tutorial.md +++ b/docs/tutorial.md @@ -92,7 +92,7 @@ plt.show() ### Select a Rank with Bi-Cross-Validation (BiCV) -$R^2X$ (in-sample fit) always increases with rank, so it cannot by itself identify an optimal rank — it can only reveal an elbow. Bi-cross-validation (BiCV) addresses this by holding out a random subset of cells *and* genes, fitting RISE on the remaining data, and scoring how well the fit predicts the held-out block. Because it is evaluated on unseen data, BiCV $R^2X$ penalizes overfitting and typically peaks (or plateaus) near the rank that generalizes best, unlike the fit $R^2X$. +$R^2X$ (in-sample fit) always increases with rank, so it cannot by itself identify an optimal rank — it can only reveal an elbow. Bi-cross-validation (BiCV) addresses this by holding out a random subset of cells *and* genes, fitting RISE on the remaining data, and scoring how well the fit predicts the held-out block. Because it is evaluated on unseen data, BiCV $R^2X$ penalizes overfitting and turns over near the rank that generalizes best, unlike the fit $R^2X$. In practice the curve more often flattens into a plateau than forms a sharp peak, so read the elbow where it stops climbing rather than trusting the single highest point -- past the plateau, neighbouring ranks can differ only in the fourth decimal. `scrise.rank_selection.bicv` exhaustively evaluates a list of candidate ranks, each with several repeated random train/test splits, and returns both the fit $R^2X$ and the BiCV $R^2X$ for every rank so you can inspect the full curve. Since PARAFAC2 fits use CANDELINC compression by default (`compress="auto"`), sweeping every rank in this way is cheap enough that there is no need to search for the best rank without evaluating every candidate. @@ -110,7 +110,7 @@ plt.tight_layout() plt.show() ``` -`bicv` returns a long-form DataFrame (columns `Rank`, `Repeat`, `Metric`, `R2X`, plus the per-trial diagnostics below) suitable for further analysis as well as plotting. Each BiCV trial holds out `held_out_cell_frac` of the cells within each condition and `held_out_gene_frac` of the genes (both default to 0.2); increase `n_repeats` for a smoother, less noisy BiCV curve at the cost of more compute. +`bicv` returns a long-form DataFrame (columns `Rank`, `Repeat`, `Metric`, `R2X`, plus the per-trial diagnostics below) suitable for further analysis as well as plotting. Each BiCV trial holds out `held_out_cell_frac` of the cells within each condition and `held_out_gene_frac` of the genes. Both default to 0.5, the half-rows-and-half-columns split recommended by [Owen and Perry](https://arxiv.org/abs/0908.2062), who report that "in simulated examples we find that a method leaving out half the rows and half the columns performs well". Increase `n_repeats` for a smoother, less noisy BiCV curve at the cost of more compute. Each BiCV row also carries the diagnostics needed to audit the curve rather than just plot it: diff --git a/scrise/rank_selection.py b/scrise/rank_selection.py index f928d284..a87ee5d3 100644 --- a/scrise/rank_selection.py +++ b/scrise/rank_selection.py @@ -8,6 +8,14 @@ predicts the held-out (test-cell x test-gene) block. Unlike the ordinary in-sample fit R2X (which increases monotonically with rank), the BiCV R2X penalizes overfitting and typically peaks near the "true" rank of the data. + +The held-out cells are scored through their own projections, whose scale has +to be corrected for the size of the held-out slice -- see +:func:`_holdout_scale`. Both held-out fractions default to one half, the +split Owen and Perry recommend for bi-cross-validation +(https://arxiv.org/abs/0908.2062); it is also the split at which the +uncorrected score happened to be unbiased, so it is the setting under which +older results are comparable to current ones. """ import warnings @@ -88,6 +96,34 @@ def _cell_loadings( return Z +def _holdout_scale( + cond_train: np.ndarray, cond_test: np.ndarray, n_cond: int +) -> np.ndarray: + """Per-held-out-cell factor correcting `A` for the held-out slice's size. + + PARAFAC2 writes slice ``k`` as ``P_k (B diag(a_k)) C^T`` with ``P_k`` + orthonormal, so the reconstruction's Frobenius norm is + ``||(B diag(a_k)) C^T||``, independent of how many cells the slice holds: + all of a slice's energy sits in ``a_k``. ``a_k`` is fit against + ``n_train_k`` cells, and a slice's energy grows like ``sqrt(n_k)``, so + reusing it unchanged for ``n_test_k`` held-out cells overstates the + predicted block by ``sqrt(n_train_k / n_test_k)`` -- a factor of two at + the default split, which drove the held-out R2X negative and made the + curve fall as the fit improved. + + Correcting it needs only the cell counts, which is exact when the two + halves of a condition have comparable per-cell energy -- true here, + since the split is random within each condition. + """ + scale = np.ones(cond_test.size) + train_counts = np.bincount(cond_train, minlength=n_cond) + test_counts = np.bincount(cond_test, minlength=n_cond) + for i in range(n_cond): + if test_counts[i] and train_counts[i]: + scale[cond_test == i] = np.sqrt(test_counts[i] / train_counts[i]) + return scale[:, np.newaxis] + + # Nonzeros per row block when streaming column moments. Bounds the per-nonzero # temporaries to a few hundred MB regardless of how large the matrix is. _MOMENT_CHUNK_NNZ = 50_000_000 @@ -231,8 +267,10 @@ def _bicv_trial( cond_slices_test = condition_slices(cond_test, n_cond) P_test, _ = project_data(W_test, [A, B, C], cond_slices_test) - # Score the held-out block. + # Score the held-out block. `A` carries the training slice's energy, so the + # held-out loadings need rescaling for the held-out slice's size. L = _cell_loadings(P_test, B, A, cond_test, n_cond) + L *= _holdout_scale(cond_train, cond_test, n_cond) LtY = np.asarray( rmatmul(np.ascontiguousarray(L.T), X_mat[test_cell_mask]), dtype=np.float64 )[:, test_gene_idx] @@ -266,8 +304,8 @@ def bicv( X: anndata.AnnData | None = None, ranks: Sequence[int] | None = None, n_repeats: int = 3, - held_out_cell_frac: float = 0.2, - held_out_gene_frac: float = 0.2, + held_out_cell_frac: float = 0.5, + held_out_gene_frac: float = 0.5, random_state: int | None = None, tolerance: float = 1e-6, max_iter: int = 200, @@ -298,10 +336,14 @@ def bicv( n_repeats : int, optional (default: 3) Number of independent random cell/gene splits per rank. Higher values give a less noisy BiCV estimate but take longer. - held_out_cell_frac : float, optional (default: 0.2) + held_out_cell_frac : float, optional (default: 0.5) Fraction of cells held out per condition in each BiCV trial. - held_out_gene_frac : float, optional (default: 0.2) - Fraction of genes held out in each BiCV trial. + held_out_gene_frac : float, optional (default: 0.5) + Fraction of genes held out in each BiCV trial. Both default to a + half-and-half split, following Owen and Perry, who report that + "in simulated examples we find that a method leaving out half the + rows and half the columns performs well" + (https://arxiv.org/abs/0908.2062). random_state : int, optional Random seed for reproducibility. tolerance : float, optional (default: 1e-6) diff --git a/scrise/tests/test_rank_selection.py b/scrise/tests/test_rank_selection.py index 960f1c30..3155ea59 100644 --- a/scrise/tests/test_rank_selection.py +++ b/scrise/tests/test_rank_selection.py @@ -13,6 +13,7 @@ from ..rank_selection import ( _bicv_trial, _cell_loadings, + _holdout_scale, _split_cells_by_condition, _split_genes, _test_block_moments, @@ -325,8 +326,8 @@ def test_reported_seed_reproduces_its_own_trial(): replayed = _bicv_trial( X, rank=3, - held_out_cell_frac=0.2, - held_out_gene_frac=0.2, + held_out_cell_frac=0.5, + held_out_gene_frac=0.5, seed=int(target["Seed"]), tolerance=1e-6, max_iter=60, @@ -388,8 +389,12 @@ def dense(mat): sel = cond_test == i if not np.any(sel): continue + n_train_i = int(np.sum(cond_train == i)) + # `A[i]` carries the training slice's energy; rescale it for this + # slice's size (see `_holdout_scale`). + scale = np.sqrt(int(sel.sum()) / n_train_i) if n_train_i else 1.0 actual = X_test_test[sel] - recon = ((P_test[i] @ B) * A[i]) @ C_test.T + recon = (((P_test[i] @ B) * A[i]) * scale) @ C_test.T ss_res += float(np.sum((actual - recon) ** 2)) ss_tot += float(np.sum(actual**2)) return 1.0 - ss_res / ss_tot @@ -500,3 +505,119 @@ def test_block_moments_stream_across_more_than_one_row_block(monkeypatch): block = dense[cell_mask] np.testing.assert_allclose(sums, block.sum(axis=0), rtol=1e-12) np.testing.assert_allclose(squares, np.sum(block**2, axis=0), rtol=1e-12) + + +# Held-out slice scaling + + +def _exact_low_rank( + rank_true=4, noise=0.0, n_cond=6, cells=(90, 160), n_genes=90, seed=0 +): + """Data that IS exactly rank `rank_true`, with unequal cells per condition. + + `means` is zero so the target stays exactly low rank -- mean-centering + would add a rank-one term and put a ceiling below 1.0 on the fit. + """ + rng = np.random.default_rng(seed) + C = rng.normal(size=(n_genes, rank_true)) + B = rng.normal(size=(rank_true, rank_true)) + A = np.abs(rng.normal(size=(n_cond, rank_true))) + 1.0 + blocks, cond = [], [] + for k in range(n_cond): + n_k = int(rng.integers(*cells)) + P, _ = np.linalg.qr(rng.normal(size=(n_k, rank_true))) + blocks.append(((P @ B) * A[k]) @ C.T) + cond.append(np.full(n_k, k)) + mat = np.concatenate(blocks, axis=0) + if noise: + mat = mat + noise * np.std(mat) * rng.normal(size=mat.shape) + adata = anndata.AnnData(X=mat) + adata.obs["condition_unique_idxs"] = np.concatenate(cond) + adata.var["means"] = np.zeros(n_genes) + return adata + + +def test_holdout_scale_is_the_square_root_of_the_cell_count_ratio(): + cond_train = np.repeat([0, 1, 2], [80, 40, 10]) + cond_test = np.repeat([0, 1, 2], [20, 40, 90]) + + scale = _holdout_scale(cond_train, cond_test, 3) + + assert scale.shape == (cond_test.size, 1) + np.testing.assert_allclose(scale[cond_test == 0, 0], np.sqrt(20 / 80)) + np.testing.assert_allclose(scale[cond_test == 1, 0], np.sqrt(40 / 40)) + np.testing.assert_allclose(scale[cond_test == 2, 0], np.sqrt(90 / 10)) + + +def test_holdout_scale_is_one_for_an_even_split(): + """The split the defaults use, and the only one that needed no correction.""" + cond_train = np.repeat([0, 1], [50, 30]) + cond_test = np.repeat([0, 1], [50, 30]) + np.testing.assert_allclose(_holdout_scale(cond_train, cond_test, 2), 1.0) + + +def test_holdout_scale_leaves_conditions_it_cannot_compare_alone(): + """A condition absent from either side keeps a factor of one.""" + cond_train = np.array([0, 0, 0]) + cond_test = np.array([0, 2, 2]) + scale = _holdout_scale(cond_train, cond_test, 3) + np.testing.assert_allclose(scale[cond_test == 2, 0], 1.0) + + +def test_exactly_low_rank_data_scores_near_one_at_its_own_rank(): + """The regression that matters: in-sample ~1 must not come with held-out < 0. + + Before the slice-size correction this returned about -0.03 while the + training block scored 0.9996 -- the model recovered the structure and the + held-out score still said it was worse than predicting the mean. + """ + adata = _exact_low_rank(rank_true=4) + trial = _bicv_trial(adata, 4, 0.5, 0.5, seed=0, tolerance=1e-8, max_iter=300) + + assert trial["Train Block R2X"] > 0.99 + assert trial["BiCV R2X"] > 0.95 + + +@pytest.mark.parametrize("frac", [0.2, 0.35, 0.5, 0.7]) +def test_score_barely_moves_with_the_held_out_fraction(frac): + """`held_out_cell_frac` sets how much is held out, not what the score is. + + The bug made this swing from -2.85 (frac 0.1) to 0.90 (frac 0.5) on the + same data at the same rank, because the error was sqrt(n_train/n_test). + """ + adata = _exact_low_rank(rank_true=4) + trial = _bicv_trial(adata, 4, frac, frac, seed=0, tolerance=1e-8, max_iter=300) + assert trial["BiCV R2X"] > 0.9 + + +def test_bicv_recovers_a_known_rank(): + """The point of the metric: peak (or plateau onset) at the true rank.""" + adata = _exact_low_rank(rank_true=3, noise=0.5, n_genes=60) + results = bicv(adata, [1, 2, 3, 5, 8], n_repeats=2, random_state=0, max_iter=150) + means = results[results["Metric"] == "BiCV R2X"].groupby("Rank")["R2X"].mean() + + # Climbs up to the true rank, then stops climbing. + assert means.loc[3] > means.loc[2] > means.loc[1] + assert means.loc[5] <= means.loc[3] + 1e-3 + assert means.loc[8] <= means.loc[3] + 1e-3 + + +def test_pure_noise_does_not_score_positive(): + """Nothing generalises from noise; a positive score would mean leakage.""" + rng = np.random.default_rng(0) + cond = np.repeat(np.arange(5), 70) + adata = anndata.AnnData(X=rng.normal(size=(cond.size, 60))) + adata.obs["condition_unique_idxs"] = cond + adata.var["means"] = np.zeros(60) + + trial = _bicv_trial(adata, 5, 0.5, 0.5, seed=0, tolerance=1e-8, max_iter=200) + assert trial["BiCV R2X"] < 0.02 + + +def test_held_out_fractions_default_to_one_half(): + """Owen and Perry's recommended split; see the module docstring.""" + import inspect + + params = inspect.signature(bicv).parameters + assert params["held_out_cell_frac"].default == 0.5 + assert params["held_out_gene_frac"].default == 0.5 From 38ddb41f6dc281b6b5babef0d6a87c173f90ed95 Mon Sep 17 00:00:00 2001 From: fishidaho Date: Fri, 11 Sep 2026 18:05:14 -0700 Subject: [PATCH 2/2] cleaing docs --- docs/tutorial.md | 2 +- scrise/rank_selection.py | 31 ++----------------------------- 2 files changed, 3 insertions(+), 30 deletions(-) diff --git a/docs/tutorial.md b/docs/tutorial.md index 4b79d7ce..fbba5aa4 100644 --- a/docs/tutorial.md +++ b/docs/tutorial.md @@ -92,7 +92,7 @@ plt.show() ### Select a Rank with Bi-Cross-Validation (BiCV) -$R^2X$ (in-sample fit) always increases with rank, so it cannot by itself identify an optimal rank — it can only reveal an elbow. Bi-cross-validation (BiCV) addresses this by holding out a random subset of cells *and* genes, fitting RISE on the remaining data, and scoring how well the fit predicts the held-out block. Because it is evaluated on unseen data, BiCV $R^2X$ penalizes overfitting and turns over near the rank that generalizes best, unlike the fit $R^2X$. In practice the curve more often flattens into a plateau than forms a sharp peak, so read the elbow where it stops climbing rather than trusting the single highest point -- past the plateau, neighbouring ranks can differ only in the fourth decimal. +$R^2X$ (in-sample fit) always increases with rank, so it cannot by itself identify an optimal rank. Bi-cross-validation (BiCV) addresses this by holding out a random subset of cells *and* genes, fitting RISE on the remaining data, and scoring how well the fit predicts the held-out block. Because it is evaluated on unseen data, BiCV $R^2X$ penalizes overfitting and turns over near the rank that generalizes best, unlike the fit $R^2X$. `scrise.rank_selection.bicv` exhaustively evaluates a list of candidate ranks, each with several repeated random train/test splits, and returns both the fit $R^2X$ and the BiCV $R^2X$ for every rank so you can inspect the full curve. Since PARAFAC2 fits use CANDELINC compression by default (`compress="auto"`), sweeping every rank in this way is cheap enough that there is no need to search for the best rank without evaluating every candidate. diff --git a/scrise/rank_selection.py b/scrise/rank_selection.py index a87ee5d3..89d77f74 100644 --- a/scrise/rank_selection.py +++ b/scrise/rank_selection.py @@ -8,14 +8,6 @@ predicts the held-out (test-cell x test-gene) block. Unlike the ordinary in-sample fit R2X (which increases monotonically with rank), the BiCV R2X penalizes overfitting and typically peaks near the "true" rank of the data. - -The held-out cells are scored through their own projections, whose scale has -to be corrected for the size of the held-out slice -- see -:func:`_holdout_scale`. Both held-out fractions default to one half, the -split Owen and Perry recommend for bi-cross-validation -(https://arxiv.org/abs/0908.2062); it is also the split at which the -uncorrected score happened to be unbiased, so it is the setting under which -older results are comparable to current ones. """ import warnings @@ -99,22 +91,7 @@ def _cell_loadings( def _holdout_scale( cond_train: np.ndarray, cond_test: np.ndarray, n_cond: int ) -> np.ndarray: - """Per-held-out-cell factor correcting `A` for the held-out slice's size. - - PARAFAC2 writes slice ``k`` as ``P_k (B diag(a_k)) C^T`` with ``P_k`` - orthonormal, so the reconstruction's Frobenius norm is - ``||(B diag(a_k)) C^T||``, independent of how many cells the slice holds: - all of a slice's energy sits in ``a_k``. ``a_k`` is fit against - ``n_train_k`` cells, and a slice's energy grows like ``sqrt(n_k)``, so - reusing it unchanged for ``n_test_k`` held-out cells overstates the - predicted block by ``sqrt(n_train_k / n_test_k)`` -- a factor of two at - the default split, which drove the held-out R2X negative and made the - curve fall as the fit improved. - - Correcting it needs only the cell counts, which is exact when the two - halves of a condition have comparable per-cell energy -- true here, - since the split is random within each condition. - """ + """Per-held-out-cell factor correcting `A` for the held-out slice's size.""" scale = np.ones(cond_test.size) train_counts = np.bincount(cond_train, minlength=n_cond) test_counts = np.bincount(cond_test, minlength=n_cond) @@ -339,11 +316,7 @@ def bicv( held_out_cell_frac : float, optional (default: 0.5) Fraction of cells held out per condition in each BiCV trial. held_out_gene_frac : float, optional (default: 0.5) - Fraction of genes held out in each BiCV trial. Both default to a - half-and-half split, following Owen and Perry, who report that - "in simulated examples we find that a method leaving out half the - rows and half the columns performs well" - (https://arxiv.org/abs/0908.2062). + Fraction of genes held out in each BiCV trial. random_state : int, optional Random seed for reproducibility. tolerance : float, optional (default: 1e-6)