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
+          python -m pip install psutil
+          python -m pip install lamindb
+          python -m pip install bionty
+          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 }}

pyproject.toml

@@ -112,6 +112,9 @@ tutorials = [
     "scvi-tools[optional]",
     "squidpy",
 ]
+dataloaders = [
+    "scdataloader"
+]
 
 all = ["scvi-tools[dev,docs,tutorials]"]
 

src/scvi/data/_utils.py

@@ -10,6 +10,8 @@
 import scipy.sparse as sp_sparse
 from anndata import AnnData
 
+from scvi.utils import attrdict
+
 try:
     # anndata >= 0.10
     from anndata.experimental import CSCDataset, CSRDataset
@@ -162,6 +164,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/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

src/scvi/model/_amortizedlda.py

@@ -61,7 +61,8 @@ class AmortizedLDA(PyroSviTrainMixin, BaseModelClass):
 
     def __init__(
         self,
-        adata: AnnData,
+        adata: AnnData | None = None,
+        registry: dict | None = None,
         n_topics: int = 20,
         n_hidden: int = 128,
         cell_topic_prior: float | Sequence[float] | None = None,

src/scvi/model/_autozi.py

@@ -104,7 +104,8 @@ class AUTOZI(VAEMixin, UnsupervisedTrainingMixin, BaseModelClass):
 
     def __init__(
         self,
-        adata: AnnData,
+        adata: AnnData | None = None,
+        registry: dict | None = None,
         n_hidden: int = 128,
         n_latent: int = 10,
         n_layers: int = 1,

src/scvi/model/_condscvi.py

@@ -67,7 +67,8 @@ class CondSCVI(RNASeqMixin, VAEMixin, UnsupervisedTrainingMixin, BaseModelClass)
 
     def __init__(
         self,
-        adata: AnnData,
+        adata: AnnData | None = None,
+        registry: dict | None = None,
         n_hidden: int = 128,
         n_latent: int = 5,
         n_layers: int = 2,

src/scvi/model/_jaxscvi.py

@@ -59,7 +59,8 @@ class JaxSCVI(JaxTrainingMixin, BaseModelClass):
 
     def __init__(
         self,
-        adata: AnnData,
+        adata: AnnData | None = None,
+        registry: dict | None = None,
         n_hidden: int = 128,
         n_latent: int = 10,
         dropout_rate: float = 0.1,

src/scvi/model/_linear_scvi.py

@@ -78,7 +78,8 @@ class LinearSCVI(RNASeqMixin, VAEMixin, UnsupervisedTrainingMixin, BaseModelClas
 
     def __init__(
         self,
-        adata: AnnData,
+        adata: AnnData | None = None,
+        registry: dict | None = None,
         n_hidden: int = 128,
         n_latent: int = 10,
         n_layers: int = 1,

src/scvi/model/_multivi.py

@@ -140,9 +140,10 @@ class MULTIVI(VAEMixin, UnsupervisedTrainingMixin, BaseModelClass, ArchesMixin):
 
     def __init__(
         self,
-        adata: AnnData,
-        n_genes: int,
-        n_regions: int,
+        adata: AnnData | None = None,
+        registry: dict | None = None,
+        n_genes: int | None = None,
+        n_regions: int | None = None,
         modality_weights: Literal["equal", "cell", "universal"] = "equal",
         modality_penalty: Literal["Jeffreys", "MMD", "None"] = "Jeffreys",
         n_hidden: int | None = None,

src/scvi/model/_peakvi.py

@@ -92,7 +92,8 @@ class PEAKVI(ArchesMixin, VAEMixin, UnsupervisedTrainingMixin, BaseModelClass):
 
     def __init__(
         self,
-        adata: AnnData,
+        adata: AnnData | None = None,
+        registry: dict | None = None,
         n_hidden: int | None = None,
         n_latent: int | None = None,
         n_layers_encoder: int = 2,

src/scvi/model/_scanvi.py

@@ -118,7 +118,8 @@ class SCANVI(RNASeqMixin, VAEMixin, ArchesMixin, BaseMinifiedModeModelClass):
 
     def __init__(
         self,
-        adata: AnnData,
+        adata: AnnData | None = None,
+        registry: dict | None = None,
         n_hidden: int = 128,
         n_latent: int = 10,
         n_layers: int = 1,
@@ -128,23 +129,29 @@ def __init__(
         linear_classifier: bool = False,
         **model_kwargs,
     ):
-        super().__init__(adata)
+        super().__init__(adata, registry)
         scanvae_model_kwargs = dict(model_kwargs)
 
         self._set_indices_and_labels()
 
         # ignores unlabeled catgegory
         n_labels = self.summary_stats.n_labels - 1
-        n_cats_per_cov = (
-            self.adata_manager.get_state_registry(REGISTRY_KEYS.CAT_COVS_KEY).n_cats_per_key
-            if REGISTRY_KEYS.CAT_COVS_KEY in self.adata_manager.data_registry
-            else None
-        )
+        if adata is not None:
+            n_cats_per_cov = (
+                self.adata_manager.get_state_registry(REGISTRY_KEYS.CAT_COVS_KEY).n_cats_per_key
+                if REGISTRY_KEYS.CAT_COVS_KEY in self.adata_manager.data_registry
+                else None
+            )
+        else:
+            # custom datamodule
+            n_cats_per_cov = self.summary_stats[f"n_{REGISTRY_KEYS.CAT_COVS_KEY}"]
+            if n_cats_per_cov == 0:
+                n_cats_per_cov = None
 
         n_batch = self.summary_stats.n_batch
-        use_size_factor_key = REGISTRY_KEYS.SIZE_FACTOR_KEY in self.adata_manager.data_registry
+        use_size_factor_key = self.registry_["setup_args"][f"{REGISTRY_KEYS.SIZE_FACTOR_KEY}_key"]
         library_log_means, library_log_vars = None, None
-        if not use_size_factor_key and self.minified_data_type is None:
+        if self.adata is not None and not use_size_factor_key and self.minified_data_type is None:
             library_log_means, library_log_vars = _init_library_size(self.adata_manager, n_batch)
 
         self.module = self._module_cls(
@@ -187,6 +194,7 @@ def from_scvi_model(
         unlabeled_category: str,
         labels_key: str | None = None,
         adata: AnnData | None = None,
+        registry: dict | None = None,
         **scanvi_kwargs,
     ):
         """Initialize scanVI model with weights from pretrained :class:`~scvi.model.SCVI` model.
@@ -203,6 +211,8 @@ def from_scvi_model(
             Value used for unlabeled cells in `labels_key` used to setup AnnData with scvi.
         adata
             AnnData object that has been registered via :meth:`~scvi.model.SCANVI.setup_anndata`.
+        registry
+            Registry of the datamodule used to train scANVI model.
         scanvi_kwargs
             kwargs for scANVI model
         """
@@ -237,7 +247,7 @@ def from_scvi_model(
             # validate new anndata against old model
             scvi_model._validate_anndata(adata)
 
-        scvi_setup_args = deepcopy(scvi_model.adata_manager.registry[_SETUP_ARGS_KEY])
+        scvi_setup_args = deepcopy(scvi_model.registry[_SETUP_ARGS_KEY])
         scvi_labels_key = scvi_setup_args["labels_key"]
         if labels_key is None and scvi_labels_key is None:
             raise ValueError(
@@ -250,7 +260,8 @@ def from_scvi_model(
             unlabeled_category=unlabeled_category,
             **scvi_setup_args,
         )
-        scanvi_model = cls(adata, **non_kwargs, **kwargs, **scanvi_kwargs)
+
+        scanvi_model = cls(adata, scvi_model.registry, **non_kwargs, **kwargs, **scanvi_kwargs)
         scvi_state_dict = scvi_model.module.state_dict()
         scanvi_model.module.load_state_dict(scvi_state_dict, strict=False)
         scanvi_model.was_pretrained = True
@@ -259,7 +270,7 @@ def from_scvi_model(
 
     def _set_indices_and_labels(self):
         """Set indices for labeled and unlabeled cells."""
-        labels_state_registry = self.adata_manager.get_state_registry(REGISTRY_KEYS.LABELS_KEY)
+        labels_state_registry = self.get_state_registry(REGISTRY_KEYS.LABELS_KEY)
         self.original_label_key = labels_state_registry.original_key
         self.unlabeled_category_ = labels_state_registry.unlabeled_category
 
@@ -479,12 +490,15 @@ def setup_anndata(
             NumericalJointObsField(REGISTRY_KEYS.CONT_COVS_KEY, continuous_covariate_keys),
         ]
         # register new fields if the adata is minified
-        adata_minify_type = _get_adata_minify_type(adata)
-        if adata_minify_type is not None:
-            anndata_fields += cls._get_fields_for_adata_minification(adata_minify_type)
-        adata_manager = AnnDataManager(fields=anndata_fields, setup_method_args=setup_method_args)
-        adata_manager.register_fields(adata, **kwargs)
-        cls.register_manager(adata_manager)
+        if adata:
+            adata_minify_type = _get_adata_minify_type(adata)
+            if adata_minify_type is not None:
+                anndata_fields += cls._get_fields_for_adata_minification(adata_minify_type)
+            adata_manager = AnnDataManager(
+                fields=anndata_fields, setup_method_args=setup_method_args
+            )
+            adata_manager.register_fields(adata, **kwargs)
+            cls.register_manager(adata_manager)
 
     @staticmethod
     def _get_fields_for_adata_minification(