11-cython-mcmc-taking-it-to-11.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimized Cython MCMC implementation\n",
    "\n",
    "The implementation in the previous section still used `numpy.random` calls, which go through Python.\n",
    "\n",
    "In this notebook, we'll demonstrate creating a separate Cython package called `cython_mcmc` that uses another Cython package named `mt19937` for faster random number generation.\n",
    "\n",
    "We will see a significant speedup for the MCMC sampler as a result of our efforts."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## First step -- compile external packages and run perfomance tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "cd ./rng\n",
    "python ./setup.py develop\n",
    "cython -a ./rng/mt19937.pyx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "cd cython_mcmc\n",
    "python ./setup.py develop\n",
    "cython -a ./cython_mcmc/mcmc.pyx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "import numpy as np\n",
    "from scipy.stats import norm\n",
    "\n",
    "from cython_mcmc import mcmc\n",
    "\n",
    "np.random.seed(123)\n",
    "data = np.random.randn(20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit mcmc.log_sampler(data, samples=15000, mu_init=1.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "posterior = mcmc.log_sampler(data, samples=15000, mu_init=1.0)\n",
    "plt.plot(posterior);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def calc_posterior_analytical(data, x, mu_0, sigma_0):\n",
    "    sigma = 1.\n",
    "    n = len(data)\n",
    "    mu_post = (mu_0 / sigma_0**2 + data.sum() / sigma**2) / (1. / sigma_0**2 + n / sigma**2)\n",
    "    sigma_post = (1. / sigma_0**2 + n / sigma**2)**-1\n",
    "    return norm(mu_post, np.sqrt(sigma_post)).pdf(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = plt.subplot()\n",
    "\n",
    "sns.distplot(posterior[500:], ax=ax, label='estimated posterior')\n",
    "x = np.linspace(-.7, .9, 500)\n",
    "post = calc_posterior_analytical(data, x, 0, 1)\n",
    "ax.plot(x, post, 'g', label='analytic posterior')\n",
    "_ = ax.set(xlabel='mu', ylabel='belief');\n",
    "ax.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Standard error in $\\mu$ shrinks as more data is collected"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "data_2000 = np.random.randn(2000)\n",
    "posterior_2000 = mcmc.log_sampler(data_2000, samples=150000, mu_init=1.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = plt.subplot()\n",
    "\n",
    "sns.distplot(posterior_2000[500::5], ax=ax, label='estimated posterior')\n",
    "x = np.linspace(-.1, .1, 500)\n",
    "post = calc_posterior_analytical(data_2000, x, 0, 1)\n",
    "ax.plot(x, post, 'g', label='analytic posterior')\n",
    "_ = ax.set(xlabel='mu', ylabel='belief');\n",
    "ax.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `log_sampler()` using external library"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%html\n",
    "./cython_mcmc/cython_mcmc/mcmc.html"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Worth noting\n",
    "\n",
    "* `from [...] cimport RandomState`\n",
    "* `from [...] import RandomState`\n",
    "* `sample_norm()` and `accept_p()` implementations.\n",
    "* `norm_logpdf()` implementation uses raw C buffers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `RandomState` extension type\n",
    "\n",
    "Cython-level compile-time interface defined in `mt19937.pxd`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ./rng/rng/mt19937.pxd | nl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%html\n",
    "./rng/rng/mt19937.html"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Putting it all together in `setup.py` with `cythonize()`\n",
    "\n",
    "* We use a `setup.py` script in the `cython_mcmc` package to compile everything together.\n",
    "* We specify all the `mcmc.pyx` external Cython and C dependencies in an `Extension()` object.\n",
    "* We use the `Cython.Build.cythonize()` command to pull everything together and compile things for us.\n",
    "* The `python setup.py develop` command is what kicks things off for us, and produces the `mcmc.so` shared object file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%cat ./cython_mcmc/setup.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}