Skip to content

Instantly share code, notes, and snippets.

@aflaxman
Created October 14, 2010 18:10
Show Gist options
  • Star 4 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save aflaxman/626689 to your computer and use it in GitHub Desktop.
Save aflaxman/626689 to your computer and use it in GitHub Desktop.
*.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
Copy link
Author

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

@kforeman
Copy link

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