def test_c2st_sre_variants_on_linearGaussian( num_dim: int, num_trials: int, prior_str: str, method_str: str, set_seed, ): """Test c2st accuracy of inference with SRE on linear Gaussian model. Args: num_dim: parameter dimension of the gaussian model prior_str: one of "gaussian" or "uniform" set_seed: fixture for manual seeding """ x_o = zeros(num_trials, num_dim) num_samples = 500 num_simulations = 2500 if num_trials == 1 else 35000 # `likelihood_mean` will be `likelihood_shift + theta`. likelihood_shift = -1.0 * ones(num_dim) likelihood_cov = 0.8 * eye(num_dim) 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 = utils.BoxUniform(-2.0 * ones(num_dim), 2.0 * ones(num_dim)) def simulator(theta): return linear_gaussian(theta, likelihood_shift, likelihood_cov) simulator, prior = prepare_for_sbi(simulator, prior) kwargs = dict( prior=prior, classifier="resnet", show_progress_bars=False, ) inference = SNRE_B(**kwargs) if method_str == "sre" else AALR(**kwargs) # Should use default `num_atoms=10` for SRE; `num_atoms=2` for AALR theta, x = simulate_for_sbi( simulator, prior, num_simulations, simulation_batch_size=50 ) _ = inference.append_simulations(theta, x).train() posterior = inference.build_posterior().set_default_x(x_o) samples = posterior.sample( sample_shape=(num_samples,), mcmc_method="slice_np_vectorized", mcmc_parameters={"thin": 3, "num_chains": 5}, ) # Get posterior samples. if prior_str == "gaussian": gt_posterior = true_posterior_linear_gaussian_mvn_prior( x_o, likelihood_shift, likelihood_cov, prior_mean, prior_cov ) target_samples = gt_posterior.sample((num_samples,)) else: target_samples = samples_true_posterior_linear_gaussian_uniform_prior( x_o, likelihood_shift, likelihood_cov, prior=prior, num_samples=num_samples ) # Check performance based on c2st accuracy. check_c2st( samples, target_samples, alg=f"sre-{prior_str}-{method_str}-{num_trials}trials" ) map_ = posterior.map(num_init_samples=1_000, init_method="prior") # Checks for log_prob() if prior_str == "gaussian" and method_str == "aalr": # For the Gaussian prior, we compute the KLd between ground truth and # posterior. We can do this only if the classifier_loss was as described in # Hermans et al. 2020 ('aalr') since Durkan et al. 2020 version only allows # evaluation up to a constant. # For the Gaussian prior, we compute the KLd between ground truth and posterior dkl = get_dkl_gaussian_prior( posterior, x_o, likelihood_shift, likelihood_cov, prior_mean, prior_cov ) max_dkl = 0.15 assert ( dkl < max_dkl ), f"KLd={dkl} is more than 2 stds above the average performance." assert ((map_ - gt_posterior.mean) ** 2).sum() < 0.5 if prior_str == "uniform": # Check whether the returned probability outside of the support is zero. posterior_prob = get_prob_outside_uniform_prior(posterior, prior, num_dim) assert ( posterior_prob == 0.0 ), "The posterior probability outside of the prior support is not zero" assert ((map_ - ones(num_dim)) ** 2).sum() < 0.5
def test_c2st_posterior_ensemble_on_linearGaussian(inference_method): """Test whether NeuralPosteriorEnsemble infers well a simple example with available ground truth. """ num_trials = 1 num_dim = 2 x_o = zeros(num_trials, num_dim) num_samples = 1000 num_simulations = 4000 if inference_method == "SNRE_A" else 2000 # likelihood_mean will be likelihood_shift+theta 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) gt_posterior = true_posterior_linear_gaussian_mvn_prior( x_o, likelihood_shift, likelihood_cov, prior_mean, prior_cov) target_samples = gt_posterior.sample((num_samples, )) simulator, prior = prepare_for_sbi( lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov), prior) # train ensemble components ensemble_size = 2 posteriors = [ infer(simulator, prior, inference_method, num_simulations) for i in range(ensemble_size) ] # create ensemble posterior = NeuralPosteriorEnsemble(posteriors) posterior.set_default_x(x_o) # test sampling and evaluation. if inference_method == "SNLE_A" or inference_method == "SNRE_A": samples = posterior.sample((num_samples, ), num_chains=20, method="slice_np_vectorized") else: samples = posterior.sample((num_samples, )) _ = posterior.potential(samples) # Compute the c2st and assert it is near chance level of 0.5. check_c2st(samples, target_samples, alg="{} posterior ensemble".format(inference_method)) map_ = posterior.map(init_method=samples, show_progress_bars=False) assert ((map_ - gt_posterior.mean)**2).sum() < 0.5 # Checks for log_prob() # For the Gaussian prior, we compute the KLd between ground truth and posterior. # This step is skipped for NLE since the probabilities are not normalised. if "snpe" in inference_method.lower() or "snre" in inference_method.lower( ): dkl = get_dkl_gaussian_prior( posterior, x_o[0], likelihood_shift, likelihood_cov, prior_mean, prior_cov, num_samples=num_samples, ) max_dkl = 0.15 assert ( dkl < max_dkl ), f"D-KL={dkl} is more than 2 stds above the average performance." # test individual log_prob and map posterior.log_prob(samples, individually=True)