Source code for mantispy.metrics._evaluate

"""The integration benchmark: scib-metrics' standard panel drawn as a mantispy heatmap, with a mAP block."""

from __future__ import annotations

import warnings
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any

import numpy as np
import pandas as pd

from mantispy._core._pca import pca_solver
from mantispy._core._reduce import get_matrix
from mantispy._core.frames import as_frame
from mantispy._core.masks import control_annotated, reference_mask
from mantispy.metrics._common import embedding

if TYPE_CHECKING:
    from anndata import AnnData
    from matplotlib.axes import Axes

#: The pair-argument keys a caller must fully specify to bypass a preset.
_PAIR_KEYS = ("pos_sameby", "pos_diffby", "neg_sameby", "neg_diffby")

#: Tokens in ``metrics=`` that select the mAP rather than a scib metric.
_MAP_TOKENS = frozenset({"map"})

#: obsm key of the unintegrated PCA scib's PC-regression comparison reads, on mantispy's own copy.
_UNINTEGRATED = "_mantispy_unintegrated_pca"


def _has_scib() -> bool:
    """Whether scib-metrics, the engine behind the benchmark panel, can be imported."""
    try:
        import scib_metrics  # noqa: F401
    except ImportError:
        return False
    return True


def _map_settings(
    map_mode: str, label_key: str, batch_key: str, map_kwargs: dict[str, Any] | None
) -> dict[str, list[str]]:
    """The four pair arguments for the mAP, from the ``replicability``/``cross_plate`` preset plus overrides.

    Raises:
        ValueError: ``map_mode`` is neither preset and ``map_kwargs`` did not fully specify the four pair arguments.
    """
    from mantispy.tl._map import MODES, _resolve_mode

    if map_mode in MODES and map_mode not in ("activity", "consistency"):
        settings = _resolve_mode(map_mode, label=label_key, batch=batch_key)
        if map_kwargs is not None:
            settings.update({key: list(value) for key, value in map_kwargs.items()})
        return settings
    if map_kwargs is None or not set(_PAIR_KEYS) <= set(map_kwargs):
        raise ValueError(
            f"map_mode must be 'replicability' or 'cross_plate', or pass map_kwargs with all of {_PAIR_KEYS}; "
            f"got {map_mode!r}"
        )
    return {key: list(map_kwargs[key]) for key in _PAIR_KEYS}


def _map_meta(adata: AnnData, settings: dict[str, list[str]]) -> tuple[pd.DataFrame, np.ndarray | None]:
    """The pair metadata frame and the treated-row mask, built once so every representation reuses them."""
    from mantispy._core.frames import as_frame

    obs = as_frame(adata.obs)
    needed = sorted({column for group in settings.values() for column in group})
    meta = pd.DataFrame({column: obs[column].reset_index(drop=True).to_numpy() for column in needed})
    treated = None
    if control_annotated(adata):
        # reference_mask reads Metadata_Control_Type for the negcon class (spec 10.3); the controls are
        # dropped from the mAP, so the treated rows are everything that is not a negative control.
        treated = ~reference_mask(adata, "negcon")
        meta = meta[treated].reset_index(drop=True)
    return meta, treated


def _rep_map(
    adata: AnnData,
    use_rep: str,
    settings: dict[str, list[str]],
    prepared: tuple[pd.DataFrame, np.ndarray | None] | None = None,
) -> float:
    """Mean of the per-group mean average precision over ``obsm[use_rep]``, controls left out."""
    from mantispy.tl._map import _score_map

    meta, treated = prepared if prepared is not None else _map_meta(adata, settings)
    features = embedding(adata, use_rep).astype(np.float32)
    if treated is not None:
        features = features[treated]

    # compute_null=False skips the permutation null: only the mean point estimate is read here, not the p-values.
    table, _ = _score_map(meta, features, settings, null_size=0, threshold=0.05, seed=0, warn=False, compute_null=False)
    return float(table["score"].mean())


def _scib_metric_specs(scib_names: list[str] | None) -> dict[str, Any]:
    """The ``Benchmarker`` metric-spec kwargs that run only the requested scib metrics.

    ``None`` returns ``{}`` so Benchmarker keeps its own default panel (it rejects an explicit ``(None, None)``).
    Otherwise every flag not named is turned off, so a caller pays only for the metrics it asked for.

    Raises:
        ValueError: A name is neither a bio-conservation nor a batch-correction metric.
    """
    if scib_names is None:
        return {}
    import dataclasses

    from scib_metrics.benchmark import BatchCorrection, BioConservation

    bio_fields = {field.name for field in dataclasses.fields(BioConservation)}
    batch_fields = {field.name for field in dataclasses.fields(BatchCorrection)}
    requested = set(scib_names)
    if unknown := requested - bio_fields - batch_fields:
        raise ValueError(
            f"unknown metric(s) {sorted(unknown)}; choose from {sorted(bio_fields | batch_fields)} or 'map'"
        )
    return {
        "bio_conservation_metrics": BioConservation(**{field: field in requested for field in bio_fields}),
        "batch_correction_metrics": BatchCorrection(**{field: field in requested for field in batch_fields}),
    }


def _benchmark(
    adata: AnnData,
    *,
    reps: Sequence[str],
    label_key: str,
    batch_key: str,
    min_max_scale: bool,
    scib_names: list[str] | None = None,
    svd_solver: str | None = None,
) -> pd.DataFrame:
    """scib-metrics' public ``get_results`` frame for the representations, kept as a seam a test can read directly.

    scib-metrics is imported here and only here, and only the public ``get_results`` is read, never the private
    ``_results``, so scib's own Total stays the number scib computes.
    """
    import anndata as ad
    import scanpy as sc
    from scib_metrics.benchmark import Benchmarker

    # scib writes its unintegrated PCA into obsm["X_pca"], so it gets a copy holding only what it reads, and the
    # baseline is computed here the way scib computes it, under a key no representation can share.
    slim = ad.AnnData(get_matrix(adata), obs=as_frame(adata.obs)[list(dict.fromkeys((batch_key, label_key)))].copy())
    for rep in reps:
        slim.obsm[rep] = adata.obsm[rep]
    solver = pca_solver(slim.n_obs, slim.n_vars, svd_solver, fallback="arpack")
    sc.tl.pca(slim, svd_solver=solver, mask=None, key_added=_UNINTEGRATED)
    # progress_bar=False keeps the tqdm bars out of executed notebooks and test logs.
    bm = Benchmarker(
        slim,
        batch_key=batch_key,
        label_key=label_key,
        embedding_obsm_keys=list(reps),
        pre_integrated_embedding_obsm_key=_UNINTEGRATED,
        n_jobs=-1,
        progress_bar=False,
        **_scib_metric_specs(scib_names),
    )
    bm.benchmark()
    return bm.get_results(min_max_scale=min_max_scale)


[docs] def evaluate_integration( adata: AnnData, *, reps: Sequence[str] = ("X_pca",), label_key: str = "Metadata_Perturbation", batch_key: str = "Metadata_Batch", min_max_scale: bool = False, metrics: Sequence[str] | None = None, map_mode: str = "replicability", map_kwargs: dict[str, Any] | None = None, svd_solver: str | None = None, ax: Axes | None = None, ) -> pd.DataFrame: """Score one or more representations against a batch and draw the integration benchmark as a heatmap. The numbers come from scib-metrics, the field's implementation of the integration panel, which mantispy reads through its public ``get_results`` rather than reimplements. mantispy renders its own heatmap from them and adds the cross-replicate mean average precision (mAP) as its own block. The columns are grouped left to right into bio conservation, batch correction, retrieval (mAP) and the aggregate scores, each block under its own header. The bio-conservation metrics measure whether a representation keeps the biology (cLISI is how label-pure a well's neighbourhood is); the batch-correction metrics whether it mixes the batches (iLISI is how batch-mixed that neighbourhood is, PCR comparison how much less of the variance the batch explains after correction). ``Total`` is scib's own weighted score, ``0.4`` batch correction and ``0.6`` bio conservation, left exactly as scib computes it. ``Total+mAP`` is mantispy's: it folds mAP into the bio group as one more bio signal, ``bio' = mean(bio metrics + mAP)``, then reweights ``0.4`` batch correction and ``0.6`` bio'. The mAP is measured over the treated wells only (controls left out), while scib's bio, batch and ``Total`` columns use every well, so ``Total`` and ``Total+mAP`` are not strictly apples-to-apples on a control-heavy screen. The return type is uniform across install states, so the numbers are always reachable: - With ``mantispy[integration]`` (scib-metrics), the frame holds every metric, the mAP column and both totals; the heatmap has all four blocks. - Without it, the frame and heatmap hold the per-representation mAP alone, and a warning notes that scib-metrics adds the full panel. This function computes no native PC-regression of its own. To audit what a representation spends its variance on besides the label, such as the cell count or the plate position, whose better direction is context-dependent and so is deliberately not in this table, use :func:`~mantispy.metrics.pc_regression` for a single covariate or :func:`~mantispy.metrics.batch_variance_explained` to stack several. Args: adata: Object holding the representations in ``obsm`` and the label and batch in ``obs``. reps: ``obsm`` keys to compare, e.g. ``("X_pca", "X_pca_harmony")``. label_key: ``obs`` column with the biological grouping. batch_key: ``obs`` column with the nuisance grouping. min_max_scale: Colour and score each metric column scaled across the representations, as scib does by default. Off by default so a single representation still has honest absolute values to colour by. metrics: Which metrics to run, to control cost. ``None`` runs the full scib panel plus the mAP. A list names exactly what to run: ``"map"`` for the mAP, and any scib metric flag (``"ilisi_knn"``, ``"clisi_knn"``, ``"kbet_per_label"``, ``"pcr_comparison"``, ``"graph_connectivity"``, ``"isolated_labels"``, ``"silhouette_label"``, ...). ``metrics=["map"]`` skips scib entirely, the cheapest option. The full panel is expensive: expect several minutes and many GB on tens of thousands of cells, since kBET and the clustering metrics dominate. map_mode: The pairing for mAP, ``"replicability"`` (the Arevalo ``mAP-nonrep`` default) or ``"cross_plate"``. See ``_map_settings``. map_kwargs: Overrides for any of the four pair arguments, on top of ``map_mode``. svd_solver: How the unintegrated baseline PCA is decomposed. ``None`` eigendecomposes the feature covariance when the screen has many more wells than features and otherwise uses arpack; any scanpy solver name forces one. The baseline matches either way. ax: Axes to draw the heatmap on, or ``None`` for a new figure. Returns: The numeric results, one row per representation and one column per metric, the mAP and both totals. Raises: ImportError: ``metrics`` asks for a scib metric and scib-metrics is not installed. ValueError: The scib panel is requested with no features in ``adata.X``; the mAP is requested but ``label_key`` has no replicates or is confounded with ``batch_key``; or ``metrics`` names an unknown metric or selects nothing runnable. KeyError: ``obs`` is missing ``label_key`` or ``batch_key``. Notes: The mAP runs on numba's threads, one per core by default. Cap them with ``numba.set_num_threads``. """ from mantispy.pl._evaluation import _integration_heatmap have_scib = _has_scib() reps = tuple(dict.fromkeys(reps)) # a duplicate rep would draw two identically-labelled rows missing = [key for key in (label_key, batch_key) if key not in adata.obs.columns] if missing: raise KeyError(f"obs is missing the column(s) the benchmark needs: {missing}") # metrics=None runs the full default panel plus the mAP; a list fully specifies it, where "map" selects the # mAP and every other name a scib metric (validated against scib's flags inside _benchmark). if metrics is None: scib_names, use_map, use_scib = None, True, have_scib else: names = [str(metric) for metric in metrics] scib_names = [name for name in names if name not in _MAP_TOKENS] use_map, use_scib = len(scib_names) < len(names), bool(scib_names) if use_scib and not have_scib: raise ImportError(f'metric(s) {scib_names} need scib-metrics; install pip install "mantispy[integration]".') if not use_map and not use_scib: raise ValueError( f"metrics={metrics!r} selected nothing to run given the installed extras; pass 'map' and/or " "scib metric names." ) if use_map: # The mAP drops the controls; value_counts drops NaN labels too. treated_labels = adata.obs[label_key] if control_annotated(adata): treated_labels = treated_labels[~reference_mask(adata, "negcon")] counts = treated_labels.value_counts() if counts.empty: raise ValueError( f"evaluate_integration found no treated profiles to score: obs[{label_key!r}] is empty or all " "missing once the controls are dropped." ) if int(counts.max()) < 2: raise ValueError( "evaluate_integration needs replicate profiles across batches, and this object looks like one row " f"per perturbation (no value of obs[{label_key!r}] has two or more profiles). It scores the " "integration of well-level replicates; for consensus profiles use mt.tl.map(mode='consistency') " "or mt.metrics.known_relationships instead." ) settings = _map_settings(map_mode, label_key, batch_key, map_kwargs) prepared = _map_meta(adata, settings) try: map_values = [_rep_map(adata, rep, settings, prepared) for rep in reps] except ValueError as error: if "negative pair" not in str(error).lower(): raise raise ValueError( f"evaluate_integration could not form negative pairs for the mAP, which usually means " f"label_key={label_key!r} is confounded with batch_key={batch_key!r} (e.g. one {label_key} per " f"{batch_key}), so no profile has a within-label, cross-batch negative. Score a label that is not " "nested within the batch, or check the design." ) from error if not use_scib: if metrics is None: # scib genuinely not installed (an explicit map-only request chose this silently) warnings.warn( 'scib-metrics is not installed, so only mAP is computed. Install with pip install "mantispy[integration]" for the full panel.', UserWarning, stacklevel=2, ) frame = pd.DataFrame( {"mean_average_precision": map_values}, index=pd.Index(list(reps), name="representation"), ) _integration_heatmap(frame, {"Retrieval": ["mean_average_precision"]}, ax=ax) return frame if adata.n_vars == 0: raise ValueError( "evaluate_integration's scib panel needs an unintegrated X with features (it computes a PCA " f"baseline), but adata has {adata.n_vars} var(s). Pass metrics=['map'] for retrieval only, or give " "the object its measured features as X." ) if min_max_scale and len(reps) == 1: warnings.warn( "min_max_scale=True needs more than one representation to scale against, so every scaled column is constant here.", UserWarning, stacklevel=2, ) results = _benchmark( adata, reps=reps, label_key=label_key, batch_key=batch_key, min_max_scale=min_max_scale, scib_names=scib_names, svd_solver=svd_solver, ) metric_type = results.loc["Metric Type"] data = results.drop(index="Metric Type").astype(float).loc[list(reps)] bio_cols = [column for column in data.columns if metric_type[column] == "Bio conservation"] batch_cols = [column for column in data.columns if metric_type[column] == "Batch correction"] aggregate_cols = ["Batch correction", "Bio conservation", "Total"] blocks: dict[str, list[str]] = {"Bio conservation": bio_cols, "Batch correction": batch_cols} if use_map: map_column = pd.Series(map_values, index=data.index) if min_max_scale: # Fold the mAP in on the same worst->0, best->1 scale as the bio metrics it joins, so Total+mAP and # the heatmap colour are not a scaled/raw mix. span = map_column.max() - map_column.min() map_column = (map_column - map_column.min()) / span data["mean_average_precision"] = map_column blocks["Retrieval"] = ["mean_average_precision"] if bio_cols: bio_prime = (data[bio_cols].sum(axis=1) + data["mean_average_precision"]) / (len(bio_cols) + 1) data["Total+mAP"] = 0.4 * data["Batch correction"] + 0.6 * bio_prime aggregate_cols = [*aggregate_cols, "Total+mAP"] else: warnings.warn( "scib-metrics returned no bio-conservation columns, so Total+mAP is left out. Check the scib-metrics version.", UserWarning, stacklevel=2, ) blocks["Aggregate"] = aggregate_cols data = data[[column for columns in blocks.values() for column in columns]] _integration_heatmap(data, blocks, ax=ax) return data