add view registry tests

scverse · Feb 5, 2022 · 1f9ccc5 · 1f9ccc5
1 parent 953cb17
commit 1f9ccc5
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 2 deletions.
diff --git a/scvi/model/_totalvi.py b/scvi/model/_totalvi.py
@@ -1221,7 +1221,8 @@ def setup_anndata(
             ProteinObsmField(
                 REGISTRY_KEYS.PROTEIN_EXP_KEY,
                 protein_expression_obsm_key,
-                batch_field.attr_key,
+                use_batch_mask=True,
+                batch_key=batch_field.attr_key,
                 colnames_uns_key=protein_names_uns_key,
                 is_count_data=True,
             ),

diff --git a/tests/dataset/test_anndata.py b/tests/dataset/test_anndata.py
@@ -361,6 +361,25 @@ def test_anntorchdataset_getitem():
         assert type(value) == np.ndarray
 
 
+def test_view_registry():
+    adata = synthetic_iid()
+    adata.obs["cont1"] = np.random.normal(size=(adata.shape[0],))
+    adata.obs["cont2"] = np.random.normal(size=(adata.shape[0],))
+    adata.obs["cat1"] = np.random.randint(0, 5, size=(adata.shape[0],))
+    adata.obs["cat2"] = np.random.randint(0, 5, size=(adata.shape[0],))
+    adata_manager = generic_setup_adata_manager(
+        adata,
+        batch_key="batch",
+        labels_key="labels",
+        protein_expression_obsm_key="protein_expression",
+        protein_names_uns_key="protein_names",
+        continuous_covariate_keys=["cont1", "cont2"],
+        categorical_covariate_keys=["cat1", "cat2"],
+    )
+    adata_manager.view_registry()
+    adata_manager.view_registry(hide_state_registries=True)
+
+
 def test_saving(save_path):
     save_path = os.path.join(save_path, "tmp_adata.h5ad")
     adata = synthetic_iid()

diff --git a/tests/dataset/utils.py b/tests/dataset/utils.py
@@ -54,7 +54,8 @@ def generic_setup_adata_manager(
             ProteinObsmField(
                 REGISTRY_KEYS.PROTEIN_EXP_KEY,
                 protein_expression_obsm_key,
-                batch_field.attr_key,
+                use_batch_mask=True,
+                batch_key=batch_field.attr_key,
                 colnames_uns_key=protein_names_uns_key,
                 is_count_data=True,
             )

diff --git a/tests/models/test_models.py b/tests/models/test_models.py
@@ -109,6 +109,9 @@ def test_scvi(save_path):
 
     # tests __repr__
     print(model)
+    # test view_registry
+    model.view_anndata_setup()
+    model.view_anndata_setup(hide_state_registries=True)
 
     assert model.is_trained is True
     z = model.get_latent_representation()
@@ -120,13 +123,21 @@ def test_scvi(save_path):
     model.get_normalized_expression(transform_batch="batch_1")
 
     adata2 = synthetic_iid()
+    # test view_registry with different anndata before transfer setup
+    with pytest.raises(ValueError):
+        model.view_anndata_setup(adata=adata2)
+        model.view_anndata_setup(adata=adata2, hide_state_registries=True)
+    # test get methods with different anndata
     model.get_elbo(adata2)
     model.get_marginal_ll(adata2, n_mc_samples=3)
     model.get_reconstruction_error(adata2)
     latent = model.get_latent_representation(adata2, indices=[1, 2, 3])
     assert latent.shape == (3, n_latent)
     denoised = model.get_normalized_expression(adata2)
     assert denoised.shape == adata.shape
+    # test view_registry with different anndata after transfer setup
+    model.view_anndata_setup(adata=adata2)
+    model.view_anndata_setup(adata=adata2, hide_state_registries=True)
 
     denoised = model.get_normalized_expression(
         adata2, indices=[1, 2, 3], transform_batch="batch_1"