Pass manager to data loader (#1280)

* pass manager to data loader

* address comment

scvi/data/anndata/_compat.py

@@ -1,13 +1,13 @@
 from anndata import AnnData
 
 from . import _constants
-from ._manager import AnnDataManager
 from .fields import (
     CategoricalJointObsField,
     CategoricalObsField,
     LayerField,
     NumericalJointObsField,
 )
+from .manager import AnnDataManager
 
 
 def manager_from_setup_dict(

scvi/data/anndata/fields/_layer_field.py

@@ -34,7 +34,7 @@ def __init__(
             else _constants._ADATA_ATTRS.LAYERS
         )
         self._attr_key = layer
-        self._is_count_data = is_count_data
+        self.is_count_data = is_count_data
 
     @property
     def registry_key(self):
@@ -56,7 +56,7 @@ def validate_field(self, adata: AnnData) -> None:
         super().validate_field(adata)
         x = self.get_field(adata)
 
-        if self._is_count_data and not _check_nonnegative_integers(x):
+        if self.is_count_data and not _check_nonnegative_integers(x):
             logger_data_loc = (
                 "adata.X" if self.attr_key is None else f"adata.layers[{self.attr_key}]"
             )

scvi/data/anndata/_manager.py → scvi/data/anndata/manager.py

@@ -70,14 +70,14 @@ def add_field(self, field: Type[BaseAnnDataField]) -> None:
         ), "Fields have been frozen. Create a new AnnDataManager object for additional fields."
         self.fields.add(field)
 
-    def _register_fields(
+    def register_fields(
         self,
         adata: AnnData,
         source_setup_dict: Optional[dict] = None,
         **transfer_kwargs
     ):
         """
-        Helper function with registers each field associated with this instance.
+        Registers each field associated with this instance with the AnnData object.
 
         Either registers or transfers the setup from `source_setup_dict` if passed in.
 
@@ -114,10 +114,6 @@ def _register_fields(
 
         self._assign_uuid()
 
-    def register_fields(self, adata: AnnData):
-        """Registers each field associated with this instance with the AnnData object."""
-        return self._register_fields(adata)
-
     def transfer_setup(
         self, adata_target: AnnData, source_setup_dict: Optional[dict] = None, **kwargs
     ) -> AnnDataManager:
@@ -143,7 +139,7 @@ def transfer_setup(
         )
         fields = self.fields
         new_adata_manager = self.__class__(fields)
-        new_adata_manager._register_fields(adata_target, setup_dict, **kwargs)
+        new_adata_manager.register_fields(adata_target, setup_dict, **kwargs)
         return new_adata_manager
 
     def get_adata_uuid(self) -> UUID:

scvi/dataloaders/_ann_dataloader.py

@@ -2,11 +2,12 @@
 import logging
 from typing import Optional, Union
 
-import anndata
 import numpy as np
 import torch
 from torch.utils.data import DataLoader
 
+from scvi.data.anndata.manager import AnnDataManager
+
 from ._anntorchdataset import AnnTorchDataset
 
 logger = logging.getLogger(__name__)
@@ -91,8 +92,8 @@ class AnnDataLoader(DataLoader):
 
     Parameters
     ----------
-    adata
-        An anndata objects
+    adata_manager
+        AnnDataManager object that has been created via setup_anndata.
     shuffle
         Whether the data should be shuffled
     indices
@@ -109,7 +110,7 @@ class AnnDataLoader(DataLoader):
 
     def __init__(
         self,
-        adata: anndata.AnnData,
+        adata_manager: AnnDataManager,
         shuffle=False,
         indices=None,
         batch_size=128,
@@ -118,11 +119,11 @@ def __init__(
         **data_loader_kwargs,
     ):
 
-        if "_scvi" not in adata.uns.keys():
+        if adata_manager.adata is None:
             raise ValueError("Please run setup_anndata() on your anndata object first.")
 
         if data_and_attributes is not None:
-            data_registry = adata.uns["_scvi"]["data_registry"]
+            data_registry = adata_manager.get_data_registry()
             for key in data_and_attributes.keys():
                 if key not in data_registry.keys():
                     raise ValueError(
@@ -131,7 +132,9 @@ def __init__(
                         )
                     )
 
-        self.dataset = AnnTorchDataset(adata, getitem_tensors=data_and_attributes)
+        self.dataset = AnnTorchDataset(
+            adata_manager.adata, getitem_tensors=data_and_attributes
+        )
 
         sampler_kwargs = {
             "batch_size": batch_size,

scvi/dataloaders/_anntorchdataset.py

@@ -58,7 +58,6 @@ def setup_getitem(self):
         ----------
         getitem_tensors:
             Either a list of keys in the scvi data registry to return when getitem is called
-            or
 
         Examples
         --------

scvi/dataloaders/_concat_dataloader.py

@@ -2,9 +2,10 @@
 from typing import List, Optional, Union
 
 import numpy as np
-from anndata import AnnData
 from torch.utils.data import DataLoader
 
+from scvi.data.anndata.manager import AnnDataManager
+
 from ._ann_dataloader import AnnDataLoader
 
 
@@ -14,8 +15,8 @@ class ConcatDataLoader(DataLoader):
 
     Parameters
     ----------
-    adata
-        AnnData object that has been registered via setup_anndata.
+    adata_manager
+        AnnDataManager object that has been created via setup_anndata.
     indices_list
         List where each element is a list of indices in the adata to load
     shuffle
@@ -32,7 +33,7 @@ class ConcatDataLoader(DataLoader):
 
     def __init__(
         self,
-        adata: AnnData,
+        adata_manager: AnnDataManager,
         indices_list: List[List[int]],
         shuffle: bool = False,
         batch_size: int = 128,
@@ -44,7 +45,7 @@ def __init__(
         for indices in indices_list:
             self.dataloaders.append(
                 AnnDataLoader(
-                    adata,
+                    adata_manager,
                     indices=indices,
                     shuffle=shuffle,
                     batch_size=batch_size,