diff --git a/docs/tutorial.md b/docs/tutorial.md index 041fdefe..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 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. 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. @@ -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..89d77f74 100644 --- a/scrise/rank_selection.py +++ b/scrise/rank_selection.py @@ -88,6 +88,19 @@ 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.""" + 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 +244,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 +281,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,9 +313,9 @@ 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) + held_out_gene_frac : float, optional (default: 0.5) Fraction of genes held out in each BiCV trial. random_state : int, optional Random seed for reproducibility. 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