.gitignore

*.pyc
*.png
*~

si_fit.py

""" 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.py

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

Now the question is "does it mix?"

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

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