def test_mdn_with_1D_uniform_prior():
"""
Note, we have this test because for 1D uniform priors, mdn log prob evaluation
results in batch_size x batch_size return. This is probably because Uniform does
not allow for event dimension > 1 and somewhere in pyknos it is used as if this was
possible.
Casting to BoxUniform solves it.
"""
num_dim = 1
x_o = torch.tensor([[1.0]])
num_samples = 100
# likelihood_mean will be likelihood_shift+theta
likelihood_shift = -1.0 * ones(num_dim)
likelihood_cov = 0.3 * eye(num_dim)
prior = Uniform(low=torch.zeros(num_dim), high=torch.ones(num_dim))
def simulator(theta: Tensor) -> Tensor:
return linear_gaussian(theta, likelihood_shift, likelihood_cov)
simulator, prior = prepare_for_sbi(simulator, prior)
infer = SNPE(simulator, prior, density_estimator="mdn")
posterior = infer(num_simulations=100,
training_batch_size=50).set_default_x(x_o)
samples = posterior.sample((num_samples, ))
log_probs = posterior.log_prob(samples)
assert log_probs.shape == torch.Size([num_samples])
def infer(
simulator: Callable, prior, method: str, num_simulations: int, num_workers: int = 1
) -> NeuralPosterior:
r"""
Return posterior distribution by running simulation-based inference.
This function provides a simple interface to run sbi. Inference is run for a single
round and hence the returned posterior $p(\theta|x)$ can be sampled and evaluated
for any $x$ (i.e. it is amortized).
The scope of this function is limited to the most essential features of sbi. For
more flexibility (e.g. multi-round inference, different density estimators) please
use the flexible interface described here:
https://www.mackelab.org/sbi/tutorial/02_flexible_interface/
Args:
simulator: A function that takes parameters $\theta$ and maps them to
simulations, or observations, `x`, $\mathrm{sim}(\theta)\to x$. Any
regular Python callable (i.e. function or class with `__call__` method)
can be used.
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.
method: What inference method to use. Either of SNPE, SNLE or SNRE.
num_simulations: Number of simulation calls. More simulations means a longer
runtime, but a better posterior estimate.
num_workers: Number of parallel workers to use for simulations.
Returns: Posterior over parameters conditional on observations (amortized).
"""
try:
method_fun: Callable = getattr(sbi.inference, method.upper())
except AttributeError:
raise NameError(
"Method not available. `method` must be one of 'SNPE', 'SNLE', 'SNRE'."
)
simulator, prior = prepare_for_sbi(simulator, prior)
inference = method_fun(prior)
theta, x = simulate_for_sbi(
simulator=simulator,
proposal=prior,
num_simulations=num_simulations,
num_workers=num_workers,
)
_ = inference.append_simulations(theta, x).train()
posterior = inference.build_posterior()
return posterior
def test_inference_with_user_sbi_problems(user_simulator: Callable,
user_prior):
"""
Test inference with combinations of user defined simulators, priors and x_os.
"""
infer = SNPE_C(
*prepare_for_sbi(user_simulator, user_prior),
density_estimator="maf",
simulation_batch_size=1,
show_progress_bars=False,
)
# Run inference.
_ = infer(num_rounds=1, num_simulations_per_round=100, max_num_epochs=2)
def test_prepare_sbi_problem(simulator: Callable, prior):
"""Test user interface by passing different kinds of simulators, prior and data.
Args:
simulator: simulator function
prior: prior as defined by the user (pytorch, scipy, custom)
x_shape: shape of data as defined by the user.
"""
simulator, prior = prepare_for_sbi(simulator, prior)
# check batch sims and type
n_batch = 1
assert simulator(prior.sample((n_batch, ))).shape[0] == n_batch
assert isinstance(simulator(prior.sample((1, ))), Tensor)
assert prior.sample().dtype == torch.float32
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()
