Repository: sbi-dev/sbi

Issue:
Issues with `PermutationInvariantEmbedding`: batch dimension for `x_o` is required
hi there

just curious why we cannot have SNPE for all iid sample

like SNLE or even traditional MLE？

thank you


Discussion:
**The basic problem**

During training, you only have 1 `x` for every `theta` (because you simulate only once for each `theta`). So, the neural network learns to approximate `p(theta | x)` (with `x` being a single x here). 

Now, after training, you want to have `p(theta | x_1, ..., x_N)`, but the neural net was never trained for many `x`, so SNPE does not know what to do.

**How does SNLE do it**

SNLE has the same "issue", it is trained on pairs of theta and a single corresponding x. The neural net then approximates `p(x | theta)`. However, one can neatly multiply these terms to obtain `p(x1,...xN | theta) = p(x1 | theta)*...*p(xN | theta)` to obtain the joint likelihood and plug it into MCMC.

**Workaround for SNPE**

The only way of having iid trials with SNPE is to simulate `N` times for each `theta` (`N` being the number of trials). You can then concatenate the `N` simulation outputs for each theta and, after training, also concatenate the iid dataset into a 1D vector. Alternatively, instead of concatenating the simulation outputs, you can also use the [permuatation-invariant embedding](https://github.com/mackelab/sbi/blob/fda9d4eced103dc640f039f72b1d9791fd5c9bf3/sbi/neural_nets/embedding_nets.py#L124), but this is not released yet so you will have to rely on the bleeding edge github version of `sbi`. Either way, this has the disadvantage that one has to simulate `N` times for each `theta`.

Hope this helps
Michael
Dear Michael,

Thank you a lot for your detail explanation. As you suggested, I tested N iid for each theta in simulator, then concatenate to 1D vector. I tested the example in tutorial and it works!

However, I am not sure if this may lose some information in more complicated situation: e.g. 
if for each trial the return x is a vector with length m
then after concatenation it would be a 1-D vector with length N*m. How the neural network could know that it is N trials with output length m, but not N/2 trials with output length 2*m (or other combination)?

I am just wondering if this would be helpful to arrange that to a matrix with N*m so that the neural network could learn more easily? Or actually it doesn't matter?

Thank you.

Best wishes
Tim


Dear Tim, good to hear that this works.

In general, the neural network does not know if it is `N` trials of length `m` or `N/2` with length `2m`. However, as for the amount of "information" the two scenarios are identical and thus the neural network should eventually converge to the same thing.

However, as you already mention, there can be differences in how easy it is to learn for the neural network. Therefore, I would recommend using the `PermutationInvariatEmbedding` as follows:

1) Install github version of sbi (or wait until we publish sbi v0.20.0 in a few days)
```
git clone https://github.com/mackelab/sbi.git
cd sbi
pip install -e .
```

2) Use the permutation invariant embedding:
```python
from sbi.neural_nets.embedding_nets import PermutationInvariantEmbedding, FCEmbedding
from sbi.utils import posterior_nn

single_trial_hidden_dim = 100
single_trial_net = FCEmbedding(input_dim=dimensionality_of_x, output_dim=single_trial_hidden_dim)
embedding_net = PermutationInvariantEmbedding(single_trial_net, trial_net_output_dim=single_trial_hidden_dim)
net = posterior_nn("maf", embedding_net=embedding_net)
inference = SNPE(prior, density_estimator=net)
```

The `PermutationInvariantEmbedding` explicitly handles each trial separately in the `single_trial_net` and then combines the trials.
Thank you, and I will try this!

Just wonder if there is any plan or timeline that this would be released? Thank you!
We are trying to make it happen this week (Friday), but no guarantees.
`sbi v0.20.0` is out.
```
pip install sbi --upgrade
```
Thank you, I am trying it.

Am I correct that the simulated data for training, x, also need to be in 3d?

e.g. if I did 20000 simulations, with 100 trials and 10 columns in each trial, then x should be in the form

(20000,100,10)?

I am following the example using simulate_for_sbi and it automatically convert that to 20000, 100*10, and got an error.

Also just wonder if I will miss anything if I just use for loop to create the corresponding theta and x?
is there any particular importance for generating prepare_for_sbi and simualte_for_sbi object, except for convenience to use?

P.S. I tested and after reshape x, it is working, thanks!
Hi! You are right, the shape should be `(20000,100,10)`, see also [here](https://github.com/mackelab/sbi/blob/4d92736f5a05ee7a3f33e01655dcd2b03a0da956/sbi/neural_nets/embedding_nets.py#L205).

And indeed, `simulate_for_sbi` will break things -- please do **not** use `simulate_for_sbi` for now, we will fix this soon. `prepare_for_sbi` and `simulate_for_sbi` are only for convenience, they do not have to be used.
Thanks!

following the above code, I have an inference object. Then I use:

inference = inference.append_simulations(theta, x) 
density_estimator = inference.train()  
 posterior = inference.build_posterior(density_estimator)
x_oo = x[0,:]
posterior_samples = posterior.sample((10000,), x=x_oo)

i.e. when I put the first simulated data to get the posterior, I still got an error that the input should be 1d.

Am I doing anything wrong? Thank you.

Hm, maybe try `x_oo = x[:1, :]` such that the batch dimension gets preserved?
yes, this is working. thanks!
For `SNPE` we are enabling iid inference by treating the trial dimension as an additional data dimension. Thus, the simulator needs to take care of this additional dimension as well, e.g., it needs to return `(num_trials, dim_single_trial_data)` for a single simulation. 

Maybe I am missing something, but if the simulator returns `(num_trials, dim_single_trial_data)` (or `(1, num_trials, dim_single_trial_data)` after calling `process_simulator`), then `simulate_for_sbi` should work just fine. 

When passing `x_o` for sampling the posterior, one needs to pass it with a leading batch dimension, e.g., `(1, num_trials, dim_single_trial_data)`. Otherwise it is unclear whether it's a multiple trials, or a multi-dimensional `x`. 



Resolving patch:
diff --git a/sbi/utils/user_input_checks.py b/sbi/utils/user_input_checks.py
--- a/sbi/utils/user_input_checks.py
+++ b/sbi/utils/user_input_checks.py
@@ -272,9 +272,12 @@ def check_for_possibly_batched_x_shape(x_shape):
                 has an inferred batch shape larger than one. This is not supported in
                 some sbi methods for reasons depending on the scenario:
 
-                    - in case you want to evaluate or sample conditioned on several xs
-                    e.g., (p(theta | [x1, x2, x3])), this is not supported yet except
-                    when using likelihood based SNLE and SNRE.
+                    - in case you want to evaluate or sample conditioned on several iid
+                      xs e.g., (p(theta | [x1, x2, x3])), this is fully supported only
+                      for likelihood based SNLE and SNRE. For SNPE it is supported only 
+                      for a fixed number of trials and using an appropriate embedding 
+                      net, i.e., by treating the trials as additional data dimension. In
+                      that case, make sure to pass xo with a leading batch dimensionen.
 
                     - in case you trained with a single round to do amortized inference
                     and now you want to evaluate or sample a given theta conditioned on
@@ -569,6 +572,10 @@ def process_x(
 
     x = atleast_2d(torch.as_tensor(x, dtype=float32))
 
+    # If x_shape is provided, we can fix a missing batch dim for >1D data.
+    if x_shape is not None and len(x_shape) > len(x.shape):
+        x = x.unsqueeze(0)
+
     input_x_shape = x.shape
     if not allow_iid_x:
         check_for_possibly_batched_x_shape(input_x_shape)


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
@@ -5,7 +5,7 @@
 from torch import eye, ones, zeros
 
 from sbi import utils as utils
-from sbi.inference import SNLE, SNPE, SNRE
+from sbi.inference import SNLE, SNPE, SNRE, simulate_for_sbi
 from sbi.neural_nets.embedding_nets import (
     CNNEmbedding,
     FCEmbedding,
@@ -96,9 +96,9 @@ def test_iid_embedding_api(num_trials, num_dim):
 
 
 @pytest.mark.slow
-@pytest.mark.parametrize("num_trials", [10, 50])
+@pytest.mark.parametrize("num_trials", [1, 10, 50])
 @pytest.mark.parametrize("num_dim", [2])
-@pytest.mark.parametrize("method", ["SNPE", "SNLE", "SNRE"])
+@pytest.mark.parametrize("method", ("SNPE",))
 def test_iid_inference(num_trials, num_dim, method):
     """Test accuracy in Gaussian linear simulator with iid trials.
 
@@ -112,51 +112,39 @@ def test_iid_inference(num_trials, num_dim, method):
 
     # 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.
+    # simulate iid x.
+    def simulator(theta, num_trials=num_trials):
+        iid_theta = theta.reshape(theta.shape[0], 1, num_dim).repeat(1, num_trials, 1)
+        return torch.randn_like(iid_theta) + iid_theta
+
+    theta, x = simulate_for_sbi(simulator, prior, num_simulations=num_thetas)
+
+    # 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,
+    )
 
-        theta = prior.sample((num_simulations,))
-        x = torch.randn_like(theta) + theta
-        x_o = zeros(1, num_dim)
+    density_estimator = posterior_nn("maf", embedding_net=embedding_net)
 
-        if method == "SNLE":
-            inference = SNLE(prior, density_estimator=likelihood_nn("maf"))
-        elif method == "SNRE":
-            inference = SNRE(prior, classifier=classifier_nn("resnet"))
-        else:
-            raise NameError
+    inference = SNPE(prior, density_estimator=density_estimator)
 
     # get reference samples from true posterior
     num_samples = 1000
+    # define x_o without batch dim to test handling below.
+    x_o = zeros(num_trials, num_dim)
     reference_samples = true_posterior_linear_gaussian_mvn_prior(
         x_o.squeeze(),
         likelihood_shift=torch.zeros(num_dim),
@@ -172,26 +160,16 @@ def test_iid_inference(num_trials, num_dim, method):
     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)
+    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[0]), :]
+        samples = posterior.sample((num_samples,), x=trial_permutet_x_o)
+        check_c2st(samples, reference_samples, alg=method + " permuted")
 
 
-@pytest.mark.parametrize(
-    "input_shape",
-    [
-        (32,),
-        (32, 32),
-        (32, 64),
-    ],
-)
+@pytest.mark.parametrize("input_shape", [(32,), (32, 32), (32, 64)])
 @pytest.mark.parametrize("num_channels", (1, 2, 3))
 def test_1d_and_2d_cnn_embedding_net(input_shape, num_channels):
     import torch
diff --git a/tests/user_input_checks_test.py b/tests/user_input_checks_test.py
--- a/tests/user_input_checks_test.py
+++ b/tests/user_input_checks_test.py
@@ -192,14 +192,19 @@ def test_process_prior(prior):
 
 
 @pytest.mark.parametrize(
-    "x, x_shape",
+    "x, x_shape, allow_iid",
     (
-        (ones(3), torch.Size([1, 3])),
-        (ones(1, 3), torch.Size([1, 3])),
+        (ones(3), torch.Size([1, 3]), False),
+        (ones(1, 3), torch.Size([1, 3]), False),
+        (ones(10, 3), torch.Size([1, 10, 3]), False),  # 2D data / iid SNPE
+        pytest.param(
+            ones(10, 3), None, False, marks=pytest.mark.xfail
+        ),  # 2D data / iid SNPE without x_shape
+        (ones(10, 10), torch.Size([1, 10]), True),  # iid likelihood based
     ),
 )
-def test_process_x(x, x_shape):
-    process_x(x, x_shape)
+def test_process_x(x, x_shape, allow_iid):
+    process_x(x, x_shape, allow_iid_x=allow_iid)
 
 
 @pytest.mark.parametrize(