Hierarchical disease dynamics model

sir.py

# =======================================
# = Hierarchical disease dynamics model =
# =======================================

# Methods taken from: Zipkin, E.F., Jennelle, C.S. and Cooch, E.G. 2010. A primer on the
# application of Markov chains to the study of wildlife disease dynamics. Methods in Ecology
# and Evolution. 1: 192-198

from numpy import *
from pymc import *

# Sample size
N = 30
# Time intervals
T = 50

"""
Create a matrix y, with the encounter history data with 0 (uninfected), 1
(infected), 2 (dead) for males and females.
"""
P_true = array([[[0.90, 0.05, 0.05],
                [0.10, 0.70, 0.20],
                [0, 0, 1]],
                [[0.85, 0.1, 0.05],
                [0.05, 0.75, 0.20],
                [0, 0, 1]]])

y = []
sex = []
for i in range(N):
    individual = [rcategorical([0.7, 0.3])]
    sex.append(int(rbernoulli(0.5)))
    for t in range(1,T):
        individual.append(rcategorical(P_true[sex[-1],individual[-1]]))

    y.append(individual)
y = array(y)
sex = array(sex)


# Create identity matrix for calculating life expectancies
Id = eye(2)

# Probability of infection
alpha0 = Normal('alpha0', mu=0.0, tau=0.01, value=0.0)
alpha1 = Normal('alpha1', mu=0.0, tau=0.01, value=0.0)
pi01 = Lambda('pi01', lambda a0=alpha0, a1=alpha1: [invlogit(a0), invlogit(a0+a1)])

# Probability of recovery
beta0 = Normal('beta0', mu=0.0, tau=0.01, value=0.0)
beta1 = Normal('beta1', mu=0.0, tau=0.01, value=0.0)
pi10 = Lambda('pi10', lambda b0=beta0, b1=beta1: [invlogit(b0), invlogit(b0+b1)])

# Probabilities of survival
gamma0 = Normal('gamma0', mu=0.0, tau=0.01, value=0.0)
gamma1 = Normal('gamma1', mu=0.0, tau=0.01, value=0.0)
s = Lambda('s', lambda g0=gamma0, g1=gamma1: [invlogit(g0), invlogit(g0+g1)])

@deterministic
def P(pi01=pi01, pi10=pi10, s=s):
    """Transition matrix"""
    return array([[[s[0]*(1-pi01[i]), s[0]*pi01[i], 1-s[0]],
            [s[1]*pi10[i], s[1]*(1-pi10[i]), 1-s[1]],
            [0, 0, 1]] for i in range(2)])

@observed
def encounters(value=y, P=P):
    """Categorical likelihood for encounters"""
    return sum([[categorical_like(value[i,t], P[sex[i],value[i,t-1]]) for t in range(1,T)] for i in range(N)])

# Calculated nodes

"""Probability of moving from susceptible (0) to infected (1) over course of
study period and expected time to first transition"""
Pr01 = Lambda('Pr01', lambda P=P: P[:,0,1]/(1-P[:,0,0]))
Ex01 = Lambda('Ex01', lambda P=P: 1/(1-P[:,0,0]))

"""Probability of moving from infected (1) to susceptible (0) over course of
study period and expected time to first transition"""
Pr10 = Lambda('Pr10', lambda P=P: P[:,1,0]/(1-P[:,1,1]))
Ex10 = Lambda('Ex10', lambda P=P: 1/(1-P[:,1,1]))

# Life expectancy
Q = Lambda('Q', lambda P=P: P[:, :2, :2])
IminusQ = Lambda('IminusQ', lambda Q=Q: [Id-q for q in Q])

InvIminusQ = Lambda('InvIminusQ', lambda IQ=IminusQ: [inverse(iq) for iq in IQ])
W = Lambda('W', lambda IIQ=InvIminusQ, IQ=IminusQ: [[iiq[i,0]+iq[i,1] for i in range(2)] for iiq,iq in zip(IIQ, IQ)])