def test_snl_on_linearGaussian_api(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 = distributions.MultivariateNormal(
loc=torch.zeros(num_dim), covariance_matrix=torch.eye(num_dim))
true_observation = torch.zeros(num_dim)
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-np",
)
posterior = infer(num_rounds=1, num_simulations_per_round=1000)
posterior.sample(num_samples=num_samples, num_chains=1)
def test_training_and_mcmc_on_device(method, model, device):
"""Test training on devices.
This test does not check training speeds.
"""
device = process_device(device)
num_dim = 2
num_samples = 10
num_simulations = 500
max_num_epochs = 5
x_o = zeros(1, num_dim)
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)
if method == SNPE:
kwargs = dict(density_estimator=utils.posterior_nn(model=model), )
mcmc_kwargs = dict(
sample_with_mcmc=True,
mcmc_method="slice_np",
)
elif method == SNLE:
kwargs = dict(density_estimator=utils.likelihood_nn(model=model), )
mcmc_kwargs = dict(mcmc_method="slice")
elif method == SNRE:
kwargs = dict(classifier=utils.classifier_nn(model=model), )
mcmc_kwargs = dict(mcmc_method="slice_np_vectorized", )
else:
raise ValueError()
inferer = method(prior, show_progress_bars=False, device=device, **kwargs)
proposals = [prior]
# Test for two rounds.
for r in range(2):
theta, x, = simulate_for_sbi(simulator,
proposal=prior,
num_simulations=num_simulations)
_ = inferer.append_simulations(theta,
x).train(training_batch_size=100,
max_num_epochs=max_num_epochs)
posterior = inferer.build_posterior(**mcmc_kwargs).set_default_x(x_o)
proposals.append(posterior)
proposals[-1].sample(sample_shape=(num_samples, ), x=x_o, **mcmc_kwargs)
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 __init__(
self,
prior: Optional[Distribution] = None,
density_estimator: Union[str, Callable] = "maf",
device: str = "cpu",
logging_level: Union[int, str] = "WARNING",
summary_writer: Optional[SummaryWriter] = None,
show_progress_bars: bool = True,
):
r"""Base class for Sequential Neural Likelihood Estimation methods.
Args:
prior: A probability distribution that expresses prior knowledge about the
parameters, e.g. which ranges are meaningful for them. Any
object with `.log_prob()`and `.sample()` (for example, a PyTorch
distribution) can be used.
density_estimator: If it is a string, use a pre-configured network of the
provided type (one of nsf, maf, mdn, made). Alternatively, a function
that builds a custom neural network can be provided. The function will
be called with the first batch of simulations (theta, x), which can
thus be used for shape inference and potentially for z-scoring. It
needs to return a PyTorch `nn.Module` implementing the density
estimator. The density estimator needs to provide the methods
`.log_prob` and `.sample()`.
See docstring of `NeuralInference` class for all other arguments.
"""
super().__init__(
prior=prior,
device=device,
logging_level=logging_level,
summary_writer=summary_writer,
show_progress_bars=show_progress_bars,
)
# As detailed in the docstring, `density_estimator` is either a string or
# a callable. The function creating the neural network is attached to
# `_build_neural_net`. It will be called in the first round and receive
# thetas and xs as inputs, so that they can be used for shape inference and
# potentially for z-scoring.
check_estimator_arg(density_estimator)
if isinstance(density_estimator, str):
self._build_neural_net = utils.likelihood_nn(
model=density_estimator)
else:
self._build_neural_net = density_estimator
# SNLE-specific summary_writer fields.
self._summary.update({"mcmc_times": []}) # type: ignore
def __init__(
self,
prior,
density_estimator: Union[str, Callable] = "maf",
device: str = "cpu",
logging_level: Union[int, str] = "WARNING",
summary_writer: Optional[SummaryWriter] = None,
show_progress_bars: bool = True,
**unused_args
):
r"""Base class for Sequential Neural Likelihood Estimation methods.
Args:
density_estimator: If it is a string, use a pre-configured network of the
provided type (one of nsf, maf, mdn, made). Alternatively, a function
that builds a custom neural network can be provided. The function will
be called with the first batch of simulations (theta, x), which can
thus be used for shape inference and potentially for z-scoring. It
needs to return a PyTorch `nn.Module` implementing the density
estimator. The density estimator needs to provide the methods
`.log_prob` and `.sample()`.
unused_args: Absorbs additional arguments. No entries will be used. If it
is not empty, we warn. In future versions, when the new interface of
0.14.0 is more mature, we will remove this argument.
See docstring of `NeuralInference` class for all other arguments.
"""
super().__init__(
prior=prior,
device=device,
logging_level=logging_level,
summary_writer=summary_writer,
show_progress_bars=show_progress_bars,
**unused_args
)
# As detailed in the docstring, `density_estimator` is either a string or
# a callable. The function creating the neural network is attached to
# `_build_neural_net`. It will be called in the first round and receive
# thetas and xs as inputs, so that they can be used for shape inference and
# potentially for z-scoring.
check_estimator_arg(density_estimator)
if isinstance(density_estimator, str):
self._build_neural_net = utils.likelihood_nn(model=density_estimator)
else:
self._build_neural_net = density_estimator
# SNLE-specific summary_writer fields.
self._summary.update({"mcmc_times": []}) # type: ignore
def __init__(
self,
simulator: Callable,
prior,
true_observation: Tensor,
density_estimator: Optional[nn.Module],
simulation_batch_size: int = 1,
summary_writer: SummaryWriter = None,
device: torch.device = None,
mcmc_method: str = "slice-np",
):
r"""Sequential Neural Likelihood
Implementation of
_Sequential Neural Likelihood: Fast Likelihood-free Inference with Autoregressive Flows_ by Papamakarios et al., AISTATS 2019, https://arxiv.org/abs/1805.07226
Args:
density_estimator: Conditional density estimator $q(x|\theta)$, a nn.Module with `log_prob` and `sample` methods
"""
super().__init__(
simulator,
prior,
true_observation,
simulation_batch_size,
device,
summary_writer,
)
if density_estimator is None:
density_estimator = utils.likelihood_nn(
model="maf", prior=self._prior, context=self._true_observation,
)
# create neural posterior which can sample()
self._neural_posterior = Posterior(
algorithm_family="snl",
neural_net=density_estimator,
prior=prior,
context=true_observation,
mcmc_method=mcmc_method,
get_potential_function=PotentialFunctionProvider(),
)
# XXX why not density_estimator.train(True)???
self._neural_posterior.neural_net.train(True)
# SNL-specific summary_writer fields
self._summary.update({"mcmc_times": []})
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(),
)
sample_kwargs = {"sample_with_mcmc": True}
num_trials = 1
elif method == SNLE:
net_provider = utils.likelihood_nn(model="mdn",
embedding_net=embedding())
sample_kwargs = {}
num_trials = 2
else:
net_provider = utils.classifier_nn(
model="mlp",
embedding_net_x=embedding(),
)
sample_kwargs = {
"mcmc_method": "slice_np_vectorized",
"mcmc_parameters": {
"num_chains": 2
},
}
num_trials = 2
inference = method(prior, net_provider, show_progress_bars=False)
theta, x = simulate_for_sbi(simulator, prior, num_simulations)
_ = inference.append_simulations(theta, x).train(training_batch_size=100,
max_num_epochs=10)
x_o = simulator(theta_o.repeat(num_trials, 1))
posterior = inference.build_posterior(**sample_kwargs).set_default_x(x_o)
posterior.log_prob(
posterior.sample((num_samples, ), show_progress_bars=False))
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)
x_o = simulator(theta_o)
if method == SNPE:
kwargs = dict(
density_estimator=utils.posterior_nn(
model="mdn",
embedding_net=embedding(),
),
sample_with_mcmc=True,
)
elif method == SNLE:
kwargs = dict(density_estimator=utils.likelihood_nn(
model="mdn", embedding_net=embedding()))
else:
kwargs = dict(density_estimator=utils.classifier_nn(
model="mlp",
embedding_net_x=embedding(),
))
infer = method(
simulator,
prior,
1,
1,
show_progress_bars=False,
mcmc_method="slice_np",
**kwargs,
)
posterior = infer(num_simulations=num_simulations,
training_batch_size=100,
max_num_epochs=10).set_default_x(x_o)
posterior.log_prob(
posterior.sample((num_samples, ), show_progress_bars=False))
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_embedding_net_api(method, num_dim: int, embedding_net: str):
"""Tests the API when using a preconfigured embedding net."""
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 = utils.BoxUniform(-2.0 * ones(num_dim), 2.0 * ones(num_dim))
theta = prior.sample((1000, ))
x = linear_gaussian(theta, likelihood_shift, likelihood_cov)
if embedding_net == "mlp":
embedding = FCEmbedding(input_dim=num_dim)
else:
raise NameError
if method == "SNPE":
density_estimator = posterior_nn("maf", embedding_net=embedding)
inference = SNPE(prior,
density_estimator=density_estimator,
show_progress_bars=False)
elif method == "SNLE":
density_estimator = likelihood_nn("maf", embedding_net=embedding)
inference = SNLE(prior,
density_estimator=density_estimator,
show_progress_bars=False)
elif method == "SNRE":
classifier = classifier_nn("resnet", embedding_net_x=embedding)
inference = SNRE(prior,
classifier=classifier,
show_progress_bars=False)
else:
raise NameError
_ = inference.append_simulations(theta, x).train(max_num_epochs=5)
posterior = inference.build_posterior().set_default_x(x_o)
s = posterior.sample((1, ))
_ = posterior.potential(s)
def test_snl_posterior_correction(mcmc_method: str, 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:
mcmc_method: which mcmc method to use for sampling
prior_str: use gaussian or uniform prior
"""
num_dim = 2
num_samples = 30
true_observation = torch.zeros((1, num_dim))
if prior_str == "gaussian":
prior = distributions.MultivariateNormal(
loc=torch.zeros(num_dim), covariance_matrix=torch.eye(num_dim))
else:
prior = utils.BoxUniform(-1.0 * torch.ones(num_dim),
torch.ones(num_dim))
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-np",
)
posterior = infer(num_rounds=1, num_simulations_per_round=1000)
posterior.sample(num_samples=num_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 __init__(
self,
simulator: Callable,
prior,
num_workers: int = 1,
simulation_batch_size: int = 1,
density_estimator: Union[str, Callable] = "maf",
mcmc_method: str = "slice_np",
mcmc_parameters: Optional[Dict[str, Any]] = None,
device: str = "cpu",
logging_level: Union[int, str] = "WARNING",
summary_writer: Optional[SummaryWriter] = None,
show_progress_bars: bool = True,
show_round_summary: bool = False,
):
r"""Base class for Sequential Neural Likelihood Estimation methods.
Args:
density_estimator: If it is a string, use a pre-configured network of the
provided type (one of nsf, maf, mdn, made). Alternatively, a function
that builds a custom neural network can be provided. The function will
be called with the first batch of simulations (theta, x), which can
thus be used for shape inference and potentially for z-scoring. It
needs to return a PyTorch `nn.Module` implementing the density
estimator. The density estimator needs to provide the methods
`.log_prob` and `.sample()`.
mcmc_method: Method used for MCMC sampling, one of `slice_np`, `slice`, `hmc`, `nuts`.
Currently defaults to `slice_np` for a custom numpy implementation of
slice sampling; select `hmc`, `nuts` or `slice` for Pyro-based sampling.
mcmc_parameters: Dictionary overriding the default parameters for MCMC.
The following parameters are supported: `thin` to set the thinning
factor for the chain, `warmup_steps` to set the initial number of
samples to discard, `num_chains` for the number of chains, `init_strategy`
for the initialisation strategy for chains; `prior` will draw init
locations from prior, whereas `sir` will use Sequential-Importance-
Resampling using `init_strategy_num_candidates` to find init
locations.
See docstring of `NeuralInference` class for all other arguments.
"""
super().__init__(
simulator=simulator,
prior=prior,
num_workers=num_workers,
simulation_batch_size=simulation_batch_size,
device=device,
logging_level=logging_level,
summary_writer=summary_writer,
show_progress_bars=show_progress_bars,
show_round_summary=show_round_summary,
)
# As detailed in the docstring, `density_estimator` is either a string or
# a callable. The function creating the neural network is attached to
# `_build_neural_net`. It will be called in the first round and receive
# thetas and xs as inputs, so that they can be used for shape inference and
# potentially for z-scoring.
check_estimator_arg(density_estimator)
if isinstance(density_estimator, str):
self._build_neural_net = utils.likelihood_nn(
model=density_estimator)
else:
self._build_neural_net = density_estimator
self._posterior = None
self._sample_with_mcmc = True
self._mcmc_method = mcmc_method
self._mcmc_parameters = mcmc_parameters
# SNLE-specific summary_writer fields.
self._summary.update({"mcmc_times": []}) # type: ignore
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_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_training_and_mcmc_on_device(method, model, data_device, training_device):
"""Test training on devices.
This test does not check training speeds.
"""
training_device = process_device(training_device)
num_dim = 2
num_samples = 10
num_simulations = 500
max_num_epochs = 5
x_o = zeros(1, num_dim).to(data_device)
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)
if method in [SNPE_A, SNPE_C]:
kwargs = dict(
density_estimator=utils.posterior_nn(model=model),
)
mcmc_kwargs = (
dict(
sample_with="mcmc",
mcmc_method="slice_np",
)
if method == SNPE_C
else {}
)
elif method == SNLE:
kwargs = dict(
density_estimator=utils.likelihood_nn(model=model),
)
mcmc_kwargs = dict(sample_with="mcmc", mcmc_method="slice")
elif method in (SNRE_A, SNRE_B):
kwargs = dict(
classifier=utils.classifier_nn(model=model),
)
mcmc_kwargs = dict(
sample_with="mcmc",
mcmc_method="slice_np_vectorized",
)
else:
raise ValueError()
inferer = method(prior, show_progress_bars=False, device=training_device, **kwargs)
proposals = [prior]
# Test for two rounds.
for _ in range(2):
theta, x = simulate_for_sbi(simulator, prior, num_simulations)
theta, x = theta.to(data_device), x.to(data_device)
_ = inferer.append_simulations(theta, x).train(
training_batch_size=100, max_num_epochs=max_num_epochs
)
posterior = inferer.build_posterior(**mcmc_kwargs).set_default_x(x_o)
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,), x=x_o, **mcmc_kwargs)
proposals[-1].log_prob(samples)
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_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
