Fixes error in get losses functions

docs/release_notes/index.md

@@ -49,6 +49,11 @@ is available in the [commit logs](https://github.com/scverse/scvi-tools/commits/
     by initializing with {class}`scvi.train.SaveCheckpoint` {pr}`2317`.
 -   {attr}`scvi.settings.dl_num_workers` is now correctly applied as the default
     `num_workers` in {class}`scvi.dataloaders.AnnDataLoader` {pr}`2322`.
+-   Add argument `return_mean` to {meth}`scvi.model.base.VAEMixin.get_reconstruction_error`
+    and {meth}`scvi.model.base.VAEMixin.get_elbo` to allow computation
+    without averaging across cells {pr}`2362`.
+-   Add support for setting `weights="importance"` in
+    {meth}`scvi.model.SCANVI.differential_expression` {pr}`2362`.
 
 #### Fixed
 
@@ -57,6 +62,7 @@ is available in the [commit logs](https://github.com/scverse/scvi-tools/commits/
 -   Fix bug in {class}`scvi.module.SCANVAE` where classifier probabilities
     were interpreted as logits. This is backwards compatible as loading older
     models will use the old code path {pr}`2301`.
+-   Fix {meth}`scvi.module.VAE.marginal_ll` when `n_mc_samples_per_pass=1` {pr}`2362`.
 -   Fix bug in {class}`scvi.external.GIMVI` where `batch_size` was not
     properly used in inference methods {pr}`2366`.
 -   Fix error message formatting in {meth}`scvi.data.fields.LayerField.transfer_field` {pr}`2368`.

scvi/external/scar/_module.py

@@ -343,7 +343,7 @@ def loss(
                 generative_outputs["pl"],
             ).sum(dim=1)
         else:
-            kl_divergence_l = 0.0
+            kl_divergence_l = torch.zeros_like(kl_divergence_z)
 
         # need to add the ambient rate and scale to the distribution for the loss
         px = generative_outputs["px"]

scvi/model/base/_log_likelihood.py

@@ -1,8 +1,9 @@
 """File for computing log likelihood of the data."""
+import numpy as np
 import torch
 
 
-def compute_elbo(vae, data_loader, feed_labels=True, **kwargs):
+def compute_elbo(vae, data_loader, feed_labels=True, return_mean=True, **kwargs):
     """Computes the ELBO.
 
     The ELBO is the reconstruction error + the KL divergences
@@ -13,22 +14,34 @@ def compute_elbo(vae, data_loader, feed_labels=True, **kwargs):
     It still gives good insights on the modeling of the data, and is fast to compute.
     """
     # Iterate once over the data and compute the elbo
-    elbo = 0
+    if return_mean:
+        elbo = 0
+    else:
+        elbo = np.array([])
     for tensors in data_loader:
         _, _, scvi_loss = vae(tensors, **kwargs)
 
-        recon_loss = scvi_loss.reconstruction_loss_sum
-        kl_local = scvi_loss.kl_local_sum
-        elbo += (recon_loss + kl_local).item()
+        recon_loss = np.sum(
+            [np.array(i) for i in scvi_loss.reconstruction_loss.values()], axis=0
+        )
+        kl_local = np.sum([np.array(i) for i in scvi_loss.kl_local.values()], axis=0)
 
-    kl_global = scvi_loss.kl_global_sum
+        if return_mean:
+            elbo += (recon_loss + kl_local).sum(0).item()
+        else:
+            elbo = np.concatenate((elbo, recon_loss + kl_local), axis=0)
+
+    kl_global = np.sum([np.array(i) for i in scvi_loss.kl_global.values()], axis=0)
     n_samples = len(data_loader.indices)
-    elbo += kl_global
-    return elbo / n_samples
+    if return_mean:
+        elbo += kl_global
+        return elbo / n_samples
+    else:
+        return elbo + kl_global / n_samples
 
 
 # do each one
-def compute_reconstruction_error(vae, data_loader, **kwargs):
+def compute_reconstruction_error(vae, data_loader, return_mean=True, **kwargs):
     """Computes log p(x/z), which is the reconstruction error.
 
     Differs from the marginal log likelihood, but still gives good
@@ -44,13 +57,24 @@ def compute_reconstruction_error(vae, data_loader, **kwargs):
         else:
             rec_loss_dict = losses.reconstruction_loss
         for key, value in rec_loss_dict.items():
-            if key in log_lkl:
-                log_lkl[key] += torch.sum(value).item()
-            else:
-                log_lkl[key] = torch.sum(value).item()
+            if return_mean:
+                if key in log_lkl:
+                    if return_mean:
+                        log_lkl[key] += torch.sum(value).item()
+                    else:
+                        log_lkl[key].append(value)
+                else:
+                    if return_mean:
+                        log_lkl[key] = torch.sum(value).item()
+                    else:
+                        log_lkl[key] = value
 
     n_samples = len(data_loader.indices)
     for key, _ in log_lkl.items():
-        log_lkl[key] = log_lkl[key] / n_samples
+        if return_mean:
+            log_lkl[key] = log_lkl[key] / n_samples
+        else:
+            log_lkl[key] = torch.cat(log_lkl[key], dim=0)
+
         log_lkl[key] = -log_lkl[key]
     return log_lkl

scvi/model/base/_vaemixin.py

@@ -23,6 +23,7 @@ def get_elbo(
         adata: Optional[AnnData] = None,
         indices: Optional[Sequence[int]] = None,
         batch_size: Optional[int] = None,
+        return_mean: bool = True,
     ) -> float:
         """Return the ELBO for the data.
 
@@ -38,12 +39,15 @@ def get_elbo(
             Indices of cells in adata to use. If `None`, all cells are used.
         batch_size
             Minibatch size for data loading into model. Defaults to `scvi.settings.batch_size`.
+        return_mean
+            If False, return the ELBO for each observation.
+            Otherwise, return the mean ELBO.
         """
         adata = self._validate_anndata(adata)
         scdl = self._make_data_loader(
             adata=adata, indices=indices, batch_size=batch_size
         )
-        elbo = compute_elbo(self.module, scdl)
+        elbo = compute_elbo(self.module, scdl, return_mean=return_mean)
         return -elbo
 
     @torch.inference_mode()
@@ -75,7 +79,7 @@ def get_marginal_ll(
             Minibatch size for data loading into model. Defaults to `scvi.settings.batch_size`.
         return_mean
             If False, return the marginal log likelihood for each observation.
-            Otherwise, return the mmean arginal log likelihood.
+            Otherwise, return the mean marginal log likelihood.
         """
         adata = self._validate_anndata(adata)
         if indices is None:
@@ -98,7 +102,7 @@ def get_marginal_ll(
                     )
                 )
             if not return_mean:
-                return torch.cat(log_lkl, 0)
+                return torch.cat(log_lkl, dim=0)
             else:
                 return np.mean(log_lkl)
         else:
@@ -114,6 +118,7 @@ def get_reconstruction_error(
         adata: Optional[AnnData] = None,
         indices: Optional[Sequence[int]] = None,
         batch_size: Optional[int] = None,
+        return_mean: Optional[bool] = True,
     ) -> float:
         r"""Return the reconstruction error for the data.
 
@@ -129,12 +134,17 @@ def get_reconstruction_error(
             Indices of cells in adata to use. If `None`, all cells are used.
         batch_size
             Minibatch size for data loading into model. Defaults to `scvi.settings.batch_size`.
+        return_mean
+            If False, return the reconstruction loss for each observation.
+            Otherwise, return the mean reconstruction loss.
         """
         adata = self._validate_anndata(adata)
         scdl = self._make_data_loader(
             adata=adata, indices=indices, batch_size=batch_size
         )
-        reconstruction_error = compute_reconstruction_error(self.module, scdl)
+        reconstruction_error = compute_reconstruction_error(
+            self.module, scdl, return_mean=return_mean
+        )
         return reconstruction_error
 
     @torch.inference_mode()

scvi/module/_autozivae.py

@@ -396,7 +396,7 @@ def loss(
                 Normal(local_library_log_means, torch.sqrt(local_library_log_vars)),
             ).sum(dim=1)
         else:
-            kl_divergence_l = 0.0
+            kl_divergence_l = torch.zeros_like(kl_divergence_z)
 
         # KL divergence wrt Bernoulli parameters
         kl_divergence_bernoulli = self.compute_global_kl_divergence()

scvi/module/_scanvae.py

@@ -332,8 +332,7 @@ def loss(
                 Normal(local_library_log_means, torch.sqrt(local_library_log_vars)),
             ).sum(dim=1)
         else:
-            kl_divergence_l = 0.0
-
+            kl_divergence_l = torch.zeros_like(loss_z1_weight)
         if is_labelled:
             loss = reconst_loss + loss_z1_weight + loss_z1_unweight
             kl_locals = {

scvi/module/_totalvae.py

@@ -631,7 +631,7 @@ def loss(
                 Normal(local_library_log_means, torch.sqrt(local_library_log_vars)),
             ).sum(dim=1)
         else:
-            kl_div_l_gene = 0.0
+            kl_div_l_gene = torch.zeros_like(kl_div_z)
 
         kl_div_back_pro_full = kl(
             Normal(py_["back_alpha"], py_["back_beta"]), self.back_mean_prior

scvi/module/_vae.py

@@ -473,7 +473,7 @@ def loss(
                 generative_outputs["pl"],
             ).sum(dim=1)
         else:
-            kl_divergence_l = torch.tensor(0.0, device=x.device)
+            kl_divergence_l = torch.zeros_like(kl_divergence_z)
 
         reconst_loss = -generative_outputs["px"].log_prob(x).sum(-1)
 
@@ -609,6 +609,8 @@ def marginal_ll(
                 q_l_x = ql.log_prob(library).sum(dim=-1)
 
                 log_prob_sum += p_l - q_l_x
+            if n_mc_samples_per_pass == 1:
+                log_prob_sum = log_prob_sum.unsqueeze(0)
 
             to_sum.append(log_prob_sum)
         to_sum = torch.cat(to_sum, dim=0)

tests/model/test_scanvi.py

@@ -237,6 +237,17 @@ def test_multiple_covariates_scanvi():
     m.get_latent_representation()
     m.get_elbo()
     m.get_marginal_ll(n_mc_samples=3)
+    # m.get_marginal_ll(adata, return_mean=True, n_mc_samples=6, n_mc_samples_per_pass=1)
+    m.differential_expression(
+        idx1=np.arange(50), idx2=51 + np.arange(50), mode="vanilla", weights="uniform"
+    )
+    m.differential_expression(
+        idx1=np.arange(50),
+        idx2=51 + np.arange(50),
+        mode="vanilla",
+        weights="importance",
+        importance_weighting_kwargs={"n_mc_samples": 10, "n_mc_samples_per_pass": 1},
+    )
     m.get_reconstruction_error()
     m.get_normalized_expression(n_samples=1)
     m.get_normalized_expression(n_samples=2)

tests/model/test_scvi.py

@@ -211,8 +211,11 @@ def test_scvi(n_latent: int = 5):
     assert z.shape == (adata.shape[0], n_latent)
     assert len(model.history["elbo_train"]) == 2
     model.get_elbo()
+    model.get_elbo(return_mean=False)
     model.get_marginal_ll(n_mc_samples=3)
+    model.get_marginal_ll(n_mc_samples=3, return_mean=False)
     model.get_reconstruction_error()
+    model.get_reconstruction_error(return_mean=False)
     model.get_normalized_expression(transform_batch="batch_1")
     model.get_normalized_expression(n_samples=2)