seir_template.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SEIR model template\n",
    "\n",
    " * To modify with you own data or parameters edit the cells underneath the \"**user input:**\" header below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### some preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import datetime\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pystan\n",
    "import logging\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "from dateutil.parser import parse\n",
    "\n",
    "data = {'prior_only':[0,0,0,0]}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# user input:\n",
    "\n",
    "The model requires three data types that are read from CSV or Excel files:\n",
    " 1. `cases`: the (cumulative) number of confirmed cases of COVID-19.\n",
    " 2. `deaths`: the (cumulative) number of deaths due to COVID-19.\n",
    " 3. `hospitalizations`: the (current) number of hospitalized (ICU and non-ICU) patients and the (current) number of ICU patients due to COVID-19.\n",
    "\n",
    "For each of the data types above, the variable `datafileinfo` contains the information required to find the data. For each datatype, it must contain the following entries:\n",
    " * `filename`: The data file name. Comma-separated values (CSV) and Excel files are accepted. See [example file](data/santa_cruz_cases.xlsx). If a single file contains more than one data type (like confirmed cases and deaths), the same file name can be used for those data types.\n",
    " * `columnname` (`columnnames` for `hospitalizations`): The column name for that data type in the file. Hospitalizations requires *two* column names for (1) the of number of hospitalized (ICU and non-ICU) patients and (2) the number of ICU patients.\n",
    " * `datecolumnname`: The column name of the associated date column in the file.\n",
    "\n",
    "**Note:** If you are unsure if the data was read in correctly, check the first plot below, it should contain all types of data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the names of the data files\n",
    "datafileinfo = {\n",
    "    'cases': {\n",
    "        'filename': 'data/sc_data_agestructured_binmode6.csv', \n",
    "        'columnnames': ('age 0 to 24', 'age 25 to 50', 'age 51 to 99'), \n",
    "        'datecolumnname': 'date'\n",
    "    },\n",
    "    'deaths': {\n",
    "        'filename': 'data/sc_deaths.csv', \n",
    "        'columnnames': ('deaths',), \n",
    "        'datecolumnname': 'date',\n",
    "    },\n",
    "    'hospitalizations': {\n",
    "        'filename': 'data/sc_hospitalization.csv', \n",
    "        'columnnames': ('positive-patients', 'positive-icu'), \n",
    "        'datecolumnname':'date',\n",
    "    },\n",
    "    'immunizations': {\n",
    "        'filename': 'data/sc_immunizations_binmode6_extended.csv', \n",
    "        'columnnames': ('age 0 to 24', 'age 25 to 50', 'age 51 to 99'), \n",
    "        'datecolumnname':'date',\n",
    "    },\n",
    "}\n",
    "\n",
    "age_group_desc = datafileinfo['cases']['columnnames']\n",
    "\n",
    "# the start date at which to start the simulation\n",
    "startdate_str = '2021-05-01'\n",
    "startdate = parse(startdate_str)\n",
    "\n",
    "# the date up to which to include observation for fitting (typically today)\n",
    "rewinddate = datetime.datetime.now()\n",
    "\n",
    "# the number of days to run the model for\n",
    "data['nt'] = max(100, (rewinddate - parse(startdate_str)).days + 61)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model permits the specification of multiple age groups of the susceptiple population and the fraction of infected who require non-ICU hospitalization (`frac_hospmild`), ICU hospitalization with recovery (`frac_iculive`), the mortality rate of infected (`frac_mort`) and the fraction of asymptomatic cases (`frac_asym`) for each age group. The remaining cases are assumed to be \"mild\" cases who recover at home. \n",
    "\n",
    "The fraction of asymptomatic cases (`frac_asym`) need to specified for each age group.\n",
    "\n",
    "A log-linear age-dependence is assumed for non-ICU hospitalization, ICU hospitalization with recovery, and mortality, leading to an exponential increase in these outcomes with age. The model estimates slope and intercept (at age 45, roughly the center of the age spread) of the log-linear relationship. Both slope an intercept use normal distribution as priors. The parameters `mu_frac_hospmild_i45` and `sigma_frac_hospmild_i45` specify the mean and standard deviation of the intercept for non-ICU hospitalization, and `mu_frac_hospmild_slope` and `sigma_frac_hospmild_slope` specify mean and standard deviation for its slope. Mean and standard deviation for the other outcomes are specified analogously. \n",
    "\n",
    "The fraction of infectious cases are modeled using a Dirichlet distribution. The parameter `alpha_multiplier` adjusts the variance of the prior distribution for each fraction thereby allowing fractions to be estimated by the model (for lower values of `alpha_multiplier`, e.g. 100) or to remain fixed (higher values of `alpha_multiplier`, e.g. 1000). It is not recommended to estimate fractions if `nage` (above) is greater than 1, that is if multiple age groups are used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the number of age groups\n",
    "data['nage'] = 3\n",
    "\n",
    "# Note: The inputs below must be lists or arrays, even if just one age group is specified\n",
    "# (the enclosing \"[]\" are not optional).\n",
    "\n",
    "# age bin edges: 0, 25, 51, 100\n",
    "data['npop'] = np.array([89724, 91178, 92862])\n",
    "data['age'] = np.array([20, 38, 65])\n",
    "\n",
    "# fraction of asympotmatic for each age group\n",
    "data['frac_asym'] = np.full(data['nage'], fill_value=0.178)\n",
    "\n",
    "'''\n",
    "posterior estimates from model simulation that generated restart\n",
    "                      mean se_mean     sd   2.5%    25%    50%    75%  97.5%  n_eff   Rhat\n",
    "frac_hospmild_i45    -2.13  1.5e-3   0.04   -2.2  -2.16  -2.13   -2.1  -2.04    754    1.0\n",
    "frac_hospmild_slope   0.03  2.1e-4 9.0e-3   0.01   0.02   0.03   0.04   0.05   1805    1.0\n",
    "frac_iculive_i45     -3.48  5.4e-3   0.17  -3.84  -3.59  -3.47  -3.36  -3.15   1021   1.01\n",
    "frac_iculive_slope    0.03  2.7e-4 9.9e-3   0.01   0.02   0.03   0.04   0.05   1334    1.0\n",
    "frac_mort_i45        -2.77  6.6e-4   0.02   -2.8  -2.78  -2.77  -2.76  -2.74    573    1.0\n",
    "frac_mort_slope       0.04  2.8e-4 9.9e-3   0.02   0.03   0.04   0.05   0.06   1216    1.0\n",
    "'''\n",
    "\n",
    "# mean and standard deviation of the intercept (at age 45) for non-ICU hospitalization\n",
    "data['mu_frac_hospmild_i45'] = -2.13\n",
    "data['sigma_frac_hospmild_i45'] = 0.04\n",
    "\n",
    "# mean and standard deviation of the slope for non-ICU hospitalization\n",
    "data['mu_frac_hospmild_slope'] = 0.03\n",
    "data['sigma_frac_hospmild_slope'] = 9.0e-3\n",
    "\n",
    "# mean and standard deviation of the intercept (at age 45) for ICU hospitalization with recovery\n",
    "data['mu_frac_iculive_i45'] = -3.48\n",
    "data['sigma_frac_iculive_i45'] = 0.17\n",
    "\n",
    "# mean and standard deviation of the slope for ICU hospitalization with recovery\n",
    "data['mu_frac_iculive_slope'] = 0.03\n",
    "data['sigma_frac_iculive_slope'] = 9.9e-3\n",
    "\n",
    "# mean and standard deviation of the intercept (at age 45) for mortality\n",
    "data['mu_frac_mort_i45'] = -2.77\n",
    "data['sigma_frac_mort_i45'] = 0.02\n",
    "\n",
    "# mean and standard deviation of the slope for mortality\n",
    "data['mu_frac_mort_slope'] = 0.04\n",
    "data['sigma_frac_mort_slope'] = 9.9e-3\n",
    "\n",
    "# a parameter modifying the prior uncertainty in the fractions above (set to 1000.0 for `nage` > 1)\n",
    "data['alpha_multiplier'] = 1000.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following parameters specify the prior distribution for the model parameters. Most of the parameters are modeled by a normal distribution for which a mean ($\\mu$, `mu_*`) and standard deviation ($\\sigma$, `sigma_*`) need to be specified.\n",
    "\n",
    "Currently, the initial conditions at the start date of the simulation assume zero infectious cases and a small number of exposed individuals. The prior distribution for the number of initially exposed individuals is modeled using an exponential distribution with a parameter `lambda_ini_exposed`. 1/`lambda_ini_exposed` is the mean of the distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# rounded from the output of the notebook that generated these values for \n",
    "# Santa Cruz county and the 2021-05-01 restart date\n",
    "\n",
    "# mean duration in \"exposed\" stage\n",
    "data['mu_duration_lat'] = 4.0\n",
    "# standard deviation (sd) of duration in \"exposed\" stage\n",
    "data['sigma_duration_lat'] = 1.0\n",
    "# mean duration in \"infectious\" stage for asymptomatic cases\n",
    "data['mu_duration_rec_asym'] = 7.0\n",
    "# sd of duration in \"infectious\" stage for asymptomatic cases\n",
    "data['sigma_duration_rec_asym'] = 3.0\n",
    "# mean duration in \"infectious\" stage for mild cases\n",
    "data['mu_duration_rec_mild'] = 5.0\n",
    "# sd of duration in \"infectious\" stage for mild cases\n",
    "data['sigma_duration_rec_mild'] = 2.0\n",
    "# mean duration in \"infectious\" stage for hospitalized cases\n",
    "data['mu_duration_pre_hosp'] = 7.0\n",
    "# sd of duration in \"infectious\" stage for hospitalized cases\n",
    "data['sigma_duration_pre_hosp'] = 1.0\n",
    "# mean duration in hospital for non-ICU cases\n",
    "data['mu_duration_hosp_mod'] = 8.0\n",
    "# sd of duration in hospital for non-ICU cases\n",
    "data['sigma_duration_hosp_mod'] = 1.0\n",
    "# mean duration in hospital for ICU cases\n",
    "data['mu_duration_hosp_icu'] = 7.0\n",
    "# sd of duration in hospital for ICU cases\n",
    "data['sigma_duration_hosp_icu'] = 1.0\n",
    "\n",
    "# lambda parameter for initial conditions of \"exposed\"\n",
    "data['lambda_ini_exposed'] = 0.3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following 4 parameters are related to the fitting of observations in the model. Each of the 4 observation types (confirmed cases (`Iobs`), non-ICU + ICU hospitalizations (`H`), ICU hospitalizations (`Hicu`), and deaths (`Rmort`)) has its own observation uncertainty, and the misfit between model and observations varies between the types. The standard deviation of the model-observation misfit is assumed to have an exponential distribution with a parameter $\\lambda$ (`lambda`). The 4 parameters below specify the value of $\\lambda$ for each observation type. Because $1/\\lambda$ is the mean of an exponential distribution, a larger value of a specific $\\lambda$ presumes a tighter fit to the corresponding observation type and lowering the value loosens the fit.\n",
    "\n",
    "The current default presumes a tighter fit to the mortality data (larger value of `lambda_Rmort`) and a looser fit to confirmed case count (lower value of `lambda_Iobs`) with the hospitalization data in between."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# lambda parameter for Iobs observation misfit\n",
    "data['lambda_Iobs'] = 0.03\n",
    "# lambda parameter for Hmod observation misfit\n",
    "data['lambda_H'] = 3.0\n",
    "# lambda parameter for Hicu observation misfit\n",
    "data['lambda_Hicu'] = 1.0\n",
    "# lambda parameter for Rmort observation misfit\n",
    "data['lambda_Rmort'] = 10.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next set of parameters are related to the fraction of infectious who have been tested, which links newly infectious in the model to the confirmed cases in the data. We assume time-dependence in the prior, e.g. to account for an increase in testing.\n",
    "\n",
    "With time-dependence, the time-evolution of the fraction tested is prescribed using a spline curve of degree `p_fractest`. Using `p_fractest = 1` (like below), uses linear interpolation in time to describe the time-evolution of the fraction tested prior: At a set of time points, normal distributions are prescribed (using mean and standard deviation of the fraction of newly infectious who get tested) which are interpolated between the knots. Use `create_spline_plot = True` to create a plot of the interpolation result before running the Stan model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the order of the splines used (set to 1 for linear interpolation)\n",
    "data['p_fractest'] = 1\n",
    "\n",
    "# the time-points at which to prescribe a new distribution\n",
    "# Note: The first entry must be the start day, the last must be the last day of the simulation\n",
    "# (if this is not the case they are automatically corrected).\n",
    "dates_fractest_str = [startdate_str, '2021-06-01', '2021-07-01', '2021-08-01', '2021-09-01']\n",
    "\n",
    "# mean fraction of tested infectious (one entry for every entry in dates_fractest_str above)\n",
    "data['mu_fractest'] = np.array([0.8,0.8,0.8,0.8,0.8]) / 0.8 * 0.3\n",
    "# sd fraction of tested infectious(one entry for every entry in dates_fractest_str above)\n",
    "data['sigma_fractest'] = np.array([0.05,0.05,0.05,0.05,0.05]) / 0.8 * 0.3\n",
    "\n",
    "# If create_spline_plot is set to true, a plot of the interpolation result\n",
    "# for the mean of the prior distribution) is created below the data plot.\n",
    "create_spline_plot = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next set of parameters is used to account for a background level of \"imported\" infectious individuals. These imports represent both infectious individuals entering the model region (visitors to Santa Cruz county) as well as infectious individuals encountered outside the model region by individuals from the model region (Santa Cruz county residents visiting other places).\n",
    "\n",
    "Just like the fraction tested above, imports are modeled using splines using the same time-points. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the order of the splines used (set to 1 for linear interpolation)\n",
    "data['p_imports'] = 1\n",
    "\n",
    "# mean of number of imports (one entry for every entry in dates_fractest_str above)\n",
    "data['mu_imports'] = np.array([2.0, 2.0, 2.0, 2.0, 2.0])\n",
    "\n",
    "# standard deviation of number of imports (one entry for every entry in dates_fractest_str above)\n",
    "data['sigma_imports'] = np.array([1.0, 1.0, 1.0, 1.0, 1.0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model also requires the specification of the distribution for the initial value of $\\beta$ using a normal distribution with mean `mu_beta1` and standard deviation `sigma_beta1`.\n",
    "\n",
    "The remaining values for $\\beta$ are estimated based on the input data every `dknot` days (using an AR(1) process and Spline interpolation). The spline interpolation offers two transition modes to extrapolate the values of $\\beta$ beyond the current date after which no data is available. Setting `splinemode` to 1 will allow the spline to be estimated freely until the current date. The more conservative option `splinemode` = 2 assumes a constant value for $\\beta$ in the last `dknot` days before the current date.    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "estimates from model simulation that generated restart\n",
    "mean beta: [0.08871386 0.14129522 0.10782254]\n",
    "sd beta:   [0.03295604 0.0419682  0.04301105]\n",
    "'''\n",
    "\n",
    "data['mu_beta1'] = np.array([0.088714, 0.141295, 0.107823])\n",
    "data['sigma_beta1'] = 0.039\n",
    "\n",
    "# distance in time between the knots used to construct the splines\n",
    "data['dknot'] = 10\n",
    "\n",
    "# itoday_1based is the (1-based) time index corresponding to the current date\n",
    "itoday_1based = (rewinddate-parse(startdate_str)).days + 1\n",
    "# specify the date up to when beta is estimated\n",
    "data['itoday'] = itoday_1based\n",
    "\n",
    "# spline mode (must be 1 or 2)\n",
    "# splinemode = 1: estimate beta up to today\n",
    "# splinemode = 2: estimate beta up to dknot days before today, then assume constant value up to today\n",
    "data['splinemode'] = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Interventions permit the modification of $\\beta$ past the current date. `ninter` is the number of interventions (allowed to be zero for no interventions). Each intervention has a start time (`t_inter`) and a length (in days) over which $\\beta$ is adjusted. Each intervention leads to a multiplication of the value of $\\beta$ before the intervention with a random number drawn from a normal distribution with mean `mu_beta_inter` and standard deviation `sigma_beta_inter`. That is, an intervention with `mu_beta_inter` set to 1.2 will lead to an average increase in $\\beta$ by 20%, 0.8 to an average decrease by 20%; a value of 1.0 will lead to no mean increase but a higher spread (less certainty) in the value of $\\beta$ after the intervention."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# number of interventions\n",
    "data['ninter'] = 1\n",
    "\n",
    "# Note: The inputs below must be lists or arrays, even if one intervention is specified\n",
    "# (the enclosing \"[]\" are not optional for single interventions).\n",
    "\n",
    "# start time of each interventions\n",
    "data['t_inter'] = [itoday_1based+5]\n",
    "# length of each intervention\n",
    "data['len_inter'] = [10]\n",
    "# mean change in beta through intervention\n",
    "data['mu_beta_inter'] = [1.0]\n",
    "# sd change in beta through intervention\n",
    "data['sigma_beta_inter'] = [0.2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Options for restart:** For a model start with a non-zero number of individuals that are infectious, hospitalized, vaccinated etc., the model requires the specification of distributions for the number of individuals in each compartment at the beginning of the simulation.\n",
    "\n",
    "A lognormal distribution is used to estimate the initial number of exposed individuals (compartment 2); normal distributions are used for the initial conditions of compartments 3 to 10. Both distributions require the specification of $\\mu$ and $\\sigma$ values, one for each age group. These are specified via the `mu_x_ini_2to10` and `sigma_x_ini_2to10` parameters (which contain `2to10` in their name to indicate that they specify the initial conditions for compartments 2 to 10).\n",
    "\n",
    "The paramater `V_ini` specifies the total number of immunized (vaccinated) inidividuals in each age group. To compute the number of vaccinated (previously recovered) individuals (compartment 11), the model creates an estimate for the number of vaccinated (previously susceptible; compartment 10) and subtracs that number from `V_ini`.\n",
    "\n",
    "To obtain the initial estimate for the number of susceptible individuals (compartment 1), the model subtracts the sum of individuals in compartments 2 through 11 from the total population (specified by the parameter `npop`, see above)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# estimates from model simulation that generated restart\n",
    "\n",
    "data['mu_x_ini_2to10'] = np.array((\n",
    "      (    3.444,     3.498,     2.810), # exposed (lognorm)\n",
    "      (   10.911,    11.752,     6.262), # infectious (norm)\n",
    "      (   34.910,    37.238,    18.962), # recovering at home (norm)\n",
    "      (    0.100,     0.392,     1.521), # pre-hospital (norm)\n",
    "      (    0.107,     0.407,     1.419), # hospitalized (non-ICU) (norm)\n",
    "      (    0.015,     0.080,     0.468), # ICU (norm)\n",
    "      (19168.830, 19860.837,  7321.913), # recovered (norm)\n",
    "      (    3.159,    23.345,   178.211), # deceased (norm)\n",
    "      ( 4326.259, 20355.454, 45601.950), # vaccinated (previously susceptible) (norm)\n",
    "   )).T\n",
    "\n",
    "data['sigma_x_ini_2to10'] = np.array((\n",
    "      (    0.324,     0.259,     0.394), # exposed (lognorm)\n",
    "      (    6.673,     7.091,     4.102), # infectious (norm)\n",
    "      (   13.892,    14.510,     9.086), # recovering at home (norm)\n",
    "      (    0.025,     0.076,     0.430), # pre-hospital (norm)\n",
    "      (    0.019,     0.056,     0.300), # hospitalized (non-ICU) (norm)\n",
    "      (    0.003,     0.013,     0.110), # ICU (norm)\n",
    "      (  713.950,   751.292,   285.387), # recovered (norm)\n",
    "      (    0.076,     0.429,     0.658), # deceased (norm)\n",
    "      (   47.368,   364.511,   453.400), # vaccinated (previously susceptible) (norm)\n",
    "   )).T\n",
    "\n",
    "data['V_ini'] = np.array([5575.0, 29761.0, 56952.0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Advanced options for performing multiple simulations (leave unchanged if not needed)**. The parameter `stan_file` specifies Stan code file used in the simulation. `datachanges` allows running multiple simulations with different parameters (optional)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data['knots_age'] = data['age']\n",
    "data['nknots_age'] = len(data['knots_age'])\n",
    "\n",
    "stan_file = {'SEIR':'stan_code/seir_v12b.continue-v2.stan'}\n",
    "\n",
    "# define alternative parametrizations here (optional, comment out if not needed)\n",
    "datachanges = {\n",
    "    'v12b':None,\n",
    "    #'v12b_14+20':{'ninter':1, 't_inter':[itoday_1based], 'len_inter':[14], 'mu_beta_inter':[1.2], 'sigma_beta_inter':[0.1]},\n",
    "    #'v12b_14+50':{'ninter':1, 't_inter':[itoday_1based], 'len_inter':[14], 'mu_beta_inter':[1.5], 'sigma_beta_inter':[0.1]},\n",
    "    #'v12b_14+100':{'ninter':1, 't_inter':[itoday_1based], 'len_inter':[14], 'mu_beta_inter':[2.0], 'sigma_beta_inter':[0.1]},\n",
    "}\n",
    "\n",
    "# a brief description of each parametrization (optional, comment out if not needed)\n",
    "desc = {\n",
    "    'v12b':'v12b',\n",
    "    #'v12b_14+20':'v12b 14d +20%',\n",
    "    #'v12b_14+50':'v12b 14d +50%',\n",
    "    #'v12b_14+100':'v12b 14d +100%',\n",
    "}\n",
    "\n",
    "# if running different Stan code versions provide a key to a stan_file entry for each datachanges entry\n",
    "# (optional, comment out if not needed)\n",
    "datachange_stan_key = {\n",
    "    'v12b':'SEIR',\n",
    "    #'v12b_14+20':'SEIR',\n",
    "    #'v12b_14+50':'SEIR',\n",
    "    #'v12b_14+100':'SEIR',\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## end of user input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_data(filename, colname_date='date', colnames_data=None):\n",
    "    if os.path.splitext(filename)[1] in ('.xls','.xlsx'):\n",
    "        df = pd.read_excel(filename)\n",
    "    elif os.path.splitext(filename)[1] == '.csv':\n",
    "        df = pd.read_csv(filename, comment='#')\n",
    "    else:\n",
    "        raise ValueError('Data file extension must be \".csv\", \".xlsx\" or \".xls\".')\n",
    "    return np.array([pd.to_datetime(d) for d in df[colname_date].values]), df.loc[:,colnames_data].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# note that the reference date is one day ahead of the start date\n",
    "# (day 1 when Stan starts simulation is 2020-02-17)\n",
    "refdate = parse(startdate_str) - datetime.timedelta(days=1)\n",
    "itoday = itoday_1based-1\n",
    "\n",
    "obs = {}\n",
    "tobs_date = {}\n",
    "\n",
    "# read in case and death data\n",
    "\n",
    "for name, key in (('cases', 'I'), ('deaths', 'Rmort')):\n",
    "    tobs_date[key], obs[key] = read_data(filename=datafileinfo[name]['filename'],\n",
    "                                         colnames_data=datafileinfo[name]['columnnames'],\n",
    "                                         colname_date=datafileinfo[name]['datecolumnname']\n",
    "                                        )\n",
    "    tobs = np.array([(d-refdate).days for d in tobs_date[key]])\n",
    "\n",
    "    itobs = tobs < itoday_1based\n",
    "    itobs &= tobs > 0\n",
    "    if name == 'cases':\n",
    "        ndays = 7\n",
    "        itobs &= tobs <= (rewinddate - startdate).days - ndays\n",
    "    tobs = tobs[itobs]\n",
    "    obs[key] = obs[key][itobs,...]\n",
    "    tobs_date[key] = tobs_date[key][itobs]\n",
    "\n",
    "    # update Stan input with data read from files\n",
    "    data['nobs_'+key] = len(tobs)\n",
    "    data['tobs_'+key] = tobs\n",
    "    if obs[key].ndim == 2:\n",
    "        if obs[key].shape[1] == 1:\n",
    "            obs[key] = obs[key][:,0]\n",
    "        else:\n",
    "            obs[key] = obs[key].T\n",
    "    data['obs_'+key] = obs[key]\n",
    "\n",
    "# read in immunization data\n",
    "\n",
    "tobs_date['vac'], tmp = read_data(filename=datafileinfo['immunizations']['filename'],\n",
    "                                  colnames_data=datafileinfo['immunizations']['columnnames'],\n",
    "                                  colname_date=datafileinfo['immunizations']['datecolumnname']\n",
    "                                 )\n",
    "\n",
    "tobs = np.array([(d-refdate).days for d in tobs_date['vac']])\n",
    "\n",
    "itobs = tobs > 0\n",
    "\n",
    "data['vac'] = tmp[itobs,:].T\n",
    "data['tvac'] = tobs[itobs]\n",
    "data['nvac'] = len(data['tvac'])\n",
    "\n",
    "# read in hospitalization data\n",
    "    \n",
    "tobs_date['H'], tmp = read_data(filename=datafileinfo['hospitalizations']['filename'],\n",
    "                                colnames_data=datafileinfo['hospitalizations']['columnnames'],\n",
    "                                colname_date=datafileinfo['hospitalizations']['datecolumnname']\n",
    "                               )\n",
    "obs['H'] = tmp[:,0]\n",
    "obs['Hicu'] = tmp[:,1]\n",
    "\n",
    "tobs = np.array([(d-refdate).days for d in tobs_date['H']])\n",
    "\n",
    "itobs = tobs < itoday_1based\n",
    "itobs &= tobs > 0\n",
    "tobs = tobs[itobs]\n",
    "obs['H'] = obs['H'][itobs]\n",
    "obs['Hicu'] = obs['Hicu'][itobs]\n",
    "tobs_date['H'] = tobs_date['H'][itobs]\n",
    "\n",
    "# update Stan input with data read from files\n",
    "data.update({'nobs_H':len(tobs), 'tobs_H':tobs, 'obs_H':obs['H'], 'obs_Hicu':obs['Hicu']})\n",
    "\n",
    "# add additional parameters\n",
    "\n",
    "data['knots_fractest'] = np.array([(parse(d)-refdate).days for d in dates_fractest_str])\n",
    "data['knots_fractest'][0] = 1\n",
    "data['knots_fractest'][-1] = data['nt']\n",
    "data['nknots_fractest'] = len(data['knots_fractest'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### prepare plotting and plot the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "variable_names = ['S','E','Iasym', 'Imild', 'Ipreh', 'Hmod', 'Hicu', 'Rlive', 'Rmort', 'Vsus', 'Vrec']\n",
    "variable_index = {v:i for i,v in enumerate(variable_names)}\n",
    "variable_desc = {'S':'susceptible',\n",
    "                 'E':'exposed',\n",
    "                 'I':'infectious',\n",
    "                 'Iasym':'infectious',\n",
    "                 'Iobs':'observed infectious (cumulative)',\n",
    "                 'Imild':'recovering at home',\n",
    "                 'Ipreh':'pre-hospital',\n",
    "                 'Hmod':'hospitalized (non-ICU)',\n",
    "                 'Hicu':'ICU',\n",
    "                 'Rlive':'recovered',\n",
    "                 'rectotal':'recovered (includes immunized)',\n",
    "                 'Rmort':'deceased',\n",
    "                 'hospitalized':'hospitalized (ICU + non-ICU)',\n",
    "                 'Vsus':'vaccinated (previously susceptible)',\n",
    "                 'Vrec':'vaccinated (previously recovered)',\n",
    "                }\n",
    "variable_colors = {'S':'tab:blue',\n",
    "                   'E':'tab:orange',\n",
    "                   'I':'tab:red',\n",
    "                   'Iobs':'tab:gray',\n",
    "                   'Iasym':'lightcoral',\n",
    "                   'Imild':'tab:olive',\n",
    "                   'Ipreh':'tab:purple',\n",
    "                   'Hmod':'lightsteelblue',\n",
    "                   'Hicu':'#2a4e6c', # a 'darksteelblue'\n",
    "                   'Rlive':'tab:green',\n",
    "                   'rectotal':'tab:green',\n",
    "                   'Rmort':'tab:brown',\n",
    "                   'hospitalized':'steelblue',\n",
    "                   'Vsus':'olive',\n",
    "                   'Vrec':'lawngreen',\n",
    "                  }\n",
    "\n",
    "for k in ('hospitalized', 'Hicu'):\n",
    "    variable_desc[k+'+sigma'] = variable_desc[k]\n",
    "    variable_colors[k+'+sigma'] = variable_colors[k]\n",
    "\n",
    "age_colors = [mpl.cm.plasma_r((i+1)/(data['nage'])) for i in range(data['nage'])]\n",
    "#age_colors = ['C{}'.format(i) for i in range(data['nage'])]\n",
    "\n",
    "create_data_plot = False\n",
    "if create_data_plot:\n",
    "    fig, axs = plt.subplots(nrows=2, figsize=(14,12), sharex=True)\n",
    "    ax = axs[0]\n",
    "    for iage in range(data['nage']):\n",
    "        ax.plot(tobs_date['I'], obs['I'][iage,:], ls='none', marker='o', markersize=8, color=age_colors[iage], label='number of confirmed cases {}'.format(age_group_desc[iage]))\n",
    "    ax.set_title('the data', size=20)\n",
    "\n",
    "    ax = axs[1]\n",
    "    ax.plot(tobs_date['Rmort'], obs['Rmort'], ls='none', marker='D', markersize=8, color=variable_colors['Rmort'], label='number of confirmed deaths')\n",
    "    ax.plot(tobs_date['H'], obs['H'], ls='none', marker='H', markersize=8, color=variable_colors['hospitalized'], label='number of confirmed hospitalized')\n",
    "    ax.plot(tobs_date['H'], obs['Hicu'], ls='none', marker='^', markersize=8, color=variable_colors['Hicu'], label='number of confirmed ICU')\n",
    "\n",
    "    for ax in axs.flat:\n",
    "        ax.set_ylim(bottom=0)\n",
    "        ax.axvline(rewinddate, color='black')\n",
    "        ax.grid(True)\n",
    "        ax.legend()\n",
    "\n",
    "if create_spline_plot:\n",
    "    from scipy.interpolate import BSpline\n",
    "    p = data['p_fractest']\n",
    "    knots_ext = np.empty(len(data['knots_fractest'])+2*p)\n",
    "    if p > 0:\n",
    "        knots_ext[:p] = data['knots_fractest'][0]\n",
    "        knots_ext[p:-p] = data['knots_fractest']\n",
    "        knots_ext[-p:] = data['knots_fractest'][-1]\n",
    "    else:\n",
    "        knots_ext[:] = data['knots_fractest']\n",
    "    t = np.arange(1,data['nt']+1)\n",
    "    t_date = np.array([refdate+datetime.timedelta(days=float(tt)) for tt in t])\n",
    "    bsp = BSpline(knots_ext, data['mu_fractest'], p)\n",
    "\n",
    "    fig,ax = plt.subplots(figsize=(14,8))\n",
    "    tmp = ax.plot([refdate+datetime.timedelta(days=float(tt)) for tt in knots_ext[p:p+len(data['mu_fractest'])]], data['mu_fractest'], marker='o', ls='none', markersize=9)\n",
    "    ax.plot(t_date, bsp(t), ls=':', color=tmp[0].get_color(), lw=2)\n",
    "    ax.grid(True)\n",
    "    ax.set_title('spline-interpolated \"mu_fractest\" (for prior distribution of fraction tested)', size=20)\n",
    "\n",
    "    ax.axvline(refdate+datetime.timedelta(days=1), color='black')\n",
    "    ax.axvline(rewinddate.date(), color='black')\n",
    "None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### set `mu_Iobs_ini` based on data values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data['mu_Iobs_ini'] = obs['I'][:,0]\n",
    "print('mu_Iobs_ini:', data['mu_Iobs_ini'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## compile and run Stan model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# avoid unnecessary recompilation\n",
    "if 'models' not in globals():\n",
    "    models = {}\n",
    "if 'fits' not in globals():\n",
    "    fits = {}\n",
    "if 'datachanges' not in globals():\n",
    "    datachanges = {}\n",
    "#if 'ref' not in datachanges:\n",
    "#    datachanges['ref'] = None\n",
    "if 'desc' not in globals():\n",
    "    desc = {'ref':'reference simulation'}\n",
    "for name in datachanges:\n",
    "    if name not in desc:\n",
    "        desc[name] = name\n",
    "\n",
    "# allow stan_file to be dictionary containing different versions of code\n",
    "# if it is not a dictionary, convert it to one\n",
    "if not isinstance(stan_file, dict):\n",
    "    stan_file = {'SEIR':stan_file}\n",
    "    datachange_stan_key = {name:'SEIR' for name in datachanges}\n",
    "else:\n",
    "    if 'datachange_stan_key' not in globals():\n",
    "        if len(stan_file) == 1:\n",
    "            datachange_stan_key = {name:list(stan_file.keys())[0] for name in datachanges}\n",
    "        else:\n",
    "            logging.warning('''Use\n",
    "   datachange_stan_key[key_datachange] = key_stan_file\n",
    "to associate the data change entry with with key \"key_datachange\" to the stan_file entry with key \"key_stan_file\".''')\n",
    "            raise ValueError('A dictionary \"datachange_stan_key\" must be used to associate data change entries with Stan file versions.')\n",
    "\n",
    "stan_codes = {}\n",
    "for stanname,filename in stan_file.items():\n",
    "    with open(filename) as f:\n",
    "        stan_codes[stanname] = f.read()\n",
    "\n",
    "refit_all = False\n",
    "#refit_required = {}\n",
    "\n",
    "for name in datachanges:\n",
    "    if name in models and models[name].model_code == stan_codes[datachange_stan_key[name]]:\n",
    "        print('{}: unchanged model code, not recompiling'.format(datachange_stan_key[name]))\n",
    "    else:\n",
    "        models[name] = pystan.StanModel(file=stan_file[datachange_stan_key[name]], model_name=datachange_stan_key[name], obfuscate_model_name=False)\n",
    "\n",
    "def get_max_rhat(fit):\n",
    "    s = fit.summary()\n",
    "    irhat = s['summary_colnames'].index(\"Rhat\")\n",
    "    return np.nanmax(s['summary'][:,irhat])\n",
    "\n",
    "num_tries = 1\n",
    "datasave = {}\n",
    "for name,datachange in datachanges.items():\n",
    "    datasave[name] = data.copy()\n",
    "    if datachange is not None:\n",
    "        datasave[name].update(datachange)\n",
    "\n",
    "    if refit_all or name not in fits or get_max_rhat(fits[name]) >= 1.1:\n",
    "        for itry in range(num_tries):\n",
    "            cfit = models[name].sampling(data=datasave[name], iter=2000, chains=4)\n",
    "            max_rhat = get_max_rhat(cfit)\n",
    "            if name not in fits or max_rhat < get_max_rhat(fits[name]):\n",
    "                fits[name] = cfit\n",
    "            if max_rhat < 1.1:\n",
    "                break\n",
    "            logging.warning('{}: in try {}/{} found Rhat={:.3f}, trying again'.format(name,itry+1,num_tries,max_rhat))\n",
    "    else:\n",
    "        print('found fit for \"{}\", not re-running model'.format(name))\n",
    "    print('dataset \"{}\"'.format(name))\n",
    "    # just print a subset of the variables (skip the long ones)\n",
    "    print('\\n'.join(x for x in fits[name].__str__().split('\\n') if '[' not in x or 'sigma_' in x or 'beta1' in x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "quantiles = {}\n",
    "for name,fit in fits.items():\n",
    "    x_noage = np.sum(fit['x'][:,:,:,:], axis=1)\n",
    "    quantiles[name] = {v:np.percentile(x_noage[:,variable_index[v],:], q=(5,25,50,75,95), axis=0) for v in variable_names}\n",
    "    quantiles[name]['hospitalized'] = np.percentile(fit['hospitalized'], q=(5,25,50,75,95), axis=0)\n",
    "    quantiles[name]['Iobs'] = np.percentile(fit['obs_I_sim'], q=(5,25,50,75,95), axis=0)\n",
    "    quantiles[name]['I'] = np.percentile(x_noage[:,variable_index['Iasym'],:]+\n",
    "                                         x_noage[:,variable_index['Imild'],:]+\n",
    "                                         x_noage[:,variable_index['Ipreh'],:], q=(5,25,50,75,95), axis=0)\n",
    "    quantiles[name]['rectotal'] = np.percentile(x_noage[:,variable_index['Rlive'],:]+\n",
    "                                                x_noage[:,variable_index['Vrec'],:]+\n",
    "                                                x_noage[:,variable_index['Vsus'],:], q=(5,25,50,75,95), axis=0)\n",
    "    # new: testing new approach\n",
    "    n = x_noage.shape[0]\n",
    "    quantiles[name]['hospitalized+sigma'] = np.percentile(x_noage[:,variable_index['Hmod'],:] +\n",
    "                                                          x_noage[:,variable_index['Hicu'],:] +\n",
    "                                                          np.random.normal(scale=fit['sigma_H'][:,None]),\n",
    "                                                          q=(5,25,50,75,95), axis=0)\n",
    "    quantiles[name]['Hicu+sigma'] = np.percentile(x_noage[:,variable_index['Hicu'],:] +\n",
    "                                                  np.random.normal(scale=fit['sigma_Hicu'][:,None]),\n",
    "                                                  q=(5,25,50,75,95), axis=0)\n",
    "t_date = np.array([refdate+datetime.timedelta(days=float(tt+1)) for tt in range(data['nt'])])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## report some estimates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def print_report(fit, times, variables=('hospitalized'), title=None):\n",
    "    print('=== {}'.format(title))\n",
    "    for itime,time in enumerate(times):\n",
    "        if isinstance(time,int):\n",
    "            it = time\n",
    "        elif isinstance(time,str):\n",
    "            it = (parse(time)-refdate).days - 1\n",
    "        else:\n",
    "            it = (time-refdate).days - 1\n",
    "        delta_day = it - (rewinddate-refdate).days + 1\n",
    "\n",
    "        if delta_day == 0:\n",
    "            textday = 'today'\n",
    "        elif delta_day == 1:\n",
    "            textday = 'tomorrow'\n",
    "        elif delta_day < 0:\n",
    "            textday = '{} days ago'.format(-delta_day)\n",
    "        else:\n",
    "            textday = 'in {} days'.format(delta_day)\n",
    "\n",
    "        if it >= t_date.size:\n",
    "            print('   beyond {} ({}):'.format(t_date[-1].date(), textday))\n",
    "            print('    ! simulation not long enough')\n",
    "            continue\n",
    "        print('   {} ({}):'.format(t_date[it].date(), textday))\n",
    "        for v in variables:\n",
    "            if v in variable_index:\n",
    "                x_noage = np.sum(fit['x'][:,:,:,:], axis=1)\n",
    "                q = np.percentile(x_noage[:,variable_index[v],it], q=(5,25,50,75,95), axis=0)\n",
    "            elif v == 'Iobs':\n",
    "                q = np.percentile(fit['obs_I_sim'][:,it], q=(5,25,50,75,95), axis=0)\n",
    "            elif v == 'I':\n",
    "                x_noage = np.sum(fit['x'][:,:,:,:], axis=1)\n",
    "                q = np.percentile(x_noage[:,variable_index['Iasym'],it]+\n",
    "                                  x_noage[:,variable_index['Imild'],it]+\n",
    "                                  x_noage[:,variable_index['Ipreh'],it], q=(5,25,50,75,95), axis=0)\n",
    "            else:\n",
    "                q = np.percentile(fit[v][:,it], q=(5,25,50,75,95), axis=0)\n",
    "            if v in variable_index or v == 'Iobs' or v == 'hospitalized' or v == 'I':\n",
    "                print('      \"{}\" 50% credible interval: [{:.0f},{:.0f}]'.format(variable_desc[v],q[1],q[3]))\n",
    "                print('       {}  90% credible interval: [{:.0f},{:.0f}] with a median estimate of {:.0f}'.format(' '*len(variable_desc[v]),q[0],q[-1],q[2]))\n",
    "            else:\n",
    "                print('      \"{}\" 50% credible interval: [{:.2f},{:.2f}]'.format(variable_desc[v],q[1],q[3]))\n",
    "                print('       {}  90% credible interval: [{:.2f},{:.2f}] with a median estimate of {:.0f}'.format(' '*len(variable_desc[v]),q[0],q[-1],q[2]))\n",
    "\n",
    "\n",
    "today = rewinddate\n",
    "times = [today, today+datetime.timedelta(days=1), today+datetime.timedelta(days=7)]\n",
    "\n",
    "for name in fits:\n",
    "    print_report(fits[name],\n",
    "                 times=times,\n",
    "                 variables=('hospitalized','I','Hicu','Rmort'), title=desc[name])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## plot results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib as mpl\n",
    "\n",
    "mpl.rcParams['legend.fontsize'] = 16\n",
    "mpl.rcParams['axes.titlesize'] = 26\n",
    "mpl.rcParams['figure.figsize'] = (24,12)\n",
    "mpl.rcParams['xtick.labelsize'] = 16\n",
    "mpl.rcParams['ytick.labelsize'] = 16\n",
    "\n",
    "def add_xaxis_datestyle(ax, fontsize_small=16, fontsize_large=20):\n",
    "    ax.xaxis.set_major_locator(mpl.dates.MonthLocator())\n",
    "    ax.xaxis.set_minor_locator(mpl.dates.WeekdayLocator(mpl.dates.MONDAY))\n",
    "    ax.xaxis.set_major_formatter(mpl.dates.DateFormatter('%b'))\n",
    "    ax.xaxis.set_minor_formatter(mpl.dates.DateFormatter('%d'))\n",
    "    ax.grid(b=True, which='minor', linestyle=':')\n",
    "    ax.tick_params(axis='x', which='major', pad=15, labelsize=fontsize_large)\n",
    "    ax.tick_params(axis='x', which='minor', labelsize=fontsize_small)\n",
    "    ax.tick_params(axis='y', which='major', labelsize=fontsize_small)\n",
    "\n",
    "def add_intervention_info(ax, **kwargs):\n",
    "    for iinter in range(data['ninter']):\n",
    "        ax.axvline(refdate+datetime.timedelta(days=data['t_inter'][iinter]), color='tab:green')\n",
    "        ax.axvline(refdate+datetime.timedelta(days=data['t_inter'][iinter]+data['len_inter'][iinter]), color='tab:red')\n",
    "        ax.text(refdate+datetime.timedelta(days=data['t_inter'][iinter]+0.5*data['len_inter'][iinter]), 0.5, 'intervention {}'.format(iinter+1),\n",
    "                color='tab:green', transform=ax.get_xaxis_transform(), ha='center', va='center', rotation=90, size=20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## quantile plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import matplotlib.dates as mdates\n",
    "\n",
    "def create_quantile_plot(name, q, variables=('hospitalized',), title=None, daterange=None, show_data=True,\n",
    "                         show_50p=True, show_90p=True, legendargs=None, ylim_bottom=0.0):\n",
    "    fig,ax = plt.subplots()\n",
    "\n",
    "    today = rewinddate\n",
    "    if daterange is None:\n",
    "        index = slice(None)\n",
    "    else:\n",
    "        index = np.logical_and(t_date >= daterange[0], t_date <= daterange[1])\n",
    "\n",
    "    for v in variables:\n",
    "        if q[v].ndim == 2:\n",
    "            if show_90p:\n",
    "                ax.fill_between(t_date[index], q[v][0,index], q[v][4,index], color=variable_colors[v], alpha=0.25)\n",
    "            if show_50p:\n",
    "                ax.fill_between(t_date[index], q[v][1,index], q[v][3,index], color=variable_colors[v], alpha=0.5)\n",
    "            ax.plot(t_date[index], q[v][2,index], color=variable_colors[v], label=variable_desc[v])\n",
    "        else:\n",
    "            for iage in range(data['nage']):\n",
    "                if show_90p:\n",
    "                    ax.fill_between(t_date[index], q[v][0,iage,index], q[v][4,iage,index], color=age_colors[iage], alpha=0.25)\n",
    "                if show_50p:\n",
    "                    ax.fill_between(t_date[index], q[v][1,iage,index], q[v][3,iage,index], color=age_colors[iage], alpha=0.5)\n",
    "                ax.plot(t_date[index], q[v][2,iage,index], color=age_colors[iage], label=variable_desc[v]+' '+age_group_desc[iage])\n",
    "\n",
    "    if show_data:\n",
    "        if 'Rmort' in variables:\n",
    "            if daterange is None:\n",
    "                index = slice(None)\n",
    "            else:\n",
    "                index = np.logical_and(tobs_date['Rmort'] >= daterange[0], tobs_date['Rmort'] <= daterange[1])\n",
    "            ax.plot(tobs_date['Rmort'][index], obs['Rmort'][index], ls='none', marker='D', markersize=10, color=variable_colors['Rmort'], label='number of deceased', zorder=4)\n",
    "        if 'hospitalized' in variables or 'hospitalized+sigma' in variables:\n",
    "            if daterange is None:\n",
    "                index = slice(None)\n",
    "            else:\n",
    "                index = np.logical_and(tobs_date['H'] >= daterange[0], tobs_date['H'] <= daterange[1])\n",
    "            ax.plot(tobs_date['H'][index], obs['H'][index], ls='none', marker='H', markersize=10, color=variable_colors['hospitalized'], label='number of confirmed hospitalized', zorder=4)\n",
    "        if 'Hicu' in variables or 'Hicu+sigma' in variables:\n",
    "            if daterange is None:\n",
    "                index = slice(None)\n",
    "            else:\n",
    "                index = np.logical_and(tobs_date['H'] >= daterange[0], tobs_date['H'] <= daterange[1])\n",
    "            ax.plot(tobs_date['H'][index], obs['Hicu'][index], ls='none', marker='^', markersize=10, color=variable_colors['Hicu'], label='number of confirmed ICU', zorder=4)\n",
    "        if 'I' in variables or 'Iobs' in variables:\n",
    "            if daterange is None:\n",
    "                index = slice(None)\n",
    "            else:\n",
    "                index = np.logical_and(tobs_date['I'] >= daterange[0], tobs_date['I'] <= daterange[1])\n",
    "            for iage in range(data['nage']):\n",
    "                ax.plot(tobs_date['I'][index], obs['I'][iage,index], ls='none', marker='o', color=age_colors[iage], label='number of confirmed cases '+age_group_desc[iage])\n",
    "\n",
    "    if ylim_bottom is not None:\n",
    "        ax.set_ylim(bottom=ylim_bottom) #, top =55000)\n",
    "\n",
    "    if daterange is None or (today >= daterange[0] and today <= daterange[1]):\n",
    "        ax.axvline(today.date(), color='black')\n",
    "\n",
    "    if daterange is not None:\n",
    "        ax.set_xlim(daterange)\n",
    "    add_xaxis_datestyle(ax)\n",
    "    if legendargs is None:\n",
    "        ax.legend()\n",
    "    else:\n",
    "        ax.legend(**legendargs)\n",
    "\n",
    "    ax.grid(True)\n",
    "    if title is not None:\n",
    "        ax.set_title(title)\n",
    "    return ax\n",
    "\n",
    "# customize date range variables etc. here\n",
    "\n",
    "for name in quantiles:\n",
    "    ax = create_quantile_plot(name, quantiles[name], variables=('Iobs',), daterange=(t_date[0], t_date[-1]), title=desc[name])\n",
    "\n",
    "None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## rate plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "daterange_export = (datetime.datetime(2021,5,1), datetime.datetime(2021,7,15))\n",
    "\n",
    "today = rewinddate\n",
    "# add 1st or 15th\n",
    "min_lookahead = 20\n",
    "for td in range(min_lookahead,50):\n",
    "    d = today+datetime.timedelta(days=td)\n",
    "    if d.day == 1 or d.day == 15:\n",
    "        daterange_export = (d-datetime.timedelta(days=100), d)\n",
    "        break\n",
    "\n",
    "def create_rate_quantile_plot(fit, variables=('hospitalized',), title=None, daterange=None, show_50p=True, legendargs=None):\n",
    "    for v in variables:\n",
    "        if v in variable_index:\n",
    "            prop = np.sum(fit['x'], axis=1)[:,variable_index[v],:]\n",
    "        elif v == 'Iobs':\n",
    "            prop = fit['obs_I_sim']\n",
    "        elif v == 'I':\n",
    "            x_noage = np.sum(fit['x'][:,:,:,:], axis=1)\n",
    "            prop = x_noage[:,variable_index['Iasym'],:]+x_noage[:,variable_index['Imild'],:]+x_noage[:,variable_index['Ipreh'],:]\n",
    "        else:\n",
    "            prop = fit['hospitalized']\n",
    "\n",
    "    qq = np.percentile(prop[:,1:]-prop[:,:-1], axis=0, q=(5,25,50,75,95))\n",
    "    x = t_date[:-1] + 0.5 * (t_date[1:]-t_date[:-1])\n",
    "\n",
    "    today = rewinddate\n",
    "    if daterange is None:\n",
    "        index = slice(None)\n",
    "    else:\n",
    "        index = np.logical_and(x >= daterange[0]-datetime.timedelta(days=1), x <= daterange[1]+datetime.timedelta(days=1))\n",
    "\n",
    "    fig,ax = plt.subplots()\n",
    "    for v in variables:\n",
    "        ax.fill_between(x[index], qq[0,index], qq[-1,index], color=variable_colors[v], alpha=0.25)\n",
    "        if show_50p:\n",
    "            ax.fill_between(x[index], qq[1,index], qq[-2,index], color=variable_colors[v], alpha=0.5)\n",
    "        if v == 'Rmort':\n",
    "            label = 'daily death rate'\n",
    "        elif v == 'hospitalized':\n",
    "            label = 'daily hospitalization rate'\n",
    "        else:\n",
    "            label = '{} rate (d$^{{-1}}$)'.format(variable_desc[v])\n",
    "\n",
    "        ax.plot(x[index], qq[2,index], color=variable_colors[v], label=label)\n",
    "\n",
    "    if daterange is None or (today >= daterange[0] and today <= daterange[1]):\n",
    "        ax.axvline(today.date(), color='black')\n",
    "\n",
    "    if daterange is not None:\n",
    "        ax.set_xlim(daterange)\n",
    "    add_xaxis_datestyle(ax)\n",
    "    if legendargs is None:\n",
    "        ax.legend()\n",
    "    else:\n",
    "        ax.legend(**legendargs)\n",
    "\n",
    "    ax.grid(True)\n",
    "    if title is not None:\n",
    "        ax.set_title(title)\n",
    "    return ax\n",
    "\n",
    "#for name in fits:\n",
    "#    ax = create_rate_quantile_plot(fits[name], variables=('Rmort',), title=desc[name], daterange=daterange_export)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## estimated fractions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "fate_colors = {'hospmild':'tab:blue', 'iculive':'tab:brown', 'mort':'tab:red'}\n",
    "fate_desc = {\n",
    "    'hospmild':'non-ICU hospital',\n",
    "    'iculive':'ICU (live)',\n",
    "    'mort':'death',\n",
    "}\n",
    "\n",
    "for name, fit in fits.items():\n",
    "    if 'frac_hospmild_i45' in fit.flatnames:\n",
    "        x = np.linspace(3, 97, 100)\n",
    "        n = fits[name]['frac_hospmild_i45'].size\n",
    "\n",
    "        fig, ax = plt.subplots()\n",
    "        for prop in ('hospmild', 'iculive', 'mort'):\n",
    "            inter = fits[name]['frac_{}_i45'.format(prop)]\n",
    "            slope = fits[name]['frac_{}_slope'.format(prop)]\n",
    "\n",
    "            tmp = np.empty((n, x.size))\n",
    "            for i in range(n):\n",
    "                tmp[i,:] = 10**((x-45.0)*slope[i] + inter[i])\n",
    "\n",
    "            qq = np.percentile(tmp, q=(5,25,50,75,95), axis=0)\n",
    "            fc = ax.fill_between(x, qq[0,:], qq[-1,:], alpha=0.25, color=fate_colors[prop])\n",
    "            ax.fill_between(x, qq[1,:], qq[-2,:], alpha=0.5, color=fc.get_facecolor()[0])\n",
    "            ax.plot(x, qq[2,:], label=fate_desc[prop], lw=2, color=fc.get_facecolor()[0][:-1])\n",
    "\n",
    "            x2 = np.array([3, 97])\n",
    "            slope = data['mu_frac_'+prop+'_slope']\n",
    "            inter = data['mu_frac_'+prop+'_i45']\n",
    "            ax.plot(x2, 10**(slope*(x2-45.0) + inter), color=fate_colors[prop], lw=5, ls=':', label='prior')\n",
    "\n",
    "        ax.legend()\n",
    "        ax.set_yscale('log')\n",
    "        ax.grid(True)\n",
    "        ax.set_xlabel('age', size=25)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## plots for website / reports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import matplotlib.patheffects as patheffects\n",
    "path_effects=(patheffects.withStroke(linewidth=3,foreground='white',alpha=0.75),)\n",
    "\n",
    "# temporarily overwrite\n",
    "daterange_export = (daterange_export[0], rewinddate + datetime.timedelta(days=28))\n",
    "daterange_plot = (daterange_export[0]-datetime.timedelta(days=1), daterange_export[-1]+datetime.timedelta(days=1))\n",
    "\n",
    "for name in quantiles:\n",
    "    if name in ('v12b', 'v12b_14+20', 'v12b_14+50', 'v12b_14+100'):\n",
    "        modelname = {'v12b':'SEIR_v12b',\n",
    "                     'v12b_14+20':'SEIR_v12b_scenario14+20',\n",
    "                     'v12b_14+50':'SEIR_v12b_scenario14+50',\n",
    "                     'v12b_14+100':'SEIR_v12b_scenario14+100'}[name]\n",
    "        title_prefix = {'v12b':'',\n",
    "                        'v12b_14+20':'scenario: $\\\\beta$+20% (14 days) ',\n",
    "                        'v12b_14+50':'scenario: $\\\\beta$+50% (14 days) ',\n",
    "                        'v12b_14+100':'scenario: $\\\\beta$+100% (14 days) ',\n",
    "                       }[name]\n",
    "\n",
    "        #\n",
    "        # hospitalizations\n",
    "        #\n",
    "\n",
    "        ax = create_quantile_plot(name, quantiles[name], variables=('hospitalized+sigma',), \n",
    "                                  title=None, show_50p=True, show_90p=True, show_data=True, \n",
    "                                  legendargs={'fontsize':20, 'loc':'upper right'})\n",
    "        ax.set_ylim([0,25])\n",
    "\n",
    "        fig = ax.get_figure()\n",
    "        fig.patch.set_alpha(0)\n",
    "        ax.set_ylabel('number hospitalized', size=20)\n",
    "        ax.set_title(title_prefix + 'Santa Cruz County COVID-19 Hospitalization Projections', size=26)\n",
    "        ax.set_xlim(left=t_date[7], right=daterange_export[1])\n",
    "        if ax.get_ylim()[1] < 7:\n",
    "            ax.set_ylim(top=7)\n",
    "\n",
    "        text_right = True\n",
    "        if text_right:\n",
    "            ax.text(0.99,0.01,'data source: CDPH, last updated: {}'.format(rewinddate.date()), transform=ax.transAxes, \n",
    "                    va='bottom', ha='right', size=14, color='tab:gray', path_effects=path_effects)\n",
    "        else:\n",
    "            ax.text(0.01,0.01,'data source: CDPH, last updated: {}'.format(rewinddate.date()), transform=ax.transAxes, \n",
    "                    va='bottom', ha='left', size=14, color='tab:gray', path_effects=path_effects)\n",
    "            #ax.text(0.99,0.95,'433 total acute beds (licensed and surge)', transform=ax.transAxes, va='top', ha='right', size=26, color=variable_colors['Hicu'], path_effects=path_effects)\n",
    "\n",
    "        #\n",
    "        # ICU\n",
    "        #\n",
    "\n",
    "        ax = create_quantile_plot(name, quantiles[name], variables=('Hicu+sigma',), daterange=daterange_plot,\n",
    "                                  title=None, show_50p=True, show_90p=True, show_data=True,\n",
    "                                  legendargs={'fontsize':20, 'loc':'upper right'})\n",
    "        #ax.set_ylim([0,15])\n",
    "\n",
    "        fig = ax.get_figure()\n",
    "        fig.patch.set_alpha(0)\n",
    "        ax.set_ylabel('number of ICU patients', size=20)\n",
    "        ax.set_title(title_prefix + 'Santa Cruz County COVID-19 ICU Hospitalization Projections', size=26)\n",
    "        ax.set_xlim(daterange_export)\n",
    "        if ax.get_ylim()[1] < 7:\n",
    "            ax.set_ylim(top=7)\n",
    "\n",
    "        if text_right:\n",
    "            ax.text(0.99,0.01,'data source: CDPH, last updated: {}'.format(rewinddate.date()), transform=ax.transAxes,\n",
    "                    va='bottom', ha='right', size=14, color='tab:gray', path_effects=path_effects)\n",
    "        else:\n",
    "            ax.text(0.01,0.01,'data source: CDPH, last updated: {}'.format(rewinddate.date()), transform=ax.transAxes,\n",
    "                    va='bottom', ha='left', size=14, color='tab:gray', path_effects=path_effects)\n",
    "\n",
    "        #\n",
    "        # Rt plots\n",
    "        #\n",
    "\n",
    "        def _insert_crossings(x, y, yval=1.0):\n",
    "            crossings = (y[1:] - yval) * (y[:-1] - yval) < 0\n",
    "\n",
    "            num_crossings = np.sum(crossings)\n",
    "            ind_crossings = np.flatnonzero(crossings) + 1\n",
    "\n",
    "            y_new = np.empty(y.size + num_crossings)\n",
    "            x_new = np.empty(x.size + num_crossings, dtype=x.dtype)\n",
    "\n",
    "            i0 = 0\n",
    "            i0_new = 0\n",
    "            for i, i1 in enumerate(ind_crossings):\n",
    "                i1_new = i1 + i\n",
    "\n",
    "                y_new[i0_new:i1_new] = y[i0:i1]\n",
    "                y_new[i1_new] = yval\n",
    "\n",
    "                alpha = (1.0 - y[i1-1]) / (y[i1] - y[i1-1])\n",
    "                x_new[i0_new:i1_new] = x[i0:i1]\n",
    "                #x_new[i1_new] = alpha * x[i1] + (1.0-alpha) * x[i1-1]\n",
    "                x_new[i1_new] = x[i1] + (1.0-alpha) * (x[i1-1] - x[i1])\n",
    "\n",
    "                i0 = i1\n",
    "                i0_new = i1_new + 1\n",
    "\n",
    "            y_new[i0_new:] = y[i0:]\n",
    "            x_new[i0_new:] = x[i0:]\n",
    "\n",
    "            return x_new, y_new\n",
    "\n",
    "        def plot_rt(ax, t_date, qq):\n",
    "            # bit of a hack to get inserts where either qq[0,:] or qq[-1,:] passes through zero\n",
    "            x, y0 = _insert_crossings(t_date, qq[0,:])\n",
    "            y1, _ = _insert_crossings(qq[-1,:], qq[0,:])\n",
    "\n",
    "            x, _ = _insert_crossings(x, y1)\n",
    "            y0, y1 = _insert_crossings(y0, y1)\n",
    "\n",
    "            ax.fill_between(x, np.maximum(y0,1), np.maximum(y1,1), color='darkorange', alpha=0.25)\n",
    "            ax.fill_between(x, np.minimum(y0,1), np.minimum(y1,1), color='darkgreen', alpha=0.25)\n",
    "\n",
    "            x, y = _insert_crossings(t_date, qq[2,:])\n",
    "\n",
    "            ax.plot(x, np.maximum(y,1), color='darkorange', label='$R_t$ estimate')\n",
    "            ax.plot(x, np.minimum(y,1), color='darkgreen')\n",
    "\n",
    "            ax.axhline(1.0, color='black')\n",
    "\n",
    "\n",
    "        fit = fits[name]\n",
    "        if 'Rt[1]' in fit.flatnames:\n",
    "            qq = np.percentile(fit['Rt'], axis=0, q=(5,25,50,75,95))\n",
    "\n",
    "            #\n",
    "            # Rt standard plot\n",
    "            #\n",
    "\n",
    "            fig, ax = plt.subplots()\n",
    "            fig.patch.set_alpha(0)\n",
    "            ax.set_title(title_prefix + 'Santa Cruz County COVID-19 Effective Reproductive Number $R_t$', size=26)\n",
    "            ax.set_ylabel('estimated value of $R_t$ with 90% credible interval', size=20)\n",
    "\n",
    "            plot_rt(ax, t_date, qq)\n",
    "\n",
    "            ax.grid(True)\n",
    "            add_xaxis_datestyle(ax)\n",
    "            #add_intervention_info(ax)\n",
    "\n",
    "            ax.axhline(1.0, color='black')\n",
    "            ax.set_ylim(bottom=0.0)\n",
    "            ax.set_xlim(left=t_date[14], right=rewinddate)\n",
    "\n",
    "            ax.text(0.99,0.01,'last updated: {}'.format(rewinddate.date()), transform=ax.transAxes,\n",
    "                    va='bottom', ha='right', size=14, color='tab:gray', path_effects=path_effects)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## model checks & diagnostics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name,fit in fits.items():\n",
    "    '''\n",
    "    if 'fractested[1]' in fit.flatnames:\n",
    "        qq = np.percentile(fit['fractested'], axis=0, q=(5,25,50,75,95))\n",
    "\n",
    "        fig,ax = plt.subplots()\n",
    "\n",
    "        ax.fill_between(tobs_date['I'], qq[0,:], qq[-1,:], color='tab:grey', alpha=0.25)\n",
    "        ax.fill_between(tobs_date['I'], qq[1,:], qq[-2,:], color='tab:grey', alpha=0.25)\n",
    "        ax.plot(tobs_date['I'], qq[2,:], color='tab:grey', label='estimated fraction of infectious that have been tested (cumulative number)')\n",
    "\n",
    "        ax.set_ylim(bottom=0.0)\n",
    "        if ax.get_ylim()[1] > 1.0:\n",
    "            ax.axhline(1.0, color='darkred', ls='--')\n",
    "        ax.set_title(desc[name])\n",
    "        add_xaxis_datestyle(ax)\n",
    "        ax.grid(True)\n",
    "        ax.legend()\n",
    "    '''\n",
    "    if 'fractest[1]' in fit.flatnames:\n",
    "        qq = np.percentile(fit['fractest'], axis=0, q=(5,25,50,75,95))\n",
    "\n",
    "        fig,ax = plt.subplots()\n",
    "\n",
    "        ax.fill_between(t_date, qq[0,:], qq[-1,:], color='tab:grey', alpha=0.25)\n",
    "        ax.fill_between(t_date, qq[1,:], qq[-2,:], color='tab:grey', alpha=0.25)\n",
    "        ax.plot(t_date, qq[2,:], color='tab:grey', label='estimated fraction of infectious that have been tested')\n",
    "\n",
    "        ax.set_ylim(bottom=0.0)\n",
    "        if ax.get_ylim()[1] > 1.0:\n",
    "            ax.axhline(1.0, color='darkred', ls='--')\n",
    "        ax.set_title(desc[name])\n",
    "        add_xaxis_datestyle(ax)\n",
    "        ax.grid(True)\n",
    "        ax.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nage = data['nage']\n",
    "\n",
    "if 'frac_hosp' in data:\n",
    "    frac_I_agg = np.zeros((len(fits)+1,nage,5))\n",
    "    # prior goes in first column\n",
    "    frac_I_agg[0,:,0] = data['frac_asym']\n",
    "    frac_I_agg[0,:,2] = data['frac_hosp'] - data['frac_icu']\n",
    "    frac_I_agg[0,:,3] = data['frac_icu'] - data['frac_mort']\n",
    "    frac_I_agg[0,:,4] = data['frac_mort']\n",
    "    frac_I_agg[0,:,1] = 1.0 - np.sum(frac_I_agg[0,:,:], axis=1)\n",
    "\n",
    "    for i,name in enumerate(fits):\n",
    "        frac_I_agg[i+1,:,:] = np.mean(fits[name]['frac_I'], axis=0)\n",
    "    names = ['prior',*list(fits.keys())]\n",
    "else:\n",
    "    frac_I_agg = np.zeros((len(fits),nage,5))\n",
    "    for i,name in enumerate(fits):\n",
    "        frac_I_agg[i,:,:] = np.mean(fits[name]['frac_I'], axis=0)\n",
    "    names = list(fits.keys())\n",
    "\n",
    "cmap = plt.get_cmap('RdYlGn_r')\n",
    "colors = [cmap(i/4) for i in range(5)]\n",
    "desc['prior'] = 'prior'\n",
    "labels = ['asymptomatic','mild','hospital (non-ICU)','ICU (non-fatal)', 'ICU (fatal)']\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.set_title('Fractionation of Infectious Individuals')\n",
    "bottom=np.zeros(5)\n",
    "for i,name in enumerate(names):\n",
    "    x = i + np.linspace(0, 1, nage+2)[:-1]\n",
    "    for iage in range(nage):\n",
    "        height = frac_I_agg[i,iage,:]\n",
    "        bottom[1:] = np.cumsum(height)[:-1]\n",
    "        ax.bar(x=x[iage], height=height, width=x[1]-x[0], bottom=bottom, color=colors)\n",
    "        for iy in range(len(height)):\n",
    "            if height[iy] > 0.1:\n",
    "                #ax.text(x[iage],bottom[iy]+0.5*height[iy], '{}\\n{:.1%}'.format(labels[iy],height[iy]), size=20, va='center', ha='center')\n",
    "                ax.text(x[iage],bottom[iy]+0.5*height[iy], '{:.1%}'.format(height[iy]), rotation=90, size=20, va='center', ha='center')\n",
    "\n",
    "ax.set(xticks=np.arange(len(names))+0.25, xticklabels=[desc[name] for name in names])\n",
    "ax.yaxis.grid(True)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.set_title('Fractionation of Infectious Individuals (hospitalizations only)')\n",
    "bottom=np.zeros(3)\n",
    "for i,name in enumerate(names):\n",
    "    x = i + np.linspace(0, 1, nage+2)[:-1]\n",
    "    for iage in range(frac_I_agg.shape[1]):\n",
    "        height = frac_I_agg[i,iage,2:]\n",
    "        bottom[1:] = np.cumsum(height)[:-1]\n",
    "        ax.bar(x=x[iage], height=height, width=x[1]-x[0], bottom=bottom, color=colors[2:])\n",
    "        for iy in range(len(height)):\n",
    "            if height[iy] > 0.01:\n",
    "                #ax.text(x[iage],bottom[iy]+0.5*height[iy], '{}\\n{:.1%}'.format(labels[iy],height[iy]), size=20, va='center', ha='center')\n",
    "                ax.text(x[iage],bottom[iy]+0.5*height[iy], '{:.1%}'.format(height[iy]), rotation=90, size=20, va='center', ha='center')\n",
    "\n",
    "ax.set(xticks=np.arange(len(names))+0.25, xticklabels=[desc[name] for name in names])\n",
    "ax.yaxis.grid(True)\n",
    "\n",
    "None"
   ]
  },
  {
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}