Skip to content

Instantly share code, notes, and snippets.

Created Oct 14, 2010
What would you like to do?
""" Script to fit the SI model defined in
Also produce some plots, to be used in a short healthyalgorithms blog
post about attaching statistical models to system dynamics models
from pylab import *
from pymc import *
# load the model
import si_model as mod
reload(mod) # this reload streamlines interactive debugging
# fit the model with mcmc
mc = MCMC(mod)
mc.use_step_method(AdaptiveMetropolis, [mod.beta, mod.gamma, mod.SI_0])
mc.sample(iter=200000, burn=100000, thin=20, verbose=1)
# plot the estimated size of S and I over time
subplots_adjust(hspace=0, right=1)
subplot(2, 1, 1)
title('SI model with uncertainty')
plot(mod.susceptible_data, 's', mec='black', color='black', alpha=.9)
plot(mod.S.stats()['mean'], color='red', linewidth=2)
plot(mod.S.stats()['95% HPD interval'], color='red',
linewidth=1, linestyle='dotted')
ylabel('S (people)')
subplot(2, 1, 2)
plot(mod.infected_data, 's', mec='black', color='black', alpha=.9)
plot(mod.I.stats()['mean'], color='blue', linewidth=2)
plot(mod.I.stats()['95% HPD interval'], color='blue',
linewidth=1, linestyle='dotted')
xlabel('Time (days)')
ylabel('I (people)')
# plot the joint distribution of the epidemic parameters
title('Joint Distribution of Epidemic Parameters')
plot([mod.beta.random() for i in range(1000)],
[mod.gamma.random() for i in range(1000)],
linestyle='none', marker='o', color='blue', mec='blue',
alpha=.5, label='Prior', zorder=-100)
plot(mod.beta.trace(), mod.gamma.trace(),
linestyle='none', marker='o', color='green', mec='green',
alpha=.5, label='Posterior', zorder=-99)
import scipy.stats
gkde = scipy.stats.gaussian_kde([mod.beta.trace(), mod.gamma.trace()])
x,y = mgrid[0:40:.05, 0:1.1:.05]
z = array(gkde.evaluate([x.flatten(),y.flatten()])).reshape(x.shape)
contour(x, y, z, linewidths=1, alpha=.5, cmap=cm.Greys)
ylabel(r'$\gamma$', fontsize=18, rotation=0)
xlabel(r'$\beta$', fontsize=18)
axis([-1, 41, -.05, 1.05])
# plot the autocorrelation of the MCMC trace for each var as evidence of mixing
def my_corr(tr, with_dots=False):
acorr(tr, detrend=mlab.detrend_mean, usevlines=True)
if with_dots:
acorr(tr, detrend=mlab.detrend_mean, usevlines=False)
axis([-11, 11,-.1, 1.1])
axes([.05, .5, .475, .5])
my_corr(mod.beta.trace(), with_dots=True)
text(-10, .9, r'$\beta$')
axes([.525, .5, .475, .5])
my_corr(mod.gamma.trace(), with_dots=True)
text(-10, .9, r'$\gamma$')
for i in range(10):
axes([.05 + i*.095, .25, .095, .25])
text(-10, .9, '$S_%d$'%i)
if i == 0:
axes([.05 + i*.095, 0., .095, .25])
text(-10, .9, '$I_%d$'%i)
if i == 0:
## SI model
from pymc import *
from numpy import *
#observed data
T = 10
susceptible_data = array([999,997,996,994,993,992,990,989,986,984])
infected_data = array([1,2,5,6,7,8,9,11,13,15])
# stochastic priors
beta = Uniform('beta', 0., 40., value=1.)
gamma = Uniform('gamma', 0., 1., value=.001)
SI_0 = Uninformative('SI_0', value=[999., 1.])
# deterministic compartmental model
def SI(SI_0=SI_0, beta=beta, gamma=gamma):
S = zeros(T)
I = zeros(T)
S[0] = SI_0[0]
I[0] = SI_0[1]
for i in range(1,T):
S[i] = S[i-1] - 0.05*beta*S[i-1]*I[i-1]/(S[i-1]+I[i-1])
I[i] = max(0., I[i-1] + 0.05*beta*S[i-1]*I[i-1]/(S[i-1]+I[i-1]) - gamma*I[i-1])
return S, I
S = Lambda('S', lambda SI=SI: SI[0])
I = Lambda('I', lambda SI=SI: SI[1])
# data likelihood
A = Poisson('A', mu=S, value=susceptible_data, observed=True)
B = Poisson('B', mu=I, value=infected_data, observed=True)

This comment has been minimized.

Copy link
Owner Author

@aflaxman aflaxman commented Oct 14, 2010

Now the question is "does it mix?"


This comment has been minimized.

Copy link
Owner Author

@aflaxman aflaxman commented Oct 17, 2010

the answer is "yes", but you need to use the adaptive metropolis step method, and you need to wait for long enough.


This comment has been minimized.

Copy link

@kforeman kforeman commented Oct 19, 2010

Now that we know it mixes, the question is "will it blend"?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment