feat(dataloaders): Custom dataloader registry support

.github/workflows/test_linux.yml

@@ -53,7 +53,7 @@ jobs:
           DISPLAY: :42
           COLUMNS: 120
         run: |
-          coverage run -m pytest -v --color=yes
+          coverage run -m pytest -v --color=yes -m "not custom_dataloader"
           coverage report
 
       - uses: codecov/codecov-action@v4

.github/workflows/test_linux_custom_dataloader.yml

@@ -0,0 +1,89 @@
+name: test (custom dataloaders)
+
+on:
+  push:
+    branches: [main, "[0-9]+.[0-9]+.x"]
+  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 "custom_dataloader" or "all tests" or if scheduled or manually triggered
+    if: >-
+      (
+        contains(github.event.pull_request.labels.*.name, 'custom_dataloader') ||
+        contains(github.event.pull_request.labels.*.name, 'all tests') ||
+        contains(github.event_name, 'schedule') ||
+        contains(github.event_name, 'workflow_dispatch')
+      )
+
+    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.11"]
+
+    name: integration
+
+    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] @ ."
+          python -m pip install scdataloader
+          python -m pip install cellxgene-census
+          python -m pip install tiledbsoma
+          python -m pip install s3fs
+          python -m pip install torchdata==0.9.0
+          python -m pip install psutil
+          python -m pip install lamindb
+          python -m pip install bionty==0.51.0
+          python -m pip install biomart
+
+      - name: Install Specific Branch of Repository
+        env:
+          GH_TOKEN: ${{ secrets.GH_TOKEN }}
+        run: |
+          git config --global url."https://${GH_TOKEN}:[email protected]/".insteadOf "https://github.com/"
+          git clone --single-branch --branch ebezzi/census-scvi-datamodule https://github.com/ori-kron-wis/cellxgene-census.git
+          git clone --single-branch --branch main https://github.com/jkobject/scDataLoader.git
+
+      - name: Run specific custom dataloader pytest
+        env:
+          MPLBACKEND: agg
+          PLATFORM: ${{ matrix.os }}
+          DISPLAY: :42
+          COLUMNS: 120
+        run: |
+          coverage run -m pytest tests/dataloaders/test_custom_dataloader.py -v --color=yes --custom-dataloader-tests
+          coverage report
+
+      - uses: codecov/codecov-action@v4
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}

CHANGELOG.md

@@ -14,6 +14,7 @@ to [Semantic Versioning]. Full commit history is available in the
     representation learning in single-cell RNA sequencing data {pr}`3015`, {pr}`3091`.
 - Add {class}`scvi.external.RESOLVI` for bias correction in single-cell resolved spatial
     transcriptomics {pr}`3144`.
+- Add support for using Lamin custom dataloaders with {class}`scvi.model.SCVI`, {pr}`2932`.
 
 #### Fixed
 

pyproject.toml

@@ -95,9 +95,11 @@ regseq = ["biopython>=1.81", "genomepy"]
 scanpy = ["scanpy>=1.10", "scikit-misc"]
 # for convinient files sharing
 pooch = ["pooch"]
+# for custom dataloders
+dataloaders = ["lamindb","biomart","bionty","cellxgene_lamin"]
 
 optional = [
-    "scvi-tools[autotune,aws,hub,pooch,regseq,scanpy]"
+    "scvi-tools[autotune,aws,hub,pooch,regseq,scanpy,dataloaders]"
 ]
 tutorials = [
     "cell2location",

src/scvi/data/_utils.py

@@ -16,6 +16,7 @@
 from torch import as_tensor, sparse_csc_tensor, sparse_csr_tensor
 
 from scvi import REGISTRY_KEYS, settings
+from scvi.utils import attrdict
 
 from . import _constants
 
@@ -150,6 +151,14 @@ def _set_data_in_registry(
         setattr(adata, attr_name, attribute)
 
 
+def _get_summary_stats_from_registry(registry: dict) -> attrdict:
+    summary_stats = {}
+    for field_registry in registry[_constants._FIELD_REGISTRIES_KEY].values():
+        field_summary_stats = field_registry[_constants._SUMMARY_STATS_KEY]
+        summary_stats.update(field_summary_stats)
+    return attrdict(summary_stats)
+
+
 def _verify_and_correct_data_format(adata: AnnData, attr_name: str, attr_key: str | None):
     """Check data format and correct if necessary.
 

src/scvi/dataloaders/__init__.py

@@ -3,6 +3,7 @@
 
 from ._ann_dataloader import AnnDataLoader
 from ._concat_dataloader import ConcatDataLoader
+from ._custom_dataloders import MappedCollectionDataModule
 from ._data_splitting import (
     DataSplitter,
     DeviceBackedDataSplitter,
@@ -20,4 +21,5 @@
     "DataSplitter",
     "SemiSupervisedDataSplitter",
     "BatchDistributedSampler",
+    "MappedCollectionDataModule",
 ]

src/scvi/dataloaders/_custom_dataloders.py

@@ -0,0 +1,194 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import psutil
+from lightning.pytorch import LightningDataModule
+from torch.utils.data import DataLoader
+
+import scvi
+
+if TYPE_CHECKING:
+    import lamindb as ln
+    import numpy as np
+
+
+class MappedCollectionDataModule(LightningDataModule):
+    def __init__(
+        self,
+        collection: ln.Collection,
+        batch_key: str | None = None,
+        label_key: str | None = None,
+        batch_size: int = 128,
+        **kwargs,
+    ):
+        self._batch_size = batch_size
+        self._batch_key = batch_key
+        self._label_key = label_key
+        self._parallel = kwargs.pop("parallel", True)
+        # here we initialize MappedCollection to use in a pytorch DataLoader
+        self._dataset = collection.mapped(
+            obs_keys=self._batch_key, parallel=self._parallel, **kwargs
+        )
+        # need by scvi and lightning.pytorch
+        self._log_hyperparams = False
+        self.allow_zero_length_dataloader_with_multiple_devices = False
+
+    def close(self):
+        self._dataset.close()
+
+    def setup(self, stage):
+        pass
+
+    def train_dataloader(self):
+        return self._create_dataloader(shuffle=True)
+
+    def inference_dataloader(self):
+        """Dataloader for inference with `on_before_batch_transfer` applied."""
+        dataloader = self._create_dataloader(shuffle=False, batch_size=4096)
+        return self._InferenceDataloader(dataloader, self.on_before_batch_transfer)
+
+    def _create_dataloader(self, shuffle, batch_size=None):
+        if self._parallel:
+            num_workers = psutil.cpu_count() - 1
+            worker_init_fn = self._dataset.torch_worker_init_fn
+        else:
+            num_workers = 0
+            worker_init_fn = None
+        if batch_size is None:
+            batch_size = self._batch_size
+        return DataLoader(
+            self._dataset,
+            batch_size=batch_size,
+            shuffle=shuffle,
+            num_workers=num_workers,
+            worker_init_fn=worker_init_fn,
+        )
+
+    @property
+    def n_obs(self) -> int:
+        return self._dataset.n_obs
+
+    @property
+    def var_names(self) -> int:
+        return self._dataset.var_joint
+
+    @property
+    def n_vars(self) -> int:
+        return self._dataset.n_vars
+
+    @property
+    def n_batch(self) -> int:
+        if self._batch_key is None:
+            return 1
+        return len(self._dataset.encoders[self._batch_key])
+
+    @property
+    def n_labels(self) -> int:
+        if self._label_key is None:
+            return 1
+        return len(self._dataset.encoders[self._label_key])
+
+    @property
+    def labels(self) -> np.ndarray:
+        return self._dataset[self._label_key]
+
+    @property
+    def registry(self) -> dict:
+        return {
+            "scvi_version": scvi.__version__,
+            "model_name": "SCVI",
+            "setup_args": {
+                "layer": None,
+                "batch_key": self._batch_key,
+                "labels_key": self._label_key,
+                "size_factor_key": None,
+                "categorical_covariate_keys": None,
+                "continuous_covariate_keys": None,
+            },
+            "field_registries": {
+                "X": {
+                    "data_registry": {"attr_name": "X", "attr_key": None},
+                    "state_registry": {
+                        "n_obs": self.n_obs,
+                        "n_vars": self.n_vars,
+                        "column_names": self.var_names,
+                    },
+                    "summary_stats": {"n_vars": self.n_vars, "n_cells": self.n_obs},
+                },
+                "batch": {
+                    "data_registry": {"attr_name": "obs", "attr_key": "_scvi_batch"},
+                    "state_registry": {
+                        "categorical_mapping": self.batch_keys,
+                        "original_key": self._batch_key,
+                    },
+                    "summary_stats": {"n_batch": self.n_batch},
+                },
+                "labels": {
+                    "data_registry": {"attr_name": "obs", "attr_key": "_scvi_labels"},
+                    "state_registry": {
+                        "categorical_mapping": self.label_keys,
+                        "original_key": self._label_key,
+                        "unlabeled_category": "unlabeled",
+                    },
+                    "summary_stats": {"n_labels": self.n_labels},
+                },
+                "size_factor": {
+                    "data_registry": {},
+                    "state_registry": {},
+                    "summary_stats": {},
+                },
+                "extra_categorical_covs": {
+                    "data_registry": {},
+                    "state_registry": {},
+                    "summary_stats": {"n_extra_categorical_covs": 0},
+                },
+                "extra_continuous_covs": {
+                    "data_registry": {},
+                    "state_registry": {},
+                    "summary_stats": {"n_extra_continuous_covs": 0},
+                },
+            },
+            "setup_method_name": "setup_anndata",
+        }
+
+    @property
+    def batch_keys(self) -> int:
+        if self._batch_key is None:
+            return None
+        return self._dataset.encoders[self._batch_key]
+
+    @property
+    def label_keys(self) -> int:
+        if self._label_key is None:
+            return None
+        return self._dataset.encoders[self._label_key]
+
+    def on_before_batch_transfer(
+        self,
+        batch,
+        dataloader_idx,
+    ):
+        X_KEY: str = "X"
+        BATCH_KEY: str = "batch"
+        LABEL_KEY: str = "labels"
+
+        return {
+            X_KEY: batch["X"].float(),
+            BATCH_KEY: batch[self._batch_key][:, None] if self._batch_key is not None else None,
+            LABEL_KEY: 0,
+        }
+
+    class _InferenceDataloader:
+        """Wrapper to apply `on_before_batch_transfer` during iteration."""
+
+        def __init__(self, dataloader, transform_fn):
+            self.dataloader = dataloader
+            self.transform_fn = transform_fn
+
+        def __iter__(self):
+            for batch in self.dataloader:
+                yield self.transform_fn(batch, dataloader_idx=None)
+
+        def __len__(self):
+            return len(self.dataloader)

src/scvi/dataloaders/_data_splitting.py

@@ -386,7 +386,8 @@ class is :class:`~scvi.dataloaders.SemiSupervisedDataLoader`,
 
     def __init__(
         self,
-        adata_manager: AnnDataManager,
+        adata_manager: AnnDataManager | None = None,
+        datamodule: pl.LightningDataModule | None = None,
         train_size: float | None = None,
         validation_size: float | None = None,
         shuffle_set_split: bool = True,

src/scvi/external/resolvi/_model.py

@@ -98,7 +98,8 @@ class RESOLVI(
 
     def __init__(
         self,
-        adata: AnnData,
+        adata: AnnData | None,
+        registry: dict | None = None,
         n_hidden: int = 32,
         n_hidden_encoder: int = 128,
         n_latent: int = 10,

src/scvi/external/stereoscope/_model.py

@@ -53,7 +53,8 @@ class RNAStereoscope(UnsupervisedTrainingMixin, BaseModelClass):
 
     def __init__(
         self,
-        sc_adata: AnnData,
+        sc_adata: AnnData | None = None,
+        registry: dict | None = None,
         **model_kwargs,
     ):
         super().__init__(sc_adata)

src/scvi/external/stereoscope/_module.py

@@ -140,6 +140,7 @@ def __init__(
         n_spots: int,
         sc_params: tuple[np.ndarray],
         prior_weight: Literal["n_obs", "minibatch"] = "n_obs",
+        **model_kwargs,
     ):
         super().__init__()
         # unpack and copy parameters