def test_api_sre_on_linearGaussian(num_dim: int):
"""Test inference API of SRE with linear Gaussian model.
Avoids intense computation for fast testing of API etc.
Args:
num_dim: parameter dimension of the Gaussian model
"""
prior = MultivariateNormal(loc=zeros(num_dim), covariance_matrix=eye(num_dim))
simulator, prior = prepare_for_sbi(diagonal_linear_gaussian, prior)
inference = SNRE_B(
classifier="resnet",
show_progress_bars=False,
)
theta, x = simulate_for_sbi(simulator, prior, 1000, simulation_batch_size=50)
ratio_estimator = inference.append_simulations(theta, x).train(max_num_epochs=5)
for num_trials in [1, 2]:
x_o = zeros(num_trials, num_dim)
potential_fn, theta_transform = ratio_estimator_based_potential(
ratio_estimator=ratio_estimator, prior=prior, x_o=x_o
)
posterior = MCMCPosterior(
potential_fn=potential_fn,
theta_transform=theta_transform,
proposal=prior,
num_chains=2,
)
posterior.sample(sample_shape=(10,))
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_api_snpe_c_posterior_correction(sample_with, mcmc_method, prior_str):
"""Test that leakage correction applied to sampling works, with both MCMC and
rejection.
"""
num_dim = 2
x_o = zeros(1, num_dim)
# 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)
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 = SNPE_C(prior, show_progress_bars=False)
theta, x = simulate_for_sbi(simulator, prior, 1000)
posterior_estimator = inference.append_simulations(theta, x).train()
potential_fn, theta_transform = posterior_estimator_based_potential(
posterior_estimator, prior, x_o)
if sample_with == "mcmc":
posterior = MCMCPosterior(
potential_fn=potential_fn,
theta_transform=theta_transform,
proposal=prior,
method=mcmc_method,
)
elif sample_with == "rejection":
posterior = RejectionPosterior(
potential_fn=potential_fn,
proposal=prior,
theta_transform=theta_transform,
)
# Posterior should be corrected for leakage even if num_rounds just 1.
samples = posterior.sample((10, ))
# Evaluate the samples to check correction factor.
_ = posterior.log_prob(samples)
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
prior_mean = zeros(num_dim)
prior_cov = eye(num_dim)
prior = MultivariateNormal(loc=prior_mean, covariance_matrix=prior_cov)
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)
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)
posterior = MCMCPosterior(
proposal=prior,
potential_fn=potential_fn,
theta_transform=theta_transform,
thin=3,
)
posterior.sample(sample_shape=(num_samples, ))
def test_mnle_api(sampler):
# Generate mixed data.
num_simulations = 100
theta = torch.rand(num_simulations, 2)
x = torch.cat(
(torch.rand(num_simulations,
1), torch.randint(0, 2, (num_simulations, 1))),
dim=1,
)
# Train and infer.
prior = BoxUniform(torch.zeros(2), torch.ones(2))
x_o = x[0]
# Build estimator manually.
density_estimator = likelihood_nn(model="mnle", **dict(tail_bound=2.0))
trainer = MNLE(density_estimator=density_estimator)
mnle = trainer.append_simulations(theta, x).train(max_num_epochs=1)
# Test different samplers.
posterior = trainer.build_posterior(prior=prior, sample_with=sampler)
posterior.set_default_x(x_o)
if sampler == "vi":
posterior.train()
posterior.sample((1, ), show_progress_bars=False)
# MNLE should work with the default potential as well.
potential_fn, parameter_transform = likelihood_estimator_based_potential(
mnle, prior, x_o)
posterior = MCMCPosterior(
potential_fn,
proposal=prior,
theta_transform=parameter_transform,
init_strategy="proposal",
)
posterior.sample((1, ), show_progress_bars=False)
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_inference_with_2d_x(embedding, method):
num_dim = 2
num_samples = 10
num_simulations = 100
prior = utils.BoxUniform(zeros(num_dim), torch.ones(num_dim))
simulator, prior = prepare_for_sbi(simulator_2d, prior)
theta_o = torch.ones(1, num_dim)
if method == SNPE:
net_provider = utils.posterior_nn(
model="mdn",
embedding_net=embedding(),
)
num_trials = 1
elif method == SNLE:
net_provider = utils.likelihood_nn(model="mdn",
embedding_net=embedding())
num_trials = 2
else:
net_provider = utils.classifier_nn(
model="mlp",
z_score_theta="structured", # Test that structured z-scoring works.
embedding_net_x=embedding(),
)
num_trials = 2
if method == SNRE:
inference = method(classifier=net_provider, show_progress_bars=False)
else:
inference = method(density_estimator=net_provider,
show_progress_bars=False)
theta, x = simulate_for_sbi(simulator, prior, num_simulations)
estimator = inference.append_simulations(theta,
x).train(training_batch_size=100,
max_num_epochs=10)
x_o = simulator(theta_o.repeat(num_trials, 1))
if method == SNLE:
potential_fn, theta_transform = likelihood_estimator_based_potential(
estimator, prior, x_o)
elif method == SNPE:
potential_fn, theta_transform = posterior_estimator_based_potential(
estimator, prior, x_o)
elif method == SNRE:
potential_fn, theta_transform = ratio_estimator_based_potential(
estimator, prior, x_o)
else:
raise NotImplementedError
posterior = MCMCPosterior(
potential_fn=potential_fn,
theta_transform=theta_transform,
proposal=prior,
method="slice_np_vectorized",
num_chains=2,
)
posterior.potential(
posterior.sample((num_samples, ), show_progress_bars=False))
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_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_api_snl_sampling_methods(sampling_method: str, prior_str: str,
init_strategy: str):
"""Runs SNL on linear Gaussian and tests sampling from posterior via mcmc.
Args:
mcmc_method: which mcmc method to use for sampling
prior_str: use gaussian or uniform prior
"""
num_dim = 2
num_samples = 10
num_trials = 2
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
if sampling_method == "rejection":
sample_with = "rejection"
elif ("slice" in sampling_method or "nuts" in sampling_method
or "hmc" in sampling_method):
sample_with = "mcmc"
else:
sample_with = "vi"
if prior_str == "gaussian":
prior = MultivariateNormal(loc=zeros(num_dim),
covariance_matrix=eye(num_dim))
else:
prior = utils.BoxUniform(-1.0 * ones(num_dim), ones(num_dim))
simulator, prior = prepare_for_sbi(diagonal_linear_gaussian, prior)
inference = SNLE(show_progress_bars=False)
theta, x = simulate_for_sbi(simulator,
prior,
num_simulations,
simulation_batch_size=1000)
likelihood_estimator = inference.append_simulations(
theta, x).train(max_num_epochs=5)
potential_fn, theta_transform = likelihood_estimator_based_potential(
prior=prior, likelihood_estimator=likelihood_estimator, x_o=x_o)
if sample_with == "rejection":
posterior = RejectionPosterior(potential_fn=potential_fn,
proposal=prior)
elif ("slice" in sampling_method or "nuts" in sampling_method
or "hmc" in sampling_method):
posterior = MCMCPosterior(
potential_fn,
proposal=prior,
theta_transform=theta_transform,
method=sampling_method,
thin=3,
num_chains=num_chains,
init_strategy=init_strategy,
)
else:
posterior = VIPosterior(potential_fn,
theta_transform=theta_transform,
vi_method=sampling_method)
posterior.train(max_num_iters=10)
posterior.sample(sample_shape=(num_samples, ))
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_training_and_mcmc_on_device(
method,
model,
data_device,
mcmc_method,
training_device,
prior_device,
prior_type="gaussian",
):
"""Test training on devices.
This test does not check training speeds.
"""
num_dim = 2
num_samples = 10
num_simulations = 100
max_num_epochs = 5
x_o = zeros(1, num_dim).to(data_device)
likelihood_shift = -1.0 * ones(num_dim).to(prior_device)
likelihood_cov = 0.3 * eye(num_dim).to(prior_device)
if prior_type == "gaussian":
prior_mean = zeros(num_dim).to(prior_device)
prior_cov = eye(num_dim).to(prior_device)
prior = MultivariateNormal(loc=prior_mean, covariance_matrix=prior_cov)
else:
prior = BoxUniform(
low=-2 * torch.ones(num_dim),
high=2 * torch.ones(num_dim),
device=prior_device,
)
def simulator(theta):
return linear_gaussian(theta, likelihood_shift, likelihood_cov)
training_device = process_device(training_device)
if method in [SNPE_A, SNPE_C]:
kwargs = dict(
density_estimator=utils.posterior_nn(model=model, num_transforms=2)
)
elif method == SNLE:
kwargs = dict(
density_estimator=utils.likelihood_nn(model=model, num_transforms=2)
)
elif method in (SNRE_A, SNRE_B):
kwargs = dict(classifier=utils.classifier_nn(model=model))
else:
raise ValueError()
inferer = method(show_progress_bars=False, device=training_device, **kwargs)
proposals = [prior]
# Test for two rounds.
for _ in range(2):
theta, x = simulate_for_sbi(simulator, proposals[-1], num_simulations)
theta, x = theta.to(data_device), x.to(data_device)
estimator = inferer.append_simulations(theta, x).train(
training_batch_size=100, max_num_epochs=max_num_epochs
)
if method == SNLE:
potential_fn, theta_transform = likelihood_estimator_based_potential(
estimator, prior, x_o
)
elif method == SNPE_A or method == SNPE_C:
potential_fn, theta_transform = posterior_estimator_based_potential(
estimator, prior, x_o
)
elif method == SNRE_A or method == SNRE_B:
potential_fn, theta_transform = ratio_estimator_based_potential(
estimator, prior, x_o
)
else:
raise ValueError
if mcmc_method == "rejection":
posterior = RejectionPosterior(
proposal=prior,
potential_fn=potential_fn,
device=training_device,
)
elif mcmc_method == "direct":
posterior = DirectPosterior(
posterior_estimator=estimator, prior=prior
).set_default_x(x_o)
else:
posterior = MCMCPosterior(
potential_fn=potential_fn,
theta_transform=theta_transform,
proposal=prior,
method=mcmc_method,
device=training_device,
)
proposals.append(posterior)
# Check for default device for inference object
weights_device = next(inferer._neural_net.parameters()).device
assert torch.device(training_device) == weights_device
samples = proposals[-1].sample(sample_shape=(num_samples,))
proposals[-1].potential(samples)
def test_sample_conditional():
"""
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`.
`eval_conditional_density` is itself tested in `sbiutils_test.py`. Here, we use
a bimodal posterior to test the conditional.
"""
num_dim = 3
dim_to_sample_1 = 0
dim_to_sample_2 = 2
x_o = zeros(1, num_dim)
likelihood_shift = -1.0 * ones(num_dim)
likelihood_cov = 0.1 * eye(num_dim)
prior = utils.BoxUniform(-2.0 * ones(num_dim), 2.0 * ones(num_dim))
def simulator(theta):
if torch.rand(1) > 0.5:
return linear_gaussian(theta, likelihood_shift, likelihood_cov)
else:
return linear_gaussian(theta, -likelihood_shift, likelihood_cov)
# Test whether SNPE works properly with structured z-scoring.
net = utils.posterior_nn("maf", z_score_x="structured", hidden_features=20)
simulator, prior = prepare_for_sbi(simulator, prior)
inference = SNPE_C(prior, density_estimator=net, show_progress_bars=False)
# We need a pretty big dataset to properly model the bimodality.
theta, x = simulate_for_sbi(simulator, prior, 10000)
posterior_estimator = inference.append_simulations(
theta, x).train(max_num_epochs=60)
posterior = DirectPosterior(
prior=prior,
posterior_estimator=posterior_estimator).set_default_x(x_o)
samples = posterior.sample((50, ))
# Evaluate the conditional density be drawing samples and smoothing with a Gaussian
# kde.
potential_fn, theta_transform = posterior_estimator_based_potential(
posterior_estimator, prior=prior, x_o=x_o)
(
conditioned_potential_fn,
restricted_tf,
restricted_prior,
) = conditonal_potential(
potential_fn=potential_fn,
theta_transform=theta_transform,
prior=prior,
condition=samples[0],
dims_to_sample=[dim_to_sample_1, dim_to_sample_2],
)
mcmc_posterior = MCMCPosterior(
potential_fn=conditioned_potential_fn,
theta_transform=restricted_tf,
proposal=restricted_prior,
)
cond_samples = mcmc_posterior.sample((500, ))
_ = analysis.pairplot(
cond_samples,
limits=[[-2, 2], [-2, 2], [-2, 2]],
figsize=(2, 2),
diag="kde",
upper="kde",
)
limits = [[-2, 2], [-2, 2], [-2, 2]]
density = gaussian_kde(cond_samples.numpy().T, bw_method="scott")
X, Y = np.meshgrid(
np.linspace(limits[0][0], limits[0][1], 50),
np.linspace(limits[1][0], limits[1][1], 50),
)
positions = np.vstack([X.ravel(), Y.ravel()])
sample_kde_grid = np.reshape(density(positions).T, X.shape)
# Evaluate the conditional with eval_conditional_density.
eval_grid = analysis.eval_conditional_density(
posterior,
condition=samples[0],
dim1=dim_to_sample_1,
dim2=dim_to_sample_2,
limits=torch.tensor([[-2, 2], [-2, 2], [-2, 2]]),
)
# Compare the two densities.
sample_kde_grid = sample_kde_grid / np.sum(sample_kde_grid)
eval_grid = eval_grid / torch.sum(eval_grid)
error = np.abs(sample_kde_grid - eval_grid.numpy())
max_err = np.max(error)
assert max_err < 0.0027
def test_c2st_sre_on_linearGaussian():
"""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.
"""
theta_dim = 3
x_dim = 2
discard_dims = theta_dim - x_dim
num_samples = 1000
num_simulations = 2100
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)
simulator, prior = prepare_for_sbi(
lambda theta: linear_gaussian(
theta, likelihood_shift, likelihood_cov, num_discarded_dims=discard_dims
),
prior,
)
inference = SNRE_B(
classifier="resnet",
show_progress_bars=False,
)
theta, x = simulate_for_sbi(
simulator, prior, num_simulations, simulation_batch_size=100
)
ratio_estimator = inference.append_simulations(theta, x).train()
num_trials = 1
x_o = zeros(num_trials, x_dim)
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,
)
potential_fn, theta_transform = ratio_estimator_based_potential(
ratio_estimator=ratio_estimator, prior=prior, x_o=x_o
)
posterior = MCMCPosterior(
potential_fn=potential_fn,
theta_transform=theta_transform,
proposal=prior,
thin=5,
num_chains=2,
)
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"snre-{num_trials}trials")
def test_api_sre_sampling_methods(sampling_method: str, prior_str: str):
"""Test leakage correction both for MCMC and rejection sampling.
Args:
mcmc_method: which mcmc method to use for sampling
prior_str: one of "gaussian" or "uniform"
"""
num_dim = 2
num_samples = 10
num_trials = 2
num_simulations = 2100
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
if sampling_method == "rejection":
sample_with = "rejection"
elif (
"slice" in sampling_method
or "nuts" in sampling_method
or "hmc" in sampling_method
):
sample_with = "mcmc"
else:
sample_with = "vi"
if prior_str == "gaussian":
prior = MultivariateNormal(loc=zeros(num_dim), covariance_matrix=eye(num_dim))
else:
prior = utils.BoxUniform(-1.0 * ones(num_dim), ones(num_dim))
simulator, prior = prepare_for_sbi(diagonal_linear_gaussian, prior)
inference = SNRE_B(classifier="resnet", show_progress_bars=False)
theta, x = simulate_for_sbi(
simulator, prior, num_simulations, simulation_batch_size=50
)
ratio_estimator = inference.append_simulations(theta, x).train(max_num_epochs=5)
potential_fn, theta_transform = ratio_estimator_based_potential(
ratio_estimator=ratio_estimator, prior=prior, x_o=x_o
)
if sample_with == "rejection":
posterior = RejectionPosterior(potential_fn=potential_fn, proposal=prior)
elif (
"slice" in sampling_method
or "nuts" in sampling_method
or "hmc" in sampling_method
):
posterior = MCMCPosterior(
potential_fn,
proposal=prior,
theta_transform=theta_transform,
method=sampling_method,
thin=3,
num_chains=num_chains,
)
else:
posterior = VIPosterior(
potential_fn=potential_fn,
theta_transform=theta_transform,
vi_method=sampling_method,
)
posterior.train(max_num_iters=10)
posterior.sample(sample_shape=(num_samples,))
def test_c2st_sre_variants_on_linearGaussian(
num_dim: int,
num_trials: int,
prior_str: str,
method_str: str,
):
"""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"
"""
x_o = zeros(num_trials, num_dim)
num_samples = 500
num_simulations = 2600 if num_trials == 1 else 40500
# `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(
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
)
ratio_estimator = inference.append_simulations(theta, x).train()
potential_fn, theta_transform = ratio_estimator_based_potential(
ratio_estimator=ratio_estimator, prior=prior, x_o=x_o
)
posterior = MCMCPosterior(
potential_fn=potential_fn,
theta_transform=theta_transform,
proposal=prior,
method="slice_np_vectorized",
thin=5,
num_chains=5,
)
samples = posterior.sample(sample_shape=(num_samples,))
# 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="proposal")
# 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