{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yj2YE5zAqZHf"
   },
   "source": [
    "# Markov Chain Monte Carlo examples\n",
    "Our goal is to simulate a normal distribution, where\n",
    "$$P(x|\\mu,\\sigma^2)=\\frac{1}{\\sqrt{2\\pi\\sigma^2}}\\exp\\left[-\\frac{(x-\\mu)^2}{2\\sigma^2}\\right]$$\n",
    "We'll do this with a variety of Markov chain Monte Carlo techniques to see how it all works."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "LpzhpPz6qV6j"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy.special as sp\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "rng = np.random.default_rng()\n",
    "from numba import jit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pEXeA--6qvwf"
   },
   "outputs": [],
   "source": [
    "def Pnorm(x, mu=0.0, sigma2=1.0):\n",
    "  return np.exp(-0.5*(x-mu)**2/sigma2)/np.sqrt(2*np.pi*sigma2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rAd7ixXDrZ02"
   },
   "outputs": [],
   "source": [
    "@jit(nopython=True)\n",
    "def MC_norm(x0, xstates, delta, log_u, mu, sigma2):\n",
    "  \"\"\"Does MCMC, starting from x0 for N(mu, sigma2), storing in xstates\n",
    "  where delta is the random (symmetric) displacements, and log_u is\n",
    "  -np.log(u) for u from a uniform distribution\"\"\"\n",
    "  Nstates = len(xstates)\n",
    "  x = x0\n",
    "  accept = 0\n",
    "  for i in range(Nstates):\n",
    "    d = delta[i]\n",
    "    xp = x + d\n",
    "    dE = 0.5*(d*d + 2*(x-mu)*d)\n",
    "    if sigma2*log_u[i] > dE:\n",
    "      x = xp\n",
    "      accept += 1\n",
    "    xstates[i] = x\n",
    "  return accept/Nstates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nmw77O4esnna"
   },
   "outputs": [],
   "source": [
    "mu, sigma2 = 1.0, 1.0\n",
    "\n",
    "x0 = 0.0\n",
    "Nstates = 2**16\n",
    "Delta = 1.0\n",
    "\n",
    "states = np.zeros(Nstates)\n",
    "delta = rng.uniform(low=-Delta, high=Delta, size=Nstates)\n",
    "log_u = -np.log(rng.uniform(size=Nstates))\n",
    "accept = MC_norm(x0, states, delta, log_u, mu, sigma2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "umJ9EYTDs_9F"
   },
   "outputs": [],
   "source": [
    "accept"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kVxuNCTZtBQl"
   },
   "outputs": [],
   "source": [
    "plt.hist(states, bins=Nstates//1024, density=True);\n",
    "x = np.sort(states)\n",
    "plt.plot(x, Pnorm(x, mu, sigma2), '--');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "uhDxykjwtEqD"
   },
   "outputs": [],
   "source": [
    "x = np.sort(states)\n",
    "y = np.linspace(0, 1, len(x))\n",
    "plt.plot(x, y);\n",
    "plt.plot(x, 0.5*(sp.erf((x-mu)/np.sqrt(2*sigma2))+1), '--');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Tux-XNd9iOeq"
   },
   "outputs": [],
   "source": [
    "np.mean(states), np.var(states)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7wplw9q7uRCP"
   },
   "outputs": [],
   "source": [
    "def acf_calc(v):\n",
    "  \"\"\"Returns the auto-correlation function for the sequence v\"\"\"\n",
    "  v_ave = np.mean(v)\n",
    "  v_len = len(v)\n",
    "  return np.correlate(v-v_ave, v-v_ave, \"full\")[(v_len-1):] / \\\n",
    "   (np.var(v)*np.array([v_len - x for x in range(v_len)]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4quUZYTctfVU"
   },
   "outputs": [],
   "source": [
    "plt.plot(acf_calc(states)[:128]);\n",
    "plt.axhline(0.0);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "P3foZpY5vSOD"
   },
   "outputs": [],
   "source": [
    "Delta_list = [0.01, 0.02, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0, 20.0, 50.0, 100.0]\n",
    "autocorr_list = []\n",
    "accept_list = []\n",
    "for Delta in Delta_list:\n",
    "  states = np.zeros(Nstates)\n",
    "  delta = rng.uniform(low=-Delta, high=Delta, size=Nstates)\n",
    "  log_u = -np.log(rng.uniform(size=Nstates))\n",
    "  accept = MC_norm(x0, states, delta, log_u, mu, sigma2)\n",
    "  autocorr = acf_calc(states)\n",
    "  t_min, t_zero = 0, 0\n",
    "  for t in range(len(autocorr)-1):\n",
    "    if autocorr[t] < 0.0 and t_zero == 0:\n",
    "      t_zero = t\n",
    "    if autocorr[t] < autocorr[t+1] and t_min == 0:\n",
    "      t_min = t\n",
    "  autocorr_list.append(min(t_zero, t_min))\n",
    "  accept_list.append(accept)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-Kkp6kuFUKlZ"
   },
   "outputs": [],
   "source": [
    "plt.plot(Delta_list, autocorr_list, 'o-');\n",
    "plt.xscale('log');\n",
    "plt.yscale('log');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GUAAyEK5UYBo"
   },
   "outputs": [],
   "source": [
    "plt.plot(Delta_list, accept_list, 'o-');\n",
    "plt.xscale('log');\n",
    "plt.ylim(0, 1);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yOn726yqUqrJ"
   },
   "outputs": [],
   "source": [
    "# Use a normal distribution for delta instead\n",
    "\n",
    "states = np.zeros(Nstates)\n",
    "delta = rng.normal(size=Nstates)\n",
    "log_u = -np.log(rng.uniform(size=Nstates))\n",
    "accept = MC_norm(x0, states, delta, log_u, mu, sigma2)\n",
    "autocorr = acf_calc(states)\n",
    "t_min, t_zero = 0, 0\n",
    "for t in range(len(autocorr)-1):\n",
    "  if autocorr[t] < 0.0 and t_zero == 0:\n",
    "    t_zero = t\n",
    "  if autocorr[t] < autocorr[t+1] and t_min == 0:\n",
    "    t_min = t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ct-S2E78QKEE"
   },
   "outputs": [],
   "source": [
    "accept, t_zero, t_min"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Js52e5VvQLwl"
   },
   "outputs": [],
   "source": [
    "min(autocorr_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "poKtoWwNQasY"
   },
   "outputs": [],
   "source": [
    "# Use a Cauchy distribution for delta instead\n",
    "\n",
    "states = np.zeros(Nstates)\n",
    "delta = rng.standard_cauchy(size=Nstates)\n",
    "log_u = -np.log(rng.uniform(size=Nstates))\n",
    "accept = MC_norm(x0, states, delta, log_u, mu, sigma2)\n",
    "autocorr = acf_calc(states)\n",
    "t_min, t_zero = 0, 0\n",
    "for t in range(len(autocorr)-1):\n",
    "  if autocorr[t] < 0.0 and t_zero == 0:\n",
    "    t_zero = t\n",
    "  if autocorr[t] < autocorr[t+1] and t_min == 0:\n",
    "    t_min = t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BbKCJzv-QeTf"
   },
   "outputs": [],
   "source": [
    "accept, t_zero, t_min"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Chq0a_ywQg_1"
   },
   "outputs": [],
   "source": [
    "x = np.sort(states)\n",
    "y = np.linspace(0, 1, len(x))\n",
    "plt.plot(x, y);\n",
    "plt.plot(x, 0.5*(sp.erf((x-mu)/np.sqrt(2*sigma2))+1), '--');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "AejGJxgxQli3"
   },
   "outputs": [],
   "source": [
    "plt.plot(acf_calc(states)[:128]);\n",
    "plt.axhline(0.0);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5dlg7wG_Q8xL"
   },
   "source": [
    "What happens if we use an asymmetric jump distribution for $T(s\\to s')$ but don't account for it the acceptance?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Qh_KtflAQrgE"
   },
   "outputs": [],
   "source": [
    "states = np.zeros(Nstates)\n",
    "Delta = 1.0\n",
    "delta = rng.triangular(-Delta, Delta, Delta, size=Nstates)\n",
    "log_u = -np.log(rng.uniform(size=Nstates))\n",
    "accept = MC_norm(x0, states, delta, log_u, mu, sigma2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Y0EE_8dcRmOg"
   },
   "outputs": [],
   "source": [
    "x = np.sort(states)\n",
    "y = np.linspace(0, 1, len(x))\n",
    "plt.plot(x, y);\n",
    "plt.plot(x, 0.5*(sp.erf((x-mu)/np.sqrt(2*sigma2))+1), '--');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4WRHfGgiR3i9"
   },
   "outputs": [],
   "source": [
    "@jit(nopython=True)\n",
    "def MC_norm_asymm(x0, xstates, delta, deltaE, log_u, mu, sigma2):\n",
    "  \"\"\"Does MCMC, starting from x0 for N(mu, sigma2), storing in xstates\n",
    "  where delta is the random (asymmetric) displacements, with \"energy\n",
    "  correction\" deltaE, and log_u is -np.log(u) for u from a uniform\n",
    "  distribution\"\"\"\n",
    "  Nstates = len(xstates)\n",
    "  x = x0\n",
    "  accept = 0\n",
    "  for i in range(Nstates):\n",
    "    d = delta[i]\n",
    "    xp = x + d\n",
    "    dE = 0.5*(d*d + 2*(x-mu)*d) + deltaE[i]\n",
    "    if sigma2*log_u[i] > dE:\n",
    "      x = xp\n",
    "      accept += 1\n",
    "    xstates[i] = x\n",
    "  return accept/Nstates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "K4P_ZUncecNs"
   },
   "outputs": [],
   "source": [
    "states = np.zeros(Nstates)\n",
    "Delta = 1.0\n",
    "delta = rng.triangular(-Delta, Delta, Delta, size=Nstates)\n",
    "deltaE = -sigma2*np.log((Delta-delta)/(Delta+delta))\n",
    "log_u = -np.log(rng.uniform(size=Nstates))\n",
    "accept = MC_norm_asymm(x0, states, delta, deltaE, log_u, mu, sigma2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xa3zRsChemvm"
   },
   "outputs": [],
   "source": [
    "x = np.sort(states)\n",
    "y = np.linspace(0, 1, len(x))\n",
    "plt.plot(x, y);\n",
    "plt.plot(x, 0.5*(sp.erf((x-mu)/np.sqrt(2*sigma2))+1), '--');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "K8r8TZ-3kBRw"
   },
   "source": [
    "What happens if we loop until we get a move we accept?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rWxQf-1shU-r"
   },
   "outputs": [],
   "source": [
    "states = np.zeros(Nstates)\n",
    "x = x0\n",
    "Delta = 1.0\n",
    "for i in range(Nstates):\n",
    "  accept = False\n",
    "  while not accept:\n",
    "    d = rng.uniform(-Delta, Delta)\n",
    "    xp = x + d\n",
    "    dE = 0.5*(d*d + 2*(x-mu)*d)\n",
    "    if sigma2*log_u[i] > dE:\n",
    "      x = xp\n",
    "      accept = True\n",
    "  states[i] = x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FZq_8BDikApZ"
   },
   "outputs": [],
   "source": [
    "x = np.sort(states)\n",
    "y = np.linspace(0, 1, len(x))\n",
    "plt.plot(x, y);\n",
    "plt.plot(x, 0.5*(sp.erf((x-mu)/np.sqrt(2*sigma2))+1), '--');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "meMhVRp5kmkj"
   },
   "source": [
    "Let's look at doing Force Bias MC, where we sample our displacement from an exponential in the range $[-\\Delta, \\Delta]$ with a PDF given by\n",
    "$$T(\\delta) = \\frac{\\alpha}{2\\sinh(\\alpha\\Delta)}\\exp(\\alpha \\delta)$$\n",
    "where $\\alpha = -(x-\\mu)/\\sigma^2$ is the derivative of log probability at $x$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qNKPwrf2kd3p"
   },
   "outputs": [],
   "source": [
    "@jit(nopython=True)\n",
    "def MC_norm_forcebias(x0, xstates, Delta, u1, u2, mu, sigma2):\n",
    "  \"\"\"Does MCMC, starting from x0 for N(mu, sigma2), storing in xstates,\n",
    "  where Delta is the range over which we sample displacements, and\n",
    "  where u1 and u2 are uniformly distributed random numbers to be used\n",
    "  to generate displacements and generate acceptance\"\"\"\n",
    "  Nstates = len(xstates)\n",
    "  x = x0\n",
    "  accept = 0\n",
    "  for i in range(Nstates):\n",
    "    alpha = -(x-mu)/sigma2\n",
    "    expap = np.exp(alpha*Delta)\n",
    "    expam = 1.0/expap\n",
    "    d = np.log(u1[i]*(expap-expam) + expam)/alpha\n",
    "    xp = x + d\n",
    "    dE = 0.5*(d*d + 2*(x-mu)*d)/sigma2\n",
    "    prob_accept = np.exp(dE)*(xp-mu)/(x-mu) * \\\n",
    "      np.sinh((x-mu)*Delta/sigma2)/np.sinh((xp-mu)*Delta/sigma2)\n",
    "    # print(f'{x} {d} {alpha} {dE} {prob_accept}')\n",
    "    if u2[i] < prob_accept:\n",
    "      x = xp\n",
    "      accept += 1\n",
    "    xstates[i] = x\n",
    "  return accept/Nstates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yswTb3ZknKFK"
   },
   "outputs": [],
   "source": [
    "states = np.zeros(Nstates)\n",
    "Delta = 2.5\n",
    "accept = MC_norm_forcebias(x0, states, Delta,\n",
    "                           rng.uniform(size=Nstates), rng.uniform(size=Nstates),\n",
    "                           mu, sigma2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bEU1Ljx-nYHt"
   },
   "outputs": [],
   "source": [
    "accept"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xSSOiuxendcR"
   },
   "outputs": [],
   "source": [
    "plt.hist(states, density=True, bins=64);\n",
    "plt.plot(x, Pnorm(x, mu, sigma2), '--');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "id0wvAf-w-Hl"
   },
   "outputs": [],
   "source": [
    "x = np.sort(states)\n",
    "y = np.linspace(0, 1, len(x))\n",
    "plt.plot(x, y);\n",
    "plt.plot(x, 0.5*(sp.erf((x-mu)/np.sqrt(2*sigma2))+1), '--');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6EqwR3Rf3MDE"
   },
   "outputs": [],
   "source": [
    "autocorr = acf_calc(states)\n",
    "t_min, t_zero = 0, 0\n",
    "for t in range(len(autocorr)-1):\n",
    "  if autocorr[t] < 0.0 and t_zero == 0:\n",
    "    t_zero = t\n",
    "  if autocorr[t] < autocorr[t+1] and t_min == 0:\n",
    "    t_min = t\n",
    "print(t_min, t_zero)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ogm1GSDM3WNl"
   },
   "outputs": [],
   "source": [
    "plt.plot(autocorr[:32]);\n",
    "plt.axhline(0.0);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VSWhqnNz3zOo"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "authorship_tag": "ABX9TyOZEj7mVqf6NFD9vjtlj79T",
   "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.13.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}