notebooks/Sampling_MCMC.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run ../widgets/config_check.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# <api>\n",
    "import numpy as np\n",
    "from scipy.stats import norm, multivariate_normal\n",
    "from matplotlib import pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "# Load classes and functions from the previous parts\n",
    "from jupyter_cms.loader import load_notebook\n",
    "smpl_intro = load_notebook('./Sampling_Intro.ipynb')\n",
    "smpl_rej = load_notebook('./Sampling_Rejection.ipynb')\n",
    "\n",
    "import os\n",
    "path = os.getcwd()\n",
    "s = '/'\n",
    "pardir = s.join(path.split(s)[:-1])\n",
    "# Load widgets\n",
    "smpl_widgets = load_notebook(str(pardir + '/widgets/2D_sampling_widget.ipynb'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Markov chain Monte-Carlo (MCMC)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "The previous sampling methods produced independent samples from the target distribution. Markov Chain Monte Carlo (MCMC) methods instead produce a sequence of dependent samples which nevertheless has a desired marginal distribution.\n",
    "\n",
    "A sequence of random variable $X_1, X_2, \\ldots$ is called a Markov chain if their joint density can be factorized as follows:\n",
    "$$ p(x_1, x_2, \\ldots) = p(x_1) \\prod_{i=1}^{\\infty} p(x_{i+1} | x_1, \\ldots, x_{i}) = p(x_1) \\prod_{i=1}^{\\infty} p(x_{i+1} | x_i) $$\n",
    "Thus, each variable is independent of its history given its direct predecessor. A Markov chain is called **_homogeneous_** if its transition density $p(x_{i+1}|x_i)$ is the same for all $i \\in \\mathbb{N}$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Detailed Balance\n",
    "\n",
    "MCMC constructs a Markov chain with a desired stationary distribution $p(x)$. The stationary distribution $p^*(x)$ of a homogeneous Markov chain with transition density $p(x' | x)$ has the property that\n",
    "$$ p^*(x') = \\int p(x' | x) p^*(x) dx, $$\n",
    "i.e. the distribution is unchanged by the action of the Markov chain.\n",
    "\n",
    "A sufficient condition to ensure that $p(x)$ is invariant is that the transition probability satisfies **_detailed balance_**:\n",
    "$$ p^*(x) p(x' | x) = p^*(x') p(x | x') $$\n",
    "Note that this is not tautological as it requires that $p^*$ is the same on both sides.\n",
    "\n",
    "Then,\n",
    "$$ \\int p(x' | x) p^*(x) dx = \\int p^*(x') p(x | x') dx = p^*(x') \\int p(x | x') dx = p^*(x') $$\n",
    "A Markov chain which satisfies detailed balance is **_reversible_**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ideally, we want that the chain converges to its target distribution $p^*(x)$ independent of where it started. This is ensured by the following technical conditions:\n",
    "\n",
    "* **Ergodic**: A Markov chain is called *ergodic* if every state $x'$ can be reached from any other state $x$, i.e. $\\forall x,x' \\; \\exists n \\mbox{ s.t. } p(X_n = x' | X_1 = x) > 0$\n",
    "\n",
    "  This condition is also called *irreducibility* and implies the existence of a unique invariant measure $p^*(x)$. Further, from the ergodic theorem we get the important property that time averages converge to ensemble averages, i.e.\n",
    "  $$ \\frac{1}{T} \\sum_{i=1}^T f(x_i) \\underset{T \\to \\infty}{\\longrightarrow} \\mathbb[E]_{p^*}[f]  $$\n",
    "\n",
    "* **Aperiodic**: A Markov chain is called *aperiodic* if the period of all its state is 1. The period of state $x$ is defined as\n",
    "$$ k = \\mathtt{gcd}\\{n > 1: p(X_n = x | X_1 = x) > 0\\} $$\n",
    "where $\\mathtt{gcd}$ is the greatest common divisor. Thus, in an aperiodic chain each state can be revisited at irregular intervals.\n",
    "\n",
    "When the Markov chain is irreducible and aperiodic it is called *mixing* and it converges to its invariant distribution from any starting distribution. In practice it is usually easy to construct construct a mixing chain, e.g. by chosing a transition density which has full support, but the mixing time can be very long."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Metropolis-Hastings\n",
    "\n",
    "The Metropolis-Hastings algorithm draws a new sample from $p(x)$ starting from a sample $x$ as follows:\n",
    "\n",
    "* Draw $x'$ from a proposal distribution $q(x'|x)$\n",
    "* With probability $$a_{x,x'} = \\min\\{1, \\frac{p(x') q(x|x')}{p(x) q(x'|x)}\\}$$ accept $x'$ as the new sample. Otherwise stay at $x$ and return it as the new sample.\n",
    "\n",
    "**Exercise:** Show that this Markov chain satisfies detailed balance with respect to $p(x)$\n",
    "\n",
    "Note that the acceptance condition simplifies when the proposal distribution is symmetric, i.e. $q(x'|x) = q(x|x')$. In this case, the algorithm is also known as the _Metropolis algorithm_ with the intuitive acceptance probability $\\min\\{1, \\frac{p(x')}{p(x)}\\}$.\n",
    "\n",
    "The algorithm also works if the normalization constant of $p(x)$ is unavailable as it cancels in the acceptance condition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# <api>\n",
    "class Proposer (object):\n",
    "    \"\"\"\n",
    "    Wraps two functions needed by a proposer q, i.e.\n",
    "      Draw a new value y ~ q.propose(x)\n",
    "      Compute log transition probability q.log_trans_prob(x,y) \n",
    "    \"\"\"\n",
    "    def __init__(self, propose, log_trans_prob):\n",
    "        self.propose = propose\n",
    "        self.log_trans_prob = log_trans_prob\n",
    "        \n",
    "    def propose(self, x):\n",
    "        return self.propose(x)\n",
    "    \n",
    "    def log_trans_prob(self, x, y):\n",
    "        return self.log_trans_prob(x, y)\n",
    "\n",
    "class MetropolisHastings (smpl_intro.Sampling):\n",
    "    def __init__(self, log_p, q, x):\n",
    "        \"\"\"\n",
    "        q is assumed to be a proposer and log_p computes log p(x)\n",
    "        \"\"\"\n",
    "        self.x = x # Current sample\n",
    "        self.log_p = log_p\n",
    "        self.q = q\n",
    "        self.num_samples = 0\n",
    "        self.accepted = 0\n",
    "        \n",
    "        \n",
    "    def __str__ (self):\n",
    "        return \"Metropolis Hastings: Accepted %d out of %d samples\" % (self.accepted, self.num_samples)\n",
    "        \n",
    "    def sample (self):\n",
    "        self.num_samples += 1\n",
    "        # Propose new candidate\n",
    "        x_prime = self.q.propose(self.x)\n",
    "        a = self.log_p(x_prime) + self.q.log_trans_prob(x_prime, self.x) \\\n",
    "            - self.log_p(self.x) - self.q.log_trans_prob(self.x, x_prime)\n",
    "        u = np.random.uniform()\n",
    "        if np.log(u) < a:\n",
    "            self.accepted += 1\n",
    "            self.x = x_prime\n",
    "        return self.x    \n",
    "    \n",
    "    \n",
    "    def reset_counters(self):\n",
    "        self.num_samples = 0\n",
    "        self.accepted = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "sampling = MetropolisHastings(log_p=lambda x: np.log(smpl_rej.p(x)), \n",
    "                              q=Proposer(lambda x: norm(loc=x, scale=0.5).rvs(),\n",
    "                                         lambda x, y: norm(loc=x, scale=0.5).logpdf(y)),\n",
    "                              x=-3.0)\n",
    "\n",
    "smpl_intro.show_sampling(sampling, plotter=smpl_intro.gauss_hist, N=2500)\n",
    "print(sampling)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using again the more difficult 2-dimensional example, we can illustrate how the sampler moves around the target density:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "sampling = MetropolisHastings(log_p=smpl_rej.p2d.logpdf, \\\n",
    "                              q=Proposer(lambda x: multivariate_normal(mean=x, cov=1e2*np.eye(2)).rvs(),\n",
    "                                         lambda x,y: 0), # Proposal is symmetric\n",
    "                              x=[1.5,0])\n",
    "samples = [sampling.sample() for _ in range(1000)]\n",
    "samples\n",
    "\n",
    "plt.plot(np.array(samples)[:,0], np.array(samples)[:,1], 'r-')\n",
    "plt.contour(smpl_rej.X, smpl_rej.Y, smpl_rej.p2d.pdf(smpl_rej.XY))\n",
    "plt.xlabel('x', fontsize=15)\n",
    "plt.ylabel('y', fontsize=15)\n",
    "plt.title('Sampling a 2D Gaussian via Metropolis-Hastings', fontsize=15);\n",
    "print(sampling)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This example also illustrates the importance of choosing a suitable proposal distribution $q(x'|x)$. If $q$ is tightly concentrated almost all proposals are accepted, but the chain moves very slowly. Contrarily, if $q$ is very broad the acceptance rate is low since most proposals fall into regions of low probability and the chain moves again slowly. In both cases, the mixing time of the chain is very long and we would need many samples to estimate reliable averages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# <api>\n",
    "def mixing_demo (q_var):\n",
    "    n = len(q_var)\n",
    "    plt.figure(figsize=(14, 5))\n",
    "    for i in range(n):\n",
    "        q = Proposer(lambda x: multivariate_normal(mean=x, cov=np.sqrt(q_var[i])*np.eye(2)).rvs(),\n",
    "                     lambda x,y: 0) # Proposal is symmetric\n",
    "        sampling = MetropolisHastings(log_p=smpl_rej.p2d.logpdf, q=q, x=[1.5,0])\n",
    "        samples = [sampling.sample() for _ in range(1000)]\n",
    "        plt.subplot(1,n,i+1)\n",
    "        \n",
    "        plt.title('Sdev[q] = ' + str(np.sqrt(q_var[i])), fontsize = 15)\n",
    "        plt.plot(np.array(samples)[:,0], 'b-')\n",
    "        plt.xlabel('Number of samples', fontsize=15)\n",
    "        plt.ylabel('Sample value', fontsize=15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mixing_demo(np.array([0.001, 0.1, 10])**2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Is there an optimal proposal or acceptance rate?\n",
    "\n",
    "* Obviously, choosing $q(x'|x) = p(x')$ would be optimal as it is always accepted and produces independent samples ...\n",
    "* When the proposal is a circular Gaussian an optimal acceptance rate of 0.234 has been derived.\n",
    "\n",
    "As a rule of thumb, the acceptance rate is often adjusted to around 0.2. But as an information theoretic argument by MacKay suggests, Metropolis-Hastings sampling is always rather inefficient as it takes a random walk across the space. In general, when the probable states of $p(x)$ cover a length scale of $L$ and the proposal takes steps of size $\\epsilon$, we need at least $(\\frac{L}{\\epsilon})^2$ steps to obtain an (almost) independent sample. This is based on the observation that a random walk/diffusion with step size/standard deviation $\\epsilon$ covers a distance of about $\\sqrt{T} \\epsilon$ after $T$ steps.\n",
    "\n",
    "Thus, many algorithms try to utilize proposals adjusted to the target distribution in order to improve mixing, speeding up sampling.\n",
    "\n",
    "**Exercise:** Illustrate the effect of the proposal width when sampling from a multi-modal target distribution."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convergence diagnostics\n",
    "\n",
    "Most MCMC algorithms, e.g. Metropolis-Hastings, produce a Markov chain sampling from the desired target distribution. Thus, all samples produced are correct in principal, but the intial samples can depend strongly on the initialization especially if the chain mixes slowly. It is common practice to drop an initial part of the samples, the so called *burn-in* period.\n",
    "\n",
    "In general, it would be desirable to have diagnostic tools in order to access the quality of samples produced and detect sampling problems. In particular, one wants to know how well the chain mixes and how many (almost) inpedendent have been produced. In the literature, several diagnostics have been proposed ... here, some of them are listed:\n",
    "\n",
    "* Traceplot: Always look at your samples as problems often show up visually (see mixing example above)\n",
    "* Autocorrelation: Similar to importance sampling an estimate of the effective sample size can be computed. In the case of MCMC, the sample size is reduced as the samples are dependent and thus empirical averages are less reliable when taken over dependent samples. A commonly used measure of effective sample size is defined as follows:\n",
    "$$ N_{eff} = \\frac{N}{1 + 2 \\sum_{\\tau=1}^{\\infty} \\rho_{\\tau}}$$\n",
    "where $\\rho_{\\tau}$ denotes the correlation coefficient at lag $\\tau$.\n",
    "\n",
    "  This is motivated from the sample mean of a stationary process $X_1, X_2, \\ldots$. Then, the sample mean $\\hat{\\mu}_N = \\frac{1}{N} \\sum_{i=1}^N X_i$ has mean $$\\mathbb{E}[\\hat{\\mu}_N] = \\mathbb{E}[X]$$ and variance $$\\mathbb{V}ar[\\hat{\\mu}_N] = \\frac{1}{N} \\sum_{\\tau=-N}^N (1 - \\frac{|\\tau|}{N}) \\gamma_{\\tau}$$ where $\\gamma_{\\tau} = \\mathbb{C}ov[X_t, X_{t+\\tau}]$.\n",
    "  \n",
    "  Now, when $\\sum_{\\tau} |\\gamma_{\\tau}| < \\infty$, we see that the variance converges as \n",
    "  $$N \\mathbb{V}ar \\to_{\\infty} \\sum_{\\tau = -\\infty}^{\\infty} \\gamma_{\\tau} = \\sigma^2 \\sum_{\\tau = -\\infty}^{\\infty} \\rho_{\\tau}$$\n",
    "  and compared to the case of independent samples the variance is increased by a factor of $1 + 2 \\sum_{\\tau=1}^{\\infty} \\rho_{\\tau}$.\n",
    "  \n",
    "  In practice, the autocorrelation function has to be estimated which is often done by fitting an AR(1) model to the produced samples. Obviously, plotting the autocorrelation function is also useful to access the quality of MCMC samples.\n",
    "* Convergence: To diagnose convergence, i.e. check if all samples can be considered as being drawn from the same distribution, the following two methods are commonly used:\n",
    "    * Geweke proposed to compare the mean of two non-overlapping parts of the produced samples, e.g. the first 10% compared to the last 50%. Convergence is then diagnosed if a t-test cannot reject that the two means are the same.\n",
    "    * Gelman and Rubin proposed a more sophisticated method which assumes that $m$ MCMC-chains have been simulated. Then, they define the *scale reduction factor*\n",
    "$$ \\hat{R} = \\sqrt{\\frac{\\hat{\\sigma}^2}{W}} $$\n",
    "This is based on the observation that the variance can be estimated in two ways: As the mean of the variances within each chain, $W$ or as the variance of all samples $\\hat{\\sigma}^2$. When the chain has not converged $W$ tends to underestimate the true variance as not all of the sampling space has been explored, while $\\hat{\\sigma}^2$ tends to overstate the variance if the initial conditions are chosen at arbitrary, over-dispersed positions. Thus, convergence is detected by $\\hat{R}$ being close to one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "widget = smpl_widgets.SamplingWidget(method='MHS', target='MN')\n",
    "widget.show()"
   ]
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}