feat: Add RESOLVI graph dataloader integration and benchmark CI

.github/workflows/test_linux_benchmark.yml

@@ -0,0 +1,75 @@
+name: test (benchmark)
+
+on:
+  pull_request:
+    branches: [main, "[0-9]+.[0-9]+.x"]
+    types: [labeled, synchronize, opened]
+  schedule:
+    - cron: "0 10 * * *" # runs at 10:00 UTC (03:00 PST) every day
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  test:
+    # if PR has label "benchmark tests" or "all tests" or if scheduled or manually triggered or on push
+    if: >-
+      (
+        contains(github.event.pull_request.labels.*.name, 'benchmark tests') ||
+        contains(github.event.pull_request.labels.*.name, 'all tests') ||
+        (contains(github.event_name, 'schedule') && github.repository == 'scverse/scvi-tools') ||
+        contains(github.event_name, 'workflow_dispatch') ||
+        contains(github.event_name, 'push')
+      )
+
+    runs-on: ${{ matrix.os }}
+
+    defaults:
+      run:
+        shell: bash -e {0} # -e to fail on error
+
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest]
+        python: ["3.13"]
+
+    permissions:
+      id-token: write
+
+    name: unit
+
+    env:
+      OS: ${{ matrix.os }}
+      PYTHON: ${{ matrix.python }}
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python }}
+          cache: "pip"
+          cache-dependency-path: "**/pyproject.toml"
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip wheel uv
+          python -m uv pip install --system "scvi-tools[tests] @ ."
+
+      - name: Run pytest
+        env:
+          MPLBACKEND: agg
+          PLATFORM: ${{ matrix.os }}
+          DISPLAY: :42
+          COLUMNS: 120
+        run: |
+          coverage run -m pytest -v --color=yes --benchmark
+          coverage report
+
+      - uses: codecov/codecov-action@v4
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          flags: benchmark

.pre-commit-config.yaml

@@ -1,6 +1,7 @@
 fail_fast: false
 default_language_version:
   python: python3
+  node: 20.20.2
 default_stages:
   - pre-commit
   - pre-push

CHANGELOG.md

@@ -11,6 +11,7 @@ to [Semantic Versioning]. The full commit history is available in the [commit lo
 
 - Add support for Python 3.14, {pr}`3563`.
 - Add support for Pandas3, {pr}`3638`.
+- Add graph-aware dataloading support for {class}`scvi.external.RESOLVI` with `torch_geometric` {pr}`37XX`.
 
 #### Fixed
 

pyproject.toml

@@ -77,7 +77,7 @@ docs = [
     "myst-nb",
     "sphinx-autodoc-typehints",
 ]
-docsbuild = ["scvi-tools[docs,autotune,hub,jax]","mlx"]
+docsbuild = ["scvi-tools[docs,autotune,hub,jax,diagvi]","mlx"]
 
 # scvi.autotune #TODO remove ray[tune] constraint once solved
 autotune = ["hyperopt>=0.2", "ray[tune]; python_version < '3.14'", "scib-metrics", "muon"]
@@ -97,9 +97,11 @@ jax = ["jax<0.10.0", "jaxlib", "optax", "numpyro", "flax"] #TODO: unpin once it
 dataloaders = ["lamindb>=1.12.1", "cellxgene-census", "tiledbsoma", "tiledbsoma_ml", "torchdata"]
 # for mlflow
 mlflow = ["mlflow","psutil","GPUtil","nvidia-ml-py"]
+# for diagvi
+diagvi = ["torch_geometric", "geomloss"]
 
 optional = [
-    "scvi-tools[autotune,mlflow,hub,jax,file_sharing,regseq,parallel,interpretability]",
+    "scvi-tools[autotune,mlflow,hub,jax,file_sharing,regseq,parallel,interpretability,diagvi]",
     "igraph","leidenalg","pynndescent",
 ]
 tutorials = [
@@ -131,6 +133,7 @@ omit = [
 testpaths = ["tests"]
 xfail_strict = true
 markers = [
+    "benchmark: mark benchmark/performance tests",
     "internet: mark tests that requires internet access",
     "optional: mark optional tests, usually take more time",
     "private: mark tests that uses private keys, like HF",
@@ -140,6 +143,7 @@ markers = [
     "dataloader: mark tests that are used to check data loaders",
     "jax: mark test as jax related",
     "mlflow: mark test for mlflow",
+    "diagvi: mark test for diagvi and torch_geometric based models",
 ]
 
 [tool.ruff]

src/scvi/dataloaders/__init__.py

@@ -10,6 +10,7 @@
     DeviceBackedDataSplitter,
     SemiSupervisedDataSplitter,
 )
+from ._graph_dataloader import GraphDataLoader, GraphDataSplitter
 from ._samplers import BatchDistributedSampler
 from ._semi_dataloader import SemiSupervisedDataLoader
 
@@ -19,10 +20,12 @@
     "CollectionAdapter",
     "ConcatDataLoader",
     "DeviceBackedDataSplitter",
-    "SemiSupervisedDataLoader",
     "DataSplitter",
-    "SemiSupervisedDataSplitter",
+    "GraphDataLoader",
+    "GraphDataSplitter",
     "BatchDistributedSampler",
     "MappedCollectionDataModule",
+    "SemiSupervisedDataLoader",
+    "SemiSupervisedDataSplitter",
     "TileDBDataModule",
 ]

src/scvi/dataloaders/_graph_dataloader.py

@@ -0,0 +1,227 @@
+"""Graph-aware dataloaders for spatial single-cell models."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+import torch
+from torch.utils.data import default_convert
+
+from scvi import REGISTRY_KEYS
+from scvi.data import AnnTorchDataset
+from scvi.dataloaders._ann_dataloader import AnnDataLoader
+from scvi.dataloaders._data_splitting import DataSplitter
+
+if TYPE_CHECKING:
+    from scvi.data import AnnDataManager
+
+
+def _as_torch_tensor(array: np.ndarray | torch.Tensor) -> torch.Tensor:
+    """Convert AnnTorchDataset output to a torch tensor without changing sparse layout."""
+    if isinstance(array, np.ndarray):
+        return torch.from_numpy(array)
+    return array
+
+
+class _GraphBatchConverter:
+    """Convert an AnnTorchDataset batch into a PyG Data object."""
+
+    def __init__(
+        self,
+        full_adata_manager: AnnDataManager,
+        neighbor_indices_key: str,
+        edge_obsm_keys: list[str],
+        load_sparse_neighbor_tensor: bool,
+        load_neighbor_expression: bool,
+    ):
+        self.neighbor_indices_key = neighbor_indices_key
+        self.edge_obsm_keys = edge_obsm_keys
+        self.load_neighbor_expression = load_neighbor_expression
+        if load_neighbor_expression:
+            self._full_dataset = AnnTorchDataset(
+                full_adata_manager,
+                getitem_tensors=[REGISTRY_KEYS.X_KEY],
+                load_sparse_tensor=load_sparse_neighbor_tensor,
+            )
+
+    def __call__(self, batch: dict[str, np.ndarray | torch.Tensor]):
+        try:
+            from torch_geometric.data import Data
+        except ImportError as error:
+            raise ImportError(
+                "torch_geometric is required for GraphDataLoader. "
+                "Install it with: pip install torch_geometric"
+            ) from error
+
+        batch = default_convert(batch)
+        ind_neighbors = batch[self.neighbor_indices_key].long()
+        n_obs, n_neighbors = ind_neighbors.shape
+
+        x = _as_torch_tensor(batch[REGISTRY_KEYS.X_KEY])
+
+        center_idx = torch.arange(n_obs, dtype=torch.long).repeat_interleave(n_neighbors)
+        neighbor_idx = torch.arange(n_obs * n_neighbors, dtype=torch.long)
+        edge_index = torch.stack([center_idx, neighbor_idx], dim=0)
+
+        edge_attrs = []
+        for key in self.edge_obsm_keys:
+            vals = batch[key].float()
+            edge_attrs.append(vals.reshape(n_obs * n_neighbors, -1))
+        edge_attr = torch.cat(edge_attrs, dim=1) if edge_attrs else None
+
+        data_kwargs = dict(batch)
+        data_kwargs.update(
+            {
+                "x": x,
+                "edge_index": edge_index,
+                "edge_attr": edge_attr,
+                "distances_n": batch.get("distance_neighbor"),
+            }
+        )
+        if self.load_neighbor_expression:
+            flat_neighbors = ind_neighbors.cpu().numpy().ravel()
+            data_kwargs["x_n"] = _as_torch_tensor(
+                self._full_dataset[flat_neighbors][REGISTRY_KEYS.X_KEY]
+            )
+        return Data(**data_kwargs)
+
+
+class GraphDataLoader(AnnDataLoader):
+    """DataLoader that yields mini-batches as :class:`torch_geometric.data.Data` objects.
+
+    Each batch contains center cells and their pre-fetched spatial neighbors. Neighbor
+    expression is looked up from ``full_adata_manager`` so neighbors outside the current
+    train/validation/test split are intentionally allowed, matching existing RESOLVI behavior.
+
+    Parameters
+    ----------
+    adata_manager
+        :class:`~scvi.data.AnnDataManager` for the split being loaded.
+    full_adata_manager
+        :class:`~scvi.data.AnnDataManager` for all observations. Used for neighbor expression
+        lookup, including cross-split neighbors.
+    indices
+        Observation indices to load from ``adata_manager``.
+    neighbor_indices_key
+        Registry key containing neighbor indices, shape ``[N, K]``.
+    edge_obsm_keys
+        Registry keys to flatten and concatenate into ``edge_attr``. Each key must have shape
+        ``[N, K]`` or ``[N, K, D]``. Defaults to ``["distance_neighbor"]``.
+    load_sparse_neighbor_tensor
+        If ``True``, loads sparse neighbor expression as sparse torch tensors. This avoids
+        densifying neighbor expression on the CPU before device transfer.
+    load_neighbor_expression
+        If ``False``, omits ``x_n`` and leaves neighbor expression gathering to the model. This is
+        useful when a model keeps a device-resident expression cache.
+    **kwargs
+        Forwarded to :class:`~scvi.dataloaders.AnnDataLoader`.
+    """
+
+    def __init__(
+        self,
+        adata_manager: AnnDataManager,
+        full_adata_manager: AnnDataManager,
+        indices: list[int] | list[bool] | None = None,
+        neighbor_indices_key: str = "index_neighbor",
+        edge_obsm_keys: list[str] | None = None,
+        load_sparse_neighbor_tensor: bool = True,
+        load_neighbor_expression: bool = True,
+        **kwargs,
+    ):
+        if "collate_fn" in kwargs:
+            raise ValueError("GraphDataLoader uses its own collate_fn to build graph batches.")
+        if kwargs.pop("iter_ndarray", False):
+            raise ValueError("GraphDataLoader does not support `iter_ndarray=True`.")
+
+        super().__init__(adata_manager, indices=indices, **kwargs)
+        self.neighbor_indices_key = neighbor_indices_key
+        self.edge_obsm_keys = (
+            list(edge_obsm_keys) if edge_obsm_keys is not None else ["distance_neighbor"]
+        )
+        self.load_sparse_neighbor_tensor = load_sparse_neighbor_tensor
+        self.load_neighbor_expression = load_neighbor_expression
+        self._graph_batch_converter = _GraphBatchConverter(
+            full_adata_manager,
+            neighbor_indices_key=self.neighbor_indices_key,
+            edge_obsm_keys=self.edge_obsm_keys,
+            load_sparse_neighbor_tensor=load_sparse_neighbor_tensor,
+            load_neighbor_expression=load_neighbor_expression,
+        )
+        self.collate_fn = self._graph_batch_converter
+
+
+class GraphDataSplitter(DataSplitter):
+    """DataSplitter that creates :class:`GraphDataLoader` instances.
+
+    Parameters
+    ----------
+    load_sparse_neighbor_tensor
+        Forwarded to :class:`GraphDataLoader`.
+    load_neighbor_expression
+        Forwarded to :class:`GraphDataLoader`.
+    """
+
+    def __init__(
+        self,
+        adata_manager: AnnDataManager,
+        neighbor_indices_key: str = "index_neighbor",
+        edge_obsm_keys: list[str] | None = None,
+        load_sparse_neighbor_tensor: bool = True,
+        load_neighbor_expression: bool = True,
+        **kwargs,
+    ):
+        super().__init__(adata_manager, **kwargs)
+        self.neighbor_indices_key = neighbor_indices_key
+        self.edge_obsm_keys = (
+            list(edge_obsm_keys) if edge_obsm_keys is not None else ["distance_neighbor"]
+        )
+        self.load_sparse_neighbor_tensor = load_sparse_neighbor_tensor
+        self.load_neighbor_expression = load_neighbor_expression
+
+    def _make_graph_dataloader(
+        self,
+        indices: np.ndarray,
+        shuffle: bool,
+        drop_last: bool,
+    ) -> GraphDataLoader:
+        return GraphDataLoader(
+            self.adata_manager,
+            full_adata_manager=self.adata_manager,
+            indices=indices,
+            shuffle=shuffle,
+            drop_last=drop_last,
+            load_sparse_tensor=self.load_sparse_tensor,
+            pin_memory=self.pin_memory,
+            neighbor_indices_key=self.neighbor_indices_key,
+            edge_obsm_keys=self.edge_obsm_keys,
+            load_sparse_neighbor_tensor=self.load_sparse_neighbor_tensor,
+            load_neighbor_expression=self.load_neighbor_expression,
+            **self.data_loader_kwargs,
+        )
+
+    def train_dataloader(self) -> GraphDataLoader:
+        """Create graph train dataloader."""
+        return self._make_graph_dataloader(
+            self.train_idx,
+            shuffle=True,
+            drop_last=self.drop_last,
+        )
+
+    def val_dataloader(self) -> GraphDataLoader | None:
+        """Create graph validation dataloader."""
+        if len(self.val_idx) > 0:
+            return self._make_graph_dataloader(
+                self.val_idx,
+                shuffle=False,
+                drop_last=False,
+            )
+
+    def test_dataloader(self) -> GraphDataLoader | None:
+        """Create graph test dataloader."""
+        if len(self.test_idx) > 0:
+            return self._make_graph_dataloader(
+                self.test_idx,
+                shuffle=False,
+                drop_last=False,
+            )