{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Modeling and Simulation in Python\n",
    "\n",
    "Chapter 14\n",
    "\n",
    "Copyright 2017 Allen Downey\n",
    "\n",
    "License: [Creative Commons Attribution 4.0 International](https://creativecommons.org/licenses/by/4.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "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 from previous chapters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "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 /= np.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": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_func(state, t, system):\n",
    "    \"\"\"Update the SIR model.\n",
    "    \n",
    "    state: State (s, i, r)\n",
    "    t: time\n",
    "    system: System object\n",
    "    \n",
    "    returns: State (sir)\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": 4,
   "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",
    "    unpack(system)\n",
    "    \n",
    "    frame = TimeFrame(columns=init.index)\n",
    "    frame.row[t0] = init\n",
    "    \n",
    "    for t in linrange(t0, t_end):\n",
    "        frame.row[t+1] = update_func(frame.row[t], t, system)\n",
    "    \n",
    "    return frame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "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": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sweep_beta(beta_array, gamma):\n",
    "    \"\"\"Sweep a range of values for beta.\n",
    "    \n",
    "    beta_array: array of beta values\n",
    "    gamma: recovery rate\n",
    "    \n",
    "    returns: SweepSeries that maps from beta to total infected\n",
    "    \"\"\"\n",
    "    sweep = SweepSeries()\n",
    "    for beta in beta_array:\n",
    "        system = make_system(beta, gamma)\n",
    "        results = run_simulation(system, update_func)\n",
    "        sweep[system.beta] = calc_total_infected(results)\n",
    "    return sweep"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SweepFrame\n",
    "\n",
    "The following sweeps two parameters and stores the results in a `SweepFrame`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sweep_parameters(beta_array, gamma_array):\n",
    "    \"\"\"Sweep a range of values for beta and gamma.\n",
    "    \n",
    "    beta_array: array of infection rates\n",
    "    gamma_array: array of recovery rates\n",
    "    \n",
    "    returns: SweepFrame with one row for each beta\n",
    "             and one column for each gamma\n",
    "    \"\"\"\n",
    "    frame = SweepFrame(columns=gamma_array)\n",
    "    for gamma in gamma_array:\n",
    "        frame[gamma] = sweep_beta(beta_array, gamma)\n",
    "    return frame"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's what the results look like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta_array = linspace(0.1, 0.9, 11)\n",
    "gamma_array = linspace(0.1, 0.7, 4)\n",
    "frame = sweep_parameters(beta_array, gamma_array)\n",
    "frame.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here's how we can plot the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "for gamma in gamma_array:\n",
    "    label = 'gamma = ' + str(gamma)\n",
    "    plot(frame[gamma], label=label)\n",
    "    \n",
    "decorate(xlabel='Contacts per day (beta)',\n",
    "         ylabel='Fraction infected',\n",
    "         loc='upper left')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's often useful to separate the code that generates results from the code that plots the results, so we can run the simulations once, save the results, and then use them for different analysis, visualization, etc."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Contact number"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After running `sweep_parameters`, we have a `SweepFrame` with one row for each value of `beta` and one column for each value of `gamma`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "frame.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following loop shows how we can loop through the columns and rows of the `SweepFrame`.  With 11 rows and 4 columns, there are 44 elements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "for gamma in frame.columns:\n",
    "    series = frame[gamma]\n",
    "    for beta in series.index:\n",
    "        frac_infected = series[beta]\n",
    "        print(beta, gamma, frac_infected)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can wrap that loop in a function and plot the results.  For each element of the `SweepFrame`, we have `beta`, `gamma`, and `frac_infected`, and we plot `beta/gamma` on the x-axis and `frac_infected` on the y-axis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_sweep_frame(frame):\n",
    "    \"\"\"Plot the values from a SweepFrame.\n",
    "    \n",
    "    For each (beta, gamma), compute the contact number,\n",
    "    beta/gamma\n",
    "    \n",
    "    frame: SweepFrame with one row per beta, one column per gamma\n",
    "    \"\"\"\n",
    "    for gamma in frame.columns:\n",
    "        series = frame[gamma]\n",
    "        for beta in series.index:\n",
    "            frac_infected = series[beta]\n",
    "            plot(beta/gamma, frac_infected, 'ro')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's what it looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_sweep_frame(frame)\n",
    "\n",
    "decorate(xlabel='Contact number (beta/gamma)',\n",
    "         ylabel='Fraction infected',\n",
    "         legend=False)\n",
    "\n",
    "savefig('figs/chap06-fig03.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It turns out that the ratio `beta/gamma`, called the \"contact number\" is sufficient to predict the total number of infections; we don't have to know `beta` and `gamma` separately.\n",
    "\n",
    "We can see that in the previous plot: when we plot the fraction infected versus the contact number, the results fall close to a curve."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the book we figured out the relationship between $c$ and $s_{\\infty}$ analytically.  Now we can compute it for a range of values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "s_inf_array = linspace(0.0001, 0.9999, 101);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "c_array = log(s_inf_array) / (s_inf_array - 1);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`total_infected` is the change in $s$ from the beginning to the end."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "frac_infected = 1 - s_inf_array\n",
    "frac_infected_series = Series(frac_infected, index=c_array);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can plot the analytic results and compare them to the simulations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_sweep_frame(frame)\n",
    "plot(frac_infected_series, label='Analysis')\n",
    "\n",
    "decorate(xlabel='Contact number (c)',\n",
    "         ylabel='Fraction infected')\n",
    "\n",
    "savefig('figs/chap06-fig04.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The agreement is generally good, except for values of `c` less than 1."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:**  If we didn't know about contact numbers, we might have explored other possibilities, like the difference between `beta` and `gamma`, rather than their ratio.\n",
    "\n",
    "Write a version of `plot_sweep_frame`, called `plot_sweep_frame_difference`, that plots the fraction infected versus the difference `beta-gamma`.\n",
    "\n",
    "What do the results look like, and what does that imply? "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Suppose you run a survey at the end of the semester and find that 26% of students had the Freshman Plague at some point.\n",
    "\n",
    "What is your best estimate of `c`?\n",
    "\n",
    "Hint: if you print `frac_infected_series`, you can read off the answer. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Alternative solution\n",
    "\n",
    "\"\"\"We can use `np.interp` to look up `s_inf` and\n",
    "estimate the corresponding value of `c`, but it only\n",
    "works if the index of the series is sorted in ascending\n",
    "order.  So we have to use `sort_index` first.\n",
    "\"\"\"\n",
    "\n",
    "frac_infected_series.sort_index(inplace=True)\n",
    "np.interp(0.26, frac_infected_series, frac_infected_series.index)"
   ]
  },
  {
   "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}