Modeling and Simulation in Python

Chapter 13

Copyright 2017 Allen Downey

License: Creative Commons Attribution 4.0 International

In [1]:
# Configure Jupyter so figures appear in the notebook
%matplotlib inline

# Configure Jupyter to display the assigned value after an assignment
%config InteractiveShell.ast_node_interactivity='last_expr_or_assign'

# import functions from the module
from modsim import *

Code from previous chapters

make_system, plot_results, and calc_total_infected are unchanged.

In [2]:
def make_system(beta, gamma):
    """Make a system object for the SIR model.
    beta: contact rate in days
    gamma: recovery rate in days
    returns: System object
    init = State(S=89, I=1, R=0)
    init /= np.sum(init)

    t0 = 0
    t_end = 7 * 14

    return System(init=init, t0=t0, t_end=t_end,
                  beta=beta, gamma=gamma)

In [3]:
def plot_results(S, I, R):
    """Plot the results of a SIR model.
    S: TimeSeries
    I: TimeSeries
    R: TimeSeries
    plot(S, '--', label='Susceptible')
    plot(I, '-', label='Infected')
    plot(R, ':', label='Recovered')
    decorate(xlabel='Time (days)',
             ylabel='Fraction of population')

In [4]:
def calc_total_infected(results):
    """Fraction of population infected during the simulation.
    results: DataFrame with columns S, I, R
    returns: fraction of population
    return get_first_value(results.S) - get_last_value(results.S)

Exercise: Write a slope function for the SIR model and test it with run_ode_solver.

In [5]:
# Solution

def slope_func(state, t, system):
    """Update the SIR model.
    state: State (s, i, r)
    t: time
    system: System object
    returns: State (sir)
    s, i, r = state

    infected = system.beta * i * s    
    recovered = system.gamma * i
    dSdt = -infected
    dIdt = infected - recovered
    dRdt = recovered
    return dSdt, dIdt, dRdt

In [6]:
system = make_system(0.333, 0.25)
slope_func(system.init, 0, system)

(-0.003658888888888889, 0.0008811111111111112, 0.002777777777777778)

In [7]:
results, details = run_ode_solver(system, slope_func, max_step=3)

sol None
t_events []
nfev 212
njev 0
nlu 0
status 0
message The solver successfully reached the end of the...
success True

In [8]:
plot_results(results.S, results.I, results.R)

In [ ]: