Tutorial 4 - Building an SEI3R Model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The SEI3R Model\n",
    "\n",
    "Another extension of SIR, designed during the paper we read for this work, https://arxiv.org/pdf/2003.13221.pdf, Planning as Inference in Epedimiological Control Models. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model Dynamics\n",
    "\n",
    "\n",
    "A member of the susceptible population $S$ moves to the exposed $E$ after being exposed to an infections person, where \"exposure\" is defined as the previously susceptible person contracting the illness. \n",
    "\n",
    "After an incubation period, a random duration parameterized by $\\alpha$, the individual develops a mild infection $I_1$, from here, the indivdual either moves to recover $R$, or progress to a severe infection $I_2$. At $I_2$ they again either recover, or progress to critical infection $I_3$. From $I_3$, the critically ill individual will either die or recover. \n",
    "\n",
    "The summed population of all compartments at time $t$ is $N_t$. For the purposes of are work, we keep this value fixed. \n",
    "\n",
    "The dynamics of this system are defined by the following equations:\n",
    "\n",
    "\n",
    "$$\\frac{d}{dt}S_t = - (1 - u) \\frac{1}{N_t} \\sum_{i=1}^3 \\beta_i I_{i,t}S_t$$\n",
    "\n",
    "$$\\frac{d}{dt}E = (1 - u) \\frac{1}{N_t} \\sum_{i=1}^3 \\beta_i I_{i,t}S_t - \\alpha E_t $$\n",
    "\n",
    "$$\\frac{d}{dt} I_{1,t} = \\alpha E_t -p_1 I_{1,t} - \\gamma_1 I_{1,t}$$\n",
    "\n",
    "\n",
    "$$\\frac{d}{dt} I_{2,t} =  p_1 I_{1,t} -p_2 I_{2,t} - \\gamma_2 I_{2,t}$$\n",
    "\n",
    "\n",
    "$$\\frac{d}{dt} I_{3,t} =  p_2 I_{2,t} -p_3 I_{3,t} - \\gamma_3 I_{3,t}$$\n",
    "\n",
    "\n",
    "$$\\frac{d}{dt} R_{t} = \\gamma_1 I_{1,t} + \\gamma_2 I_{2,t} +  \\gamma_3 I_{3,t}$$\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For simulations with this model, we initialize with $0.01 \\%$ of the population as exposed, and the remaining $99.99 \\%$ as susceptible. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import torch\n",
    "import pyro\n",
    "import pyro.distributions as dist\n",
    "from pyro.contrib.epidemiology import CompartmentalModel, binomial_dist, infection_dist\n",
    "\n",
    "%matplotlib inline\n",
    "#assert pyro.__version__.startswith('1.5.0') # I have 1.5.1, hopefully not a problem to comment this out. \n",
    "torch.set_default_dtype(torch.double)  # Required for MCMC inference.\n",
    "pyro.enable_validation(True)  # Always a good idea.\n",
    "smoke_test = ('CI' in os.environ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# check https://github.com/SwapneelM/covid/blob/master/SEIR/seir.py for help. \n",
    "class SimpleSEI3RModel(CompartmentalModel):\n",
    "    \"\"\"\n",
    "    Susceptible-Exposed-Infected1-Infected2-Infected3-Recovered-(Dead) model.\n",
    "\n",
    "    This is a stochastic discrete-time discrete-state model with seven\n",
    "    compartments: \"S\" for susceptible, \"E\" for exposed, \"I\" for infected (with 3 degrees of illness), \"D\"\n",
    "    for deceased individuals, and \"R\" for recovered individuals (the recovered\n",
    "    individuals are implicit: ``R = population - S - E - I1 - I2 - I3 - D``) with\n",
    "    transitions ``S -> E -> I1 -> R``,  ``S -> E -> I1 -> I2 -> R``,  ``S -> E -> I1 -> -> I2 -> I3 -> R``\n",
    "    and ``I3 -> D``.\n",
    "\n",
    "    Because the transitions are not simple linear succession, this model\n",
    "    implements a custom :meth:`compute_flows()` method.\n",
    "\n",
    "    :param int population: Total ``population = S + E + I3 + R + D``.\n",
    "    :param float incubation_time: Mean incubation time (duration in state\n",
    "        ``E``). Must be greater than 1.\n",
    "    :param float recovery_time: Mean recovery time (duration in state\n",
    "        ``I``). Must be greater than 1.\n",
    "    :param float mortality_rate: Portion of infections resulting in death.\n",
    "        Must be in the open interval ``(0, 1)``.\n",
    "    :param iterable data: Time series of new observed infections. Each time\n",
    "        step is Binomial distributed between 0 and the number of ``S -> E``\n",
    "        transitions. This allows false negative but no false positives.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, population, incubation_time, recovery_time,\n",
    "                 mortality_rate, data):\n",
    "        compartments = (\"S\", \"E\", \"I\", \"I2\", \"I3\", \"R\")  # D is implicit.\n",
    "        duration = len(data)\n",
    "        super().__init__(compartments, duration, population)\n",
    "\n",
    "        assert isinstance(incubation_time, float)\n",
    "        assert incubation_time > 1\n",
    "        self.incubation_time = incubation_time\n",
    "\n",
    "        assert isinstance(recovery_time, float)\n",
    "        assert recovery_time > 1\n",
    "        self.recovery_time = recovery_time\n",
    "\n",
    "        assert isinstance(mortality_rate, float)\n",
    "        assert 0 < mortality_rate < 1\n",
    "        self.mortality_rate = mortality_rate\n",
    "\n",
    "        self.data = data\n",
    "\n",
    "    def global_model(self):\n",
    "        tau_e = self.incubation_time\n",
    "        tau_i = self.recovery_time\n",
    "        mu = self.mortality_rate\n",
    "        R0 = pyro.sample(\"R0\", dist.LogNormal(0., 1.))\n",
    "        rho = pyro.sample(\"rho\", dist.Beta(10, 10))\n",
    "        return R0, tau_e, tau_i, mu, rho\n",
    "\n",
    "    def initialize(self, params):\n",
    "        # Start with a single infection ( at stage 1) .\n",
    "        # pretty ambiguous. Check with paper see if they do something different. \n",
    "        return {\"S\": self.population - 1, \"E\": 0, \"I\": 1, \"I2\":0, \"I3\":0, \"R\": 0}\n",
    "\n",
    "    def transition(self, params, state, t):\n",
    "        \"\"\"\n",
    "        How we go from one compartment to another. After definining changes we update compartments. \n",
    "        \n",
    "        \"\"\"\n",
    "        R0, tau_e, tau_i, mu, rho = params # add the other ones from paper\n",
    "        p1 = .5\n",
    "        p2 = .5\n",
    "        gamma1 = .2\n",
    "        gamma2 = .3\n",
    "        gamma3 = .5\n",
    "        kappa = 0.005\n",
    "\n",
    "        # Sample flows between compartments.\n",
    "        S2E = pyro.sample(\"S2E_{}\".format(t),\n",
    "                          infection_dist(individual_rate=R0 / tau_i,\n",
    "                                         num_susceptible=state[\"S\"],\n",
    "                                         num_infectious=state[\"I\"],\n",
    "                                         population=self.population))\n",
    "        \n",
    "        E2I = pyro.sample(\"E2I_{}\".format(t),\n",
    "                          binomial_dist(state[\"E\"], 1 / tau_e))\n",
    "        \n",
    "        I2I2 = pyro.sample(\"I2I2_{}\".format(t),\n",
    "                          binomial_dist(state[\"I\"], p1))\n",
    "        \n",
    "        I2R = pyro.sample(\"I2R_{}\".format(t),\n",
    "                          binomial_dist(state[\"I\"], gamma1))\n",
    "       \n",
    "        I22I3 = pyro.sample(\"I22I3_{}\".format(t),\n",
    "                          binomial_dist(state[\"I2\"], p2))\n",
    "        \n",
    "        I22R = pyro.sample(\"I22R_{}\".format(t),\n",
    "                          binomial_dist(state[\"I2\"], gamma2))\n",
    "        \n",
    "        \n",
    "        I32R = pyro.sample(\"I32R_{}\".format(t),\n",
    "                          binomial_dist(state[\"I3\"], gamma3))\n",
    "        \n",
    "        I32D = pyro.sample(\"I32D_{}\".format(t),\n",
    "                          binomial_dist(state[\"I3\"], kappa))\n",
    "        \n",
    "        \n",
    "        # Of the 1/tau_i expected recoveries-or-deaths, a portion mu die and\n",
    "        # the remaining recover. Alternatively we could model this with a\n",
    "        # Multinomial distribution I2_ and extract the two components I2D and\n",
    "        # I2R, however the Multinomial distribution does not currently\n",
    "        # implement overdispersion or moment matching.\n",
    "        \n",
    "\n",
    "        # update compartments. \n",
    "        state[\"S\"] = state[\"S\"] - S2E\n",
    "        state[\"E\"] = state[\"E\"] + S2E - E2I\n",
    "        state[\"I\"] = state[\"I\"] + E2I - I2R - I2I2\n",
    "        state[\"I2\"] = state[\"I2\"] + I2I2 - I22R - I22I3\n",
    "        state[\"I3\"] = state[\"I3\"] + I22I3 - I32R - I32D\n",
    "        state[\"R\"] = state[\"R\"] + I32R + I2R + I22R\n",
    "        \n",
    "        \n",
    "        \n",
    "        \n",
    "\n",
    "        # Condition on observations. \n",
    "        t_is_observed = isinstance(t, slice) or t < self.duration\n",
    "        pyro.sample(\"obs_{}\".format(t),\n",
    "                    binomial_dist(S2E, rho),\n",
    "                    obs=self.data[t] if t_is_observed else None)\n",
    "\n",
    "    def compute_flows(self, prev, curr, t):\n",
    "        \"\"\"\n",
    "        Having a tough time making these line up. Need to choose one to define first? \n",
    "        \n",
    "        Should a new variable be implicit? \n",
    "        \"\"\"\n",
    "        \n",
    "        \n",
    "        # \n",
    "        \n",
    "        # First part, flow through first parts\n",
    "        S2E = prev[\"S\"] - curr[\"S\"]  # S can only go to E.\n",
    "        E2I = prev[\"E\"] - curr[\"E\"] + S2E\n",
    "        \n",
    "        \n",
    "        # Secondly, flow through I leaving to go to I2 or R\n",
    "        \n",
    "        I2 = prev[\"I\"] - curr[\"I\"] + E2I \n",
    "        \n",
    "        \n",
    "        \n",
    "        # Next is I2I2, I22R\n",
    "        \n",
    "        # Here is where I get stuck. I cannot define either of these without specifying one first. \n",
    "        I2R = prev[\"I\"] - curr[\"I\"] + E2I - I2I2\n",
    "\n",
    "        I2I2 = prev[\"I\"] - curr[\"I\"] + E2I - I2R\n",
    "\n",
    "        \n",
    "\n",
    "        I32D = curr[\"D\"] - prev[\"D\"]  # D can only have come from I3. \n",
    "        II32R = prev[\"I3\"] - curr[\"I3\"] + I22I3 - I32D\n",
    "\n",
    "        #again, stuck I22I3 is also defined weirdly\n",
    "        \n",
    "        \n",
    "        # flow is change in I2 plus incoming minus recovered subset. \n",
    "        I22I3 = prev[\"I2\"] - curr[\"I2\"] + I2I2 - I22R \n",
    "        \n",
    "        I22R = prev[\"I2\"] - curr[\"I2\"] + I2I2 - I22I3\n",
    "        \n",
    "        \n",
    "\n",
    "        \n",
    "        \n",
    "        return {\n",
    "            \"S2E_{}7\".format(t): S2E,\n",
    "            \"E2I_{}\".format(t): E2I,\n",
    "            \"I2I2_{}\".format(t): I2I2,\n",
    "            \"I22I3_{}\".format(t): I22I3,\n",
    "\n",
    "            \"I32D_{}\".format(t): I32D,\n",
    "            \"I2R_{}\".format(t): I2R,\n",
    "            \"I22R_{}\".format(t): I22R,\n",
    "            \"I32R_{}\".format(t): I32R,\n",
    "\n",
    "        }\n",
    "\n"
   ]
  }
 ],
 "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}