Created
April 8, 2014 13:15
-
-
Save fonnesbeck/10122382 to your computer and use it in GitHub Desktop.
Sao Paulo SIR model in PyMC 2
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
from pymc import * | |
# Deterministic duration of infection (2 weeks) | |
delta = 14. | |
gamma = 1./delta | |
def sir(): | |
# Early case-detection probability | |
theta = Beta('theta', 1, 1, value=[0.5, 0.9]) | |
@potential | |
def constrain_theta(theta=theta): | |
# Ensure late detection higher than early detection | |
if theta[0] > theta[1]: | |
return -np.inf | |
return 0 | |
# Changepoint for detection | |
switch = DiscreteUniform('switch', 0, periods, value=12) | |
@deterministic | |
def p_case(s=switch, t=theta): | |
# Detection rates | |
out = np.empty(periods) | |
out[:s] = theta[0] | |
out[s:] = theta[1] | |
return out | |
# Unknown true number of cases | |
cases = DiscreteUniform('cases', I, I+100, value=I+10) | |
# Observation likelihood | |
case_obs = Binomial('case_obs', cases, p_case, value=I, observed=True) | |
# True numbers of susceptibles | |
susceptibles = Lambda('susceptibles', lambda cases=cases: np.sum(cases) - np.cumsum(cases)) | |
# Constrain to expected range of R0 for measles | |
R0 = Normal('R0', 15, 4, value=15) | |
# Calculate initial transmission rate from R0 | |
beta = Lambda('beta', lambda R0=R0: R0 * gamma) | |
# Effective reproduction number | |
Rt = Lambda('Rt', lambda R0=R0, s=susceptibles: R0 * s / N) | |
# Force of infection | |
lam = Lambda('lam', lambda beta=beta, cases=cases: beta * cases) | |
# 2-week infection probabilities | |
p = Lambda('p', lambda lam=lam: 1. - np.exp(-lam / N)) | |
# Probability of observation a function of infection and detection | |
p_obs = Lambda('p_obs', lambda p=p, p_case=p_case: (p*p_case)[:-1]) | |
# Infection likelihood | |
new_cases = Binomial('new_cases', susceptibles[:-1], p_obs, value=I[1:], observed=True) | |
return locals() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment