Improve docs of MNPE and EnsemblePosterior

Author: michaeldeistler

docs/sbi.rst

@@ -57,6 +57,7 @@ Training
 
    sbi.inference.NPE_C
    sbi.inference.NPE_A
+   sbi.inference.MNPE
    sbi.inference.FMPE
    sbi.inference.NPSE
    sbi.inference.NLE_A

sbi/inference/posteriors/ensemble_posterior.py

@@ -19,8 +19,9 @@
 class EnsemblePosterior(NeuralPosterior):
     r"""Wrapper for bundling together different posterior instances into an ensemble.
 
-    This class creates a posterior ensemble from a set of N different, already trained
-    posterior estimators :math:`p_{i}(\theta|x_o)`, where :math:`i \in \{1,...,N\}`.
+    This class creates a posterior ensemble from a set of :math:`N` different, already
+    trained posterior estimators :math:`p_{i}(\theta \mid x_o)`, where
+    :math:`i \in \{1, \ldots, N\}`.
 
     It can wrap all posterior classes available in ``sbi`` and even a mixture of
     different posteriors, i.e. obtained via SNLE and SNPE at the same time, since it
@@ -30,32 +31,38 @@ class EnsemblePosterior(NeuralPosterior):
 
     So far, ``log_prob()``, ``sample()`` and ``map()`` functionality are supported.
 
-    Attributes:
-        posteriors: List of the posterior estimators making up the ensemble.
-        num_components: Number of posterior estimators.
-        weights: Weight of each posterior distribution. If none are provided each
-            posterior is weighted with 1/N.
-        priors: Prior distributions of all posterior components.
-        theta_transform: If passed, this transformation will be applied during the
-            optimization performed when obtaining the map. It does not affect the
-            .sample() and .log_prob() methods.
-        device: device to host the posterior distribution.
-
     Example:
     --------
 
     ::
 
         import torch
-        from joblib import Parallel, delayed
-        from sbi.examples.minimal import simple
+        from sbi.inference import NPE, EnsemblePosterior
+
+        theta = prior.sample((100,))
+        x = simulate(theta)
 
         n_ensembles = 10
-        posteriors = Parallel(n_jobs=-1)(delayed(simple)() for i in range(n_ensembles))
+        posteriors = []
+        for _ in range(n_ensembles):
+            inference = NPE()
+            inference.append_simulations(theta, x).train()
+            posteriors.append(inference.build_posterior())
 
         ensemble = EnsemblePosterior(posteriors)
         ensemble.set_default_x(torch.zeros((3,)))
         ensemble.sample((1,))
+
+    Attributes:
+        posteriors: List of the posterior estimators making up the ensemble.
+        num_components: Number of posterior estimators.
+        weights: Weight of each posterior distribution. If none are provided each
+            posterior is weighted with 1/N.
+        priors: Prior distributions of all posterior components.
+        theta_transform: If passed, this transformation will be applied during the
+            optimization performed when obtaining the map. It does not affect the
+            .sample() and .log_prob() methods.
+        device: device to host the posterior distribution.
     """
 
     def __init__(

sbi/inference/trainers/nle/mnle.py

@@ -20,13 +20,13 @@
 
 
 class MNLE(LikelihoodEstimatorTrainer):
-    """Mixed Neural Likelihood Estimation (MNLE) [1].
+    """Method that can infer parameters given discrete and continuous data (Mixed NLE).
 
     Like NLE, but designed to be applied to data with mixed types, e.g., continuous
     data and discrete data like they occur in decision-making experiments
     (reation times and choices).
 
-    [1] Flexible and efficient simulation-based inference for models of
+    Flexible and efficient simulation-based inference for models of
     decision-making, Boelts et al. 2021,
     https://www.biorxiv.org/content/10.1101/2021.12.22.473472v2
     """

sbi/inference/trainers/npe/mnpe.py

@@ -21,12 +21,45 @@
 
 
 class MNPE(NPE_C):
-    r"""Mixed Neural Posterior Estimation (MNPE).
+    r"""Method that can infer discrete and continuous parameters (Mixed NPE).
 
-    Like NPE-C, but designed to be applied to data with mixed types, e.g.,
-    continuous parameters and discrete parameters like they occur in models with
-    switching components. The emebedding net will only operate on the continuous
-    parameters, note this to design the dimension of the embedding net.
+    MNPE (Mixed Neural Posterior Estimation) is similar to NPE: it directly
+    estimates a distribution over parameters given data. Unlike NPE, it is designed to
+    be applied to parmaeters with mixed types, i.e., continuous and discrete parameters.
+    This can occur, for example, in models with switching components. The emebedding
+    net will only operate on the continuous parameters, note this to design the
+    dimension of the embedding net.
+
+    Example:
+    --------
+
+    ::
+
+        import torch
+        from sbi.inference import MNPE
+
+        dim_theta_discrete = 3
+        dim_theta_continuous = 2
+        dim_theta = 5
+        dim_x = 50
+
+        num_sims = 100
+
+        discrete_theta = torch.randint(low=0, high=2, size=(100, dim_theta_discrete))
+        continuous_theta = torch.randn((num_sims, dim_theta_discrete))
+
+        # Important: the theta must have all continuous paramters first, and
+        # discrete parameters after this.
+        theta = torch.cat([continuous_theta, discrete_theta], dim=1)
+        x = torch.randn((num_sims, dim_x))
+
+        inference = MNPE()
+        _ = inference.append_simulations(theta, x).train()
+
+        posterior = inference.build_posterior()
+
+        x_o = torch.randn((1, dim_x))
+        samples = posterior.sample((100,), x=x_o)
     """
 
     def __init__(