[codex] Use cuGraph for Leiden clustering

CHANGELOG.md

@@ -15,6 +15,7 @@ and this project adheres to [Semantic Versioning][].
 - Initial release with GPU-accelerated metrics using CuPy/RAPIDS.
 - All metrics from scib-metrics ported: silhouette, LISI, kBET, NMI/ARI, graph connectivity, isolated labels, PCR comparison, BRAS.
 - CuPy-based utility functions: cdist, pdist_squareform, PCA, KMeans, silhouette_samples, Simpson index.
+- cuGraph-backed Leiden clustering for GPU-accelerated NMI/ARI computation.
 - NeighborsResults dataclass with sparse graph properties.
 - Unit tests validating correctness against scikit-learn.
 - Benchmark test suite comparing runtime against scib-metrics.

README.md

@@ -20,9 +20,11 @@
 [badge-zulip]: https://img.shields.io/badge/zulip-join_chat-brightgreen.svg
 [link-zulip]: https://scverse.zulipchat.com/
 
-GPU-accelerated metrics for benchmarking single-cell integration outputs using RAPIDS (cuML, CuPy).
+GPU-accelerated metrics for benchmarking single-cell integration outputs using RAPIDS (cuGraph, cuDF, CuPy).
 
-This package provides the same metrics as [scib-metrics](https://github.com/YosefLab/scib-metrics) but replaces JAX with [RAPIDS](https://rapids.ai/) (CuPy, cuML) for GPU acceleration. All implementations leverage CuPy for device-level computation on NVIDIA GPUs.
+This package provides the same metrics as [scib-metrics](https://github.com/YosefLab/scib-metrics)
+but replaces JAX with [RAPIDS](https://rapids.ai/) for GPU acceleration. Implementations use RAPIDS
+libraries such as CuPy and cuGraph for device-level computation on NVIDIA GPUs.
 
 ## Metrics
 

pyproject.toml

@@ -18,12 +18,13 @@ urls.Home-page = "https://github.com/maarten-devries/scib-rapids"
 dependencies = [
     "anndata",
     "cupy-cuda12x>=13.0",
+    "cudf-cu12",
+    "cugraph-cu12",
     "numpy",
     "pandas",
     "scipy",
     "scikit-learn",
     "pynndescent",
-    "igraph>0.9.0",
     "umap-learn>=0.5.0",
 ]
 

src/scib_rapids/metrics/_nmi_ari.py

@@ -1,8 +1,6 @@
 import logging
-import random
 import warnings
 
-import igraph
 import numpy as np
 from scipy.sparse import spmatrix
 from sklearn.metrics.cluster import adjusted_rand_score, normalized_mutual_info_score
@@ -20,14 +18,57 @@ def _compute_clustering_kmeans(X: np.ndarray, n_clusters: int) -> np.ndarray:
     return kmeans.labels_
 
 
+def _create_cugraph_from_sparse(connectivity_graph: spmatrix):
+    try:
+        import cudf
+        import cugraph
+    except ImportError as err:
+        raise ImportError(
+            "Leiden clustering requires RAPIDS cuGraph and cuDF. "
+            "Install the CUDA 12 packages with `pip install cugraph-cu12 cudf-cu12`."
+        ) from err
+
+    coo_graph = connectivity_graph.tocoo(copy=False)
+    df = cudf.DataFrame(
+        {
+            "source": coo_graph.row.astype(np.int32, copy=False),
+            "destination": coo_graph.col.astype(np.int32, copy=False),
+            "weight": coo_graph.data.astype(np.float32, copy=False),
+        }
+    )
+
+    graph = cugraph.Graph()
+    graph.from_cudf_edgelist(
+        df,
+        source="source",
+        destination="destination",
+        weight="weight",
+        vertices=cudf.Series(np.arange(connectivity_graph.shape[0], dtype=np.int32)),
+    )
+    return graph
+
+
 def _compute_clustering_leiden(connectivity_graph: spmatrix, resolution: float, seed: int) -> np.ndarray:
-    rng = random.Random(seed)
-    igraph.set_random_number_generator(rng)
-    g = igraph.Graph.Weighted_Adjacency(connectivity_graph, mode="directed")
-    g.to_undirected(mode="each")
-    clustering = g.community_leiden(objective_function="modularity", weights="weight", resolution=resolution)
-    clusters = clustering.membership
-    return np.asarray(clusters)
+    if connectivity_graph.shape[0] != connectivity_graph.shape[1]:
+        raise ValueError("connectivity_graph must be square")
+
+    if connectivity_graph.nnz == 0:
+        return np.arange(connectivity_graph.shape[0])
+
+    try:
+        from cugraph import leiden
+    except ImportError as err:
+        raise ImportError(
+            "Leiden clustering requires RAPIDS cuGraph. Install it with `pip install cugraph-cu12`."
+        ) from err
+
+    g = _create_cugraph_from_sparse(connectivity_graph)
+    clusters, _ = leiden(g, resolution=resolution, random_state=seed)
+    clusters = clusters.to_pandas()
+    labels = np.arange(connectivity_graph.shape[0])
+    vertices = clusters["vertex"].to_numpy(dtype=np.intp)
+    labels[vertices] = clusters["partition"].to_numpy()
+    return labels
 
 
 def _compute_nmi_ari_cluster_labels(
@@ -86,7 +127,7 @@ def nmi_ari_cluster_labels_leiden(
     resolution
         Resolution parameter (used if optimize_resolution is False).
     n_jobs
-        Number of jobs for parallelization.
+        Deprecated. Ignored by the GPU Leiden backend.
     seed
         Seed for reproducibility.
 
@@ -96,17 +137,11 @@ def nmi_ari_cluster_labels_leiden(
     """
     conn_graph = X.knn_graph_connectivities
     if optimize_resolution:
+        if n_jobs != 1:
+            warnings.warn("`n_jobs` is ignored by the GPU Leiden backend.", stacklevel=2)
         n = 10
         resolutions = np.array([2 * x / n for x in range(1, n + 1)])
-        try:
-            from joblib import Parallel, delayed
-
-            out = Parallel(n_jobs=n_jobs)(
-                delayed(_compute_nmi_ari_cluster_labels)(conn_graph, labels, r, seed=seed) for r in resolutions
-            )
-        except ImportError:
-            warnings.warn("Using for loop over clustering resolutions. `pip install joblib` for parallelization.")
-            out = [_compute_nmi_ari_cluster_labels(conn_graph, labels, r, seed=seed) for r in resolutions]
+        out = [_compute_nmi_ari_cluster_labels(conn_graph, labels, r, seed=seed) for r in resolutions]
         nmi_ari = np.array(out)
         nmi_ind = np.argmax(nmi_ari[:, 0])
         nmi, ari = nmi_ari[nmi_ind, :]

tests/test_cross_validate.py

@@ -1,4 +1,4 @@
-"""Cross-validation tests: verify scib-rapids produces the same results as scib-metrics.
+"""Cross-validation tests: verify scib-rapids produces comparable results to scib-metrics.
 
 Both backends use float32 internally (JAX defaults to float32, CuPy casts to float32),
 so small numerical differences (< 1e-4) are expected from different operation ordering
@@ -9,7 +9,8 @@
     1. Generate deterministic test data and save to a temp directory as .npy files.
     2. For each metric, run a small Python script in each venv (scib-metrics / scib-rapids)
        that loads the data, computes the metric, and saves the result as .npy.
-    3. Compare the two results with appropriate tolerances.
+    3. Compare the two results with appropriate tolerances, except for metrics that
+       intentionally use a different backend.
 
 Usage:
     pytest tests/test_cross_validate.py -v -s
@@ -92,9 +93,7 @@ def _run_in_venv(python: str, script: str, timeout: int = 120) -> str:
     )
     if result.returncode != 0:
         raise RuntimeError(
-            f"Script failed (rc={result.returncode}).\n"
-            f"--- stdout ---\n{result.stdout}\n"
-            f"--- stderr ---\n{result.stderr}"
+            f"Script failed (rc={result.returncode}).\n--- stdout ---\n{result.stdout}\n--- stderr ---\n{result.stderr}"
         )
     return result.stdout.strip()
 
@@ -167,7 +166,10 @@ def _compare(
         r = rapids_res[k]
         m = metrics_res[k]
         np.testing.assert_allclose(
-            r, m, atol=atol, rtol=rtol,
+            r,
+            m,
+            atol=atol,
+            rtol=rtol,
             err_msg=f"{label} key={k}",
         )
 
@@ -176,6 +178,7 @@ def _compare(
 # Individual metric comparison tests
 # ---------------------------------------------------------------------------
 
+
 class TestCrossValidation:
     """Verify numerical agreement between scib-rapids and scib-metrics.
 
@@ -273,17 +276,18 @@ def test_pcr_comparison(self, shared_data):
 
     # --- NMI/ARI ---
     # KMeans uses different RNGs (JAX vs NumPy), so exact match is not expected.
-    # Leiden clustering is deterministic given the same graph.
+    # Leiden uses cuGraph in scib-rapids and igraph in scib-metrics, so exact
+    # cluster agreement is not expected.
 
     def test_nmi_ari_leiden(self, shared_data):
         code = """\
 result = lib.nmi_ari_cluster_labels_leiden(nn, labels, optimize_resolution=False, resolution=1.0)
 results = {"nmi": result["nmi"], "ari": result["ari"]}
 """
         r = _run_metric(RAPIDS_PYTHON, "scib_rapids", shared_data, code)
-        m = _run_metric(METRICS_PYTHON, "scib_metrics", shared_data, code)
-        # Leiden is deterministic — should be exact
-        _compare(r, m, atol=1e-10, rtol=1e-10, label="nmi_ari_leiden")
+        for key, value in r.items():
+            assert np.isfinite(value), f"nmi_ari_leiden key={key}"
+            assert 0 <= value <= 1, f"nmi_ari_leiden key={key}"
 
     # --- Utility functions ---
 

tests/test_metrics.py

@@ -4,14 +4,19 @@
 output types, shapes, and value ranges on toy data.
 """
 
+import sys
+
 import numpy as np
+import pandas as pd
 import pytest
+import scipy.sparse as sp
 from sklearn.datasets import make_blobs
 from sklearn.metrics import silhouette_samples as sk_silhouette_samples
 from sklearn.metrics.pairwise import pairwise_distances_argmin
 from sklearn.neighbors import NearestNeighbors
 
 import scib_rapids
+from scib_rapids.metrics import _nmi_ari
 from scib_rapids.nearest_neighbors import NeighborsResults
 from tests.utils.data import dummy_x_labels, dummy_x_labels_batch
 
@@ -90,13 +95,55 @@ def test_nmi_ari_cluster_labels_kmeans():
 
 
 def test_nmi_ari_cluster_labels_leiden():
+    pytest.importorskip("cugraph")
     X, labels = dummy_x_labels(symmetric_positive=True, x_is_neighbors_results=True)
     out = scib_rapids.nmi_ari_cluster_labels_leiden(X, labels, optimize_resolution=False, resolution=0.1)
     nmi, ari = out["nmi"], out["ari"]
     assert isinstance(nmi, float)
     assert isinstance(ari, float)
 
 
+def test_compute_clustering_leiden_uses_cugraph(monkeypatch):
+    calls = {}
+
+    class FakeGraph:
+        def from_cudf_edgelist(self, df, **kwargs):
+            calls["edgelist"] = df
+            calls["edgelist_kwargs"] = kwargs
+
+    class FakeClusters:
+        def to_pandas(self):
+            return pd.DataFrame({"vertex": [1, 0], "partition": [7, 3]})
+
+    def fake_leiden(graph, **kwargs):
+        calls["graph"] = graph
+        calls["leiden_kwargs"] = kwargs
+        return FakeClusters(), 0.0
+
+    monkeypatch.setitem(
+        sys.modules,
+        "cudf",
+        type("FakeCudf", (), {"DataFrame": staticmethod(pd.DataFrame), "Series": staticmethod(pd.Series)}),
+    )
+    monkeypatch.setitem(
+        sys.modules,
+        "cugraph",
+        type("FakeCugraph", (), {"Graph": FakeGraph, "leiden": staticmethod(fake_leiden)}),
+    )
+
+    connectivity = sp.csr_matrix(np.array([[0.0, 0.25], [0.5, 0.0]], dtype=np.float32))
+    labels = _nmi_ari._compute_clustering_leiden(connectivity, resolution=0.4, seed=123)
+
+    np.testing.assert_array_equal(labels, np.array([3, 7]))
+    assert calls["edgelist_kwargs"].pop("vertices").tolist() == [0, 1]
+    assert calls["edgelist_kwargs"] == {
+        "source": "source",
+        "destination": "destination",
+        "weight": "weight",
+    }
+    assert calls["leiden_kwargs"] == {"resolution": 0.4, "random_state": 123}
+
+
 def test_kbet():
     X, _, batch = dummy_x_labels_batch(x_is_neighbors_results=True)
     acc_rate, stats, pvalues = scib_rapids.kbet(X, batch)