Module5-BayesianThinkingAndMCMC/M5B_Markov_Chain_Monte_Carlo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Markov Chain Monte Carlo Method\n",
    "<html>\n",
    "    <summary></summary>\n",
    "         <div> <p></p> </div>\n",
    "         <div style=\"font-size: 20px; width: 800px;\"> \n",
    "              <h1>\n",
    "               <left>Markov Chain Monte Carlo Methods</left>\n",
    "              </h1>\n",
    "              <p><left>============================================================================</left> </p>\n",
    "<pre>Course: BIOM 421, Spring 2024\n",
    "Instructor: Brian Munsky\n",
    "Authors: Drs. Kaan Öcal, Huy Vo, Brian Munsky\n",
    "Contact Info: munsky@colostate.edu\n",
    "</pre>\n",
    "         </div>\n",
    "    </p>\n",
    "\n",
    "</html>\n",
    "\n",
    "<details>\n",
    "  <summary>Copyright info</summary>\n",
    "\n",
    "```\n",
    "Copyright 2024 Brian Munsky\n",
    "\n",
    "Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:\n",
    "\n",
    "1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.\n",
    "\n",
    "2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.\n",
    "\n",
    "3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.\n",
    "\n",
    "THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n",
    "```\n",
    "<details>\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kfORGQnMZfb4"
   },
   "outputs": [],
   "source": [
    "# Import packages\n",
    "import numpy as np\n",
    "import scipy as sp\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm import tqdm        # fancy progress bars"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction\n",
    "\n",
    "In the last lecture we will learn how to perform Bayesian inference in practice to estimate model parameters from data. This will involve a new technique, Markov Chain Monte Carlo (MCMC), which is the current state-of-the-art approach to deal with the complicated probability distributions arising in Bayesian inference, statistical physics and machine learning. Our goal in this module is to understand the idea behind MCMC, to apply it to a simple example or two, and to combine it with the Chemical Master Equation to do some inference.\n",
    "\n",
    "## Contents\n",
    "* Monte Carlo methods\n",
    "* Markov Chains\n",
    "* Markov Chain Monte Carlo (MCMC)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zSMFxe4jnqS7"
   },
   "source": [
    "# Monte Carlo methods"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pSLRwQdM9Xy_"
   },
   "source": [
    "![monte_carlo.jpg](https://upload.wikimedia.org/wikipedia/commons/thumb/4/4b/D8E_1411_%288525313935%29.jpg/640px-D8E_1411_%288525313935%29.jpg)\n",
    "\n",
    "*Source: Wikimedia*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wHvEepNJzw27"
   },
   "source": [
    "## What is a Monte Carlo method?\n",
    "\n",
    "A Monte Carlo method, named after a famous casino in Monaco, is any algorithm that uses randomly generated numbers to solve a numerical problem.\n",
    "\n",
    "For example, if we try to estimate the mean of a probability distribution $p(x)$, we can draw a few samples from $p(x)$ and compute their average. This will give a rough estimate of the mean, and we can make our estimate more precise by drawing more samples from $p(x)$. Note that our estimate of the mean can fluctuate randomly, even though the true mean of $p(x)$ is a deterministic quantity - never confuse the empirical mean with the true mean.\n",
    "\n",
    "Monte Carlo methods are often easier to implement than other methods, and they are usually flexible enough to work for many related problems with small modifications. Their main disadvantage is that the outputs are random by nature, and getting enough samples to obtain reliable estimates can be very time-consuming for some problems.\n",
    "\n",
    "A typical Monte Carlo method goes like this:\n",
    "1. Fix a probability distribution to sample from\n",
    "2. Generate a sequence of random samples from that probability distribution\n",
    "3. Perform a deterministic computation on each sample\n",
    "4. Combine your results to arrive at the answer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Yh5s7ts3xV-A"
   },
   "source": [
    "## **Example:** Computing the number $\\pi$ using Monte Carlo\n",
    "\n",
    "To demonstrate the flexibility of Monte Carlo methods we can try to compute $\\pi$ using random sampling. In the figure below we see a circle with radius $r = 1$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 316
    },
    "id": "98_UhdHxoxza",
    "jupyter": {
     "source_hidden": true
    },
    "outputId": "c33a6e79-55d3-40dd-da09-f78d2ef6fa14"
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(3,3))\n",
    "\n",
    "circle = plt.Circle((0, 0), 1)\n",
    "square = plt.Rectangle((-1, -1), 2, 2, color=\"orange\")\n",
    "\n",
    "ax.add_patch(square)\n",
    "ax.add_patch(circle)\n",
    "\n",
    "ax.set_xlim(-1, 1)\n",
    "ax.set_ylim(-1, 1)\n",
    "\n",
    "ax.set_xlabel(\"x\")\n",
    "ax.set_ylabel(\"y\");\n",
    "\n",
    "ax.set_aspect(\"equal\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XLyrfbNpoqQQ"
   },
   "source": [
    "From the formula $A = \\pi r^2$ we know that the circle has area $\\pi$. The square containing the circle has area $4$, so the circle fills out a fraction $\\pi/4 \\approx 0.79$ of the square. In other words, if we sample a point at random from the square, then with probability $\\pi/4$ it will lie in the circle.\n",
    "\n",
    "We can now estimate $\\pi$ by repeatedly sampling points from the square and counting how many of them lie in the circle. This is a typical example of a Monte Carlo method:\n",
    "1. Our probability distribution is uniform on the square.\n",
    "2. Sample points from that distribution so that we get even coverage of the square.\n",
    "3. For each point, check if it lies inside the unit circle.\n",
    "4. Return the fraction of points that lie inside the unit circle, times $4$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "fgtCgGmDygaP",
    "outputId": "1d9dd2d5-80ee-4ac3-8265-58b1b51230b4"
   },
   "outputs": [],
   "source": [
    "# Estimating $\\pi$ using random sampling\n",
    "N = 10_000 # number of samples to use.\n",
    "\n",
    "# Draw random points in the square by sampling their x and y coordinates\n",
    "points = np.random.uniform(-1, 1, size=(N, 2))\n",
    "\n",
    "# The interior of the circle is given by the equation x^2 + y^2 <= 1\n",
    "# (Pythagoras' Theorem)\n",
    "inside = (points[:,0] ** 2 + points[:,1] ** 2) <= 1\n",
    "pi_est = 4 * np.mean(inside)\n",
    "\n",
    "print(f\"π is approximately {pi_est}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 316
    },
    "id": "oJSKsjL108Qc",
    "jupyter": {
     "source_hidden": true
    },
    "outputId": "db6eae22-a16c-41e1-b8c9-84eecca28d8d"
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(3,3))\n",
    "\n",
    "ax.scatter(points[inside,0], points[inside,1], s=1)\n",
    "ax.scatter(points[~inside,0], points[~inside,1], s=1)    # ~ negates the condition\n",
    "\n",
    "ax.set_xlim(-1, 1)\n",
    "ax.set_ylim(-1, 1)\n",
    "\n",
    "ax.set_xlabel(\"x\")\n",
    "ax.set_ylabel(\"y\");\n",
    "\n",
    "ax.set_aspect(\"equal\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Gtfs7Qks3q5s"
   },
   "source": [
    "We expect that our approximation will get better the more samples we use. In fact, we can compute the variance of our estimate (we secretly are using the binomial distribution):\n",
    "\n",
    "$$ \\mathrm{Var}(\\hat \\pi) \\approx 0.674 \\cdot \\frac 1 {N} $$\n",
    "\n",
    "The variance decreases in proportion to $N$, which tells us that our estimate will get more accurate the larger $N$ is. More precisely, since the variance measures the typical deviation squared, we predict that the actual error decreases as $1 / \\sqrt{N}$. We can verify this empirically:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 362
    },
    "id": "y7kTTQ1E1-UW",
    "outputId": "e3effc10-e2a2-4fd6-8612-c253479a056b"
   },
   "outputs": [],
   "source": [
    "# Estimating the Monte Carlo error\n",
    "# Sample a lot of points at once\n",
    "N = 1_000_000\n",
    "points = np.random.uniform(-1, 1, size=(N, 2))\n",
    "\n",
    "# Check which points are inside the circle\n",
    "inside = (points[:,0] ** 2 + points[:,1] ** 2) <= 1\n",
    "\n",
    "# Compute a running estimate of π\n",
    "pi_est = 4 * np.cumsum(inside) / range(1, N+1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1,2, sharex=True, figsize=(6, 3))\n",
    "\n",
    "# Plot the estimate vs. π\n",
    "ax[0].plot(range(N), pi_est)\n",
    "ax[0].set_ylim([3.1,3.2])\n",
    "ax[0].set_xlabel(\"Number of samples\")\n",
    "ax[0].set_ylabel(\"Estimate for $\\\\pi$\")\n",
    "ax[0].axhline(np.pi, color=\"black\")\n",
    "\n",
    "# Plot the magnitude of the error\n",
    "ax[1].plot(range(N), abs(pi_est - np.pi))\n",
    "ax[1].plot([1, N], np.sqrt(0.674) / np.sqrt([1., N]))   # the error should decrease as 1/sqrt(N)\n",
    "ax[1].set_xscale(\"log\")\n",
    "ax[1].set_xlim([100, N])\n",
    "ax[1].set_yscale(\"log\")\n",
    "ax[1].set_xlabel(\"Number of samples\")\n",
    "ax[1].set_ylabel(\"|Estimate error|\")\n",
    "ax[1].legend([\"Computed\", \"Predicted\"])\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "X-JCRe7Y4ct1"
   },
   "source": [
    "In this example we can compute the error because we know the true value of $\\pi = 3.14159$... In general we don't exactly know the quantity we are trying to estimate, and it can be hard to detect when we have sufficient samples to be confident about our answer.\n",
    "\n",
    "This is a simple example of a Monte Carlo method where the distribution we need to sample from (uniform) is very tractable. In Bayesian inference, the distribution we want to sample from (the posterior) is often very complex. The standard way to sample from such complex distributions, and our goal for this lecture, is using Markov Chain Monte Carlo (MCMC). In the next section we will talk about Markov chains, the essential idea behind MCMC, before finally moving on to applications."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9Fu5vz3KUkDB"
   },
   "source": [
    "## Exercises\n",
    "\n",
    "1. **a.** In the example, each sampled point has a $\\pi / 4$ chance of being in the circle. Show that the total number of points lying in the circle follows the binomial distribution $\\textrm{Bin}(N, \\pi/4)$.\n",
    "\n",
    "    **b.** Using the properties of the binomial distribution, compute the variance of our estimate $\\hat \\pi$ and check that it agrees with the equation above."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QjV33hK39_A-"
   },
   "source": [
    "## A note on formulas vs. samples\n",
    "\n",
    "There are two ways to compute the posterior in Bayesian inference:\n",
    "* Find a formula for the posterior\n",
    "* Obtain representative samples for the posterior\n",
    "\n",
    "The former approach might seem more sensible since we are used to understanding probability distributions based on their formulae. For example, the normal distribution is defined by its pdf,\n",
    "\n",
    "$$ p(x) = \\frac{1}{\\sqrt{2 \\pi}} e^{-x^2/2} $$\n",
    "\n",
    "But actually deriving useful information from such a formula can be hard. Even computing the mean and variance (or the normalisation constant $\\sqrt{2\\pi}$) for the normal distribution is not trivial, and the normal distribution is one of the simplest distributions we know. The posteriors we encounter in real life often have very complicated formulae, and doing any calculations based on these is going to be difficult.\n",
    "\n",
    "Perhaps counterintuitively, having a formula for a probability distribution does not mean we know how to sample from it, as we will see in the exercises.\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Using the pdf for the normal distribution, compute its mean and variance. Can you compute the mean of $\\sqrt{x}$?\n",
    "\n",
    "    *Hint: The normal distribution is symmetric about $0$, which allows you to determine the mean. For the variance, use integration by parts. Where do we need the normalisation constant $\\sqrt{2\\pi}$? For $\\sqrt{x}$, good luck.*\n",
    "\n",
    "2. *(For later)*. The Poisson distribution with rate $\\lambda$ has pdf\n",
    "\n",
    "$$ \\mathrm{Poi}(n; \\lambda) = e^{-\\lambda} \\frac{\\lambda^n}{n!} $$\n",
    "\n",
    "Think about how you would sample from the Poisson distribution - there are various possibilities. Carefully analyse to what extent knowledge of the pdf helps you in this task, versus special properties of the Poisson distribution."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gy-0xAbK-NM7"
   },
   "source": [
    "# Markov chains"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yR4TkRrl7dv3"
   },
   "source": [
    "## Definition\n",
    "\n",
    "A *Markov chain* is a memoryless random walk (in some space $E$).\n",
    "\n",
    "What does this mean? Our random walk starts with an initial position in our space $E$, which we call $X_0$. This $X_0$ may be pre-specified, or it may itself be random. In either case, once we have fixed $X_0$, we move step by step to positions $X_1$, $X_2$, $X_3$, $\\ldots$ in our space $E$, which comprise our walk.\n",
    "\n",
    "The individual positions we visit are random, so we cannot predict where we will eventually end up. The probability law for each step, however, is fixed and satisfies the following\n",
    "\n",
    "**Rule:** Given our current position $X_n$, the next position $X_{n+1}$ is independent of the positions before $X_n$.\n",
    "\n",
    "This is what memoryless means: where we go next does not depend on the past, only on the present.\n",
    "\n",
    "To define a Markov chain we need to define the initial distribution $p(X_0)$, as well as the *transition distributions* $p(X_{n+1} | X_n = x)$ for all possible positions $x$. Since $X_{n+1}$ only depends on $X_n$, but not any of the previous positions, this is all the information we need."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LOy-hGy8CFxA"
   },
   "source": [
    "## Marginal distributions\n",
    "\n",
    "Given the initial position $X_0$, we cannot predict where we will end up after $n$ steps, that is, we do not know $X_n$. We can nevertheless compute the probability distribution of $X_n$ for each $n$. These are called the **marginal distributions** of the Markov chain. The initial distribution $p(X_0)$ is a marginal distribution.\n",
    "\n",
    "The marginal distribution of $X_1$ is given by the following formula\n",
    "\n",
    "$$ p(X_1 = x) = \\sum_{y \\in E} p(X_1 = x \\, | \\, X_0 = y) \\, p(X_0 = y) $$\n",
    "\n",
    "This follows from the law of total probability: to get the probability of ending up at $x$, we just take the probability of moving to $x$ from any starting position $y$, times the probability of actually starting at $y$, and sum the results for all possible $y$.\n",
    "\n",
    "The marginal distribution of $X_2$ can be computed the same way if we know the distribution of $X_1$, and once we know the distribution of $X_2$, we can compute the marginal distribution of $X_3$, then $X_4$, etc. The marginal distributions can all be obtained from the same general formula, which we will call the update equation:\n",
    "\n",
    "$$ p(X_{n+1} = x) = \\sum_{y \\in E} p(X_{n+1} = x \\, | \\, X_{n} = y) \\, p(X_{n} = y) $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mdUJ66FeDUq2"
   },
   "source": [
    "## Exercises\n",
    "These questions assume some knowledge of linear algebra. Assume the states of the Markov chain are numbered $1, 2, \\ldots, k$.\n",
    "1. Show that any probability distribution on the states can be expressed as a $k$-dimensional vector $\\vec p$, where $p_i = p(i)$.\n",
    "2. Show that the transition distributions can be expressed as a $k\\times k$ matrix $\\bf T$, where $T_{x,y} = p(X_{n+1} = x \\, | \\, X_n = y)$.\n",
    "3. If the distribution of $X_n$ is given by the vector $\\vec p$, then the distribution of the next state $X_{n+1}$ is given by the matrix-vector product ${\\bf T} \\cdot \\vec p$.\n",
    "4. If the initial distribution is given by $\\vec p$, then the distribution of the $n$-th state $X_n$ is given by ${\\bf T}^n \\cdot \\vec p = {\\bf T} \\cdot {\\bf T} \\cdot \\ldots \\cdot {\\bf T} \\cdot \\vec p$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JYXQ12gS-o2I"
   },
   "source": [
    "## **Example:** Brownian motion\n",
    "\n",
    "Brownian motion, or the random movement of small particles in a solution, is an example of a Markov chain: if we measure the particle at fixed time intervals, then its path will be random, and knowing where the particle was in the past generally does not tell us anything about where it will move next.\n",
    "\n",
    "Here the space is $\\mathbb{R}^3$, our standard $3$-dimensional space. Between observations the molecule moves from its current position $X_n$ to a random new position $X_{n+1}$. How is $X_{n+1}$ distributed? For standard Brownian motion we make the assumption that $X_{n+1}$ equals $X_n$ plus a normally distributed displacement:\n",
    "\n",
    "$$ X_{n+1} = X_n + \\sqrt{D} \\cdot \\epsilon_n $$\n",
    "\n",
    "Here $\\epsilon_n$ follows the standard normal distribution. This is justified by the Central Limit Theorem, which says that summing up many small independent random displacements gives us a normal distribution. The diffusion parameter $D$ determines the scale of the displacement and measures the typical distance moved between steps.\n",
    "\n",
    "If the particle starts at the origin, then after $n$ steps the distribution of the particle will be normal (with variance $n \\cdot D$). Thus the marginal distributions of Brownian motion are normal, centered around the initial position."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "13iZCPT_YXqv",
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "def colorline(ax, x, y, colors=None, **kwargs):\n",
    "    \"\"\"\n",
    "    Fancy version of plt.plot, taken from\n",
    "    http://matplotlib.org/examples/pylab_examples/multicolored_line.html\n",
    "\n",
    "    Plot a colored line with coordinates x and y.\n",
    "    \"\"\"\n",
    "\n",
    "    # Default colors equally spaced on [0,1]:\n",
    "    if colors is None:\n",
    "        colors = np.linspace(0.0, 1.0, len(x))\n",
    "\n",
    "    # Special case if a single number:\n",
    "    if not hasattr(colors, \"__iter__\"):  # to check for numerical input -- this is a hack\n",
    "        colors = np.array([colors])\n",
    "\n",
    "    colors = np.asarray(colors)\n",
    "\n",
    "    points = np.array([x, y]).T.reshape(-1, 1, 2)\n",
    "    segments = np.concatenate([points[:-1], points[1:]], axis=1)\n",
    "    lc = matplotlib.collections.LineCollection(segments, array=colors, **kwargs)\n",
    "\n",
    "    ax.add_collection(lc)\n",
    "\n",
    "    return lc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 335
    },
    "id": "vwGwNIRhAeAL",
    "outputId": "d2960a88-9d4d-4dff-872f-49329b278057"
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, sharex=True, sharey=True, figsize=(6, 3))\n",
    "\n",
    "nsteps = 50\n",
    "Ds = [ 0.1, 1 ]\n",
    "\n",
    "for (ax, D) in zip(axes, Ds):\n",
    "  # Simulate Brownian motion for `nsteps` steps\n",
    "    Xs = np.zeros((nsteps+1, 2))\n",
    "    for n in range(nsteps):\n",
    "        Xs[n+1,:] = Xs[n,:] + np.sqrt(D) * np.random.randn(2)\n",
    "\n",
    "    colorline(ax, Xs[:,0], Xs[:,1])\n",
    "\n",
    "    ax.set_title(f\"D = {D}\")\n",
    "\n",
    "    ax.set_xlim(-8, 8)\n",
    "    ax.set_ylim(-8, 8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wKhoV9ZGaPp9"
   },
   "source": [
    "### Exercises\n",
    "\n",
    "Consider a particle undergoing Brownian motion with diffusion constant $D$.\n",
    "1. Show that the particle does not drift in any direction: $\\mathbb{E}[X_n] = 0$.\n",
    "2. Show that the average square displacement between steps is given by $\\mathbb{E}[\\left|X_{n+1} - X_n\\right|^2] = 3D$.\n",
    "\n",
    "*Hint: By Pythagoras' Theorem, for any vectors $\\vec a$ and $\\vec b$ in 3 dimensions,*\n",
    "\n",
    "$$\\left|\\vec a - \\vec b\\right|^2 = (a_x - b_x)^2 + (a_y - b_y)^2 + (a_z - b_z)^2 $$\n",
    "\n",
    "3. Show that the marginal distribution of $X_n$ is normal with variance $n \\cdot D$. How does this distribution change as $n$ becomes larger?\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Dda47ZCAaGrq"
   },
   "source": [
    "Markov chains are useful because they are easy to simulate, but can give rise to very complex marginal distributions. This provides an easy way to sample complex distributions: simply construct a Markov chain whose marginal distributions approach the desired target distribution. This is the essence of Markov chain Monte Carlo. In order to better understand MCMC we will study marginal distributions in some more detail."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2ric5UB3NQXq"
   },
   "source": [
    "## Stationary distributions\n",
    "\n",
    "In this section we focus on the marginal distributions $p(X_n)$ and how they evolve in time. If we start a Markov chain with initial distribution $X_0$, then the distribution of each subsequent step $X_n$ can be computed iteratively using the update equation above. It often happens that the marginal distributions $p(X_n)$ become more and more similar with each step, and eventually converge to a limiting distribution $\\pi$. This is called the **stationary**, or equilibrium distribution of the Markov chain.\n",
    "\n",
    "We can write $\\pi = p(X_\\infty)$, which is not mathematically rigorous, but captures the intuition."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8QQW5ANUT0lD"
   },
   "source": [
    "### Example: Brownian motion\n",
    "\n",
    "Let us illustrate stationary distributions using Brownian motion. If we place a small particle in a beaker of solution, then its location is well determined at first. But as time passes, the particle diffuses away and its position becomes harder to predict. If we wait long enough, the particle can be anywhere in the beaker - the  equilibrium distribution is uniform within the container. We can observe this by placing a droplet of ink in water: at first the droplet is very localised, but it slowly spreads out until eventually the ink is evenly distributed in the container, reflecting the uniform distribution of the individual ink molecules in the water.\n",
    "\n",
    "**Note:** The idealised Brownian motion we discussed above does not take place in a bounded region and does not have a stationary distribution. The marginal distributions simply become more and more spread out."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YSl2vc0gT78v"
   },
   "source": [
    "## An equation for the stationary distribution\n",
    "The stationary distribution $\\pi$ has an important property, namely\n",
    "\n",
    "**Stationarity:** If $X_n$ is distributed according to $\\pi$, then $X_{n+1}$ is also distributed according to $\\pi$.\n",
    "\n",
    "This makes sense: the distributions of the $X_n$ and $X_{n+1}$ approach each other more and more as $n$ becomes larger, so in the limit the distribution of $X_{\\infty+1}$ is exactly the same as $X_{\\infty}$.\n",
    "\n",
    "We can write this using our update equation as\n",
    "$$ \\pi(x) = \\sum_{y \\in E} p(X_{n+1} = x \\, | \\, X_{n} = y) \\, \\pi(y) $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jaV7B76AWdha"
   },
   "source": [
    "## Probability flow and detailed balance\n",
    "\n",
    "In general it may be difficult to compute the stationary distribution $\\pi$ using the last equation. If we consider Brownian motion in a container, it is not obvious that our droplet of ink will eventually be uniformly dispersed through the container, no matter its shape. But physically this is what happens, as can be observed by experiment.\n",
    "\n",
    "To understand the stationary distribution we can consider the motion of ink in the container. Brownian motion is symmetric, that is, the probability for a particle to move from position $x$ to $y$ in one step equals the probability of moving from $y$ to $x$. If the ink is uniformly distributed in the container, that means that the *total amount* of ink moving from $x$ to $y$ is the same as that moving from $y$ to $x$. In other words, while individual ink molecules will jump around from place to place, on a macroscopic scale these jumps cancel out exactly and no net movement occurs. If there is more ink in one region of the container, then more molecules will move out from that region and the ink distribution will change.\n",
    "\n",
    "This argument uses the probability flow $F(x \\leftarrow y)$, which measures the total amount of probability moving from state $y$ to $x$ and is given by\n",
    "\n",
    "$$ F(x \\leftarrow y) = p(X_n = y) \\, P(X_{n+1} = x \\, | \\, X_n = y) $$\n",
    "\n",
    "We can interpret the update equation in terms of the probability flow as follows: the probability that $X_n = x$ is the sum of the probability flowing into $x$ at step $n$, starting from the distribution $X_{n-1}$ at the previous step.\n",
    "\n",
    "We say that **detailed balance** occurs if the probability flux from any state $x$ to any other state $y$ is equal to the probability flux from $y$ to $x$:\n",
    "\n",
    "$$ F(x \\leftarrow y) = F(y \\leftarrow x) $$\n",
    "\n",
    "A distribution for which detailed balance holds is a stationary distribution of the Markov chain: the net probability flow between any two states is $0$.\n",
    "\n",
    "**Important:** Not all stationary distributions satisfy detailed balance.\n",
    "\n",
    "Detailed balance does not always hold, but when it does, it is usually easy to verify and makes a lot of calculations much simpler."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wPqQKXdCVpSt"
   },
   "source": [
    "# Markov Chain Monte Carlo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rDFWkO25VxOh"
   },
   "source": [
    "## Overview\n",
    "\n",
    "The idea behind Markov chain Monte Carlo is simple. Given an arbitrary target distribution $\\pi(x)$, we construct a Markov chain whose stationary distribution is $\\pi(x)$. We can then sample from $\\pi(x)$ by picking an initial value $X_0$ and simulating the Markov chain until we have reached the stationary distribution. We can formalise this as follows\n",
    "\n",
    "**Input:** Target distribution $\\pi(x)$.\n",
    "\n",
    "**Output:** Samples from the target distribution.\n",
    "\n",
    "**Steps:**\n",
    "1. Construct a Markov chain whose stationary distribution is $\\pi(x)$\n",
    "2. Pick an arbitrary initial value $X_0$\n",
    "3. Simulate steps $X_1, X_2, X_3, \\ldots$ until $p(X_N) \\approx p(X_\\infty) = \\pi$\n",
    "4. Return $X_N, X_{N+1}, X_{N+2}, \\ldots$\n",
    "\n",
    "The number of steps required until $p(X_N)$ resembles $\\pi$ is called the burn-in period. The samples during the burn-in period, depending on the initial value $X_0$, will not be distributed correctly and are discarded. Once the burn-in is complete, subsequent samples will be approximately distributed according to the target distribution.\n",
    "\n",
    "**Note:** The stationary distribution is an *asymptotic* limit, ie. $p(X_n)$ only approximates the target distribution for large $n$. Our algorithm is therefore not exact, but we can make the difference between $p(X_n)$ and $\\pi$ as small as desired by picking a suitably large burn-in period.\n",
    "\n",
    "**Note:** Each sample returned from the algorithm is (approximately) sampled from $\\pi(x)$, but the samples are not *independent*. Because we are simulating a random walk, each sample will be similar to the last sample. We may need a lot of samples to fully explore the target distribution.\n",
    "\n",
    "The main question we are left with is that of constructing the right Markov chain. There are a variety of ways to do this, but the most widespread one is the Metropolis-Hastings algorithm.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "r9SEKsSYWICV"
   },
   "source": [
    "## Exercises\n",
    "1. This exercise again uses some linear algebra. Assume the states of the Markov chain are numbered $1, 2, \\ldots, k$. If $\\bf T$ is the transition matrix, and the stationary distribution is represented by the vector $\\vec \\pi$, show that $\\vec \\pi = {\\bf T} \\cdot \\vec \\pi$. That is, $\\vec \\pi$ is an eigenvector of the transition matrix with eigenvalue $1$.\n",
    "\n",
    "2. Show that any distribution $\\pi(x)$ for which detailed balance holds is a stationary distribution.\n",
    "\n",
    "3. Show that for Brownian motion in a container, detailed balance holds for the uniform distribution, hence that the uniform distribution is the stationary distribution.\n",
    "\n",
    "4. *(Bonus)* Can you find a Markov chain with a stationary distribution that does not satisfy detailed balance?\n",
    "\n",
    "    *Hint: Look for a Markov chain with $3$ states.*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8Fb_8ZP9NA88"
   },
   "source": [
    "## The Metropolis-Hastings algorithm\n",
    "\n",
    "Assume we are given a target distribution $\\pi(x)$ and *any* Markov chain with transition probabilities $p(x \\, | \\, y)$. In general, $\\pi(x)$ will not be the stationary distribution of this Markov chain. But by tweaking the transition probabilities a little bit we can always arrange for that to be the case.\n",
    "\n",
    "**Notation:** In the following we denote the original transition properties by $q(x \\, | \\, y) = p(x \\, | \\, y)$. This notation is almost universal in the MCMC literature.\n",
    "\n",
    "The solution uses [Maxwell's demon](https://en.wikipedia.org/wiki/Maxwell%27s_demon), a small, fiendish creature that follows our particle through the Markov chain and gatekeeps its movement. Whenever the particle wants to move from state $x$ to another state $y$, there is a certain chance the demon blocks the move and the particle stays at $x$. If this happens we say that the move has been rejected, otherwise it is accepted and the particle moves to $y$.\n",
    "\n",
    "Denote by $\\alpha(y \\, | \\, x)$ the probability that a move from $x$ to another state $y$ is accepted by the demon. The transition probability from $x$ to $y$ changes from  $q(y \\, | \\, x)$ to $\\alpha(y \\, | \\, x) \\cdot q(y \\, | \\, x)$.\n",
    "\n",
    "Our target distribution satisfies detailed balance for this modified Markov chain if\n",
    "\n",
    "$$ \\frac{\\pi(x)}{\\pi(y)} = \\frac{q(x \\, | \\, y) \\cdot \\alpha(x \\, | \\, y)}{q(y \\, | \\, x) \\cdot \\alpha(y \\, | \\, x)} $$\n",
    "\n",
    "We can rearrange this as\n",
    "\n",
    "$$ \\frac{\\alpha(y \\, | \\, x)}{\\alpha(x \\, | \\, y)} = \\frac{\\pi(y)}{\\pi(x)} \\cdot \\frac{q(x \\, | \\, y)}{q(y \\, | \\, x)} $$\n",
    "\n",
    "Call the right-hand side $r(y, x)$. We can choose any set of acceptance probabilities between $x$ and $y$ as long as they are between $0$ and $1$ and satisfy this equation. The perhaps simplest way to achieve this is the following:\n",
    "* Choose the larger one of $\\alpha(y \\, | \\, x)$ and $\\alpha(x \\, | \\, y)$ to be $1$, and the other one to make the ratio work out.\n",
    "\n",
    "A more direct, but slightly complicated definition is:\n",
    "* If $r(y, x) \\leq 1$, let $\\alpha(y \\, | \\, x) = r(y, x)$ and $\\alpha(x \\, | \\, y) = 1$\n",
    "* If $r(y, x) \\geq 1$, let $\\alpha(y \\, | \\, x) = 1$ and $\\alpha(x \\, | \\, y) = 1 / r(y, x)$\n",
    "\n",
    "The quantity $\\alpha$ is called the acceptance ratio.\n",
    "\n",
    "In pseudocode, if we start at position $x$, to perform the next step we do the following:\n",
    "1. Sample a value $y$ from $q(y \\, | \\, x)$\n",
    "2. Compute the Metropolis-Hastings ratio\n",
    "$$\\alpha = \\frac{\\pi(y)}{\\pi(x)} \\cdot \\frac{q(x \\, | \\, y)}{q(y \\, | \\, x)}$$\n",
    "3. If $\\alpha > 1$, move to $y$. If $\\alpha \\leq 1$, move to $y$ with probability $\\alpha$, otherwise stay at $x$.\n",
    "\n",
    "**Note:** Since each step is subject to rejection by the demon, the Metropolis-Hastings algorithm will often remain in one same state for many steps. This duplication of states is expected behaviour.\n",
    "\n",
    "**Important:** The Metropolis-Hastings algorithm only uses the probability ratio $\\pi(y) / \\pi(x)$. We therefore do not have to know the normalising constant of $\\pi(x)$ as it cancels out. This makes Metropolis-Hastings very useful for Bayesian inference, where we can often compute the unnormalised posterior (prior times likelihood), but not the normalisation constant."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Stq4mFusgt4G"
   },
   "source": [
    "## Exercises\n",
    "\n",
    "1. Show that the Metropolis-Hastings algorithm describes a Markov chain that has $\\pi(x)$ as its stationary distribution.\n",
    "\n",
    "3. Show that for symmetric proposals with $q(y \\, | \\, x) = q(x \\, | \\, y)$ the acceptance ratio simplifies to\n",
    "\n",
    "$$ \\alpha = \\frac{\\pi(y)}{\\pi(x)} $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XFV_eP5_iQ-Z"
   },
   "source": [
    "## Metropolis-Hastings in action\n",
    "\n",
    "In this section we will revisit two examples - the coin toss, and the gene expression model - to see how MCMC provides an alternative way to compute the posteriors. We can write a utility function, `metropolis_hastings`, that does most of the work for us, and just plug in the numbers for the two examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "6uIQrXgSinjt"
   },
   "outputs": [],
   "source": [
    "# Generic implementation of Metropolis-Hastings\n",
    "def simulate_step(log_target, proposal_width, x):\n",
    "  \"\"\"\n",
    "  Perform one step of the MH algorithm starting at `x`, with a normal distribution\n",
    "  for the proposal. `log_target` is the logarithm of the target distribution\n",
    "  and `proposal_width` is the width of the proposal distribution. Returns the next\n",
    "  sample.\n",
    "  \"\"\"\n",
    "  # propose a new value\n",
    "  y = x + proposal_width * np.random.randn(len(x))\n",
    "\n",
    "  log_alpha = log_target(y) - log_target(x)     # symmetric proposal, so the q-terms cancel\n",
    "\n",
    "  if log_alpha > 0:\n",
    "    return y           # accept\n",
    "\n",
    "  alpha = np.exp(log_alpha)\n",
    "\n",
    "  p = np.random.rand()\n",
    "  if p < alpha:        # this happens with probability alpha\n",
    "    return y           # accept\n",
    "  else:\n",
    "    return x           # reject\n",
    "\n",
    "def metropolis_hastings(log_target, proposal_width, nsteps, x0):\n",
    "  \"\"\"\n",
    "  Perform MCMC by running the Metropolis-Hastings algorithm for `nsteps` steps,\n",
    "  starting at `x0`. `log_target` is the logarithm of the target distribution\n",
    "  and `proposal_width` is the width of the proposal distribution. Returns a matrix\n",
    "  of samples.\n",
    "  \"\"\"\n",
    "\n",
    "  ret = np.zeros((nsteps, len(x0)))\n",
    "  ret[0] = x0\n",
    "\n",
    "  for i in tqdm(range(1, nsteps)):\n",
    "    ret[i] = simulate_step(log_target, proposal_width, ret[i-1])\n",
    "\n",
    "  return ret"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "I7Aqpzw3qR1a"
   },
   "source": [
    "## Example: Flipping coins (again)\n",
    "\n",
    "We can use our helper function to perform MCMC for us. We choose $X_0 = 0.5$ for our initial position, although any number between $0$ and $1$ will be fine. The width of our proposal distribution is $0.1$ - while any positive number will work in theory, very small or very large values will result in slow convergence. We already have the (unnormalised) posterior density, and finally choose a burn-in period of $1000$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# Estimating the fairness of a coin based on repeated coin flips\n",
    "data_ct = [ 1, 1, 1, 1, 0, 1, 1 ]  # 1 for heads, 0 for tails\n",
    "n_heads = np.sum(data_ct)\n",
    "\n",
    "# Prior for the coin toss example\n",
    "def prior_ct(h):\n",
    "  return sp.stats.beta(4,4).pdf(h)\n",
    "\n",
    "# Likelihood function\n",
    "def likelihood_ct(data, h):\n",
    "  n_heads = sum(data)\n",
    "  n_tails = len(data) - n_heads\n",
    "\n",
    "  return h ** n_heads * (1-h) ** n_tails\n",
    "\n",
    "# Posterior\n",
    "# We will ignore the normalisation constant for now; it does not appear in the plot\n",
    "def posterior_unnormalised_ct(h, data):\n",
    "  return likelihood_ct(data, h) * prior_ct(h)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 104,
     "referenced_widgets": [
      "ee68c4885b5e4f879ba773c6fecb2ed1",
      "100f44b61d7b4872a5b83d04f52ac412",
      "d6903a72fab74f6cb3f61c54bda7e28a",
      "b2618b11fa1d4d549802c522a69532e0",
      "ba04e7b7c3584db29a4f3a407ae4dfa7",
      "49e7b48b057d4d998ce07d7cbbc79a6c",
      "617d171bce9841848d057fc43977e7f3",
      "2b74cbedb1224308a07a8fec3539a215",
      "18ef5c4162844f419803aec0b31f3859",
      "db4cf2cf35e84a828a87646a42af0965",
      "3e7fb2df9be540fea7f31f09b5980383"
     ]
    },
    "id": "pLsJiFWtoOPw",
    "outputId": "baf9c5ef-bbd8-4ce2-a16e-0bfeed7928e0"
   },
   "outputs": [],
   "source": [
    "def logposterior_unnormalised_ct(x):\n",
    "  h = x[0]\n",
    "\n",
    "  return np.log(posterior_unnormalised_ct(h, data_ct))\n",
    "\n",
    "nsamples = 10000\n",
    "burnin = 1000\n",
    "post_ct_mcmc = metropolis_hastings(logposterior_unnormalised_ct, 0.1, nsamples + burnin, np.array([0.5]))\n",
    "post_ct_mcmc = post_ct_mcmc[burnin:]        # We assume a burn-in period of 100\n",
    "\n",
    "# This is the actual posterior\n",
    "post_ct_exact = sp.stats.beta(4 + n_heads, 4 + len(data_ct) - n_heads).rvs(nsamples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 341
    },
    "id": "j7gL-wszkhoO",
    "jupyter": {
     "source_hidden": true
    },
    "outputId": "80e0b7a7-dda5-4245-d382-1bb94244d562"
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 2, width_ratios=[1, 2], figsize=(9, 3))\n",
    "ax[0].hist(post_ct_mcmc[:,0], range=(0, 1), bins=30, alpha=0.4, label=\"MCMC\")\n",
    "ax[0].hist(post_ct_exact, range=(0, 1), bins=30, alpha=0.4, label=\"Exact\")\n",
    "\n",
    "ax[0].set_xlabel(\"Head probability $h$\")\n",
    "ax[0].set_ylabel(\"Posterior probability\")\n",
    "ax[0].set_yticks([])\n",
    "ax[0].set_xlim(0, 1)\n",
    "\n",
    "ax[0].legend(loc=\"upper left\")\n",
    "\n",
    "ax[1].plot(post_ct_mcmc[:,0], alpha=0.5)\n",
    "ax[1].plot(post_ct_exact, alpha=0.4)\n",
    "\n",
    "ax[1].set_xlim(0, nsamples)\n",
    "ax[1].set_ylim(0, 1)\n",
    "\n",
    "ax[1].set_xlabel(\"Iterations\")\n",
    "ax[1].set_ylabel(\"Heads probability $h$\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZQDT2lG5rwa_"
   },
   "source": [
    "## Example: Gene expression\n",
    "\n",
    "To demonstrate the flexibility of MCMC we only have to change a few lines in this example. We only change the target density, the initial value (set to the prior mean) and in order to get somewhat faster convergence we set the proposal width to $0.3$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_ge = [ 20, 20, 19, 16, 15, 22, 17, 27, 17, 17, 21, 21, 16, 22, 25, 22, 23,\n",
    "            20, 21, 16, 16, 18, 16, 21, 17, 21, 25, 16, 15, 23 ]\n",
    "\n",
    "deg = 0.5     # degradation rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "b-ZFr8NoyYvv"
   },
   "outputs": [],
   "source": [
    "# Estimating the mRNA production rate based on snapshot data\n",
    "# Prior for the gene expression model\n",
    "def prior_ge(h):\n",
    "  return sp.stats.gamma(1, scale=5).pdf(h)\n",
    "\n",
    "# Likelihood\n",
    "def likelihood_ge(data, sigma):\n",
    "  ret = 1\n",
    "  for n in data:\n",
    "    ret *= sp.stats.poisson(sigma / deg).pmf(n)\n",
    "\n",
    "  return ret\n",
    "\n",
    "# Posterior (unnormalised, again)\n",
    "def posterior_unnormalised_ge(h, data):\n",
    "  return likelihood_ge(data, h) * prior_ge(h)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 70,
     "referenced_widgets": [
      "049765dc735049d5bf889de32251e48f",
      "1b7068940c35468e966731073f017dda",
      "c651e36e5e4b46bd8ebc4d14c770f88d",
      "5afcfb0381974f5d99708ea1f006925d",
      "d2ab3b0ac723472696899acc83c5b6b4",
      "f9e02edc69f34ac0b37b8bad4f110bf7",
      "49c28c2216844683a4c2b0d45ecfd66d",
      "3133d18172c149799c3624dab85007ff",
      "7d2e0db750734a008b6182ed56205edb",
      "26b0c079536f46a686c2925c66970fbe",
      "e11e8bf077ed444aab3fa660730103b4"
     ]
    },
    "id": "i9ZOENHgoukk",
    "outputId": "f3d00797-2b23-47ff-d1f5-cd434f302635"
   },
   "outputs": [],
   "source": [
    "def logposterior_unnormalised_ge(x):\n",
    "  rho = x[0]\n",
    "  return np.log(posterior_unnormalised_ge(rho, data_ge))\n",
    "\n",
    "nsamples = 1000\n",
    "burnin = 1000\n",
    "post_ge_mcmc = metropolis_hastings(logposterior_unnormalised_ge, 0.1, nsamples + burnin, np.array([5.]))\n",
    "post_ge_mcmc = post_ge_mcmc[burnin:]        # We assume a burn-in period of 100\n",
    "\n",
    "# This is the actual posterior\n",
    "post_ge_exact = sp.stats.gamma(1 + np.sum(data_ge), scale = 5 * deg / (5 * len(data_ge) + 1)).rvs(nsamples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 337
    },
    "id": "5aaiMKwNuJGp",
    "jupyter": {
     "source_hidden": true
    },
    "outputId": "d1743fe5-202c-40ef-bb33-e6bfd33f36b3"
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 2, width_ratios=[1, 2], figsize=(9, 3))\n",
    "\n",
    "ax[0].hist(post_ge_mcmc[:,0], range=(0, 20), bins=50, alpha=0.4, label=\"MCMC\")\n",
    "ax[0].hist(post_ge_exact, range=(0, 20), bins=50, alpha=0.4, label=\"Exact\")\n",
    "\n",
    "ax[0].set_xlabel(\"Production rate $\\\\rho$\")\n",
    "ax[0].set_ylabel(\"Posterior probability\")\n",
    "ax[0].set_yticks([])\n",
    "ax[0].set_xlim(0, 20)\n",
    "\n",
    "ax[0].legend(loc=\"upper left\")\n",
    "\n",
    "ax[1].plot(post_ge_mcmc[:,0], alpha=0.4)\n",
    "ax[1].plot(post_ge_exact, alpha=0.4)\n",
    "\n",
    "ax[1].set_xlim(0, nsamples)\n",
    "ax[1].set_ylim(0, 20)\n",
    "\n",
    "ax[1].set_xlabel(\"Iterations\")\n",
    "ax[1].set_ylabel(\"Production rate $\\\\rho$\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3vNgwVaM4PSb"
   },
   "source": [
    "## MCMC & The Chemical Master Equation\n",
    "\n",
    "We just illustrated a typical workflow when estimating model parameters:\n",
    "1. Define a quantitative model with some unknown parameters $\\theta$, and fix a prior $p(\\theta)$\n",
    "2. Collect experimental data\n",
    "3. Write down the likelihood function $p(\\mathcal{Data} \\, | \\, \\theta)$\n",
    "4. Optimise the likelihood to get parameter estimates **OR** compute posterior using MCMC\n",
    "\n",
    "Focussing on point 3., we were able to compute the likelihood function explicitly in the above example. For most reaction networks we cannot do that, and we need to find another way to compute the likelihoods.\n",
    "\n",
    "The experimental data consists of the number of molecules measured per cell at different time points. To compute the likelihood we thus need to know the probability of observing a fixed number of molecules in a cell. But this is precisely what the Chemical Master Equation computes. In short:\n",
    "\n",
    "**The Chemical Master Equation can be solved to compute the likelihood function.**\n",
    "\n",
    "Instead of using the Poisson solution in the last example, we can solve the Chemical Master Equation numerically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 49,
     "referenced_widgets": [
      "51f8b5381f85413493d48ad699c5250e",
      "82ee332afb9040588f98b93a7904c147",
      "c95476fe29a840afa12555d45d348e0a",
      "224ba782404f48418d1d09bbdb3cbdab",
      "2607ff2291014dd78caf9f5756289850",
      "f077de29ed984d69bb51014af09b7bb8",
      "d49771598e014eea87b7d29fd557c2df",
      "6f83308bb78942c886682db51c10e7d6",
      "cb09e596f4c9413aae8a38e7b3702289",
      "eaa2862229a44e1ba2e7ee24a808884a",
      "c24bec8e0263408186599c0880a3256f"
     ]
    },
    "id": "Mx2tPeOD8vhx",
    "outputId": "6fb2ff2c-3315-40ef-8894-52a0a85004f6"
   },
   "outputs": [],
   "source": [
    "def likelihood_ge_fsp(rho, data, Nmax=100):\n",
    "  # Build FSP matrix. We choose a truncation to 100 states.\n",
    "  A = build_fsp_matrix(rho, Nmax)\n",
    "\n",
    "  # Build initial conditions\n",
    "  u0 = np.zeros(Nmax)\n",
    "  u0[0] = 1                       # Assign probability 1 to the state with 0 mRNA\n",
    "\n",
    "  # Compute steady-state solution\n",
    "  # We approximate this by solving the system for large enough t\n",
    "  ut = solve_odes(A, u0, 100)\n",
    "\n",
    "  # The probability of observing n molecules in a cell is now ut[n].\n",
    "  ret = 1\n",
    "  for n in data:\n",
    "    ret *= ut[n]\n",
    "\n",
    "  return ret\n",
    "\n",
    "def build_fsp_matrix(rho, Nmax):\n",
    "  ret = np.zeros((Nmax, Nmax))\n",
    "\n",
    "  for i in range(Nmax):\n",
    "    if i+1 < Nmax:\n",
    "      ret[i+1,i] = rho          # mRNA production going from i to i+1\n",
    "    if i > 0:\n",
    "      ret[i-1,i] = i * deg      # mRNA degradation going from i to i-1\n",
    "\n",
    "    ret[i,i] = -(rho + i * deg)\n",
    "\n",
    "  return ret\n",
    "\n",
    "def solve_odes(A, u0, t):\n",
    "  return sp.linalg.expm(A * t) @ u0\n",
    "\n",
    "# Unnormalised posterior, computed using the FSP\n",
    "def logposterior_unnormalised_ge_fsp(x):\n",
    "  rho = x[0]\n",
    "  return np.log(prior_ge(rho)) + np.log(likelihood_ge_fsp(rho, data_ge))\n",
    "\n",
    "# Sample from the posterior using MCMC\n",
    "post_ge_mcmc = metropolis_hastings(logposterior_unnormalised_ge_fsp, 0.1, nsamples + burnin, np.array([5.]))\n",
    "post_ge_mcmc = post_ge_mcmc[burnin:]        # We assume a burn-in period of 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 320
    },
    "id": "uhiiapHH_QWl",
    "jupyter": {
     "source_hidden": true
    },
    "outputId": "99e07728-b844-4093-b3c6-ba74b24f5bee"
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 2, width_ratios=[1, 2], figsize=(9, 3))\n",
    "\n",
    "ax[0].hist(post_ge_mcmc[:,0], range=(0, 20), bins=50, alpha=0.4, label=\"MCMC\")\n",
    "ax[0].hist(post_ge_exact, range=(0, 20), bins=50, alpha=0.4, label=\"Exact\")\n",
    "\n",
    "ax[0].set_xlabel(\"Production rate $\\\\rho$\")\n",
    "ax[0].set_ylabel(\"Posterior probability\")\n",
    "ax[0].set_yticks([])\n",
    "ax[0].set_xlim(0, 20)\n",
    "\n",
    "ax[0].legend(loc=\"upper left\")\n",
    "\n",
    "ax[1].plot(post_ge_mcmc[:,0], alpha=0.4)\n",
    "ax[1].plot(post_ge_exact, alpha=0.4)\n",
    "\n",
    "ax[1].set_xlim(0, nsamples)\n",
    "ax[1].set_ylim(0, 20)\n",
    "\n",
    "ax[1].set_xlabel(\"Iterations\")\n",
    "ax[1].set_ylabel(\"Production rate $\\\\rho$\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "p-OvY6aJ_i2n"
   },
   "source": [
    "We get the same results as before. This version, however, takes much longer to run. This is because we have to solve the Chemical Master Equation once at each step of our MCMC chain, and doing so requires a fair bit of matrix computation on behalf of the computer. A system with $100$ states is not a big deal for modern computers, but more complex systems can cause a fair amount of frustration."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kRUlH0F7QrHT"
   },
   "source": [
    "## Exercises\n",
    "\n",
    "1. Write a function `metropolis_hastings(log_target, proposal_width, nsteps, x_0)`  that runs the Metropolis-Hastings algorithm for `nsteps` iterations starting at `x_0` and returns the sequence of states visited. The argument `log_target(x)` is a function that computes the log probability of the target distribution at its argument `x`. The proposal distribution at a position `x` is the normal distribution scaled by the factor `proposal_width`, that is, we start with Brownian motion with diffusion constant `proposal_width ** 2`.\n",
    "\n",
    "2. **a.** Using the above function, write a program that produces samples from the exponential distribution $\\textrm{Exp}(1)$. You can choose any `proposal_width` you like. How long is the burn-in period? Verify that your code works correctly by comparing a histogram of the returned samples with the exponential distribution.\n",
    "\n",
    "   **b.** Plot a time trace of your samples, and compare it with samples obtained directly using `scipy`. Can you see the difference between the samples obtained using MCMC and independent samples?\n",
    "\n",
    "   **c.** What do you observe if you choose a very small `proposal_width`? What about a very large `proposal_width`? Can you think of possible explanations for this behaviour?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5uZgTZ4cU33_"
   },
   "source": [
    "# The world of MCMC\n",
    "\n",
    "MCMC is commonly used in fields such as statistical inference, machine learning and statistical mechanics, and is likely to play an increasingly important role in the age of big data and big CPU/GPU. The original Metropolis-Hastings algorithm was developed in the 1950s, but there has been a lot of work on improving MCMC algorithms since. Some of these that are useful for Bayesian inference are:\n",
    "\n",
    "## Particle filters\n",
    "\n",
    "A lot of biological data is time-series data, ie. it consists of observations $n_1, n_2, n_3, \\ldots$ at different times $t_1, t_2, t_3, \\ldots$. There are many specialised MCMC methods for such problems, and they are generally known as [Sequential Monte Carlo (SMC)](https://arxiv.org/abs/2007.11936) methods.\n",
    "\n",
    "## Complex likelihoods\n",
    "\n",
    "Metropolis-Hastings requires us to compute the likelihood once per step, but in practice only a fraction of steps are accepted. This is wasteful, particularly when likelihood evaluations are expensive - the CME is a good example. One idea called [delayed acceptance](https://www.tandfonline.com/doi/abs/10.1198/106186005X76983) is to use a cheap approximation to the likelihood to eliminate samples that are likely to be rejected. We can also use machine learning to \"learn\" the likelihood function from a limited number of evaluations. Finally, we can sometimes use stochastic approximations of the likelihood that are cheaper to compute (ie., using Monte Carlo within MCMC) - this is called [pseudo-marginal](https://www.sciencedirect.com/science/article/pii/S0022519320301107) MCMC.\n",
    "\n",
    "## No likelihoods\n",
    "\n",
    "Sometimes likelihoods are so complicated that they are effectively uncomputable. In such cases, we can often probe the likelihood function using stochastic simulations to get an idea of how the system behaves for a certain $x$. This field is generally called simulator-based inference, or likelihood-free inference. [Approximate Bayesian Computation (ABC)](https://journals.plos.org/ploscompbiol/article?id=10.1371/journal.pcbi.1002803) is perhaps the most well-known method in this category, and it involves repeatedly simulating a system at various parameters to find those most consistent with the data. [Synthetic likelihoods](https://doi.org/10.1038/nature09319) are a more recent approach that has gained traction.\n",
    "\n",
    "## Efficient exploration\n",
    "\n",
    "Metropolis-Hastings relies on a random walk, often Brownian motion, to explore the posterior. As we know from biology, this can be a very slow way to move around, which has inspired a lot of research. [Hamiltonian Monte Carlo (HMC)](https://arxiv.org/abs/1701.02434) uses a clever combination of physical and statistical ideas to get around this limitation and perform massive jumps between samples, and is currently the gold standard for standard MCMC.\n",
    "\n",
    "## High-dimensional MCMC\n",
    "\n",
    "Some posteriors involve sampling thousands of parameters at once, even if we only care about a select few. Exploring high-dimensional spaces effectively is very difficult, and Brownian motion is particularly bad at this. [Gibbs sampling](https://towardsdatascience.com/gibbs-sampling-explained-b271f332ed8d) is a commonly used alternative for high-dimensional problems, but Hamiltonian Monte Carlo also scales fantastically to high dimensions.\n",
    "\n",
    "## Adaptive MCMC\n",
    "\n",
    "The performance of most MCMC algorithms is quite sensitive to tuning. Trying to find the right settings by hand for complex or high-dimensional posteriors is cumbersome. MCMC algorithms that adapt to the target distribution on the fly can yield massive improvements in speed while saving researchers a lot of time.\n",
    "\n",
    "## Probabilistic programming\n",
    "\n",
    "Probabilistic programming languages allow users to define a stochastic model in a straightforward and intuitive way, with the computer taking over the complicated business of converting this into a likelihood function and writing a MCMC sampler. [Stan](https://mc-stan.org/), [Pyro](https://pyro.ai/) (for Python) or [Turing.jl](https://turing.ml) (for Julia) are modern probabilistic programming languages that allow users to use state-of-the art adaptive Hamiltonian Monte Carlo sampling without having to fret over the details."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}