diff --git a/CHANGELOG.md b/CHANGELOG.md index e3fed9d..1db6db8 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/README.md b/README.md index aeb6551..650ddbb 100644 --- a/README.md +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index a0dbfcc..3ce45c8 100644 --- a/pyproject.toml +++ b/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", ] diff --git a/src/scib_rapids/metrics/_nmi_ari.py b/src/scib_rapids/metrics/_nmi_ari.py index 044a0cd..67b4621 100644 --- a/src/scib_rapids/metrics/_nmi_ari.py +++ b/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, :] diff --git a/tests/test_cross_validate.py b/tests/test_cross_validate.py index 12c52bf..bd579d9 100644 --- a/tests/test_cross_validate.py +++ b/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,7 +276,8 @@ 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 = """\ @@ -281,9 +285,9 @@ def test_nmi_ari_leiden(self, shared_data): 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 --- diff --git a/tests/test_metrics.py b/tests/test_metrics.py index 2d9286b..f18ba4e 100644 --- a/tests/test_metrics.py +++ b/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,6 +95,7 @@ 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"] @@ -97,6 +103,47 @@ def test_nmi_ari_cluster_labels_leiden(): 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)