Repository: sbi-dev/sbi

Issue:
Tests for Permutation Embedding



Discussion:

Resolving patch:
diff --git a/sbi/inference/snpe/snpe_base.py b/sbi/inference/snpe/snpe_base.py
--- a/sbi/inference/snpe/snpe_base.py
+++ b/sbi/inference/snpe/snpe_base.py
@@ -113,7 +113,9 @@ def append_simulations(
             NeuralInference object (returned so that this function is chainable).
         """
 
-        is_valid_x, num_nans, num_infs = handle_invalid_x(x, True)  # Hardcode to True
+        is_valid_x, num_nans, num_infs = handle_invalid_x(
+            x, exclude_invalid_x=True
+        )  # Hardcode to True
 
         x = x[is_valid_x]
         theta = theta[is_valid_x]
diff --git a/sbi/neural_nets/embedding_nets.py b/sbi/neural_nets/embedding_nets.py
--- a/sbi/neural_nets/embedding_nets.py
+++ b/sbi/neural_nets/embedding_nets.py
@@ -1,7 +1,8 @@
-from typing import List, Tuple, Union
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Affero General Public License v3, see <https://www.gnu.org/licenses/>.
 
 import torch
-from torch import Size, Tensor, nn
+from torch import Tensor, nn
 
 
 class FCEmbedding(nn.Module):
@@ -10,7 +11,7 @@ def __init__(
         input_dim: int,
         output_dim: int = 20,
         num_layers: int = 2,
-        num_hiddens: int = 20,
+        num_hiddens: int = 40,
     ):
         """Fully-connected multi-layer neural network to be used as embedding network.
 
@@ -47,7 +48,7 @@ class CNNEmbedding(nn.Module):
     def __init__(
         self,
         input_dim: int,
-        output_dim: int,
+        output_dim: int = 20,
         num_fully_connected: int = 2,
         num_hiddens: int = 120,
         out_channels_cnn_1: int = 10,
@@ -61,6 +62,7 @@ def __init__(
 
         Args:
             input_dim: Dimensionality of input.
+            output_dim: Dimensionality of the output.
             num_conv: Number of convolutional layers.
             num_fully_connected: Number fully connected layer, minimum of 2.
             num_hiddens: Number of hidden dimensions in fully-connected layers.
@@ -128,36 +130,38 @@ class PermutationInvariantEmbedding(nn.Module):
 
     def __init__(
         self,
-        single_trial_net: nn.Module,
-        input_dim: int,
-        output_dim: int,
-        num_fully_connected: int = 2,
-        num_hiddens: int = 20,
+        trial_net: nn.Module,
+        trial_net_output_dim: int,
         combining_operation: str = "mean",
+        num_layers: int = 2,
+        num_hiddens: int = 40,
+        output_dim: int = 20,
     ):
         """Permutation invariant multi-layer NN.
-            The single_trial_net is applied to each "trial" of the input
-            and is combined by the combining_operation (mean or sum).
+
+        The trial_net is applied to each "trial" of the input
+        and is combined by the combining_operation (mean or sum) to construct a
+        permutation invariant embedding across iid trials.
+        This embedding is embedded again using an additional fully connected net.
 
         Args:
-            single_trial_net: Network to process one trial, the combining_operation is
-                applied to its output. Taskes as input (batch, input_dim).
+            trial_net: Network to process one trial. The combining_operation is
+                applied to its output. Takes as input (batch, input_dim), where
+                input_dim is the dimensionality of a single trial. Produces output
+                (batch, latent_dim).
                 Remark: This network should be large enough as it acts on all (iid)
                 inputs seperatley and needs enough capacity to process the information
                 of all inputs.
-            input_dim: Dimensionality of input to the fully connected layers
-                (output_dimension of single_trial_net).
-            output_dim: Dimensionality of the output.
-            num_fully_connected: Number of fully connected layer, minimum of 2.
-            num_hiddens: Number of hidden dimensions in fully-connected layers.
+            trial_net_output_dim: Dimensionality of the output of the trial_net / input
+                to the fully connected layers.
             combining_operation: How to combine the permutational dimensions, one of
                 'mean' or 'sum'.
+            num_layers: Number of fully connected layer, minimum of 2.
+            num_hiddens: Number of hidden dimensions in fully-connected layers.
+            output_dim: Dimensionality of the output.
         """
         super().__init__()
-        self.single_trial_net = single_trial_net
-        self.input_dim = input_dim
-        self.output_dim = output_dim
-        self.num_hiddens = num_hiddens
+        self.trial_net = trial_net
         self.combining_operation = combining_operation
 
         if combining_operation not in ["sum", "mean"]:
@@ -165,9 +169,9 @@ def __init__(
 
         # construct fully connected layers
         self.fc_subnet = FCEmbedding(
-            input_dim=input_dim,
+            input_dim=trial_net_output_dim,
             output_dim=output_dim,
-            num_layers=num_fully_connected,
+            num_layers=num_layers,
             num_hiddens=num_hiddens,
         )
 
@@ -180,9 +184,9 @@ def forward(self, x: Tensor) -> Tensor:
         """
         batch, permutation_dim, _ = x.shape
 
-        iid_embeddings = self.single_trial_net(
-            x.view(batch * permutation_dim, -1)
-        ).view(batch, permutation_dim, -1)
+        iid_embeddings = self.trial_net(x.view(batch * permutation_dim, -1)).view(
+            batch, permutation_dim, -1
+        )
 
         if self.combining_operation == "mean":
             e = iid_embeddings.mean(dim=1)


Test patch:
diff --git a/tests/embedding_net_test.py b/tests/embedding_net_test.py
--- a/tests/embedding_net_test.py
+++ b/tests/embedding_net_test.py
@@ -1,18 +1,18 @@
-# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
-# under the Affero General Public License v3, see <https://www.gnu.org/licenses/>.
-
 from __future__ import annotations
 
 import pytest
 import torch
 from torch import eye, ones, zeros
 
-from sbi import analysis as analysis
 from sbi import utils as utils
 from sbi.inference import SNLE, SNPE, SNRE
-from sbi.neural_nets.embedding_nets import FCEmbedding
-from sbi.simulators.linear_gaussian import linear_gaussian
+from sbi.neural_nets.embedding_nets import FCEmbedding, PermutationInvariantEmbedding
+from sbi.simulators.linear_gaussian import (
+    linear_gaussian,
+    true_posterior_linear_gaussian_mvn_prior,
+)
 from sbi.utils import classifier_nn, likelihood_nn, posterior_nn
+from tests.test_utils import check_c2st
 
 
 @pytest.mark.parametrize("method", ["SNPE", "SNLE", "SNRE"])
@@ -35,7 +35,7 @@ def test_embedding_net_api(method, num_dim: int, embedding_net: str):
     if embedding_net == "mlp":
         embedding = FCEmbedding(input_dim=num_dim)
     else:
-        raise NameError
+        raise NameError(f"{embedding_net} not supported.")
 
     if method == "SNPE":
         density_estimator = posterior_nn("maf", embedding_net=embedding)
@@ -58,3 +58,123 @@ def test_embedding_net_api(method, num_dim: int, embedding_net: str):
 
     s = posterior.sample((1,))
     _ = posterior.potential(s)
+
+
+@pytest.mark.parametrize("num_trials", [1, 2])
+@pytest.mark.parametrize("num_dim", [1, 2])
+def test_iid_embedding_api(num_trials, num_dim):
+
+    prior = utils.BoxUniform(-2.0 * ones(num_dim), 2.0 * ones(num_dim))
+
+    num_thetas = 1000
+    theta = prior.sample((num_thetas,))
+
+    # simulate iid x.
+    iid_theta = theta.reshape(num_thetas, 1, num_dim).repeat(1, num_trials, 1)
+    x = torch.randn_like(iid_theta) + iid_theta
+    x_o = zeros(1, num_trials, num_dim)
+
+    output_dim = 5
+    single_trial_net = FCEmbedding(input_dim=num_dim, output_dim=output_dim)
+    embedding_net = PermutationInvariantEmbedding(
+        single_trial_net,
+        trial_net_output_dim=output_dim,
+    )
+
+    density_estimator = posterior_nn("maf", embedding_net=embedding_net)
+    inference = SNPE(prior, density_estimator=density_estimator)
+
+    _ = inference.append_simulations(theta, x).train(max_num_epochs=5)
+    posterior = inference.build_posterior().set_default_x(x_o)
+
+    s = posterior.sample((1,))
+    _ = posterior.potential(s)
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize("num_trials", [10, 50])
+@pytest.mark.parametrize("num_dim", [2])
+@pytest.mark.parametrize("method", ["SNPE", "SNLE", "SNRE"])
+def test_iid_inference(num_trials, num_dim, method):
+    """Test accuracy in Gaussian linear simulator with iid trials.
+
+    Tests permutation invariance of NPE iid embeddings, and MCMC iid-trials inference
+    with NLE and NRE trained on single trials.
+    """
+
+    prior = torch.distributions.MultivariateNormal(
+        torch.zeros(num_dim), torch.eye(num_dim)
+    )
+
+    # Scale number of training samples with num_trials.
+    num_thetas = 1000 + 100 * num_trials
+    # Likelihood-based methods train on single trials.
+    num_simulations = num_thetas
+
+    if method == "SNPE":  # SNPE needs embedding and iid trials during training.
+        theta = prior.sample((num_thetas,))
+        # simulate iid x.
+        iid_theta = theta.reshape(num_thetas, 1, num_dim).repeat(1, num_trials, 1)
+        x = torch.randn_like(iid_theta) + iid_theta
+        x_o = zeros(1, num_trials, num_dim)
+
+        # embedding
+        latent_dim = 10
+        single_trial_net = FCEmbedding(
+            input_dim=num_dim,
+            num_hiddens=40,
+            num_layers=2,
+            output_dim=latent_dim,
+        )
+        embedding_net = PermutationInvariantEmbedding(
+            single_trial_net,
+            trial_net_output_dim=latent_dim,
+            # NOTE: post-embedding is not needed really.
+            num_layers=1,
+            num_hiddens=10,
+            output_dim=10,
+        )
+
+        density_estimator = posterior_nn("maf", embedding_net=embedding_net)
+
+        inference = SNPE(prior, density_estimator=density_estimator)
+    else:  # likelihood-based methods: single-trial training without embeddings.
+
+        theta = prior.sample((num_simulations,))
+        x = torch.randn_like(theta) + theta
+        x_o = zeros(1, num_dim)
+
+        if method == "SNLE":
+            inference = SNLE(prior, density_estimator=likelihood_nn("maf"))
+        elif method == "SNRE":
+            inference = SNRE(prior, classifier=classifier_nn("resnet"))
+        else:
+            raise NameError
+
+    # get reference samples from true posterior
+    num_samples = 1000
+    reference_samples = true_posterior_linear_gaussian_mvn_prior(
+        x_o.squeeze(),
+        likelihood_shift=torch.zeros(num_dim),
+        likelihood_cov=torch.eye(num_dim),
+        prior_cov=prior.covariance_matrix,
+        prior_mean=prior.loc,
+    ).sample((num_samples,))
+
+    # training
+    _ = inference.append_simulations(theta, x).train()
+
+    # inference
+    posterior = inference.build_posterior().set_default_x(x_o)
+    samples = posterior.sample((num_samples,))
+
+    if method == "SNPE":
+        check_c2st(samples, reference_samples, alg=method)
+        # permute and test again
+        num_repeats = 2
+        for _ in range(num_repeats):
+            trial_permutet_x_o = x_o[:, torch.randperm(x_o.shape[1]), :]
+            samples = posterior.sample((num_samples,), x=trial_permutet_x_o)
+            check_c2st(samples, reference_samples, alg=method + " permuted")
+    else:
+        check_c2st(samples, reference_samples, alg=method)