{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Modeling and Simulation in Python\n",
    "\n",
    "Chapter 13\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 from previous chapters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`make_system`, `plot_results`, and `calc_total_infected` are unchanged."
   ]
  },
  {
   "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 /= 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": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_results(S, I, R):\n",
    "    \"\"\"Plot the results of a SIR model.\n",
    "    \n",
    "    S: TimeSeries\n",
    "    I: TimeSeries\n",
    "    R: TimeSeries\n",
    "    \"\"\"\n",
    "    plot(S, '--', label='Susceptible')\n",
    "    plot(I, '-', label='Infected')\n",
    "    plot(R, ':', label='Recovered')\n",
    "    decorate(xlabel='Time (days)',\n",
    "             ylabel='Fraction of population')"
   ]
  },
  {
   "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 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",
    "    init, t0, t_end = system.init, system.t0, system.t_end\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": 7,
   "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",
    "    beta, gamma = system.beta, system.gamma\n",
    "    s, i, r = state\n",
    "\n",
    "    infected = beta * i * s    \n",
    "    recovered = 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": "markdown",
   "metadata": {},
   "source": [
    "### Sweeping beta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make a range of values for `beta`, with constant `gamma`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta_array = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0 , 1.1]\n",
    "gamma = 0.2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the simulation once for each value of `beta` and print total infections."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "for beta in beta_array:\n",
    "    system = make_system(beta, gamma)\n",
    "    results = run_simulation(system, update_func)\n",
    "    print(system.beta, calc_total_infected(results))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Wrap that loop in a function and return a `SweepSeries` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "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": [
    "Sweep `beta` and plot the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "infected_sweep = sweep_beta(beta_array, gamma)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "label = 'gamma = ' + str(gamma)\n",
    "plot(infected_sweep, label=label)\n",
    "\n",
    "decorate(xlabel='Contact rate (beta)',\n",
    "         ylabel='Fraction infected')\n",
    "\n",
    "savefig('figs/chap13-fig01.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sweeping gamma"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the same array of values for `beta`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta_array"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now an array of values for `gamma`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "gamma_array = [0.2, 0.4, 0.6, 0.8]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For each value of `gamma`, sweep `beta` and plot the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(7, 4))\n",
    "\n",
    "for gamma in gamma_array:\n",
    "    infected_sweep = sweep_beta(beta_array, gamma)\n",
    "    label = 'gamma = ' + str(gamma)\n",
    "    plot(infected_sweep, label=label)\n",
    "    \n",
    "decorate(xlabel='Contact rate (beta)',\n",
    "         ylabel='Fraction infected',\n",
    "         loc='upper left')\n",
    "\n",
    "plt.legend(bbox_to_anchor=(1.02, 1.02))\n",
    "plt.tight_layout()\n",
    "savefig('figs/chap13-fig02.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:**  Suppose the infectious period for the Freshman Plague is known to be 2 days on average, and suppose during one particularly bad year, 40% of the class is infected at some point.  Estimate the time between contacts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "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": 19,
   "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 `SweepFrame` look like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "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": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "for gamma in gamma_array:\n",
    "    label = 'gamma = ' + str(gamma)\n",
    "    plot(frame[gamma], label=label)\n",
    "    \n",
    "decorate(xlabel='Contact rate (beta)',\n",
    "         ylabel='Fraction infected',\n",
    "         title='',\n",
    "         loc='upper left')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also plot one line for each value of `beta`, although there are a lot of them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(7, 4))\n",
    "\n",
    "\n",
    "for beta in [1.1, 0.9, 0.7, 0.5, 0.3]:\n",
    "    label = 'beta = ' + str(beta)\n",
    "    plot(frame.row[beta], label=label)\n",
    "    \n",
    "decorate(xlabel='Recovery rate (gamma)',\n",
    "         ylabel='Fraction infected')\n",
    "\n",
    "plt.legend(bbox_to_anchor=(1.02, 1.02))\n",
    "plt.tight_layout()\n",
    "savefig('figs/chap13-fig03.pdf')"
   ]
  },
  {
   "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.\n",
    "\n",
    "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": 23,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "contour(frame)\n",
    "\n",
    "decorate(xlabel='Recovery rate (gamma)',\n",
    "         ylabel='Contact rate (beta)',\n",
    "         title='Fraction infected, contour plot')\n",
    "\n",
    "savefig('figs/chap13-fig04.pdf')"
   ]
  },
  {
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}