def test_multi_round_snpe_on_linearGaussian_based_on_mmd(algorithm_str: str):
"""Test whether APT infers well a simple example where ground truth is available."""
num_dim = 3
true_observation = torch.zeros((1, num_dim))
num_samples = 100
prior = distributions.MultivariateNormal(
loc=torch.zeros(num_dim), covariance_matrix=torch.eye(num_dim))
target_samples = get_true_posterior_samples_linear_gaussian_mvn_prior(
true_observation, num_samples=num_samples)
neural_net = utils.posterior_nn(
model="maf",
prior=prior,
context=true_observation,
)
snpe_common_args = dict(
simulator=linear_gaussian,
true_observation=true_observation,
density_estimator=neural_net,
prior=prior,
z_score_obs=True,
use_combined_loss=False,
retrain_from_scratch_each_round=False,
discard_prior_samples=False,
)
if algorithm_str == "snpe_b":
infer = SnpeB(simulation_batch_size=10, **snpe_common_args)
elif algorithm_str == "snpe_c":
infer = SnpeC(
num_atoms=10,
simulation_batch_size=50,
sample_with_mcmc=False,
**snpe_common_args,
)
# Run inference.
num_rounds, num_simulations_per_round = 2, 1000
posterior = infer(num_rounds=num_rounds,
num_simulations_per_round=num_simulations_per_round)
# Draw from posterior.
samples = posterior.sample(num_samples)
# Compute the mmd, and check if larger than expected.
mmd = utils.unbiased_mmd_squared(target_samples, samples)
max_mmd = 0.02
print(f"mmd for {algorithm_str} is {mmd}.")
assert (mmd < max_mmd
), f"MMD={mmd} is more than 2 stds above the average performance."
samples = posterior.sample(num_samples)
def test_nonlinearGaussian_based_on_mmd():
simulator = non_linear_gaussian
# Ground truth parameters as specified in 'Sequential Neural Likelihood' paper.
ground_truth_parameters = torch.tensor([-0.7, -2.9, -1.0, -0.9, 0.6])
# ground truth observation using same seed as 'Sequential Neural Likelihood' paper.
ground_truth_observation = torch.tensor([
-0.97071232,
-2.94612244,
-0.44947218,
-3.42318484,
-0.13285634,
-3.36401699,
-0.85367595,
-2.42716377,
])
# assume batch dims
parameter_dim = ground_truth_parameters.shape[0]
observation_dim = ground_truth_observation.shape[0]
prior = BoxUniform(
low=-3 * torch.ones(parameter_dim),
high=3 * torch.ones(parameter_dim),
)
infer = SnpeC(
simulator=simulator,
true_observation=ground_truth_observation[None, ],
prior=prior,
num_atoms=-1,
z_score_obs=True,
use_combined_loss=False,
retrain_from_scratch_each_round=False,
discard_prior_samples=False,
)
num_rounds, num_simulations_per_round = 2, 1000
posterior = infer(num_rounds=num_rounds,
num_simulations_per_round=num_simulations_per_round)
samples = posterior.sample(1000)
# Sample from (analytically tractable) target distribution.
target_samples = get_ground_truth_posterior_samples_nonlinear_gaussian(
num_samples=1000)
# Compute and check if MMD is larger than expected.
mmd = utils.unbiased_mmd_squared(target_samples.float(), samples.float())
max_mmd = 0.16 # mean mmd plus 2 stds.
assert mmd < max_mmd, f"MMD={mmd} larger than mean plus 2 stds."
def test_snl_on_linearGaussian_based_on_mmd(num_dim: int, prior_str: str):
"""Test SNL on linear Gaussian, comparing to ground truth posterior via MMD.
NOTE: The MMD threshold is calculated based on a number of test runs and taking the mean plus 2 stds.
Args:
num_dim: parameter dimension of the gaussian model
prior_str: one of "gaussian" or "uniform"
"""
true_observation = torch.zeros((1, num_dim))
num_samples = 200
if prior_str == "gaussian":
prior = distributions.MultivariateNormal(
loc=torch.zeros(num_dim), covariance_matrix=torch.eye(num_dim))
target_samples = get_true_posterior_samples_linear_gaussian_mvn_prior(
true_observation, num_samples=num_samples)
else:
prior = utils.BoxUniform(-1.0 * torch.ones(num_dim),
torch.ones(num_dim))
target_samples = get_true_posterior_samples_linear_gaussian_uniform_prior(
true_observation, num_samples=num_samples, prior=prior)
neural_likelihood = utils.likelihood_nn(
model="maf",
prior=prior,
context=true_observation,
)
infer = SNL(
simulator=linear_gaussian,
prior=prior,
true_observation=true_observation,
density_estimator=neural_likelihood,
mcmc_method="slice-np",
)
posterior = infer(num_rounds=1, num_simulations_per_round=1000)
samples = posterior.sample(num_samples=num_samples)
# compute the mmd
mmd = utils.unbiased_mmd_squared(target_samples, samples)
# check if mmd is larger than expected
# NOTE: the mmd is calculated based on a number of test runs
max_mmd = 0.02
assert (mmd < max_mmd
), f"MMD={mmd} is more than 2 stds above the average performance."
def test_multi_round_snl_on_linearGaussian_based_on_mmd():
"""Test SNL on linear Gaussian, comparing to ground truth posterior via MMD.
NOTE: The MMD threshold is calculated based on a number of test runs and taking the mean plus 2 stds.
"""
num_dim = 3
true_observation = torch.zeros((1, num_dim))
num_samples = 200
prior = distributions.MultivariateNormal(
loc=torch.zeros(num_dim), covariance_matrix=torch.eye(num_dim))
target_samples = get_true_posterior_samples_linear_gaussian_mvn_prior(
true_observation, num_samples=num_samples)
neural_likelihood = utils.likelihood_nn(
model="maf",
prior=prior,
context=true_observation,
)
infer = SNL(
simulator=linear_gaussian,
prior=prior,
true_observation=true_observation,
density_estimator=neural_likelihood,
simulation_batch_size=50,
mcmc_method="slice",
)
posterior = infer(num_rounds=2, num_simulations_per_round=1000)
samples = posterior.sample(num_samples=500)
mmd = utils.unbiased_mmd_squared(target_samples, samples)
# check if mmd is larger than expected
# NOTE: the mmd is calculated based on a number of test runs
max_mmd = 0.02
assert (mmd < max_mmd
), f"MMD={mmd} is more than 2 stds above the average performance."
def test_apt_on_linearGaussian_based_on_mmd(num_dim: int, prior_str: str,
algorithm_str: str,
simulation_batch_size: int):
"""Test whether APT infers well a simple example where ground truth is available."""
true_observation = torch.zeros(num_dim)
num_samples = 100
if prior_str == "gaussian":
prior = distributions.MultivariateNormal(
loc=torch.zeros(num_dim), covariance_matrix=torch.eye(num_dim))
target_samples = get_true_posterior_samples_linear_gaussian_mvn_prior(
true_observation, num_samples=num_samples)
else:
prior = utils.BoxUniform(-1.0 * torch.ones(num_dim),
torch.ones(num_dim))
target_samples = get_true_posterior_samples_linear_gaussian_uniform_prior(
true_observation, num_samples=num_samples, prior=prior)
neural_net = utils.posterior_nn(model="maf",
prior=prior,
context=true_observation)
snpe_common_args = dict(
simulator=linear_gaussian,
true_observation=true_observation,
density_estimator=neural_net,
prior=prior,
z_score_obs=True,
simulation_batch_size=simulation_batch_size,
use_combined_loss=False,
retrain_from_scratch_each_round=False,
discard_prior_samples=False,
)
if algorithm_str == "snpe_b":
infer = SnpeB(**snpe_common_args)
elif algorithm_str == "snpe_c":
infer = SnpeC(num_atoms=-1, sample_with_mcmc=False, **snpe_common_args)
# Run inference.
num_rounds, num_simulations_per_round = 1, 1000
posterior = infer(
num_rounds=num_rounds,
num_simulations_per_round=num_simulations_per_round,
)
# Draw from posterior.
samples = posterior.sample(num_samples)
# Compute the mmd, and check if larger than expected
mmd = utils.unbiased_mmd_squared(target_samples, samples)
max_mmd = 0.03
print(f"mmd for {algorithm_str} is {mmd}.")
assert (mmd < max_mmd
), f"MMD={mmd} is more than 2 stds above the average performance."
# Checks for log_prob()
if prior_str == "gaussian":
# For the Gaussian prior, we compute the KLd between ground truth and posterior.
dkl = test_utils.get_dkl_gaussian_prior(posterior, true_observation,
num_dim)
max_dkl = 0.05 if num_dim == 1 else 0.8
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 = test_utils.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,
) = test_utils.get_normalization_uniform_prior(posterior, prior,
true_observation)
# 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 summarize(
summary_writer,
summary,
round_,
true_observation,
parameter_bank,
observation_bank,
simulator,
posterior_samples_acceptance_rate=None,
):
# get ground truth if available
try:
(
_,
prior,
ground_truth_parameters,
ground_truth_observation,
) = simulators.get_simulator_prior_and_groundtruth(simulator.name)
# Update summaries.
except:
pass
try:
mmd = utils.unbiased_mmd_squared(
parameter_bank[-1],
simulator.get_ground_truth_posterior_samples(num_samples=1000),
)
summary["mmds"].append(mmd.item())
except:
pass
try:
# Median |x - x0| for most recent round.
median_observation_distance = torch.median(
torch.sqrt(
torch.sum(
(observation_bank[-1] -
true_observation.reshape(1, -1))**2,
dim=-1,
)))
summary["median_observation_distances"].append(
median_observation_distance.item())
summary_writer.add_scalar(
tag="median_observation_distance",
scalar_value=summary["median_observation_distances"][-1],
global_step=round_ + 1,
)
except:
pass
try:
# KDE estimate of negative log prob true parameters using
# parameters from most recent round.
negative_log_prob_true_parameters = -utils.gaussian_kde_log_eval(
samples=parameter_bank[-1],
query=ground_truth_parameters.reshape(1, -1),
)
summary["negative_log_probs_true_parameters"].append(
negative_log_prob_true_parameters.item())
summary_writer.add_scalar(
tag="negative_log_prob_true_parameters",
scalar_value=summary["negative_log_probs_true_parameters"][-1],
global_step=round_ + 1,
)
except:
pass
try:
# Rejection sampling acceptance rate
summary["rejection_sampling_acceptance-rates"].append(
posterior_samples_acceptance_rate)
summary_writer.add_scalar(
tag="rejection_sampling_acceptance_rate",
scalar_value=summary["rejection_sampling_acceptance_rates"][-1],
global_step=round_ + 1,
)
except:
pass
try:
# Plot most recently sampled parameters.
parameters = utils.tensor2numpy(parameter_bank[-1])
figure = utils.plot_hist_marginals(
data=parameters,
ground_truth=utils.tensor2numpy(ground_truth_parameters).reshape(
-1),
lims=simulator.parameter_plotting_limits,
)
summary_writer.add_figure(tag="posterior_samples",
figure=figure,
global_step=round_ + 1)
except:
pass
# Write quantities using SummaryWriter.
summary_writer.add_scalar(
tag="epochs_trained",
scalar_value=summary["epochs"][-1],
global_step=round_ + 1,
)
summary_writer.add_scalar(
tag="best_validation_log_prob",
scalar_value=summary["best_validation_log_probs"][-1],
global_step=round_ + 1,
)
if summary["mmds"]:
summary_writer.add_scalar(
tag="mmd",
scalar_value=summary["mmds"][-1],
global_step=round_ + 1,
)
summary_writer.flush()
return summary_writer, summary
def test_sre_on_linearGaussian_based_on_mmd(num_dim: int, prior_str: str,
classifier_loss: str):
"""Test MMD accuracy of inference with SRE on linear Gaussian model.
NOTE: The mmd threshold is calculated based on a number of test runs and taking the mean plus 2 stds.
Args:
num_dim: parameter dimension of the gaussian model
prior_str: one of "gaussian" or "uniform"
"""
true_observation = torch.zeros(num_dim)
num_samples = 300
if prior_str == "gaussian":
prior = distributions.MultivariateNormal(
loc=torch.zeros(num_dim), covariance_matrix=torch.eye(num_dim))
target_samples = get_true_posterior_samples_linear_gaussian_mvn_prior(
true_observation[None, ], num_samples=num_samples)
else:
prior = utils.BoxUniform(-1.0 * torch.ones(num_dim),
torch.ones(num_dim))
target_samples = get_true_posterior_samples_linear_gaussian_uniform_prior(
true_observation[None, ], num_samples=num_samples, prior=prior)
classifier = utils.classifier_nn(
"resnet",
prior=prior,
context=true_observation,
)
num_atoms = 2 if classifier_loss == "aalr" else -1
infer = SRE(
simulator=linear_gaussian,
prior=prior,
true_observation=true_observation,
num_atoms=num_atoms,
classifier=classifier,
classifier_loss=classifier_loss,
simulation_batch_size=50,
mcmc_method="slice-np",
)
posterior = infer(num_rounds=1, num_simulations_per_round=1000)
samples = posterior.sample(num_samples=num_samples)
# Check if mmd is larger than expected.
mmd = utils.unbiased_mmd_squared(target_samples, samples)
max_mmd = 0.045
assert (mmd < max_mmd
), f"MMD={mmd} is more than 2 stds above the average performance."
# Checks for log_prob()
if prior_str == "gaussian" and classifier_loss == "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. 2019 ('aalr') since Durkan et al. 2019 version only allows
# evaluation up to a constant.
# For the Gaussian prior, we compute the KLd between ground truth and posterior
dkl = test_utils.get_dkl_gaussian_prior(posterior, true_observation,
num_dim)
max_dkl = 0.05 if num_dim == 1 else 0.8
assert (
dkl < max_dkl
), f"KLd={dkl} is more than 2 stds above the average performance."
if prior_str == "uniform":
# Check whether the returned probability outside of the support is zero.
posterior_prob = test_utils.get_prob_outside_uniform_prior(
posterior, num_dim)
assert (
posterior_prob == 0.0
), "The posterior probability outside of the prior support is not zero"