test: update tests to work with updated build_posterior logic

Author: abelaba

tests/embedding_net_test.py

@@ -13,6 +13,7 @@
 
 from sbi import utils
 from sbi.inference import NLE, NPE, NRE, simulate_for_sbi
+from sbi.inference.posteriors.posterior_parameters import MCMCPosteriorParameters
 from sbi.neural_nets import classifier_nn, likelihood_nn, posterior_nn
 from sbi.neural_nets.embedding_nets import (
     CNNEmbedding,
@@ -81,8 +82,9 @@ def test_embedding_net_api(
 
     _ = inference.append_simulations(theta, x).train(max_num_epochs=2)
     posterior = inference.build_posterior(
-        mcmc_method="slice_np_vectorized",
-        mcmc_parameters=mcmc_params_fast,
+        posterior_parameters=MCMCPosteriorParameters(
+            method="slice_np_vectorized", **mcmc_params_fast
+        )
     ).set_default_x(x_o)
 
     s = posterior.sample((1,))

tests/inference_on_device_test.py

@@ -33,6 +33,9 @@
 )
 from sbi.inference.posteriors.importance_posterior import ImportanceSamplingPosterior
 from sbi.inference.posteriors.mcmc_posterior import MCMCPosterior
+from sbi.inference.posteriors.posterior_parameters import (
+    MCMCPosteriorParameters,
+)
 from sbi.inference.potentials.base_potential import BasePotential
 from sbi.inference.potentials.likelihood_based_potential import LikelihoodBasedPotential
 from sbi.inference.potentials.posterior_based_potential import PosteriorBasedPotential
@@ -190,9 +193,9 @@ def simulator(theta):
         # mcmc cases
         if sampling_method in ["slice_np", "slice_np_vectorized", "nuts_pymc"]:
             posterior = inferer.build_posterior(
-                sample_with="mcmc",
-                mcmc_method=sampling_method,
-                mcmc_parameters=mcmc_params_fast,
+                posterior_parameters=MCMCPosteriorParameters(
+                    method=sampling_method, **mcmc_params_fast
+                )
             )
         elif sampling_method in ["rejection", "direct"]:
             # all other cases: rejection, direct

tests/linearGaussian_snle_test.py

@@ -19,6 +19,7 @@
     VIPosterior,
     likelihood_estimator_based_potential,
 )
+from sbi.inference.posteriors.posterior_parameters import MCMCPosteriorParameters
 from sbi.neural_nets import likelihood_nn
 from sbi.simulators.linear_gaussian import (
     diagonal_linear_gaussian,
@@ -64,8 +65,9 @@ def test_api_nle_multiple_trials_and_rounds_map(
         for num_trials in [1, 3]:
             x_o = zeros((num_trials, num_dim))
             posterior = inference.build_posterior(
-                mcmc_method="slice_np_vectorized",
-                mcmc_parameters=mcmc_params_fast,
+                posterior_parameters=MCMCPosteriorParameters(
+                    method="slice_np_vectorized", **mcmc_params_fast
+                )
             ).set_default_x(x_o)
             posterior.sample(sample_shape=(num_samples,))
         proposals.append(posterior)

tests/linearGaussian_snre_test.py

@@ -23,6 +23,7 @@
     VIPosterior,
     ratio_estimator_based_potential,
 )
+from sbi.inference.posteriors.posterior_parameters import MCMCPosteriorParameters
 from sbi.neural_nets.ratio_estimators import RatioEstimator
 from sbi.simulators.linear_gaussian import (
     diagonal_linear_gaussian,
@@ -65,8 +66,9 @@ def test_api_nre_multiple_trials_and_rounds_map(
         for num_trials in [1, 3]:
             x_o = zeros((num_trials, num_dim))
             posterior = inference.build_posterior(
-                mcmc_method="slice_np_vectorized",
-                mcmc_parameters=mcmc_params_fast,
+                posterior_parameters=MCMCPosteriorParameters(
+                    method="slice_np_vectorized", **mcmc_params_fast
+                )
             ).set_default_x(x_o)
             posterior.sample(sample_shape=(num_samples,))
         proposals.append(posterior)

tests/linearGaussian_vector_field_test.py

@@ -21,6 +21,7 @@
     simulate_for_sbi,
     vector_field_estimator_based_potential,
 )
+from sbi.inference.posteriors.posterior_parameters import VectorFieldPosteriorParameters
 from sbi.neural_nets.factory import flowmatching_nn
 from sbi.simulators import linear_gaussian
 from sbi.simulators.linear_gaussian import (
@@ -103,9 +104,9 @@ def test_c2st_vector_field_on_linearGaussian(
         posterior = inference.build_posterior(
             score_estimator,
             sample_with=method,
-            vectorfield_sampling_parameters={
-                "neural_ode_backend": "zuko",
-            },
+            posterior_parameters=VectorFieldPosteriorParameters(
+                neural_ode_backend="zuko"
+            ),
         )
         posterior.set_default_x(x_o)
         samples = posterior.sample((num_samples,))

tests/mnle_test.py

@@ -10,6 +10,7 @@
 from torch.distributions import Beta, Binomial, Distribution, Gamma
 
 from sbi.inference import MNLE, MCMCPosterior
+from sbi.inference.posteriors.posterior_parameters import MCMCPosteriorParameters
 from sbi.inference.posteriors.rejection_posterior import RejectionPosterior
 from sbi.inference.posteriors.vi_posterior import VIPosterior
 from sbi.inference.potentials.base_potential import BasePotential
@@ -182,7 +183,9 @@ def test_mnle_accuracy_with_different_samplers_and_trials(
     ).sample((num_samples,), show_progress_bars=False)
 
     posterior = trainer.build_posterior(
-        prior=prior, sample_with=sampler, mcmc_parameters=mcmc_params_accurate
+        prior=prior,
+        sample_with=sampler,
+        posterior_parameters=MCMCPosteriorParameters(**mcmc_params_accurate),
     )
     posterior.set_default_x(x_o)
     if sampler == "vi":

tests/posterior_nn_test.py

@@ -21,7 +21,10 @@
     NRE_C,
     DirectPosterior,
 )
-from sbi.inference.posteriors.posterior_parameters import RejectionPosteriorParameters
+from sbi.inference.posteriors.posterior_parameters import (
+    MCMCPosteriorParameters,
+    RejectionPosteriorParameters,
+)
 from sbi.simulators.linear_gaussian import (
     diagonal_linear_gaussian,
     linear_gaussian,
@@ -195,9 +198,10 @@ def test_batched_mcmc_sample_log_prob_with_different_x(
     x_o = ones(num_dim) if x_o_batch_dim == 0 else ones(x_o_batch_dim, num_dim)
 
     posterior = inference.build_posterior(
-        sample_with="mcmc",
-        mcmc_method="slice_np_vectorized",
-        mcmc_parameters=mcmc_params_fast,
+        posterior_parameters=MCMCPosteriorParameters(
+            method="slice_np_vectorized",
+            **mcmc_params_fast,
+        )
     )
 
     samples = posterior.sample_batched(
@@ -219,9 +223,10 @@ def test_batched_mcmc_sample_log_prob_with_different_x(
         inference = snlre_method(prior=prior)
         _ = inference.append_simulations(theta, x).train()
         posterior = inference.build_posterior(
-            sample_with="mcmc",
-            mcmc_method="slice_np_vectorized",
-            mcmc_parameters=mcmc_params_fast,
+            posterior_parameters=MCMCPosteriorParameters(
+                method="slice_np_vectorized",
+                **mcmc_params_fast,
+            )
         )
 
         x_o = torch.stack([0.5 * ones(num_dim), -0.5 * ones(num_dim)], dim=0)

tests/posterior_parameters_test.py

@@ -385,3 +385,12 @@ def test_invalid_literal_field_values():
     """
 
     MCMCPosteriorParameters(method="invalid")
+
+
+def test_if_warning_raised_for_deprecated_build_posterior_parameters(get_inference):
+    """
+    Check if the build_posterior method raises a warning for deprecated parameters
+    """
+
+    with pytest.warns(DeprecationWarning):
+        get_inference.build_posterior(mcmc_parameters={})

tests/save_and_load_test.py

@@ -8,21 +8,27 @@
 
 from sbi import utils as utils
 from sbi.inference import NLE, NPE, NRE
+from sbi.inference.posteriors.posterior_parameters import (
+    DirectPosteriorParameters,
+    MCMCPosteriorParameters,
+    RejectionPosteriorParameters,
+    VIPosteriorParameters,
+)
 from sbi.inference.posteriors.vi_posterior import VIPosterior
 
 
 @pytest.mark.parametrize(
-    "inference_method, sampling_method",
+    "inference_method, posterior_parameter",
     (
-        (NPE, "direct"),
-        pytest.param(NLE, "mcmc", marks=pytest.mark.mcmc),
-        pytest.param(NRE, "mcmc", marks=pytest.mark.mcmc),
-        pytest.param(NRE, "vi", marks=pytest.mark.mcmc),
-        (NRE, "rejection"),
+        (NPE, DirectPosteriorParameters),
+        pytest.param(NLE, MCMCPosteriorParameters, marks=pytest.mark.mcmc),
+        pytest.param(NRE, MCMCPosteriorParameters, marks=pytest.mark.mcmc),
+        pytest.param(NRE, VIPosteriorParameters, marks=pytest.mark.mcmc),
+        (NRE, RejectionPosteriorParameters),
     ),
 )
 def test_picklability(
-    inference_method, sampling_method: str, tmp_path, mcmc_params_fast
+    inference_method, posterior_parameter, tmp_path, mcmc_params_fast
 ):
     num_dim = 2
     prior = utils.BoxUniform(low=-2 * torch.ones(num_dim), high=2 * torch.ones(num_dim))
@@ -33,10 +39,14 @@ def test_picklability(
 
     inference = inference_method(prior=prior)
     _ = inference.append_simulations(theta, x).train(max_num_epochs=1)
-    posterior = inference.build_posterior(
-        sample_with=sampling_method, mcmc_parameters=mcmc_params_fast
-    ).set_default_x(x_o)
-
+    if posterior_parameter is MCMCPosteriorParameters:
+        posterior = inference.build_posterior(
+            posterior_parameters=posterior_parameter(**mcmc_params_fast)
+        ).set_default_x(x_o)
+    else:
+        posterior = inference.build_posterior(
+            posterior_parameters=posterior_parameter()
+        ).set_default_x(x_o)
     # After sample and log_prob, the posterior should still be picklable
     if isinstance(posterior, VIPosterior):
         posterior.train(max_num_iters=10)

tests/sbc_test.py

@@ -14,6 +14,10 @@
 from sbi.diagnostics import check_sbc, get_nltp, run_sbc
 from sbi.inference import NLE, NPE, NPSE
 from sbi.inference.posteriors.base_posterior import NeuralPosterior
+from sbi.inference.posteriors.posterior_parameters import (
+    MCMCPosteriorParameters,
+    VIPosteriorParameters,
+)
 from sbi.simulators.linear_gaussian import linear_gaussian
 from sbi.utils import BoxUniform, MultipleIndependent
 from tests.test_utils import PosteriorPotential, TractablePosterior
@@ -106,9 +110,11 @@ def simulator(theta):
     posterior_kwargs = {}
     if method == NLE:
         posterior_kwargs = {
-            "sample_with": "mcmc" if sampler == "mcmc" else "vi",
-            "mcmc_method": "slice_np_vectorized",
-            "mcmc_parameters": mcmc_params_fast,
+            "posterior_parameters": MCMCPosteriorParameters(
+                method="slice_np_vectorized", **mcmc_params_fast
+            )
+            if sampler == "mcmc"
+            else VIPosteriorParameters()
         }
 
     posterior = train_inference_method(