refactor: add batched option to utils, refactor sbc.

Author: janfb

sbi/diagnostics/sbc.py

@@ -25,6 +25,7 @@ def run_sbc(
     reduce_fns: Union[str, Callable, List[Callable]] = "marginals",
     num_workers: int = 1,
     show_progress_bar: bool = True,
+    use_batched_sampling: bool = True,
     **kwargs,
 ):
     """Run simulation-based calibration (SBC) (parallelized across sbc runs).
@@ -47,6 +48,8 @@ def run_sbc(
         num_workers: number of CPU cores to use in parallel for running
             `num_sbc_samples` inferences.
         show_progress_bar: whether to display a progress over sbc runs.
+        use_batched_sampling: whether to use batched sampling for posterior
+            samples.
 
     Returns:
         ranks: ranks of the ground truth parameters under the inferred
@@ -81,13 +84,16 @@ def run_sbc(
 
     # Get posterior samples, batched or parallelized.
     posterior_samples = get_posterior_samples_on_batch(
-        xs, posterior, num_posterior_samples, num_workers, show_progress_bar
+        xs,
+        posterior,
+        (num_posterior_samples,),
+        num_workers,
+        show_progress_bar,
+        use_batched_sampling=use_batched_sampling,
     )
-    # for calibration methods its handy to have len(xs) in first dim.
-    posterior_samples = posterior_samples.transpose(0, 1)
 
     # take a random draw from each posterior to get data averaged posterior samples.
-    dap_samples = posterior_samples[:, 0, :]
+    dap_samples = posterior_samples[0, :, :]
     assert dap_samples.shape == (num_sbc_samples, thetas.shape[1]), "Wrong dap shape."
 
     ranks = _run_sbc(
@@ -126,8 +132,8 @@ def _run_sbc(
 
     ranks = torch.zeros((num_sbc_samples, len(reduce_fns)))
     # Iterate over all sbc samples and calculate ranks.
-    for sbc_idx, (ths, theta_i, x_i) in tqdm(
-        enumerate(zip(posterior_samples, thetas, xs)),
+    for sbc_idx, (true_theta, x_i) in tqdm(
+        enumerate(zip(thetas, xs)),
         total=num_sbc_samples,
         disable=not show_progress_bar,
         desc=f"Calculating ranks for {num_sbc_samples} sbc samples.",
@@ -139,8 +145,12 @@ def _run_sbc(
 
         # For each reduce_fn (e.g., per marginal for SBC)
         for dim_idx, reduce_fn in enumerate(reduce_fns):
+            # rank posterior samples against true parameter, reduced to 1D.
             ranks[sbc_idx, dim_idx] = (
-                (reduce_fn(ths, x_i) < reduce_fn(theta_i.unsqueeze(0), x_i))
+                (
+                    reduce_fn(posterior_samples[:, sbc_idx, :], x_i)
+                    < reduce_fn(true_theta.unsqueeze(0), x_i)
+                )
                 .sum()
                 .item()
             )

sbi/diagnostics/tarp.py

@@ -64,7 +64,7 @@ def run_tarp(
     posterior_samples = get_posterior_samples_on_batch(
         xs,
         posterior,
-        num_posterior_samples,
+        (num_posterior_samples,),
         num_workers,
         show_progress_bar=show_progress_bar,
     )

sbi/utils/diagnostics_utils.py

@@ -5,24 +5,26 @@
 
 from sbi.inference.posteriors.base_posterior import NeuralPosterior
 from sbi.inference.posteriors.vi_posterior import VIPosterior
+from sbi.sbi_types import Shape
 
 
 def get_posterior_samples_on_batch(
     xs: Tensor,
     posterior: NeuralPosterior,
-    num_samples: int,
+    sample_shape: Shape,
     num_workers: int = 1,
     show_progress_bar: bool = False,
+    use_batched_sampling: bool = True,
 ) -> Tensor:
     """Get posterior samples for a batch of xs.
 
     Args:
         xs: batch of observations.
         posterior: sbi posterior.
-        num_posterior_samples: number of samples to draw from the posterior in each sbc
-            run.
+        num_samples: number of samples to draw from the posterior for each x.
         num_workers: number of workers to use for parallelization.
         show_progress_bars: whether to show progress bars.
+        use_batched_sampling: whether to use batched sampling if possible.
 
     Returns:
         posterior_samples: of shape (num_samples, batch_size, dim_parameters).
@@ -32,35 +34,44 @@ def get_posterior_samples_on_batch(
     # Try using batched sampling when implemented.
     try:
         # has shape (num_samples, batch_size, dim_parameters)
-        posterior_samples = posterior.sample_batched(
-            (num_samples,), xs, show_progress_bars=show_progress_bar
-        )
+        if use_batched_sampling:
+            posterior_samples = posterior.sample_batched(
+                sample_shape, x=xs, show_progress_bars=show_progress_bar
+            )
+        else:
+            raise NotImplementedError
     except NotImplementedError:
         # We need a function with extra training step for new x for VIPosterior.
-        def sample_fun(posterior: NeuralPosterior, sample_shape, x: Tensor) -> Tensor:
+        def sample_fun(
+            posterior: NeuralPosterior, sample_shape: Shape, x: Tensor, seed: int = 0
+        ) -> Tensor:
             if isinstance(posterior, VIPosterior):
                 posterior.set_default_x(x)
                 posterior.train()
+            torch.manual_seed(seed)
             return posterior.sample(sample_shape, x=x, show_progress_bars=False)
 
         # Run in parallel with progress bar.
+        seeds = torch.randint(0, 2**32, (batch_size,))
         outputs = list(
             tqdm(
                 Parallel(return_as="generator", n_jobs=num_workers)(
-                    delayed(sample_fun)(posterior, (num_samples,), x=x) for x in xs
+                    delayed(sample_fun)(posterior, sample_shape, x=x, seed=s)
+                    for x, s in zip(xs, seeds)
                 ),
                 disable=not show_progress_bar,
                 total=len(xs),
-                desc=f"Sampling {batch_size} times {num_samples} posterior samples.",
+                desc=f"Sampling {batch_size} times {sample_shape} posterior samples.",
             )
-        )
-        # Transpose to sample_batched shape convention:
-        posterior_samples = torch.stack(outputs).transpose(0, 1)  # type: ignore
+        )  # (batch_size, num_samples, dim_parameters)
+        # Transpose to shape convention: (sample_shape, batch_size, dim_parameters)
+        posterior_samples = torch.stack(
+            outputs  # type: ignore
+        ).permute(1, 0, 2)
 
-    assert posterior_samples.shape[:2] == (
-        num_samples,
+    assert posterior_samples.shape[:2] == sample_shape + (
         batch_size,
-    ), f"""Expected batched posterior samples of shape {(num_samples, batch_size)} got {
-        posterior_samples.shape[:2]
-    }."""
+    ), f"""Expected batched posterior samples of shape {
+        sample_shape + (batch_size,)
+    } got {posterior_samples.shape[:2]}."""
     return posterior_samples

tests/sbc_test.py

@@ -13,7 +13,9 @@
 from sbi.analysis import sbc_rank_plot
 from sbi.diagnostics import check_sbc, get_nltp, run_sbc
 from sbi.inference import SNLE, SNPE, simulate_for_sbi
-from sbi.simulators import linear_gaussian
+from sbi.simulators.linear_gaussian import (
+    linear_gaussian,
+)
 from sbi.utils import BoxUniform, MultipleIndependent
 from sbi.utils.user_input_checks import process_prior, process_simulator
 from tests.test_utils import PosteriorPotential, TractablePosterior
@@ -91,29 +93,21 @@ def simulator(theta):
 
 
 @pytest.mark.slow
-@pytest.mark.parametrize(
-    "density_estimator",
-    [
-        pytest.param(
-            "mdn",
-            marks=pytest.mark.xfail(
-                reason="MDN batched sampling results in miscalibrated posteriors",
-                strict=True,
-            ),
-        ),
-        "maf",
-    ],
-)
+@pytest.mark.parametrize("density_estimator", ["mdn", "maf"])
 @pytest.mark.parametrize("cov_method", ("sbc", "coverage"))
 def test_consistent_sbc_results(density_estimator, cov_method):
     """Test consistent SBC results on well-trained NPE."""
 
-    num_dim = 3
-    prior = BoxUniform(low=-2 * torch.ones(num_dim), high=2 * torch.ones(num_dim))
+    num_dim = 2
+
+    likelihood_shift = -1.0 * ones(num_dim)
+    likelihood_cov = 0.3 * eye(num_dim)
+    prior_mean = zeros(num_dim)
+    prior_cov = eye(num_dim)
+    prior = MultivariateNormal(loc=prior_mean, covariance_matrix=prior_cov)
 
     def simulator(theta):
-        # linear gaussian
-        return theta + 1.0 + torch.randn_like(theta) * 0.1
+        return linear_gaussian(theta, likelihood_shift, likelihood_cov)
 
     num_simulations = 2000
     num_posterior_samples = 1000