Modeling and Simulation in Python

Chapter 13

Copyright 2017 Allen Downey

License: Creative Commons Attribution 4.0 International


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 *

Code from previous chapters

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)

Sweeping beta

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


Out[8]:
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))


0.1 0.010756340768063644
0.2 0.11898421353185373
0.3 0.5890954199973404
0.4 0.8013385277185551
0.5 0.8965769637207062
0.6 0.942929291399791
0.7 0.966299311298026
0.8 0.9781518959989762
0.9 0.9840568957948106
1.0 0.9868823507202488
1.1 0.988148177093735

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)


Out[11]:
values
0.1 0.010756
0.2 0.118984
0.3 0.589095
0.4 0.801339
0.5 0.896577
0.6 0.942929
0.7 0.966299
0.8 0.978152
0.9 0.984057
1.0 0.986882
1.1 0.988148

In [12]:
label = 'gamma = ' + str(gamma)
plot(infected_sweep, label=label)

decorate(xlabel='Contact rate (beta)',
         ylabel='Fraction infected')

savefig('figs/chap13-fig01.pdf')


Saving figure to file figs/chap13-fig01.pdf

Sweeping gamma

Using the same array of values for beta


In [13]:
beta_array


Out[13]:
[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1]

And now an array of values for gamma


In [14]:
gamma_array = [0.2, 0.4, 0.6, 0.8]


Out[14]:
[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')


Saving figure to file 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

# Sweep beta with fixed gamma
gamma = 1/2
infected_sweep = sweep_beta(beta_array, gamma)


Out[16]:
values
0.1 0.002736
0.2 0.007235
0.3 0.015929
0.4 0.038603
0.5 0.132438
0.6 0.346765
0.7 0.530585
0.8 0.661553
0.9 0.754595
1.0 0.821534
1.1 0.870219

In [17]:
# Solution

# Interpolating by eye, we can see that the infection rate passes through 0.4
# when beta is between 0.6 and 0.7
# We can use the `crossings` function to interpolate more precisely
# (although we don't know about it yet :)
beta_estimate = crossings(infected_sweep, 0.4)


Out[17]:
array([0.62548698])

In [18]:
# Solution

# Time between contacts is 1/beta
time_between_contacts = 1/beta_estimate


Out[18]:
array([1.59875429])

SweepFrame

The following sweeps two parameters and stores the results in a SweepFrame


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()


Out[20]:
0.2 0.4 0.6 0.8
0.1 0.010756 0.003642 0.002191 0.001567
0.2 0.118984 0.010763 0.005447 0.003644
0.3 0.589095 0.030185 0.010771 0.006526
0.4 0.801339 0.131563 0.020917 0.010780
0.5 0.896577 0.396409 0.046140 0.017640

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')


Saving figure to file 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')


Saving figure to file figs/chap13-fig04.pdf

In [ ]: