Thorough use of distributions to clean module-level code

docs/release_notes/v0.17.0.md

@@ -3,6 +3,7 @@
 ## Changes
 
 - Experimental MuData support for {class}`~scvi.model.TOTALVI` via the method {meth}`~scvi.model.TOTALVI.setup_mudata`. For several of the existing `AnnDataField` classes, there is now a MuData counterpart with an additional `mod_key` argument used to indicate the modality where the data lives (e.g. {class}`~scvi.data.fields.LayerField` to {class}`~scvi.data.fields.MuDataLayerField`). These modified classes are simply wrapped versions of the original `AnnDataField` code via  the new {method}`scvi.data.fields.MuDataWrapper` method [#1474].
+- Modification of the `generative` method's outputs to return prior and likelihood properties as `torch.Distribution` objects. Concerned modules are `_amortizedlda.py`, `_autozivae.py`, `multivae.py`, `_peakvae.py`, `_scanvae.py`, `_vae.py`, and `_vaec.py`. This allows facilitating the manipulation of these distributions for model training and inference [#1356].
 
 ## Breaking changes
 
@@ -12,8 +13,11 @@
 
 - [@jjhong922]
 - [@adamgayoso]
+- [@PierreBoyeau]
 
 [#1474]: https://github.com/YosefLab/scvi-tools/pull/1474
+[#1356]: https://github.com/YosefLab/scvi-tools/pull/1356
 
 [@jjhong922]: https://github.com/jjhong922
 [@adamgayoso]: https://github.com/adamgayoso
+[@pierreboyeau]: https://github.com/PierreBoyeau

scvi/distributions/__init__.py

@@ -2,6 +2,7 @@
     JaxNegativeBinomialMeanDisp,
     NegativeBinomial,
     NegativeBinomialMixture,
+    Poisson,
     ZeroInflatedNegativeBinomial,
 )
 
@@ -10,4 +11,5 @@
     "NegativeBinomialMixture",
     "ZeroInflatedNegativeBinomial",
     "JaxNegativeBinomialMeanDisp",
+    "Poisson",
 ]

scvi/distributions/_negative_binomial.py

@@ -8,7 +8,9 @@
 import torch.nn.functional as F
 from numpyro.distributions import constraints as numpyro_constraints
 from numpyro.distributions.util import promote_shapes, validate_sample
-from torch.distributions import Distribution, Gamma, Poisson, constraints
+from torch.distributions import Distribution, Gamma
+from torch.distributions import Poisson as PoissonTorch
+from torch.distributions import constraints
 from torch.distributions.utils import (
     broadcast_all,
     lazy_property,
@@ -236,6 +238,33 @@ def _gamma(theta, mu):
     return gamma_d
 
 
+class Poisson(PoissonTorch):
+    """
+    Poisson distribution.
+
+    Parameters
+    ----------
+    rate
+        rate of the Poisson distribution.
+    validate_args : optional
+        whether to validate input.
+    scale : optional
+        Normalized mean expression of the distribution.
+        This optional parameter is not used in any computations, but allows to store
+        normalization expression levels.
+
+    """
+
+    def __init__(
+        self,
+        rate: torch.Tensor,
+        validate_args: Optional[bool] = None,
+        scale: Optional[torch.Tensor] = None,
+    ):
+        super().__init__(rate=rate, validate_args=validate_args)
+        self.scale = scale
+
+
 class NegativeBinomial(Distribution):
     r"""
     Negative binomial distribution.
@@ -262,7 +291,9 @@ class NegativeBinomial(Distribution):
         Mean of the distribution.
     theta
         Inverse dispersion.
-    validate_args
+    scale : optional
+        Normalized mean expression of the distribution.
+    validate_args : optional
         Raise ValueError if arguments do not match constraints
     """
 
@@ -279,6 +310,7 @@ def __init__(
         logits: Optional[torch.Tensor] = None,
         mu: Optional[torch.Tensor] = None,
         theta: Optional[torch.Tensor] = None,
+        scale: Optional[torch.Tensor] = None,
         validate_args: bool = False,
     ):
         self._eps = 1e-8
@@ -299,6 +331,7 @@ def __init__(
             mu, theta = broadcast_all(mu, theta)
         self.mu = mu
         self.theta = theta
+        self.scale = scale
         super().__init__(validate_args=validate_args)
 
     @property
@@ -319,7 +352,7 @@ def sample(
             # Clamping as distributions objects can have buggy behaviors when
             # their parameters are too high
             l_train = torch.clamp(p_means, max=1e8)
-            counts = Poisson(
+            counts = PoissonTorch(
                 l_train
             ).sample()  # Shape : (n_samples, n_cells_batch, n_vars)
             return counts
@@ -368,6 +401,8 @@ class ZeroInflatedNegativeBinomial(NegativeBinomial):
         Inverse dispersion.
     zi_logits
         Logits scale of zero inflation probability.
+    scale : optional
+        Normalized mean expression of the distribution.
     validate_args
         Raise ValueError if arguments do not match constraints
     """
@@ -388,6 +423,7 @@ def __init__(
         mu: Optional[torch.Tensor] = None,
         theta: Optional[torch.Tensor] = None,
         zi_logits: Optional[torch.Tensor] = None,
+        scale: Optional[torch.Tensor] = None,
         validate_args: bool = False,
     ):
 
@@ -397,6 +433,7 @@ def __init__(
             logits=logits,
             mu=mu,
             theta=theta,
+            scale=scale,
             validate_args=validate_args,
         )
         self.zi_logits, self.mu, self.theta = broadcast_all(
@@ -522,7 +559,7 @@ def sample(
             # Clamping as distributions objects can have buggy behaviors when
             # their parameters are too high
             l_train = torch.clamp(p_means, max=1e8)
-            counts = Poisson(
+            counts = PoissonTorch(
                 l_train
             ).sample()  # Shape : (n_samples, n_cells_batch, n_features)
             return counts

scvi/external/gimvi/_module.py

@@ -134,6 +134,7 @@ def __init__(
             n_layers_individual=n_layers_encoder_individual,
             n_layers_shared=n_layers_encoder_shared,
             dropout_rate=dropout_rate_encoder,
+            return_dist=True,
         )
 
         self.l_encoders = ModuleList(
@@ -143,6 +144,7 @@ def __init__(
                     1,
                     n_layers=1,
                     dropout_rate=dropout_rate_encoder,
+                    return_dist=True,
                 )
                 if self.model_library_bools[i]
                 else None
@@ -197,7 +199,7 @@ def sample_from_posterior_z(
             else:
                 raise Exception("Must provide a mode when having multiple datasets")
         outputs = self.inference(x, mode)
-        qz_m = outputs["qz_m"]
+        qz_m = outputs["qz"].loc
         z = outputs["z"]
         if deterministic:
             z = qz_m
@@ -268,9 +270,9 @@ def sample_scale(
             decode_mode = mode
         inference_out = self.inference(x, mode)
         if deterministic:
-            z = inference_out["qz_m"]
-            if inference_out["ql_m"] is not None:
-                library = inference_out["ql_m"]
+            z = inference_out["qz"].loc
+            if inference_out["ql"] is not None:
+                library = inference_out["ql"].loc
             else:
                 library = inference_out["library"]
         else:
@@ -372,14 +374,14 @@ def inference(self, x: torch.Tensor, mode: Optional[int] = None) -> dict:
         if self.log_variational:
             x_ = torch.log(1 + x_)
 
-        qz_m, qz_v, z = self.z_encoder(x_, mode)
-        ql_m, ql_v, library = None, None, None
+        qz, z = self.z_encoder(x_, mode)
+        ql, library = None, None
         if self.model_library_bools[mode]:
-            ql_m, ql_v, library = self.l_encoders[mode](x_)
+            ql, library = self.l_encoders[mode](x_)
         else:
             library = torch.log(torch.sum(x, dim=1)).view(-1, 1)
 
-        return dict(qz_m=qz_m, qz_v=qz_v, z=z, ql_m=ql_m, ql_v=ql_v, library=library)
+        return dict(qz=qz, z=z, ql=ql, library=library)
 
     @auto_move_data
     def generative(
@@ -447,10 +449,8 @@ def loss(
         x = tensors[REGISTRY_KEYS.X_KEY]
         batch_index = tensors[REGISTRY_KEYS.BATCH_KEY]
 
-        qz_m = inference_outputs["qz_m"]
-        qz_v = inference_outputs["qz_v"]
-        ql_m = inference_outputs["ql_m"]
-        ql_v = inference_outputs["ql_v"]
+        qz = inference_outputs["qz"]
+        ql = inference_outputs["ql"]
         px_rate = generative_outputs["px_rate"]
         px_r = generative_outputs["px_r"]
         px_dropout = generative_outputs["px_dropout"]
@@ -466,11 +466,9 @@ def loss(
         )
 
         # KL Divergence
-        mean = torch.zeros_like(qz_m)
-        scale = torch.ones_like(qz_v)
-        kl_divergence_z = kl(Normal(qz_m, torch.sqrt(qz_v)), Normal(mean, scale)).sum(
-            dim=1
-        )
+        mean = torch.zeros_like(qz.loc)
+        scale = torch.ones_like(qz.scale)
+        kl_divergence_z = kl(qz, Normal(mean, scale)).sum(dim=1)
 
         if self.model_library_bools[mode]:
             library_log_means = getattr(self, f"library_log_means_{mode}")
@@ -483,7 +481,7 @@ def loss(
                 one_hot(batch_index, self.n_batch), library_log_vars
             )
             kl_divergence_l = kl(
-                Normal(ql_m, torch.sqrt(ql_v)),
+                ql,
                 Normal(local_library_log_means, local_library_log_vars.sqrt()),
             ).sum(dim=1)
         else:

scvi/model/_autozi.py

@@ -5,7 +5,7 @@
 import torch
 from anndata import AnnData
 from torch import logsumexp
-from torch.distributions import Beta, Normal
+from torch.distributions import Beta
 
 from scvi import REGISTRY_KEYS
 from scvi._compat import Literal
@@ -212,13 +212,12 @@ def get_marginal_ll(
                 # Distribution parameters and sampled variables
                 inf_outputs, gen_outputs, _ = self.module.forward(tensors)
 
-                px_r = gen_outputs["px_r"]
-                px_rate = gen_outputs["px_rate"]
-                px_dropout = gen_outputs["px_dropout"]
-                qz_m = inf_outputs["qz_m"]
-                qz_v = inf_outputs["qz_v"]
+                px = gen_outputs["px"]
+                px_r = px.theta
+                px_rate = px.mu
+                px_dropout = px.zi_logits
+                qz = inf_outputs["qz"]
                 z = inf_outputs["z"]
-                library = inf_outputs["library"]
 
                 # Reconstruction Loss
                 bernoulli_params_batch = self.module.reshape_bernoulli(
@@ -235,36 +234,22 @@ def get_marginal_ll(
                 )
 
                 # Log-probabilities
-                log_prob_sum = torch.zeros(qz_m.shape[0]).to(self.device)
-                p_z = (
-                    Normal(torch.zeros_like(qz_m), torch.ones_like(qz_v))
-                    .log_prob(z)
-                    .sum(dim=-1)
-                )
+                p_z = gen_outputs["pz"].log_prob(z).sum(dim=-1)
                 p_x_zld = -reconst_loss
-                log_prob_sum += p_z + p_x_zld
-
-                q_z_x = Normal(qz_m, qz_v.sqrt()).log_prob(z).sum(dim=-1)
-                log_prob_sum -= q_z_x
+                q_z_x = qz.log_prob(z).sum(dim=-1)
+                log_prob_sum = p_z + p_x_zld - q_z_x
 
                 if not self.use_observed_lib_size:
+                    ql = inf_outputs["ql"]
+                    library = inf_outputs["library"]
                     (
                         local_library_log_means,
                         local_library_log_vars,
                     ) = self.module._compute_local_library_params(batch_index)
 
-                    p_l = (
-                        Normal(
-                            local_library_log_means.to(self.device),
-                            local_library_log_vars.to(self.device).sqrt(),
-                        )
-                        .log_prob(library)
-                        .sum(dim=-1)
-                    )
-
-                    ql_m = inf_outputs["ql_m"]
-                    ql_v = inf_outputs["ql_v"]
-                    q_l_x = Normal(ql_m, ql_v.sqrt()).log_prob(library).sum(dim=-1)
+                    p_l = gen_outputs["pl"].log_prob(library).sum(dim=-1)
+
+                    q_l_x = ql.log_prob(library).sum(dim=-1)
 
                     log_prob_sum += p_l - q_l_x
 

scvi/model/_condscvi.py

@@ -139,7 +139,7 @@ def get_vamp_prior(
             x = tensors[REGISTRY_KEYS.X_KEY]
             y = tensors[REGISTRY_KEYS.LABELS_KEY]
             out = self.module.inference(x, y)
-            mean_, var_ = out["qz_m"], out["qz_v"]
+            mean_, var_ = out["qz"].loc, (out["qz"].scale ** 2)
             mean += [mean_.cpu()]
             var += [var_.cpu()]
 

scvi/model/_totalvi.py

@@ -330,14 +330,14 @@ def get_latent_library_size(
         post = self._make_data_loader(
             adata=adata, indices=indices, batch_size=batch_size
         )
+
         libraries = []
         for tensors in post:
             inference_inputs = self.module._get_inference_input(tensors)
             outputs = self.module.inference(**inference_inputs)
             if give_mean:
-                ql_m = outputs["ql_m"]
-                ql_v = outputs["ql_v"]
-                library = torch.exp(ql_m + 0.5 * ql_v)
+                ql = outputs["ql"]
+                library = torch.exp(ql.loc + 0.5 * (ql.scale**2))
             else:
                 library = outputs["library_gene"]
             libraries += [library.cpu()]