Sandbox::https://github.com/AllenDowney/ModSimPy/blob/master/chap11.py

Repository navigation

from modsim import *

def make_system(beta, gamma):

init = State(s=89, i=1, r=0)

init /= init.sum()

return System(init=init, t_end=7*14,

beta=beta, gamma=gamma)

from modsim import *

def update_func(t, state, system):

s, i, r = state.s, state.i, state.r

infected = system.beta * i * s

recovered = system.gamma * i

s -= infected

i += infected - recovered

r += recovered

return State(s=s, i=i, r=r)

from modsim import *

def plot_results(S, I, R):

S.plot(style='--', label='Susceptible')

I.plot(style='-', label='Infected')

R.plot(style=':', label='Recovered')

decorate(xlabel='Time (days)',

ylabel='Fraction of population')

from modsim import *

def run_simulation(system, update_func):

frame = TimeFrame(columns=system.init.index)

frame.loc[0] = system.init

for t in range(0, system.t_end):

frame.loc[t+1] = update_func(t, frame.loc[t], system)

return frame