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)
