Experimental Feature: Function delegation to rapids_singlecell when gpu is available

pyproject.toml

@@ -102,6 +102,12 @@ optional-dependencies.docs = [
   "sphinxcontrib-bibtex>=2.3",
   "sphinxcontrib-spelling>=7.6.2",
 ]
+optional-dependencies.gpu-cuda11 = [
+  "rapids-singlecell[rapids11]",
+]
+optional-dependencies.gpu-cuda12 = [
+  "rapids-singlecell[rapids12]",
+]
 optional-dependencies.leiden = [
   "leidenalg",
   "spatialleiden>=0.4",

src/squidpy/__init__.py

@@ -4,6 +4,7 @@
 from importlib.metadata import PackageMetadata
 
 from squidpy import datasets, experimental, gr, im, pl, read, tl
+from squidpy.settings import settings
 
 try:
     md: PackageMetadata = metadata.metadata(__name__)
@@ -15,4 +16,4 @@
 
 del metadata, md
 
-__all__ = ["datasets", "experimental", "gr", "im", "pl", "read", "tl"]
+__all__ = ["datasets", "experimental", "gr", "im", "pl", "read", "tl", "settings"]

src/squidpy/_utils.py

@@ -19,6 +19,8 @@
 import xarray as xr
 from spatialdata.models import Image2DModel, Labels2DModel
 
+from squidpy.settings import DeviceType
+
 __all__ = ["singledispatchmethod", "Signal", "SigQueue", "NDArray", "NDArrayA"]
 
 
@@ -387,3 +389,36 @@ def _ensure_dim_order(img_da: xr.DataArray, order: Literal["cyx", "yxc"] = "yxc"
         img_da = img_da.expand_dims({"c": [0]})
     # After possible expand, just transpose to target
     return img_da.transpose(*tuple(order))
+
+
+def resolve_device_arg(device: DeviceType | None) -> Literal["cpu", "gpu"]:
+    """
+    Resolve per-call device argument to actual backend.
+
+    Parameters
+    ----------
+    device
+        Per-call device setting. None uses ``settings.device``.
+
+    Returns
+    -------
+    Literal["cpu", "gpu"]
+        The resolved backend to use.
+
+    Raises
+    ------
+    RuntimeError
+        If GPU is requested but rapids-singlecell is not installed.
+    """
+    from squidpy.settings import settings
+
+    if device is None:
+        device = settings.device
+    if device == "cpu":
+        return "cpu"
+    if device == "gpu":
+        if not settings.gpu_available():
+            raise RuntimeError("GPU unavailable. Install with: pip install squidpy[gpu-cuda12] or squidpy[gpu-cuda11]")
+        return "gpu"
+    # if device == "auto"
+    return "gpu" if settings.gpu_available() else "cpu"

src/squidpy/gr/_ppatterns.py

@@ -23,7 +23,7 @@
 from squidpy._constants._constants import SpatialAutocorr
 from squidpy._constants._pkg_constants import Key
 from squidpy._docs import d, inject_docs
-from squidpy._utils import NDArrayA, Signal, SigQueue, _get_n_cores, parallelize
+from squidpy._utils import NDArrayA, Signal, SigQueue, _get_n_cores, parallelize, resolve_device_arg
 from squidpy.gr._utils import (
     _assert_categorical_obs,
     _assert_connectivity_key,
@@ -352,6 +352,7 @@ def co_occurrence(
     n_jobs: int | None = None,
     backend: str = "loky",
     show_progress_bar: bool = True,
+    device: Literal["cpu", "gpu"] | None = None,
 ) -> tuple[NDArrayA, NDArrayA] | None:
     """
     Compute co-occurrence probability of clusters.
@@ -369,6 +370,9 @@ def co_occurrence(
         Number of splits in which to divide the spatial coordinates in
         :attr:`anndata.AnnData.obsm` ``['{spatial_key}']``.
     %(parallelize)s
+    device
+        Device to use for computation. If ``None``, uses :attr:`squidpy.settings.device`.
+        Set to ``"gpu"`` to use rapids-singlecell GPU acceleration.
 
     Returns
     -------
@@ -381,6 +385,22 @@ def co_occurrence(
         - :attr:`anndata.AnnData.uns` ``['{cluster_key}_co_occurrence']['interval']`` - the distance thresholds
           computed at ``interval``.
     """
+    effective_device = resolve_device_arg(device)
+
+    if effective_device == "gpu":
+        from rapids_singlecell.squidpy import co_occurrence as rsc_co_occurrence
+
+        return rsc_co_occurrence(
+            adata,
+            cluster_key=cluster_key,
+            spatial_key=spatial_key,
+            interval=interval,
+            copy=copy,
+            n_splits=n_splits,
+            n_jobs=n_jobs,
+            backend=backend,
+            show_progress_bar=show_progress_bar,
+        )
 
     if isinstance(adata, SpatialData):
         adata = adata.table

src/squidpy/settings/__init__.py

@@ -0,0 +1,7 @@
+"""Squidpy settings and configuration."""
+
+from __future__ import annotations
+
+from squidpy.settings._settings import DeviceType, settings
+
+__all__ = ["settings", "DeviceType"]

src/squidpy/settings/_settings.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from contextvars import ContextVar
+from typing import Literal, get_args
+
+__all__ = ["settings", "DeviceType"]
+
+DeviceType = Literal["auto", "cpu", "gpu"]
+
+_device_var: ContextVar[DeviceType] = ContextVar("device", default="auto")
+
+
+class _SqSettings:
+    """Global settings for squidpy."""
+
+    @property
+    def device(self) -> DeviceType:
+        """Current compute device setting."""
+        return _device_var.get()
+
+    @device.setter
+    def device(self, value: DeviceType) -> None:
+        valid = get_args(DeviceType)
+        if value not in valid:
+            raise ValueError(f"Invalid device {value!r}. Must be one of: {valid}")
+        if value == "gpu" and not self.gpu_available():
+            raise RuntimeError(
+                "Cannot set device='gpu': rapids-singlecell not installed. "
+                "Install with: pip install squidpy[gpu-cuda12] or squidpy[gpu-cuda11]"
+            )
+        _device_var.set(value)
+
+    @staticmethod
+    def gpu_available() -> bool:
+        """
+        Check if GPU acceleration is available.
+
+        Returns
+        -------
+        bool
+            True if rapids-singlecell is installed and importable.
+        """
+        try:
+            import rapids_singlecell  # noqa: F401
+
+            return True
+        except ImportError:
+            return False
+
+
+settings = _SqSettings()

tests/test_gpu.py

@@ -0,0 +1,63 @@
+"""Tests for GPU functionality (skipped in CI without GPU)."""
+
+from __future__ import annotations
+
+import pytest
+
+from squidpy.settings import settings
+
+# Skip all tests in this module if GPU is not available
+pytestmark = pytest.mark.skipif(
+    not settings.gpu_available(),
+    reason="GPU tests require rapids-singlecell to be installed",
+)
+
+
+class TestGPUCoOccurrence:
+    """Test GPU-accelerated co_occurrence function."""
+
+    def test_co_occurrence_gpu(self, adata):
+        """Test co_occurrence with GPU device."""
+        import squidpy as sq
+
+        # Run with explicit GPU device
+        result = sq.gr.co_occurrence(
+            adata,
+            cluster_key="leiden",
+            copy=True,
+            device="gpu",
+        )
+
+        assert result is not None
+        arr, interval = result
+        assert arr.ndim == 3
+        assert arr.shape[1] == arr.shape[0] == adata.obs["leiden"].unique().shape[0]
+
+    def test_co_occurrence_gpu_vs_cpu(self, adata):
+        """Test that GPU and CPU results are approximately equal."""
+        import numpy as np
+
+        import squidpy as sq
+
+        # Run on CPU
+        cpu_result = sq.gr.co_occurrence(
+            adata,
+            cluster_key="leiden",
+            copy=True,
+            device="cpu",
+        )
+
+        # Run on GPU
+        gpu_result = sq.gr.co_occurrence(
+            adata,
+            cluster_key="leiden",
+            copy=True,
+            device="gpu",
+        )
+
+        cpu_arr, cpu_interval = cpu_result
+        gpu_arr, gpu_interval = gpu_result
+
+        # Results should be close (allow for floating point differences)
+        np.testing.assert_allclose(cpu_interval, gpu_interval, rtol=1e-5)
+        np.testing.assert_allclose(cpu_arr, gpu_arr, rtol=1e-5)

tests/test_settings.py

@@ -0,0 +1,35 @@
+"""Tests for squidpy.settings module."""
+
+from __future__ import annotations
+
+import pytest
+
+from squidpy.settings import settings
+
+
+class TestSettings:
+    """Test the settings module."""
+
+    def test_default_device(self):
+        """Test that default device is 'auto'."""
+        # Reset to default
+        settings.device = "auto"
+        assert settings.device == "auto"
+
+    def test_set_device_cpu(self):
+        """Test setting device to 'cpu'."""
+        settings.device = "cpu"
+        assert settings.device == "cpu"
+        settings.device = "auto"  # reset
+
+    def test_set_device_invalid(self):
+        """Test that invalid device raises ValueError."""
+        with pytest.raises(ValueError, match="Invalid device"):
+            settings.device = "invalid"
+
+    def test_set_device_gpu_without_rsc(self):
+        """Test that setting device to 'gpu' without rapids-singlecell raises RuntimeError."""
+        # This will fail if rapids-singlecell is not installed
+        if not settings.gpu_available():
+            with pytest.raises(RuntimeError, match="rapids-singlecell not installed"):
+                settings.device = "gpu"