{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Basic usage and comparison\n",
"\n",
"This example will run through a basic comparison of `MultiHMCGibbs` with `NUTS` and `HMCGibbs` for a simple problem where all three methods should give the same results. Our target distribution will be a 2D normal with some covariance."
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import jax.numpy as jnp\n",
"import numpyro\n",
"import numpyro.distributions as dist\n",
"import arviz\n",
"import corner\n",
"import matplotlib.pyplot as plt\n",
"\n",
"from numpyro.infer import MCMC, NUTS, HMCGibbs\n",
"from jax import random\n",
"\n",
"from MultiHMCGibbs import MultiHMCGibbs"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## The model\n",
"\n",
"As with any Numpyro sampler we start with a model function for our distribution."
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"def model():\n",
" x = numpyro.sample(\"x\", dist.Normal(0.0, 2.0))\n",
" y = numpyro.sample(\"y\", dist.Normal(0.0, 2.0))\n",
" numpyro.sample(\"obs\", dist.Normal(x + y, 1.0), obs=jnp.array([1.0]))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## NUTS\n",
"\n",
"To start we will use NUTS so we know what we are comparing to."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"rng_key = random.PRNGKey(0)\n",
"hmc_key, rng_key = random.split(rng_key)"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"divergences per chain: [0 0 0 0]\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"/Users/coleman/anaconda3/envs/lensing/lib/python3.10/site-packages/xarray/core/concat.py:544: FutureWarning: unique with argument that is not not a Series, Index, ExtensionArray, or np.ndarray is deprecated and will raise in a future version.\n",
" common_dims = tuple(pd.unique([d for v in vars for d in v.dims]))\n"
]
},
{
"data": {
"text/html": [
"\n",
"\n",
"
\n",
" \n",
" \n",
" mean \n",
" sd \n",
" hdi_3% \n",
" hdi_97% \n",
" mcse_mean \n",
" mcse_sd \n",
" ess_bulk \n",
" ess_tail \n",
" r_hat \n",
" \n",
" \n",
" \n",
" \n",
" x \n",
" 0.462 \n",
" 1.491 \n",
" -2.447 \n",
" 3.221 \n",
" 0.043 \n",
" 0.031 \n",
" 1197.0 \n",
" 1200.0 \n",
" 1.0 \n",
" \n",
" \n",
" y \n",
" 0.420 \n",
" 1.498 \n",
" -2.312 \n",
" 3.307 \n",
" 0.044 \n",
" 0.033 \n",
" 1169.0 \n",
" 1379.0 \n",
" 1.0 \n",
" \n",
" \n",
"
\n",
"
"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"hmc_kernel = NUTS(model)\n",
"mcmc = MCMC(\n",
" hmc_kernel,\n",
" num_warmup=100,\n",
" num_samples=1000,\n",
" num_chains=4,\n",
" chain_method='vectorized',\n",
" progress_bar=False\n",
")\n",
"mcmc.run(hmc_key)\n",
"\n",
"inf_data_hmc = arviz.from_numpyro(mcmc)\n",
"print(f'divergences per chain: {inf_data_hmc.sample_stats.diverging.values.sum(axis=1)}')\n",
"display(arviz.summary(inf_data_hmc))\n",
"fig = corner.corner(inf_data_hmc, color='C0')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## MultiHMCGibbs\n",
"\n",
"To use `MultiHMCGibbs` you need to create a list of HMC or NUTS kernels that wrap the same model, but each can have their own keywords such as `target_accept_prob` or `max_tree_depth`. The other argument is a list of lists containing the **free** parameters for each of the inner kernels. Internally each the sampler will:\n",
"1. Loop over the kernels in the list\n",
"2. Conditioned it on the non-free parameters\n",
"3. Re-calculate the likelihood and gradients at the new conditioned point\n",
"4. Step the kernel forward\n",
"5. Move on to the next kernel\n",
"\n",
"After stepping the last kernel all the parameters will be updated and a single Gibbs step has been taken.\n",
"\n",
"**Important**: All free parameters must be listed **exactly once** for the sampler to work. If this is not the case a `ValueError` will be raised listing what values are either duplicated, extra, or missing from the list."
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"multigibbs_key, rng_key = random.split(rng_key)"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"divergences per chain per step:\n",
" [[0 0 0 0]\n",
" [0 0 0 0]]\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"/Users/coleman/anaconda3/envs/lensing/lib/python3.10/site-packages/xarray/core/concat.py:544: FutureWarning: unique with argument that is not not a Series, Index, ExtensionArray, or np.ndarray is deprecated and will raise in a future version.\n",
" common_dims = tuple(pd.unique([d for v in vars for d in v.dims]))\n"
]
},
{
"data": {
"text/html": [
"\n",
"\n",
"
\n",
" \n",
" \n",
" mean \n",
" sd \n",
" hdi_3% \n",
" hdi_97% \n",
" mcse_mean \n",
" mcse_sd \n",
" ess_bulk \n",
" ess_tail \n",
" r_hat \n",
" \n",
" \n",
" \n",
" \n",
" x \n",
" 0.512 \n",
" 1.545 \n",
" -2.279 \n",
" 3.448 \n",
" 0.080 \n",
" 0.057 \n",
" 376.0 \n",
" 669.0 \n",
" 1.00 \n",
" \n",
" \n",
" y \n",
" 0.355 \n",
" 1.552 \n",
" -2.493 \n",
" 3.249 \n",
" 0.081 \n",
" 0.057 \n",
" 377.0 \n",
" 678.0 \n",
" 1.01 \n",
" \n",
" \n",
"
\n",
"
"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"inner_kernels = [\n",
" NUTS(model),\n",
" NUTS(model)\n",
"]\n",
"\n",
"outer_kernel = MultiHMCGibbs(\n",
" inner_kernels,\n",
" [['y'], ['x']] # first updated y, then update x\n",
")\n",
"\n",
"mcmc_gibbs = MCMC(\n",
" outer_kernel,\n",
" num_warmup=100,\n",
" num_samples=1000,\n",
" num_chains=4,\n",
" chain_method='vectorized',\n",
" progress_bar=False\n",
")\n",
"mcmc_gibbs.run(multigibbs_key)\n",
"\n",
"inf_data_gibbs = arviz.from_numpyro(mcmc_gibbs)\n",
"print(f'divergences per chain per step:\\n {inf_data_gibbs.sample_stats.diverging.values.sum(axis=1).T}')\n",
"display(arviz.summary(inf_data_gibbs))\n",
"\n",
"print('HMC (Blue)')\n",
"print('MultiHMCGibbs (Orange)')\n",
"fig = corner.corner(inf_data_gibbs, color='C1')\n",
"_ = corner.corner(inf_data_hmc, fig=fig, color='C0')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can see that the two samplers give the same results, a good indication that everything is working.\n",
"\n",
"## HMCGibbs\n",
"\n",
"For this simple model we can also do the conditioning on `y` by hand and use Numpyro's builtin `HMCGibbs` sampler. Let's see how that compares.\n",
"\n",
"Note: We will use `chian_method='sequential'` for this sampler because it does support `vectorized`. We also don't have access to the `divergences` of the HMC kernel so those will not be reported either."
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"gibbs_key, rng_key = random.split(rng_key)"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"/Users/coleman/anaconda3/envs/lensing/lib/python3.10/site-packages/xarray/core/concat.py:544: FutureWarning: unique with argument that is not not a Series, Index, ExtensionArray, or np.ndarray is deprecated and will raise in a future version.\n",
" common_dims = tuple(pd.unique([d for v in vars for d in v.dims]))\n"
]
},
{
"data": {
"text/html": [
"\n",
"\n",
"
\n",
" \n",
" \n",
" mean \n",
" sd \n",
" hdi_3% \n",
" hdi_97% \n",
" mcse_mean \n",
" mcse_sd \n",
" ess_bulk \n",
" ess_tail \n",
" r_hat \n",
" \n",
" \n",
" \n",
" \n",
" x \n",
" 0.459 \n",
" 1.445 \n",
" -2.298 \n",
" 3.091 \n",
" 0.067 \n",
" 0.047 \n",
" 472.0 \n",
" 1299.0 \n",
" 1.01 \n",
" \n",
" \n",
" y \n",
" 0.425 \n",
" 1.438 \n",
" -2.218 \n",
" 3.136 \n",
" 0.075 \n",
" 0.053 \n",
" 366.0 \n",
" 824.0 \n",
" 1.01 \n",
" \n",
" \n",
"
\n",
"
"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"def gibbs_fn(rng_key, gibbs_sites, hmc_sites):\n",
" y = hmc_sites['y']\n",
" new_x = dist.Normal(0.8 * (1-y), jnp.sqrt(0.8)).sample(rng_key)\n",
" return {'x': new_x}\n",
"\n",
"\n",
"kernel_gibbs_fn = HMCGibbs(hmc_kernel, gibbs_fn=gibbs_fn, gibbs_sites=['x'])\n",
"\n",
"mcmc_gibbs_fn = MCMC(\n",
" kernel_gibbs_fn,\n",
" num_warmup=100,\n",
" num_samples=1000,\n",
" num_chains=4,\n",
" chain_method='sequential',\n",
" progress_bar=False\n",
")\n",
"\n",
"mcmc_gibbs_fn.run(gibbs_key)\n",
"\n",
"inf_data_gibbs_fn = arviz.from_numpyro(mcmc_gibbs_fn)\n",
"display(arviz.summary(inf_data_gibbs_fn))\n",
"\n",
"print('HMC (Blue)')\n",
"print('MultiHMCGibbs (Orange)')\n",
"print('HMCGibbs (Green)')\n",
"fig = corner.corner(inf_data_gibbs, color='C1')\n",
"_ = corner.corner(inf_data_hmc, fig=fig, color='C0')\n",
"_ = corner.corner(inf_data_gibbs_fn, fig=fig, color='C2')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"As before, we see that this sampler gives the same results."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "lensing",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.9"
}
},
"nbformat": 4,
"nbformat_minor": 2
}