bump version and benchmarking changes (#69)

* bump version and benchmarking changes

* docs fix

* changelog

* fix diffusion distance (#70)

Author: adamgayoso

.readthedocs.yaml

@@ -3,7 +3,7 @@ version: 2
 build:
     os: ubuntu-20.04
     tools:
-        python: "3.8"
+        python: "3.10"
 sphinx:
     configuration: docs/conf.py
     fail_on_warning: false

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.1.0 (2022-01-03)
+
+-   Add benchmarking pipeline with plotting ([#52][] and [#69][])
+-   Fix diffusion distance computation, affecting kbet ([#70][])
+
+[#52]: https://github.com/YosefLab/scib-metrics/pull/52
+[#69]: https://github.com/YosefLab/scib-metrics/pull/69
+[#70]: https://github.com/YosefLab/scib-metrics/pull/70
+
 ## 0.0.9 (2022-12-16)
 
 -   Add kbet ([#60][])

pyproject.toml

@@ -5,7 +5,7 @@ requires = ["hatchling"]
 
 [project]
 name = "scib-metrics"
-version = "0.0.9"
+version = "0.1.0"
 description = "Accelerated and Python-only scIB metrics"
 readme = "README.md"
 requires-python = ">=3.8"
@@ -44,7 +44,7 @@ dev = [
     "pre-commit"
 ]
 doc = [
-    "sphinx>=4",
+    "sphinx>=4,<5.3",
     "sphinx-book-theme",
     "myst-nb",
     "sphinxcontrib-bibtex>=1.0.0",

src/scib_metrics/benchmark/_core.py

@@ -1,7 +1,9 @@
 import os
+import warnings
 from dataclasses import asdict, dataclass
 from enum import Enum
-from typing import Callable, List, Optional, Union
+from functools import partial
+from typing import Any, Dict, List, Optional, Union
 
 import matplotlib
 import matplotlib.pyplot as plt
@@ -18,12 +20,16 @@
 
 import scib_metrics
 
+Kwargs = Dict[str, Any]
+MetricType = Union[bool, Kwargs]
+
 _LABELS = "labels"
 _BATCH = "batch"
 _X_PRE = "X_pre"
 _METRIC_TYPE = "Metric Type"
 _AGGREGATE_SCORE = "Aggregate score"
 
+# Mapping of metric fn names to clean DataFrame column names
 metric_name_cleaner = {
     "silhouette_label": "Silhouette label",
     "silhouette_batch": "Silhouette batch",
@@ -40,34 +46,36 @@
 }
 
 
-@dataclass
+@dataclass(frozen=True)
 class BioConservation:
     """Specification of bio conservation metrics to run in the pipeline.
 
     Metrics can be included using a boolean flag. Custom keyword args can be
-    used by passing a partial callable of that metric here.
+    used by passing a dictionary here. Keyword args should not set data-related
+    parameters, such as `X` or `labels`.
     """
 
-    isolated_labels: Union[bool, Callable] = True
-    nmi_ari_cluster_labels_leiden: Union[bool, Callable] = True
-    nmi_ari_cluster_labels_kmeans: Union[bool, Callable] = False
-    silhouette_label: Union[bool, Callable] = True
-    clisi_knn: Union[bool, Callable] = True
+    isolated_labels: MetricType = True
+    nmi_ari_cluster_labels_leiden: MetricType = True
+    nmi_ari_cluster_labels_kmeans: MetricType = False
+    silhouette_label: MetricType = True
+    clisi_knn: MetricType = True
 
 
-@dataclass
+@dataclass(frozen=True)
 class BatchCorrection:
     """Specification of which batch correction metrics to run in the pipeline.
 
     Metrics can be included using a boolean flag. Custom keyword args can be
-    used by passing a partial callable of that metric here.
+    used by passing a dictionary here. Keyword args should not set data-related
+    parameters, such as `X` or `labels`.
     """
 
-    silhouette_batch: Union[bool, Callable] = True
-    ilisi_knn: Union[bool, Callable] = True
-    kbet_per_label: Union[bool, Callable] = True
-    graph_connectivity: Union[bool, Callable] = True
-    pcr_comparison: Union[bool, Callable] = True
+    silhouette_batch: MetricType = True
+    ilisi_knn: MetricType = True
+    kbet_per_label: MetricType = True
+    graph_connectivity: MetricType = True
+    pcr_comparison: MetricType = True
 
 
 class MetricAnnDataAPI(Enum):
@@ -138,6 +146,7 @@ def __init__(
         self._emb_adatas = {}
         self._neighbor_values = (15, 50, 90)
         self._prepared = False
+        self._benchmarked = False
         self._batch_key = batch_key
         self._label_key = label_key
         self._n_jobs = n_jobs
@@ -183,6 +192,12 @@ def prepare(self) -> None:
 
     def benchmark(self) -> None:
         """Run the pipeline."""
+        if self._benchmarked:
+            warnings.warn(
+                "The benchmark has already been run. Running it again will overwrite the previous results.",
+                UserWarning,
+            )
+
         if not self._prepared:
             self.prepare()
 
@@ -193,13 +208,12 @@ def benchmark(self) -> None:
         for emb_key, ad in tqdm(self._emb_adatas.items(), desc="Embeddings", position=0, colour="green"):
             pbar = tqdm(total=num_metrics, desc="Metrics", position=1, leave=False, colour="blue")
             for metric_type, metric_collection in self._metric_collection_dict.items():
-                for metric_name, use_metric in asdict(metric_collection).items():
-                    if use_metric:
-                        if isinstance(metric_name, str):
-                            metric_fn = getattr(scib_metrics, metric_name)
-                        else:
-                            # Callable in this case
-                            metric_fn = use_metric
+                for metric_name, use_metric_or_kwargs in asdict(metric_collection).items():
+                    if use_metric_or_kwargs:
+                        metric_fn = getattr(scib_metrics, metric_name)
+                        if isinstance(use_metric_or_kwargs, dict):
+                            # Kwargs in this case
+                            metric_fn = partial(metric_fn, **use_metric_or_kwargs)
                         metric_value = getattr(MetricAnnDataAPI, metric_name)(ad, metric_fn)
                         # nmi/ari metrics return a dict
                         if isinstance(metric_value, dict):
@@ -211,6 +225,8 @@ def benchmark(self) -> None:
                             self._results.loc[metric_name, _METRIC_TYPE] = metric_type
                         pbar.update(1)
 
+        self._benchmarked = True
+
     def get_results(self, min_max_scale: bool = True, clean_names: bool = True) -> pd.DataFrame:
         """Return the benchmarking results.
 
@@ -242,6 +258,7 @@ def get_results(self, min_max_scale: bool = True, clean_names: bool = True) -> p
 
         # Compute scores
         per_class_score = df.groupby(_METRIC_TYPE).mean().transpose()
+        # This is the default scIB weighting from the manuscript
         per_class_score["Total"] = 0.4 * per_class_score["Batch correction"] + 0.6 * per_class_score["Bio conservation"]
         df = pd.concat([df.transpose(), per_class_score], axis=1)
         df.loc[_METRIC_TYPE, per_class_score.columns] = _AGGREGATE_SCORE

src/scib_metrics/utils/_diffusion_nn.py

@@ -113,7 +113,12 @@ def diffusion_nn(X: csr_matrix, k: int, n_comps: int = 100):
     Neighbors graph
     """
     transitions = _compute_transitions(X)
-    _, embedding = _compute_eigen(transitions, n_comps=n_comps)
+    evals, evecs = _compute_eigen(transitions, n_comps=n_comps)
+    evals += 1e-8  # Avoid division by zero
+    # Multiscale such that the number of steps t gets "integrated out"
+    # First eigenvalue is 1, so we start at the second one
+    embedding = evecs
+    embedding[:, 1:] = (evals[1:] / (1 - evals[1:])) * embedding[:, 1:]
     nn_obj = pynndescent.NNDescent(embedding, n_neighbors=k + 1)
     neigh_inds, neigh_distances = nn_obj.neighbor_graph
     # We purposely ignore the first neighbor as it is the cell itself

tests/test_benchmarker.py

@@ -1,8 +1,5 @@
-from functools import partial
-
 import pandas as pd
 
-from scib_metrics import clisi_knn
 from scib_metrics.benchmark import BatchCorrection, Benchmarker, BioConservation
 from tests.utils.data import dummy_benchmarker_adata
 
@@ -36,7 +33,7 @@ def test_benchmarker_custom_metric_booleans():
 
 
 def test_benchmarker_custom_metric_callable():
-    bioc = BioConservation(clisi_knn=partial(clisi_knn, perplexity=10))
+    bioc = BioConservation(clisi_knn={"perplexity": 10})
     ad, emb_keys, batch_key, labels_key = dummy_benchmarker_adata()
     bm = Benchmarker(ad, batch_key, labels_key, emb_keys, bio_conservation_metrics=bioc)
     bm.benchmark()