diff --git a/CLAUDE.md b/CLAUDE.md index 8603771..45e43f1 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -51,7 +51,7 @@ AnnData inputs (predicted + real) ### Key Abstractions -- **`MetricsEvaluator`** (`src/cell_eval/_evaluator.py`) — Main programmatic entry point. Validates input AnnData objects, computes differential expression via `pdex`, and orchestrates the metric pipeline. `compute_ceiling()` estimates a per-metric data ceiling (upper bound) from the real data alone: it bootstraps each perturbation (and the control) to 2× cells, splits into two equal halves treated as real/pred, and runs the full pipeline (DE computed in-memory). Exposed via the `run --ceiling` / `--ceiling-seed` CLI flags (additive: writes `ceiling_results.csv` / `agg_ceiling_results.csv`). +- **`MetricsEvaluator`** (`src/cell_eval/_evaluator.py`) — Main programmatic entry point. Validates input AnnData objects, computes differential expression via `pdex`, and orchestrates the metric pipeline. `compute_ceiling()` estimates a per-metric data ceiling (upper bound) from the real data alone: it splits the data into two *disjoint* halves of `n/2` cells (no cell in both), runs the full pipeline (DE computed in-memory) on that self-split, and maps the reliability metrics in the explicit `SB_METRICS` list from half depth to full depth with the Spearman-Brown correction `r'=2r/(1+r)` (all other metrics — error, counts, `clustering_agreement`, `pearson_edistance` — emitted as NaN). Returns `(results, agg_results)`. Exposed via the `run --ceiling` / `--ceiling-seed` CLI flags (additive: writes `ceiling_results.csv` / `agg_ceiling_results.csv`). - **`MetricRegistry`** (`src/cell_eval/metrics/_registry.py`) — Global singleton `metrics_registry`. Metrics are registered with a name, type (`DE` or `ANNDATA_PAIR`), compute function, and best-value indicator. Supports both plain functions and class-based metrics requiring instantiation. diff --git a/README.md b/README.md index 61e8d42..43a9d92 100644 --- a/README.md +++ b/README.md @@ -84,10 +84,23 @@ This will give you metric evaluations for each perturbation individually (`resul #### Data ceiling To estimate the *maximum* achievable score on each metric given the noise inherent in the real -data, pass `--ceiling`. This is computed from the **real data only**: per perturbation (and the -control), its cells are bootstrapped to twice their count and split into two equal halves; one -half plays "real" and the other "prediction", and the full metric suite is run on that self-split. -The result is, per metric, an upper bound on how well any model could score on this dataset. +data, pass `--ceiling`. This is computed from the **real data only**: each perturbation's cells +(and the control's) are split into two *disjoint* halves of `n/2` cells (no cell in both), one half +plays "real" and the other "prediction", and the full metric suite is run on that self-split. Each +reliability metric is then mapped from half depth back to full depth by the analytical +Spearman-Brown correction `r' = 2r/(1+r)`. The result is, per metric, an unbiased upper bound on how +well any model could score on this dataset. + +A disjoint split is used rather than a bootstrap self-split: a bootstrap draws the two halves from +the same cells, so they are not independent, which biases the ceiling in *both* directions (so it is +not a reliable upper bound). The shared cells make the halves agree more than two independent +samples would (inflating it), while the duplicate cells over-call the FDR-gated DE metrics and drag +the recovery metrics (recall / overlap / AUC) down. The disjoint split is unbiased but shallow (each +half `n/2`), which the Spearman-Brown doubling corrects. + +The correction is applied only to a fixed set of reliability metrics (the `SB_METRICS` list in +`_evaluator.py`); every other metric — error metrics, unbounded counts, and reliability metrics +left off that list (`clustering_agreement`, `pearson_edistance`) — is reported as `NaN`. ```bash cell-eval run \ @@ -99,8 +112,8 @@ cell-eval run \ ``` This is *additive*: it writes the normal `results.csv` / `agg_results.csv` **and** -`ceiling_results.csv` / `agg_ceiling_results.csv`. The bootstrap is reproducible via -`--ceiling-seed` (default `0`). From python, call `compute_ceiling` on the evaluator: +`ceiling_results.csv` / `agg_ceiling_results.csv`. The split is reproducible via `--ceiling-seed` +(default `0`). From python, call `compute_ceiling` on the evaluator: ```python ceiling, ceiling_agg = evaluator.compute_ceiling(seed=0) diff --git a/pyproject.toml b/pyproject.toml index c256a3c..fb653db 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "cell-eval" -version = "0.8.1" +version = "0.8.2" description = "Evaluation metrics for single-cell perturbation predictions" readme = "README.md" authors = [ diff --git a/src/cell_eval/_cli/_run.py b/src/cell_eval/_cli/_run.py index b62a124..f73b1d8 100644 --- a/src/cell_eval/_cli/_run.py +++ b/src/cell_eval/_cli/_run.py @@ -100,14 +100,16 @@ def parse_args_run(parser: ap.ArgumentParser): "--ceiling", action="store_true", help="Additionally compute a data ceiling: a real-data-only upper bound on " - "each metric, estimated by bootstrapping the real data against itself. " - "Writes ceiling_results.csv / agg_ceiling_results.csv alongside the normal results.", + "each metric, estimated by splitting the real data into two disjoint halves " + "of n/2 cells and applying the Spearman-Brown correction (2r/(1+r)) to map " + "each reliability metric back to full depth. Writes ceiling_results.csv / " + "agg_ceiling_results.csv alongside the normal results.", ) parser.add_argument( "--ceiling-seed", type=int, default=0, - help="Random seed for the data ceiling bootstrap [default: %(default)s]", + help="Random seed for the data ceiling disjoint split [default: %(default)s]", ) parser.add_argument( "--cpm-filter", diff --git a/src/cell_eval/_evaluator.py b/src/cell_eval/_evaluator.py index c0ac1a6..f8b1b36 100644 --- a/src/cell_eval/_evaluator.py +++ b/src/cell_eval/_evaluator.py @@ -1,7 +1,6 @@ import logging import multiprocessing as mp import os -import warnings from typing import Any, Literal import anndata as ad @@ -19,6 +18,38 @@ logger = logging.getLogger(__name__) +# Metrics that receive the Spearman-Brown ceiling correction (r' = 2r/(1+r)): +# the bounded, higher-is-better reliability metrics for which doubling the depth +# is meaningful and empirically accurate. Every OTHER metric in the ceiling output +# - error metrics, unbounded counts, and reliability metrics where the doubling is +# not trustworthy (e.g. clustering_agreement, pearson_edistance) - is emitted as +# NaN. Edit this set to change which metrics are corrected (names must match the +# metric column names produced by the pipeline). +SB_METRICS = frozenset( + { + "pearson_delta", + "discrimination_score_l1", + "discrimination_score_l2", + "discrimination_score_cosine", + "overlap_at_N", + "overlap_at_50", + "overlap_at_100", + "overlap_at_200", + "overlap_at_500", + "precision_at_N", + "precision_at_50", + "precision_at_100", + "precision_at_200", + "precision_at_500", + "de_spearman_sig", + "de_spearman_lfc_sig", + "de_direction_match", + "de_sig_genes_recall", + "pr_auc", + "roc_auc", + } +) + def _available_cpus() -> int: """Return CPUs the current process is allowed to use. @@ -164,20 +195,34 @@ def compute_ceiling( ) -> tuple[pl.DataFrame, pl.DataFrame]: """Estimate a data ceiling: the maximum achievable score per metric. - Uses the real data only. For each perturbation (and the control) the - cells are bootstrapped to twice their count and split into two equal - halves; one half is treated as "real" and the other as "prediction". - Running the normal metric pipeline on that self-split yields, per metric, - an upper bound on how well any model could score given the noise inherent - in the real data. - - Outputs mirror :meth:`compute` (``ceiling_results.csv`` / - ``agg_ceiling_results.csv``). The bootstrap DE is computed in-memory and - not written to disk. The same ``pdex_kwargs`` and ``allow_discrete`` used - for the main evaluation are reused so the ceiling is directly comparable. + Uses the real data only. Each perturbation's cells (and the control's) are + split into two *disjoint* halves of ``n/2`` cells - no cell in both - and + one half is treated as "real", the other as "prediction". Running the + normal metric pipeline on that self-split measures each metric's + reliability at half depth; the Spearman-Brown correction ``r' = 2r/(1+r)`` + then maps it to full depth (``n``), an unbiased upper bound on how well any + model could score given the noise inherent in the real data. + + A *disjoint* split is used (rather than a bootstrap self-split) because a + bootstrap draws both halves from the same cells, so they are not + independent - which biases the ceiling in both directions: shared cells + make the halves over-agree (inflating it), while duplicate cells over-call + the FDR-gated DE metrics and drag the recovery metrics down. The cost of a + disjoint split is depth (each half is ``n/2``), which the Spearman-Brown + doubling corrects for. + + The correction is applied only to the reliability metrics listed in the + module-level ``SB_METRICS`` set (bounded, higher-is-better, and empirically + well-behaved under doubling). Every other metric - error metrics, unbounded + counts, and reliability metrics left off that list (``clustering_agreement``, + ``pearson_edistance``) - is emitted as ``NaN`` (no defensible ceiling). Outputs + mirror :meth:`compute` (``ceiling_results.csv`` / + ``agg_ceiling_results.csv``); the self-split DE is computed in-memory and + never written. The same ``pdex_kwargs`` and ``allow_discrete`` as the main + evaluation are reused so the ceiling is directly comparable. """ logger.info(f"Computing data ceiling (seed={seed})") - half_real, half_pred = self._bootstrap_halves(seed) + half_real, half_pred = self._disjoint_halves(seed) ceiling_pair = PerturbationAnndataPair( real=half_real, @@ -194,7 +239,7 @@ def compute_ceiling( anndata_pair=ceiling_pair, num_threads=self._num_threads, allow_discrete=self._allow_discrete, - outdir=None, # keep the bootstrap DE in-memory; never persisted + outdir=None, # keep the self-split DE in-memory; never persisted prefix=None, pdex_kwargs=dict(self._pdex_kwargs), ) @@ -208,51 +253,44 @@ def compute_ceiling( pipeline.skip_metrics(skip_metrics) pipeline.compute_de_metrics(ceiling_de) pipeline.compute_anndata_metrics(ceiling_pair) - results = pipeline.get_results() - agg_results = pipeline.get_agg_results() + + # Half-depth self-split scores -> full-depth ceiling via Spearman-Brown. + results = _spearman_brown_correct(pipeline.get_results()) + agg_results = results.drop("perturbation").describe() if write_csv: self._write_results(results, agg_results, basename) return results, agg_results - def _bootstrap_halves(self, seed: int) -> tuple[ad.AnnData, ad.AnnData]: - """Build two same-size bootstrap halves of the real data. + def _disjoint_halves(self, seed: int) -> tuple[ad.AnnData, ad.AnnData]: + """Split the real data into two *disjoint* halves of ``n/2`` cells each. - Resampling is stratified per perturbation (including the control): each - group's ``n`` cells are drawn ``2n`` times with replacement and split - into two halves of ``n`` cells. This guarantees both halves carry every - perturbation plus the control with the same per-perturbation membership - as the real data, so the resulting ``PerturbationAnndataPair`` validates - and the bootstrap DE keeps the same statistical power. + Each perturbation's cells (including the control's) are shuffled and split + without replacement into two halves of ``floor(n/2)`` cells - so no cell + appears in both halves, giving the independence a bootstrap self-split + lacks. Perturbations with fewer than 2 cells cannot be split and are + dropped from both halves. The resulting half depth (``n/2``) is corrected + back to full depth by the Spearman-Brown doubling in + :meth:`compute_ceiling`. """ real = self.anndata_pair.real pert_col = self.anndata_pair.pert_col - rng = np.random.default_rng(seed) - # Group row positions per perturbation in a single pass (`observed=True` - # keeps only perturbations actually present). `.indices` yields positional - # indices, so they can index the AnnData directly. - half_real_idx: list[np.ndarray] = [] - half_pred_idx: list[np.ndarray] = [] + a_idx: list[np.ndarray] = [] + b_idx: list[np.ndarray] = [] for _pert, idx in real.obs.groupby(pert_col, observed=True).indices.items(): - draws = rng.choice(idx, size=2 * idx.size, replace=True) - half_real_idx.append(draws[: idx.size]) - half_pred_idx.append(draws[idx.size :]) - - # Sampling with replacement duplicates obs names; anndata warns about the - # non-unique index on slice, so silence that one known-benign warning and - # make the names unique immediately afterwards. - with warnings.catch_warnings(): - warnings.filterwarnings( - "ignore", message="Observation names are not unique" - ) - half_real = real[np.concatenate(half_real_idx)].copy() - half_pred = real[np.concatenate(half_pred_idx)].copy() - half_real.obs_names_make_unique() - half_pred.obs_names_make_unique() - + perm = rng.permutation(np.asarray(idx)) + h = perm.size // 2 + if h < 1: + continue # < 2 cells: cannot form two disjoint halves + a_idx.append(perm[:h]) + b_idx.append(perm[h : 2 * h]) + + # Disjoint split has no duplicate rows, so obs names stay unique. + half_real = real[np.concatenate(a_idx)].copy() + half_pred = real[np.concatenate(b_idx)].copy() return half_real, half_pred def _write_results( @@ -281,6 +319,27 @@ def _write_results( agg_results.write_csv(agg_outpath) +def _spearman_brown_correct(results: pl.DataFrame) -> pl.DataFrame: + """Map half-depth self-split scores to the full-depth ceiling. + + Applies the Spearman-Brown prophecy ``r' = 2r/(1+r)`` - the reliability of a + test of doubled length - to the reliability metrics listed in ``SB_METRICS``. + Every other column (error metrics, unbounded counts, and reliability metrics + not in that set) is emitted as ``NaN``, since a Spearman-Brown ceiling has no + defensible meaning there. + """ + nan = float("nan") + exprs: list[pl.Expr] = [] + for col in results.columns: + if col == "perturbation": + continue + if col in SB_METRICS: + exprs.append((2.0 * pl.col(col) / (1.0 + pl.col(col))).alias(col)) + else: + exprs.append(pl.lit(nan).alias(col)) + return results.with_columns(exprs) if exprs else results + + def _build_anndata_pair( real: ad.AnnData | str, pred: ad.AnnData | str, diff --git a/tests/test_ceiling.py b/tests/test_ceiling.py new file mode 100644 index 0000000..f83d68d --- /dev/null +++ b/tests/test_ceiling.py @@ -0,0 +1,113 @@ +import shutil + +import numpy as np +import polars as pl +import pytest + +from cell_eval import MetricsEvaluator +from cell_eval._evaluator import SB_METRICS, _spearman_brown_correct +from cell_eval.data import CONTROL_VAR, PERT_COL, build_random_anndata + +OUTDIR = "TEST_OUTPUT_CEILING" + + +def test_spearman_brown_doubling_on_reliability_metrics(): + """SB doubling r' = 2r/(1+r) is applied to reliability (best_value == ONE) + metric columns.""" + df = pl.DataFrame( + { + "perturbation": ["a", "b"], + "pearson_delta": [0.5, 1.0 / 3.0], + "overlap_at_N": [0.6, 0.2], + } + ) + out = _spearman_brown_correct(df) + # 2*0.5/(1+0.5)=2/3 ; 2*(1/3)/(1+1/3)=0.5 + assert out["pearson_delta"].to_list() == pytest.approx([2 / 3, 0.5], abs=1e-6) + # 2*0.6/1.6=0.75 ; 2*0.2/1.2=1/3 + assert out["overlap_at_N"].to_list() == pytest.approx([0.75, 1 / 3], abs=1e-6) + + +def test_non_reliability_and_excluded_metrics_are_nan(): + """Error metrics, unbounded counts, and the excluded reliabilities + (clustering_agreement, pearson_edistance) are emitted as NaN.""" + df = pl.DataFrame( + { + "perturbation": ["a"], + "pearson_delta": [0.5], # reliability -> SB + "mse": [0.01], # error metric -> NaN + "mae": [0.02], # error metric -> NaN + "de_nsig_counts_real": [10.0], # unbounded count -> NaN + "clustering_agreement": [0.4], # excluded reliability -> NaN + "pearson_edistance": [0.9], # excluded reliability -> NaN + } + ) + out = _spearman_brown_correct(df) + assert out["pearson_delta"][0] == pytest.approx(2 / 3, abs=1e-6) + for col in ( + "mse", + "mae", + "de_nsig_counts_real", + "clustering_agreement", + "pearson_edistance", + ): + assert np.isnan(out[col][0]), col + # SB_METRICS is the explicit inclusion list; the excluded ones are absent + assert "pearson_delta" in SB_METRICS + assert "clustering_agreement" not in SB_METRICS + assert "pearson_edistance" not in SB_METRICS + + +def test_disjoint_halves_share_no_cells(): + adata_real = build_random_anndata() + evaluator = MetricsEvaluator( + adata_pred=adata_real.copy(), + adata_real=adata_real, + control_pert=CONTROL_VAR, + pert_col=PERT_COL, + outdir=OUTDIR, + skip_de=True, + ) + half_real, half_pred = evaluator._disjoint_halves(seed=0) + + # disjoint: no original cell (by name) appears in both halves + assert set(half_real.obs_names).isdisjoint(set(half_pred.obs_names)) + # both halves carry the control + every (splittable) perturbation + assert CONTROL_VAR in set(half_real.obs[PERT_COL].astype(str)) + assert set(half_real.obs[PERT_COL].astype(str)) == set( + half_pred.obs[PERT_COL].astype(str) + ) + shutil.rmtree(OUTDIR) + + +def test_compute_ceiling_end_to_end(): + """compute_ceiling returns (results, agg): reliability metrics are SB-corrected + and bounded in [0, 1]; error metrics come back as NaN.""" + adata_real = build_random_anndata() + evaluator = MetricsEvaluator( + adata_pred=adata_real.copy(), + adata_real=adata_real, + control_pert=CONTROL_VAR, + pert_col=PERT_COL, + outdir=OUTDIR, + skip_de=True, + ) + results, agg = evaluator.compute_ceiling( + profile="anndata", write_csv=False, break_on_error=True + ) + assert results.height > 0 + assert "perturbation" in results.columns + + assert "pearson_delta" in results.columns + pv = results["pearson_delta"].drop_nulls().to_numpy() + assert np.all(pv <= 1.0 + 1e-9) # SB doubling can never exceed 1 + + # error metrics are emitted as NaN (no defensible SB ceiling) + for col in ("mse", "mae"): + if col in results.columns: + assert np.all(np.isnan(results[col].to_numpy())) + # excluded reliability metrics are NaN too + for col in ("clustering_agreement", "pearson_edistance"): + if col in results.columns: + assert np.all(np.isnan(results[col].to_numpy())) + shutil.rmtree(OUTDIR) diff --git a/tests/test_eval.py b/tests/test_eval.py index 9e094aa..2684031 100644 --- a/tests/test_eval.py +++ b/tests/test_eval.py @@ -405,32 +405,6 @@ def _assert_results_close(a, b) -> None: assert np.allclose(num_a, num_b, equal_nan=True) -def test_ceiling_bootstrap_halves_membership(): - """Each half must carry every perturbation + control with the real counts.""" - adata_real = build_random_anndata() - evaluator = MetricsEvaluator( - adata_pred=adata_real.copy(), - adata_real=adata_real, - control_pert=CONTROL_VAR, - pert_col=PERT_COL, - outdir=OUTDIR, - skip_de=True, # membership only depends on the bootstrap, skip pdex - ) - - half_real, half_pred = evaluator._bootstrap_halves(seed=0) - - real_counts = evaluator.anndata_pair.real.obs[PERT_COL].value_counts().to_dict() - assert CONTROL_VAR in real_counts - for half in (half_real, half_pred): - half_counts = half.obs[PERT_COL].value_counts().to_dict() - # same set of perturbations (incl. control) and same per-pert membership - assert half_counts == real_counts - # sampling with replacement must yield unique obs names - assert half.obs_names.is_unique - - shutil.rmtree(OUTDIR) - - def test_eval_ceiling(): adata_real = build_random_anndata() adata_pred = adata_real.copy() @@ -441,9 +415,13 @@ def test_eval_ceiling(): pert_col=PERT_COL, outdir=OUTDIR, ) - results, agg_results = evaluator.compute_ceiling(break_on_error=True) + results, agg = evaluator.compute_ceiling(break_on_error=True) assert results.height > 0 - assert agg_results.height > 0 + assert "perturbation" in results.columns + # a reliability metric is SB-corrected; the doubling can never exceed 1 + assert "pearson_delta" in results.columns + pv = results["pearson_delta"].drop_nulls().to_numpy() + assert np.all(pv <= 1.0 + 1e-9) assert os.path.exists(f"{OUTDIR}/ceiling_results.csv") assert os.path.exists(f"{OUTDIR}/agg_ceiling_results.csv") shutil.rmtree(OUTDIR) @@ -462,7 +440,6 @@ def test_eval_ceiling_prefix(): ) evaluator.compute_ceiling(break_on_error=True) assert os.path.exists(f"{OUTDIR}/arbitrary_ceiling_results.csv") - assert os.path.exists(f"{OUTDIR}/arbitrary_agg_ceiling_results.csv") shutil.rmtree(OUTDIR) @@ -498,7 +475,7 @@ def test_eval_ceiling_pds_skips_de(): skip_de=True, ) assert evaluator.de_comparison is None - results, _ = evaluator.compute_ceiling( + results, _agg = evaluator.compute_ceiling( profile="pds", break_on_error=True, write_csv=False, @@ -520,7 +497,15 @@ def test_eval_ceiling_reproducible(): ) r1, _ = evaluator.compute_ceiling(seed=7, write_csv=False, break_on_error=True) r2, _ = evaluator.compute_ceiling(seed=7, write_csv=False, break_on_error=True) - _assert_results_close(r1, r2) + r1 = r1.sort("perturbation") + r2 = r2.sort("perturbation") + assert r1["perturbation"].to_list() == r2["perturbation"].to_list() + for col in (c for c in r1.columns if c != "perturbation"): + np.testing.assert_allclose( + r1[col].to_numpy().astype(float), + r2[col].to_numpy().astype(float), + equal_nan=True, + ) shutil.rmtree(OUTDIR) @@ -551,9 +536,8 @@ def test_eval_ceiling_does_not_clobber_de(): with open(pred_de_path, "rb") as fh: assert fh.read() == before_pred - # ceiling outputs exist, but no ceiling DE artifacts are written + # ceiling output exists, but no ceiling DE artifacts are written assert os.path.exists(f"{OUTDIR}/ceiling_results.csv") - assert os.path.exists(f"{OUTDIR}/agg_ceiling_results.csv") assert not os.path.exists(f"{OUTDIR}/ceiling_real_de.csv") assert not os.path.exists(f"{OUTDIR}/ceiling_pred_de.csv") shutil.rmtree(OUTDIR)