{ "cells": [ { "cell_type": "markdown", "id": "691d7d4f", "metadata": {}, "source": [ "# Sampling by MCMC method\n", "\n", "The goal here is to learn how to code a Metropolis, HMC and lifted Metropolis sampler \n", "\n", "Library importation" ] }, { "cell_type": "code", "execution_count": 1, "id": "83fcff75", "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "import matplotlib.pyplot as plt" ] }, { "cell_type": "markdown", "id": "24c4e9a0", "metadata": {}, "source": [ "## 1. Metropolis scheme\n", "\n", "First, let's implement a Metropolis algorithm. Each step will depend on the target distribution ```target``` (or ```logtarget```) and on the way to propose a new state (```proposal```), which we'll choose to be symmetrical ($p(x\\o x') = p(x' \\to x)$), for example by adding to $x$ a $\\delta$ drawn uniformly or according to a Gaussian. We also need to define an initialization function and a number of iterations to be performed. \n", "\n", "1. Before you start coding, how do you decide whether to accept a state?\n", "2. If rejected, which state is added to the set of samples?\n", "3. You can code the Metropolis algorithm directly for a particular target distribution (e.g. Gaussian), at least in dimension 2. It may be interesting to display the acceptance rate" ] }, { "cell_type": "code", "execution_count": 2, "id": "9c793687", "metadata": {}, "outputs": [], "source": [ "def proposal(x):\n", " scale = 1.0\n", " jump = np.random.normal(scale=1.0, size=x.shape)\n", " return x + jump\n", "\n" ] }, { "cell_type": "code", "execution_count": 3, "id": "cbd9d32e", "metadata": {}, "outputs": [], "source": [ "mu=4.\n", "scale=1.0\n", "def neglogtarget(x):\n", " return np.sum((x-mu) ** 2.0) / (2.0 * scale ** 2.0)" ] }, { "cell_type": "code", "execution_count": 4, "id": "44f12ecd", "metadata": {}, "outputs": [], "source": [ "samples = [np.array([3.0,2.0])]\n", "iterations = 100000\n", "accept = 0.0\n", "for _ in range(iterations):\n", " current = samples[-1]\n", " proposed = proposal(current)\n", " if np.random.random() < np.exp(-(neglogtarget(proposed)-neglogtarget(current))):\n", " samples.append(proposed)\n", " accept += 1.\n", " else:\n", " samples.append(current)\n", "samples = np.array(samples)\n", "print(accept/float(iterations))" ] }, { "cell_type": "markdown", "id": "8f0ec87a", "metadata": {}, "source": [ "4. Trace plot $||x||$.\n", "5. Draw a section plot" ] }, { "cell_type": "code", "execution_count": 5, "id": "c08896f7", "metadata": {}, "outputs": [], "source": [ "plt.plot(np.cumsum(np.array([sum(s**2.0)**0.5 for s in samples]))/(1.0+np.arange(len(samples))))\n", "plt.xlabel('t')\n", "plt.ylabel('')\n", " " ] }, { "cell_type": "code", "execution_count": 6, "id": "743ed37a", "metadata": {}, "outputs": [], "source": [ "plt.scatter(samples[:,0], samples[:,1], alpha=0.1)\n", "plt.xlabel('x_0')\n", "plt.ylabel('x_1')" ] }, { "cell_type": "markdown", "id": "7afc74cc", "metadata": {}, "source": [ "## 2. Hybrid or Hamiltonian MC scheme\n", "\n", "6. We now wish to code a Hamiltonian Monte Carlo:\n", "- we need to code a Newtonian integrator ```leapfrog``` which takes as parameters the number of steps ```L``` and their amplitude ```step size```.\n", "- we need to define the ```neglogtarget``` gradient.\n", "- Consider a kinetic energy $K(p)=p^2$.\n", "\n", "\n", "\n", "\n", "\n", "For this, we consider the leapfrog in the pseudocode." ] }, { "cell_type": "code", "execution_count": 7, "id": "0bf17df9", "metadata": {}, "outputs": [], "source": [ "L = 10\n", "def leapfrog(q0, p0):\n", " q = q0.copy()\n", " p = p0.copy()\n", " step_size = 0.5\n", " for i in range(L):\n", " p -= grad_neglogtarget(q) * step_size / 2\n", " q += p * step_size\n", " p -= grad_neglogtarget(q) * step_size / 2\n", " return q, p\n" ] }, { "cell_type": "code", "execution_count": 8, "id": "67cef686", "metadata": {}, "outputs": [], "source": [ "def grad_neglogtarget(x):\n", " return (x-mu) / (scale ** 2.0)" ] }, { "cell_type": "code", "execution_count": 30, "id": "2066f9a7", "metadata": {}, "outputs": [], "source": [ "samples2 = [np.array([3.0,2.0])]\n", "iterations = 100000\n", "accept = 0.0\n", "for _ in range(iterations):\n", " q0 = samples2[-1]\n", " p0 = np.random.standard_normal(size=q0.size)\n", "\n", " q_star, p_star = leapfrog(q0, p0)\n", "\n", " h0 = neglogtarget(q0) + np.dot(p0,p0) / 2.0\n", " h = neglogtarget(q_star) + np.dot(p_star,p_star) / 2.0\n", " log_accept_ratio = h - h0\n", "\n", " if np.random.random() < np.exp(-log_accept_ratio):\n", " samples2.append(q_star)\n", " accept += 1.\n", " else:\n", " samples2.append(q0)\n", "\n", "\n", "samples2 = np.array(samples2)\n", "print(accept/float(iterations))\n", "\n" ] }, { "cell_type": "markdown", "id": "490176b5", "metadata": {}, "source": [ "In comparison with Metropolis results:\n", "\n", "7. Plot the trace plot ||𝑥||\n", "8. Plot a section plot" ] }, { "cell_type": "code", "execution_count": 31, "id": "20b6b781", "metadata": {}, "outputs": [], "source": [ "plt.plot(np.cumsum(np.array([sum(s**2.0)**0.5 for s in samples]))/(1.0+np.arange(len(samples))),ls='--',label='Metropolis')\n", "plt.plot(np.cumsum(np.array([sum(s**2.0)**0.5 for s in samples2]))/(1.0+np.arange(len(samples2))),label='HMC')\n", "plt.xlabel('t')\n", "plt.ylabel('')\n", "plt.legend() " ] }, { "cell_type": "code", "execution_count": 32, "id": "859eef1a", "metadata": { "scrolled": true }, "outputs": [], "source": [ "plt.scatter(samples2[:,0], samples2[:,1], alpha=0.1)\n", "plt.xlabel('x_0')\n", "plt.ylabel('x_1')" ] }, { "cell_type": "markdown", "id": "b1e5b298", "metadata": {}, "source": [ "9. Draw a trajectory for HMC and Metropolis (10 points, 100 points)" ] }, { "cell_type": "code", "execution_count": 33, "id": "a01fb6c1", "metadata": {}, "outputs": [], "source": [ "plt.plot(samples2[-100:][:,0],samples2[-100:][:,1],'o-',label='HMC')\n", "plt.plot(samples[-100:][:,0],samples[-100:][:,1],'o-',label='Metropolis')\n", "plt.xlabel('x_0')\n", "plt.ylabel('x_1')\n", "plt.legend() " ] }, { "cell_type": "markdown", "id": "087e92f5", "metadata": {}, "source": [ "10. Check that the distributions of $x_0$ and/or $x_1$ match." ] }, { "cell_type": "code", "execution_count": 34, "id": "20bcf5b3", "metadata": {}, "outputs": [], "source": [ "n1,b1,v1=plt.hist(samples[:,0],bins=1000,histtype='step',label='Metropolis',cumulative=True)\n", "n2,b2,v2=plt.hist(samples2[:,0],bins=1000,histtype='step',label='HMC',cumulative=True)\n", "plt.xlabel('x_0')\n", "plt.legend() " ] }, { "cell_type": "markdown", "id": "957f64bb", "metadata": {}, "source": [ "11. What happens if we choose to sample a very steep or very wide Gaussian? What happens if we play with ```L``` or ```step_size```?" ] }, { "cell_type": "markdown", "id": "3145567e", "metadata": {}, "source": [ "## 3. Lifted Metropolis scheme\n", "\n", "12. We now wish to code a lifted Metropolis Monte Carlo:\n", "- we need to extend the state space by a new lifting variable $v$, which stands for the direction. Which distribution can we choose for $v$?\n", "- we need to implement a redirection scheme at event, as a flip of $v$.\n" ] }, { "cell_type": "code", "execution_count": 35, "id": "a85304ea", "metadata": {}, "outputs": [], "source": [ "def lifted_proposal(x,v):\n", " return x + v\n" ] }, { "cell_type": "code", "execution_count": 36, "id": "2aef4d34", "metadata": {}, "outputs": [], "source": [ "def redirection_flip(x,v):\n", " return -v\n", "\n", "def redirection_reflect(x,v):\n", " grad_x = grad_neglogtarget(x)\n", " n_x = grad_x / np.dot(grad_x,grad_x) ** 0.5\n", " v_reflect = v - 2 * np.dot(v, n_x) * n_x\n", " new_proposed = lifted_proposal(current, -v_reflect)\n", " old_proposed = lifted_proposal(current,v)\n", " p_acc_new = min(1,np.exp(-(neglogtarget(new_proposed)-neglogtarget(x))))\n", " p_acc_old = min(1,np.exp(-(neglogtarget(old_proposed)-neglogtarget(x))))\n", " p_acc_v = (1.0-p_acc_new)/(1.0-p_acc_old)\n", " if np.random.random() < p_acc_v:\n", " return v_reflect\n", " else:\n", " return -v\n", " " ] }, { "cell_type": "code", "execution_count": 37, "id": "ce2e43f4", "metadata": {}, "outputs": [], "source": [ "samples3 = [np.array([3.0,2.0])]\n", "dim = samples3[0].shape\n", "iterations = 100000\n", "accept = 0.0\n", "for _ in range(iterations//10): \n", " v = np.random.normal(scale=1.0, size=dim)\n", " for __ in range(10):\n", " current = samples3[-1]\n", " proposed = lifted_proposal(current,v)\n", " if np.random.random() < np.exp(-(neglogtarget(proposed)-neglogtarget(current))):\n", " samples3.append(proposed)\n", " accept += 1.\n", " else:\n", " v = redirection_flip(current,v)\n", " samples3.append(current)\n", "samples3 = np.array(samples3)\n", "print(accept/float(iterations))\n", "\n", "samples4 = [np.array([3.0,2.0])]\n", "dim = samples4[0].shape\n", "iterations = 100000\n", "accept = 0.0\n", "for _ in range(iterations//10): \n", " v = np.random.normal(scale=1.0, size=dim)\n", " for __ in range(10):\n", " current = samples4[-1]\n", " proposed = lifted_proposal(current,v)\n", " if np.random.random() < np.exp(-(neglogtarget(proposed)-neglogtarget(current))):\n", " samples4.append(proposed)\n", " accept += 1.\n", " else:\n", " v = redirection_reflect(current,v)\n", " samples4.append(current)\n", "samples4 = np.array(samples4)\n", "print(accept/float(iterations))" ] }, { "cell_type": "markdown", "id": "e44ecdc8", "metadata": {}, "source": [ "In comparison with Metropolis and HMC results:\n", "\n", "13. Plot the trace plot ||𝑥||\n", "14. Plot a section plot" ] }, { "cell_type": "code", "execution_count": null, "id": "f006bf9e", "metadata": {}, "outputs": [], "source": [ "plt.plot(np.cumsum(np.array([sum(s**2.0)**0.5 for s in samples]))/(1.0+np.arange(len(samples))),ls='--',label='Metropolis')\n", "plt.plot(np.cumsum(np.array([sum(s**2.0)**0.5 for s in samples2]))/(1.0+np.arange(len(samples2))),label='HMC')\n", "plt.plot(np.cumsum(np.array([sum(s**2.0)**0.5 for s in samples3]))/(1.0+np.arange(len(samples3))),label='Flip Lifted Met')\n", "plt.plot(np.cumsum(np.array([sum(s**2.0)**0.5 for s in samples4]))/(1.0+np.arange(len(samples4))),label='Reflect Lifted Met')\n", "plt.xlabel('t')\n", "plt.ylabel('')\n", "plt.legend() " ] }, { "cell_type": "code", "execution_count": null, "id": "dadc05cd", "metadata": {}, "outputs": [], "source": [ "plt.scatter(samples[:,0][::100], samples[:,1][::100], alpha=0.1,label='Metropolis')\n", "plt.scatter(samples2[:,0][::100], samples2[:,1][::100], alpha=0.1,label='HMC')\n", "plt.scatter(samples3[:,0][::100], samples3[:,1][::100], alpha=0.1,label='Flip Lifted Met')\n", "plt.scatter(samples4[:,0][::100], samples4[:,1][::100], alpha=0.1,label='Reflect Lifted Met')\n", "plt.xlabel('x_0')\n", "plt.ylabel('x_1')\n", "plt.legend() " ] }, { "cell_type": "markdown", "id": "fd075c03", "metadata": {}, "source": [ "15. Draw a trajectory for Lifted Metropolis, HMC and Metropolis (10 points, 100 points)" ] }, { "cell_type": "code", "execution_count": null, "id": "dbc4b906", "metadata": {}, "outputs": [], "source": [ "plt.plot(samples[-100:][:,0],samples[-100:][:,1],'o-',label='Metropolis')\n", "plt.plot(samples2[-100:][:,0],samples2[-100:][:,1],'o-',label='HMC')\n", "plt.plot(samples3[-100:][:,0],samples3[-100:][:,1],'o-',label='Flip Lifted Met')\n", "plt.plot(samples4[-100:][:,0],samples4[-100:][:,1],'o-',label='Reflect Lifted Met')\n", "plt.xlabel('x_0')\n", "plt.ylabel('x_1')\n", "plt.legend() " ] }, { "cell_type": "markdown", "id": "4008c751", "metadata": {}, "source": [ "10. Check that the distributions of $x_0$ and/or $x_1$ match." ] }, { "cell_type": "code", "execution_count": null, "id": "5f5897d7", "metadata": {}, "outputs": [], "source": [ "n1,b1,v1=plt.hist(samples[:,0],bins=1000,histtype='step',label='Metropolis',cumulative=True)\n", "n2,b2,v2=plt.hist(samples2[:,0],bins=1000,histtype='step',label='HMC',cumulative=True)\n", "n3,b3,v3=plt.hist(samples3[:,0],bins=1000,histtype='step',label='Flip Lifted Met',cumulative=True)\n", "n3,b3,v3=plt.hist(samples4[:,0],bins=1000,histtype='step',label='Reflect Lifted Met',cumulative=True)\n", "plt.xlabel('x_0')\n", "plt.legend() " ] }, { "cell_type": "markdown", "id": "15fa0a70", "metadata": {}, "source": [ "## 4. ECMC/PDMP sampling (continuous-time lifted dynamics)" ] }, { "cell_type": "code", "execution_count": null, "id": "637305fe", "metadata": {}, "outputs": [], "source": [ "def get_next_event(x,v):\n", " DeltaEStar = - np.log(np.random.random()) \n", " a = np.sum(v ** 2.0) / (2.0 * scale ** 2.0)\n", " initial_min_dist = np.dot(v, grad_neglogtarget(x))/ (2.0 * a)\n", " E0 = neglogtarget(x)\n", " Emin = E0 - a * initial_min_dist ** 2.0\n", " if initial_min_dist > 0.0:\n", " DeltaEStar += E0\n", " else:\n", " DeltaEStar += Emin\n", " time_to_next_event = - initial_min_dist + ((DeltaEStar - Emin) / a) ** 0.5 \n", " return time_to_next_event\n", " \n", " \n", "def redirection_ecmc(x,v): \n", " grad_x = grad_neglogtarget(x)\n", " n_x = grad_x / np.dot(grad_x,grad_x) ** 0.5\n", " v_perp = v - np.dot(v, n_x) * n_x \n", " v_perp /= np.dot(v_perp,v_perp) ** 0.5\n", " a = np.random.random()\n", " v_new = -(1.0-a**2.0)**0.5 * n_x + a * v_perp\n", " return v_new" ] }, { "cell_type": "code", "execution_count": null, "id": "ab90a64f", "metadata": {}, "outputs": [], "source": [ "samples5 = [np.array([3.0,2.0])]\n", "dim = samples5[0].shape\n", "iterations = 100000\n", "av_events = 0.0\n", "t_sample_0 = 0.5\n", "v = np.random.normal(scale=1.0, size=dim)\n", "v /= np.dot(v,v) ** 0.5\n", "current = samples5[-1].copy()\n", "for _ in range(iterations): \n", " t_sample = t_sample_0 \n", " while len(samples5) < iterations:\n", " av_events += 1\n", " t_ev = get_next_event(current,v)\n", " if t_ev < t_sample:\n", " current += v * t_ev\n", " t_sample -= t_ev\n", " v = redirection_ecmc(current,v) \n", " else:\n", " while t_sample < t_ev and len(samples5)')\n", "plt.legend() " ] }, { "cell_type": "code", "execution_count": null, "id": "9fbe544c", "metadata": {}, "outputs": [], "source": [ "plt.plot(samples[-100:][:,0],samples[-100:][:,1],'o-',label='Metropolis')\n", "plt.plot(samples2[-100:][:,0],samples2[-100:][:,1],'o-',label='HMC')\n", "plt.plot(samples3[-100:][:,0],samples3[-100:][:,1],'o-',label='Flip Lifted Met')\n", "plt.plot(samples4[-100:][:,0],samples4[-100:][:,1],'o-',label='Reflect Lifted Met')\n", "plt.plot(samples5[-100:][:,0],samples5[-100:][:,1],'o-',label='ECMC')\n", "plt.xlabel('x_0')\n", "plt.ylabel('x_1')\n", "plt.legend() " ] }, { "cell_type": "code", "execution_count": null, "id": "33734604", "metadata": {}, "outputs": [], "source": [ "n1,b1,v1=plt.hist(samples[:,0],bins=1000,histtype='step',label='Metropolis',cumulative=True,density=True)\n", "n2,b2,v2=plt.hist(samples2[:,0],bins=1000,histtype='step',label='HMC',cumulative=True,density=True)\n", "n3,b3,v3=plt.hist(samples3[:,0],bins=1000,histtype='step',label='Flip Lifted Met',cumulative=True,density=True)\n", "n4,b4,v4=plt.hist(samples4[:,0],bins=1000,histtype='step',label='Reflect Lifted Met',cumulative=True,density=True)\n", "n5,b5,v5=plt.hist(samples5[:,0],bins=1000,histtype='step',label='ECMC',cumulative=True,density=True)\n", "plt.xlabel('x_0')\n", "plt.legend() " ] }, { "cell_type": "markdown", "id": "1a78e3ea", "metadata": {}, "source": [ "## 5. Assessing convergence by autocorrelation analysis" ] }, { "cell_type": "code", "execution_count": null, "id": "cfea65b3", "metadata": {}, "outputs": [], "source": [ "import scipy.signal" ] }, { "cell_type": "code", "execution_count": null, "id": "be93d40f", "metadata": {}, "outputs": [], "source": [ "def plot_aco(data_list,obs,ax,label_list):\n", " for iter in range(len(data_list)):\n", " samples_foraco = data_list[iter] \n", " if obs=='norm':\n", " dat = np.array([sum(s**2.0)**0.5 for s in samples_foraco])\n", " elif obs=='x0':\n", " dat = np.array([s[0] for s in samples_foraco])\n", " elif obs=='x1':\n", " dat = np.array([s[1] for s in samples_foraco])\n", " elif obs=='x':\n", " dat = np.array([s for s in samples_foraco]) \n", " label = label_list[iter]\n", " mean = np.mean(dat)\n", " dat -= mean\n", " lim_max = (len(dat)+1)//2\n", " aco = scipy.signal.correlate(dat, dat, 'same', method='fft')[len(dat)-lim_max:len(dat)]\n", " if obs=='x':\n", " aco = sum(aco.T)\n", " aco /= np.arange(len(dat),len(dat)-lim_max,-1)\n", " aco /= aco[0]\n", " aco_x = 1.0 * np.arange(0,lim_max)\n", " if label=='HMC':\n", " aco_x *= L\n", " if label=='ECMC':\n", " aco_x *= av_events\n", " ax.plot(aco_x, aco,label=label_list[iter])\n", " ax.axhline(y=0.) \n", " return\n" ] }, { "cell_type": "code", "execution_count": null, "id": "17531820", "metadata": {}, "outputs": [], "source": [ "obss = ['norm','x0','x1','x']\n", "fig = plt.figure(figsize=(12,12))\n", "ax1 = fig.add_subplot(221)\n", "ax2 = fig.add_subplot(222)\n", "ax3 = fig.add_subplot(223)\n", "ax4 = fig.add_subplot(224)\n", "axes = [ax1,ax2,ax3,ax4]\n", "for _ in range(len(axes)):\n", " ax, obs = axes[_],obss[_]\n", " ax.title.set_text(obs + ' Autocorr')\n", " plot_aco([samples,samples2,samples3,samples4,samples5],obs,ax,['Metropolis','HMC','Flip Lifted Metro','Reflect Lifted Metro','ECMC'])\n", " ax.legend()\n", " ax.set_xlim(0,100)\n", "plt.show()\n", "\n" ] } ], "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.8.8" } }, "nbformat": 4, "nbformat_minor": 5 }