def _set_maybe_z_scored_prior(self) -> None:
r"""
Compute and store potentially standardized prior (if z-scoring was requested).
This function is highly similar to `SNPE_C._set_maybe_z_scored_prior()`.
The proposal posterior is:
$p(\theta|x) = 1/Z * q(\theta|x) * prop(\theta) / p(\theta)$
Let's denote z-scored theta by `a`: a = (theta - mean) / std
Then $p'(a|x) = 1/Z_2 * q'(a|x) * prop'(a) / p'(a)$
The ' indicates that the evaluation occurs in standardized space. The constant
scaling factor has been absorbed into $Z_2$.
From the above equation, we see that we need to evaluate the prior **in
standardized space**. We build the standardized prior in this function.
The standardize transform that is applied to the samples theta does not use
the exact prior mean and std (due to implementation issues). Hence, the z-scored
prior will not be exactly have mean=0 and std=1.
"""
if self.z_score_theta:
# Default to cpu to avoid mismatch with prior (always on cpu) below.
scale = self._neural_net._transform._transforms[0]._scale.cpu()
shift = self._neural_net._transform._transforms[0]._shift.cpu()
# Following the definition of the linear transform in
# `standardizing_transform` in `sbiutils.py`:
# shift=-mean / std
# scale=1 / std
# Solving these equations for mean and std:
estim_prior_std = 1 / scale
estim_prior_mean = -shift * estim_prior_std
# Compute the discrepancy of the true prior mean and std and the mean and
# std that was empirically estimated from samples.
# N(theta|m,s) = N((theta-m_e)/s_e|(m-m_e)/s_e, s/s_e)
# Above: m,s are true prior mean and std. m_e,s_e are estimated prior mean
# and std (estimated from samples and used to build standardize transform).
almost_zero_mean = (self._prior.mean -
estim_prior_mean) / estim_prior_std
almost_one_std = torch.sqrt(self._prior.variance) / estim_prior_std
if isinstance(self._prior, MultivariateNormal):
self._maybe_z_scored_prior = MultivariateNormal(
almost_zero_mean, torch.diag(almost_one_std))
else:
range_ = torch.sqrt(almost_one_std * 3.0)
self._maybe_z_scored_prior = utils.BoxUniform(
almost_zero_mean - range_, almost_zero_mean + range_)
else:
self._maybe_z_scored_prior = self._prior
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_inference_with_rejection_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
num_simulations = 2000
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], 500)
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()
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_z_scoring_warning(snpe_method: type):
# Create data with large variance.
num_dim = 2
theta = torch.ones(100, num_dim)
x = torch.rand(100, num_dim)
x[:50] += 1e7
# Make sure a warning is raised because z-scoring will map these data to duplicate
# data points.
with pytest.warns(UserWarning, match="Z-scoring these simulation outputs"):
snpe_method(utils.BoxUniform(zeros(num_dim),
ones(num_dim))).append_simulations(
theta, x).train(max_num_epochs=1)
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_api_snl_sampling_methods(sampling_method: str, prior_str: str,
set_seed):
"""Runs SNL on linear Gaussian and tests sampling from posterior via mcmc.
Args:
mcmc_method: which mcmc method to use for sampling
prior_str: use gaussian or uniform prior
set_seed: fixture for manual seeding
"""
num_dim = 2
num_samples = 10
num_trials = 2
# HMC with uniform prior needs good likelihood.
num_simulations = 10000 if sampling_method == "hmc" else 1000
x_o = zeros((num_trials, num_dim))
# Test for multiple chains is cheap when vectorized.
num_chains = 3 if sampling_method == "slice_np_vectorized" else 1
if sampling_method == "rejection":
sample_with = "rejection"
else:
sample_with = "mcmc"
if prior_str == "gaussian":
prior = MultivariateNormal(loc=zeros(num_dim),
covariance_matrix=eye(num_dim))
else:
prior = utils.BoxUniform(-1.0 * ones(num_dim), ones(num_dim))
simulator, prior = prepare_for_sbi(diagonal_linear_gaussian, prior)
inference = SNL(prior, show_progress_bars=False)
theta, x = simulate_for_sbi(simulator,
prior,
num_simulations,
simulation_batch_size=50)
_ = inference.append_simulations(theta, x).train(max_num_epochs=5)
posterior = inference.build_posterior(
sample_with=sample_with,
mcmc_method=sampling_method).set_default_x(x_o)
posterior.sample(
sample_shape=(num_samples, ),
x=x_o,
mcmc_parameters={
"thin": 3,
"num_chains": num_chains
},
)
def get_simulator(N, g, K):
_W_eigs = get_W_eigs_np(g, K)
num_dim = 2 * N * RANK
prior = utils.BoxUniform(low=-1.0 * torch.ones(num_dim),
high=1.0 * torch.ones(num_dim))
def simulator(params):
params = params.numpy()
U = np.reshape(params[:(RANK * N)], (N, RANK))
V = np.reshape(params[(RANK * N):], (N, RANK))
x = _W_eigs(U, V)
return x
return simulator, prior
def test_api_snpe_c_posterior_correction(sample_with, mcmc_method, prior_str):
"""Test that leakage correction applied to sampling works, with both MCMC and
rejection.
"""
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))
simulator, prior = prepare_for_sbi(
lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov),
prior)
inference = SNPE_C(prior, show_progress_bars=False)
theta, x = simulate_for_sbi(simulator, prior, 1000)
posterior_estimator = inference.append_simulations(theta, x).train()
potential_fn, theta_transform = posterior_estimator_based_potential(
posterior_estimator, prior, x_o)
if sample_with == "mcmc":
posterior = MCMCPosterior(
potential_fn=potential_fn,
theta_transform=theta_transform,
proposal=prior,
method=mcmc_method,
)
elif sample_with == "rejection":
posterior = RejectionPosterior(
potential_fn=potential_fn,
proposal=prior,
theta_transform=theta_transform,
)
# Posterior should be corrected for leakage even if num_rounds just 1.
samples = posterior.sample((10, ))
# Evaluate the samples to check correction factor.
_ = posterior.log_prob(samples)
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_apt_posterior_correction(sample_with_mcmc, mcmc_method, prior):
"""Test that leakage correction applied to sampling works, with both MCMC and rejection."""
num_dim = 2
if prior == "gaussian":
prior = distributions.MultivariateNormal(
loc=torch.zeros(num_dim), covariance_matrix=torch.eye(num_dim))
else:
prior = utils.BoxUniform(low=-1.0 * torch.ones(num_dim),
high=torch.ones(num_dim))
true_observation = torch.zeros((1, num_dim))
neural_net = utils.posterior_nn(
model="maf",
prior=prior,
context=true_observation,
)
infer = SnpeC(
simulator=linear_gaussian,
true_observation=true_observation,
density_estimator=neural_net,
prior=prior,
num_atoms=-1,
z_score_obs=True,
simulation_batch_size=50,
use_combined_loss=False,
retrain_from_scratch_each_round=False,
discard_prior_samples=False,
sample_with_mcmc=sample_with_mcmc,
mcmc_method=mcmc_method,
)
# 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 samples from posterior (should be corrected for leakage)
# even if just num_rounds=1.
samples = posterior.sample(10)
# Evaluate the samples to check correction factor.
posterior.log_prob(samples)
def flat_prior(self, low = None, high = None, log = True):
"""
BoxUniform prior
"""
if low == None:
if log:
low = torch.tensor([-7, .5])
else:
low = torch.tensor([1e-7, 10**.5])
if high == None:
if log:
high = torch.tensor([-2.5, 20])
else:
high = torch.tensor([10**-2.5, 20])
return utils.BoxUniform(low = low, high = high)
def test_inference_with_nan_simulator(method: type, exclude_invalid_x: bool,
percent_nans: float, 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
num_simulations = 2000
def linear_gaussian_nan(theta,
likelihood_shift=likelihood_shift,
likelihood_cov=likelihood_cov):
x = linear_gaussian(theta, likelihood_shift, likelihood_cov)
# Set nan randomly.
x[torch.rand(x.shape) < (percent_nans * 1.0 /
x.shape[1])] = 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)
inference = method(prior)
theta, x = simulate_for_sbi(simulator, prior, num_simulations)
_ = inference.append_simulations(
theta, x).train(exclude_invalid_x=exclude_invalid_x)
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"{method}")
def test_restricted_prior_log_prob(prior):
"""Test whether the log-prob method of the restricted prior works appropriately."""
def simulator(theta):
perturbed_theta = theta + 0.5 * torch.randn(2)
perturbed_theta[theta[:, 0] < 0.8] = torch.as_tensor(
[float("nan"), float("nan")])
return perturbed_theta
if prior == "uniform":
prior = utils.BoxUniform(-2 * torch.ones(2), 2 * torch.ones(2))
else:
prior = MultivariateNormal(torch.zeros(2), torch.eye(2))
theta, x = simulate_for_sbi(simulator, prior, 1000)
restriction_estimator = RestrictionEstimator(prior=prior)
restriction_estimator.append_simulations(theta, x)
_ = restriction_estimator.train(max_num_epochs=40)
restricted_prior = restriction_estimator.restrict_prior()
def integrate_grid(distribution):
resolution = 500
range_ = 4
x = torch.linspace(-range_, range_, resolution)
y = torch.linspace(-range_, range_, resolution)
X, Y = torch.meshgrid(x, y)
xy = torch.stack([X, Y])
xy = torch.reshape(xy, (2, resolution**2)).T
dist_on_grid = torch.exp(distribution.log_prob(xy))
integral = torch.sum(dist_on_grid) / resolution**2 * (2 * range_)**2
return integral
integal_restricted = integrate_grid(restricted_prior)
error = torch.abs(integal_restricted - torch.as_tensor(1.0))
assert error < 0.01, "The restricted prior does not integrate to one."
theta = prior.sample((10_000, ))
restricted_prior_probs = torch.exp(restricted_prior.log_prob(theta))
valid_thetas = restricted_prior.predict(theta).bool()
assert torch.all(restricted_prior_probs[valid_thetas] > 0.0
), "Accepted theta have zero probability."
assert torch.all(restricted_prior_probs[torch.logical_not(valid_thetas)] ==
0.0), "Rejected theta has non-zero probablity."
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 infer_and_save_posterior(device: str, simulator_type: str,
simulation_type: str, params: Dict) -> None:
parameter_prior = sbi_utils.BoxUniform(
low=torch.tensor([0.0, 0.0]).type(torch.FloatTensor),
high=torch.tensor([4.0, 4.0]).type(torch.FloatTensor))
inferrer = inference_class.TimeSeriesInference(
parameter_prior=parameter_prior, device=device)
inferrer.load_simulator(dirname=ARTIFACT_PATH,
simulator_type=simulator_type,
simulation_type=simulation_type)
batch_size = params.pop("batch_size")
learning_rate = params.pop("learning_rate")
hidden_features = params.pop("hidden_features")
num_transforms = params.pop("num_transforms")
inferrer.infer_snpe_posterior(
embedding_net_conf=params,
batch_size=batch_size,
learning_rate=learning_rate,
hidden_features=hidden_features,
num_transforms=num_transforms,
)
inferrer.save_inference(ARTIFACT_PATH)
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_sampling_methods(mcmc_method: str, prior_str: str, set_seed):
"""Runs SNL on linear Gaussian and tests sampling from posterior via mcmc.
Args:
mcmc_method: which mcmc method to use for sampling
prior_str: use gaussian or uniform prior
set_seed: fixture for manual seeding
"""
num_dim = 2
num_samples = 10
x_o = zeros((1, num_dim))
if prior_str == "gaussian":
prior = MultivariateNormal(loc=zeros(num_dim),
covariance_matrix=eye(num_dim))
else:
prior = utils.BoxUniform(-1.0 * ones(num_dim), ones(num_dim))
simulator, prior = prepare_for_sbi(diagonal_linear_gaussian, prior)
inference = SNL(
prior,
show_progress_bars=False,
)
theta, x = simulate_for_sbi(simulator,
prior,
200,
simulation_batch_size=50)
_ = inference.append_simulations(theta, x).train(max_num_epochs=5)
posterior = inference.build_posterior(
mcmc_method=mcmc_method).set_default_x(x_o)
posterior.sample(sample_shape=(num_samples, ),
x=x_o,
mcmc_parameters={"thin": 3})
def test_sre_posterior_correction(mcmc_method: str, prior_str: str):
"""Test leakage correction both for MCMC and rejection sampling.
Args:
mcmc_method: which mcmc method to use for sampling
prior_str: one of "gaussian" or "uniform"
"""
num_dim = 2
if prior_str == "gaussian":
prior = distributions.MultivariateNormal(
loc=torch.zeros(num_dim), covariance_matrix=torch.eye(num_dim))
else:
prior = utils.BoxUniform(low=-1.0 * torch.ones(num_dim),
high=torch.ones(num_dim))
true_observation = torch.zeros(num_dim)
classifier = utils.classifier_nn(
"resnet",
prior=prior,
context=true_observation,
)
infer = SRE(
simulator=linear_gaussian,
prior=prior,
true_observation=true_observation,
classifier=classifier,
simulation_batch_size=50,
mcmc_method=mcmc_method,
)
posterior = infer(num_rounds=1, num_simulations_per_round=1000)
samples = posterior.sample(num_samples=50)
def sample(self, sample_shape: torch.Size):
"""
Return samples from the toy density.
We first sample from a box uniform and then compute their L1-distance to a
hyperplane in the 8D parameter space. We then accept with probability
(1.-distance). If the distance is larger than 1.0, we never accept.
"""
num_dim = 3
all_samples = torch.empty(0, num_dim)
num_accepted = 0
while num_accepted < sample_shape[0]:
proposed_samples = utils.BoxUniform(
-2 * torch.ones(num_dim),
2 * torch.ones(num_dim)).sample(sample_shape)
vec_prior_samples = proposed_samples * self.normal_vec
dist_to_zero = torch.abs(torch.sum(vec_prior_samples, dim=1))
accept_or_not = (self.noise_factor * torch.rand(dist_to_zero.shape)
> dist_to_zero)
accepted_samples = proposed_samples[accept_or_not]
num_accepted += accepted_samples.shape[0]
all_samples = torch.cat((all_samples, accepted_samples), dim=0)
return all_samples
def test_api_sre_sampling_methods(mcmc_method: str, prior_str: str, set_seed):
"""Test leakage correction both for MCMC and rejection sampling.
Args:
mcmc_method: which mcmc method to use for sampling
prior_str: one of "gaussian" or "uniform"
set_seed: fixture for manual seeding
"""
num_dim = 2
x_o = zeros(num_dim)
if prior_str == "gaussian":
prior = MultivariateNormal(loc=zeros(num_dim), covariance_matrix=eye(num_dim))
else:
prior = utils.BoxUniform(low=-1.0 * ones(num_dim), high=ones(num_dim))
simulator, prior = prepare_for_sbi(diagonal_linear_gaussian, prior)
inference = SRE(prior, classifier="resnet", show_progress_bars=False,)
theta, x = simulate_for_sbi(simulator, prior, 200, simulation_batch_size=50)
_ = inference.append_simulations(theta, x).train(max_num_epochs=5)
posterior = inference.build_posterior(mcmc_method=mcmc_method)
posterior.sample(sample_shape=(10,), x=x_o)
def main() -> None:
# Parse command line arguments
parser = argparse.ArgumentParser(description="CNN Tuning")
parser.add_argument("--max_evals",
required=False,
type=int,
default=5,
help="Maximum number of sampled sets of parameters ")
args, *_ = parser.parse_known_args()
mlflow.set_experiment(f"snpe-fully-padded-cnn-timeseries-tuning")
# Initialize the model and load context - needs to be done whether using local data or doing transforms
print("\t Initialize inference object")
parameter_prior = sbi_utils.BoxUniform(
low=torch.tensor([0.0, 0.0]).type(torch.FloatTensor),
high=torch.tensor([4.0, 4.0]).type(torch.FloatTensor))
inferrer = TimeSeriesInference(parameter_prior=parameter_prior,
device="gpu")
inferrer.load_simulator(dirname=ARTIFACT_PATH,
simulator_type="double_rho",
simulation_type="timeseries")
print("\t Tuning model")
tuning(inferrer=inferrer, max_evals=args.max_evals)
import torch
import pandas as pd
import sys
import csv
from vit_pytorch.efficient import ViT
from linformer import Linformer
from vit_pytorch import ViT as ViT_modified
from collections import Counter, OrderedDict
import torch.nn as nn
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
Prior= utils.BoxUniform(low=torch.tensor([0.1, -1.5, -6.0, -3.5, -3.5, 2.0, 0.7, 300.0, 0., 0., 0.,\
1., 0. ]), \
high=torch.tensor([1.6, 1.5, 3.0, 4.5, 4.5, 5.5, 2.0, 2300.0, 1., 1., 1.,\
2., 1. ]), device='cuda')
def build_den():
class VisualTrans(nn.Module):
def __init__(self, file_path):
super(VisualTrans, self).__init__()
self.file_path = file_path
self.model = ViT_modified(
n_classes=1,
image_size=(1,
962), # image size is a tuple of (height, width)
patch_size=(1, 13), # patch size is a tuple of (height, width)
def prior_(self):
self.prior = utils.BoxUniform(low=self.sim.min, high=self.sim.max)
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
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."
# 5: "joint_6_dryfriction",
# 6: "joint_7_dryfriction",
# }
# dp_mapping = create_damping_dryfriction_domain_param_map_wamjsc()
# Behavioral policy
policy_hparam = dict()
policy = TimePolicy(env_real.spec, wam_jsp_7dof_sin, env_real.dt)
# Prior
dp_nom = env_sim.get_nominal_domain_param()
prior_hparam = dict(
low=to.tensor([dp_nom[name] * 0.5 for name in dp_mapping.values()]),
high=to.tensor([dp_nom[name] * 1.5 for name in dp_mapping.values()]),
)
prior = sbiutils.BoxUniform(**prior_hparam)
# Time series embedding
embedding_hparam = dict(
downsampling_factor=1,
# len_rollouts=env_sim.max_steps,
# recurrent_network_type=nn.RNN,
# only_last_output=True,
# hidden_size=20,
# num_recurrent_layers=1,
# output_size=1,
)
embedding = create_embedding(BayesSimEmbedding.name, env_sim.spec,
**embedding_hparam)
# Posterior (normalizing flow)
% (N, num_sims, num_batch, num_transforms, num_atoms, g, K, rs))
save_path = os.path.join(base_path, save_dir)
if not os.path.exists(save_path):
os.makedirs(save_path)
if os.path.exists(os.path.join(base_path, save_dir, "optim.pkl")):
print("SNPE optimization already run. Exitting.")
exit()
_W_eigs = get_W_eigs_np(g, K)
M = 1000
RANK = 2
num_dim = 2 * N * RANK
prior = utils.BoxUniform(low=-1.0 * torch.ones(num_dim),
high=1.0 * torch.ones(num_dim))
def simulator(params):
params = params.numpy()
U = np.reshape(params[:(RANK * N)], (N, RANK))
V = np.reshape(params[(RANK * N):], (N, RANK))
x = _W_eigs(U, V)
return x
simulator, prior = prepare_for_sbi(simulator, prior)
density_estimator_build_fun = posterior_nn(
model="maf",
hidden_features=50,
num_transforms=num_transforms,
def run_snpe(total_runs=10, num_generation=6, seed=46):
theta_true = np.log([1.0, 0.005, 1.0])
torch.manual_seed(seed)
num_workers = 16 #for parallel execution of simulations
Ndata = 1000
use_embedding_net = True
use_mcmc = False #becomes very slow, but can handle leakage
result_posterior = []
store_time = []
a0, b0 = np.log(0.002), np.log(2)
prior = utils.BoxUniform(low=torch.tensor([a0, a0, a0]),
high=torch.tensor([b0, b0, b0]))
simulator_sbi, prior = prepare_for_sbi(wrapper, prior)
x_o = np.load("target_original_shape_ts.npy")
x_o = torch.tensor(x_o)
#x_o = torch.flatten(x_o)#.flatten()reshape((2,51))
#NN for summary statistic
embedding_net = SummaryNet_large()
try:
for run in range(total_runs):
print(f"starting run {run}")
theta_store = []
time_ticks = []
posteriors = []
proposal = prior
if use_embedding_net:
neural_posterior = utils.posterior_nn(
model='maf',
embedding_net=embedding_net,
hidden_features=10,
num_transforms=2)
else:
neural_posterior = utils.posterior_nn(model='maf',
hidden_features=10,
num_transforms=2)
inference = SNPE_C(prior=prior, density_estimator=neural_posterior)
for i in range(num_generation):
print(f"starting round {i}")
time_begin = time.time()
theta, x = simulate_for_sbi(simulator_sbi,
proposal,
num_simulations=Ndata,
num_workers=num_workers)
mask = torch.tensor(np.invert(np.isnan(x.numpy())[:, 0, 0]))
x = x[mask, :, :]
theta = theta[mask, :]
density_estimator = inference.append_simulations(
theta, x, proposal=proposal).train()
posterior = inference.build_posterior(
density_estimator, sample_with_mcmc=use_mcmc)
print("building post done")
posteriors.append(posterior)
proposal = posterior.set_default_x(x_o)
posterior_samples = posterior.sample((Ndata, ), x=x_o).numpy()
print("Post samples done")
theta_store.append(posterior_samples)
time_ticks.append(time.time() - time_begin)
result_posterior.append(theta_store)
store_time.append(time_ticks)
except KeyboardInterrupt:
return np.asarray(result_posterior), np.asarray(store_time), posteriors
return np.asarray(result_posterior), np.asarray(store_time), posteriors
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."