Skip to content

Instantly share code, notes, and snippets.

@aflaxman aflaxman/.gitignore
Created Oct 14, 2010

Embed
What would you like to do?
*.pyc
*.png
*~
""" Script to fit the SI model defined in si_model.py
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
figure(1)
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')
yticks([950,975,1000,1025])
ylabel('S (people)')
axis([-1,10,925,1050])
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')
xticks(range(1,10,2))
xlabel('Time (days)')
yticks([0,15,30])
ylabel('I (people)')
axis([-1,10,-1,35])
savefig('SI.png')
# plot the joint distribution of the epidemic parameters
figure(2)
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)
legend()
axis([-1, 41, -.05, 1.05])
savefig('param_dist.png')
# plot the autocorrelation of the MCMC trace for each var as evidence of mixing
figure(3)
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)
xticks([])
yticks([])
axis([-11, 11,-.1, 1.1])
axes([.05, .5, .475, .5])
my_corr(mod.beta.trace(), with_dots=True)
yticks([0,.5,1.])
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])
my_corr(mod.S.trace()[:,i])
text(-10, .9, '$S_%d$'%i)
if i == 0:
yticks([0,.5,1.])
axes([.05 + i*.095, 0., .095, .25])
my_corr(mod.I.trace()[:,i])
text(-10, .9, '$I_%d$'%i)
if i == 0:
yticks([0,.5,1.])
savefig('acorr.png')
## 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
@deterministic
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)
@aflaxman

This comment has been minimized.

Copy link
Owner Author

aflaxman commented Oct 14, 2010

Now the question is "does it mix?"

@aflaxman

This comment has been minimized.

Copy link
Owner Author

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.

@kforeman

This comment has been minimized.

Copy link

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
You can’t perform that action at this time.