Repository: sbi-dev/sbi

Issue:
Simulation-based calibration with Variational Inference
Hello,

I am encountering an error when trying to apply simulation-based calibration with a Variational Inference posterior: `VIPosterior` (I tried `MCMCPosterior` and it runs fine).
The error occurs when I call the function `run_sbc`:
```
AttributeError: The variational posterior was not fit on the specified `default_x` tensor([[...]]). Please train using `posterior.train()`.
```

To help diagnose the issue, I have prepared a minimal code example that reproduces the error:
```python
import torch
from sbi import analysis as analysis
from sbi import utils as utils
from sbi.analysis import run_sbc
from sbi.inference import (SNLE, VIPosterior,
                           likelihood_estimator_based_potential,
                           prepare_for_sbi, simulate_for_sbi)

num_dim = 3
prior = utils.BoxUniform(low=-2 * torch.ones(num_dim), high=2 * torch.ones(num_dim))

len_x_o = 10
x_o = torch.zeros(len_x_o, num_dim,)+ 0.01*torch.randn(len_x_o, num_dim,)

def linear_gaussian(theta):
    return theta + 1.0 + torch.randn_like(theta) * 0.1

simulator, prior = prepare_for_sbi(linear_gaussian, prior)
inference = SNLE(prior=prior, density_estimator="nsf")
theta, x = simulate_for_sbi(simulator, prior, num_simulations=200)

likelihood_estimator = inference.append_simulations(theta, x).train(show_train_summary=True)
potential_fn, parameter_transform = likelihood_estimator_based_potential(likelihood_estimator, prior, x_o)
posterior = VIPosterior(potential_fn=potential_fn, prior=prior, theta_transform=parameter_transform)
posterior.train()
posterior = posterior.set_default_x(x_o)
```

```python
thetas = prior.sample((1_000,))
xs = simulator(thetas)

ranks, dap_samples = run_sbc(
    thetas, xs, posterior, num_posterior_samples=1_000
)
```

I suspect that the reason for this error is due to the posterior not being trained after the introduction of each new vector in the simulated dataset `xs`.


Discussion:
Hello @hdahmou, thanks for creating this issue. Indeed, this seems to be a bug: the `run_sbc` method should call `posterior.train()` for each new "`xi in xs`". 
I will make a PR with a fix soon! 

Resolving patch:
diff --git a/sbi/analysis/sbc.py b/sbi/analysis/sbc.py
--- a/sbi/analysis/sbc.py
+++ b/sbi/analysis/sbc.py
@@ -10,6 +10,7 @@
 
 from sbi.inference import DirectPosterior
 from sbi.inference.posteriors.base_posterior import NeuralPosterior
+from sbi.inference.posteriors.vi_posterior import VIPosterior
 from sbi.simulators.simutils import tqdm_joblib
 from sbi.utils.metrics import c2st
 
@@ -169,6 +170,11 @@ def sbc_on_batch(
     ranks = torch.zeros((thetas.shape[0], len(reduce_fns)))
 
     for idx, (tho, xo) in enumerate(zip(thetas, xs)):
+        # VI posterior needs to be trained on the current xo.
+        if isinstance(posterior, VIPosterior):
+            posterior.set_default_x(xo)
+            posterior.train()
+
         # Draw posterior samples and save one for the data average posterior.
         ths = posterior.sample((num_posterior_samples,), x=xo, show_progress_bars=False)
 
diff --git a/sbi/inference/posteriors/direct_posterior.py b/sbi/inference/posteriors/direct_posterior.py
--- a/sbi/inference/posteriors/direct_posterior.py
+++ b/sbi/inference/posteriors/direct_posterior.py
@@ -12,7 +12,7 @@
     posterior_estimator_based_potential,
 )
 from sbi.samplers.rejection.rejection import accept_reject_sample
-from sbi.types import Shape, TorchTransform
+from sbi.types import Shape
 from sbi.utils import check_prior, match_theta_and_x_batch_shapes, within_support
 from sbi.utils.torchutils import ensure_theta_batched
 
@@ -51,8 +51,8 @@ def __init__(
             x_shape: Shape of a single simulator output. If passed, it is used to check
                 the shape of the observed data and give a descriptive error.
             enable_transform: Whether to transform parameters to unconstrained space
-                during MAP optimization. When False, an identity transform will be returned
-                for `theta_transform`.
+                during MAP optimization. When False, an identity transform will be
+                returned for `theta_transform`.
         """
         # Because `DirectPosterior` does not take the `potential_fn` as input, it
         # builds it itself. The `potential_fn` and `theta_transform` are used only for
diff --git a/sbi/inference/posteriors/vi_posterior.py b/sbi/inference/posteriors/vi_posterior.py
--- a/sbi/inference/posteriors/vi_posterior.py
+++ b/sbi/inference/posteriors/vi_posterior.py
@@ -243,7 +243,8 @@ def set_q(
         move_all_tensor_to_device(q, self._device)
         assert isinstance(
             q, Distribution
-        ), "Something went wrong when initializing the variational distribution. Please create an issue on github https://github.com/mackelab/sbi/issues"
+        ), """Something went wrong when initializing the variational distribution.
+            Please create an issue on github https://github.com/mackelab/sbi/issues"""
         check_variational_distribution(q, self._prior)
         self._q = q
 


Test patch:
diff --git a/tests/sbc_test.py b/tests/sbc_test.py
--- a/tests/sbc_test.py
+++ b/tests/sbc_test.py
@@ -17,8 +17,10 @@
 
 @pytest.mark.parametrize("reduce_fn_str", ("marginals", "posterior_log_prob"))
 @pytest.mark.parametrize("prior", ("boxuniform", "independent"))
-@pytest.mark.parametrize("method", (SNPE, SNLE))
-def test_running_sbc(method, prior, reduce_fn_str, model="mdn"):
+@pytest.mark.parametrize(
+    "method, sampler", ((SNPE, None), (SNLE, "mcmc"), (SNLE, "vi"))
+)
+def test_running_sbc(method, prior, reduce_fn_str, sampler, model="mdn"):
     """Tests running inference and then SBC and obtaining nltp."""
 
     num_dim = 2
@@ -46,7 +48,12 @@ def simulator(theta):
     _ = inferer.append_simulations(theta, x).train(
         training_batch_size=100, max_num_epochs=max_num_epochs
     )
-    posterior = inferer.build_posterior()
+    if method == SNLE:
+        posterior_kwargs = {"sample_with": "mcmc" if sampler == "mcmc" else "vi"}
+    else:
+        posterior_kwargs = {}
+
+    posterior = inferer.build_posterior(**posterior_kwargs)
 
     thetas = prior.sample((num_sbc_runs,))
     xs = simulator(thetas)