{
"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": [
{
"data": {
"text/html": [
"\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" 0.1 | \n",
" 0.3 | \n",
" 0.5 | \n",
" 0.7 | \n",
"
\n",
" \n",
" \n",
" \n",
" | 0.10 | \n",
" 0.084693 | \n",
" 0.005444 | \n",
" 0.002736 | \n",
" 0.001827 | \n",
"
\n",
" \n",
" | 0.18 | \n",
" 0.708623 | \n",
" 0.015914 | \n",
" 0.006118 | \n",
" 0.003783 | \n",
"
\n",
" \n",
" | 0.26 | \n",
" 0.900780 | \n",
" 0.055380 | \n",
" 0.011639 | \n",
" 0.006427 | \n",
"
\n",
" \n",
" | 0.34 | \n",
" 0.956888 | \n",
" 0.267864 | \n",
" 0.022115 | \n",
" 0.010191 | \n",
"
\n",
" \n",
" | 0.42 | \n",
" 0.977045 | \n",
" 0.524563 | \n",
" 0.047816 | \n",
" 0.015946 | \n",
"
\n",
" \n",
"
\n",
"