diff --git a/.gitignore b/.gitignore index 3e77b28..f16e0f6 100644 --- a/.gitignore +++ b/.gitignore @@ -25,3 +25,5 @@ docs/notebooks/data/ # IDEs /.idea/ +/docs/superpowers/ +/.superpowers/ diff --git a/CHANGELOG.md b/CHANGELOG.md index 4437db9..2e3fea4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,15 @@ and this project adheres to [Semantic Versioning][]. [keep a changelog]: https://keepachangelog.com/en/1.0.0/ [semantic versioning]: https://semver.org/spec/v2.0.0.html +## 0.6.1 + +### Added + +- Add `scib_metrics.perturbation`: mandatory mean/additive/linear baseline predictors and an + evaluation surface (delta correlation, DE-gene rank recovery, Systema-style shared/specific + decomposition, combination-additivity) for perturbation-response prediction, built on + `pertpy` (optional `scib-metrics[perturbation]` extra). + ## 0.6.0 (2026-07-28) ### Added diff --git a/docs/api.md b/docs/api.md index b117ef2..6ba874f 100644 --- a/docs/api.md +++ b/docs/api.md @@ -20,6 +20,33 @@ from scib_metrics.benchmark import Benchmarker benchmark.BatchCorrection ``` +## Perturbation-prediction evaluation + +Import as: + +``` +from scib_metrics.perturbation import PerturbationBenchmarker +``` + +```{eval-rst} +.. module:: scib_metrics.perturbation +.. currentmodule:: scib_metrics + +.. autosummary:: + :toctree: generated + + perturbation.MeanBaseline + perturbation.AdditiveBaseline + perturbation.LinearBaseline + perturbation.PerturbationBaselines + perturbation.PerturbationMetrics + perturbation.PerturbationBenchmarker + perturbation.delta_correlation + perturbation.de_rank_recovery + perturbation.systema_decomposition + perturbation.combination_additivity +``` + ## Metrics Import as: diff --git a/pyproject.toml b/pyproject.toml index dd131da..484b323 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,7 @@ dependencies = [ "umap-learn>=0.5", ] optional-dependencies.parallel = [ "joblib" ] +optional-dependencies.perturbation = [ "pertpy" ] optional-dependencies.tutorial = [ "adjusttext", # missing liger dependency "goatools", # missing liger dependency @@ -60,6 +61,7 @@ test = [ "coverage>=7.10", "harmonypy", "joblib", + "pertpy", "pytest", "pytest-cov", # For VS Code's coverage functionality "scib>=1.1.4", diff --git a/src/scib_metrics/perturbation/__init__.py b/src/scib_metrics/perturbation/__init__.py new file mode 100644 index 0000000..d5e6551 --- /dev/null +++ b/src/scib_metrics/perturbation/__init__.py @@ -0,0 +1,27 @@ +from scib_metrics.perturbation._baselines import ( + AdditiveBaseline, + BasePerturbationPredictor, + LinearBaseline, + MeanBaseline, +) +from scib_metrics.perturbation._core import PerturbationBaselines, PerturbationBenchmarker, PerturbationMetrics +from scib_metrics.perturbation._metrics import ( + combination_additivity, + de_rank_recovery, + delta_correlation, + systema_decomposition, +) + +__all__ = [ + "AdditiveBaseline", + "BasePerturbationPredictor", + "LinearBaseline", + "MeanBaseline", + "PerturbationBaselines", + "PerturbationBenchmarker", + "PerturbationMetrics", + "combination_additivity", + "de_rank_recovery", + "delta_correlation", + "systema_decomposition", +] diff --git a/src/scib_metrics/perturbation/_baselines.py b/src/scib_metrics/perturbation/_baselines.py new file mode 100644 index 0000000..e967799 --- /dev/null +++ b/src/scib_metrics/perturbation/_baselines.py @@ -0,0 +1,220 @@ +"""Naive baseline predictors for perturbation-response evaluation. + +All predictors operate in delta-space: `.predict()` returns the predicted expression +*delta relative to control*, not an absolute expression profile. This matches +`pertpy.tools.PerturbationSpace.compute_control_diff`'s convention and the metric +functions in `scib_metrics.perturbation._metrics`, which all consume deltas. +""" + +from __future__ import annotations + +import warnings +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any, Self + +import numpy as np +from sklearn.linear_model import Ridge + +from scib_metrics.perturbation._utils import import_pertpy + +if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + + from anndata import AnnData + + from scib_metrics._types import NdArray + + +class BasePerturbationPredictor(ABC): + """Base class for perturbation-response baseline predictors.""" + + @abstractmethod + def fit( + self, + adata_train: AnnData, + target_col: str = "perturbation", + reference_key: str = "control", + perturbation_encodings: Mapping[str, NdArray] | None = None, + ) -> Self: + """Fit the baseline on training data. + + Parameters + ---------- + adata_train + Cell-level AnnData. `adata_train.obs[target_col]` holds the perturbation label + of each cell; `reference_key` marks control cells. + target_col + `.obs` column name holding the perturbation label. + reference_key + Perturbation label marking control cells. + perturbation_encodings + Optional feature vector per perturbation name. Required by `LinearBaseline`, + ignored by `MeanBaseline` and `AdditiveBaseline`. + + Returns + ------- + `self`. + """ + + @abstractmethod + def predict(self, perturbations: Sequence[str]) -> NdArray: + """Predict expression deltas (relative to control) for the requested perturbations. + + Parameters + ---------- + perturbations + Names of the perturbations to predict for. + + Returns + ------- + Array of shape `(len(perturbations), n_genes)`. + """ + + +def _pseudobulk_control_diff(pt, adata: AnnData, target_col: str, reference_key: str) -> tuple[AnnData, Any]: + """Pseudobulk `adata` by `target_col` (mean mode) and subtract the control mean in place. + + Note: `pertpy.tools.PerturbationSpace` is not exported at the `pt.tl` top level (it exists + only as the internal base class of `PseudobulkSpace`/`CentroidSpace`/etc.) — always call + `compute_control_diff`/`add`/`subtract` on a `PseudobulkSpace` (or other concrete space) + instance, never `pt.tl.PerturbationSpace()` directly (that raises `AttributeError`). + + Returns + ------- + Tuple of `(diffed pseudobulk AnnData, the PseudobulkSpace instance used)` — callers that + also need `.add()` (e.g. `AdditiveBaseline`) reuse the same instance rather than creating + a second one. + """ + ps = pt.tl.PseudobulkSpace() + pseudobulk = ps.compute(adata, target_col=target_col, mode="mean") + ps.compute_control_diff(pseudobulk, target_col=target_col, reference_key=reference_key, copy=False) + return pseudobulk, ps + + +class MeanBaseline(BasePerturbationPredictor): + """Predicts the mean training-perturbation delta for every requested perturbation. + + Deliberately naive: this is the "did your model beat just guessing the average + perturbed profile" floor from the perturbation-evaluation literature. It is blind to + which perturbation is being requested. + """ + + def fit( + self, + adata_train: AnnData, + target_col: str = "perturbation", + reference_key: str = "control", + perturbation_encodings: Mapping[str, NdArray] | None = None, + ) -> Self: + pt = import_pertpy() + pseudobulk, _ = _pseudobulk_control_diff(pt, adata_train, target_col, reference_key) + is_control = pseudobulk.obs[target_col].astype(str).to_numpy() == reference_key + deltas = np.asarray(pseudobulk.X)[~is_control] + if deltas.shape[0] == 0: + raise ValueError(f"No non-control perturbations found in `adata_train.obs[{target_col!r}]`.") + self.mean_delta_ = deltas.mean(axis=0) + return self + + def predict(self, perturbations: Sequence[str]) -> NdArray: + return np.tile(self.mean_delta_, (len(perturbations), 1)) + + +class AdditiveBaseline(BasePerturbationPredictor): + """Predicts combination perturbations additively from their trained single components. + + For a requested perturbation whose name decomposes via `sep` into components that were + each present as a training perturbation (e.g. `"A+B"` when `A` and `B` were trained on), + predicts `delta(A) + delta(B)` via `pertpy.tools.PerturbationSpace.add`. For a requested + perturbation with no such decomposition, falls back to the same global mean-delta + behavior as `MeanBaseline` — this is the expected, literature-standard degenerate case: + predicting a single unseen perturbation from a global average shift is mathematically + identical to `MeanBaseline` when there is no combination structure to exploit. + """ + + def __init__(self, sep: str = "+") -> None: + self.sep = sep + + def fit( + self, + adata_train: AnnData, + target_col: str = "perturbation", + reference_key: str = "control", + perturbation_encodings: Mapping[str, NdArray] | None = None, + ) -> Self: + pt = import_pertpy() + self._diffed_adata, self._ps = _pseudobulk_control_diff(pt, adata_train, target_col, reference_key) + self._target_col = target_col + self._reference_key = reference_key + is_control = self._diffed_adata.obs[target_col].astype(str).to_numpy() == reference_key + deltas = np.asarray(self._diffed_adata.X)[~is_control] + if deltas.shape[0] == 0: + raise ValueError(f"No non-control perturbations found in `adata_train.obs[{target_col!r}]`.") + self.mean_delta_ = deltas.mean(axis=0) + return self + + def predict(self, perturbations: Sequence[str]) -> NdArray: + trained_names = set(self._diffed_adata.obs_names.astype(str)) + rows = [] + for perturbation in perturbations: + components = perturbation.split(self.sep) + if len(components) > 1 and all(c in trained_names for c in components): + # `ensure_consistency=False` is correct here: `self._diffed_adata` was already + # differenced against control in `fit()`. pertpy's `.add()` still warns about this + # (it can't see that we pre-diffed) — the warning is a known false positive for + # our exact usage, so it's suppressed rather than left to spam every `.predict()` call. + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + combined = self._ps.add( + self._diffed_adata, + perturbations=components, + reference_key=self._reference_key, + ensure_consistency=False, + target_col=self._target_col, + ) + rows.append(np.asarray(combined.X)[-1]) + else: + rows.append(self.mean_delta_) + return np.stack(rows, axis=0) + + +class LinearBaseline(BasePerturbationPredictor): + """Ridge regression from a perturbation encoding to its expression delta. + + Unlike `MeanBaseline` and `AdditiveBaseline`, this baseline can generalize to a held-out + perturbation with no training-set relationship to any trained perturbation, as long as a + feature encoding is supplied for it. + """ + + def __init__(self, **ridge_kwargs) -> None: + self.ridge_kwargs = ridge_kwargs + + def fit( + self, + adata_train: AnnData, + target_col: str = "perturbation", + reference_key: str = "control", + perturbation_encodings: Mapping[str, NdArray] | None = None, + ) -> Self: + if not perturbation_encodings: + raise ValueError("`LinearBaseline` requires `perturbation_encodings`.") + pt = import_pertpy() + pseudobulk, _ = _pseudobulk_control_diff(pt, adata_train, target_col, reference_key) + names = pseudobulk.obs_names.astype(str).tolist() + encodings, deltas = [], [] + for name, delta in zip(names, np.asarray(pseudobulk.X), strict=True): + if name == reference_key or name not in perturbation_encodings: + continue + encodings.append(np.asarray(perturbation_encodings[name])) + deltas.append(delta) + if not encodings: + raise ValueError("None of the training perturbations have a matching entry in `perturbation_encodings`.") + self.model_ = Ridge(**self.ridge_kwargs).fit(np.stack(encodings), np.stack(deltas)) + self.perturbation_encodings_ = perturbation_encodings + return self + + def predict(self, perturbations: Sequence[str]) -> NdArray: + missing = [p for p in perturbations if p not in self.perturbation_encodings_] + if missing: + raise ValueError(f"No encoding supplied for perturbations: {missing}.") + encodings = np.stack([np.asarray(self.perturbation_encodings_[p]) for p in perturbations]) + return self.model_.predict(encodings) diff --git a/src/scib_metrics/perturbation/_core.py b/src/scib_metrics/perturbation/_core.py new file mode 100644 index 0000000..dbfb877 --- /dev/null +++ b/src/scib_metrics/perturbation/_core.py @@ -0,0 +1,381 @@ +"""Orchestrator for the perturbation-prediction evaluation surface.""" + +from __future__ import annotations + +import os +import warnings +from dataclasses import dataclass, fields +from typing import TYPE_CHECKING, Any + +import anndata +import matplotlib as mpl +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +from anndata import AnnData +from plottable import ColumnDefinition, Table +from plottable.cmap import normed_cmap +from sklearn.preprocessing import MinMaxScaler + +from scib_metrics.perturbation._baselines import AdditiveBaseline, LinearBaseline, MeanBaseline +from scib_metrics.perturbation._metrics import ( + combination_additivity, + de_rank_recovery, + delta_correlation, + systema_decomposition, +) +from scib_metrics.perturbation._utils import import_pertpy + +if TYPE_CHECKING: + from scib_metrics._types import NdArray + +Kwargs = dict[str, Any] +MetricType = bool | Kwargs + +_BASELINE_CLASSES = {"mean": MeanBaseline, "additive": AdditiveBaseline, "linear": LinearBaseline} + + +@dataclass(frozen=True) +class PerturbationBaselines: + """Specification of which baseline predictors to run in the pipeline. + + Baselines can be included using a boolean flag. Custom keyword args (passed to the + predictor's constructor) can be used by passing a dictionary here. + """ + + mean: MetricType = True + additive: MetricType = True + linear: MetricType = True + + +@dataclass(frozen=True) +class PerturbationMetrics: + """Specification of which perturbation-evaluation metrics to run in the pipeline.""" + + delta_correlation: MetricType = True + de_rank_recovery: MetricType = False + systema_decomposition: MetricType = True + combination_additivity: MetricType = True + ground_truth_significance: bool = False + + +class PerturbationBenchmarker: + """Benchmarking pipeline for perturbation-response prediction. + + Runs mandatory naive baselines (mean / additive / linear) and, optionally, the user's own + model predictions, through a shared set of evaluation metrics, so every run carries an + honest floor. + + Parameters + ---------- + adata_train + Cell-level AnnData. `adata_train.obs[target_col]` holds the perturbation label of each + cell; `reference_key` marks control cells. + adata_test + Held-out cell-level AnnData (not pre-averaged), `adata_test.obs[target_col]` + identifying which held-out perturbation each row belongs to. + predictions + The user's own model predictions to compare against the baselines, as + `{model_name: deltas}` where `deltas` has shape `(n_held_out_perturbations, n_genes)` + aligned to the sorted, deduplicated, non-control values of + `adata_test.obs[target_col]`. + target_col + `.obs` column name holding the perturbation label, in both `adata_train` and `adata_test`. + reference_key + Perturbation label marking control cells, in both `adata_train` and `adata_test`. + perturbation_encodings + Optional feature vector per perturbation name, covering both trained and held-out + names. Required if `baselines.linear` is enabled. + true_de_gene_indices + Optional true top-DE gene indices per held-out perturbation. Required if + `metrics.de_rank_recovery` is enabled. + baselines + Specification of which baseline predictors to run. + metrics + Specification of which metrics to run. + """ + + def __init__( + self, + adata_train: AnnData, + adata_test: AnnData, + predictions: dict[str, NdArray] | None = None, + target_col: str = "perturbation", + reference_key: str = "control", + perturbation_encodings: dict[str, NdArray] | None = None, + true_de_gene_indices: dict[str, NdArray] | None = None, + baselines: PerturbationBaselines = PerturbationBaselines(), + metrics: PerturbationMetrics = PerturbationMetrics(), + ) -> None: + self.adata_train = adata_train + self.adata_test = adata_test + self.predictions = predictions or {} + self.target_col = target_col + self.reference_key = reference_key + self.perturbation_encodings = perturbation_encodings + self.true_de_gene_indices = true_de_gene_indices + self.baselines = baselines + self.metrics = metrics + self._results: pd.DataFrame | None = None + self._combination_additivity: pd.DataFrame | None = None + self._ground_truth_significance: pd.DataFrame | None = None + + def benchmark(self) -> None: + """Fit and run every enabled baseline and the user's predictions through every enabled metric.""" + held_out = [ + p for p in sorted(self.adata_test.obs[self.target_col].astype(str).unique()) if p != self.reference_key + ] + if not held_out: + raise ValueError(f"No held-out perturbations found in `adata_test.obs[{self.target_col!r}]`.") + + baseline_names, predicted_deltas = [], {} + for field in fields(self.baselines): + flag = getattr(self.baselines, field.name) + if not flag: + continue + kwargs = flag if isinstance(flag, dict) else {} + predictor = _BASELINE_CLASSES[field.name](**kwargs) + predictor.fit( + self.adata_train, + target_col=self.target_col, + reference_key=self.reference_key, + perturbation_encodings=self.perturbation_encodings, + ) + predicted_deltas[field.name] = predictor.predict(held_out) + baseline_names.append(field.name) + + # Check for collisions between user predictions and enabled baselines + collisions = set(self.predictions) & set(baseline_names) + if collisions: + raise ValueError( + f"`predictions` keys collide with enabled baseline names: {sorted(collisions)}. " + "Rename your prediction(s) or disable the corresponding baseline(s)." + ) + + for name, preds in self.predictions.items(): + predicted_deltas[name] = np.asarray(preds) + + true_deltas = self._compute_true_deltas(held_out) + gene_indices = [self.true_de_gene_indices[p] for p in held_out] if self.true_de_gene_indices else None + # Only restrict `delta_correlation` to `true_de_gene_indices` when the caller actually + # opted into `de_rank_recovery`; merely supplying `true_de_gene_indices` (e.g. because + # `de_rank_recovery` is wanted) should not silently change what `delta_correlation` means. + delta_corr_gene_indices = gene_indices if self.metrics.de_rank_recovery else None + + run_systema_decomposition = self.metrics.systema_decomposition and len(held_out) >= 2 + if self.metrics.systema_decomposition and not run_systema_decomposition: + warnings.warn( + "`metrics.systema_decomposition` is enabled but fewer than 2 held-out perturbations " + f"are present ({len(held_out)} found); skipping `systema_decomposition`.", + stacklevel=2, + ) + + rows: dict[str, dict[str, float | bool]] = {} + for name, preds in predicted_deltas.items(): + row: dict[str, float | bool] = {"is_baseline": name in baseline_names} + if self.metrics.delta_correlation: + kwargs = self.metrics.delta_correlation if isinstance(self.metrics.delta_correlation, dict) else {} + row["delta_correlation"] = delta_correlation( + preds, true_deltas, gene_indices=delta_corr_gene_indices, **kwargs + )["mean"] + if self.metrics.de_rank_recovery: + if not self.true_de_gene_indices: + raise ValueError("`metrics.de_rank_recovery` requires `true_de_gene_indices`.") + kwargs = self.metrics.de_rank_recovery if isinstance(self.metrics.de_rank_recovery, dict) else {} + k = kwargs.get("k", 50) + row["de_rank_recovery"] = de_rank_recovery(preds, gene_indices, k=k)["mean"] + if run_systema_decomposition: + decomposition = systema_decomposition(preds, true_deltas) + row["systema_shared"] = decomposition["shared"] + row["systema_specific"] = decomposition["specific"] + rows[name] = row + self._results = pd.DataFrame.from_dict(rows, orient="index") + + if self.metrics.combination_additivity: + kwargs = ( + self.metrics.combination_additivity if isinstance(self.metrics.combination_additivity, dict) else {} + ) + pt = import_pertpy() + combined = anndata.concat([self.adata_train, self.adata_test]) + pseudobulk = pt.tl.PseudobulkSpace().compute(combined, target_col=self.target_col, mode="mean") + self._combination_additivity = combination_additivity( + pseudobulk, target_col=self.target_col, reference_key=self.reference_key, **kwargs + ) + + if self.metrics.ground_truth_significance: + self._ground_truth_significance = self._compute_ground_truth_significance() + + def _compute_true_deltas(self, held_out: list[str]) -> NdArray: + pt = import_pertpy() + test_labels = self.adata_test.obs[self.target_col].astype(str) + if self.reference_key in test_labels.unique(): + source = self.adata_test + else: + # `adata_test` is documented (see the class docstring) as holding only the + # held-out perturbation(s), with no requirement that it carry its own control + # cells — and the reference synthetic train/test split (test = only the "A+B" + # cells) confirms this is the expected shape. But `compute_control_diff` needs a + # `reference_key` group to diff against; pseudobulking `adata_test` alone in that + # case raises `ValueError: Reference key control not found in perturbation` + # (confirmed while running this task's tests). Borrow control cells from + # `adata_train`, which always has them (every baseline's `.fit()` requires it), + # and diff the held-out perturbations against that instead. + train_labels = self.adata_train.obs[self.target_col].astype(str) + is_control = train_labels.to_numpy() == self.reference_key + if not is_control.any(): + raise ValueError( + f"No {self.reference_key!r} cells found in `adata_test` or `adata_train`; " + "cannot compute true deltas." + ) + source = anndata.concat([self.adata_train[is_control], self.adata_test]) + ps = pt.tl.PseudobulkSpace() + pseudobulk = ps.compute(source, target_col=self.target_col, mode="mean") + ps.compute_control_diff(pseudobulk, target_col=self.target_col, reference_key=self.reference_key, copy=False) + obs_names = pseudobulk.obs_names.astype(str) + return np.stack([np.asarray(pseudobulk.X)[obs_names.get_loc(p)] for p in held_out]) + + def _compute_ground_truth_significance(self) -> pd.DataFrame: + test_labels = self.adata_test.obs[self.target_col].astype(str) + if self.reference_key not in test_labels.unique(): + warnings.warn( + f"No {self.reference_key!r} cells found in `adata_test`; skipping the ground-truth " + "significance diagnostic.", + stacklevel=2, + ) + return pd.DataFrame(columns=["distance", "pvalue", "significant", "pvalue_adj", "significant_adj"]) + pt = import_pertpy() + # `Distance`/`DistanceTest`'s AnnData-based methods read from `.obsm["X_pca"]` by default + # when neither `layer_key` nor `obsm_key` is given (confirmed against the installed pertpy: + # omitting both raises `KeyError: 'X_pca'` on an AnnData with only `.X` populated). Route + # through a layer instead so this works on raw/normalized expression directly. + adata_for_test = self.adata_test.copy() + adata_for_test.layers["_ground_truth_significance_expression"] = adata_for_test.X + distance_test = pt.tl.DistanceTest(metric="edistance", layer_key="_ground_truth_significance_expression") + return distance_test(adata_for_test, groupby=self.target_col, contrast=self.reference_key) + + def get_combination_additivity(self) -> pd.DataFrame: + """Return the combination-additivity diagnostic. + + Scores how well an additive model predicts combination perturbations (e.g. `"A+B"`) + from their singles (`"A"`, `"B"`), via `pertpy.tools.PerturbationSpace.evaluate_combinations` + over the pseudobulked union of `adata_train` and `adata_test`. + + Returns + ------- + DataFrame indexed by combination name with `"distance"`, `"predicted_magnitude"` and + `"measured_magnitude"` columns. + """ + if self._combination_additivity is None: + raise RuntimeError( + "Combination additivity was not computed. Set " + "`metrics=PerturbationMetrics(combination_additivity=True)` (the default) and call " + "`.benchmark()`." + ) + return self._combination_additivity + + def get_results(self, min_max_scale: bool = False) -> pd.DataFrame: + """Return the benchmarking results. + + Parameters + ---------- + min_max_scale + Whether to min-max scale the score columns (excludes `is_baseline`). + + Returns + ------- + DataFrame indexed by predictor name, with an `is_baseline` flag column and one column + per enabled metric. + """ + if self._results is None: + raise RuntimeError("Call `.benchmark()` before `.get_results()`.") + results = self._results.copy() + if min_max_scale: + score_cols = [c for c in results.columns if c != "is_baseline"] + results[score_cols] = MinMaxScaler().fit_transform(results[score_cols]) + return results + + def get_ground_truth_significance(self) -> pd.DataFrame: + """Return the ground-truth significance diagnostic. + + Tests whether each held-out perturbation's *true* cells are significantly different + from control, via `pertpy.tools.DistanceTest`. This scores the test data itself, not + any predictor — use it to filter out held-out perturbations with no real signal before + trusting scores against them. + + Returns + ------- + DataFrame with `"distance"`, `"pvalue"`, `"significant"`, `"pvalue_adj"` and + `"significant_adj"` columns, indexed by perturbation name. + """ + if self._ground_truth_significance is None: + raise RuntimeError( + "Ground-truth significance was not computed. Set " + "`metrics=PerturbationMetrics(ground_truth_significance=True)` and call `.benchmark()`." + ) + return self._ground_truth_significance + + def plot_results_table(self, min_max_scale: bool = False, show: bool = True, save_dir: str | None = None) -> Table: + """Plot the benchmarking results as a table, with baseline rows visually distinguished. + + Parameters + ---------- + min_max_scale + Whether to min-max scale the score columns. + show + Whether to show the plot. + save_dir + Directory to save the plot to. If `None`, the plot is not saved. + + Returns + ------- + The `plottable.Table` instance. + """ + df = self.get_results(min_max_scale=min_max_scale) + is_baseline = df["is_baseline"] + plot_df = df.drop(columns="is_baseline").astype(np.float64) + plot_df["Predictor"] = [f"{name} (baseline)" if is_baseline[name] else str(name) for name in plot_df.index] + + cmap_fn = lambda col_data: normed_cmap(col_data, cmap=mpl.cm.PRGn, num_stds=2.5) + score_cols = [c for c in plot_df.columns if c != "Predictor"] + # "Predictor" values can be as long as " (baseline)"; size the column to the + # longest actual value rather than a fixed width, so labels like "additive (baseline)" + # aren't truncated. + predictor_width = max(2.0, 0.14 * plot_df["Predictor"].str.len().max()) + column_definitions = [ + ColumnDefinition("Predictor", width=predictor_width, textprops={"ha": "left", "weight": "bold"}), + ] + column_definitions += [ + ColumnDefinition( + col, + # Score column names (e.g. "delta_correlation", "systema_specific") are wider + # than the circular value markers below them; wrapping onto a second line at + # the first underscore, as `Benchmarker.plot_results_table` does at its first + # space, keeps adjacent headers from overlapping. + title=col.replace("_", "\n", 1), + width=1, + textprops={"ha": "center", "bbox": {"boxstyle": "circle", "pad": 0.25}}, + cmap=cmap_fn(plot_df[col]), + formatter="{:.2f}", + ) + for col in score_cols + ] + with mpl.rc_context({"svg.fonttype": "none"}): + fig, ax = plt.subplots(figsize=(len(score_cols) * 1.25 + predictor_width, 3 + 0.3 * len(plot_df))) + table = Table( + plot_df, + cell_kw={"linewidth": 0, "edgecolor": "k"}, + column_definitions=column_definitions, + ax=ax, + row_dividers=True, + footer_divider=True, + textprops={"fontsize": 10, "ha": "center"}, + row_divider_kw={"linewidth": 1, "linestyle": (0, (1, 5))}, + col_label_divider_kw={"linewidth": 1, "linestyle": "-"}, + column_border_kw={"linewidth": 1, "linestyle": "-"}, + index_col="Predictor", + ).autoset_fontcolors(colnames=score_cols) + if show: + plt.show() + if save_dir is not None: + fig.savefig(os.path.join(save_dir, "perturbation_results.svg"), facecolor=ax.get_facecolor(), dpi=300) + return table diff --git a/src/scib_metrics/perturbation/_metrics.py b/src/scib_metrics/perturbation/_metrics.py new file mode 100644 index 0000000..578ce19 --- /dev/null +++ b/src/scib_metrics/perturbation/_metrics.py @@ -0,0 +1,182 @@ +"""Evaluation metrics for perturbation-response prediction. + +All metrics consume expression *deltas* (predicted and true, relative to control), matching +the output convention of `scib_metrics.perturbation._baselines`. +""" + +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING, Literal + +import numpy as np +import pandas as pd + +from scib_metrics.perturbation._utils import import_pertpy + +if TYPE_CHECKING: + from collections.abc import Sequence + + +def delta_correlation( + predicted_deltas: np.ndarray, + true_deltas: np.ndarray, + gene_indices: Sequence[np.ndarray] | None = None, + method: Literal["pearson", "spearman"] = "pearson", +) -> dict[str, np.ndarray | float]: + """Per-perturbation correlation between predicted and true expression deltas. + + Parameters + ---------- + predicted_deltas + Array of shape `(n_perturbations, n_genes)`. + true_deltas + Array of shape `(n_perturbations, n_genes)`, aligned row-for-row with `predicted_deltas`. + gene_indices + Optional per-perturbation gene index arrays (length `n_perturbations`) restricting the + correlation to a caller-supplied gene subset (e.g. top-DE genes) for that perturbation. + method + `"pearson"` or `"spearman"`. + + Returns + ------- + Dict with `"per_perturbation"` (array of shape `(n_perturbations,)`) and `"mean"` (float). + """ + predicted_deltas = np.asarray(predicted_deltas) + true_deltas = np.asarray(true_deltas) + if predicted_deltas.shape != true_deltas.shape: + raise ValueError("`predicted_deltas` and `true_deltas` must have the same shape.") + pt = import_pertpy() + metric_name = "pearson_distance" if method == "pearson" else "spearman_distance" + distance = pt.tl.Distance(metric=metric_name) + scores = np.empty(predicted_deltas.shape[0]) + for i in range(predicted_deltas.shape[0]): + pred_i, true_i = predicted_deltas[i], true_deltas[i] + if gene_indices is not None: + idx = np.asarray(gene_indices[i]) + pred_i, true_i = pred_i[idx], true_i[idx] + scores[i] = 1 - distance(pred_i[None, :], true_i[None, :]) + return {"per_perturbation": scores, "mean": float(scores.mean())} + + +def de_rank_recovery( + predicted_deltas: np.ndarray, + true_de_gene_indices: Sequence[np.ndarray], + k: int, +) -> dict[str, np.ndarray | float]: + """Recall@k of top predicted-delta-magnitude genes against a true DE gene set. + + Parameters + ---------- + predicted_deltas + Array of shape `(n_perturbations, n_genes)`. + true_de_gene_indices + Per-perturbation array of the true top-DE gene indices, length `n_perturbations`. + k + Number of top genes (by predicted delta magnitude) considered "predicted as DE". + + Returns + ------- + Dict with `"per_perturbation"` (recall@k per perturbation) and `"mean"` (float). + """ + predicted_deltas = np.asarray(predicted_deltas) + if predicted_deltas.shape[0] != len(true_de_gene_indices): + raise ValueError("`predicted_deltas` and `true_de_gene_indices` must have the same length.") + scores = np.empty(predicted_deltas.shape[0]) + for i, true_idx in enumerate(true_de_gene_indices): + true_idx = np.asarray(true_idx) + predicted_top_k = np.argsort(-np.abs(predicted_deltas[i]))[:k] + n_recovered = np.intersect1d(predicted_top_k, true_idx).shape[0] + scores[i] = n_recovered / true_idx.shape[0] + return {"per_perturbation": scores, "mean": float(scores.mean())} + + +def systema_decomposition(predicted_deltas: np.ndarray, true_deltas: np.ndarray) -> dict[str, float]: + """Score shared (systematic) vs. perturbation-specific components of the effect separately. + + Requires at least 2 held-out perturbations to estimate the shared component; a predictor + that only reproduces the systematic shift will score well on `"shared"` but poorly on + `"specific"`. + + Parameters + ---------- + predicted_deltas + Array of shape `(n_perturbations, n_genes)`, `n_perturbations >= 2`. + true_deltas + Array of shape `(n_perturbations, n_genes)`, aligned row-for-row with `predicted_deltas`. + + Returns + ------- + Dict with `"shared"` and `"specific"` correlation scores. + """ + predicted_deltas = np.asarray(predicted_deltas) + true_deltas = np.asarray(true_deltas) + if predicted_deltas.shape != true_deltas.shape: + raise ValueError("`predicted_deltas` and `true_deltas` must have the same shape.") + if predicted_deltas.shape[0] < 2: + raise ValueError("`systema_decomposition` requires at least 2 held-out perturbations.") + pt = import_pertpy() + distance = pt.tl.Distance(metric="pearson_distance") + + true_shared = true_deltas.mean(axis=0) + predicted_shared = predicted_deltas.mean(axis=0) + shared_score = 1 - distance(predicted_shared[None, :], true_shared[None, :]) + if np.isnan(shared_score): + shared_score = 0.0 + + true_specific = (true_deltas - true_shared).ravel() + predicted_specific = (predicted_deltas - predicted_shared).ravel() + specific_score = 1 - distance(predicted_specific[None, :], true_specific[None, :]) + if np.isnan(specific_score): + specific_score = 0.0 + + return {"shared": float(shared_score), "specific": float(specific_score)} + + +def combination_additivity( + adata, + target_col: str = "perturbation", + reference_key: str = "control", + combinations: Sequence[str] | None = None, + sep: str = "+", +) -> pd.DataFrame: + """Score how well an additive model predicts combination perturbations. + + Thin wrapper over `pertpy.tools.PerturbationSpace.evaluate_combinations`. `adata` must be + perturbation-level (one observation per perturbation, e.g. from + `pertpy.tools.PseudobulkSpace.compute`). + + Parameters + ---------- + adata + Perturbation-level AnnData (one observation per perturbation). + target_col + `.obs` column identifying each perturbation. + reference_key + Control perturbation subtracted to obtain effects. + combinations + Combination names to evaluate. If `None`, every `obs_name` containing `sep` whose + components are all present as singles is used. + sep + Separator between components in combination names. + + Returns + ------- + DataFrame indexed by combination with `"distance"`, `"predicted_magnitude"` and + `"measured_magnitude"` columns, or an empty DataFrame with those columns if no evaluable + combination is present. + """ + pt = import_pertpy() + names = adata.obs_names.astype(str) + if combinations is None: + combinations = [n for n in names if sep in n and all(part in names for part in n.split(sep))] + if not combinations: + warnings.warn( + f"No combination-named perturbations found in `adata.obs_names` (using separator {sep!r}); " + "returning an empty result.", + stacklevel=2, + ) + return pd.DataFrame(columns=["distance", "predicted_magnitude", "measured_magnitude"]) + return pt.tl.PseudobulkSpace().evaluate_combinations( + adata, combinations=combinations, target_col=target_col, reference_key=reference_key, sep=sep + ) diff --git a/src/scib_metrics/perturbation/_utils.py b/src/scib_metrics/perturbation/_utils.py new file mode 100644 index 0000000..ec220fd --- /dev/null +++ b/src/scib_metrics/perturbation/_utils.py @@ -0,0 +1,37 @@ +"""Internal helpers for the optional pertpy-backed perturbation module.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from types import ModuleType + + +def import_pertpy() -> ModuleType: + """Import and return the `pertpy` package, raising a clear error if it's missing. + + `pertpy.tools._coda._sccoda` unconditionally calls + `jax.config.update("jax_enable_x64", True)` at import time — a side effect of + importing `pertpy` at all, not of using any Sccoda functionality — which silently + flips JAX's global float precision for the rest of the process. scib-metrics' own + jax-based code (e.g. `scib_metrics.utils._kmeans`) assumes the default (float32) + precision and breaks under x64, so this restores whatever the setting was + immediately before importing `pertpy`. + + Returns + ------- + The imported `pertpy` module. + """ + import jax + + was_x64_enabled = jax.config.jax_enable_x64 + try: + import pertpy + except ImportError as e: + raise ImportError( + "The `scib_metrics.perturbation` module requires `pertpy`. " + "Install it with `pip install scib-metrics[perturbation]`." + ) from e + jax.config.update("jax_enable_x64", was_x64_enabled) + return pertpy diff --git a/tests/perturbation/__init__.py b/tests/perturbation/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/perturbation/_synthetic.py b/tests/perturbation/_synthetic.py new file mode 100644 index 0000000..e6108ff --- /dev/null +++ b/tests/perturbation/_synthetic.py @@ -0,0 +1,60 @@ +"""Synthetic perturbation AnnData for testing scib_metrics.perturbation.""" + +from __future__ import annotations + +import numpy as np +from anndata import AnnData + + +def make_synthetic_perturbation_adata( + n_genes: int = 20, + n_cells_per_group: int = 30, + seed: int = 0, + target_col: str = "perturbation", + reference_key: str = "control", +) -> AnnData: + """Build a synthetic cell-level AnnData with a control group and perturbations A, B, A+B. + + Ground truth: perturbation `A` shifts genes `[0:5]` by `+3`, `B` shifts genes `[5:10]` + by `+2`, and `A+B` is exactly the sum of the two shifts (perfectly additive), so an + additive predictor should recover it exactly while a predictor blind to combination + structure should not. + + Parameters + ---------- + n_genes + Number of genes. + n_cells_per_group + Number of cells per perturbation group. + seed + Random seed for the noise added on top of each group's shift. + target_col + `.obs` column name holding the perturbation label. + reference_key + Perturbation label marking control cells. + + Returns + ------- + AnnData of shape `(4 * n_cells_per_group, n_genes)` with groups + `[reference_key, "A", "B", "A+B"]`. + """ + rng = np.random.default_rng(seed) + shift_a = np.zeros(n_genes) + shift_a[0:5] = 3.0 + shift_b = np.zeros(n_genes) + shift_b[5:10] = 2.0 + shifts = { + reference_key: np.zeros(n_genes), + "A": shift_a, + "B": shift_b, + "A+B": shift_a + shift_b, + } + rows, labels = [], [] + for name, shift in shifts.items(): + base = rng.normal(loc=10.0, scale=1.0, size=(n_cells_per_group, n_genes)) + rows.append(base + shift) + labels.extend([name] * n_cells_per_group) + adata = AnnData(X=np.concatenate(rows, axis=0).astype(np.float32)) + adata.obs[target_col] = labels + adata.obs[target_col] = adata.obs[target_col].astype("category") + return adata diff --git a/tests/perturbation/conftest.py b/tests/perturbation/conftest.py new file mode 100644 index 0000000..02a1656 --- /dev/null +++ b/tests/perturbation/conftest.py @@ -0,0 +1,23 @@ +"""Neutralize a pertpy import side effect before any test in this directory collects. + +`pertpy.tools._coda._sccoda` unconditionally calls `jax.config.update("jax_enable_x64", +True)` at import time, regardless of whether Sccoda is ever used. Every test file in this +directory triggers that import via `pytest.importorskip("pertpy")` at collection time, which +would otherwise leak jax's global float precision into unrelated, already-collected test +modules in the same pytest process (e.g. `tests/test_metrics.py::test_kmeans`, whose +`jax.lax.while_loop` carry state is not written to tolerate float64). Importing pertpy once +here and restoring the prior setting, before pytest collects any test module in this +directory, neutralizes the mutation for the rest of the session (the import is cached, so +later `import pertpy` / `pytest.importorskip("pertpy")` calls are no-ops that don't +re-trigger it). +""" + +try: + import jax + + _was_x64_enabled = jax.config.jax_enable_x64 + import pertpy # noqa: F401 + + jax.config.update("jax_enable_x64", _was_x64_enabled) +except ImportError: + pass diff --git a/tests/perturbation/test_baselines.py b/tests/perturbation/test_baselines.py new file mode 100644 index 0000000..0067aee --- /dev/null +++ b/tests/perturbation/test_baselines.py @@ -0,0 +1,73 @@ +import numpy as np +import pytest + +pytest.importorskip("pertpy") + +from scib_metrics.perturbation._baselines import AdditiveBaseline, LinearBaseline, MeanBaseline +from tests.perturbation._synthetic import make_synthetic_perturbation_adata + + +def test_mean_baseline_broadcasts_same_prediction_regardless_of_identity(): + adata = make_synthetic_perturbation_adata() + baseline = MeanBaseline().fit(adata, target_col="perturbation", reference_key="control") + predictions = baseline.predict(["A", "B", "unseen_name"]) + assert predictions.shape == (3, 20) + np.testing.assert_allclose(predictions[0], predictions[1]) + np.testing.assert_allclose(predictions[0], predictions[2]) + + +def test_mean_baseline_raises_when_reference_key_missing(): + adata = make_synthetic_perturbation_adata() + with pytest.raises(ValueError): + MeanBaseline().fit(adata, target_col="perturbation", reference_key="not_a_real_label") + + +def test_additive_baseline_recovers_perfect_combination(): + adata = make_synthetic_perturbation_adata() + train = adata[adata.obs["perturbation"].isin(["control", "A", "B"])].copy() + baseline = AdditiveBaseline().fit(train, target_col="perturbation", reference_key="control") + predicted = baseline.predict(["A+B"])[0] + + full_pseudobulk_a = adata[adata.obs["perturbation"] == "A"].X.mean(axis=0) + full_pseudobulk_b = adata[adata.obs["perturbation"] == "B"].X.mean(axis=0) + full_pseudobulk_control = adata[adata.obs["perturbation"] == "control"].X.mean(axis=0) + true_combination_delta = (full_pseudobulk_a - full_pseudobulk_control) + ( + full_pseudobulk_b - full_pseudobulk_control + ) + np.testing.assert_allclose(predicted, true_combination_delta, atol=1e-4) + + +def test_additive_baseline_falls_back_to_mean_for_non_decomposable_names(): + adata = make_synthetic_perturbation_adata() + train = adata[adata.obs["perturbation"].isin(["control", "A", "B"])].copy() + baseline = AdditiveBaseline().fit(train, target_col="perturbation", reference_key="control") + predicted = baseline.predict(["totally_unseen"])[0] + np.testing.assert_allclose(predicted, baseline.mean_delta_) + + +def test_linear_baseline_requires_perturbation_encodings(): + adata = make_synthetic_perturbation_adata() + with pytest.raises(ValueError, match="perturbation_encodings"): + LinearBaseline().fit(adata, target_col="perturbation", reference_key="control") + + +def test_linear_baseline_generalizes_to_unseen_encoding(): + adata = make_synthetic_perturbation_adata() + train = adata[adata.obs["perturbation"].isin(["control", "A", "B"])].copy() + encodings = {"A": np.array([1.0, 0.0]), "B": np.array([0.0, 1.0]), "A+B": np.array([1.0, 1.0])} + baseline = LinearBaseline().fit( + train, target_col="perturbation", reference_key="control", perturbation_encodings=encodings + ) + predicted = baseline.predict(["A+B"])[0] + assert predicted.shape == (20,) + + +def test_linear_baseline_raises_for_missing_encoding_at_predict_time(): + adata = make_synthetic_perturbation_adata() + train = adata[adata.obs["perturbation"].isin(["control", "A", "B"])].copy() + encodings = {"A": np.array([1.0, 0.0]), "B": np.array([0.0, 1.0])} + baseline = LinearBaseline().fit( + train, target_col="perturbation", reference_key="control", perturbation_encodings=encodings + ) + with pytest.raises(ValueError, match="No encoding supplied"): + baseline.predict(["C"]) diff --git a/tests/perturbation/test_core.py b/tests/perturbation/test_core.py new file mode 100644 index 0000000..73d0e79 --- /dev/null +++ b/tests/perturbation/test_core.py @@ -0,0 +1,232 @@ +import numpy as np +import pandas as pd +import pytest + +pytest.importorskip("pertpy") + +from scib_metrics.perturbation._core import PerturbationBaselines, PerturbationBenchmarker, PerturbationMetrics +from tests.perturbation._synthetic import make_synthetic_perturbation_adata + + +def _make_train_test_split(): + adata = make_synthetic_perturbation_adata() + train = adata[adata.obs["perturbation"].isin(["control", "A", "B"])].copy() + test = adata[adata.obs["perturbation"] == "A+B"].copy() + return train, test + + +def _perturbation_encodings(): + # `LinearBaseline` (Task 5) hard-requires a `perturbation_encodings` entry for every + # trained and held-out name whenever it's enabled (`PerturbationBaselines.linear=True` + # by default) -- see its docstring ("Required if `baselines.linear` is enabled."). A + # multi-hot encoding over the base perturbations ("A", "B") that composes additively for + # combinations ("A+B" = "A" | "B") is the natural minimal encoding for this synthetic + # fixture, and lets `LinearBaseline` generalize to the held-out combination. + return {"A": np.array([1.0, 0.0]), "B": np.array([0.0, 1.0]), "A+B": np.array([1.0, 1.0])} + + +def test_benchmarker_runs_default_baselines_and_metrics(): + train, test = _make_train_test_split() + benchmarker = PerturbationBenchmarker(train, test, perturbation_encodings=_perturbation_encodings()) + benchmarker.benchmark() + results = benchmarker.get_results(min_max_scale=False) + assert isinstance(results, pd.DataFrame) + assert set(results.index) == {"mean", "additive", "linear"} + assert results.loc["mean", "is_baseline"] + assert "delta_correlation" in results.columns + assert "systema_shared" not in results.columns # only 1 held-out perturbation: A+B + + +def test_benchmarker_includes_user_predictions_alongside_baselines(): + train, test = _make_train_test_split() + + user_predictions = {"my_model": np.zeros((1, 20))} + benchmarker = PerturbationBenchmarker( + train, test, predictions=user_predictions, perturbation_encodings=_perturbation_encodings() + ) + benchmarker.benchmark() + results = benchmarker.get_results(min_max_scale=False) + assert "my_model" in results.index + assert not results.loc["my_model", "is_baseline"] + + +def test_benchmarker_can_disable_a_baseline(): + train, test = _make_train_test_split() + benchmarker = PerturbationBenchmarker(train, test, baselines=PerturbationBaselines(linear=False)) + benchmarker.benchmark() + results = benchmarker.get_results(min_max_scale=False) + assert "linear" not in results.index + + +def test_benchmarker_raises_before_get_results_without_benchmark(): + train, test = _make_train_test_split() + benchmarker = PerturbationBenchmarker(train, test) + with pytest.raises(RuntimeError): + benchmarker.get_results() + + +def test_benchmarker_de_rank_recovery_requires_true_de_gene_indices(): + train, test = _make_train_test_split() + benchmarker = PerturbationBenchmarker( + train, + test, + metrics=PerturbationMetrics(de_rank_recovery=True), + perturbation_encodings=_perturbation_encodings(), + ) + with pytest.raises(ValueError, match="true_de_gene_indices"): + benchmarker.benchmark() + + +def test_benchmarker_raises_on_prediction_baseline_collision(): + train, test = _make_train_test_split() + user_predictions = {"mean": np.zeros((1, 20))} + benchmarker = PerturbationBenchmarker( + train, + test, + predictions=user_predictions, + perturbation_encodings=_perturbation_encodings(), + ) + with pytest.raises(ValueError, match="collide with enabled baseline names"): + benchmarker.benchmark() + + +def test_get_ground_truth_significance_requires_flag_enabled(): + train, test = _make_train_test_split() + benchmarker = PerturbationBenchmarker(train, test, baselines=PerturbationBaselines(linear=False)) + benchmarker.benchmark() + with pytest.raises(RuntimeError): + benchmarker.get_ground_truth_significance() + + +def test_get_ground_truth_significance_warns_without_control_cells_in_test(): + train, test = _make_train_test_split() # test split has no control cells + benchmarker = PerturbationBenchmarker( + train, + test, + baselines=PerturbationBaselines(linear=False), + metrics=PerturbationMetrics(ground_truth_significance=True), + ) + with pytest.warns(UserWarning, match="skipping"): + benchmarker.benchmark() + result = benchmarker.get_ground_truth_significance() + assert result.empty + + +def test_get_ground_truth_significance_runs_with_control_cells_in_test(): + adata = make_synthetic_perturbation_adata() + train = adata[adata.obs["perturbation"].isin(["control", "A", "B"])].copy() + test = adata[adata.obs["perturbation"].isin(["control", "A+B"])].copy() + benchmarker = PerturbationBenchmarker( + train, + test, + baselines=PerturbationBaselines(linear=False), + metrics=PerturbationMetrics(ground_truth_significance=True), + ) + benchmarker.benchmark() + result = benchmarker.get_ground_truth_significance() + assert "pvalue" in result.columns + + +def test_get_combination_additivity_raises_before_benchmark(): + train, test = _make_train_test_split() + benchmarker = PerturbationBenchmarker(train, test, perturbation_encodings=_perturbation_encodings()) + with pytest.raises(RuntimeError): + benchmarker.get_combination_additivity() + + +def test_get_combination_additivity_finds_combination_present_only_in_test(): + # train only has the trained singles ("A", "B"); the combination "A+B" is held out in + # `test`. `combination_additivity` must pseudobulk the *union* of train and test to ever + # see a combination-named perturbation at all. + train, test = _make_train_test_split() + benchmarker = PerturbationBenchmarker(train, test, perturbation_encodings=_perturbation_encodings()) + benchmarker.benchmark() + result = benchmarker.get_combination_additivity() + assert not result.empty + assert "A+B" in result.index + + +def test_delta_correlation_unaffected_by_true_de_gene_indices_when_de_rank_recovery_disabled(): + train, test = _make_train_test_split() + gene_indices = {"A+B": np.arange(5)} + + benchmarker_without = PerturbationBenchmarker( + train, + test, + perturbation_encodings=_perturbation_encodings(), + metrics=PerturbationMetrics(de_rank_recovery=False), + ) + benchmarker_without.benchmark() + results_without = benchmarker_without.get_results(min_max_scale=False) + + benchmarker_with = PerturbationBenchmarker( + train, + test, + perturbation_encodings=_perturbation_encodings(), + true_de_gene_indices=gene_indices, + metrics=PerturbationMetrics(de_rank_recovery=False), + ) + benchmarker_with.benchmark() + results_with = benchmarker_with.get_results(min_max_scale=False) + + pd.testing.assert_series_equal( + results_without["delta_correlation"], results_with["delta_correlation"], check_names=True + ) + + +def test_systema_decomposition_warns_when_fewer_than_two_held_out_perturbations(): + # The default fixture's test split holds out only "A+B" -- a single perturbation -- + # while `metrics.systema_decomposition` defaults to `True`. + train, test = _make_train_test_split() + benchmarker = PerturbationBenchmarker(train, test, perturbation_encodings=_perturbation_encodings()) + with pytest.warns(UserWarning, match="systema_decomposition"): + benchmarker.benchmark() + results = benchmarker.get_results(min_max_scale=False) + assert "systema_shared" not in results.columns + assert "systema_specific" not in results.columns + + +def test_plot_results_table_returns_a_table(): + from plottable import Table + + train, test = _make_train_test_split() + benchmarker = PerturbationBenchmarker(train, test, baselines=PerturbationBaselines(linear=False)) + benchmarker.benchmark() + table = benchmarker.plot_results_table(show=False) + assert isinstance(table, Table) + + +def test_perturbation_benchmarker(): + # Mirrors `tests/test_benchmarker.py::test_benchmarker`: run the full default pipeline + # (all baselines, all default-on metrics) end-to-end on synthetic data and plot it, the + # same way that test exercises `Benchmarker` with `BatchCorrection()`/`BioConservation()`. + # + # Unlike `_make_train_test_split()` (which holds out only "A+B" and is used by the other, + # narrower tests in this file), this split holds out *two* perturbations -- "B" and "A+B" + # -- so `systema_decomposition` (needs >=2 held-out perturbations) actually runs instead + # of warning and skipping, and "A+B" is a real combination in the train+test union so + # `combination_additivity` finds a non-empty result instead of warning and returning empty. + adata = make_synthetic_perturbation_adata() + train = adata[adata.obs["perturbation"].isin(["control", "A"])].copy() + test = adata[adata.obs["perturbation"].isin(["B", "A+B"])].copy() + encodings = {"A": np.array([1.0, 0.0]), "B": np.array([0.0, 1.0]), "A+B": np.array([1.0, 1.0])} + + benchmarker = PerturbationBenchmarker(train, test, perturbation_encodings=encodings) + benchmarker.benchmark() + + results = benchmarker.get_results() + assert isinstance(results, pd.DataFrame) + assert set(results.index) == {"mean", "additive", "linear"} + assert "delta_correlation" in results.columns + assert "systema_shared" in results.columns + assert "systema_specific" in results.columns + + combinations = benchmarker.get_combination_additivity() + assert "A+B" in combinations.index + + # `show=True` opens a real, blocking GUI window (confirmed identical to + # `Benchmarker.plot_results_table`'s behavior) -- unreliable to depend on inside a test + # runner. `save_dir` sidesteps that: it writes a real, inspectable SVG regardless of + # backend/runner quirks. Open /tmp/perturbation_results.svg after running this test to + # visually verify the plot (colors, baseline-row labeling, etc.). + benchmarker.plot_results_table(show=False, save_dir="/tmp") diff --git a/tests/perturbation/test_metrics.py b/tests/perturbation/test_metrics.py new file mode 100644 index 0000000..21c0aca --- /dev/null +++ b/tests/perturbation/test_metrics.py @@ -0,0 +1,102 @@ +import numpy as np +import pytest + +pytest.importorskip("pertpy") + +import pandas as pd + +from scib_metrics.perturbation._metrics import ( + combination_additivity, + de_rank_recovery, + delta_correlation, + systema_decomposition, +) +from tests.perturbation._synthetic import make_synthetic_perturbation_adata + + +def test_delta_correlation_perfect_prediction_scores_one(): + rng = np.random.default_rng(0) + true_deltas = rng.normal(size=(5, 20)) + result = delta_correlation(true_deltas, true_deltas) + np.testing.assert_allclose(result["per_perturbation"], 1.0, atol=1e-6) + assert result["mean"] == pytest.approx(1.0, abs=1e-6) + + +def test_delta_correlation_anti_correlated_prediction_scores_near_minus_one(): + rng = np.random.default_rng(0) + true_deltas = rng.normal(size=(5, 20)) + result = delta_correlation(-true_deltas, true_deltas) + np.testing.assert_allclose(result["per_perturbation"], -1.0, atol=1e-6) + + +def test_delta_correlation_restricts_to_gene_indices(): + true_deltas = np.array([[1.0, 2.0, 100.0]]) + predicted_deltas = np.array([[1.0, 2.0, -100.0]]) + result = delta_correlation(predicted_deltas, true_deltas, gene_indices=[np.array([0, 1])]) + assert result["mean"] == pytest.approx(1.0, abs=1e-6) + + +def test_delta_correlation_raises_on_shape_mismatch(): + with pytest.raises(ValueError): + delta_correlation(np.zeros((2, 3)), np.zeros((2, 4))) + + +def test_de_rank_recovery_perfect_when_top_k_matches(): + predicted_deltas = np.array([[5.0, 0.1, -4.0, 0.2, 0.3]]) + true_de_gene_indices = [np.array([0, 2])] + result = de_rank_recovery(predicted_deltas, true_de_gene_indices, k=2) + assert result["per_perturbation"][0] == pytest.approx(1.0) + assert result["mean"] == pytest.approx(1.0) + + +def test_de_rank_recovery_zero_when_top_k_disjoint(): + predicted_deltas = np.array([[5.0, 0.1, -4.0, 0.2, 0.3]]) + true_de_gene_indices = [np.array([1, 3])] + result = de_rank_recovery(predicted_deltas, true_de_gene_indices, k=2) + assert result["per_perturbation"][0] == pytest.approx(0.0) + + +def test_de_rank_recovery_raises_on_length_mismatch(): + with pytest.raises(ValueError): + de_rank_recovery(np.zeros((2, 5)), [np.array([0])], k=1) + + +def test_systema_decomposition_perfect_prediction_scores_one_on_both(): + rng = np.random.default_rng(0) + true_deltas = rng.normal(size=(4, 20)) + result = systema_decomposition(true_deltas, true_deltas) + assert result["shared"] == pytest.approx(1.0, abs=1e-6) + assert result["specific"] == pytest.approx(1.0, abs=1e-6) + + +def test_systema_decomposition_shared_only_prediction_scores_low_on_specific(): + rng = np.random.default_rng(0) + true_deltas = rng.normal(size=(4, 20)) + shared_only_prediction = np.tile(true_deltas.mean(axis=0), (4, 1)) + result = systema_decomposition(shared_only_prediction, true_deltas) + assert result["shared"] == pytest.approx(1.0, abs=1e-6) + assert result["specific"] < 0.5 + + +def test_systema_decomposition_raises_with_fewer_than_two_perturbations(): + with pytest.raises(ValueError): + systema_decomposition(np.zeros((1, 5)), np.zeros((1, 5))) + + +def test_combination_additivity_scores_perfect_additivity(): + adata = make_synthetic_perturbation_adata() + pt = pytest.importorskip("pertpy") + pseudobulk = pt.tl.PseudobulkSpace().compute(adata, target_col="perturbation", mode="mean") + result = combination_additivity(pseudobulk, target_col="perturbation", reference_key="control") + assert isinstance(result, pd.DataFrame) + assert result.loc["A+B", "distance"] == pytest.approx(0.0, abs=0.1) + + +def test_combination_additivity_returns_empty_when_no_combinations_present(): + pt = pytest.importorskip("pertpy") + adata = make_synthetic_perturbation_adata() + singles_only = adata[adata.obs["perturbation"].isin(["control", "A", "B"])].copy() + pseudobulk = pt.tl.PseudobulkSpace().compute(singles_only, target_col="perturbation", mode="mean") + result = combination_additivity(pseudobulk, target_col="perturbation", reference_key="control") + assert result.empty + assert list(result.columns) == ["distance", "predicted_magnitude", "measured_magnitude"] diff --git a/tests/perturbation/test_public_api.py b/tests/perturbation/test_public_api.py new file mode 100644 index 0000000..0e70fb0 --- /dev/null +++ b/tests/perturbation/test_public_api.py @@ -0,0 +1,31 @@ +def test_public_api_is_importable(): + from scib_metrics.perturbation import ( + AdditiveBaseline, + BasePerturbationPredictor, + LinearBaseline, + MeanBaseline, + PerturbationBaselines, + PerturbationBenchmarker, + PerturbationMetrics, + combination_additivity, + de_rank_recovery, + delta_correlation, + systema_decomposition, + ) + + assert all( + callable(obj) + for obj in ( + AdditiveBaseline, + BasePerturbationPredictor, + LinearBaseline, + MeanBaseline, + PerturbationBaselines, + PerturbationBenchmarker, + PerturbationMetrics, + combination_additivity, + de_rank_recovery, + delta_correlation, + systema_decomposition, + ) + ) diff --git a/tests/perturbation/test_utils.py b/tests/perturbation/test_utils.py new file mode 100644 index 0000000..93052bb --- /dev/null +++ b/tests/perturbation/test_utils.py @@ -0,0 +1,17 @@ +import sys +from unittest.mock import patch + +import pytest + +from scib_metrics.perturbation._utils import import_pertpy + + +def test_import_pertpy_returns_module(): + pertpy = pytest.importorskip("pertpy") + assert import_pertpy() is pertpy + + +def test_import_pertpy_raises_clear_error_when_missing(): + with patch.dict(sys.modules, {"pertpy": None}): + with pytest.raises(ImportError, match=r"scib-metrics\[perturbation\]"): + import_pertpy() diff --git a/tests/test_benchmarker.py b/tests/test_benchmarker.py index 69a21fc..c9e7402 100644 --- a/tests/test_benchmarker.py +++ b/tests/test_benchmarker.py @@ -19,7 +19,7 @@ def test_benchmarker(): bm.benchmark() results = bm.get_results() assert isinstance(results, pd.DataFrame) - bm.plot_results_table() + bm.plot_results_table(show=False) def test_benchmarker_default(): @@ -33,7 +33,7 @@ def test_benchmarker_default(): bm.benchmark() results = bm.get_results() assert isinstance(results, pd.DataFrame) - bm.plot_results_table() + bm.plot_results_table(show=False) def test_benchmarker_custom_metric_booleans(): @@ -91,7 +91,7 @@ def test_benchmarker_custom_near_neighs(): bm.benchmark() results = bm.get_results() assert isinstance(results, pd.DataFrame) - bm.plot_results_table() + bm.plot_results_table(show=False) @pytest.mark.parametrize("solver", ["arpack", "randomized"]) @@ -101,4 +101,4 @@ def test_benchmarker_different_solvers(solver): bm.benchmark() results = bm.get_results() assert isinstance(results, pd.DataFrame) - bm.plot_results_table() + bm.plot_results_table(show=False)