fix: slow tests. (#1209)

* fix: tsnpe sir test

* refactor: sample conditional test.

* fix: changes to fix remaining slow tests.

tests/lc2st_test.py

@@ -116,10 +116,10 @@ def test_lc2st_true_negativ_rate(method):
     num_runs = 100
     confidence_level = 0.95
 
-    # bad estimator :small training and num_epochs
+    # use little training data and num_epochs to obtain "bad" estimator.
     # (no convergence to the true posterior)
-    num_train = 1_000
-    num_epochs = 5
+    num_train = 100
+    num_epochs = 2
 
     num_cal = 1_000
     num_eval = 10_000

tests/linearGaussian_snpe_test.py

@@ -356,7 +356,7 @@ def simulator(theta):
     elif method_str.startswith("tsnpe"):
         sample_method = "rejection" if method_str == "tsnpe_rejection" else "sir"
         inference = SNPE_C(**creation_args)
-        theta = prior.sample((900,))
+        theta = prior.sample((1000,))
         x = simulator(theta)
         posterior_estimator = inference.append_simulations(theta, x).train()
         posterior1 = DirectPosterior(
@@ -493,39 +493,45 @@ def simulator(theta):
 @pytest.mark.mcmc
 def test_sample_conditional(mcmc_params_accurate: dict):
     """
-    Test whether sampling from the conditional gives the same results as evaluating.
+    Test whether sampling from the conditional gives the same results as
+    evaluating.
+
+    This compares samples that get smoothed with a Gaussian kde to evaluating
+    the conditional log-probability with `eval_conditional_density`.
 
-    This compares samples that get smoothed with a Gaussian kde to evaluating the
-    conditional log-probability with `eval_conditional_density`.
+    `eval_conditional_density` is itself tested in `sbiutils_test.py`. Here, we
+    use a bimodal posterior to test the conditional.
 
-    `eval_conditional_density` is itself tested in `sbiutils_test.py`. Here, we use
-    a bimodal posterior to test the conditional.
+    NOTE: The comparison between conditional log_probs obtained from the MCMC
+    posterior and from analysis.eval_conditional_density can be gamed by
+    underfitting the posterior estimator, i.e., by using a small number of
+    simulations.
     """
 
     num_dim = 3
     dim_to_sample_1 = 0
     dim_to_sample_2 = 2
-    num_simulations = 5000
+    num_simulations = 5500
     num_conditional_samples = 1000
     num_conditions = 50
 
     x_o = zeros(1, num_dim)
 
     likelihood_shift = -1.0 * ones(num_dim)
-    likelihood_cov = 0.1 * eye(num_dim)
+    likelihood_cov = 0.05 * eye(num_dim)
 
     prior = utils.BoxUniform(-2.0 * ones(num_dim), 2.0 * ones(num_dim))
 
+    # TODO: janfb does not see how this setup results in a bi-model posterior.
     def simulator(theta):
         batch_size, _ = theta.shape
         # create -1 1 mask for bimodality
         mask = torch.ones(batch_size, 1)
         # set mask to -1 randomly across the batch
         mask = mask * 2 * (torch.rand(batch_size, 1) > 0.5) - 1
-        if torch.rand(1) > 0.5:
-            return linear_gaussian(theta, likelihood_shift, likelihood_cov)
-        else:
-            return linear_gaussian(theta, -likelihood_shift, likelihood_cov)
+
+        # Sample bi-modally by applying a 1-(-1) mask to the likelihood shift.
+        return linear_gaussian(theta, mask * likelihood_shift, likelihood_cov)
 
     # Test whether SNPE works properly with structured z-scoring.
     net = posterior_nn("maf", z_score_x="structured", hidden_features=20)
@@ -572,6 +578,7 @@ def simulator(theta):
 
     limits = [[-2, 2], [-2, 2], [-2, 2]]
 
+    # Fit a Gaussian KDE to the conditional samples and get log-probs.
     density = gaussian_kde(cond_samples.numpy().T, bw_method="scott")
 
     X, Y = np.meshgrid(
@@ -581,7 +588,7 @@ def simulator(theta):
     positions = np.vstack([X.ravel(), Y.ravel()])
     sample_kde_grid = np.reshape(density(positions).T, X.shape)
 
-    # Evaluate the conditional with eval_conditional_density.
+    # Get conditional log probs eval_conditional_density.
     eval_grid = analysis.eval_conditional_density(
         posterior,
         condition=samples[0],
@@ -598,6 +605,7 @@ def simulator(theta):
 
     max_err = np.max(error)
     assert max_err < 0.0027
+    print(f"Max error: {max_err}")
 
 
 def test_mdn_conditional_density(num_dim: int = 3, cond_dim: int = 1):

tests/linearGaussian_snre_test.py

@@ -154,7 +154,7 @@ def test_c2st_snre_variants_on_linearGaussian_with_multiple_trials(
     """
 
     num_dim = 2
-    num_simulations = 1750
+    num_simulations = 2000
     num_samples = 500
     x_o = zeros(num_trials, num_dim)
 

tests/mnle_test.py

@@ -271,7 +271,7 @@ def sim_wrapper(theta):
     potential_fn = MixedLikelihoodBasedPotential(estimator, proposal, x_o)
 
     conditioned_potential_fn = ConditionedPotential(
-        potential_fn, condition=theta_o, dims_to_sample=[0, 1], allow_iid_x=True
+        potential_fn, condition=theta_o, dims_to_sample=[0, 1]
     )
 
     # True posterior samples

tests/multiprocessing_test.py

@@ -26,11 +26,11 @@ def slow_linear_gaussian(theta):
 
 
 @pytest.mark.slow
-@pytest.mark.parametrize("num_workers", [4])
+@pytest.mark.parametrize("num_workers", [2])  # GitHub Actions has 2 cores
 @pytest.mark.parametrize("sim_batch_size", ((1, 10, 100)))
 def test_benchmarking_parallel_simulation(sim_batch_size, num_workers):
     """Test whether joblib is faster than serial processing."""
-    num_simulations = 100
+    num_simulations = 1000
     theta = torch.zeros(num_simulations, 2)
     show_pbar = True
 

tests/sbc_test.py

@@ -92,8 +92,7 @@ def test_consistent_sbc_results(method, model="mdn"):
     num_dim = 2
     prior = BoxUniform(-torch.ones(num_dim), torch.ones(num_dim))
 
-    num_simulations = 1000
-    max_num_epochs = 20
+    num_simulations = 2000
     num_sbc_runs = 100
 
     likelihood_shift = -1.0 * ones(num_dim)
@@ -107,9 +106,7 @@ def simulator(theta):
     theta = prior.sample((num_simulations,))
     x = simulator(theta)
 
-    _ = inferer.append_simulations(theta, x).train(
-        training_batch_size=100, max_num_epochs=max_num_epochs
-    )
+    _ = inferer.append_simulations(theta, x).train(training_batch_size=100)
 
     posterior = inferer.build_posterior()
     num_posterior_samples = 1000