def test_iaf():
# test for substitute logic for exposed methods `sample_posterior` and `get_transforms`
N, dim = 3000, 3
data = random.normal(random.PRNGKey(0), (N, dim))
true_coefs = np.arange(1., dim + 1.)
logits = np.sum(true_coefs * data, axis=-1)
labels = dist.Bernoulli(logits=logits).sample(random.PRNGKey(1))
def model(data, labels):
coefs = numpyro.sample('coefs', dist.Normal(np.zeros(dim),
np.ones(dim)))
offset = numpyro.sample('offset', dist.Uniform(-1, 1))
logits = offset + np.sum(coefs * data, axis=-1)
return numpyro.sample('obs', dist.Bernoulli(logits=logits), obs=labels)
adam = optim.Adam(0.01)
rng_init = random.PRNGKey(1)
guide = AutoIAFNormal(model)
svi = SVI(model, guide, elbo, adam)
svi_state = svi.init(rng_init,
model_args=(data, labels),
guide_args=(data, labels))
params = svi.get_params(svi_state)
x = random.normal(random.PRNGKey(0), (dim + 1, ))
rng = random.PRNGKey(1)
actual_sample = guide.sample_posterior(rng, params)
actual_output = guide.get_transform(params)(x)
flows = []
for i in range(guide.num_flows):
if i > 0:
flows.append(constraints.PermuteTransform(
np.arange(dim + 1)[::-1]))
arn_init, arn_apply = AutoregressiveNN(
dim + 1, [dim + 1, dim + 1],
permutation=np.arange(dim + 1),
skip_connections=guide._skip_connections,
nonlinearity=guide._nonlinearity)
arn = partial(arn_apply, params['auto_arn__{}$params'.format(i)])
flows.append(InverseAutoregressiveTransform(arn))
transform = constraints.ComposeTransform(flows)
rng_seed, rng_sample = random.split(rng)
expected_sample = guide.unpack_latent(
transform(dist.Normal(np.zeros(dim + 1), 1).sample(rng_sample)))
expected_output = transform(x)
assert_allclose(actual_sample['coefs'], expected_sample['coefs'])
assert_allclose(
actual_sample['offset'],
constraints.biject_to(constraints.interval(-1, 1))(
expected_sample['offset']))
assert_allclose(actual_output, expected_output)
def main(args):
jax_config.update('jax_platform_name', args.device)
print("Start vanilla HMC...")
vanilla_samples = mcmc(args.num_warmup, args.num_samples, init_params=np.array([2., 0.]),
potential_fn=dual_moon_pe, progbar=True)
opt_init, opt_update, get_params = optimizers.adam(0.001)
rng_guide, rng_init, rng_train = random.split(random.PRNGKey(1), 3)
guide = AutoIAFNormal(rng_guide, dual_moon_model, get_params, hidden_dims=[args.num_hidden])
svi_init, svi_update, _ = svi(dual_moon_model, guide, elbo, opt_init, opt_update, get_params)
opt_state, _ = svi_init(rng_init)
def body_fn(val, i):
opt_state_, rng_ = val
loss, opt_state_, rng_ = svi_update(i, rng_, opt_state_)
return (opt_state_, rng_), loss
print("Start training guide...")
(last_state, _), losses = lax.scan(body_fn, (opt_state, rng_train), np.arange(args.num_iters))
print("Finish training guide. Extract samples...")
guide_samples = guide.sample_posterior(random.PRNGKey(0), last_state,
sample_shape=(args.num_samples,))
transform = guide.get_transform(last_state)
unpack_fn = guide.unpack_latent
_, potential_fn, constrain_fn = initialize_model(random.PRNGKey(0), dual_moon_model)
transformed_potential_fn = make_transformed_pe(potential_fn, transform, unpack_fn)
transformed_constrain_fn = lambda x: constrain_fn(unpack_fn(transform(x))) # noqa: E731
init_params = np.zeros(guide.latent_size)
print("\nStart NeuTra HMC...")
zs = mcmc(args.num_warmup, args.num_samples, init_params, potential_fn=transformed_potential_fn)
print("Transform samples into unwarped space...")
samples = vmap(transformed_constrain_fn)(zs)
summary(tree_map(lambda x: x[None, ...], samples))
# make plots
# IAF guide samples (for plotting)
iaf_base_samples = dist.Normal(np.zeros(2), 1.).sample(random.PRNGKey(0), (1000,))
iaf_trans_samples = vmap(transformed_constrain_fn)(iaf_base_samples)['x']
x1 = np.linspace(-3, 3, 100)
x2 = np.linspace(-3, 3, 100)
X1, X2 = np.meshgrid(x1, x2)
P = np.clip(np.exp(-dual_moon_pe(np.stack([X1, X2], axis=-1))), a_min=0.)
fig = plt.figure(figsize=(12, 16), constrained_layout=True)
gs = GridSpec(3, 2, figure=fig)
ax1 = fig.add_subplot(gs[0, 0])
ax2 = fig.add_subplot(gs[0, 1])
ax3 = fig.add_subplot(gs[1, 0])
ax4 = fig.add_subplot(gs[1, 1])
ax5 = fig.add_subplot(gs[2, 0])
ax6 = fig.add_subplot(gs[2, 1])
ax1.plot(np.log(losses[1000:]))
ax1.set_title('Autoguide training log loss (after 1000 steps)')
ax2.contourf(X1, X2, P, cmap='OrRd')
sns.kdeplot(guide_samples['x'][:, 0].copy(), guide_samples['x'][:, 1].copy(), n_levels=30, ax=ax2)
ax2.set(xlim=[-3, 3], ylim=[-3, 3],
xlabel='x0', ylabel='x1', title='Posterior using AutoIAFNormal guide')
sns.scatterplot(iaf_base_samples[:, 0], iaf_base_samples[:, 1], ax=ax3, hue=iaf_trans_samples[:, 0] < 0.)
ax3.set(xlim=[-3, 3], ylim=[-3, 3],
xlabel='x0', ylabel='x1', title='AutoIAFNormal base samples (True=left moon; False=right moon)')
ax4.contourf(X1, X2, P, cmap='OrRd')
sns.kdeplot(vanilla_samples[:, 0].copy(), vanilla_samples[:, 1].copy(), n_levels=30, ax=ax4)
ax4.plot(vanilla_samples[-50:, 0], vanilla_samples[-50:, 1], 'bo-', alpha=0.5)
ax4.set(xlim=[-3, 3], ylim=[-3, 3],
xlabel='x0', ylabel='x1', title='Posterior using vanilla HMC sampler')
sns.scatterplot(zs[:, 0], zs[:, 1], ax=ax5, hue=samples['x'][:, 0] < 0.,
s=30, alpha=0.5, edgecolor="none")
ax5.set(xlim=[-5, 5], ylim=[-5, 5],
xlabel='x0', ylabel='x1', title='Samples from the warped posterior - p(z)')
ax6.contourf(X1, X2, P, cmap='OrRd')
sns.kdeplot(samples['x'][:, 0].copy(), samples['x'][:, 1].copy(), n_levels=30, ax=ax6)
ax6.plot(samples['x'][-50:, 0], samples['x'][-50:, 1], 'bo-', alpha=0.2)
ax6.set(xlim=[-3, 3], ylim=[-3, 3],
xlabel='x0', ylabel='x1', title='Posterior using NeuTra HMC sampler')
plt.savefig("neutra.pdf")
plt.close()
def main(args):
print("Start vanilla HMC...")
nuts_kernel = NUTS(dual_moon_model)
mcmc = MCMC(nuts_kernel, args.num_warmup, args.num_samples)
mcmc.run(random.PRNGKey(0))
mcmc.print_summary()
vanilla_samples = mcmc.get_samples()['x'].copy()
adam = optim.Adam(0.01)
# TODO: it is hard to find good hyperparameters such that IAF guide can learn this model.
# We will use BNAF instead!
guide = AutoIAFNormal(dual_moon_model,
num_flows=2,
hidden_dims=[args.num_hidden, args.num_hidden])
svi = SVI(dual_moon_model, guide, adam, AutoContinuousELBO())
svi_state = svi.init(random.PRNGKey(1))
print("Start training guide...")
last_state, losses = lax.scan(lambda state, i: svi.update(state),
svi_state, np.zeros(args.num_iters))
params = svi.get_params(last_state)
print("Finish training guide. Extract samples...")
guide_samples = guide.sample_posterior(
random.PRNGKey(0), params,
sample_shape=(args.num_samples, ))['x'].copy()
transform = guide.get_transform(params)
_, potential_fn, constrain_fn = initialize_model(random.PRNGKey(2),
dual_moon_model)
transformed_potential_fn = partial(transformed_potential_energy,
potential_fn, transform)
transformed_constrain_fn = lambda x: constrain_fn(transform(x)
) # noqa: E731
print("\nStart NeuTra HMC...")
nuts_kernel = NUTS(potential_fn=transformed_potential_fn)
mcmc = MCMC(nuts_kernel, args.num_warmup, args.num_samples)
init_params = np.zeros(guide.latent_size)
mcmc.run(random.PRNGKey(3), init_params=init_params)
mcmc.print_summary()
zs = mcmc.get_samples()
print("Transform samples into unwarped space...")
samples = vmap(transformed_constrain_fn)(zs)
print_summary(tree_map(lambda x: x[None, ...], samples))
samples = samples['x'].copy()
# make plots
# guide samples (for plotting)
guide_base_samples = dist.Normal(np.zeros(2),
1.).sample(random.PRNGKey(4), (1000, ))
guide_trans_samples = vmap(transformed_constrain_fn)(
guide_base_samples)['x']
x1 = np.linspace(-3, 3, 100)
x2 = np.linspace(-3, 3, 100)
X1, X2 = np.meshgrid(x1, x2)
P = np.exp(DualMoonDistribution().log_prob(np.stack([X1, X2], axis=-1)))
fig = plt.figure(figsize=(12, 16), constrained_layout=True)
gs = GridSpec(3, 2, figure=fig)
ax1 = fig.add_subplot(gs[0, 0])
ax2 = fig.add_subplot(gs[0, 1])
ax3 = fig.add_subplot(gs[1, 0])
ax4 = fig.add_subplot(gs[1, 1])
ax5 = fig.add_subplot(gs[2, 0])
ax6 = fig.add_subplot(gs[2, 1])
ax1.plot(np.log(losses[1000:]))
ax1.set_title('Autoguide training log loss (after 1000 steps)')
ax2.contourf(X1, X2, P, cmap='OrRd')
sns.kdeplot(guide_samples[:, 0], guide_samples[:, 1], n_levels=30, ax=ax2)
ax2.set(xlim=[-3, 3],
ylim=[-3, 3],
xlabel='x0',
ylabel='x1',
title='Posterior using AutoIAFNormal guide')
sns.scatterplot(guide_base_samples[:, 0],
guide_base_samples[:, 1],
ax=ax3,
hue=guide_trans_samples[:, 0] < 0.)
ax3.set(
xlim=[-3, 3],
ylim=[-3, 3],
xlabel='x0',
ylabel='x1',
title='AutoIAFNormal base samples (True=left moon; False=right moon)')
ax4.contourf(X1, X2, P, cmap='OrRd')
sns.kdeplot(vanilla_samples[:, 0],
vanilla_samples[:, 1],
n_levels=30,
ax=ax4)
ax4.plot(vanilla_samples[-50:, 0],
vanilla_samples[-50:, 1],
'bo-',
alpha=0.5)
ax4.set(xlim=[-3, 3],
ylim=[-3, 3],
xlabel='x0',
ylabel='x1',
title='Posterior using vanilla HMC sampler')
sns.scatterplot(zs[:, 0],
zs[:, 1],
ax=ax5,
hue=samples[:, 0] < 0.,
s=30,
alpha=0.5,
edgecolor="none")
ax5.set(xlim=[-5, 5],
ylim=[-5, 5],
xlabel='x0',
ylabel='x1',
title='Samples from the warped posterior - p(z)')
ax6.contourf(X1, X2, P, cmap='OrRd')
sns.kdeplot(samples[:, 0], samples[:, 1], n_levels=30, ax=ax6)
ax6.plot(samples[-50:, 0], samples[-50:, 1], 'bo-', alpha=0.2)
ax6.set(xlim=[-3, 3],
ylim=[-3, 3],
xlabel='x0',
ylabel='x1',
title='Posterior using NeuTra HMC sampler')
plt.savefig("neutra.pdf")
plt.close()
def main(args):
jax_config.update('jax_platform_name', args.device)
print("Start vanilla HMC...")
nuts_kernel = NUTS(potential_fn=dual_moon_pe)
mcmc = MCMC(nuts_kernel, args.num_warmup, args.num_samples)
mcmc.run(random.PRNGKey(11), init_params=np.array([2., 0.]))
vanilla_samples = mcmc.get_samples()
adam = optim.Adam(0.001)
rng_init, rng_train = random.split(random.PRNGKey(1), 2)
guide = AutoIAFNormal(dual_moon_model, hidden_dims=[args.num_hidden], skip_connections=True)
svi = SVI(dual_moon_model, guide, elbo, adam)
svi_state = svi.init(rng_init)
print("Start training guide...")
last_state, losses = lax.scan(lambda state, i: svi.update(state), svi_state, np.zeros(args.num_iters))
params = svi.get_params(last_state)
print("Finish training guide. Extract samples...")
guide_samples = guide.sample_posterior(random.PRNGKey(0), params,
sample_shape=(args.num_samples,))
transform = guide.get_transform(params)
unpack_fn = guide.unpack_latent
_, potential_fn, constrain_fn = initialize_model(random.PRNGKey(0), dual_moon_model)
transformed_potential_fn = make_transformed_pe(potential_fn, transform, unpack_fn)
transformed_constrain_fn = lambda x: constrain_fn(unpack_fn(transform(x))) # noqa: E731
init_params = np.zeros(guide.latent_size)
print("\nStart NeuTra HMC...")
# TODO: exlore why neutra samples are not good
# Issue: https://github.com/pyro-ppl/numpyro/issues/256
nuts_kernel = NUTS(potential_fn=transformed_potential_fn)
mcmc = MCMC(nuts_kernel, args.num_warmup, args.num_samples)
mcmc.run(random.PRNGKey(10), init_params=init_params)
zs = mcmc.get_samples()
print("Transform samples into unwarped space...")
samples = vmap(transformed_constrain_fn)(zs)
summary(tree_map(lambda x: x[None, ...], samples))
# make plots
# IAF guide samples (for plotting)
iaf_base_samples = dist.Normal(np.zeros(2), 1.).sample(random.PRNGKey(0), (1000,))
iaf_trans_samples = vmap(transformed_constrain_fn)(iaf_base_samples)['x']
x1 = np.linspace(-3, 3, 100)
x2 = np.linspace(-3, 3, 100)
X1, X2 = np.meshgrid(x1, x2)
P = np.clip(np.exp(-dual_moon_pe(np.stack([X1, X2], axis=-1))), a_min=0.)
fig = plt.figure(figsize=(12, 16), constrained_layout=True)
gs = GridSpec(3, 2, figure=fig)
ax1 = fig.add_subplot(gs[0, 0])
ax2 = fig.add_subplot(gs[0, 1])
ax3 = fig.add_subplot(gs[1, 0])
ax4 = fig.add_subplot(gs[1, 1])
ax5 = fig.add_subplot(gs[2, 0])
ax6 = fig.add_subplot(gs[2, 1])
ax1.plot(np.log(losses[1000:]))
ax1.set_title('Autoguide training log loss (after 1000 steps)')
ax2.contourf(X1, X2, P, cmap='OrRd')
sns.kdeplot(guide_samples['x'][:, 0].copy(), guide_samples['x'][:, 1].copy(), n_levels=30, ax=ax2)
ax2.set(xlim=[-3, 3], ylim=[-3, 3],
xlabel='x0', ylabel='x1', title='Posterior using AutoIAFNormal guide')
sns.scatterplot(iaf_base_samples[:, 0], iaf_base_samples[:, 1], ax=ax3, hue=iaf_trans_samples[:, 0] < 0.)
ax3.set(xlim=[-3, 3], ylim=[-3, 3],
xlabel='x0', ylabel='x1', title='AutoIAFNormal base samples (True=left moon; False=right moon)')
ax4.contourf(X1, X2, P, cmap='OrRd')
sns.kdeplot(vanilla_samples[:, 0].copy(), vanilla_samples[:, 1].copy(), n_levels=30, ax=ax4)
ax4.plot(vanilla_samples[-50:, 0], vanilla_samples[-50:, 1], 'bo-', alpha=0.5)
ax4.set(xlim=[-3, 3], ylim=[-3, 3],
xlabel='x0', ylabel='x1', title='Posterior using vanilla HMC sampler')
sns.scatterplot(zs[:, 0], zs[:, 1], ax=ax5, hue=samples['x'][:, 0] < 0.,
s=30, alpha=0.5, edgecolor="none")
ax5.set(xlim=[-5, 5], ylim=[-5, 5],
xlabel='x0', ylabel='x1', title='Samples from the warped posterior - p(z)')
ax6.contourf(X1, X2, P, cmap='OrRd')
sns.kdeplot(samples['x'][:, 0].copy(), samples['x'][:, 1].copy(), n_levels=30, ax=ax6)
ax6.plot(samples['x'][-50:, 0], samples['x'][-50:, 1], 'bo-', alpha=0.2)
ax6.set(xlim=[-3, 3], ylim=[-3, 3],
xlabel='x0', ylabel='x1', title='Posterior using NeuTra HMC sampler')
plt.savefig("neutra.pdf")
plt.close()