{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy import stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "lengths = np.array([195, 185])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# lets create a grid of our two parameters\n",
    "mu = np.linspace(150, 250)\n",
    "sigma = np.linspace(0, 15)[::-1]\n",
    "\n",
    "mm, ss = np.meshgrid(mu, sigma)  # just broadcasted parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/miniconda3/lib/python3.7/site-packages/scipy/stats/_distn_infrastructure.py:1662: RuntimeWarning: divide by zero encountered in true_divide\n",
      "  x = np.asarray((x - loc)/scale, dtype=dtyp)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Text(0, 0.5, '$\\\\sigma$')"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 7724 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "likelihood = stats.norm(mm, ss).pdf(lengths[0]) * stats.norm(mm, ss).pdf(lengths[1])\n",
    "aspect = mm.max() / ss.max() / 3\n",
    "extent = [mm.min(), mm.max(), ss.min(), ss.max()]\n",
    "# extent = left right bottom top\n",
    "plt.imshow(likelihood, cmap='Reds', aspect=aspect, extent=extent) \n",
    "plt.xlabel(r'$\\mu$')\n",
    "plt.ylabel(r'$\\sigma$')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0, 0.5, '$\\\\sigma$')"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 9560 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "prior = stats.norm(200, 15).pdf(mm) * stats.cauchy(0, 10).pdf(ss)\n",
    "\n",
    "plt.imshow(prior, cmap='Greens', aspect=aspect, extent=extent)\n",
    "plt.xlabel(r'$\\mu$')\n",
    "plt.ylabel(r'$\\sigma$')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0, 0.5, '$\\\\sigma$')"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 7188 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "unnormalized_posterior = prior * likelihood\n",
    "posterior = unnormalized_posterior / np.nan_to_num(unnormalized_posterior).sum()\n",
    "plt.imshow(posterior, cmap='Blues', aspect=aspect, extent=extent)\n",
    "plt.xlabel(r'$\\mu$')\n",
    "plt.ylabel(r'$\\sigma$')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0, 0.5, '$\\\\sigma$')"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 7188 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.imshow(unnormalized_posterior, cmap='Blues', aspect=aspect, extent=extent)\n",
    "plt.xlabel(r'$\\mu$')\n",
    "plt.ylabel(r'$\\sigma$')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pymc3 as pm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Auto-assigning NUTS sampler...\n",
      "Initializing NUTS using jitter+adapt_diag...\n",
      "Sequential sampling (1 chains in 1 job)\n",
      "NUTS: [sigma, mu]\n",
      "100%|██████████| 10500/10500 [00:07<00:00, 1402.54it/s]\n",
      "The acceptance probability does not match the target. It is 0.8838815681167924, but should be close to 0.8. Try to increase the number of tuning steps.\n",
      "Only one chain was sampled, this makes it impossible to run some convergence checks\n"
     ]
    }
   ],
   "source": [
    "with pm.Model():\n",
    "    # priors\n",
    "    mu = pm.Normal('mu', mu=200, sd=15)\n",
    "    sigma = pm.HalfCauchy('sigma', 10)\n",
    "    \n",
    "    # likelihood\n",
    "    observed = pm.Normal('observed', mu=mu, sd=sigma, observed=lengths)\n",
    "    # sample\n",
    "    trace = pm.sample(draws=10000, chains=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/miniconda3/lib/python3.7/site-packages/matplotlib/axes/_base.py:3604: MatplotlibDeprecationWarning: \n",
      "The `ymin` argument was deprecated in Matplotlib 3.0 and will be removed in 3.2. Use `bottom` instead.\n",
      "  alternative='`bottom`', obj_type='argument')\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 40708 characters]\n",
      "text/plain": [
       "<Figure size 864x288 with 4 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "pm.traceplot(trace);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0, 15.0)"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 93112 characters]\n",
      "text/plain": [
       "<Figure size 1152x432 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, axes = plt.subplots(2, sharex=True, sharey=True, figsize = (16, 6))\n",
    "axes[0].imshow(posterior, cmap='Blues', extent=extent, aspect=1)\n",
    "axes[0].set_ylabel('$\\sigma$')\n",
    "axes[1].scatter(trace['mu'], trace['sigma'], alpha=0.01)\n",
    "axes[1].set_ylabel('$\\sigma$')\n",
    "axes[0].set_title('True posterior')\n",
    "axes[1].set_title('Sampled $\\\\theta$')\n",
    "plt.xlabel('$\\mu$')\n",
    "plt.xlim(150, mm.max())\n",
    "plt.ylim(0, ss.max())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyro\n",
    "import pyro.distributions as dist\n",
    "import pyro.optim\n",
    "from pyro.infer import SVI, Trace_ELBO\n",
    "import torch\n",
    "import torch.distributions.constraints as constraints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def model():\n",
    "    # priors\n",
    "    mu = pyro.sample('mu', dist.Normal(loc=torch.tensor(200.), \n",
    "                                       scale=torch.tensor(15.)))\n",
    "    sigma = pyro.sample('sigma', dist.HalfCauchy(scale=torch.tensor(10.)))\n",
    "    \n",
    "    # likelihood\n",
    "    with pyro.plate('plate', size=2):\n",
    "        pyro.sample(f'obs', dist.Normal(loc=mu, scale=sigma), \n",
    "                    obs=torch.tensor([195., 185.]))\n",
    "    \n",
    "def guide():\n",
    "    # variational parameters\n",
    "    var_mu = pyro.param('var_mu', torch.tensor(180.))\n",
    "    var_mu_sig = pyro.param('var_mu_sig', torch.tensor(5.), \n",
    "                             constraint=constraints.positive)\n",
    "    var_sig = pyro.param('var_sig', torch.tensor(5.))\n",
    "    \n",
    "    # factorized distribution\n",
    "    pyro.sample('mu', dist.Normal(loc=var_mu, scale=var_mu_sig))\n",
    "    pyro.sample('sigma', dist.Chi2(var_sig))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[iteration    1] loss: 20.5714\n",
      "[iteration  101] loss: 12.0774\n",
      "[iteration  201] loss: 11.7402\n",
      "[iteration  301] loss: 9.8967\n",
      "[iteration  401] loss: 28.1990\n",
      "[iteration  501] loss: 13.2090\n",
      "[iteration  601] loss: 11.4618\n",
      "[iteration  701] loss: 14.9521\n",
      "[iteration  801] loss: 11.2410\n",
      "[iteration  901] loss: 9.7651\n",
      "[iteration 1001] loss: 9.7140\n",
      "[iteration 1101] loss: 9.5427\n",
      "[iteration 1201] loss: 9.0695\n",
      "[iteration 1301] loss: 9.0411\n",
      "[iteration 1401] loss: 10.2914\n",
      "[iteration 1501] loss: 8.6590\n",
      "[iteration 1601] loss: 9.7012\n",
      "[iteration 1701] loss: 12.4810\n",
      "[iteration 1801] loss: 8.0311\n",
      "[iteration 1901] loss: 9.6596\n",
      "[iteration 2001] loss: 10.1730\n",
      "[iteration 2101] loss: 8.9640\n",
      "[iteration 2201] loss: 9.5495\n",
      "[iteration 2301] loss: 9.6852\n",
      "[iteration 2401] loss: 10.8672\n",
      "[iteration 2501] loss: 8.9980\n",
      "[iteration 2601] loss: 8.3449\n",
      "[iteration 2701] loss: 9.4454\n",
      "[iteration 2801] loss: 9.4535\n",
      "[iteration 2901] loss: 9.0612\n",
      "[iteration 3001] loss: 8.6264\n",
      "[iteration 3101] loss: 8.5423\n",
      "[iteration 3201] loss: 7.1836\n",
      "[iteration 3301] loss: 8.9224\n",
      "[iteration 3401] loss: 8.8739\n",
      "[iteration 3501] loss: 9.0185\n",
      "[iteration 3601] loss: 38.2326\n",
      "[iteration 3701] loss: 8.8873\n",
      "[iteration 3801] loss: 9.0606\n",
      "[iteration 3901] loss: 8.0937\n",
      "[iteration 4001] loss: 9.1192\n",
      "[iteration 4101] loss: 7.1300\n",
      "[iteration 4201] loss: 9.4013\n",
      "[iteration 4301] loss: 8.7548\n",
      "[iteration 4401] loss: 9.2406\n",
      "[iteration 4501] loss: 9.3022\n",
      "[iteration 4601] loss: 8.9886\n",
      "[iteration 4701] loss: 9.1033\n",
      "[iteration 4801] loss: 9.0818\n",
      "[iteration 4901] loss: 8.5367\n"
     ]
    }
   ],
   "source": [
    "pyro.clear_param_store()\n",
    "pyro.enable_validation(True)\n",
    "\n",
    "svi = SVI(model, guide, \n",
    "          optim=pyro.optim.ClippedAdam({\"lr\":0.01}), \n",
    "          loss=Trace_ELBO())\n",
    "\n",
    "# do gradient steps\n",
    "c = 0\n",
    "for step in range(5000):\n",
    "    c += 1\n",
    "    loss = svi.step()\n",
    "    if step % 100 == 0:\n",
    "        print(\"[iteration {:>4}] loss: {:.4f}\".format(c, loss))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "190.1922149658203 5.435718059539795\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "10.607619285583496"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(pyro.param('var_mu').item(), pyro.param('var_mu_sig').item())\n",
    "pyro.param('var_sig').item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma = dist.Chi2(pyro.param('var_sig')).sample((10000,)).numpy()\n",
    "mu = dist.Normal(pyro.param('var_mu'), pyro.param('var_mu_sig')).sample((10000,)).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0, 15.0)"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 113292 characters]\n",
      "text/plain": [
       "<Figure size 1152x432 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, axes = plt.subplots(2, sharex=True, sharey=True, figsize = (16, 6))\n",
    "axes[0].imshow(posterior, cmap='Blues', extent=extent, aspect=1)\n",
    "axes[0].set_ylabel('$\\sigma$')\n",
    "axes[1].scatter(mu, sigma, alpha=0.05)\n",
    "axes[1].set_ylabel('$\\sigma$')\n",
    "axes[0].set_title('True posterior $P(\\\\theta|D)$')\n",
    "axes[1].set_title('Approximated posterior $Q(\\\\theta)$')\n",
    "plt.xlabel('$\\mu$')\n",
    "plt.xlim(150, mm.max())\n",
    "plt.ylim(0, ss.max())"
   ]
  }
 ],
 "metadata": {
  "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": 2
}
