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 modsim.py module
from modsim import *
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)
Out[6]:
In [7]:
results, details = run_ode_solver(system, slope_func, max_step=3)
details
Out[7]:
In [8]:
plot_results(results.S, results.I, results.R)
In [ ]: