def mcmc(num_warmup,
num_samples,
init_params,
sampler='hmc',
constrain_fn=None,
print_summary=True,
**sampler_kwargs):
"""
Convenience wrapper for MCMC samplers -- runs warmup, prints
diagnostic summary and returns a collections of samples
from the posterior.
:param num_warmup: Number of warmup steps.
:param num_samples: Number of samples to generate from the Markov chain.
:param init_params: Initial parameters to begin sampling. The type can
must be consistent with the input type to `potential_fn`.
:param sampler: currently, only `hmc` is implemented (default).
:param constrain_fn: Callable that converts a collection of unconstrained
sample values returned from the sampler to constrained values that
lie within the support of the sample sites.
:param print_summary: Whether to print diagnostics summary for
each sample site. Default is ``True``.
:param `**sampler_kwargs`: Sampler specific keyword arguments.
- *HMC*: Refer to :func:`~numpyro.mcmc.hmc` and
:func:`~numpyro.mcmc.hmc.init_kernel` for accepted arguments. Note
that all arguments must be provided as keywords.
:return: collection of samples from the posterior.
"""
if sampler == 'hmc':
if constrain_fn is None:
constrain_fn = identity
potential_fn = sampler_kwargs.pop('potential_fn')
kinetic_fn = sampler_kwargs.pop('kinetic_fn', None)
algo = sampler_kwargs.pop('algo', 'NUTS')
progbar = sampler_kwargs.pop('progbar', True)
init_kernel, sample_kernel = hmc(potential_fn, kinetic_fn, algo)
hmc_state = init_kernel(init_params,
num_warmup,
progbar=progbar,
**sampler_kwargs)
samples = fori_collect(num_samples,
sample_kernel,
hmc_state,
transform=lambda x: constrain_fn(x.z),
progbar=progbar,
diagnostics_fn=get_diagnostics_str,
progbar_desc='sample')
if print_summary:
summary(samples)
return samples
else:
raise ValueError('sampler: {} not recognized'.format(sampler))
def run_diagnostics(mcmc):
"""Extract diagnostic metrics from a fitted MCMC model: the minimum effective sample size
and the maximum Gelman-Rubin test value.
Args:
mcmc: a fitted MCMC object
Returns:
metrics: a dictionary of the metrics
"""
if not isinstance(mcmc, MCMC):
raise ValueError(
"The argument to run_diagnostics() must be a fitted MCMC object")
summary_dict = summary(mcmc._states["z"])
min_ess = np.inf
max_rhat = 0
for _, stats_dict in summary_dict.items():
min_ess = min(min_ess, stats_dict["n_eff"].min())
max_rhat = max(max_rhat, stats_dict["r_hat"].max())
metrics = {"min_ess": min_ess, "max_rhat": max_rhat}
return metrics
def main(args):
jax_config.update('jax_platform_name', args.device)
print('Simulating data...')
(transition_prior, emission_prior, transition_prob, emission_prob,
supervised_categories, supervised_words,
unsupervised_words) = simulate_data(
random.PRNGKey(1),
num_categories=args.num_categories,
num_words=args.num_words,
num_supervised_data=args.num_supervised,
num_unsupervised_data=args.num_unsupervised,
)
print('Starting inference...')
zs = run_inference(transition_prior, emission_prior, supervised_categories,
supervised_words, unsupervised_words, random.PRNGKey(2),
args)
summary(zs)
print_results(zs, transition_prob, emission_prob)
def __init__(self, mcmc, priors):
from numpyro.diagnostics import print_summary, summary
from operator import attrgetter
""" Class for dealing with posterior from numpyro
:param fit: fit object from numpyro
:param priors: list of prior classes used for fit
"""
self.nsrc=priors[0].nsrc
self.samples=mcmc.get_samples()
self.samples['src_f']=np.swapaxes(self.samples['src_f'],1,2)
# get summary statistics. Code based on numpyro print_summary
prob = 0.9
exclude_deterministic = True
sites = mcmc._states[mcmc._sample_field]
if isinstance(sites, dict) and exclude_deterministic:
state_sample_field = attrgetter(mcmc._sample_field)(mcmc._last_state)
# XXX: there might be the case that state.z is not a dictionary but
# its postprocessed value `sites` is a dictionary.
# TODO: in general, when both `sites` and `state.z` are dictionaries,
# they can have different key names, not necessary due to deterministic
# behavior. We might revise this logic if needed in the future.
if isinstance(state_sample_field, dict):
sites = {k: v for k, v in mcmc._states[mcmc._sample_field].items()
if k in state_sample_field}
stats_summary = summary(sites, prob=prob)
diverge = mcmc.get_extra_fields()['diverging']
self.Rhat = {'src_f': stats_summary['src_f']['r_hat'],
'sigma_conf': stats_summary['sigma_conf']['r_hat'],
'bkg': stats_summary['bkg']['r_hat']}
self.n_eff = {'src_f': stats_summary['src_f']['n_eff'],
'sigma_conf': stats_summary['sigma_conf']['n_eff'],
'bkg': stats_summary['bkg']['n_eff']}
self.divergences=diverge
print("Number of divergences: {}".format(np.sum(diverge)))
if len(priors) < 2:
self.samples['bkg']=self.samples['bkg'][:,None]
self.samples['sigma_conf'] = self.samples['sigma_conf'][:, None]
def __init__(self, mcmc, priors, sed_prior):
from numpyro.diagnostics import summary
import jax.numpy as jnp
from operator import attrgetter
""" Class for dealing with posterior from numpyro
:param fit: fit object from numpyro
:param priors: list of prior classes used for fit
"""
self.nsrc = priors[0].nsrc
self.samples = mcmc.get_samples()
self.samples['src_f'] = jnp.power(
10.0, sed_prior.emulator['net_apply'](sed_prior.emulator['params'],
self.samples['params']))
#self.samples['src_f']=jnp.power(10.0,sed_prior.emulator['net_apply'](sed_prior.emulator['params'],jnp.vstack((self.samples['sfr'][None,:],self.samples['agn_frac'][None,:],self.samples['redshift'][None,:])).T))
self.samples['src_f'] = np.swapaxes(self.samples['src_f'], 1, 2)
self.samples['sigma_conf'] = np.zeros_like(self.samples['bkg'])
# get summary statistics. Code based on numpyro print_summary
prob = 0.9
exclude_deterministic = True
sites = mcmc._states[mcmc._sample_field]
if isinstance(sites, dict) and exclude_deterministic:
state_sample_field = attrgetter(mcmc._sample_field)(
mcmc._last_state)
# XXX: there might be the case that state.z is not a dictionary but
# its postprocessed value `sites` is a dictionary.
# TODO: in general, when both `sites` and `state.z` are dictionaries,
# they can have different key names, not necessary due to deterministic
# behavior. We might revise this logic if needed in the future.
if isinstance(state_sample_field, dict):
sites = {
k: v
for k, v in mcmc._states[mcmc._sample_field].items()
if k in state_sample_field
}
stats_summary = summary(sites, prob=prob)
diverge = mcmc.get_extra_fields()['diverging']
self.Rhat = [stats_summary[i]['r_hat'] for i in stats_summary.keys()]
self.n_eff = [stats_summary[i]['n_eff'] for i in stats_summary.keys()]
self.divergences = diverge
print("Number of divergences: {}".format(np.sum(diverge)))
def run_inference(model, args, rng_key, X, Y):
start = time.time()
kernel = NUTS(model)
mcmc = MCMC(
kernel,
num_warmup=args.num_warmup,
num_samples=args.num_samples,
num_chains=args.num_chains,
progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True,
)
mcmc.run(rng_key, X, Y)
mcmc.print_summary(exclude_deterministic=False)
samples = mcmc.get_samples()
summary_dict = summary(samples, group_by_chain=False)
print("\nMCMC elapsed time:", time.time() - start)
return summary_dict
def mcmc(num_warmup, num_samples, init_params, num_chains=1, sampler='hmc',
constrain_fn=None, print_summary=True, **sampler_kwargs):
"""
Convenience wrapper for MCMC samplers -- runs warmup, prints
diagnostic summary and returns a collections of samples
from the posterior.
:param num_warmup: Number of warmup steps.
:param num_samples: Number of samples to generate from the Markov chain.
:param init_params: Initial parameters to begin sampling. The type can
must be consistent with the input type to `potential_fn`.
:param sampler: currently, only `hmc` is implemented (default).
:param constrain_fn: Callable that converts a collection of unconstrained
sample values returned from the sampler to constrained values that
lie within the support of the sample sites.
:param print_summary: Whether to print diagnostics summary for
each sample site. Default is ``True``.
:param `**sampler_kwargs`: Sampler specific keyword arguments.
- *HMC*: Refer to :func:`~numpyro.mcmc.hmc` and
:func:`~numpyro.mcmc.hmc.init_kernel` for accepted arguments. Note
that all arguments must be provided as keywords.
:return: collection of samples from the posterior.
.. testsetup::
import jax
from jax import random
import jax.numpy as np
import numpyro.distributions as dist
from numpyro.handlers import sample
from numpyro.hmc_util import initialize_model
from numpyro.mcmc import hmc
from numpyro.util import fori_collect
.. doctest::
>>> true_coefs = np.array([1., 2., 3.])
>>> data = random.normal(random.PRNGKey(2), (2000, 3))
>>> dim = 3
>>> labels = dist.Bernoulli(logits=(true_coefs * data).sum(-1)).sample(random.PRNGKey(3))
>>>
>>> def model(data, labels):
... coefs_mean = np.zeros(dim)
... coefs = sample('beta', dist.Normal(coefs_mean, np.ones(3)))
... intercept = sample('intercept', dist.Normal(0., 10.))
... return sample('y', dist.Bernoulli(logits=(coefs * data + intercept).sum(-1)), obs=labels)
>>>
>>> init_params, potential_fn, constrain_fn = initialize_model(random.PRNGKey(0), model,
... data, labels)
>>> num_warmup, num_samples = 1000, 1000
>>> samples = mcmc(num_warmup, num_samples, init_params,
... potential_fn=potential_fn,
... constrain_fn=constrain_fn) # doctest: +SKIP
warmup: 100%|██████████| 1000/1000 [00:09<00:00, 109.40it/s, 1 steps of size 5.83e-01. acc. prob=0.79]
sample: 100%|██████████| 1000/1000 [00:00<00:00, 1252.39it/s, 1 steps of size 5.83e-01. acc. prob=0.85]
mean sd 5.5% 94.5% n_eff Rhat
coefs[0] 0.96 0.07 0.85 1.07 455.35 1.01
coefs[1] 2.05 0.09 1.91 2.20 332.00 1.01
coefs[2] 3.18 0.13 2.96 3.37 320.27 1.00
intercept -0.03 0.02 -0.06 0.00 402.53 1.00
"""
sequential_chain = False
if xla_bridge.device_count() < num_chains:
sequential_chain = True
warnings.warn('There are not enough devices to run parallel chains: expected {} but got {}.'
' Chains will be drawn sequentially. If you are running `mcmc` in CPU,'
' consider to disable XLA intra-op parallelism by setting the environment'
' flag "XLA_FLAGS=--xla_force_host_platform_device_count={}".'
.format(num_chains, xla_bridge.device_count(), num_chains))
progbar = sampler_kwargs.pop('progbar', True)
if num_chains > 1:
progbar = False
if sampler == 'hmc':
if constrain_fn is None:
constrain_fn = identity
potential_fn = sampler_kwargs.pop('potential_fn')
kinetic_fn = sampler_kwargs.pop('kinetic_fn', None)
algo = sampler_kwargs.pop('algo', 'NUTS')
if num_chains > 1:
rngs = sampler_kwargs.pop('rng', vmap(PRNGKey)(np.arange(num_chains)))
else:
rng = sampler_kwargs.pop('rng', PRNGKey(0))
init_kernel, sample_kernel = hmc(potential_fn, kinetic_fn, algo)
if progbar:
hmc_state = init_kernel(init_params, num_warmup, progbar=progbar, rng=rng,
**sampler_kwargs)
samples_flat = fori_collect(0, num_samples, sample_kernel, hmc_state,
transform=lambda x: constrain_fn(x.z),
progbar=progbar,
diagnostics_fn=get_diagnostics_str,
progbar_desc='sample')
samples = tree_map(lambda x: x[np.newaxis, ...], samples_flat)
else:
def single_chain_mcmc(rng, init_params):
hmc_state = init_kernel(init_params, num_warmup, run_warmup=False, rng=rng,
**sampler_kwargs)
samples = fori_collect(num_warmup, num_warmup + num_samples, sample_kernel, hmc_state,
transform=lambda x: constrain_fn(x.z),
progbar=progbar)
return samples
if num_chains == 1:
samples_flat = single_chain_mcmc(rng, init_params)
samples = tree_map(lambda x: x[np.newaxis, ...], samples_flat)
else:
if sequential_chain:
samples = lax.map(lambda args: single_chain_mcmc(*args), (rngs, init_params))
else:
samples = pmap(single_chain_mcmc)(rngs, init_params)
samples_flat = tree_map(lambda x: np.reshape(x, (-1,) + x.shape[2:]), samples)
if print_summary:
summary(samples)
return samples_flat
else:
raise ValueError('sampler: {} not recognized'.format(sampler))
def print_summary(self):
if 'z' not in self._samples:
raise ValueError('No latent samples `z` collected. Pass `z` to `collect_fields` arg.')
summary(self._samples['z'])
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 mcmc(num_warmup, num_samples, init_params, sampler='hmc',
constrain_fn=None, print_summary=True, **sampler_kwargs):
"""
Convenience wrapper for MCMC samplers -- runs warmup, prints
diagnostic summary and returns a collections of samples
from the posterior.
:param num_warmup: Number of warmup steps.
:param num_samples: Number of samples to generate from the Markov chain.
:param init_params: Initial parameters to begin sampling. The type can
must be consistent with the input type to `potential_fn`.
:param sampler: currently, only `hmc` is implemented (default).
:param constrain_fn: Callable that converts a collection of unconstrained
sample values returned from the sampler to constrained values that
lie within the support of the sample sites.
:param print_summary: Whether to print diagnostics summary for
each sample site. Default is ``True``.
:param `**sampler_kwargs`: Sampler specific keyword arguments.
- *HMC*: Refer to :func:`~numpyro.mcmc.hmc` and
:func:`~numpyro.mcmc.hmc.init_kernel` for accepted arguments. Note
that all arguments must be provided as keywords.
:return: collection of samples from the posterior.
.. testsetup::
import jax
from jax import random
import jax.numpy as np
import numpyro.distributions as dist
from numpyro.handlers import sample
from numpyro.hmc_util import initialize_model
from numpyro.mcmc import hmc
from numpyro.util import fori_collect
.. doctest::
>>> true_coefs = np.array([1., 2., 3.])
>>> data = random.normal(random.PRNGKey(2), (2000, 3))
>>> dim = 3
>>> labels = dist.Bernoulli(logits=(true_coefs * data).sum(-1)).sample(random.PRNGKey(3))
>>>
>>> def model(data, labels):
... coefs_mean = np.zeros(dim)
... coefs = sample('beta', dist.Normal(coefs_mean, np.ones(3)))
... intercept = sample('intercept', dist.Normal(0., 10.))
... return sample('y', dist.Bernoulli(logits=(coefs * data + intercept).sum(-1)), obs=labels)
>>>
>>> init_params, potential_fn, constrain_fn = initialize_model(random.PRNGKey(0), model,
... data, labels)
>>> num_warmup, num_samples = 1000, 1000
>>> samples = mcmc(num_warmup, num_samples, init_params,
... potential_fn=potential_fn,
... constrain_fn=constrain_fn) # doctest: +SKIP
warmup: 100%|██████████| 1000/1000 [00:09<00:00, 109.40it/s, 1 steps of size 5.83e-01. acc. prob=0.79]
sample: 100%|██████████| 1000/1000 [00:00<00:00, 1252.39it/s, 1 steps of size 5.83e-01. acc. prob=0.85]
mean sd 5.5% 94.5% n_eff Rhat
coefs[0] 0.96 0.07 0.85 1.07 455.35 1.01
coefs[1] 2.05 0.09 1.91 2.20 332.00 1.01
coefs[2] 3.18 0.13 2.96 3.37 320.27 1.00
intercept -0.03 0.02 -0.06 0.00 402.53 1.00
"""
if sampler == 'hmc':
if constrain_fn is None:
constrain_fn = identity
potential_fn = sampler_kwargs.pop('potential_fn')
kinetic_fn = sampler_kwargs.pop('kinetic_fn', None)
algo = sampler_kwargs.pop('algo', 'NUTS')
progbar = sampler_kwargs.pop('progbar', True)
init_kernel, sample_kernel = hmc(potential_fn, kinetic_fn, algo)
hmc_state = init_kernel(init_params, num_warmup, progbar=progbar, **sampler_kwargs)
samples = fori_collect(num_samples, sample_kernel, hmc_state,
transform=lambda x: constrain_fn(x.z),
progbar=progbar,
diagnostics_fn=get_diagnostics_str,
progbar_desc='sample')
if print_summary:
summary(samples)
return samples
else:
raise ValueError('sampler: {} not recognized'.format(sampler))
def print_summary(self, prob=0.9):
summary(self._states['z'], prob=prob)
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()