def test_predictive(num_samples, parallel):
model, data, true_probs = beta_bernoulli()
init_params, potential_fn, transforms, _ = initialize_model(
model, model_args=(data, ))
nuts_kernel = NUTS(potential_fn=potential_fn, transforms=transforms)
mcmc = MCMC(nuts_kernel, 100, initial_params=init_params, warmup_steps=100)
mcmc.run(data)
samples = mcmc.get_samples()
with ignore_experimental_warning():
with optional(pytest.warns(UserWarning), num_samples
not in (None, 100)):
predictive_samples = predictive(model,
samples,
num_samples=num_samples,
return_sites=["beta", "obs"],
parallel=parallel)
# check shapes
assert predictive_samples["beta"].shape == (100, 5)
assert predictive_samples["obs"].shape == (100, 1000, 5)
# check sample mean
assert_close(predictive_samples["obs"].reshape([-1, 5]).mean(0),
true_probs,
rtol=0.1)
def init_guide(self):
nuts_kernel = NUTS(self.model)
num_samples = 3
mcmc = MCMC(nuts_kernel, num_samples=num_samples, warmup_steps=7)
mcmc.run()
hmc_samples = {
k: v.detach().cpu().numpy()
for k, v in mcmc.get_samples().items()
}
# Get latent coordinates samples from mcmc #
gp_mean = torch.tensor(hmc_samples['gp_mean']).reshape(
num_samples, self.V_net, self.H_dim, self.K_net, self.n_dir,
self.n_w)
gp_coord_demean = torch.tensor(
hmc_samples[f'gp_coord_demean']).reshape(num_samples, self.V_net,
self.H_dim, self.K_net,
self.n_dir, self.n_w,
self.T_net)
# keep only last sample
gp_mean_ini = gp_mean[num_samples - 1]
gp_coord_demean_ini = gp_coord_demean[num_samples - 1]
return (gp_mean_ini, gp_coord_demean_ini)
def test_null_model_with_hook(kernel, model, jit, num_chains):
num_warmup, num_samples = 10, 10
initial_params, potential_fn, transforms, _ = initialize_model(
model, num_chains=num_chains)
iters = []
hook = partial(_hook, iters)
mp_context = "spawn" if "CUDA_TEST" in os.environ else None
kern = kernel(potential_fn=potential_fn,
transforms=transforms,
jit_compile=jit)
mcmc = MCMC(kern,
num_samples=num_samples,
warmup_steps=num_warmup,
num_chains=num_chains,
initial_params=initial_params,
hook_fn=hook,
mp_context=mp_context)
mcmc.run()
samples = mcmc.get_samples()
assert samples == {}
if num_chains == 1:
expected = [("Warmup", i)
for i in range(num_warmup)] + [("Sample", i)
for i in range(num_samples)]
assert iters == expected
def hmc(matrix_factorization_poisson, num_samples=1000, warmup_steps=200):
nuts_kernel = NUTS(matrix_factorization_poisson)
mcmc = MCMC(nuts_kernel, num_samples=1000, warmup_steps=200)
mcmc.run(anime_matrix_train.values, k=k)
hmc_samples = {k: v.detach().cpu().numpy()
for k, v in mcmc.get_samples().items()}
return hmc_samples
def test_gamma_poisson(hyperpriors):
def model(data):
with pyro.plate("latent_dim", data.shape[1]):
alpha = pyro.sample(
"alpha", dist.HalfCauchy(1.)) if hyperpriors else torch.tensor(
[1., 1.])
beta = pyro.sample(
"beta", dist.HalfCauchy(1.)) if hyperpriors else torch.tensor(
[1., 1.])
gamma_poisson = GammaPoissonPair()
rate = pyro.sample("rate", gamma_poisson.latent(alpha, beta))
with pyro.plate("data", data.shape[0]):
pyro.sample("obs", gamma_poisson.conditional(rate), obs=data)
true_rate = torch.tensor([3., 10.])
num_samples = 100
data = dist.Poisson(rate=true_rate).sample(
sample_shape=(torch.Size((100, ))))
hmc_kernel = NUTS(collapse_conjugate(model),
jit_compile=True,
ignore_jit_warnings=True)
mcmc = MCMC(hmc_kernel, num_samples=num_samples, warmup_steps=50)
mcmc.run(data)
samples = mcmc.get_samples()
posterior = posterior_replay(model, samples, data, num_samples=num_samples)
assert_equal(posterior["rate"].mean(0), true_rate, prec=0.3)
def test_beta_binomial_hmc():
num_samples = 1000
total = 10
counts = dist.Binomial(total, 0.3).sample()
concentration1 = torch.tensor(0.5)
concentration0 = torch.tensor(1.5)
prior = dist.Beta(concentration1, concentration0)
likelihood = dist.Beta(1 + counts, 1 + total - counts)
posterior = dist.Beta(concentration1 + counts,
concentration0 + total - counts)
def model():
prob = pyro.sample("prob", prior)
pyro.sample("counts", dist.Binomial(total, prob), obs=counts)
reparam_model = poutine.reparam(model,
{"prob": ConjugateReparam(likelihood)})
kernel = HMC(reparam_model)
samples = MCMC(kernel, num_samples, warmup_steps=0).run()
pred = Predictive(reparam_model, samples, num_samples=num_samples)
trace = pred.get_vectorized_trace()
samples = trace.nodes["prob"]["value"]
assert_close(samples.mean(), posterior.mean, atol=0.01)
assert_close(samples.std(), posterior.variance.sqrt(), atol=0.01)
def nuts_sampling(self, x_data, y_data, num_samples, warmup_steps):
nuts_kernel = NUTS(self.model, target_accept_prob=0.99)
mcmc = MCMC(nuts_kernel,
num_samples=num_samples,
warmup_steps=warmup_steps)
mcmc.run(x_data, y_data)
self.posterior_samples = mcmc.get_samples()
def test_mcmc_interface(num_draws, group_by_chain, num_chains):
num_samples = 2000
data = torch.tensor([1.0])
initial_params, _, transforms, _ = initialize_model(normal_normal_model, model_args=(data,),
num_chains=num_chains)
kernel = PriorKernel(normal_normal_model)
mcmc = MCMC(kernel=kernel, num_samples=num_samples, warmup_steps=100, num_chains=num_chains,
mp_context="spawn", initial_params=initial_params, transforms=transforms)
mcmc.run(data)
samples = mcmc.get_samples(num_draws, group_by_chain=group_by_chain)
# test sample shape
expected_samples = num_draws if num_draws is not None else num_samples
if group_by_chain:
expected_shape = (mcmc.num_chains, expected_samples, 1)
elif num_draws is not None:
# FIXME: what is the expected behavior of num_draw is not None and group_by_chain=False?
expected_shape = (expected_samples, 1)
else:
expected_shape = (mcmc.num_chains * expected_samples, 1)
assert samples['y'].shape == expected_shape
# test sample stats
if group_by_chain:
samples = {k: v.reshape((-1,) + v.shape[2:]) for k, v in samples.items()}
sample_mean = samples['y'].mean()
sample_std = samples['y'].std()
assert_close(sample_mean, torch.tensor(0.0), atol=0.05)
assert_close(sample_std, torch.tensor(1.0), atol=0.05)
def test_bernoulli_latent_model(jit):
def model(data):
y_prob = pyro.sample("y_prob", dist.Beta(1.0, 1.0))
y = pyro.sample("y", dist.Bernoulli(y_prob))
with pyro.plate("data", data.shape[0]):
z = pyro.sample("z", dist.Bernoulli(0.65 * y + 0.1))
pyro.sample("obs", dist.Normal(2.0 * z, 1.0), obs=data)
pyro.sample("nuisance", dist.Bernoulli(0.3))
N = 2000
y_prob = torch.tensor(0.3)
y = dist.Bernoulli(y_prob).sample(torch.Size((N, )))
z = dist.Bernoulli(0.65 * y + 0.1).sample()
data = dist.Normal(2.0 * z, 1.0).sample()
hmc_kernel = HMC(
model,
trajectory_length=1,
max_plate_nesting=1,
jit_compile=jit,
ignore_jit_warnings=True,
)
mcmc = MCMC(hmc_kernel, num_samples=600, warmup_steps=200)
mcmc.run(data)
samples = mcmc.get_samples()
assert_equal(samples["y_prob"].mean(0), y_prob, prec=0.06)
def test_num_chains(num_chains, cpu_count, default_init_params, monkeypatch):
monkeypatch.setattr(torch.multiprocessing, "cpu_count", lambda: cpu_count)
data = torch.tensor([1.0])
initial_params, _, transforms, _ = initialize_model(normal_normal_model,
model_args=(data, ),
num_chains=num_chains)
if default_init_params:
initial_params = None
kernel = PriorKernel(normal_normal_model)
available_cpu = max(1, cpu_count - 1)
mp_context = "spawn"
with optional(pytest.warns(UserWarning), available_cpu < num_chains):
mcmc = MCMC(
kernel,
num_samples=10,
warmup_steps=10,
num_chains=num_chains,
initial_params=initial_params,
transforms=transforms,
mp_context=mp_context,
)
mcmc.run(data)
assert mcmc.num_chains == num_chains
if mcmc.num_chains == 1 or available_cpu < num_chains:
assert isinstance(mcmc.sampler, _UnarySampler)
else:
assert isinstance(mcmc.sampler, _MultiSampler)
def test_logistic_regression(step_size, trajectory_length, num_steps,
adapt_step_size, adapt_mass_matrix, full_mass):
dim = 3
data = torch.randn(2000, dim)
true_coefs = torch.arange(1., dim + 1.)
labels = dist.Bernoulli(logits=(true_coefs * data).sum(-1)).sample()
def model(data):
coefs_mean = pyro.param('coefs_mean', torch.zeros(dim))
coefs = pyro.sample('beta', dist.Normal(coefs_mean, torch.ones(dim)))
y = pyro.sample('y',
dist.Bernoulli(logits=(coefs * data).sum(-1)),
obs=labels)
return y
hmc_kernel = HMC(model,
step_size=step_size,
trajectory_length=trajectory_length,
num_steps=num_steps,
adapt_step_size=adapt_step_size,
adapt_mass_matrix=adapt_mass_matrix,
full_mass=full_mass)
mcmc = MCMC(hmc_kernel,
num_samples=500,
warmup_steps=100,
disable_progbar=True)
mcmc.run(data)
samples = mcmc.get_samples()['beta']
assert_equal(rmse(true_coefs, samples.mean(0)).item(), 0.0, prec=0.1)
def test_nuts_conjugate_gaussian(fixture, num_samples, warmup_steps,
expected_means, expected_precs, mean_tol,
std_tol):
pyro.get_param_store().clear()
nuts_kernel = NUTS(fixture.model)
mcmc = MCMC(nuts_kernel, num_samples, warmup_steps)
mcmc.run(fixture.data)
samples = mcmc.get_samples()
for i in range(1, fixture.chain_len + 1):
param_name = 'loc_' + str(i)
latent = samples[param_name]
latent_loc = latent.mean(0)
latent_std = latent.std(0)
expected_mean = torch.ones(fixture.dim) * expected_means[i - 1]
expected_std = 1 / torch.sqrt(
torch.ones(fixture.dim) * expected_precs[i - 1])
# Actual vs expected posterior means for the latents
logger.debug('Posterior mean (actual) - {}'.format(param_name))
logger.debug(latent_loc)
logger.debug('Posterior mean (expected) - {}'.format(param_name))
logger.debug(expected_mean)
assert_equal(rmse(latent_loc, expected_mean).item(),
0.0,
prec=mean_tol)
# Actual vs expected posterior precisions for the latents
logger.debug('Posterior std (actual) - {}'.format(param_name))
logger.debug(latent_std)
logger.debug('Posterior std (expected) - {}'.format(param_name))
logger.debug(expected_std)
assert_equal(rmse(latent_std, expected_std).item(), 0.0, prec=std_tol)
def test_gaussian_mixture_model(jit):
K, N = 3, 1000
def gmm(data):
mix_proportions = pyro.sample("phi", dist.Dirichlet(torch.ones(K)))
with pyro.plate("num_clusters", K):
cluster_means = pyro.sample(
"cluster_means", dist.Normal(torch.arange(float(K)), 1.))
with pyro.plate("data", data.shape[0]):
assignments = pyro.sample("assignments",
dist.Categorical(mix_proportions))
pyro.sample("obs",
dist.Normal(cluster_means[assignments], 1.),
obs=data)
return cluster_means
true_cluster_means = torch.tensor([1., 5., 10.])
true_mix_proportions = torch.tensor([0.1, 0.3, 0.6])
cluster_assignments = dist.Categorical(true_mix_proportions).sample(
torch.Size((N, )))
data = dist.Normal(true_cluster_means[cluster_assignments], 1.0).sample()
nuts_kernel = NUTS(gmm,
max_plate_nesting=1,
jit_compile=jit,
ignore_jit_warnings=True)
mcmc = MCMC(nuts_kernel, num_samples=300, warmup_steps=100)
mcmc.run(data)
samples = mcmc.get_samples()
assert_equal(samples["phi"].mean(0).sort()[0],
true_mix_proportions,
prec=0.05)
assert_equal(samples["cluster_means"].mean(0).sort()[0],
true_cluster_means,
prec=0.2)
def test_beta_binomial(hyperpriors):
def model(data):
with pyro.plate("plate_0", data.shape[-1]):
alpha = pyro.sample(
"alpha", dist.HalfCauchy(1.)) if hyperpriors else torch.tensor(
[1., 1.])
beta = pyro.sample(
"beta", dist.HalfCauchy(1.)) if hyperpriors else torch.tensor(
[1., 1.])
beta_binom = BetaBinomialPair()
with pyro.plate("plate_1", data.shape[-2]):
probs = pyro.sample("probs", beta_binom.latent(alpha, beta))
with pyro.plate("data", data.shape[0]):
pyro.sample("binomial",
beta_binom.conditional(
probs=probs, total_count=total_count),
obs=data)
true_probs = torch.tensor([[0.7, 0.4], [0.6, 0.4]])
total_count = torch.tensor([[1000, 600], [400, 800]])
num_samples = 80
data = dist.Binomial(
total_count=total_count,
probs=true_probs).sample(sample_shape=(torch.Size((10, ))))
hmc_kernel = NUTS(collapse_conjugate(model),
jit_compile=True,
ignore_jit_warnings=True)
mcmc = MCMC(hmc_kernel, num_samples=num_samples, warmup_steps=50)
mcmc.run(data)
samples = mcmc.get_samples()
posterior = posterior_replay(model, samples, data, num_samples=num_samples)
assert_equal(posterior["probs"].mean(0), true_probs, prec=0.05)
def run_default_mcmc(
data,
kernel,
num_samples,
warmup_steps=None,
initial_params=None,
num_chains=1,
hook_fn=None,
mp_context=None,
transforms=None,
num_draws=None,
group_by_chain=False,
):
mcmc = MCMC(
kernel=kernel,
num_samples=num_samples,
warmup_steps=warmup_steps,
initial_params=initial_params,
num_chains=num_chains,
hook_fn=hook_fn,
mp_context=mp_context,
transforms=transforms,
)
mcmc.run(data)
return mcmc.get_samples(num_draws,
group_by_chain=group_by_chain), mcmc.num_chains
def test_sequential_consistent(monkeypatch):
# test if there is no stuff left from the previous chain
monkeypatch.setattr(torch.multiprocessing, 'cpu_count', lambda: 1)
class FirstKernel(NUTS):
def setup(self, warmup_steps, *args, **kwargs):
self._chain_id = 0 if '_chain_id' not in self.__dict__ else 1
pyro.set_rng_seed(self._chain_id)
super(FirstKernel, self).setup(warmup_steps, *args, **kwargs)
class SecondKernel(NUTS):
def setup(self, warmup_steps, *args, **kwargs):
self._chain_id = 1 if '_chain_id' not in self.__dict__ else 0
pyro.set_rng_seed(self._chain_id)
super(SecondKernel, self).setup(warmup_steps, *args, **kwargs)
data = torch.tensor([1.0])
kernel = FirstKernel(normal_normal_model)
mcmc = MCMC(kernel, num_samples=100, warmup_steps=100, num_chains=2)
mcmc.run(data)
samples1 = mcmc.get_samples(group_by_chain=True)
kernel = SecondKernel(normal_normal_model)
mcmc = MCMC(kernel, num_samples=100, warmup_steps=100, num_chains=2)
mcmc.run(data)
samples2 = mcmc.get_samples(group_by_chain=True)
assert_close(samples1["y"][0], samples2["y"][1])
assert_close(samples1["y"][1], samples2["y"][0])
def nuts(data, model, iter=None, warmup=None, num_chains=None):
assert type(data) == dict
assert type(model) == Model
iter, warmup, num_chains = apply_default_hmc_args(iter, warmup, num_chains)
nuts_kernel = NUTS(model.fn, jit_compile=False, adapt_step_size=True)
mcmc = MCMC(nuts_kernel,
num_samples=iter,
warmup_steps=warmup,
num_chains=num_chains)
mcmc.run(**data)
samples = mcmc.get_samples()
transformed_samples = run_model_on_samples_and_data(
model.fn, samples, data)
def loc(d):
# Optimization: For the data used for inference, values for
# `mu` are already computed and available from
# `transformed_samples`.
if d == data:
return transformed_samples['mu']
else:
# TODO: This computes more than is necessary. (i.e. It
# build additional tensors we immediately throw away.)
# This is minor, but might be worth addressing eventually.
return run_model_on_samples_and_data(model.fn, samples, d)['mu']
all_samples = dict(samples, **transformed_samples)
return Samples(all_samples, lambda name: all_samples[name], loc)
def pyro_mcmc(
self,
num_samples: int,
potential_function: Callable,
context: Tensor,
mcmc_method: str = "slice",
thin: int = 10,
warmup_steps: int = 200,
num_chains: Optional[int] = 1,
):
"""Return samples obtained using Pyro's HMC, NUTS or slice kernels.
Args:
num_samples: desired number of samples
potential_fn: defining the potential function as a callable **class** makes
it picklable for Pyro's MCMC to use it across chains in parallel, even
if the potential function requires evaluating a neural network. context:
conditioning observation
mcmc_method: one of "hmc", "nuts" or "slice" (default "slice")
thin: thinning (subsampling) factor (default 10)
warmup_steps: initial samples to discard (defaults to 200)
num_chains: whether to sample in parallel. If None, will use all
CPUs except one (default 1)
Returns: tensor of shape num_samples x parameter dimension
"""
if num_chains is None:
num_chains = mp.cpu_count - 1
# XXX move outside function, and assert inside; remember return to train
# Always sample in eval mode.
self.neural_net.eval()
kernels = dict(slice=Slice, hmc=HMC, nuts=NUTS)
try:
kernel = kernels[mcmc_method](potential_fn=potential_function)
except KeyError:
raise ValueError("`mcmc_method` not one of 'slice', 'hmc', 'nuts'.")
initial_params = self._prior.sample((num_chains,))
sampler = MCMC(
kernel=kernel,
num_samples=(thin * num_samples) // num_chains + num_chains,
warmup_steps=warmup_steps,
initial_params={"": initial_params},
num_chains=num_chains,
mp_context="fork",
)
sampler.run()
samples = next(iter(sampler.get_samples().values())).reshape(
-1, len(self._prior.mean) # len(prior.mean) = dim of theta
)
samples = samples[::thin][:num_samples]
assert samples.shape[0] == num_samples
return samples
def test_gaussian_hmm(num_steps):
dim = 4
def model(data):
initialize = pyro.sample("initialize", dist.Dirichlet(torch.ones(dim)))
with pyro.plate("states", dim):
transition = pyro.sample("transition",
dist.Dirichlet(torch.ones(dim, dim)))
emission_loc = pyro.sample(
"emission_loc", dist.Normal(torch.zeros(dim), torch.ones(dim)))
emission_scale = pyro.sample(
"emission_scale",
dist.LogNormal(torch.zeros(dim), torch.ones(dim)))
x = None
with ignore_jit_warnings([("Iterating over a tensor", RuntimeWarning)
]):
for t, y in pyro.markov(enumerate(data)):
x = pyro.sample(
"x_{}".format(t),
dist.Categorical(
initialize if x is None else transition[x]),
infer={"enumerate": "parallel"})
pyro.sample("y_{}".format(t),
dist.Normal(emission_loc[x], emission_scale[x]),
obs=y)
def _get_initial_trace():
guide = AutoDelta(
poutine.block(model,
expose_fn=lambda msg: not msg["name"].startswith("x")
and not msg["name"].startswith("y")))
elbo = TraceEnum_ELBO(max_plate_nesting=1)
svi = SVI(model, guide, optim.Adam({"lr": .01}), elbo)
for _ in range(100):
svi.step(data)
return poutine.trace(guide).get_trace(data)
def _generate_data():
transition_probs = torch.rand(dim, dim)
emissions_loc = torch.arange(dim, dtype=torch.Tensor().dtype)
emissions_scale = 1.
state = torch.tensor(1)
obs = [dist.Normal(emissions_loc[state], emissions_scale).sample()]
for _ in range(num_steps):
state = dist.Categorical(transition_probs[state]).sample()
obs.append(
dist.Normal(emissions_loc[state], emissions_scale).sample())
return torch.stack(obs)
data = _generate_data()
nuts_kernel = NUTS(model,
max_plate_nesting=1,
jit_compile=True,
ignore_jit_warnings=True)
if num_steps == 30:
nuts_kernel.initial_trace = _get_initial_trace()
mcmc = MCMC(nuts_kernel, num_samples=5, warmup_steps=5)
mcmc.run(data)
def test_model_with_potential_fn():
init_params = {"z": torch.tensor(0.)}
def potential_fn(params):
return params["z"]
mcmc = MCMC(kernel=HMC(potential_fn=potential_fn),
num_samples=10,
warmup_steps=10,
initial_params=init_params)
mcmc.run()
def _pyro_mcmc(
self,
num_samples: int,
potential_function: Callable,
mcmc_method: str = "slice",
thin: int = 10,
warmup_steps: int = 200,
num_chains: Optional[int] = 1,
show_progress_bars: bool = True,
):
r"""Return samples obtained using Pyro HMC, NUTS or slice kernels.
Args:
num_samples: Desired number of samples.
potential_function: A callable **class**. A class, but not a function,
is picklable for Pyro MCMC to use it across chains in parallel,
even when the potential function requires evaluating a neural network.
mcmc_method: One of `hmc`, `nuts` or `slice`.
thin: Thinning (subsampling) factor.
warmup_steps: Initial number of samples to discard.
num_chains: Whether to sample in parallel. If None, use all but one CPU.
show_progress_bars: Whether to show a progressbar during sampling.
Returns: Tensor of shape (num_samples, shape_of_single_theta).
"""
num_chains = mp.cpu_count - 1 if num_chains is None else num_chains
# Go into eval mode for evaluating during sampling
self.net.eval()
kernels = dict(slice=Slice, hmc=HMC, nuts=NUTS)
sampler = MCMC(
kernel=kernels[mcmc_method](potential_fn=potential_function),
num_samples=(thin * num_samples) // num_chains + num_chains,
warmup_steps=warmup_steps,
initial_params={"": self._mcmc_init_params},
num_chains=num_chains,
mp_context="fork",
disable_progbar=not show_progress_bars,
)
sampler.run()
samples = next(iter(sampler.get_samples().values())).reshape(
-1,
len(self._prior.mean) # len(prior.mean) = dim of theta
)
samples = samples[::thin][:num_samples]
assert samples.shape[0] == num_samples
# Back to training mode
self.net.train(True)
return samples
def init_guide(self):
nuts_kernel = NUTS(self.model)
mcmc = MCMC(nuts_kernel, num_samples=1, warmup_steps=30)
mcmc.run()
hmc_samples = {k: v.detach().cpu().numpy() for k, v in mcmc.get_samples().items()}
# Get latent coordinates from mcmc #
gp_coord_mcmc = np.array([[hmc_samples[f'latent_coord_{v}_{h}'] for h in range(self.H_dim)] for v in range(self.V_net)]).transpose(0,1,3,2)
self.gp_coord_mcmc = torch.from_numpy(gp_coord_mcmc)
def run_inference(data, gen_model, ode_model, method, iterations=10000, num_particles=1, num_samples=1000, warmup_steps=500, init_scale=0.1,
seed=12, lr=0.5, return_sites="_RETURN"):
torch_data = torch.tensor(data, dtype=torch.float)
if isinstance(ode_model, ForwardSensManualJacobians) or \
isinstance(ode_model, ForwardSensTorchJacobians):
ode_op = ForwardSensOp
elif isinstance(ode_model, AdjointSensManualJacobians) or \
isinstance(ode_model, AdjointSensTorchJacobians):
ode_op = AdjointSensOp
else:
raise ValueError('Unknown sensitivity solver: Use "Forward" or "Adjoint"')
model = gen_model(ode_op, ode_model)
pyro.set_rng_seed(seed)
pyro.clear_param_store()
if method == 'VI':
guide = AutoMultivariateNormal(model, init_scale=init_scale)
optim = AdagradRMSProp({"eta": lr})
if num_particles == 1:
svi = SVI(model, guide, optim, loss=Trace_ELBO())
else:
svi = SVI(model, guide, optim, loss=Trace_ELBO(num_particles=num_particles,
vectorize_particles=True))
loss_trace = []
t0 = timer.time()
for j in range(iterations):
loss = svi.step(torch_data)
loss_trace.append(loss)
if j % 500 == 0:
print("[iteration %04d] loss: %.4f" % (j + 1, np.mean(loss_trace[max(0, j - 1000):j + 1])))
t1 = timer.time()
print('VI time: ', t1 - t0)
predictive = Predictive(model, guide=guide, num_samples=num_samples,
return_sites=return_sites) # "ode_params", "scale",
vb_samples = predictive(torch_data)
return vb_samples
elif method == 'NUTS':
nuts_kernel = NUTS(model, adapt_step_size=True, init_strategy=init_to_median)
# mcmc = MCMC(nuts_kernel, num_samples=iterations, warmup_steps=warmup_steps, num_chains=2)
mcmc = MCMC(nuts_kernel, num_samples=iterations, warmup_steps=warmup_steps, num_chains=1)
t0 = timer.time()
mcmc.run(torch_data)
t1 = timer.time()
print('NUTS time: ', t1 - t0)
hmc_samples = {k: v.detach().cpu().numpy() for k, v in mcmc.get_samples().items()}
return hmc_samples
else:
raise ValueError('Unknown method: Use "NUTS" or "VI"')
def sample_posterior_mcmc(self, num_samples, thin=10):
"""
Samples from posterior for true observation q(theta | x0) using MCMC.
:param num_samples: Number of samples to generate.
:param mcmc_method: Which MCMC method to use ['metropolis-hastings', 'slice', 'hmc', 'nuts']
:param thin: Generate (num_samples * thin) samples in total, then select every
'thin' sample.
:return: torch.Tensor of shape [num_samples, parameter_dim]
"""
# Always sample in eval mode.
self._neural_posterior.eval()
if self._mcmc_method == "slice-np":
self.posterior_sampler.gen(20) # Burn-in for 200 samples
samples = torch.Tensor(self.posterior_sampler.gen(num_samples))
else:
if self._mcmc_method == "slice":
kernel = Slice(potential_function=self._potential_function)
elif self._mcmc_method == "hmc":
kernel = HMC(potential_fn=self._potential_function)
elif self._mcmc_method == "nuts":
kernel = NUTS(potential_fn=self._potential_function)
else:
raise ValueError(
"'mcmc_method' must be one of ['slice', 'hmc', 'nuts']."
)
num_chains = mp.cpu_count() - 1
initial_params = self._prior.sample((num_chains,))
sampler = MCMC(
kernel=kernel,
num_samples=(thin * num_samples) // num_chains + num_chains,
warmup_steps=200,
initial_params={"": initial_params},
num_chains=num_chains,
mp_context="spawn",
)
sampler.run()
samples = next(iter(sampler.get_samples().values())).reshape(
-1, self._simulator.parameter_dim
)
samples = samples[::thin][:num_samples]
assert samples.shape[0] == num_samples
# Back to training mode.
self._neural_posterior.train()
return samples
def sample_posterior(self, num_samples, thin=1):
"""
Samples from posterior for true observation q(theta | x0) ~ q(x0 | theta) p(theta)
using most recent likelihood estimate q(x0 | theta) with MCMC.
:param num_samples: Number of samples to generate.
:param thin: Generate (num_samples * thin) samples in total, then select every
'thin' sample.
:return: torch.Tensor of shape [num_samples, parameter_dim]
"""
# Always sample in eval mode.
self._neural_likelihood.eval()
if self._mcmc_method == "slice-np":
self.posterior_sampler.gen(20)
samples = torch.Tensor(self.posterior_sampler.gen(num_samples))
else:
if self._mcmc_method == "slice":
kernel = Slice(potential_function=self._potential_function)
elif self._mcmc_method == "hmc":
kernel = HMC(potential_fn=self._potential_function)
elif self._mcmc_method == "nuts":
kernel = NUTS(potential_fn=self._potential_function)
else:
raise ValueError(
"'mcmc_method' must be one of ['slice', 'hmc', 'nuts']."
)
num_chains = mp.cpu_count() - 1
# TODO: decide on way to initialize chain
initial_params = self._prior.sample((num_chains,))
sampler = MCMC(
kernel=kernel,
num_samples=num_samples // num_chains + num_chains,
warmup_steps=200,
initial_params={"": initial_params},
num_chains=num_chains,
)
sampler.run()
samples = next(iter(sampler.get_samples().values())).reshape(
-1, self._simulator.parameter_dim
)
samples = samples[:num_samples].to(device)
assert samples.shape[0] == num_samples
# Back to training mode.
self._neural_likelihood.train()
return samples
class bay_reg_mcmc:
def __init__(self,model, num_samples=1000, warmup_steps=100, cuda=False):
self.model=model
self.num_samples=num_samples
self.warmup_steps=warmup_steps
self.cuda = cuda
def fit(self,x_train,y_train, verbose=False):
if verbose:
disable_progbar=False
else:
disable_progbar=True
if self.cuda == True:
x_train=torch.tensor(x_train).cuda()
y_train=torch.tensor(y_train).cuda()
torch.set_default_tensor_type(torch.cuda.FloatTensor)
else:
x_train=torch.tensor(x_train)
y_train=torch.tensor(y_train)
self.posterior=MCMC(kernel=NUTS(self.model), num_samples=self.num_samples, warmup_steps=self.warmup_steps, disable_progbar=disable_progbar)
self.posterior=self.posterior
self.posterior.run(x_train,y_train)
def predict(self, x_test, num_samples=300, percentiles=(5.0, 50.0, 95.0)):
x_test=torch.tensor(x_test)
samples=self.posterior.get_samples()
#posterior_predictive=predictive(self.model, samples, x_test, num_samples=num_samples,return_sites=("w","b","_RETURN"))
posterior_predictive= Predictive(self.model, posterior_samples=samples, return_sites=("_RETURN",'obs')).forward(x_test)
#confidence intervall
convidence_intervalls=posterior_predictive['_RETURN'].detach().cpu().numpy()
y_pred=np.percentile(convidence_intervalls, percentiles, axis=0).T
y_median_conv=y_pred[:,1].reshape(-1,1)
y_lower_upper_quantil_conv=np.concatenate((y_pred[:,1].reshape(-1,1)-y_pred[:,0].reshape(-1,1),y_pred[:,2].reshape(-1,1)-y_pred[:,1].reshape(-1,1)),axis=1)
#predictive intervall
predictive_intervalls=posterior_predictive['obs'].detach().cpu().numpy()
y_pred=np.percentile(predictive_intervalls, percentiles, axis=0).T
y_median_pred=y_pred[:,1].reshape(-1,1)
y_lower_upper_quantil_pred=np.concatenate((y_pred[:,1].reshape(-1,1)-y_pred[:,0].reshape(-1,1),y_pred[:,2].reshape(-1,1)-y_pred[:,1].reshape(-1,1)),axis=1)
return(y_median_conv, y_lower_upper_quantil_conv, y_median_pred, y_lower_upper_quantil_pred)
def test_():
# if torch.cuda.is_available():
# device = torch.device("cuda")
# torch.set_default_tensor_type("torch.cuda.FloatTensor")
# else:
# input("CUDA not available, do you wish to continue?")
# device = torch.device("cpu")
# torch.set_default_tensor_type("torch.FloatTensor")
loc = torch.Tensor([0, 0])
covariance_matrix = torch.Tensor([[1, 0.99], [0.99, 1]])
likelihood = distributions.MultivariateNormal(
loc=loc, covariance_matrix=covariance_matrix)
bound = 1.5
low, high = -bound * torch.ones(2), bound * torch.ones(2)
prior = distributions.Uniform(low=low, high=high)
# def potential_function(inputs_dict):
# parameters = next(iter(inputs_dict.values()))
# return -(likelihood.log_prob(parameters) + prior.log_prob(parameters).sum())
prior = distributions.Uniform(low=-5 * torch.ones(4),
high=2 * torch.ones(4))
from nsf import distributions as distributions_
likelihood = distributions_.LotkaVolterraOscillating()
potential_function = PotentialFunction(likelihood, prior)
# kernel = Slice(potential_function=potential_function)
from pyro.infer.mcmc import HMC, NUTS
# kernel = HMC(potential_fn=potential_function)
kernel = NUTS(potential_fn=potential_function)
num_chains = 3
sampler = MCMC(
kernel=kernel,
num_samples=10000 // num_chains,
warmup_steps=200,
initial_params={"": torch.zeros(num_chains, 4)},
num_chains=num_chains,
)
sampler.run()
samples = next(iter(sampler.get_samples().values()))
utils.plot_hist_marginals(utils.tensor2numpy(samples),
ground_truth=utils.tensor2numpy(loc),
lims=[-6, 3])
# plt.show()
plt.savefig("/home/conor/Dropbox/phd/projects/lfi/out/mcmc.pdf")
plt.close()
def test_dirichlet_categorical(jit):
def model(data):
concentration = torch.tensor([1.0, 1.0, 1.0])
p_latent = pyro.sample('p_latent', dist.Dirichlet(concentration))
pyro.sample("obs", dist.Categorical(p_latent), obs=data)
return p_latent
true_probs = torch.tensor([0.1, 0.6, 0.3])
data = dist.Categorical(true_probs).sample(sample_shape=(torch.Size((2000,))))
hmc_kernel = HMC(model, trajectory_length=1, jit_compile=jit, ignore_jit_warnings=True)
mcmc = MCMC(hmc_kernel, num_samples=200, warmup_steps=100)
mcmc.run(data)
samples = mcmc.get_samples()
assert_equal(samples['p_latent'].mean(0), true_probs, prec=0.02)
def test_gamma_normal():
def model(data):
rate = torch.tensor([1.0, 1.0])
concentration = torch.tensor([1.0, 1.0])
p_latent = pyro.sample("p_latent", dist.Gamma(rate, concentration))
pyro.sample("obs", dist.Normal(3, p_latent), obs=data)
return p_latent
true_std = torch.tensor([0.5, 2])
data = dist.Normal(3, true_std).sample(sample_shape=(torch.Size((2000, ))))
hmc_kernel = HMC(model, num_steps=15, step_size=0.01, adapt_step_size=True)
mcmc = MCMC(hmc_kernel, num_samples=200, warmup_steps=200)
mcmc.run(data)
samples = mcmc.get_samples()
assert_equal(samples["p_latent"].mean(0), true_std, prec=0.05)
def test_gamma_beta(jit):
def model(data):
alpha_prior = pyro.sample('alpha', dist.Gamma(concentration=1., rate=1.))
beta_prior = pyro.sample('beta', dist.Gamma(concentration=1., rate=1.))
pyro.sample('x', dist.Beta(concentration1=alpha_prior, concentration0=beta_prior), obs=data)
true_alpha = torch.tensor(5.)
true_beta = torch.tensor(1.)
data = dist.Beta(concentration1=true_alpha, concentration0=true_beta).sample(torch.Size((5000,)))
nuts_kernel = NUTS(model, jit_compile=jit, ignore_jit_warnings=True)
mcmc = MCMC(nuts_kernel, num_samples=500, warmup_steps=200)
mcmc.run(data)
samples = mcmc.get_samples()
assert_equal(samples["alpha"].mean(0), true_alpha, prec=0.08)
assert_equal(samples["beta"].mean(0), true_beta, prec=0.05)
