Refactor CI tests for speed

sbi/inference/abc/abc_base.py

@@ -1,4 +1,5 @@
 """Base class for Approximate Bayesian Computation methods."""
+
 import logging
 from abc import ABC
 from typing import Callable, Union

sbi/inference/abc/mcabc.py

@@ -1,4 +1,5 @@
 """Monte-Carlo Approximate Bayesian Computation (Rejection ABC)."""
+
 from typing import Any, Callable, Dict, Optional, Tuple, Union
 
 import torch

sbi/inference/abc/smcabc.py

@@ -1,4 +1,5 @@
 """Sequential Monte Carlo Approximate Bayesian Computation."""
+
 from typing import Any, Callable, Dict, Optional, Tuple, Union
 
 import numpy as np

sbi/inference/posteriors/mcmc_posterior.py

@@ -117,9 +117,9 @@
                 v0.19.0. Instead, use e.g.,
                 `init_strategy_parameters={"num_candidate_samples": 1000}`"""
             )
-            self.init_strategy_parameters[
-                "num_candidate_samples"
-            ] = init_strategy_num_candidates
+            self.init_strategy_parameters["num_candidate_samples"] = (
+                init_strategy_num_candidates
+            )
 
         self.potential_ = self._prepare_potential(method)
 
@@ -239,9 +239,9 @@
                 v0.19.0. Instead, use e.g.,
                 `init_strategy_parameters={"num_candidate_samples": 1000}`"""
             )
-            self.init_strategy_parameters[
-                "num_candidate_samples"
-            ] = init_strategy_num_candidates
+            self.init_strategy_parameters["num_candidate_samples"] = (
+                init_strategy_num_candidates
+            )
         if sample_with is not None:
             raise ValueError(
                 f"You set `sample_with={sample_with}`. As of sbi v0.18.0, setting "
@@ -661,9 +661,9 @@
             self._posterior_sampler is not None
         ), """No samples have been generated, call .sample() first."""
 
-        sampler: Union[
-            MCMC, SliceSamplerSerial, SliceSamplerVectorized
-        ] = self._posterior_sampler
+        sampler: Union[MCMC, SliceSamplerSerial, SliceSamplerVectorized] = (
+            self._posterior_sampler
+        )
 
         # If Pyro sampler and samples not transformed, use arviz' from_pyro.
         # Exclude 'slice' kernel as it lacks the 'divergence' diagnostics key.

sbi/neural_nets/mnle.py

@@ -81,9 +81,9 @@ def build_mnle(
 
     # Set up a NSF for modelling the continuous data, conditioned on the discrete data.
     cont_nle = build_nsf(
-        batch_x=torch.log(cont_x)
-        if log_transform_x
-        else cont_x,  # log transform manually.
+        batch_x=(
+            torch.log(cont_x) if log_transform_x else cont_x
+        ),  # log transform manually.
         batch_y=torch.cat((batch_y, disc_x), dim=1),  # condition on discrete data too.
         z_score_y=z_score_y,
         z_score_x=z_score_x,

sbi/utils/user_input_checks_utils.py

@@ -166,9 +166,9 @@ def __init__(
         super().__init__(
             batch_shape=batch_shape,
             event_shape=event_shape,
-            validate_args=prior._validate_args
-            if validate_args is None
-            else validate_args,
+            validate_args=(
+                prior._validate_args if validate_args is None else validate_args
+            ),
         )
 
         self.prior = prior

tests/embedding_net_test.py

@@ -3,8 +3,9 @@
 import pytest
 import torch
 from torch import eye, ones, zeros
+from torch.distributions import MultivariateNormal
 
-from sbi import utils as utils
+from sbi import utils
 from sbi.inference import SNLE, SNPE, SNRE, simulate_for_sbi
 from sbi.neural_nets.embedding_nets import (
     CNNEmbedding,
@@ -58,16 +59,20 @@ def test_embedding_net_api(method, num_dim: int, embedding_net: str):
     else:
         raise NameError
 
-    _ = inference.append_simulations(theta, x).train(max_num_epochs=5)
-    posterior = inference.build_posterior().set_default_x(x_o)
+    _ = 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),
+    ).set_default_x(x_o)
 
     s = posterior.sample((1,))
     _ = posterior.potential(s)
 
 
 @pytest.mark.parametrize("num_trials", [1, 2])
 @pytest.mark.parametrize("num_dim", [1, 2])
-def test_iid_embedding_api(num_trials, num_dim):
+def test_embedding_api_with_multiple_trials(num_trials, num_dim):
+    """Tests the API when using iid trial-based data."""
     prior = utils.BoxUniform(-2.0 * ones(num_dim), 2.0 * ones(num_dim))
 
     num_thetas = 1000
@@ -97,7 +102,7 @@ def test_iid_embedding_api(num_trials, num_dim):
 
 @pytest.mark.slow
 def test_iid_embedding_varying_num_trials(trial_factor=50, max_num_trials=20):
-    """Test embedding net with varying number of trials."""
+    """Test inference accuracy with embeddings for varying number of trials."""
     num_dim = 2
     prior = torch.distributions.MultivariateNormal(
         torch.zeros(num_dim), torch.eye(num_dim)
@@ -241,9 +246,6 @@ def simulator(theta, num_trials=num_trials):
 @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
-    from torch.distributions import MultivariateNormal
-
     estimator_provider = posterior_nn(
         "mdn",
         embedding_net=CNNEmbedding(

tests/inference_on_device_test.py

@@ -423,9 +423,11 @@ def test_embedding_nets_integration_training_device(
             .to(data_device)
         )
 
-        with pytest.warns(
-            UserWarning
-        ) if data_device != training_device else nullcontext():
+        with (
+            pytest.warns(UserWarning)
+            if data_device != training_device
+            else nullcontext()
+        ):
             density_estimator_append = inference.append_simulations(theta, X)
 
         density_estimator_train = density_estimator_append.train(

tests/linearGaussian_snle_test.py

@@ -29,61 +29,65 @@
 
 from .test_utils import check_c2st, get_prob_outside_uniform_prior
 
+# mcmc params for fast testing.
+mcmc_parameters = {
+    "method": "slice_np_vectorized",
+    "num_chains": 20,
+    "thin": 5,
+    "warmup_steps": 50,
+}
 
-@pytest.mark.parametrize("num_dim", (1, 3))
-def test_api_snl_on_linearGaussian(num_dim: int):
-    """Test API for inference on linear Gaussian model using SNL.
 
-    Avoids expensive computations by training on few simulations and generating few
-    posterior samples.
-
-    Args:
-        num_dim: parameter dimension of the gaussian model
-    """
-    num_samples = 10
+@pytest.mark.parametrize("num_dim", (1,))  # dim 3 is tested below.
+@pytest.mark.parametrize("prior_str", ("uniform", "gaussian"))
+def test_api_snle_multiple_trials_and_rounds_map(num_dim: int, prior_str: str):
+    """Test SNLE API with 2 rounds, different priors num trials and MAP."""
+    num_rounds = 2
+    num_samples = 1
+    num_simulations = 100
 
-    prior_mean = zeros(num_dim)
-    prior_cov = eye(num_dim)
-    prior = MultivariateNormal(loc=prior_mean, covariance_matrix=prior_cov)
+    if prior_str == "gaussian":
+        prior_mean = zeros(num_dim)
+        prior_cov = eye(num_dim)
+        prior = MultivariateNormal(loc=prior_mean, covariance_matrix=prior_cov)
+    else:
+        prior = BoxUniform(-2.0 * ones(num_dim), 2.0 * ones(num_dim))
 
     simulator, prior = prepare_for_sbi(diagonal_linear_gaussian, prior)
-    density_estimator = likelihood_nn("maf", num_transforms=3)
-    inference = SNLE(density_estimator=density_estimator, show_progress_bars=False)
+    inference = SNLE(prior=prior, density_estimator="mdn", show_progress_bars=False)
 
-    theta, x = simulate_for_sbi(simulator, prior, 1000, simulation_batch_size=50)
-    likelihood_estimator = inference.append_simulations(theta, x).train(
-        training_batch_size=100
-    )
-
-    for num_trials in [1, 2]:
-        x_o = zeros((num_trials, num_dim))
-        potential_fn, theta_transform = likelihood_estimator_based_potential(
-            prior=prior, likelihood_estimator=likelihood_estimator, x_o=x_o
+    proposals = [prior]
+    for _ in range(num_rounds):
+        theta, x = simulate_for_sbi(
+            simulator,
+            proposals[-1],
+            num_simulations,
+            simulation_batch_size=num_simulations,
         )
-        posterior = MCMCPosterior(
-            proposal=prior,
-            potential_fn=potential_fn,
-            theta_transform=theta_transform,
-            thin=3,
+        inference.append_simulations(theta, x).train(
+            training_batch_size=100, max_num_epochs=2
         )
-        posterior.sample(sample_shape=(num_samples,))
-
+        for num_trials in [1, 3]:
+            x_o = zeros((num_trials, num_dim))
+            posterior = inference.build_posterior(
+                mcmc_method="slice_np_vectorized",
+                mcmc_parameters=dict(num_chains=10, thin=5, warmup_steps=10),
+            ).set_default_x(x_o)
+            posterior.sample(sample_shape=(num_samples,))
+        proposals.append(posterior)
+        posterior.map(num_iter=1)
 
-def test_c2st_snl_on_linearGaussian(density_estimator="maf"):
-    """Test whether SNL infers well a simple example with available ground truth.
 
-    This example has different number of parameters theta than number of x. This test
-    also acts as the only functional test for SNL not marked as slow.
-
-    """
+def test_c2st_snl_on_linear_gaussian_different_dims(model_str="maf"):
+    """Test SNLE on linear Gaussian task with different theta and x dims."""
 
     theta_dim = 3
     x_dim = 2
     discard_dims = theta_dim - x_dim
 
     x_o = zeros(1, x_dim)
-    num_samples = 1000
-    num_simulations = 3000
+    num_samples = 500
+    num_simulations = 1000
 
     # likelihood_mean will be likelihood_shift+theta
     likelihood_shift = -1.0 * ones(x_dim)
@@ -110,11 +114,11 @@ def test_c2st_snl_on_linearGaussian(density_estimator="maf"):
         ),
         prior,
     )
-    density_estimator = likelihood_nn(model=density_estimator, num_transforms=3)
+    density_estimator = likelihood_nn(model=model_str, num_transforms=3)
     inference = SNLE(density_estimator=density_estimator, show_progress_bars=False)
 
     theta, x = simulate_for_sbi(
-        simulator, prior, num_simulations, simulation_batch_size=50
+        simulator, prior, num_simulations, simulation_batch_size=num_simulations
     )
     likelihood_estimator = inference.append_simulations(theta, x).train()
     potential_fn, theta_transform = likelihood_estimator_based_potential(
@@ -124,14 +128,12 @@ def test_c2st_snl_on_linearGaussian(density_estimator="maf"):
         proposal=prior,
         potential_fn=potential_fn,
         theta_transform=theta_transform,
-        method="slice_np_vectorized",
-        num_chains=5,
-        thin=10,
+        **mcmc_parameters,
     )
     samples = posterior.sample((num_samples,))
 
     # Compute the c2st and assert it is near chance level of 0.5.
-    check_c2st(samples, target_samples, alg=f"snle_a-{density_estimator}")
+    check_c2st(samples, target_samples, alg=f"snle_a-{model_str}")
 
 
 @pytest.mark.slow
@@ -199,9 +201,7 @@ def test_c2st_and_map_snl_on_linearGaussian_different(num_dim: int, prior_str: s
             proposal=prior,
             potential_fn=potential_fn,
             theta_transform=theta_transform,
-            method="slice_np_vectorized",
-            thin=5,
-            num_chains=5,
+            **mcmc_parameters,
         )
 
         samples = posterior.sample(sample_shape=(num_samples,))
@@ -309,8 +309,7 @@ def test_c2st_multi_round_snl_on_linearGaussian(num_trials: int):
         proposal=prior,
         potential_fn=potential_fn,
         theta_transform=theta_transform,
-        thin=5,
-        num_chains=20,
+        **mcmc_parameters,
     )
 
     theta, x = simulate_for_sbi(
@@ -327,8 +326,7 @@ def test_c2st_multi_round_snl_on_linearGaussian(num_trials: int):
         proposal=prior,
         potential_fn=potential_fn,
         theta_transform=theta_transform,
-        thin=5,
-        num_chains=20,
+        **mcmc_parameters,
     )
 
     samples = posterior.sample(sample_shape=(num_samples,))
@@ -444,7 +442,7 @@ def test_api_snl_sampling_methods(
     num_simulations = 1000
     x_o = zeros((num_trials, num_dim))
     # Test for multiple chains is cheap when vectorized.
-    num_chains = 3 if sampling_method == "slice_np_vectorized" else 1
+    num_chains = 10 if sampling_method == "slice_np_vectorized" else 1
     if sampling_method == "rejection":
         sample_with = "rejection"
     elif (
@@ -491,7 +489,7 @@ def test_api_snl_sampling_methods(
                 proposal=prior,
                 theta_transform=theta_transform,
                 method=sampling_method,
-                thin=3,
+                thin=5,
                 num_chains=num_chains,
                 init_strategy=init_strategy,
             )