# Modeling and Simulation in Python

Chapter 11

Copyright 2017 Allen Downey

``````

In :

# 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 *

``````

### SIR implementation

We'll use a `State` object to represent the number (or fraction) of people in each compartment.

``````

In :

init = State(S=89, I=1, R=0)

``````

To convert from number of people to fractions, we divide through by the total.

``````

In :

init /= sum(init)

``````

`make_system` creates a `System` object with the given parameters.

``````

In :

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 /= sum(init)

t0 = 0
t_end = 7 * 14

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

``````

Here's an example with hypothetical values for `beta` and `gamma`.

``````

In :

tc = 3      # time between contacts in days
tr = 4      # recovery time in days

beta = 1 / tc      # contact rate in per day
gamma = 1 / tr     # recovery rate in per day

system = make_system(beta, gamma)

``````

The update function takes the state during the current time step and returns the state during the next time step.

``````

In :

def update_func(state, t, system):
"""Update the SIR model.

state: State with variables S, I, R
t: time step
system: System with beta and gamma

returns: State object
"""
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)

``````

To run a single time step, we call it like this:

``````

In :

state = update_func(init, 0, system)

``````

Now we can run a simulation by calling the update function for each time step.

``````

In :

def run_simulation(system, update_func):
"""Runs a simulation of the system.

system: System object
update_func: function that updates state

returns: State object for final state
"""
state = system.init

for t in linrange(system.t0, system.t_end):
state = update_func(state, t, system)

return state

``````

The result is the state of the system at `t_end`

``````

In :

run_simulation(system, update_func)

``````

Exercise Suppose the time between contacts is 4 days and the recovery time is 5 days. After 14 weeks, how many students, total, have been infected?

Hint: what is the change in `S` between the beginning and the end of the simulation?

``````

In :

# Solution goes here

``````

### Using TimeSeries objects

If we want to store the state of the system at each time step, we can use one `TimeSeries` object for each state variable.

``````

In :

def run_simulation(system, update_func):
"""Runs a simulation of the system.

Add three Series objects to the System: S, I, R

system: System object
update_func: function that updates state
"""
S = TimeSeries()
I = TimeSeries()
R = TimeSeries()

state = system.init
t0 = system.t0
S[t0], I[t0], R[t0] = state

for t in linrange(system.t0, system.t_end):
state = update_func(state, t, system)
S[t+1], I[t+1], R[t+1] = state

return S, I, R

``````

Here's how we call it.

``````

In :

tc = 3      # time between contacts in days
tr = 4      # recovery time in days

beta = 1 / tc      # contact rate in per day
gamma = 1 / tr     # recovery rate in per day

system = make_system(beta, gamma)
S, I, R = run_simulation(system, update_func)

``````

And then we can plot the results.

``````

In :

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

``````

Here's what they look like.

``````

In :

plot_results(S, I, R)
savefig('figs/chap11-fig01.pdf')

``````

### Using a DataFrame

Instead of making three `TimeSeries` objects, we can use one `DataFrame`.

We have to use `row` to selects rows, rather than columns. But then Pandas does the right thing, matching up the state variables with the columns of the `DataFrame`.

``````

In :

def run_simulation(system, update_func):
"""Runs a simulation of the system.

system: System object
update_func: function that updates state

returns: TimeFrame
"""
frame = TimeFrame(columns=system.init.index)
frame.row[system.t0] = system.init

for t in linrange(system.t0, system.t_end):
frame.row[t+1] = update_func(frame.row[t], t, system)

return frame

``````

Here's how we run it, and what the result looks like.

``````

In :

tc = 3      # time between contacts in days
tr = 4      # recovery time in days

beta = 1 / tc      # contact rate in per day
gamma = 1 / tr     # recovery rate in per day

system = make_system(beta, gamma)
results = run_simulation(system, update_func)
results.head()

``````

We can extract the results and plot them.

``````

In :

plot_results(results.S, results.I, results.R)

``````

## Exercises

Exercise Suppose the time between contacts is 4 days and the recovery time is 5 days. Simulate this scenario for 14 weeks and plot the results.

``````

In :

# Solution goes here

``````