{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b23edecb",
   "metadata": {},
   "source": [
    "### Exercise 57"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f199921",
   "metadata": {},
   "source": [
    "#### Import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5f1275c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.stats import norm, uniform, lognorm, poisson, gamma\n",
    "from scipy.special import gamma as gammafunc\n",
    "from scipy.stats import gaussian_kde as kde"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2050596",
   "metadata": {},
   "source": [
    "#### Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ce0d6717",
   "metadata": {},
   "outputs": [],
   "source": [
    "y = np.array([6, 8, 7, 6, 7, 4, 11, 8, 6, 3])\n",
    "n = len(y)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22713432",
   "metadata": {},
   "source": [
    "#### Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0f6bee7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "S = 10_000\n",
    "\n",
    "# Plotting\n",
    "mplot = 200\n",
    "space = np.linspace(2, 10, mplot)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2656ed2",
   "metadata": {},
   "source": [
    "#### Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d6cede51",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_log_likelihood(theta):\n",
    "    return np.sum(poisson(theta).logpmf(y))\n",
    "\n",
    "\n",
    "def get_MH_sample(S, sigma, prior):\n",
    "    sample = np.zeros(S)\n",
    "    \n",
    "    for s in range(S):\n",
    "        if s == 0:\n",
    "            sample[s] = prior.rvs()\n",
    "            continue\n",
    "        \n",
    "        else:\n",
    "            current = sample[s - 1]\n",
    "            proposed = sample[s - 1] + sigma*norm().rvs()\n",
    "            \n",
    "            if proposed < 0:\n",
    "                sample[s] = current\n",
    "                continue\n",
    "            \n",
    "            logprior_current = prior.logpdf(current)\n",
    "            logprior_proposed = prior.logpdf(proposed)\n",
    "            \n",
    "            loglik_current = get_log_likelihood(current)\n",
    "            loglik_proposed = get_log_likelihood(proposed)\n",
    "            \n",
    "            log_acceptance_prob = logprior_proposed + loglik_proposed - logprior_current - loglik_current\n",
    "            \n",
    "            U = np.random.random()\n",
    "            \n",
    "            if np.log(U) < log_acceptance_prob:\n",
    "                sample[s] = proposed\n",
    "            else:\n",
    "                sample[s] = current\n",
    "    return sample"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a64c6021",
   "metadata": {},
   "source": [
    "#### (a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e5c2ef61",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify prior\n",
    "a, b = .1, .1\n",
    "prior = gamma(a=a, scale=1/b)\n",
    "\n",
    "# Run algorithm\n",
    "sample = get_MH_sample(S, 1.0, prior)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "797d2f5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plots\n",
    "plt.plot(space, kde(sample)(space), color='black', label='MCMC')\n",
    "plt.plot(space, gamma(a=a + np.sum(y), scale=1/(b + n)).pdf(space), \n",
    "         color='black', linestyle='dashed', label='exact')\n",
    "plt.legend()\n",
    "plt.xlabel('theta')\n",
    "plt.ylabel('posterior density')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3a265a41",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Trace plots\n",
    "plt.plot(np.arange(S), sample, color='black')\n",
    "plt.xlabel('iteration')\n",
    "plt.ylabel('theta')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09e83499",
   "metadata": {},
   "source": [
    "#### (b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e40e7229",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify prior\n",
    "a, b = .5, 50\n",
    "prior = uniform(loc=a, scale=b - a)\n",
    "\n",
    "# Run algorithm\n",
    "sample = get_MH_sample(S, 1.0, prior)\n",
    "\n",
    "# Exact\n",
    "ybar = np.mean(y)\n",
    "I = gamma(a=1 + n*ybar, scale=1/n).cdf(b) - gamma(a=1 + n*ybar, scale=1/n).cdf(a)\n",
    "posterior = n**(1 + n*ybar)/(gammafunc(1 + n*ybar)) * space**(n*ybar) * np.exp(-n*space) / I"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d0f6f64f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plots\n",
    "plt.plot(space, kde(sample)(space), color='black', label='MCMC')\n",
    "plt.plot(space, posterior, color='black', linestyle='dashed', label='exact')\n",
    "plt.legend()\n",
    "plt.xlabel('theta')\n",
    "plt.ylabel('posterior density')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "215de2ea",
   "metadata": {},
   "source": [
    "#### (c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "02f44225",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify prior\n",
    "prior = lognorm(1)\n",
    "\n",
    "# Run algorithm\n",
    "sample = get_MH_sample(S, 1.0, prior)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "b9b0406d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plots\n",
    "plt.plot(space, kde(sample)(space), color='black', label='MCMC')\n",
    "plt.legend()\n",
    "plt.xlabel('theta')\n",
    "plt.ylabel('posterior density')\n",
    "plt.show()"
   ]
  }
 ],
 "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.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}