{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Modeling and Simulation in Python\n",
    "\n",
    "Chapter 12\n",
    "\n",
    "Copyright 2017 Allen Downey\n",
    "\n",
    "License: [Creative Commons Attribution 4.0 International](https://creativecommons.org/licenses/by/4.0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure Jupyter so figures appear in the notebook\n",
    "%matplotlib inline\n",
    "\n",
    "# Configure Jupyter to display the assigned value after an assignment\n",
    "%config InteractiveShell.ast_node_interactivity='last_expr_or_assign'\n",
    "\n",
    "# import functions from the modsim.py module\n",
    "from modsim import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Code\n",
    "\n",
    "Here's the code from the previous notebook that we'll need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_system(beta, gamma):\n",
    "    \"\"\"Make a system object for the SIR model.\n",
    "    \n",
    "    beta: contact rate in days\n",
    "    gamma: recovery rate in days\n",
    "    \n",
    "    returns: System object\n",
    "    \"\"\"\n",
    "    init = State(S=89, I=1, R=0)\n",
    "    init /= sum(init)\n",
    "\n",
    "    t0 = 0\n",
    "    t_end = 7 * 14\n",
    "\n",
    "    return System(init=init, t0=t0, t_end=t_end,\n",
    "                  beta=beta, gamma=gamma)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_func(state, t, system):\n",
    "    \"\"\"Update the SIR model.\n",
    "    \n",
    "    state: State with variables S, I, R\n",
    "    t: time step\n",
    "    system: System with beta and gamma\n",
    "    \n",
    "    returns: State object\n",
    "    \"\"\"\n",
    "    s, i, r = state\n",
    "\n",
    "    infected = system.beta * i * s    \n",
    "    recovered = system.gamma * i\n",
    "    \n",
    "    s -= infected\n",
    "    i += infected - recovered\n",
    "    r += recovered\n",
    "    \n",
    "    return State(S=s, I=i, R=r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_simulation(system, update_func):\n",
    "    \"\"\"Runs a simulation of the system.\n",
    "        \n",
    "    system: System object\n",
    "    update_func: function that updates state\n",
    "    \n",
    "    returns: TimeFrame\n",
    "    \"\"\"\n",
    "    frame = TimeFrame(columns=system.init.index)\n",
    "    frame.row[system.t0] = system.init\n",
    "    \n",
    "    for t in linrange(system.t0, system.t_end):\n",
    "        frame.row[t+1] = update_func(frame.row[t], t, system)\n",
    "    \n",
    "    return frame"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Given the results, we can compute metrics that quantify whatever we are interested in, like the total number of sick students, for example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_total_infected(results):\n",
    "    \"\"\"Fraction of population infected during the simulation.\n",
    "    \n",
    "    results: DataFrame with columns S, I, R\n",
    "    \n",
    "    returns: fraction of population\n",
    "    \"\"\"\n",
    "    return get_first_value(results.S) - get_last_value(results.S)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's an example.|"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta = 0.333\n",
    "gamma = 0.25\n",
    "system = make_system(beta, gamma)\n",
    "\n",
    "results = run_simulation(system, update_func)\n",
    "print(beta, gamma, calc_total_infected(results))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Write functions that take a `TimeFrame` object as a parameter and compute the other metrics mentioned in the book:\n",
    "\n",
    "1.  The fraction of students who are sick at the peak of the outbreak.\n",
    "\n",
    "2.  The day the outbreak peaks.\n",
    "\n",
    "3.  The fraction of students who are sick at the end of the semester.\n",
    "\n",
    "Note: Not all of these functions require the `System` object, but when you write a set of related functons, it is often convenient if they all take the same parameters.\n",
    "\n",
    "Hint:  If you have a `TimeSeries` called `I`, you can compute the largest value of the series like this:\n",
    "\n",
    "    I.max()\n",
    "\n",
    "And the index of the largest value like this:\n",
    "\n",
    "    I.idxmax()\n",
    "\n",
    "You can read about these functions in the `Series` [documentation](https://pandas.pydata.org/pandas-docs/stable/generated/pandas.Series.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### What if?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use this model to evaluate \"what if\" scenarios.  For example, this function models the effect of immunization by moving some fraction of the population from S to R before the simulation starts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_immunization(system, fraction):\n",
    "    \"\"\"Immunize a fraction of the population.\n",
    "    \n",
    "    Moves the given fraction from S to R.\n",
    "    \n",
    "    system: System object\n",
    "    fraction: number from 0 to 1\n",
    "    \"\"\"\n",
    "    system.init.S -= fraction\n",
    "    system.init.R += fraction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start again with the system we used in the previous sections."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "tc = 3      # time between contacts in days \n",
    "tr = 4      # recovery time in days\n",
    "\n",
    "beta = 1 / tc      # contact rate in per day\n",
    "gamma = 1 / tr     # recovery rate in per day\n",
    "\n",
    "system = make_system(beta, gamma)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And run the model without immunization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = run_simulation(system, update_func)\n",
    "calc_total_infected(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now with 10% immunization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "system2 = make_system(beta, gamma)\n",
    "add_immunization(system2, 0.1)\n",
    "results2 = run_simulation(system2, update_func)\n",
    "calc_total_infected(results2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "10% immunization leads to a drop in infections of 16 percentage points.\n",
    "\n",
    "Here's what the time series looks like for S, with and without immunization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot(results.S, '-', label='No immunization')\n",
    "plot(results2.S, '--', label='10% immunization')\n",
    "\n",
    "decorate(xlabel='Time (days)',\n",
    "         ylabel='Fraction susceptible')\n",
    "\n",
    "savefig('figs/chap05-fig02.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can sweep through a range of values for the fraction of the population who are immunized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "immunize_array = linspace(0, 1, 11)\n",
    "for fraction in immunize_array:\n",
    "    system = make_system(beta, gamma)\n",
    "    add_immunization(system, fraction)\n",
    "    results = run_simulation(system, update_func)\n",
    "    print(fraction, calc_total_infected(results))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This function does the same thing and stores the results in a `Sweep` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sweep_immunity(immunize_array):\n",
    "    \"\"\"Sweeps a range of values for immunity.\n",
    "    \n",
    "    immunize_array: array of fraction immunized\n",
    "    \n",
    "    returns: Sweep object\n",
    "    \"\"\"\n",
    "    sweep = SweepSeries()\n",
    "    \n",
    "    for fraction in immunize_array:\n",
    "        system = make_system(beta, gamma)\n",
    "        add_immunization(system, fraction)\n",
    "        results = run_simulation(system, update_func)\n",
    "        sweep[fraction] = calc_total_infected(results)\n",
    "        \n",
    "    return sweep"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's how we run it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "immunize_array = linspace(0, 1, 21)\n",
    "infected_sweep = sweep_immunity(immunize_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here's what the results look like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot(infected_sweep)\n",
    "\n",
    "decorate(xlabel='Fraction immunized',\n",
    "         ylabel='Total fraction infected',\n",
    "         title='Fraction infected vs. immunization rate',\n",
    "         legend=False)\n",
    "\n",
    "savefig('figs/chap05-fig03.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If 40% of the population is immunized, less than 4% of the population gets sick."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Logistic function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To model the effect of a hand-washing campaign, I'll use a [generalized logistic function](https://en.wikipedia.org/wiki/Generalised_logistic_function) (GLF), which is a convenient function for modeling curves that have a generally sigmoid shape.  The parameters of the GLF correspond to various features of the curve in a way that makes it easy to find a function that has the shape you want, based on data or background information about the scenario."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def logistic(x, A=0, B=1, C=1, M=0, K=1, Q=1, nu=1):\n",
    "    \"\"\"Computes the generalize logistic function.\n",
    "    \n",
    "    A: controls the lower bound\n",
    "    B: controls the steepness of the transition \n",
    "    C: not all that useful, AFAIK\n",
    "    M: controls the location of the transition\n",
    "    K: controls the upper bound\n",
    "    Q: shift the transition left or right\n",
    "    nu: affects the symmetry of the transition\n",
    "    \n",
    "    returns: float or array\n",
    "    \"\"\"\n",
    "    exponent = -B * (x - M)\n",
    "    denom = C + Q * exp(exponent)\n",
    "    return A + (K-A) / denom ** (1/nu)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following array represents the range of possible spending."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "spending = linspace(0, 1200, 21)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`compute_factor` computes the reduction in `beta` for a given level of campaign spending.\n",
    "\n",
    "`M` is chosen so the transition happens around \\$500.\n",
    "\n",
    "`K` is the maximum reduction in `beta`, 20%.\n",
    "\n",
    "`B` is chosen by trial and error to yield a curve that seems feasible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_factor(spending):\n",
    "    \"\"\"Reduction factor as a function of spending.\n",
    "    \n",
    "    spending: dollars from 0 to 1200\n",
    "    \n",
    "    returns: fractional reduction in beta\n",
    "    \"\"\"\n",
    "    return logistic(spending, M=500, K=0.2, B=0.01)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's what it looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "percent_reduction = compute_factor(spending) * 100\n",
    "\n",
    "plot(spending, percent_reduction)\n",
    "\n",
    "decorate(xlabel='Hand-washing campaign spending (USD)',\n",
    "         ylabel='Percent reduction in infection rate',\n",
    "         title='Effect of hand washing on infection rate',\n",
    "         legend=False)\n",
    "\n",
    "savefig('figs/chap05-fig04.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Modify the parameters `M`, `K`, and `B`, and see what effect they have on the shape of the curve.  Read about the [generalized logistic function on Wikipedia](https://en.wikipedia.org/wiki/Generalised_logistic_function).  Modify the other parameters and see what effect they have."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hand washing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can model the effect of a hand-washing campaign by modifying `beta`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_hand_washing(system, spending):\n",
    "    \"\"\"Modifies system to model the effect of hand washing.\n",
    "    \n",
    "    system: System object\n",
    "    spending: campaign spending in USD\n",
    "    \"\"\"\n",
    "    factor = compute_factor(spending)\n",
    "    system.beta *= (1 - factor)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start with the same values of `beta` and `gamma` we've been using."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "tc = 3      # time between contacts in days \n",
    "tr = 4      # recovery time in days\n",
    "\n",
    "beta = 1 / tc      # contact rate in per day\n",
    "gamma = 1 / tr     # recovery rate in per day\n",
    "\n",
    "beta, gamma"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can sweep different levels of campaign spending."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "spending_array = linspace(0, 1200, 13)\n",
    "\n",
    "for spending in spending_array:\n",
    "    system = make_system(beta, gamma)\n",
    "    add_hand_washing(system, spending)\n",
    "    results = run_simulation(system, update_func)\n",
    "    print(spending, system.beta, calc_total_infected(results))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's a function that sweeps a range of spending and stores the results in a `SweepSeries`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sweep_hand_washing(spending_array):\n",
    "    \"\"\"Run simulations with a range of spending.\n",
    "    \n",
    "    spending_array: array of dollars from 0 to 1200\n",
    "    \n",
    "    returns: Sweep object\n",
    "    \"\"\"\n",
    "    sweep = SweepSeries()\n",
    "    \n",
    "    for spending in spending_array:\n",
    "        system = make_system(beta, gamma)\n",
    "        add_hand_washing(system, spending)\n",
    "        results = run_simulation(system, update_func)\n",
    "        sweep[spending] = calc_total_infected(results)\n",
    "        \n",
    "    return sweep"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's how we run it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "spending_array = linspace(0, 1200, 20)\n",
    "infected_sweep = sweep_hand_washing(spending_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here's what it looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot(infected_sweep)\n",
    "\n",
    "decorate(xlabel='Hand-washing campaign spending (USD)',\n",
    "         ylabel='Total fraction infected',\n",
    "         title='Effect of hand washing on total infections',\n",
    "         legend=False)\n",
    "\n",
    "savefig('figs/chap05-fig05.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's put it all together to make some public health spending decisions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Suppose we have \\$1200 to spend on any combination of vaccines and a hand-washing campaign."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_students = 90\n",
    "budget = 1200\n",
    "price_per_dose = 100\n",
    "max_doses = int(budget / price_per_dose)\n",
    "dose_array = linrange(max_doses, endpoint=True)\n",
    "max_doses"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can sweep through a range of doses from, 0 to `max_doses`, model the effects of immunization and the hand-washing campaign, and run simulations.\n",
    "\n",
    "For each scenario, we compute the fraction of students who get sick."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "for doses in dose_array:\n",
    "    fraction = doses / num_students\n",
    "    spending = budget - doses * price_per_dose\n",
    "    \n",
    "    system = make_system(beta, gamma)\n",
    "    add_immunization(system, fraction)\n",
    "    add_hand_washing(system, spending)\n",
    "    \n",
    "    results, run_simulation(system, update_func)\n",
    "    print(doses, system.init.S, system.beta, calc_total_infected(results))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following function wraps that loop and stores the results in a `Sweep` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sweep_doses(dose_array):\n",
    "    \"\"\"Runs simulations with different doses and campaign spending.\n",
    "    \n",
    "    dose_array: range of values for number of vaccinations\n",
    "    \n",
    "    return: Sweep object with total number of infections \n",
    "    \"\"\"\n",
    "    sweep = SweepSeries()\n",
    "    \n",
    "    for doses in dose_array:\n",
    "        fraction = doses / num_students\n",
    "        spending = budget - doses * price_per_dose\n",
    "        \n",
    "        system = make_system(beta, gamma)\n",
    "        add_immunization(system, fraction)\n",
    "        add_hand_washing(system, spending)\n",
    "        \n",
    "        results = run_simulation(system, update_func)\n",
    "        sweep[doses] = calc_total_infected(results)\n",
    "\n",
    "    return sweep"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can compute the number of infected students for each possible allocation of the budget."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "infected_sweep = sweep_doses(dose_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And plot the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot(infected_sweep)\n",
    "\n",
    "decorate(xlabel='Doses of vaccine',\n",
    "         ylabel='Total fraction infected',\n",
    "         title='Total infections vs. doses',\n",
    "         legend=False)\n",
    "\n",
    "savefig('figs/chap05-fig06.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercises\n",
    "\n",
    "**Exercise:** Suppose the price of the vaccine drops to $50 per dose.  How does that affect the optimal allocation of the spending?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Suppose we have the option to quarantine infected students.  For example, a student who feels ill might be moved to an infirmary, or a private dorm room, until they are no longer infectious.\n",
    "\n",
    "How might you incorporate the effect of quarantine in the SIR model?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  }
 ],
 "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}