{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MCMC: sampler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext line_profiler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numba\n",
    "import numpy as np\n",
    "import scipy.stats"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Markov Chain Monte Carlo is a good example of a non-trivial algorithm we'd like to compute. To keep this tractable, we are just going to implement a gaussian sampler from MCMC. Don't worry about the code; we'll look at a much simpler Metropolis sampler at the end."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sampler(\n",
    "    data, samples=4, mu_init=0.5, proposal_width=0.5, mu_prior_mu=0, mu_prior_sd=1.0\n",
    "):\n",
    "    mu_current = mu_init\n",
    "    posterior = [mu_current]\n",
    "    for _ in range(samples):\n",
    "        # suggest new position\n",
    "        mu_proposal = scipy.stats.norm(mu_current, proposal_width).rvs()\n",
    "\n",
    "        # Compute likelihood by multiplying probabilities of each data point\n",
    "        likelihood_current = scipy.stats.norm(mu_current, 1).pdf(data).prod()\n",
    "        likelihood_proposal = scipy.stats.norm(mu_proposal, 1).pdf(data).prod()\n",
    "\n",
    "        # Compute prior probability of current and proposed mu\n",
    "        prior_current = scipy.stats.norm(mu_prior_mu, mu_prior_sd).pdf(mu_current)\n",
    "        prior_proposal = scipy.stats.norm(mu_prior_mu, mu_prior_sd).pdf(mu_proposal)\n",
    "\n",
    "        p_current = likelihood_current * prior_current\n",
    "        p_proposal = likelihood_proposal * prior_proposal\n",
    "\n",
    "        # Accept proposal?\n",
    "        p_accept = p_proposal / p_current\n",
    "\n",
    "        # Usually would include prior probability, which we neglect here for simplicity\n",
    "        accept = np.random.rand() < p_accept\n",
    "\n",
    "        if accept:\n",
    "            # Update position\n",
    "            mu_current = mu_proposal\n",
    "\n",
    "        posterior.append(mu_current)\n",
    "\n",
    "    return posterior"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "np.random.seed(123)\n",
    "data = np.random.randn(20)\n",
    "\n",
    "posterior = sampler(data, samples=15_000, mu_init=1.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1, 2, figsize=(12, 5))\n",
    "vals = posterior\n",
    "axs[0].plot(vals)\n",
    "axs[0].set_xlabel(\"Raw data\")\n",
    "axs[1].hist(vals[500:], bins=np.linspace(-1, 1, 100))\n",
    "axs[1].set_xlabel(\"Histogram\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I made a big, fat performance mistake. I expect it might not be the one you might expect. If I profiled the code above, I could see it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%lprun -f sampler sampler(data, samples=1_500, mu_init=1.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now do you see it?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's a much lighter weight norm:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def norm_pdf(loc, scale, x):\n",
    "    y = (x - loc) / scale\n",
    "    return np.exp(-(y**2) / 2) / np.sqrt(2 * np.pi) / scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert norm_pdf(0.4, 0.7, 0.2) == scipy.stats.norm(0.4, 0.7).pdf(0.2)\n",
    "assert scipy.stats.norm(0.3, 1).pdf(data).prod() == np.prod(norm_pdf(0.3, 1, data))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's a new version. Let's remove the list appending while we are at it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sampler(\n",
    "    data, samples=4, mu_init=0.5, proposal_width=0.5, mu_prior_mu=0, mu_prior_sd=1.0\n",
    "):\n",
    "    mu_current = mu_init\n",
    "\n",
    "    posterior = np.empty(samples + 1)\n",
    "    posterior[0] = mu_current\n",
    "\n",
    "    for i in range(samples):\n",
    "        # suggest new position\n",
    "        mu_proposal = np.random.normal(mu_current, proposal_width)\n",
    "\n",
    "        # Compute likelihood by multiplying probabilities of each data point\n",
    "        likelihood_current = np.prod(norm_pdf(mu_current, 1, data))\n",
    "        likelihood_proposal = np.prod(norm_pdf(mu_proposal, 1, data))\n",
    "\n",
    "        # Compute prior probability of current and proposed mu\n",
    "        prior_current = norm_pdf(mu_prior_mu, mu_prior_sd, mu_current)\n",
    "        prior_proposal = norm_pdf(mu_prior_mu, mu_prior_sd, mu_proposal)\n",
    "\n",
    "        p_current = likelihood_current * prior_current\n",
    "        p_proposal = likelihood_proposal * prior_proposal\n",
    "\n",
    "        # Accept proposal?\n",
    "        p_accept = p_proposal / p_current\n",
    "\n",
    "        # Usually would include prior probability, which we neglect here for simplicity\n",
    "        accept = np.random.rand() < p_accept\n",
    "\n",
    "        if accept:\n",
    "            # Update position\n",
    "            mu_current = mu_proposal\n",
    "\n",
    "        posterior[i + 1] = mu_current\n",
    "\n",
    "    return posterior"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "np.random.seed(123)\n",
    "data = np.random.randn(20)\n",
    "\n",
    "posterior = sampler(data, samples=15_000, mu_init=1.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%lprun -f sampler sampler(data, samples=15_000, mu_init=1.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Much* better. Instancing scipy distributions is *very* costly. But we'd love to be able to produce massive amounts of MC. Can we try Numba?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.vectorize([numba.float64(numba.float64, numba.float64, numba.float64)])\n",
    "def norm_pdf(loc, scale, x):\n",
    "    y = (x - loc) / scale\n",
    "    return np.exp(-(y**2) / 2) / np.sqrt(2 * np.pi) / scale"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> ### Python note\n",
    ">\n",
    "> Functions look up methods and values from their surrounding scope *when called*. This means that `norm_pdf` is now the new `norm_pdf`, even though we have not changed the surrounding function.\n",
    ">\n",
    "> You probably should not normally do this. And it makes this notebook depend on the order of execution, which is not ideal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "np.random.seed(123)\n",
    "data = np.random.randn(20)\n",
    "\n",
    "posterior = sampler(data, samples=15000, mu_init=1.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%lprun -f sampler sampler(data, samples=15_000, mu_init=1.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's go all in and do Numba start to finish:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.njit\n",
    "def norm_pdf(loc, scale, x):\n",
    "    y = (x - loc) / scale\n",
    "    return np.exp(-(y**2) / 2) / np.sqrt(2 * np.pi) / scale"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: This is mostly redefined to show that is could be done with `njit` instead of `vectorize`, `vectorize` is actually a bit simpler/faster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.njit\n",
    "def sampler(\n",
    "    data,\n",
    "    samples=4,\n",
    "    mu_init=0.5,\n",
    "    proposal_width=0.5,\n",
    "    mu_prior_mu=0,\n",
    "    mu_prior_sd=1.0,\n",
    "):\n",
    "    mu_current = mu_init\n",
    "\n",
    "    posterior = np.empty(samples + 1)\n",
    "    posterior[0] = mu_current\n",
    "\n",
    "    for i in range(samples):\n",
    "        # suggest new position\n",
    "        mu_proposal = np.random.normal(mu_current, proposal_width)\n",
    "\n",
    "        # Compute likelihood by multiplying probabilities of each data point\n",
    "        likelihood_current = np.prod(norm_pdf(mu_current, 1, data))\n",
    "        likelihood_proposal = np.prod(norm_pdf(mu_proposal, 1, data))\n",
    "\n",
    "        # Compute prior probability of current and proposed mu\n",
    "        prior_current = norm_pdf(mu_prior_mu, mu_prior_sd, mu_current)\n",
    "        prior_proposal = norm_pdf(mu_prior_mu, mu_prior_sd, mu_proposal)\n",
    "\n",
    "        p_current = likelihood_current * prior_current\n",
    "        p_proposal = likelihood_proposal * prior_proposal\n",
    "\n",
    "        # Accept proposal?\n",
    "        p_accept = p_proposal / p_current\n",
    "\n",
    "        # Usually would include prior probability, which we neglect here for simplicity\n",
    "        accept = np.random.rand() < p_accept\n",
    "\n",
    "        if accept:\n",
    "            # Update position\n",
    "            mu_current = mu_proposal\n",
    "\n",
    "        posterior[i + 1] = mu_current\n",
    "\n",
    "    return posterior"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "np.random.seed(123)\n",
    "data = np.random.randn(20)\n",
    "\n",
    "posterior = sampler(data, samples=15_000, mu_init=1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ouch! Let's try that again:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "np.random.seed(123)\n",
    "data = np.random.randn(20)\n",
    "\n",
    "posterior = sampler(data, samples=15_000, mu_init=1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sweet."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Metropolis sampler\n",
    "\n",
    "Let's look at a simpler example, and one that you might find more useful/instructive."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Metropolis sampler\n",
    "\n",
    "\n",
    "def p(x):\n",
    "    \"Any function that you want to sample. Plot will look better if it is normalized.\"\n",
    "    return 1 / (1 + x**2) / np.pi\n",
    "    # return 1/(1 + x**2) * np.sin(x)**2 / (np.pi * np.sinh(1) / np.exp(1))\n",
    "\n",
    "\n",
    "def metropolis(p, samples=50_000, sigma=1):\n",
    "    x = np.zeros(samples + 1)\n",
    "    x[0] = np.random.rand()\n",
    "\n",
    "    for i in range(samples):\n",
    "        # suggest new position\n",
    "        x_Star = np.random.normal(x[i], sigma)\n",
    "\n",
    "        # Compute alpha - the fractional chance of moving to a new point\n",
    "        alpha = p(x_Star) / p(x[i])\n",
    "\n",
    "        # Accept/reject based on alpha\n",
    "        accept = alpha > np.random.rand()\n",
    "\n",
    "        # Add the current (moved?) point\n",
    "        x[i + 1] = x_Star if accept else x[i]\n",
    "\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vals = metropolis(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "vals = metropolis(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(-10, 10, 200)\n",
    "fig, axs = plt.subplots(2, figsize=(10, 6))\n",
    "axs[0].plot(vals[:500], \"r\")\n",
    "axs[0].plot(np.arange(500, len(vals)), vals[500:], \"g\")\n",
    "\n",
    "axs[1].hist(vals[500:], bins=400, range=(-20, 20), density=True)\n",
    "axs[1].plot(x, p(x), lw=3)\n",
    "axs[1].set_xlim(-10, 10);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Try the other function, and try adding `@numba.njit` above."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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"
  },
  "vscode": {
   "interpreter": {
    "hash": "2d4e9b9c84dab3e1662173f95b81bd7f8a551068d04f5f3c42d164db7312a928"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}