def test_c2st_slice_np_on_Gaussian(num_dim: int):
"""Test MCMC on Gaussian, comparing to ground truth target via c2st.
Args:
num_dim: parameter dimension of the gaussian model
"""
warmup = 100
num_samples = 500
likelihood_shift = -1.0 * ones(num_dim)
likelihood_cov = 0.3 * eye(num_dim)
prior_mean = zeros(num_dim)
prior_cov = eye(num_dim)
x_o = zeros((1, num_dim))
target_distribution = true_posterior_linear_gaussian_mvn_prior(
x_o[0], likelihood_shift, likelihood_cov, prior_mean, prior_cov
)
target_samples = target_distribution.sample((num_samples,))
def lp_f(x):
return target_distribution.log_prob(torch.as_tensor(x, dtype=torch.float32))
sampler = SliceSampler(lp_f=lp_f, x=np.zeros((num_dim,)).astype(np.float32))
_ = sampler.gen(warmup)
samples = sampler.gen(num_samples)
samples = torch.as_tensor(samples, dtype=torch.float32)
check_c2st(samples, target_samples, alg=f"slice_np")
def mdn_inference_with_different_methods(method):
num_dim = 2
x_o = torch.tensor([[1.0, 0.0]])
num_samples = 500
# 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[0], likelihood_shift, likelihood_cov, prior_mean, prior_cov)
target_samples = gt_posterior.sample((num_samples, ))
def simulator(theta: Tensor) -> Tensor:
return linear_gaussian(theta, likelihood_shift, likelihood_cov)
infer = method(simulator, prior, density_estimator="mdn")
posterior = infer(num_simulations=1000,
training_batch_size=50).set_default_x(x_o)
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"{method}")
def test_mcabc_inference_on_linear_gaussian(
num_dim, lra=False, sass=False, sass_expansion_degree=1,
):
x_o = zeros((1, num_dim))
num_samples = 1000
# 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[0], likelihood_shift, likelihood_cov, prior_mean, prior_cov
)
target_samples = gt_posterior.sample((num_samples,))
def simulator(theta):
return linear_gaussian(theta, likelihood_shift, likelihood_cov)
infer = ABC(simulator, prior, simulation_batch_size=10000)
phat = infer(
x_o,
100000,
quantile=0.01,
lra=lra,
sass=sass,
sass_expansion_degree=sass_expansion_degree,
sass_fraction=0.33,
)
check_c2st(phat.sample((num_samples,)), target_samples, alg="MCABC")
def test_smcabc_inference_on_linear_gaussian(
num_dim, lra=False, sass=False, sass_expansion_degree=1
):
x_o = zeros((1, num_dim))
num_samples = 1000
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[0], likelihood_shift, likelihood_cov, prior_mean, prior_cov
)
target_samples = gt_posterior.sample((num_samples,))
def simulator(theta):
return linear_gaussian(theta, likelihood_shift, likelihood_cov)
infer = SMC(simulator, prior, simulation_batch_size=10000, algorithm_variant="C")
phat = infer(
x_o,
num_particles=1000,
num_initial_pop=5000,
epsilon_decay=0.5,
num_simulations=100000,
distance_based_decay=True,
lra=lra,
sass=sass,
sass_fraction=0.5,
sass_expansion_degree=sass_expansion_degree,
)
check_c2st(phat.sample((num_samples,)), target_samples, alg="SMCABC")
def test_c2st_sre_on_linearGaussian_different_dims(set_seed):
"""Test whether SRE 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 SRE not marked as slow.
Args:
set_seed: fixture for manual seeding
"""
theta_dim = 3
x_dim = 2
discard_dims = theta_dim - x_dim
x_o = ones(1, x_dim)
num_samples = 1000
likelihood_shift = -1.0 * ones(
x_dim) # likelihood_mean will be likelihood_shift+theta
likelihood_cov = 0.3 * eye(x_dim)
prior_mean = zeros(theta_dim)
prior_cov = eye(theta_dim)
prior = MultivariateNormal(loc=prior_mean, covariance_matrix=prior_cov)
target_samples = samples_true_posterior_linear_gaussian_mvn_prior_different_dims(
x_o[0],
likelihood_shift,
likelihood_cov,
prior_mean,
prior_cov,
num_discarded_dims=discard_dims,
num_samples=num_samples,
)
def simulator(theta):
return linear_gaussian(theta,
likelihood_shift,
likelihood_cov,
num_discarded_dims=discard_dims)
simulator, prior = prepare_for_sbi(simulator, prior)
inference = SRE(
prior,
classifier="resnet",
show_progress_bars=False,
)
theta, x = simulate_for_sbi(simulator,
prior,
5000,
simulation_batch_size=50)
_ = inference.append_simulations(theta, x).train()
posterior = inference.build_posterior()
samples = posterior.sample((num_samples, ),
x=x_o,
mcmc_parameters={"thin": 3})
# Compute the c2st and assert it is near chance level of 0.5.
check_c2st(samples, target_samples, alg="snpe_c")
def test_c2st_multi_round_snpe_on_linearGaussian(method_str: str, set_seed):
"""Test whether SNPE B/C infer well a simple example with available ground truth.
Args:
set_seed: fixture for manual seeding.
"""
num_dim = 2
x_o = zeros((1, num_dim))
num_samples = 1000
# 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[0], likelihood_shift, likelihood_cov, prior_mean, prior_cov)
target_samples = gt_posterior.sample((num_samples, ))
def simulator(theta):
return linear_gaussian(theta, likelihood_shift, likelihood_cov)
if method_str == "snpe_c_non_atomic":
density_estimator = utils.posterior_nn("mdn", num_components=5)
method_str = "snpe_c"
else:
density_estimator = "maf"
simulator, prior = prepare_for_sbi(simulator, prior)
creation_args = dict(
simulator=simulator,
prior=prior,
density_estimator=density_estimator,
show_progress_bars=False,
)
if method_str == "snpe_b":
infer = SNPE_B(simulation_batch_size=10, **creation_args)
posterior1 = infer(num_simulations=1000)
posterior1.set_default_x(x_o)
posterior = infer(num_simulations=1000, proposal=posterior1)
elif method_str == "snpe_c":
infer = SNPE_C(simulation_batch_size=50,
sample_with_mcmc=False,
**creation_args)
posterior = infer(num_simulations=500, num_atoms=10).set_default_x(x_o)
posterior = infer(num_simulations=1000,
num_atoms=10,
proposal=posterior).set_default_x(x_o)
samples = posterior.sample((num_samples, ))
# Compute the c2st and assert it is near chance level of 0.5.
check_c2st(samples, target_samples, alg=method_str)
def test_c2st_snl_on_linearGaussian(num_dim: int, prior_str: str, set_seed):
"""Test SNL on linear Gaussian, comparing to ground truth posterior via c2st.
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((1, num_dim))
num_samples = 500
# likelihood_mean will be likelihood_shift+theta
likelihood_shift = -1.0 * ones(num_dim)
likelihood_cov = 0.3 * 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)
gt_posterior = true_posterior_linear_gaussian_mvn_prior(
x_o[0], likelihood_shift, likelihood_cov, prior_mean, prior_cov)
target_samples = gt_posterior.sample((num_samples, ))
else:
prior = utils.BoxUniform(-2.0 * ones(num_dim), 2.0 * ones(num_dim))
target_samples = samples_true_posterior_linear_gaussian_uniform_prior(
x_o,
likelihood_shift,
likelihood_cov,
prior=prior,
num_samples=num_samples)
simulator = lambda theta: linear_gaussian(theta, likelihood_shift,
likelihood_cov)
infer = SNL(
*prepare_for_sbi(simulator, prior),
mcmc_method="slice_np",
show_progress_bars=False,
)
posterior = infer(num_rounds=1,
num_simulations_per_round=1000).set_default_x(x_o)
samples = posterior.sample(sample_shape=(num_samples, ),
mcmc_parameters={"thin": 3})
# Check performance based on c2st accuracy.
check_c2st(samples, target_samples, alg=f"snle_a-{prior_str}-prior")
# TODO: we do not have a test for SNL log_prob(). This is because the output
# TODO: density is not normalized, so KLd does not make sense.
if prior_str == "uniform":
# Check whether the returned probability outside of the support is zero.
posterior_prob = get_prob_outside_uniform_prior(posterior, num_dim)
assert (
posterior_prob == 0.0
), "The posterior probability outside of the prior support is not zero"
def test_c2st_multi_round_snl_on_linearGaussian(num_trials: int):
"""Test SNL on linear Gaussian, comparing to ground truth posterior via c2st."""
num_dim = 2
x_o = zeros((num_trials, num_dim))
num_samples = 500
num_simulations_per_round = 600 * num_trials
# 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)
inference = SNLE(show_progress_bars=False)
theta, x = simulate_for_sbi(simulator,
prior,
num_simulations_per_round,
simulation_batch_size=50)
likelihood_estimator = inference.append_simulations(theta, x).train()
potential_fn, theta_transform = likelihood_estimator_based_potential(
prior=prior, likelihood_estimator=likelihood_estimator, x_o=x_o)
posterior1 = MCMCPosterior(
proposal=prior,
potential_fn=potential_fn,
theta_transform=theta_transform,
thin=5,
num_chains=20,
)
theta, x = simulate_for_sbi(simulator,
posterior1,
num_simulations_per_round,
simulation_batch_size=50)
likelihood_estimator = inference.append_simulations(theta, x).train()
potential_fn, theta_transform = likelihood_estimator_based_potential(
prior=prior, likelihood_estimator=likelihood_estimator, x_o=x_o)
posterior = MCMCPosterior(
proposal=prior,
potential_fn=potential_fn,
theta_transform=theta_transform,
thin=5,
num_chains=20,
)
samples = posterior.sample(sample_shape=(num_samples, ))
# Check performance based on c2st accuracy.
check_c2st(samples, target_samples, alg="multi-round-snl")
def test_c2st_snl_on_linearGaussian_different_dims(set_seed):
"""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.
Args:
set_seed: fixture for manual seeding
"""
device = "cpu"
configure_default_device(device)
theta_dim = 3
x_dim = 2
discard_dims = theta_dim - x_dim
x_o = ones(1, x_dim)
num_samples = 1000
# likelihood_mean will be likelihood_shift+theta
likelihood_shift = -1.0 * ones(x_dim)
likelihood_cov = 0.3 * eye(x_dim)
prior_mean = zeros(theta_dim)
prior_cov = eye(theta_dim)
prior = MultivariateNormal(loc=prior_mean, covariance_matrix=prior_cov)
target_samples = samples_true_posterior_linear_gaussian_mvn_prior_different_dims(
x_o[0],
likelihood_shift,
likelihood_cov,
prior_mean,
prior_cov,
num_discarded_dims=discard_dims,
num_samples=num_samples,
)
simulator = lambda theta: linear_gaussian(theta,
likelihood_shift,
likelihood_cov,
num_discarded_dims=discard_dims)
infer = SNL(
*prepare_for_sbi(simulator, prior),
simulation_batch_size=50,
mcmc_method="slice_np",
show_progress_bars=False,
device=device,
)
posterior = infer(num_rounds=1,
num_simulations_per_round=5000) # type: ignore
samples = posterior.sample((num_samples, ),
x=x_o,
mcmc_parameters={"thin": 3})
# Compute the c2st and assert it is near chance level of 0.5.
check_c2st(samples, target_samples, alg="snle_a")
def test_c2st_snpe_on_linearGaussian_different_dims(set_seed):
"""Test whether SNPE B/C infer well a simple example with available ground truth.
This example has different number of parameters theta than number of x. Also
this implicitly tests simulation_batch_size=1.
Args:
set_seed: fixture for manual seeding
"""
theta_dim = 3
x_dim = 2
discard_dims = theta_dim - x_dim
x_o = zeros(1, x_dim)
num_samples = 1000
# likelihood_mean will be likelihood_shift+theta
likelihood_shift = -1.0 * ones(x_dim)
likelihood_cov = 0.3 * eye(x_dim)
prior_mean = zeros(theta_dim)
prior_cov = eye(theta_dim)
prior = MultivariateNormal(loc=prior_mean, covariance_matrix=prior_cov)
target_samples = samples_true_posterior_linear_gaussian_mvn_prior_different_dims(
x_o[0],
likelihood_shift,
likelihood_cov,
prior_mean,
prior_cov,
num_discarded_dims=discard_dims,
num_samples=num_samples,
)
def simulator(theta):
return linear_gaussian(theta,
likelihood_shift,
likelihood_cov,
num_discarded_dims=discard_dims)
simulator, prior = prepare_for_sbi(simulator, prior)
inference = SNPE_C(
prior,
density_estimator="maf",
show_progress_bars=False,
)
theta, x = simulate_for_sbi(simulator,
prior,
2000,
simulation_batch_size=1) # type: ignore
_ = inference.append_simulations(theta, x).train()
posterior = inference.build_posterior()
samples = posterior.sample((num_samples, ), x=x_o)
# Compute the c2st and assert it is near chance level of 0.5.
check_c2st(samples, target_samples, alg="snpe_c")
def test_inference_with_restriction_estimator():
# likelihood_mean will be likelihood_shift+theta
num_dim = 3
likelihood_shift = -1.0 * ones(num_dim)
likelihood_cov = 0.3 * eye(num_dim)
x_o = zeros(1, num_dim)
num_samples = 500
def linear_gaussian_nan(theta,
likelihood_shift=likelihood_shift,
likelihood_cov=likelihood_cov):
condition = theta[:, 0] < 0.0
x = linear_gaussian(theta, likelihood_shift, likelihood_cov)
x[condition] = float("nan")
return x
prior = utils.BoxUniform(-2.0 * ones(num_dim), 2.0 * ones(num_dim))
target_samples = samples_true_posterior_linear_gaussian_uniform_prior(
x_o,
likelihood_shift=likelihood_shift,
likelihood_cov=likelihood_cov,
num_samples=num_samples,
prior=prior,
)
simulator, prior = prepare_for_sbi(linear_gaussian_nan, prior)
restriction_estimator = RestrictionEstimator(prior=prior)
proposals = [prior]
num_rounds = 2
for r in range(num_rounds):
theta, x = simulate_for_sbi(simulator, proposals[-1], 1000)
restriction_estimator.append_simulations(theta, x)
if r < num_rounds - 1:
_ = restriction_estimator.train()
proposals.append(restriction_estimator.restrict_prior())
all_theta, all_x, _ = restriction_estimator.get_simulations()
# Any method can be used in combination with the `RejectionEstimator`.
inference = SNPE_C(prior=prior)
posterior_estimator = inference.append_simulations(all_theta,
all_x).train()
# Build posterior.
posterior = DirectPosterior(
prior=prior,
posterior_estimator=posterior_estimator).set_default_x(x_o)
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"{SNPE_C}")
def test_c2st_multi_round_snl_on_linearGaussian(set_seed):
"""Test SNL on linear Gaussian, comparing to ground truth posterior via c2st.
Args:
set_seed: fixture for manual seeding
"""
num_dim = 2
x_o = zeros((1, num_dim))
num_samples = 500
# 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[0], likelihood_shift, likelihood_cov, prior_mean, prior_cov)
target_samples = gt_posterior.sample((num_samples, ))
simulator = lambda theta: linear_gaussian(theta, likelihood_shift,
likelihood_cov)
simulator, prior = prepare_for_sbi(simulator, prior)
inference = SNL(
prior,
show_progress_bars=False,
)
theta, x = simulate_for_sbi(simulator,
prior,
750,
simulation_batch_size=50)
_ = inference.append_simulations(theta, x).train()
posterior1 = inference.build_posterior(mcmc_method="slice_np_vectorized",
mcmc_parameters={
"thin": 5,
"num_chains": 20
}).set_default_x(x_o)
theta, x = simulate_for_sbi(simulator,
posterior1,
750,
simulation_batch_size=50)
_ = inference.append_simulations(theta, x).train()
posterior = inference.build_posterior().copy_hyperparameters_from(
posterior1)
samples = posterior.sample(sample_shape=(num_samples, ),
mcmc_parameters={"thin": 3})
# Check performance based on c2st accuracy.
check_c2st(samples, target_samples, alg="multi-round-snl")
def test_c2st_snle_external_data_on_linearGaussian(set_seed):
"""Test whether SNPE C infers well a simple example with available ground truth.
Args:
set_seed: fixture for manual seeding
"""
num_dim = 2
device = "cpu"
configure_default_device(device)
x_o = zeros(1, num_dim)
num_samples = 1000
# 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[0], likelihood_shift, likelihood_cov, prior_mean, prior_cov)
target_samples = gt_posterior.sample((num_samples, ))
def simulator(theta):
return linear_gaussian(theta, likelihood_shift, likelihood_cov)
infer = SNL(
*prepare_for_sbi(simulator, prior),
simulation_batch_size=1000,
show_progress_bars=False,
device=device,
)
external_theta = prior.sample((1000, ))
external_x = simulator(external_theta)
infer.provide_presimulated(external_theta, external_x)
posterior = infer(
num_rounds=1,
num_simulations_per_round=1000,
training_batch_size=100,
).set_default_x(x_o)
samples = posterior.sample((num_samples, ))
# Compute the c2st and assert it is near chance level of 0.5.
check_c2st(samples, target_samples, alg="snpe_c")
def test_mcabc_inference_on_linear_gaussian(
num_dim,
lra=False,
sass=False,
sass_expansion_degree=1,
kde=False,
kde_bandwidth="cv",
):
x_o = zeros((1, num_dim))
num_samples = 1000
# 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[0], likelihood_shift, likelihood_cov, prior_mean, prior_cov)
target_samples = gt_posterior.sample((num_samples, ))
def simulator(theta):
return linear_gaussian(theta, likelihood_shift, likelihood_cov)
inferer = MCABC(simulator, prior, simulation_batch_size=10000)
phat = inferer(
x_o,
120000,
quantile=0.01,
lra=lra,
sass=sass,
sass_expansion_degree=sass_expansion_degree,
sass_fraction=0.33,
kde=kde,
kde_kwargs=dict(bandwidth=kde_bandwidth) if kde else {},
return_summary=False,
)
check_c2st(
phat.sample((num_samples, )) if kde else phat,
target_samples,
alg=f"MCABC_lra{lra}_sass{sass}_kde{kde}_{kde_bandwidth}",
)
def test_inference_with_nan_simulator(method: type, exclude_invalid_x: bool,
percent_nans: float, set_seed):
# likelihood_mean will be likelihood_shift+theta
num_dim = 3
likelihood_shift = -1.0 * ones(num_dim)
likelihood_cov = 0.3 * eye(num_dim)
x_o = zeros(1, num_dim)
num_samples = 500
num_simulations = 2000
def linear_gaussian_nan(theta,
likelihood_shift=likelihood_shift,
likelihood_cov=likelihood_cov):
x = linear_gaussian(theta, likelihood_shift, likelihood_cov)
# Set nan randomly.
x[torch.rand(x.shape) < (percent_nans * 1.0 /
x.shape[1])] = float("nan")
return x
prior = utils.BoxUniform(-2.0 * ones(num_dim), 2.0 * ones(num_dim))
target_samples = samples_true_posterior_linear_gaussian_uniform_prior(
x_o,
likelihood_shift=likelihood_shift,
likelihood_cov=likelihood_cov,
num_samples=num_samples,
prior=prior,
)
simulator, prior = prepare_for_sbi(linear_gaussian_nan, prior)
inference = method(prior)
theta, x = simulate_for_sbi(simulator, prior, num_simulations)
_ = inference.append_simulations(
theta, x).train(exclude_invalid_x=exclude_invalid_x)
posterior = inference.build_posterior().set_default_x(x_o)
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"{method}")
def test_c2st_multi_round_snl_on_linearGaussian(set_seed):
"""Test SNL on linear Gaussian, comparing to ground truth posterior via c2st.
Args:
set_seed: fixture for manual seeding
"""
num_dim = 2
x_o = zeros((1, num_dim))
num_samples = 500
# 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[0], likelihood_shift, likelihood_cov, prior_mean, prior_cov)
target_samples = gt_posterior.sample((num_samples, ))
simulator = lambda theta: linear_gaussian(theta, likelihood_shift,
likelihood_cov)
infer = SNL(
*prepare_for_sbi(simulator, prior),
simulation_batch_size=50,
mcmc_method="slice",
show_progress_bars=False,
)
posterior1 = infer(num_simulations=500).set_default_x(x_o)
posterior = infer(num_simulations=500,
proposal=posterior1).set_default_x(x_o)
samples = posterior.sample(sample_shape=(num_samples, ),
mcmc_parameters={"thin": 3})
# Check performance based on c2st accuracy.
check_c2st(samples, target_samples, alg="multi-round-snl")
def test_c2st_slice_np_vectorized_on_Gaussian(num_dim: int):
"""Test MCMC on Gaussian, comparing to ground truth target via c2st.
Args:
num_dim: parameter dimension of the gaussian model
"""
num_samples = 500
warmup = 500
num_chains = 5
likelihood_shift = -1.0 * ones(num_dim)
likelihood_cov = 0.3 * eye(num_dim)
prior_mean = zeros(num_dim)
prior_cov = eye(num_dim)
x_o = zeros((1, num_dim))
target_distribution = true_posterior_linear_gaussian_mvn_prior(
x_o[0], likelihood_shift, likelihood_cov, prior_mean, prior_cov
)
target_samples = target_distribution.sample((num_samples,))
def lp_f(x):
return target_distribution.log_prob(torch.as_tensor(x, dtype=torch.float32))
sampler = SliceSamplerVectorized(
log_prob_fn=lp_f,
init_params=np.zeros(
(
num_chains,
num_dim,
)
).astype(np.float32),
num_chains=num_chains,
)
samples = sampler.run(warmup + int(num_samples / num_chains))
samples = samples[:, warmup:, :]
samples = samples.reshape(-1, num_dim)
samples = torch.as_tensor(samples, dtype=torch.float32)
check_c2st(samples, target_samples, alg="slice_np_vectorized")
def mdn_inference_with_different_methods(method):
num_dim = 2
x_o = torch.tensor([[1.0, 0.0]])
num_samples = 500
num_simulations = 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[0], likelihood_shift, likelihood_cov, prior_mean, prior_cov)
target_samples = gt_posterior.sample((num_samples, ))
def simulator(theta: Tensor) -> Tensor:
return linear_gaussian(theta, likelihood_shift, likelihood_cov)
simulator, prior = prepare_for_sbi(simulator, prior)
inference = method(density_estimator="mdn")
theta, x = simulate_for_sbi(simulator, prior, num_simulations)
estimator = inference.append_simulations(theta, x).train()
if method == SNPE:
posterior = DirectPosterior(posterior_estimator=estimator, prior=prior)
else:
potential_fn, theta_transform = likelihood_estimator_based_potential(
likelihood_estimator=estimator, prior=prior, x_o=x_o)
posterior = MCMCPosterior(potential_fn=potential_fn,
theta_transform=theta_transform,
proposal=prior)
samples = posterior.sample((num_samples, ), x=x_o)
# Compute the c2st and assert it is near chance level of 0.5.
check_c2st(samples, target_samples, alg=f"{method}")
def test_c2st_and_map_snl_on_linearGaussian_different(num_dim: int,
prior_str: str):
"""Test SNL on linear Gaussian, comparing to ground truth posterior via c2st.
Args:
num_dim: parameter dimension of the gaussian model
prior_str: one of "gaussian" or "uniform"
"""
num_samples = 500
num_simulations = 3000
trials_to_test = [1]
# likelihood_mean will be likelihood_shift+theta
likelihood_shift = -1.0 * ones(num_dim)
# Use increased cov to avoid too small posterior cov for many trials.
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))
simulator, prior = prepare_for_sbi(
lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov),
prior)
density_estimator = likelihood_nn("maf", 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=10000)
likelihood_estimator = inference.append_simulations(theta, x).train()
# Test inference amortized over trials.
for num_trials in trials_to_test:
x_o = zeros((num_trials, num_dim))
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, ))
elif prior_str == "uniform":
target_samples = samples_true_posterior_linear_gaussian_uniform_prior(
x_o,
likelihood_shift,
likelihood_cov,
prior=prior,
num_samples=num_samples,
)
else:
raise ValueError(f"Wrong prior_str: '{prior_str}'.")
potential_fn, theta_transform = likelihood_estimator_based_potential(
prior=prior, likelihood_estimator=likelihood_estimator, x_o=x_o)
posterior = MCMCPosterior(
proposal=prior,
potential_fn=potential_fn,
theta_transform=theta_transform,
method="slice_np_vectorized",
thin=5,
num_chains=5,
)
samples = posterior.sample(sample_shape=(num_samples, ))
# Check performance based on c2st accuracy.
check_c2st(samples,
target_samples,
alg=f"snle_a-{prior_str}-prior-{num_trials}-trials")
map_ = posterior.map(num_init_samples=1_000,
init_method="proposal",
show_progress_bars=False)
# TODO: we do not have a test for SNL log_prob(). This is because the output
# TODO: density is not normalized, so KLd does not make sense.
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
else:
assert ((map_ - gt_posterior.mean)**2).sum() < 0.5
def test_c2st_snpe_on_linearGaussian(
num_dim: int,
prior_str: str,
set_seed,
):
"""Test whether SNPE C infers well a simple example with available ground truth.
Args:
set_seed: fixture for manual seeding
"""
x_o = zeros(1, num_dim)
num_samples = 1000
# likelihood_mean will be likelihood_shift+theta
likelihood_shift = -1.0 * ones(num_dim)
likelihood_cov = 0.3 * 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)
gt_posterior = true_posterior_linear_gaussian_mvn_prior(
x_o[0], likelihood_shift, likelihood_cov, prior_mean, prior_cov)
target_samples = gt_posterior.sample((num_samples, ))
else:
prior = utils.BoxUniform(-2.0 * ones(num_dim), 2.0 * ones(num_dim))
target_samples = samples_true_posterior_linear_gaussian_uniform_prior(
x_o,
likelihood_shift,
likelihood_cov,
prior=prior,
num_samples=num_samples)
def simulator(theta):
return linear_gaussian(theta, likelihood_shift, likelihood_cov)
simulator, prior = prepare_for_sbi(simulator, prior)
inference = SNPE_C(
prior,
show_progress_bars=False,
)
theta, x = simulate_for_sbi(simulator,
prior,
2000,
simulation_batch_size=1000)
_ = inference.append_simulations(theta, x).train(training_batch_size=100)
posterior = inference.build_posterior().set_default_x(x_o)
samples = posterior.sample((num_samples, ))
# Compute the c2st and assert it is near chance level of 0.5.
check_c2st(samples, target_samples, alg="snpe_c")
# Checks for log_prob()
if prior_str == "gaussian":
# For the Gaussian prior, we compute the KLd between ground truth and posterior.
dkl = get_dkl_gaussian_prior(posterior, x_o[0], likelihood_shift,
likelihood_cov, prior_mean, prior_cov)
max_dkl = 0.15
assert (
dkl < max_dkl
), f"D-KL={dkl} is more than 2 stds above the average performance."
elif prior_str == "uniform":
# Check whether the returned probability outside of the support is zero.
posterior_prob = get_prob_outside_uniform_prior(posterior, num_dim)
assert (
posterior_prob == 0.0
), "The posterior probability outside of the prior support is not zero"
# Check whether normalization (i.e. scaling up the density due
# to leakage into regions without prior support) scales up the density by the
# correct factor.
(
posterior_likelihood_unnorm,
posterior_likelihood_norm,
acceptance_prob,
) = get_normalization_uniform_prior(posterior, prior, x_o)
# The acceptance probability should be *exactly* the ratio of the unnormalized
# and the normalized likelihood. However, we allow for an error margin of 1%,
# since the estimation of the acceptance probability is random (based on
# rejection sampling).
assert (
acceptance_prob * 0.99 < posterior_likelihood_unnorm /
posterior_likelihood_norm < acceptance_prob * 1.01
), "Normalizing the posterior density using the acceptance probability failed."
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)
def test_c2st_multi_round_snpe_on_linearGaussian(method_str: str, set_seed):
"""Test whether SNPE B/C infer well a simple example with available ground truth.
Args:
set_seed: fixture for manual seeding.
"""
num_dim = 2
x_o = zeros((1, num_dim))
num_samples = 1000
# 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, ))
if method_str == "snpe_c_non_atomic":
density_estimator = utils.posterior_nn("mdn", num_components=5)
method_str = "snpe_c"
elif method_str == "snpe_a":
density_estimator = "mdn_snpe_a"
else:
density_estimator = "maf"
simulator, prior = prepare_for_sbi(
lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov),
prior)
creation_args = dict(
simulator=simulator,
prior=prior,
density_estimator=density_estimator,
show_progress_bars=False,
)
if method_str == "snpe_b":
inference = SNPE_B(**creation_args)
theta, x = simulate_for_sbi(simulator,
prior,
500,
simulation_batch_size=10)
_ = inference.append_simulations(theta, x).train()
posterior1 = inference.build_posterior().set_default_x(x_o)
theta, x = simulate_for_sbi(simulator,
posterior1,
1000,
simulation_batch_size=10)
_ = inference.append_simulations(theta, x, proposal=posterior1).train()
posterior = inference.build_posterior().set_default_x(x_o)
elif method_str == "snpe_c":
inference = SNPE_C(**creation_args)
theta, x = simulate_for_sbi(simulator,
prior,
500,
simulation_batch_size=50)
_ = inference.append_simulations(theta, x).train()
posterior1 = inference.build_posterior().set_default_x(x_o)
theta, x = simulate_for_sbi(simulator,
posterior1,
1000,
simulation_batch_size=50)
_ = inference.append_simulations(theta, x, proposal=posterior1).train()
posterior = inference.build_posterior().set_default_x(x_o)
elif method_str == "snpe_a":
inference = SNPE_A(**creation_args)
proposal = prior
final_round = False
num_rounds = 3
for r in range(num_rounds):
if r == 2:
final_round = True
theta, x = simulate_for_sbi(simulator,
proposal,
500,
simulation_batch_size=50)
inference = inference.append_simulations(theta,
x,
proposal=proposal)
_ = inference.train(max_num_epochs=200, final_round=final_round)
posterior = inference.build_posterior().set_default_x(x_o)
proposal = posterior
samples = posterior.sample((num_samples, ))
# Compute the c2st and assert it is near chance level of 0.5.
check_c2st(samples, target_samples, alg=method_str)
def test_c2st_snl_on_linearGaussian_different_dims_and_trials(
num_dim: int, prior_str: str, set_seed):
"""Test SNL on linear Gaussian, comparing to ground truth posterior via c2st.
Args:
num_dim: parameter dimension of the gaussian model
prior_str: one of "gaussian" or "uniform"
set_seed: fixture for manual seeding
"""
num_samples = 500
num_simulations = 7500
trials_to_test = [1, 5, 10]
# likelihood_mean will be likelihood_shift+theta
likelihood_shift = -1.0 * ones(num_dim)
# Use increased cov to avoid too small posterior cov for many trials.
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))
simulator, prior = prepare_for_sbi(
lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov),
prior)
inference = SNL(prior, show_progress_bars=False)
theta, x = simulate_for_sbi(simulator,
prior,
num_simulations,
simulation_batch_size=50)
_ = inference.append_simulations(theta, x).train()
# Test inference amortized over trials.
for num_trials in trials_to_test:
x_o = zeros((num_trials, num_dim))
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,
)
posterior = inference.build_posterior(
mcmc_method="slice_np_vectorized").set_default_x(x_o)
samples = posterior.sample(sample_shape=(num_samples, ),
mcmc_parameters={
"thin": 3,
"num_chains": 2
})
# Check performance based on c2st accuracy.
check_c2st(samples,
target_samples,
alg=f"snle_a-{prior_str}-prior-{num_trials}-trials")
map_ = posterior.map(num_init_samples=1_000, init_method="prior")
# TODO: we do not have a test for SNL log_prob(). This is because the output
# TODO: density is not normalized, so KLd does not make sense.
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
else:
assert ((map_ - gt_posterior.mean)**2).sum() < 0.5
def test_c2st_slice_np_parallelized(
vectorized: bool, num_workers: int, seed: Union[None, int]
):
"""Test MCMC on Gaussian, comparing to ground truth target via c2st.
Args:
num_dim: parameter dimension of the gaussian model
"""
num_dim = 2
num_samples = 500
warmup = 500
num_chains = 5
likelihood_shift = -1.0 * ones(num_dim)
likelihood_cov = 0.3 * eye(num_dim)
prior_mean = zeros(num_dim)
prior_cov = eye(num_dim)
x_o = zeros((1, num_dim))
target_distribution = true_posterior_linear_gaussian_mvn_prior(
x_o[0], likelihood_shift, likelihood_cov, prior_mean, prior_cov
)
target_samples = target_distribution.sample((num_samples,))
def lp_f(x):
return target_distribution.log_prob(torch.as_tensor(x, dtype=torch.float32))
initial_params = torch.zeros((num_chains, num_dim))
# Maybe test seeding.
if seed is not None:
torch.manual_seed(seed)
samples = slice_np_parallized(
lp_f,
initial_params,
num_samples,
thin=1,
warmup_steps=warmup,
vectorized=vectorized,
num_workers=num_workers,
show_progress_bars=False,
)
# Repeat to test seeding.
if seed is not None:
torch.manual_seed(seed)
samples_2 = slice_np_parallized(
lp_f,
initial_params,
num_samples,
thin=1,
warmup_steps=warmup,
vectorized=vectorized,
num_workers=num_workers,
)
# Test seeding.
assert torch.allclose(samples, samples_2)
samples = torch.as_tensor(samples, dtype=torch.float32).reshape(-1, num_dim)
check_c2st(
samples, target_samples, alg=f"slice_np {'vectorized' if vectorized else ''}"
)
def test_mdn_conditional_density(num_dim: int = 3, cond_dim: int = 1):
"""Test whether the conditional density infered from MDN parameters of a
`DirectPosterior` matches analytical results for MVN. This uses a n-D joint and
conditions on the last m values to generate a conditional.
Gaussian prior used for easier ground truthing of conditional posterior.
Args:
num_dim: Dimensionality of the MVM.
cond_dim: Dimensionality of the condition.
"""
assert (
num_dim > cond_dim
), "The number of dimensions needs to be greater than that of the condition!"
x_o = zeros(1, num_dim)
num_samples = 1000
num_simulations = 2700
condition = 0.1 * ones(1, num_dim)
dims = list(range(num_dim))
dims2sample = dims[-cond_dim:]
dims2condition = dims[:-cond_dim]
# 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)
joint_posterior = true_posterior_linear_gaussian_mvn_prior(
x_o[0], likelihood_shift, likelihood_cov, prior_mean, prior_cov)
joint_cov = joint_posterior.covariance_matrix
joint_mean = joint_posterior.loc
conditional_mean, conditional_cov = conditional_of_mvn(
joint_mean, joint_cov, condition[0, dims2condition])
conditional_dist_gt = MultivariateNormal(conditional_mean, conditional_cov)
conditional_samples_gt = conditional_dist_gt.sample((num_samples, ))
def simulator(theta):
return linear_gaussian(theta, likelihood_shift, likelihood_cov)
simulator, prior = prepare_for_sbi(simulator, prior)
inference = SNPE_C(density_estimator="mdn", show_progress_bars=False)
theta, x = simulate_for_sbi(simulator,
prior,
num_simulations,
simulation_batch_size=1000)
posterior_mdn = inference.append_simulations(
theta, x).train(training_batch_size=100)
conditioned_mdn = ConditionedMDN(posterior_mdn,
x_o,
condition=condition,
dims_to_sample=[0])
conditional_samples_sbi = conditioned_mdn.sample((num_samples, ))
check_c2st(
conditional_samples_sbi,
conditional_samples_gt,
alg="analytic_mdn_conditioning_of_direct_posterior",
)
def test_mnle_accuracy(sampler):
def mixed_simulator(theta):
# Extract parameters
beta, ps = theta[:, :1], theta[:, 1:]
# Sample choices and rts independently.
choices = Binomial(probs=ps).sample()
rts = InverseGamma(concentration=2 * torch.ones_like(beta),
rate=beta).sample()
return torch.cat((rts, choices), dim=1)
prior = MultipleIndependent(
[
Gamma(torch.tensor([1.0]), torch.tensor([0.5])),
Beta(torch.tensor([2.0]), torch.tensor([2.0])),
],
validate_args=False,
)
num_simulations = 2000
num_samples = 1000
theta = prior.sample((num_simulations, ))
x = mixed_simulator(theta)
# MNLE
trainer = MNLE(prior)
trainer.append_simulations(theta, x).train()
posterior = trainer.build_posterior()
mcmc_kwargs = dict(
num_chains=10,
warmup_steps=100,
method="slice_np_vectorized",
init_strategy="proposal",
)
for num_trials in [10]:
theta_o = prior.sample((1, ))
x_o = mixed_simulator(theta_o.repeat(num_trials, 1))
# True posterior samples
transform = mcmc_transform(prior)
true_posterior_samples = MCMCPosterior(
PotentialFunctionProvider(prior, atleast_2d(x_o)),
theta_transform=transform,
proposal=prior,
**mcmc_kwargs,
).sample((num_samples, ), show_progress_bars=False)
posterior = trainer.build_posterior(prior=prior, sample_with=sampler)
posterior.set_default_x(x_o)
if sampler == "vi":
posterior.train()
mnle_posterior_samples = posterior.sample(
sample_shape=(num_samples, ),
show_progress_bars=False,
**mcmc_kwargs if sampler == "mcmc" else {},
)
check_c2st(
mnle_posterior_samples,
true_posterior_samples,
alg=f"MNLE with {sampler}",
)
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_snl_on_linearGaussian():
"""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.
"""
theta_dim = 3
x_dim = 2
discard_dims = theta_dim - x_dim
x_o = zeros(1, x_dim)
num_samples = 1000
num_simulations = 3100
# likelihood_mean will be likelihood_shift+theta
likelihood_shift = -1.0 * ones(x_dim)
likelihood_cov = 0.3 * eye(x_dim)
prior_mean = zeros(theta_dim)
prior_cov = eye(theta_dim)
prior = MultivariateNormal(loc=prior_mean, covariance_matrix=prior_cov)
target_samples = samples_true_posterior_linear_gaussian_mvn_prior_different_dims(
x_o,
likelihood_shift,
likelihood_cov,
prior_mean,
prior_cov,
num_discarded_dims=discard_dims,
num_samples=num_samples,
)
simulator, prior = prepare_for_sbi(
lambda theta: linear_gaussian(theta,
likelihood_shift,
likelihood_cov,
num_discarded_dims=discard_dims),
prior,
)
density_estimator = likelihood_nn("maf", 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)
likelihood_estimator = inference.append_simulations(theta, x).train()
potential_fn, theta_transform = likelihood_estimator_based_potential(
prior=prior, likelihood_estimator=likelihood_estimator, x_o=x_o)
posterior = MCMCPosterior(
proposal=prior,
potential_fn=potential_fn,
theta_transform=theta_transform,
method="slice_np_vectorized",
num_chains=5,
thin=10,
)
samples = posterior.sample((num_samples, ))
# Compute the c2st and assert it is near chance level of 0.5.
check_c2st(samples, target_samples, alg="snle_a")
def test_smcabc_inference_on_linear_gaussian(
num_dim,
prior_type: str,
lra=False,
sass=False,
sass_expansion_degree=1,
kde=False,
kde_bandwidth="cv",
transform=False,
num_simulations=20000,
):
x_o = zeros((1, num_dim))
num_samples = 1000
likelihood_shift = -1.0 * ones(num_dim)
likelihood_cov = 0.3 * eye(num_dim)
if prior_type == "gaussian":
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[0], likelihood_shift, likelihood_cov, prior_mean, prior_cov)
target_samples = gt_posterior.sample((num_samples, ))
elif prior_type == "uniform":
prior = BoxUniform(-ones(num_dim), ones(num_dim))
target_samples = samples_true_posterior_linear_gaussian_uniform_prior(
x_o[0], likelihood_shift, likelihood_cov, prior, num_samples)
else:
raise ValueError("Wrong prior string.")
def simulator(theta):
return linear_gaussian(theta, likelihood_shift, likelihood_cov)
infer = SMC(simulator,
prior,
simulation_batch_size=10000,
algorithm_variant="C")
phat = infer(
x_o,
num_particles=1000,
num_initial_pop=5000,
epsilon_decay=0.5,
num_simulations=num_simulations,
distance_based_decay=True,
return_summary=False,
lra=lra,
sass=sass,
sass_fraction=0.5,
sass_expansion_degree=sass_expansion_degree,
kde=kde,
kde_kwargs=dict(
bandwidth=kde_bandwidth,
transform=biject_to(prior.support) if transform else None,
),
)
check_c2st(
phat.sample((num_samples, )) if kde else phat,
target_samples,
alg=
f"SMCABC-{prior_type}prior-lra{lra}-sass{sass}-kde{kde}-{kde_bandwidth}",
)
if kde:
samples = phat.sample((10, ))
phat.log_prob(samples)
def test_c2st_snpe_on_linearGaussian_different_dims():
"""Test whether SNPE B/C infer well a simple example with available ground truth.
This example has different number of parameters theta than number of x. Also
this implicitly tests simulation_batch_size=1. It also impleictly tests whether the
prior can be `None` and whether we can stop and resume training.
"""
theta_dim = 3
x_dim = 2
discard_dims = theta_dim - x_dim
x_o = zeros(1, x_dim)
num_samples = 1000
# likelihood_mean will be likelihood_shift+theta
likelihood_shift = -1.0 * ones(x_dim)
likelihood_cov = 0.3 * eye(x_dim)
prior_mean = zeros(theta_dim)
prior_cov = eye(theta_dim)
prior = MultivariateNormal(loc=prior_mean, covariance_matrix=prior_cov)
target_samples = samples_true_posterior_linear_gaussian_mvn_prior_different_dims(
x_o,
likelihood_shift,
likelihood_cov,
prior_mean,
prior_cov,
num_discarded_dims=discard_dims,
num_samples=num_samples,
)
simulator, prior = prepare_for_sbi(
lambda theta: linear_gaussian(theta,
likelihood_shift,
likelihood_cov,
num_discarded_dims=discard_dims),
prior,
)
# Test whether prior can be `None`.
inference = SNPE_C(prior=None,
density_estimator="maf",
show_progress_bars=False)
# type: ignore
theta, x = simulate_for_sbi(simulator,
prior,
2000,
simulation_batch_size=1)
inference = inference.append_simulations(theta, x)
posterior_estimator = inference.train(
max_num_epochs=10) # Test whether we can stop and resume.
posterior_estimator = inference.train(resume_training=True,
force_first_round_loss=True)
posterior = DirectPosterior(
prior=prior,
posterior_estimator=posterior_estimator).set_default_x(x_o)
samples = posterior.sample((num_samples, ))
# Compute the c2st and assert it is near chance level of 0.5.
check_c2st(samples, target_samples, alg="snpe_c")