def test_c2st_sre_on_linearGaussian_different_dims(set_seed):
"""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.
Args:
set_seed: fixture for manual seeding
"""
device = "cpu"
configure_default_device(device)
theta_dim = 3
x_dim = 2
discard_dims = theta_dim - x_dim
x_o = ones(1, x_dim)
num_samples = 1000
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)
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)
infer = SRE(
*prepare_for_sbi(simulator, prior),
classifier="resnet",
simulation_batch_size=50,
show_progress_bars=False,
device=device,
)
posterior = infer(num_rounds=1, num_simulations_per_round=5000)
samples = posterior.sample((num_samples, ),
x=x_o,
mcmc_parameters={"thin": 3})
# Compute the c2st and assert it is near chance level of 0.5.
check_c2st(samples, target_samples, alg="snpe_c")
def test_c2st_snle_external_data_on_linearGaussian(set_seed):
"""Test whether SNPE C infers well a simple example with available ground truth.
Args:
set_seed: fixture for manual seeding
"""
num_dim = 2
device = "cpu"
configure_default_device(device)
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)
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):
return linear_gaussian(theta, likelihood_shift, likelihood_cov)
infer = SNL(
*prepare_for_sbi(simulator, prior),
simulation_batch_size=1000,
show_progress_bars=False,
device=device,
)
external_theta = prior.sample((1000, ))
external_x = simulator(external_theta)
infer.provide_presimulated(external_theta, external_x)
posterior = infer(
num_rounds=1,
num_simulations_per_round=1000,
training_batch_size=100,
).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")
def __init__(
self,
simulator: Callable,
prior,
num_workers: int = 1,
simulation_batch_size: int = 1,
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 inference methods.
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.
num_workers: Number of parallel workers to use for simulations.
simulation_batch_size: Number of parameter sets that the simulator
maps to data x at once. If None, we simulate all parameter sets at the
same time. If >= 1, the simulator has to process data of shape
(simulation_batch_size, parameter_dimension).
device: torch device on which to compute, e.g. gpu or cpu.
logging_level: Minimum severity of messages to log. One of the strings
"INFO", "WARNING", "DEBUG", "ERROR" and "CRITICAL".
summary_writer: A `SummaryWriter` to control, among others, log
file location (default is `<current working directory>/logs`.)
show_progress_bars: Whether to show a progressbar during simulation and
sampling.
show_round_summary: Whether to show the validation loss and leakage after
each round.
"""
# We set the device globally by setting the default tensor type for all tensors.
assert device in (
"gpu",
"cpu",
), "Currently, only 'gpu' or 'cpu' are supported as devices."
self._device = configure_default_device(device)
self._simulator, self._prior = simulator, prior
self._show_progress_bars = show_progress_bars
self._show_round_summary = show_round_summary
self._batched_simulator = lambda theta: simulate_in_batches(
self._simulator,
theta,
simulation_batch_size,
num_workers,
self._show_progress_bars,
)
# Initialize roundwise (theta, x, prior_masks) for storage of parameters,
# simulations and masks indicating if simulations came from prior.
self._theta_roundwise, self._x_roundwise, self._prior_masks = [], [], []
# Initialize list that indicates the round from which simulations were drawn.
self._data_round_index = []
self._round = 0
# XXX We could instantiate here the Posterior for all children. Two problems:
# 1. We must dispatch to right PotentialProvider for mcmc based on name
# 2. `method_family` cannot be resolved only from `self.__class__.__name__`,
# since SRE, AALR demand different handling but are both in SRE class.
self._summary_writer = (self._default_summary_writer()
if summary_writer is None else summary_writer)
# Logging during training (by SummaryWriter).
self._summary = dict(
median_observation_distances=[],
epochs=[],
best_validation_log_probs=[],
)
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
"""
device = "cpu"
configure_default_device(device)
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)
infer = SNPE_C(
*prepare_for_sbi(simulator, prior),
simulation_batch_size=1000,
show_progress_bars=False,
sample_with_mcmc=False,
device=device,
)
posterior = infer(num_simulations=2000,
training_batch_size=100).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
device = "cpu"
configure_default_device(device)
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)
inference = SNPE_C(
*prepare_for_sbi(simulator, prior),
density_estimator=net,
simulation_batch_size=1,
show_progress_bars=True,
device=device,
)
# We need a pretty big dataset to properly model the bimodality.
posterior = inference(num_simulations=10000,
proposal=None,
max_num_epochs=50).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