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 *
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 update_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
s -= infected
i += infected - recovered
r += recovered
return State(S=s, I=i, R=r)
In [4]:
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 [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 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
In [7]:
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 the SweepFrame from the previous chapter, with one row for each value of beta and one column for each value of 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_array = [0.2, 0.4, 0.6, 0.8]
frame = sweep_parameters(beta_array, gamma_array)
frame.head()
In [9]:
frame.shape
The following loop shows how we can loop through the columns and rows of the SweepFrame. With 11 rows and 4 columns, there are 44 elements.
In [10]:
for gamma in frame.columns:
column = frame[gamma]
for beta in column.index:
frac_infected = column[beta]
print(beta, gamma, frac_infected)
Now we can wrap that loop in a function and plot the results. For each element of the SweepFrame, we have beta, gamma, and frac_infected, and we plot beta/gamma on the x-axis and frac_infected on the y-axis.
In [11]:
def plot_sweep_frame(frame):
"""Plot the values from a SweepFrame.
For each (beta, gamma), compute the contact number,
beta/gamma
frame: SweepFrame with one row per beta, one column per gamma
"""
for gamma in frame.columns:
column = frame[gamma]
for beta in column.index:
frac_infected = column[beta]
plot(beta/gamma, frac_infected, 'ro')
Here's what it looks like:
In [12]:
plot_sweep_frame(frame)
decorate(xlabel='Contact number (beta/gamma)',
ylabel='Fraction infected')
savefig('figs/chap14-fig01.pdf')
It turns out that the ratio beta/gamma, called the "contact number" is sufficient to predict the total number of infections; we don't have to know beta and gamma separately.
We can see that in the previous plot: when we plot the fraction infected versus the contact number, the results fall close to a curve.
In the book we figured out the relationship between $c$ and $s_{\infty}$ analytically. Now we can compute it for a range of values:
In [13]:
s_inf_array = linspace(0.0001, 0.9999, 101);
In [14]:
c_array = log(s_inf_array) / (s_inf_array - 1);
total_infected is the change in $s$ from the beginning to the end.
In [15]:
frac_infected = 1 - s_inf_array
frac_infected_series = Series(frac_infected, index=c_array);
Now we can plot the analytic results and compare them to the simulations.
In [16]:
plot_sweep_frame(frame)
plot(frac_infected_series, label='Analysis')
decorate(xlabel='Contact number (c)',
ylabel='Fraction infected')
savefig('figs/chap14-fig02.pdf')
The agreement is generally good, except for values of c less than 1.
Exercise: If we didn't know about contact numbers, we might have explored other possibilities, like the difference between beta and gamma, rather than their ratio.
Write a version of plot_sweep_frame, called plot_sweep_frame_difference, that plots the fraction infected versus the difference beta-gamma.
What do the results look like, and what does that imply?
In [17]:
# Solution goes here
In [18]:
# Solution goes here
In [19]:
# Solution goes here
Exercise: Suppose you run a survey at the end of the semester and find that 26% of students had the Freshman Plague at some point.
What is your best estimate of c?
Hint: if you print frac_infected_series, you can read off the answer.
In [20]:
# Solution goes here
In [21]:
# Solution goes here
In [ ]: