{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Draw samples from prior\n",
    "\n",
    "Orphaned code: needs updating to work"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_prior(i, log_scale=False, **kwargs):\n",
    "    with numpyro.handlers.seed(rng_seed=i):\n",
    "        beta, x, y, z, obs_rate, hosp_rate = SEIR_stochastic(**kwargs)\n",
    "        \n",
    "    T = len(x)\n",
    "    ts = np.arange(T)\n",
    "    \n",
    "    fig, ax1 = plt.subplots()\n",
    "    ax2 = ax1.twinx()\n",
    "    \n",
    "    l1 = ax1.plot(ts, x[:,1:])    \n",
    "    l2 = ax2.plot(ts, y, 'x', color='r')\n",
    "    l3 = ax2.plot(ts, z, 'x', color='k')\n",
    "        \n",
    "    plt.legend(l1 + l2 + l3, ['I', 'R', 'C', 'y', 'hosp'])\n",
    "    \n",
    "    ax1.set_ylabel('compartment count')\n",
    "    ax2.set_ylabel('observed cases')\n",
    "\n",
    "    if log_scale:\n",
    "        ax1.set_yscale('log')\n",
    "        ax2.set_yscale('log')\n",
    "    \n",
    "    ylim = np.array(ax1.get_ylim())\n",
    "    ax2.set_ylim(obs_rate * ylim )\n",
    "    \n",
    "    plt.show()\n",
    "\n",
    "def test_prior_predictive():\n",
    "    T = 100\n",
    "    T_future = 100\n",
    "    N = 1e6\n",
    "    predictive = Predictive(SEIR_stochastic, \n",
    "                            posterior_samples = {}, \n",
    "                            num_samples=100)\n",
    "\n",
    "    samples = predictive(PRNGKey(2), T=T, T_future=100, N=N, det_conc=200)\n",
    "\n",
    "    plot_samples(samples, T=200, plot_fields=['I', 'y'])\n",
    "    plt.title('Prior predictive')\n",
    "    plt.yscale('log') \n",
    "    plt.show()\n",
    "    \n",
    "    return samples\n",
    "    \n",
    "run_tests = True\n",
    "    \n",
    "if run_tests:\n",
    "    for i in range(5):\n",
    "        test_prior(i, log_scale=False, T=100, T_future=100, N=1e8, use_hosp=True)\n",
    "        \n",
    "#    samples = test_prior_predictive()"
   ]
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
