notebooks/Sampling_Gibbs.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run ../widgets/config_check.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# <api>\n",
    "import numpy as np\n",
    "from scipy.stats import norm, multivariate_normal\n",
    "from scipy.special import logsumexp\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')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Gibbs Sampling\n",
    "\n",
    "Gibbs sampling is a technique to sample from a multivariate distribution, e.g. $p(\\mathbf{x}) = p(x_1, \\ldots, x_D)$ where each sample is a $D$-dimensional vector $\\mathbf{x} \\in \\mathbb{R}^D$. Instead of sampling directly from the full joint distribution, Gibbs sampling considers each coordinate in turn and samples a new value from its conditional distribution, e.g.\n",
    "$$ p(x_d | x_1,\\ldots,x_{d-1},x_{d+1},\\ldots,X_D) = p(x_d | \\mathbf{x}_{-d}) $$\n",
    "\n",
    "This Gibbs sampling step can be considered as an adapted Metropolis-Hastings proposal which is always accepted. The multi-variate extension then draws on the observation that MCMC transitions can be combined, either in random or in sequential order, if each individual transition leaves the target density invariant. To see this, consider several transition densities $p_1(x'|x), \\ldots, p_K(x'|x)$ and mix them with probabilities $\\pi_k$, i.e.\n",
    "$$ p(x'|x) = \\sum_{k=1}^K \\pi_k p_k(x'|x) $$\n",
    "Then,\n",
    "$$ \\int p(x'|x) p^*(x) dx = \\int \\sum_{k=1}^K \\pi_k p_k(x'|x) p^*(x) dx = \\sum_{k=1}^K \\pi_k \\int p_k(x'|x) p^*(x) dx = \\sum_{k=1}^K \\pi_k p^*(x) = p^*(x) $$\n",
    "since each $p_k$ leaves the target density $p^*$ invariant.\n",
    "\n",
    "**Exercise**: Show that Gibbs sampling is a special type of Metropolis-Hastings sampling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# <api>\n",
    "class GibbsSampling (smpl_intro.Sampling):\n",
    "    def __init__ (self, transitions, x):\n",
    "        self.transitions = transitions\n",
    "        self.x = x\n",
    "        \n",
    "    def sample (self):\n",
    "        x = self.x[:]\n",
    "        K = len(self.transitions)\n",
    "        for i in range(K):\n",
    "            x[i] = self.transitions[i](x)\n",
    "        self.x = x\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Gibbs sampling requires that the conditional densities can be sampled from. It is most often used, when these can be computed in closed form, but it can also be combined with other methods such as slice sampling. Below is an illustration of Gibbs sampling for the 2-dimensional Gaussian example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# <api>\n",
    "def cond_prob (i, x):\n",
    "    return np.random.normal(loc=0.99*x[i], scale=np.sqrt(1-0.99**2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "sampling = GibbsSampling([lambda x: cond_prob(1,x), lambda x: cond_prob(0,x)], [1.5,0])\n",
    "\n",
    "samples = [sampling.sample() for _ in range(1000)]\n",
    "plt.plot(np.array(samples)[:,0], np.array(samples)[:,1], 'r-')\n",
    "plt.xlabel('x', fontsize=15)\n",
    "plt.ylabel('y', fontsize=15)\n",
    "plt.title('Sampling a 2D Gaussian via Gibbs sampling', fontsize=15)\n",
    "plt.contour(smpl_rej.X, smpl_rej.Y, smpl_rej.p2d.pdf(smpl_rej.XY));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Gibbs sampling is vastly popular as\n",
    "\n",
    "* it has no free parameters which need to be tuned\n",
    "* appears efficient as it moves on every sample, i.e. never rejects\n",
    "\n",
    "Gibbs sampling is most efficient if all coordinates are almost independent. In the case of high dependencies the conditional distributions are sharply peaked and Gibbs sampling moves slowly across the space. The above example nicely illustrates this effect as the length scale of each move is restricted to the standard deviation $\\sqrt{1 - 0.99^2} = 0.141$ of the conditional distribution while the actual distribution has a much larger extend."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "__Example__: Gibbs sampling for linear regression\n",
    "\n",
    "Let $\\mathbf{X}$ denote the design matrix, i.e. values of $D$ regressors (plus bias input) for $N$ training cases collected in a $N \\times D$ matrix, and $\\mathbf{t} \\in \\mathbb{R}^N$ denote the target outputs.\n",
    "\n",
    "Linear regression then models the targets as\n",
    "$$ p(\\mathbf{t} | \\mathbf{X}, \\mathbf{w}, \\tau) = \\prod_{i=1}^N \\mathcal{N}(t_i ; \\mathbf{w}^T \\mathbf{x}_i, \\tau) $$\n",
    "with weight vector $\\mathbf{w}$ and observation noise precision $\\tau = \\frac{1}{\\sigma^2}$.\n",
    "\n",
    "With conjugate priors the conditional posterior distributions can be computed analytically:\n",
    "* $\\mathbf{w}$: Conjugate prior is Gaussian with mean $\\mathbf{0}$ and covariance $\\Sigma_0$\n",
    "  \n",
    "  Posterior is also Gaussian with covariance\n",
    "  $$ \\Sigma_N = ( \\Sigma_0^{-1} + \\tau \\mathbf{X}^T \\mathbf{X} )^{-1} $$\n",
    "  and mean\n",
    "  $$ \\mu_N = \\tau \\Sigma_N \\mathbf{X}^T \\mathbf{t} $$\n",
    "* $\\tau$: Conjugate prior is Gamma with shape $\\alpha_0$ and rate $\\beta_0$\n",
    "\n",
    "  Posterior is again a Gamma distribution with shape $\\alpha_N = \\alpha_0 + \\frac{N}{2}$ and rate $\\beta_N = \\beta_0 + \\frac{1}{2} (\\mathbf{t} - \\mathbf{X} \\mathbf{w})^T (\\mathbf{t} - \\mathbf{X} \\mathbf{w})$\n",
    "  \n",
    "The program below implements the Gibbs updates according to the above formulas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# <api>\n",
    "def lm_Gibbs_w (sample, X, t, Sigma_0_inv):\n",
    "    # Compute posterior covariance\n",
    "    tau = sample[1]\n",
    "    Sigma_N = np.linalg.inv( Sigma_0_inv + tau*np.dot(X.T, X) )\n",
    "    # and mean\n",
    "    mu_N = tau*np.dot(Sigma_N, np.dot(X.T, t))\n",
    "    # Draw new weight vector for next sample\n",
    "    w_new = np.random.multivariate_normal(mu_N.squeeze(), Sigma_N, size=1)\n",
    "    return w_new.T\n",
    "\n",
    "def lm_Gibbs_tau (sample, X, t, alpha_0, beta_0):\n",
    "    # Compute posterior shape\n",
    "    N, D = X.shape\n",
    "    alpha_N = alpha_0 + 0.5*N\n",
    "    # and rate\n",
    "    w = sample[0]\n",
    "    err = t - np.dot(X, w)\n",
    "    beta_N = beta_0 + 0.5*np.dot(err.T, err).squeeze()\n",
    "    # Draw new precision for next sample\n",
    "    tau_new = np.random.gamma(alpha_N, 1.0/beta_N, size=1)\n",
    "    return tau_new"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Small demo data set\n",
    "N = 250\n",
    "x = np.hstack([np.ones((N,1)), 5+np.random.normal(size=(N,1))])\n",
    "t = (np.dot(x, np.array([-1,2])) + np.random.normal(size=(N,), scale=0.8))[:,np.newaxis]\n",
    "plt.plot(x[:,1], t, 'k.');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Hyper parameters\n",
    "N, D = x.shape\n",
    "alpha_0 = beta_0 = 0.01;\n",
    "Sigma_0_inv = np.linalg.inv( np.eye(D) )\n",
    "\n",
    "# Gibbs samples\n",
    "w = np.reshape(np.zeros(D), (D,1))\n",
    "sampling = GibbsSampling([lambda s: lm_Gibbs_w(s, x, t, Sigma_0_inv), \n",
    "                          lambda s: lm_Gibbs_tau(s, x, t, alpha_0, beta_0)],\n",
    "                         [w, 1.0])\n",
    "\n",
    "samples = [sampling.sample() for _ in range(1000)]\n",
    "\n",
    "# Plot histogramm of weights and precision\n",
    "plt.figure(figsize=(8, 8))\n",
    "ax = plt.subplot(311)\n",
    "plt.hist(list(map(lambda x: float(x[0][0]), samples)))\n",
    "plt.title('intercept')\n",
    "plt.subplot(312, sharex=ax)\n",
    "plt.hist(list(map(lambda x: float(x[0][1]), samples)))\n",
    "plt.title('w_1')\n",
    "plt.subplot(313)\n",
    "plt.hist(list(map(lambda x: 1.0/np.sqrt(float(x[1])), samples)))\n",
    "plt.title('sigma');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise**: Define a Gibbs sampler which samples each weight individually. What are the advantages/disadvantages of this approach?"
   ]
  }
 ],
 "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
}