def test_log_prob_with_different_x():
num_dim = 2
prior = MultivariateNormal(loc=zeros(num_dim),
covariance_matrix=eye(num_dim))
simulator, prior = prepare_for_sbi(diagonal_linear_gaussian, prior)
inference = SNPE_C(prior)
theta, x = simulate_for_sbi(simulator, prior, 1000)
_ = inference.append_simulations(theta, x).train()
posterior = inference.build_posterior()
_ = posterior.sample((10, ), x=ones(1, num_dim))
theta = posterior.sample((10, ), ones(1, num_dim))
posterior.log_prob(theta, x=ones(num_dim))
posterior.log_prob(theta, x=ones(num_dim))
posterior.log_prob(theta, x=ones(1, num_dim))
posterior = posterior.set_default_x(ones(1, num_dim))
posterior.log_prob(theta, x=None)
posterior.sample((10, ), x=None)
# Both must fail due to batch size of x > 1.
with pytest.raises(ValueError):
posterior.log_prob(theta, x=ones(2, num_dim))
with pytest.raises(ValueError):
posterior.sample(2, x=ones(2, num_dim))
def example_posterior():
"""Return an inferred `NeuralPosterior` for interactive examination."""
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)
prior_mean = zeros(num_dim)
prior_cov = eye(num_dim)
prior = MultivariateNormal(loc=prior_mean, covariance_matrix=prior_cov)
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,
1000,
simulation_batch_size=10,
num_workers=6)
_ = inference.append_simulations(theta, x).train()
return inference.build_posterior().set_default_x(x_o)
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_user_sbi_problems(user_simulator: Callable, user_prior):
"""
Test inference with combinations of user defined simulators, priors and x_os.
"""
simulator, prior = prepare_for_sbi(user_simulator, user_prior)
inference = SNPE_C(prior, density_estimator="maf", show_progress_bars=False,)
# Run inference.
theta, x = simulate_for_sbi(simulator, prior, 100)
_ = inference.append_simulations(theta, x).train(max_num_epochs=2)
_ = inference.build_posterior()
def test_inference_with_restriction_estimator(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
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)
rejection_estimator = RestrictionEstimator(prior=prior)
proposals = [prior]
num_rounds = 2
for r in range(num_rounds):
theta, x = simulate_for_sbi(simulator, proposals[-1], 1000)
rejection_estimator.append_simulations(theta, x)
if r < num_rounds - 1:
_ = rejection_estimator.train()
proposals.append(rejection_estimator.restrict_prior())
all_theta, all_x, _ = rejection_estimator.get_simulations()
# Any method can be used in combination with the `RejectionEstimator`.
inference = SNPE_C(prior=prior)
_ = inference.append_simulations(all_theta, all_x).train()
# Build posterior.
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"{SNPE_C}")
def test_api_snpe_c_posterior_correction(
sample_with_mcmc, mcmc_method, prior_str, set_seed
):
"""Test that leakage correction applied to sampling works, with both MCMC and
rejection.
Args:
set_seed: fixture for manual seeding
"""
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))
def simulator(theta):
return linear_gaussian(theta, likelihood_shift, likelihood_cov)
simulator, prior = prepare_for_sbi(simulator, prior)
inference = SNPE_C(
prior,
density_estimator="maf",
simulation_batch_size=50,
sample_with_mcmc=sample_with_mcmc,
mcmc_method=mcmc_method,
show_progress_bars=False,
)
theta, x = simulate_for_sbi(simulator, prior, 1000)
_ = inference.append_simulations(theta, x).train(max_num_epochs=5)
posterior = inference.build_posterior()
posterior = posterior.set_sample_with_mcmc(sample_with_mcmc).set_mcmc_method(
mcmc_method
)
# Posterior should be corrected for leakage even if num_rounds just 1.
samples = posterior.sample((10,), x=x_o)
# Evaluate the samples to check correction factor.
posterior.log_prob(samples, x=x_o)
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_sample_conditional(set_seed):
"""
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)
net = utils.posterior_nn("maf", 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)
_ = inference.append_simulations(theta, x).train(max_num_epochs=50)
posterior = inference.build_posterior().set_default_x(x_o)
samples = posterior.sample((50, ))
# Evaluate the conditional density be drawing samples and smoothing with a Gaussian
# kde.
cond_samples = posterior.sample_conditional(
(500, ),
condition=samples[0],
dims_to_sample=[dim_to_sample_1, dim_to_sample_2])
_ = utils.pairplot(
cond_samples,
limits=[[-2, 2], [-2, 2], [-2, 2]],
fig_size=(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 = utils.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.0025
