In [2]:
# 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 [3]:
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 [4]:
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 [5]:
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)
In [6]:
def run_simulation(system, update_func):
"""Runs a simulation of the system.
system: System object
update_func: function that updates state
returns: TimeFrame
"""
init, t0, t_end = system.init, system.t0, system.t_end
frame = TimeFrame(columns=init.index)
frame.row[t0] = init
for t in linrange(t0, t_end):
frame.row[t+1] = update_func(frame.row[t], t, system)
return frame
In [7]:
def update_func(state, t, system):
"""Update the SIR model.
state: State (s, i, r)
t: time
system: System object
returns: State (sir)
"""
beta, gamma = system.beta, system.gamma
s, i, r = state
infected = beta * i * s
recovered = gamma * i
s -= infected
i += infected - recovered
r += recovered
return State(S=s, I=i, R=r)
Make a range of values for beta, with constant gamma.
In [8]:
beta_array = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0 , 1.1]
gamma = 0.2
Run the simulation once for each value of beta and print total infections.
In [9]:
for beta in beta_array:
system = make_system(beta, gamma)
results = run_simulation(system, update_func)
print(system.beta, calc_total_infected(results))
Wrap that loop in a function and return a SweepSeries object.
In [10]:
def sweep_beta(beta_array, gamma):
"""Sweep a range of values for beta.
beta_array: array of beta values
gamma: recovery rate
returns: SweepSeries that maps from beta to total infected
"""
sweep = SweepSeries()
for beta in beta_array:
system = make_system(beta, gamma)
results = run_simulation(system, update_func)
sweep[system.beta] = calc_total_infected(results)
return sweep
Sweep beta and plot the results.
In [11]:
infected_sweep = sweep_beta(beta_array, gamma)
In [12]:
label = 'gamma = ' + str(gamma)
plot(infected_sweep, label=label)
decorate(xlabel='Contact rate (beta)',
ylabel='Fraction infected')
savefig('figs/chap13-fig01.pdf')
Using the same array of values for beta
In [13]:
beta_array
And now an array of values for gamma
In [14]:
gamma_array = [0.2, 0.4, 0.6, 0.8]
For each value of gamma, sweep beta and plot the results.
In [29]:
plt.figure(figsize=(7, 4))
for gamma in gamma_array:
infected_sweep = sweep_beta(beta_array, gamma)
label = 'gamma = ' + str(gamma)
plot(infected_sweep, label=label)
decorate(xlabel='Contact rate (beta)',
ylabel='Fraction infected',
loc='upper left')
plt.legend(bbox_to_anchor=(1.02, 1.02))
plt.tight_layout()
savefig('figs/chap13-fig02.pdf')
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.
In [16]:
# Solution goes here
In [17]:
# Solution goes here
In [18]:
# Solution goes here
In [19]:
def sweep_parameters(beta_array, gamma_array):
"""Sweep a range of values for beta and gamma.
beta_array: array of infection rates
gamma_array: array of recovery rates
returns: SweepFrame with one row for each beta
and one column for each gamma
"""
frame = SweepFrame(columns=gamma_array)
for gamma in gamma_array:
frame[gamma] = sweep_beta(beta_array, gamma)
return frame
Here's what the SweepFrame look like.
In [20]:
frame = sweep_parameters(beta_array, gamma_array)
frame.head()
And here's how we can plot the results.
In [21]:
for gamma in gamma_array:
label = 'gamma = ' + str(gamma)
plot(frame[gamma], label=label)
decorate(xlabel='Contact rate (beta)',
ylabel='Fraction infected',
title='',
loc='upper left')
We can also plot one line for each value of beta, although there are a lot of them.
In [28]:
plt.figure(figsize=(7, 4))
for beta in [1.1, 0.9, 0.7, 0.5, 0.3]:
label = 'beta = ' + str(beta)
plot(frame.row[beta], label=label)
decorate(xlabel='Recovery rate (gamma)',
ylabel='Fraction infected')
plt.legend(bbox_to_anchor=(1.02, 1.02))
plt.tight_layout()
savefig('figs/chap13-fig03.pdf')
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.
After running sweep_parameters, we have a SweepFrame with one row for each value of beta and one column for each value of gamma.
In [23]:
contour(frame)
decorate(xlabel='Recovery rate (gamma)',
ylabel='Contact rate (beta)',
title='Fraction infected, contour plot')
savefig('figs/chap13-fig04.pdf')
In [ ]: