post.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this post, I will demonstrate how you can use [emcee](https://emcee.readthedocs.io/en/latest/) to sample models defined using [PyMC3](https://docs.pymc.io/).\n",
    "Thomas Wiecki [wrote about how to do this this with an earlier version of PyMC](http://twiecki.github.io/blog/2013/09/23/emcee-pymc/), but I needed an update since I wanted to do a comparison and PyMC's interface has changed a lot since he wrote his post.\n",
    "This isn't necessarily something that you'll *want* to do (and I definitely don't recommend it in general), but I figured that I would post it here for posterity.\n",
    "\n",
    "For simplicity, let's use the simulated data from my [previous blog post](/posts/pymc-tensorflow):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "np.random.seed(42)\n",
    "\n",
    "true_params = np.array([0.5, -2.3, -0.23])\n",
    "\n",
    "N = 50\n",
    "t = np.linspace(0, 10, 2)\n",
    "x = np.random.uniform(0, 10, 50)\n",
    "y = x * true_params[0] + true_params[1]\n",
    "y_obs = y + np.exp(true_params[-1]) * np.random.randn(N)\n",
    "\n",
    "plt.plot(x, y_obs, \".k\", label=\"observations\")\n",
    "plt.plot(t, true_params[0] * t + true_params[1], label=\"truth\")\n",
    "plt.xlabel(\"x\")\n",
    "plt.ylabel(\"y\")\n",
    "plt.legend(fontsize=14);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we can code up the model in PyMC3 following [Jake VanderPlas' notation](http://jakevdp.github.io/blog/2014/06/14/frequentism-and-bayesianism-4-bayesian-in-python/), and sample it using PyMC3's NUTS[sic] sampler:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pymc3 as pm\n",
    "import theano.tensor as tt\n",
    "\n",
    "with pm.Model() as model:\n",
    "    logs = pm.Normal(\"logs\", mu=0.0, sigma=3.0)\n",
    "    alphaperp = pm.Normal(\"alphaperp\", mu=0.0, sigma=3.0)\n",
    "    theta = pm.Uniform(\"theta\", -np.pi, np.pi, testval=0.0)\n",
    "\n",
    "    # alpha_perp = alpha * cos(theta)\n",
    "    alpha = pm.Deterministic(\"alpha\", alphaperp / tt.cos(theta))\n",
    "\n",
    "    # beta = tan(theta)\n",
    "    beta = pm.Deterministic(\"beta\", tt.tan(theta))\n",
    "\n",
    "    # The observation model\n",
    "    mu = alpha * x + beta\n",
    "    pm.Normal(\"obs\", mu=mu, sd=tt.exp(logs), observed=y_obs)\n",
    "\n",
    "    trace = pm.sample(\n",
    "        draws=2000,\n",
    "        tune=5000,\n",
    "        chains=2,\n",
    "        cores=1,\n",
    "        target_accept=0.9,\n",
    "        return_inferencedata=True,\n",
    "        init=\"adapt_full\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we can take a look at the [corner plot](https://corner.readthedocs.io/en/latest/):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import corner\n",
    "\n",
    "names = [\"alpha\", \"beta\", \"logs\"]\n",
    "corner.corner(trace, var_names=names, truths=dict(zip(names, true_params)));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sampling the PyMC3 model using emcee\n",
    "\n",
    "To sample this using emcee, we'll need to do a little bit of bookkeeping.\n",
    "I've coded this up using version 3 of emcee that is currently available as [the master branch on GitHub](https://github.com/dfm/emcee) or as [a pre-release on PyPI](https://pypi.org/project/emcee/3.0rc1/), so you'll need to install that version to run this.\n",
    "\n",
    "To sample from this model, we need to expose the Theano method for evaluating the log probability to Python.\n",
    "There is a version of this built into PyMC3, but I also want to return the values of all the deterministic variables using the [\"blobs\" feature in emcee](https://emcee.readthedocs.io/en/latest/user/blobs/) so the function is slightly more complicated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import theano\n",
    "\n",
    "with model:\n",
    "    f = theano.function(\n",
    "        model.vars, [model.logpt] + list(model.vars) + list(model.deterministics)\n",
    "    )\n",
    "\n",
    "    def log_prob_func(params):\n",
    "        dct = model.bijection.rmap(params)\n",
    "        args = (dct[k.name] for k in model.vars)\n",
    "        results = f(*args)\n",
    "        return tuple(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now we can run the sampler:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import emcee\n",
    "\n",
    "with model:\n",
    "    # First we work out the shapes of all of the deterministic variables\n",
    "    res = pm.find_MAP()\n",
    "    vec = model.bijection.map(res)\n",
    "    initial_blobs = log_prob_func(vec)[1:]\n",
    "    dtype = [\n",
    "        (var.name, float, np.shape(b))\n",
    "        for var, b in zip(list(model.vars) + list(model.deterministics), initial_blobs)\n",
    "    ]\n",
    "\n",
    "    # Then sample as usual\n",
    "    coords = vec + 1e-5 * np.random.randn(25, len(vec))\n",
    "    nwalkers, ndim = coords.shape\n",
    "    sampler = emcee.EnsembleSampler(nwalkers, ndim, log_prob_func, blobs_dtype=dtype)\n",
    "    sampler.run_mcmc(coords, 5000, progress=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we can use this to make the same corner plot as above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "df = pd.DataFrame.from_records(sampler.get_blobs(flat=True, discard=100, thin=30))\n",
    "corner.corner(df[[\"alpha\", \"beta\", \"logs\"]], truths=true_params);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The last thing that we might want to look at is [the integrated autocorrelation time](/posts/autocorr/) for each method.\n",
    "First, as expected, the autocorrelation for PyMC3 is very short (about 1 step):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "[\n",
    "    float(emcee.autocorr.integrated_time(np.array(trace.posterior[var].values).T))\n",
    "    for var in [\"logs\", \"alphaperp\", \"theta\"]\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And, the autocorrelation for emcee is about 40 steps:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler.get_autocorr_time(discard=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want to compare these results in detail, you'll want to make sure that you take into account the fact that each step of NUTS is significantly more expensive than one step with emcee, but that's way beyond the scope of this post...\n",
    "\n",
    "*11/22/18: This post has been updated with suggestions from Thomas Wiecki. The `find_MAP` call has been removed from the PyMC sampling, and `model.bijection` is now used to map between arrays and dicts of parameters.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "a570d0109328185cd12f52c24faac3746ace769974c6e73c9f7351d594465c68"
  },
  "kernelspec": {
   "display_name": "Python 3.9.13 ('env': venv)",
   "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}