Uniform and better MCMC params for the tests

pyproject.toml

@@ -118,7 +118,8 @@ testpaths = [
 ]
 markers = [
     "slow: marks tests as slow (deselect with '-m \"not slow\"')",
-    "gpu: marks tests that require a gpu (deselect with '-m \"not gpu\"')"
+    "gpu: marks tests that require a gpu (deselect with '-m \"not gpu\"')",
+    "mcmc: marks tests that require MCMC sampling (deselect with '-m \"not mcmc\"')"
 ]
 
 # Pyright configuration

sbi/neural_nets/density_estimators/mixed_density_estimator.py

@@ -150,7 +150,8 @@ def log_prob_iid(self, x: Tensor, context: Tensor) -> Tensor:
         inference. The speed up is achieved by exploiting the fact that there are only
         finite number of possible categories in the discrete part of the dat: one can
         just calculate the log probs for each possible category (given the current batch
-        of context) and then copy those log probs into the entire batch of iid categories.
+        of context) and then copy those log probs into the entire batch of iid
+        categories.
         For example, for the drift-diffusion model, there are only two choices, but
         often 100s or 1000 trials. With this method a evaluation over trials then passes
         a batch of `2 (one per choice) * num_contexts` into the NN, whereas the normal
@@ -175,7 +176,8 @@ def log_prob_iid(self, x: Tensor, context: Tensor) -> Tensor:
         net_device = next(self.discrete_net.parameters()).device
         assert (
             net_device == x.device and x.device == context.device
-        ), f"device mismatch: net, x, context: {net_device}, {x.device}, {context.device}."
+        ), f"device mismatch: net, x, context: \
+            {net_device}, {x.device}, {context.device}."
 
         x_cont_repeated, x_disc_repeated = _separate_x(x_repeated)
         x_cont, x_disc = _separate_x(x)

sbi/utils/user_input_checks.py

@@ -610,8 +610,8 @@ def process_x(
 def prepare_for_sbi(simulator: Callable, prior) -> Tuple[Callable, Distribution]:
     """Prepare simulator and prior for usage in sbi.
 
-    NOTE: This method is deprecated as of sbi version v0.23.0. and will be removed in a future release.
-    Please use `process_prior` and `process_simulator` in the future.
+    NOTE: This method is deprecated as of sbi version v0.23.0. and will be removed in a
+    future release. Please use `process_prior` and `process_simulator` in the future.
     This is a wrapper around `process_prior` and `process_simulator` which can be
     used in isolation as well.
 
@@ -633,7 +633,8 @@ def prepare_for_sbi(simulator: Callable, prior) -> Tuple[Callable, Distribution]
     """
 
     warnings.warn(
-        "This method is deprecated as of sbi version v0.23.0. and will be removed in a future release."
+        "This method is deprecated as of sbi version v0.23.0. and will be removed in a \
+        future release."
         "Please use `process_prior` and `process_simulator` in the future.",
         DeprecationWarning,
         stacklevel=2,

tests/conftest.py

@@ -29,3 +29,15 @@ def pytest_collection_modifyitems(config, items):
         for item in items:
             if "gpu" in item.keywords:
                 item.add_marker(skip_gpu)
+
+
+@pytest.fixture(scope="function")
+def mcmc_params_accurate() -> dict:
+    """Fixture for MCMC parameters for functional tests."""
+    return dict(num_chains=20, thin=2, warmup_steps=50)
+
+
+@pytest.fixture(scope="function")
+def mcmc_params_fast() -> dict:
+    """Fixture for MCMC parameters for fast tests."""
+    return dict(num_chains=1, thin=1, warmup_steps=10)

tests/embedding_net_test.py

@@ -21,10 +21,13 @@
 from .test_utils import check_c2st
 
 
+@pytest.mark.mcmc
 @pytest.mark.parametrize("method", ["SNPE", "SNLE", "SNRE"])
 @pytest.mark.parametrize("num_dim", [1, 2])
 @pytest.mark.parametrize("embedding_net", ["mlp"])
-def test_embedding_net_api(method, num_dim: int, embedding_net: str):
+def test_embedding_net_api(
+    method, num_dim: int, embedding_net: str, mcmc_params_fast: dict
+):
     """Tests the API when using a preconfigured embedding net."""
 
     x_o = zeros(1, num_dim)
@@ -62,7 +65,7 @@ def test_embedding_net_api(method, num_dim: int, embedding_net: str):
     _ = inference.append_simulations(theta, x).train(max_num_epochs=2)
     posterior = inference.build_posterior(
         mcmc_method="slice_np_vectorized",
-        mcmc_parameters=dict(num_chains=2, warmup_steps=10, thin=5),
+        mcmc_parameters=mcmc_params_fast,
     ).set_default_x(x_o)
 
     s = posterior.sample((1,))

tests/ensemble_test.py

@@ -48,13 +48,15 @@ def simulator(theta):
     (
         (SNPE_C, 1),
         pytest.param(SNPE_C, 5, marks=pytest.mark.xfail),
-        (SNLE_A, 1),
-        (SNLE_A, 5),
-        (SNRE_A, 1),
-        (SNRE_A, 5),
+        pytest.param(SNLE_A, 1, marks=pytest.mark.mcmc),
+        pytest.param(SNLE_A, 5, marks=pytest.mark.mcmc),
+        pytest.param(SNRE_A, 1, marks=pytest.mark.mcmc),
+        pytest.param(SNRE_A, 5, marks=pytest.mark.mcmc),
     ),
 )
-def test_c2st_posterior_ensemble_on_linearGaussian(inference_method, num_trials):
+def test_c2st_posterior_ensemble_on_linearGaussian(
+    inference_method, num_trials, mcmc_params_accurate: dict
+):
     """Test whether EnsemblePosterior infers well a simple example with available
     ground truth.
     """
@@ -63,7 +65,7 @@ def test_c2st_posterior_ensemble_on_linearGaussian(inference_method, num_trials)
     ensemble_size = 2
     x_o = zeros(num_trials, num_dim)
     num_samples = 500
-    num_simulations = 2000 if inference_method == SNRE_A else 1500
+    num_simulations = 2000
 
     # likelihood_mean will be likelihood_shift+theta
     likelihood_shift = -1.0 * ones(num_dim)
@@ -100,10 +102,8 @@ def simulator(theta):
     if isinstance(inferer, (SNLE_A, SNRE_A)):
         samples = posterior.sample(
             (num_samples,),
-            num_chains=20,
             method="slice_np_vectorized",
-            thin=5,
-            warmup_steps=50,
+            **mcmc_params_accurate,
         )
     else:
         samples = posterior.sample((num_samples,))

tests/inference_on_device_test.py

@@ -45,20 +45,20 @@
     [
         (SNPE_C, "maf", "direct"),
         (SNPE_C, "mdn", "rejection"),
-        (SNPE_C, "maf", "slice_np_vectorized"),
-        (SNPE_C, "mdn", "slice"),
-        (SNLE, "nsf", "slice_np_vectorized"),
-        (SNLE, "mdn", "slice"),
+        pytest.param(SNPE_C, "maf", "slice_np_vectorized", marks=pytest.mark.mcmc),
+        pytest.param(SNPE_C, "mdn", "slice", marks=pytest.mark.mcmc),
+        pytest.param(SNLE, "nsf", "slice_np_vectorized", marks=pytest.mark.mcmc),
+        pytest.param(SNLE, "mdn", "slice", marks=pytest.mark.mcmc),
         (SNLE, "nsf", "rejection"),
         (SNLE, "maf", "importance"),
-        (SNRE_A, "mlp", "slice_np_vectorized"),
-        (SNRE_A, "mlp", "slice"),
+        pytest.param(SNRE_A, "mlp", "slice_np_vectorized", marks=pytest.mark.mcmc),
+        pytest.param(SNRE_A, "mlp", "slice", marks=pytest.mark.mcmc),
         (SNRE_B, "resnet", "rejection"),
         (SNRE_B, "resnet", "importance"),
-        (SNRE_B, "resnet", "slice"),
+        pytest.param(SNRE_B, "resnet", "slice", marks=pytest.mark.mcmc),
         (SNRE_C, "resnet", "rejection"),
         (SNRE_C, "resnet", "importance"),
-        (SNRE_C, "resnet", "nuts"),
+        pytest.param(SNRE_C, "resnet", "nuts", marks=pytest.mark.mcmc),
     ],
 )
 @pytest.mark.parametrize(
@@ -71,7 +71,13 @@
 )
 @pytest.mark.parametrize("prior_type", ["gaussian", "uniform"])
 def test_training_and_mcmc_on_device(
-    method, model, sampling_method, training_device, prior_device, prior_type
+    method,
+    model,
+    sampling_method,
+    training_device,
+    prior_device,
+    prior_type,
+    mcmc_params_fast: dict,
 ):
     """Test training on devices.
 
@@ -85,9 +91,11 @@ def test_training_and_mcmc_on_device(
 
     num_dim = 2
     num_samples = 10
-    num_simulations = 100
-    max_num_epochs = 2
+    max_num_epochs = 10
     num_rounds = 2  # test proposal sampling in round 2.
+    num_simulations_per_round = [200, num_samples]
+    # use more warmup steps to avoid Infs during MCMC in round two.
+    mcmc_params_fast["warmup_steps"] = 20
 
     x_o = zeros(1, num_dim).to(data_device)
     likelihood_shift = -1.0 * ones(num_dim).to(prior_device)
@@ -134,7 +142,7 @@ def simulator(theta):
     proposals = [prior]
 
     for _ in range(num_rounds):
-        theta = proposals[-1].sample((num_simulations,))
+        theta = proposals[-1].sample((num_simulations_per_round[_],))
         x = simulator(theta).to(data_device)
         theta = theta.to(data_device)
 
@@ -147,10 +155,7 @@ def simulator(theta):
             posterior = inferer.build_posterior(
                 sample_with="mcmc",
                 mcmc_method=sampling_method,
-                mcmc_parameters=dict(
-                    thin=10 if sampling_method == "slice_np_vectorized" else 1,
-                    num_chains=10 if sampling_method == "slice_np_vectorized" else 1,
-                ),
+                mcmc_parameters=mcmc_params_fast,
             )
         elif sampling_method in ["rejection", "direct"]:
             # all other cases: rejection, direct
@@ -331,6 +336,7 @@ def test_embedding_nets_integration_training_device(
     embedding_net_device: str,
     data_device: str,
     training_device: str,
+    mcmc_params_fast: dict,
 ) -> None:
     """Test embedding nets integration with different devices, priors and methods."""
     # add other methods
@@ -436,7 +442,7 @@ def test_embedding_nets_integration_training_device(
                 if inference_method == SNPE_A
                 else dict(
                     mcmc_method="slice_np_vectorized",
-                    mcmc_parameters=dict(thin=10, num_chains=20, warmup_steps=10),
+                    mcmc_parameters=mcmc_params_fast,
                 )
             ),
         )

tests/linearGaussian_mdn_test.py

@@ -23,16 +23,10 @@
 from tests.test_utils import check_c2st
 
 
-def test_mdn_with_snpe():
-    mdn_inference_with_different_methods(SNPE)
-
-
-@pytest.mark.slow
-def test_mdn_with_snle():
-    mdn_inference_with_different_methods(SNLE)
-
-
-def mdn_inference_with_different_methods(method):
+@pytest.mark.parametrize(
+    "method", (SNPE, pytest.param(SNLE, marks=[pytest.mark.slow, pytest.mark.mcmc]))
+)
+def test_mdn_inference_with_different_methods(method, mcmc_params_accurate: dict):
     num_dim = 2
     x_o = torch.tensor([[1.0, 0.0]])
     num_samples = 500
@@ -68,9 +62,7 @@ def simulator(theta: Tensor) -> Tensor:
             theta_transform=theta_transform,
             proposal=prior,
             method="slice_np_vectorized",
-            num_chains=20,
-            warmup_steps=50,
-            thin=5,
+            **mcmc_params_accurate,
         )
 
     samples = posterior.sample((num_samples,), x=x_o)