Backport PR scverse#1702: Quick fix in poisson sample() function for vae

meeseeksmachine · Sep 20, 2022 · e490809 · e490809
1 parent 3e2525b
commit e490809
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/scvi/module/_vae.py b/scvi/module/_vae.py
@@ -461,7 +461,7 @@ def sample(
 
         dist = generative_outputs["px"]
         if self.gene_likelihood == "poisson":
-            l_train = generative_outputs["px"].mu
+            l_train = generative_outputs["px"].rate
             l_train = torch.clamp(l_train, max=1e8)
             dist = torch.distributions.Poisson(
                 l_train

diff --git a/tests/models/test_models.py b/tests/models/test_models.py
@@ -362,6 +362,14 @@ def test_scvi(save_path):
     model = SCVI(adata, gene_likelihood="nb")
     model.get_likelihood_parameters()
 
+    # test different gene_likelihoods
+    for gene_likelihood in ["zinb", "nb", "poisson"]:
+        model = SCVI(adata, gene_likelihood=gene_likelihood)
+        model.train(1, check_val_every_n_epoch=1, train_size=0.5)
+        model.posterior_predictive_sample()
+        model.get_latent_representation()
+        model.get_normalized_expression()
+
     # test train callbacks work
     a = synthetic_iid()
     SCVI.setup_anndata(