{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Default Model Example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Preprocess dataset, build and run default model with a custom infection-to-fatality delay."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING (theano.tensor.blas): Using NumPy C-API based implementation for BLAS functions.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Set Theano Environmental Variables for Parallelisation\n"
     ]
    }
   ],
   "source": [
    "from epimodel.preprocessing.data_preprocessor import preprocess_data\n",
    "from epimodel.pymc3_models.models import DefaultModel\n",
    "from epimodel.pymc3_models.epi_params import EpidemiologicalParameters, bootstrapped_negbinom_values\n",
    "import pymc3 as pm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dropping NPI Travel Screen/Quarantine\n",
      "Dropping NPI Travel Bans\n",
      "Dropping NPI Public Transport Limited\n",
      "Dropping NPI Internal Movement Limited\n",
      "Dropping NPI Public Information Campaigns\n",
      "Dropping NPI Symptomatic Testing\n",
      "Masking invalid values\n"
     ]
    }
   ],
   "source": [
    "data = preprocess_data('../notebooks/double-entry-data/double_entry_final.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create custom infection to fatality delay using default generation interval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "example_symptom_to_fatality_delay = {\n",
    "                'mean_mean': 18,\n",
    "                'mean_sd': 1,\n",
    "                'disp_mean': 10,\n",
    "                'disp_sd': 3,\n",
    "                'source': 'made up',\n",
    "                'dist': 'negbinom',\n",
    "                'notes': 'For example purposes only'\n",
    "            }\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "ep = EpidemiologicalParameters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 250/250 [05:54<00:00,  1.42s/it]\n"
     ]
    }
   ],
   "source": [
    "infection_to_fatality_delay = bootstrapped_negbinom_values([ep.incubation_period,\n",
    "                                                            example_symptom_to_fatality_delay])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "ep.infection_to_fatality_delay = infection_to_fatality_delay[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialise model with epidemiological parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "with DefaultModel(data) as model:\n",
    "    model.build_model(**ep.get_model_build_dict())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run model\n",
    "\n",
    "The notebook features a small number of samples to make documentation compilation proceed quickly. For a serious run, use at least 1000 samples and 500 tuning steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Only 100 samples in chain.\n",
      "Auto-assigning NUTS sampler...\n",
      "Initializing NUTS using jitter+adapt_diag...\n",
      "Multiprocess sampling (4 chains in 4 jobs)\n",
      "NUTS: [DeathsDelayDisp, DeathsDelayMean, InitialSizeDeaths_log, CasesDelayDisp, CasesDelayMean, InitialSizeCases_log, PsiDeaths, PsiCases, GrowthDeathsNoise, GrowthCasesNoise, GI_sd, GI_mean, RegionLogR_noise, HyperRVar, CM_Alpha]\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "        <style>\n",
       "            /* Turns off some styling */\n",
       "            progress {\n",
       "                /* gets rid of default border in Firefox and Opera. */\n",
       "                border: none;\n",
       "                /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
       "                background-size: auto;\n",
       "            }\n",
       "            .progress-bar-interrupted, .progress-bar-interrupted::-webkit-progress-bar {\n",
       "                background: #F44336;\n",
       "            }\n",
       "        </style>\n",
       "      <progress value='800' class='' max='800' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      100.00% [800/800 13:21<00:00 Sampling 4 chains, 0 divergences]\n",
       "    </div>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Sampling 4 chains for 100 tune and 100 draw iterations (400 + 400 draws total) took 803 seconds.\n",
      "The rhat statistic is larger than 1.4 for some parameters. The sampler did not converge.\n",
      "The number of effective samples is smaller than 10% for some parameters.\n"
     ]
    }
   ],
   "source": [
    "with model.model:\n",
    "    model.trace = pm.sample(100, tune=100, cores=4, chains=4, max_treedepth=12, target_accept=0.95)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Trace variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['CM_Alpha',\n",
       " 'HyperRVar_log__',\n",
       " 'RegionLogR_noise',\n",
       " 'GI_mean',\n",
       " 'GI_sd',\n",
       " 'GrowthCasesNoise',\n",
       " 'GrowthDeathsNoise',\n",
       " 'PsiCases_log__',\n",
       " 'PsiDeaths_log__',\n",
       " 'InitialSizeCases_log',\n",
       " 'CasesDelayMean',\n",
       " 'CasesDelayDisp',\n",
       " 'InitialSizeDeaths_log',\n",
       " 'DeathsDelayMean',\n",
       " 'DeathsDelayDisp',\n",
       " 'CMReduction',\n",
       " 'HyperRVar',\n",
       " 'RegionR',\n",
       " 'PsiCases',\n",
       " 'PsiDeaths',\n",
       " 'InfectedCases',\n",
       " 'ExpectedCases',\n",
       " 'InfectedDeaths',\n",
       " 'ExpectedDeaths']"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.trace.varnames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}