Sao Paulo SIR model in PyMC 2

measles_sir.py

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