added some fixes to the custom dataloader stuff

src/scvi/data/_manager.py

@@ -192,6 +192,55 @@ def register_fields(
         self._assign_uuid()
         self._assign_most_recent_manager_uuid()
 
+    def register_data_module_fields(
+        self,
+        datamodule,
+        source_registry: dict | None = None,
+        **transfer_kwargs,
+    ):
+        """Registers each field associated with this instance with the AnnData object.
+
+        Either registers or transfers the setup from `source_setup_dict` if passed in.
+        Sets ``self.adata``.
+
+        Parameters
+        ----------
+        adata
+            AnnData object to be registered.
+        source_registry
+            Registry created after registering an AnnData using an
+            :class:`~scvi.data.AnnDataManager` object.
+        transfer_kwargs
+            Additional keywords which modify transfer behavior. Only applicable if
+            ``source_registry`` is set.
+        """
+        if self.adata is not None:
+            raise AssertionError("Existing AnnData object registered with this Manager instance.")
+
+        if source_registry is None and transfer_kwargs:
+            raise TypeError(
+                f"register_fields() got unexpected keyword arguments {transfer_kwargs} passed "
+                "without a source_registry."
+            )
+
+        self._validate_anndata_object(datamodule)
+
+        for field in self.fields:
+            self._add_field(
+                field=field,
+                adata=datamodule,
+                source_registry=source_registry,
+                **transfer_kwargs,
+            )
+
+        # Save arguments for register_fields.
+        self._source_registry = deepcopy(source_registry)
+        self._transfer_kwargs = deepcopy(transfer_kwargs)
+
+        self.adata = datamodule
+        self._assign_uuid()
+        self._assign_most_recent_manager_uuid()
+
     def _add_field(
         self,
         field: AnnDataField,

src/scvi/model/_scvi.py

@@ -140,8 +140,10 @@ def __init__(
             f"gene_likelihood: {gene_likelihood}, latent_distribution: {latent_distribution}."
         )
 
+        # in the next part we need to construct the same module no mather the way
+        # dataloader was given
         if self._module_init_on_train:
-            # Here we need to adjust given the new custom data loader
+            # Here we need to adjust given the new custom data loader like CZI case
             self.module = None
             warnings.warn(
                 "Model was initialized without `adata`. The module will be initialized when "
@@ -225,6 +227,68 @@ def setup_anndata(
         adata_manager = AnnDataManager(fields=anndata_fields, setup_method_args=setup_method_args)
         adata_manager.register_fields(adata, **kwargs)
         cls.register_manager(adata_manager)
+        # adata_manager.get_state_registry(SCVI.REGISTRY_KEYS.X_KEY).to_dict()
+        # adata_manager.registry[_constants._FIELD_REGISTRIES_KEY]
+        # pprint(adata_manager.registry)
+
+    @classmethod
+    @setup_anndata_dsp.dedent
+    def setup_datamodule(
+        cls,
+        datamodule,
+        layer: str | None = None,
+        batch_key: str | None = None,
+        labels_key: str | None = None,
+        size_factor_key: str | None = None,
+        categorical_covariate_keys: list[str] | None = None,
+        continuous_covariate_keys: list[str] | None = None,
+        **kwargs,
+    ):
+        """%(summary)s.
+
+        Parameters
+        ----------
+        %(param_adata)s
+        %(param_layer)s
+        %(param_batch_key)s
+        %(param_labels_key)s
+        %(param_size_factor_key)s
+        %(param_cat_cov_keys)s
+        %(param_cont_cov_keys)s
+        """
+        setup_method_args = cls._get_setup_method_args(**locals())
+        anndata_fields = [
+            LayerField(REGISTRY_KEYS.X_KEY, layer, is_count_data=True),
+            CategoricalObsField(REGISTRY_KEYS.BATCH_KEY, batch_key),
+            CategoricalObsField(REGISTRY_KEYS.LABELS_KEY, labels_key),
+            NumericalObsField(REGISTRY_KEYS.SIZE_FACTOR_KEY, size_factor_key, required=False),
+            CategoricalJointObsField(REGISTRY_KEYS.CAT_COVS_KEY, categorical_covariate_keys),
+            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.registry["setup_method_name"] = "setup_datamodule"
+        adata_manager.registry["setup_args"]["layer"] = datamodule.datapipe.layer_name
+        adata_manager.registry["setup_args"]["batch_key"] = datamodule.batch_keys
+        adata_manager.registry["setup_args"]["labels_key"]
+        adata_manager.registry["setup_args"]["batch_key"]
+        adata_manager.registry["setup_args"]["batch_key"]
+        adata_manager.registry["setup_args"]["batch_key"]
+        # datamodule._datapipe.obs_column_names
+        # datamodule._datapipe.obs_encoders
+        # adata_manager.register_fields(adata, **kwargs)
+        # how to etract the information we need from the datamodule
+        adata_manager.register_data_module_fields(
+            datamodule, **kwargs
+        )  # here we need a new function for data module
+
+        cls.register_manager(adata_manager)
+        # adata_manager.get_state_registry(SCVI.REGISTRY_KEYS.X_KEY).to_dict()
+        # adata_manager.registry[_constants._FIELD_REGISTRIES_KEY]
+        # pprint(adata_manager.registry)
 
     @staticmethod
     def _get_fields_for_adata_minification(

src/scvi/model/base/_base_model.py

@@ -815,6 +815,23 @@ def setup_anndata(
         on a model-specific instance of :class:`~scvi.data.AnnDataManager`.
         """
 
+    @classmethod
+    @abstractmethod
+    @setup_anndata_dsp.dedent
+    def setup_datamodule(
+        cls,
+        datamodule,
+        *args,
+        **kwargs,
+    ):
+        """%(summary)s.
+
+        Each model class deriving from this class provides parameters to this method
+        according to its needs. To operate correctly with the model initialization,
+        the implementation must call :meth:`~scvi.model.base.BaseModelClass.register_manager`
+        on a model-specific instance of :class:`~scvi.data.AnnDataManager`.
+        """
+
     @staticmethod
     def view_setup_args(dir_path: str, prefix: str | None = None) -> None:
         """Print args used to setup a saved model.

src/scvi/model/base/_training_mixin.py

@@ -102,6 +102,7 @@ def train(
                 )
 
         if datamodule is None:
+            # In the general case we enter here
             datasplitter_kwargs = datasplitter_kwargs or {}
             datamodule = self._data_splitter_cls(
                 self.adata_manager,
@@ -114,6 +115,7 @@ def train(
                 **datasplitter_kwargs,
             )
         elif self.module is None:
+            # in CZI case we enter here
             self.module = self._module_cls(
                 datamodule.n_vars,
                 n_batch=datamodule.n_batch,
@@ -122,6 +124,8 @@ def train(
                 n_cats_per_cov=getattr(datamodule, "n_cats_per_cov", None),
                 **self._module_kwargs,
             )
+        # after either of the cases we should be here with the same self.module
+        # and same datamodule
 
         plan_kwargs = plan_kwargs or {}
         training_plan = self._training_plan_cls(self.module, **plan_kwargs)

tests/dataloaders/test_custom_dataloader.py

@@ -1 +1,64 @@
 from __future__ import annotations
+
+import os
+
+import numpy as np
+import scanpy as sc
+
+import scvi
+from scvi.data import _constants, synthetic_iid
+from scvi.model import SCVI
+
+# We will now create the SCVI model object:
+# Its parameters:
+n_layers = 1
+n_latent = 10
+batch_size = 1024
+train_size = 0.9
+max_epochs = 1
+
+
+# COMAPRE TO THE ORIGINAL METHOD!!! - use the same data!!!
+# We first create a registry using the orignal way of anndata in order to compare and add
+# what is missing
+adata = synthetic_iid()
+adata.obs["size_factor"] = np.random.randint(1, 5, size=(adata.shape[0],))
+SCVI.setup_anndata(
+    adata,
+    batch_key="batch",
+    labels_key="labels",
+    size_factor_key="size_factor",
+)
+#
+model_orig = SCVI(adata, n_latent=n_latent)
+model_orig.train(1, check_val_every_n_epoch=1, train_size=0.5)
+
+# Saving the model
+save_dir = "/Users/orikr/runs/290724/"  # tempfile.TemporaryDirectory()
+model_dir = os.path.join(save_dir, "scvi_orig_model")
+model_orig.save(model_dir, overwrite=True)
+
+# Loading the model (just as a compariosn)
+model_orig_loaded = scvi.model.SCVI.load(model_dir, adata=adata)
+
+# Obtaining model outputs
+SCVI_LATENT_KEY = "X_scVI"
+latent = model_orig.get_latent_representation()
+adata.obsm[SCVI_LATENT_KEY] = latent
+# latent.shape
+
+# You can see all necessary entries and the structure at
+adata_manager = model_orig.adata_manager
+model_orig.view_anndata_setup(hide_state_registries=True)
+# adata_manager.get_state_registry(SCVI.REGISTRY_KEYS.X_KEY).to_dict()
+adata_manager.registry[_constants._FIELD_REGISTRIES_KEY]
+
+# Plot UMAP and save the figure for later check
+sc.pp.neighbors(adata, use_rep="scvi", key_added="scvi")
+sc.tl.umap(adata, neighbors_key="scvi")
+sc.pl.umap(adata, color="dataset_id", title="SCVI")
+
+# Now return and add all the registry stuff that we will need
+
+# Now add the missing stuff from the current CZI implemenation in order for us to have the exact
+# same steps like the original way (except than setup_anndata)