{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Problem set 2, due date January 26th"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from scipy.integrate import odeint\n",
    "\n",
    "plt.style.use('custom.mplstyle')\n",
    "%config InlineBackend.figure_format = 'retina'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Susceptible-Infected-Recovered (SIR) models\n",
    "___\n",
    "\n",
    "[SIR models](https://en.wikipedia.org/wiki/Compartmental_models_in_epidemiology#The_SIR_model_without_vital_dynamics) describe the spread of a pathogen in a population. The SIR model breaks down the overall population into three subsets: \n",
    "* the susceptible population $S$, which consists of folks who haven't been infected yet and can be infected\n",
    "* the infected population $I$, which can recover from the infection or infect the susceptible population\n",
    "* the recovered population $R$, which assumedly can't be reinfected. Note that sometimes this group is referred to as \"removed,\" which could also include death by infection.\n",
    "\n",
    "\n",
    "![SIR](SIR.jpg)\n",
    "\n",
    "\n",
    "The system of differential equations for this model are:\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "\\frac{dS(t)}{dt} &= -\\beta S(t) I(t) \\\\\n",
    "\\frac{dI(t)}{dt} &= \\beta S(t) I(t) - \\nu I(t) \\\\\n",
    "\\frac{dR(t)}{dt} &= \\nu I(t)\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "where $\\beta > 0$ is the rate at which a susceptible person meeting an infected individual becomes infected. Note that such encounters occur at rate $S(t) I(t)$. $\\nu > 0$ is the rate at which infectious people recover. Here, we are letting $S$, $I$, and $R$ be the **fractions of the population** which are susceptible, infectious, or recovered, respectively: $S(t) + I(t) + R(t) = 1$. Because of this relationship, we consider the equation for $dR(t)/dt$ to be superfluous, as at each time point: $R(t) = 1- S(t) -I(t)$. Additionally, by modeling the fractions of the overall population, we don't need to worry about what the size of the total population size, making our lives a bit easier. \n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The ratio of $\\beta$ and $\\nu$ constitutes what is called $R_0$, the basic reproduction number:\n",
    "\n",
    "$$\n",
    "R_0 = \\frac{\\beta}{\\nu}\n",
    "$$\n",
    "\n",
    "$R_0$ describes how many people an infectious person infects on average; you probably have heard about $R_0$ in the context of SARS-CoV-2 infections. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Question 1** (10 points): What are the units of $\\beta$, $\\nu$, and $R_0$?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**your answer:**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## $R_0$ of different pathogens\n",
    "\n",
    "For a given pathogen, the exact value of $R_0$ can change by changing the encounter rate within a popualtion. I.e., if we limit encounters, the basic reproduction number goes down. That is why lockdowns have been used to reduce infections.\n",
    "\n",
    "Different pathogens can also have different intrinsic pathogenicity and transmissibilities. For example, the Omicron variant of SARS-CoV-2 is highly transmissable (i.e., a relatively large $R_0$). \n",
    "\n",
    "Estimating $R_0$ is not straightforward and can be sensitive to data sampling.  Here are some of the $R_0$ estimates for various pathogens:\n",
    "\n",
    "\n",
    "* measles: $R_0\\simeq 15$\n",
    "* influenza 1918 pandemic: $R_0 \\simeq 2-3$\n",
    "* influenza H1N1 pandemic (2009 pandemic): $R_0 \\simeq 1.5 $\n",
    "* Ebola outbreak, West Africa (2014): $R_0 \\simeq 1.5- 2.5$\n",
    "* SARS-CoV-2 (original strain 2019-2020): $R_0 \\simeq 1.8- 3.6$\n",
    "* SARS-CoV-2 ($\\Delta$ variant): $R_0 \\simeq 6$\n",
    "* SARS-CoV-2 (Omicron variant): $R_0 \\simeq 8$ (still to be determined)\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Limiting cases of the SIR model:\n",
    "\n",
    "Note that $dS(t)/dt$ and $dI(t)/dt$ are coupled in a nonlinear way due to the presence of $\\beta S(t) I(t)$. Systems of differential equations coupled nonlinearly rarely have analytical solutions. We can play the game of looking at how the populations evolve in limiting cases which make the solution tractable.\n",
    "\n",
    "### **Question 2** (10 points): Solve the dynamics of the infectious population $I(t)$ at the beginning of an infectious spread, i.e., by setting $S\\approx 1$.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**your answer:**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Question 3** (10 points): Argue that at large times (stationary state), the fraction of the infected population $I(t \\to \\infty)$ should go to zero. In this regime, solve dynamics of the susceptible population $S(t)$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**your answer:**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Euler method and the SIR model\n",
    "\n",
    "The solutions from Questions 2,3 are accurate only in the limiting regimes. Here, we'll solve the equations irrespective of any regime using numerical integration.\n",
    "___\n",
    "\n",
    "Because $R(t) = 1 - S(t) - I(t)$, we need only solve for the evolution of $S(t)$ and $I(t)$. As before (see tutorial 2),\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "S(t + \\Delta t) &= S(t) + \\frac{dS(t)}{dt} \\Delta t + O(\\Delta t^2) \\\\\n",
    "I(t + \\Delta t) &= I(t) + \\frac{dI(t)}{dt} \\Delta t + O(\\Delta t^2)\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "Applying the approximation and plugging in the values for the derivatives, we have\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "S(t + \\Delta t) &\\approx S(t) - \\beta S(t) I(t) \\Delta t \\\\\n",
    "I(t + \\Delta t) &\\approx I(t) + \\beta S(t) I(t) \\Delta t - \\nu I(t) \\Delta t\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "\n",
    "### **Question 4** (20 points): Write a function which takes in $S(0)$, $I(0)$, $R_0$, $\\nu$, $\\Delta t$, and $T$ (the total time over which we want to track the evolution of the populations), and outputs $\\vec{t} = \\{0, \\Delta t, 2 \\Delta t, \\ldots \\}$, $\\vec{S} = \\{S(0), S(\\Delta t), S(2 \\Delta t), \\ldots\\}$, and $\\vec{I} = \\{I(0), I(\\Delta t), I(2 \\Delta t), \\ldots\\}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def euler_sir(s0, i0, r0, nu, T, dt):\n",
    "    # Your code here.\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Question 5** (5 points): Run your function, obtaining trajectories for $S$ and $I$ using the following parameters:\n",
    "- $I(0) = 10^{-6}$\n",
    "- $S(0) = 1 - I(0)$\n",
    "- $\\nu = \\frac{1}{7}$ (recovery after a week)\n",
    "- a range of $R_0 \\in [0.1, 0.5, 1, 2, 3, 5, 10]$ (average number of infections per individual)\n",
    "- $T = 100$ days\n",
    "- $\\Delta t = 0.01$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Question 6** (10 points): Plot $S(t)$, $I(t)$, and $R(t)$ vs. time as solid lines. On the same plot, also plot the analytical solutions for the two limiting cases in Questions 2, 3 as dashed lines. Display a legend with each line's label. Label the axes appropriately."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Better numerical integration and the SIR model\n",
    "___\n",
    "\n",
    "We stated above the Euler's method employs a first-order Taylor series approximation to solve differential equations. The SIR model is seemingly simple, but it is much more complicated than solving the differential equation for exponential growth. Let's now use `scipy` which has functions to integrate differential equations more accurately. Specifically let's use [scipy.integrate.odeint](https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.odeint.html). (We don't concern ourselves with the implementation here. Though for those who are curious, this function uses [multistep methods](https://en.wikipedia.org/wiki/Linear_multistep_method#Families_of_multistep_methods) whereas Euler's method uses only one step. So this `scipy` function should be more accurate.)\n",
    "\n",
    "Reading the documentation, we see an example that is illuminating and follow it. We need to define a function which outputs an array-like object which contains the derivatives.\n",
    "\n",
    "### **Question 7:** (20 points) Write a function which computes $dS(t)/dt$ and $dI(t)/dt$ taking as input $[S, I]$, the current values of $S$ and $I$; $T$, a sequence of time points; and $\\beta$ and $\\nu$. This function should return $[dS/dt, dI/dt]$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sir_ode(current_vals, T, beta, nu):\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Question 8:** (5 points) Run the `scipy` ode integrator using your function, obtaining trajectories for $S$ and $I$ using the following parameters:\n",
    "- $I(0) = 10^{-6}$\n",
    "- $S(0) = 1 - I(0)$\n",
    "- $\\nu = \\frac{1}{7}$ (recovery after a week)\n",
    "- a range of $R_0 \\in [0.1, 0.5, 1, 2, 3, 5, 10]$ (average number of infections per individual)\n",
    "- $T$ = `np.linspace(0, 100, 101)` (specifies that S(t), I(t) is evaluated at $t \\in [0, 100]$ (days) $ \\subset \\mathbb{Z}$)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Question 9:** (10 points) Plot the solutions from your implementation of the Euler method as a scatter, and plot the solutions from `scipy` as solid lines. When plotting the results from using the Euler method, plot every 100th point since we had $\\Delta t = 0.01$. Set the size of the markers of the scatter to be 5: `s=5`. Label all plotted objects. Because we now have six plotted objects, the legend is large. We move the legend outside the axes: `plt.legend(loc='center right', bbox_to_anchor=(1.3, 0.5))`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that, at least for this choice of parameters, Euler's method produces an approximation which matches the results given by the `scipy` integrator (at least visually). Let's continue to understand the trajectories of $S(t)$ and $I(t)$ using the function from `scipy` since it's much quicker and more accurate in principle."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How does $R_0$ determine $S(\\infty)$, the fraction of individuals who remain susceptible at the end of an outbreak (i.e., were not infected during the outbreak)?\n",
    "___\n",
    "\n",
    "\n",
    "Let's first find the relation between $R_0$, $I(t)$ and $S(t)$. We compute $dI/dS$ to see how $R_0$ appears in our model naturally.\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "\\frac{dI(t)}{dS(t)} &=  \\frac{dI(t)}{dt} \\left( \\frac{dS(t)}{dt} \\right)^{-1} \\\\\n",
    "&= \\frac{\\beta S(t) I(t) - \\nu I(t)}{-\\beta S(t) I(t)} \\\\\n",
    "&= -1 + \\frac{1}{R_0 S(t)}\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "This is a simple ordinary differential equation which we can solve using separation of variables.\n",
    "\n",
    "$$\n",
    "I(t) = -S(t) + \\frac{\\ln S(t)}{R_0} + C \\quad \\forall t\n",
    "$$\n",
    "\n",
    "where $C$ is some constant of integration found from initial conditions.\n",
    "\n",
    "### **Question 10:** (5 points)  Determine $C$ based on $R_0$ and the initial conditions $I(0)$, $S(0)$. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**your answer:**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Question 11:** (10 points) As we discussed in Question 3, at large times (stationary state), the fraction of the infected population $I(t \\to \\infty)$ should go to zero. By using this assumption, find a relationship between the fraction of individuals who remain susceptible at the end of an outbreak $S(t\\to\\infty)$, the basic reproduction number $R_0$, and the initial conditions $S(0)$, $I(0)$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**your answer:**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Essentially, this shows that we can determine the final state of an outbreak based on   $R_0$ and the initial conditions.\n",
    "\n",
    "### **Question 12:** (20 points) Write a numerical function which takes in an $R_0$, $S(0)$, $I(0)$, and outputs the fraction of the population which remains susceptible. To numerically solve for $S(\\infty)$ you sould use the technique introduced for solving a *transcendental equation* in the lecture for tutorial 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def stationary_s(s0, i0, r0):\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Question 13:** (5 points) Use you function from Question 12 to find $S(\\infty)$ for different $R_0$ values: \n",
    "- $I(0) = 10^{-6}$\n",
    "- $S(0) = 0.2$ \n",
    "- A range of $R_0$ = `np.logspace(-1, 1, 10)`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Question 15:** (10 points) Plot $S(\\infty)$ vs. $R_0$ on a logarithmic x-axis with a range of  $R_0$ = `np.logspace(-1, 3, 10)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Herd immunity is the idea that a disease cannot grow into an epidemic when many people in a population are immune and not susceptible. The condition for herd immunity is obtained when \n",
    "\n",
    "$$\n",
    "\\frac{dI(0)}{dt} < 0\n",
    "$$\n",
    "\n",
    "Since\n",
    "\n",
    "$$\n",
    "\\frac{dI(0)}{dt} = (\\beta S(0) - \\nu) I(0)\n",
    "$$\n",
    "\n",
    "we see that\n",
    "\n",
    "$$\n",
    "S(0) < 1/R_0\n",
    "$$\n",
    "\n",
    "is the condition for herd immunity.\n",
    "\n",
    "### **Question 16:** (10 points) How does this herd immunity threshold compared to $S(\\infty)$ (the remaining number of susceptibles)?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**your answer:**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Bonus Question:** (50 points, *due date: Feb 2nd*)     \n",
    "* Compare and plot the features of outbreaks ($S(t),\\, I(t), \\, S(\\infty)$) associated with the different pathogens listed at the beginning of this notebook. \n",
    "* How do the conditions for herd immunity compare among these pathogens?\n",
    "* Why is estimating $R_0$ nontrivial?\n",
    "* Why is the SIR model too simplistic in characterizing a pandemic? What factors are not taken into account? How can one improve upon this model?\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "biophys2022",
   "language": "python",
   "name": "biophys2022"
  },
  "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": 4
}