Spatial disease transmission model, translated to PyMC3

disease_transmission3.py

obs_date = '1997-06-15'


with Model() as june_model:
    
    n_districts, n_periods, n_age_groups = I_obs.shape
    
    ### Confirmation sub-model


    # Specify priors on age-specific means
    age_classes = np.unique(age_index)

    mu = Normal("mu", mu=0, tau=0.0001, shape=len(age_classes))
    sig = HalfCauchy('sig', 5, testval=1)
    var = sig**2
    cor = Uniform('cor', -1, 1, testval=0)

    # Build variance-covariance matrix with first-order correlation 
    # among age classes
    diag = TT.eye(len(age_classes)) * var
    Sigma = TT.fill_diagonal_offset(TT.fill_diagonal_offset(diag, var*cor, 1), var*cor, -1)

    # Age-specific probabilities of confirmation as multivariate normal 
    # random variables
    beta_age = MvNormal('beta_age', mu, inv(Sigma), shape=len(age_classes))
    p_age = Deterministic('p_age', invlogit(beta_age))
    p_confirm = invlogit(beta_age[age_index])

    # Confirmation likelihood
    lab_confirmed = Bernoulli('lab_confirmed', p_confirm, observed=y)



    '''
    Truncate data at observation period
    '''
    obs_index = all_district_data[0].index <= obs_date
    I_obs_t = np.array([I_dist[obs_index] for I_dist in I_obs])  
                        

    # Index for observation date, used to index out values of interest 
    # from the model.
    t_obs = obs_index.sum() - 1
    

    # Binomial confirmation process: confirm by age
    I_age = Binomial('I_age', n=I_obs_t, p=p_age, shape=I_obs_t.shape)

    # Aggregate by age
    I = TT.sum(I_age, 2)
        
    # Estimage age distribution from observed distribution of infecteds to date
    age_dist, _age_dist = june_model.TransformedVar(
        'age_dist', Dirichlet.dist(TT.constant([1]*n_age_groups), shape=n_age_groups),
        simplextransform)
    
    I_age_like = Potential('I_age_like', Multinomial.dist(TT.sum(I_age), 
                                                       age_dist,
                                                       shape=n_age_groups).logp(TT.sum(I_age, (0,1))))
    

    # Weakly-informative prior on proportion susceptible being 
    # between 0 and 0.07
    p_susceptible = Beta('p_susceptible', 2, 100)
    
    # Estimated total susceptibles by district
    S_0 = Binomial('S_0', n=N.values.astype(int), p=p_susceptible, shape=N.shape, 
                   testval=(0.04*N.values).astype(int))
    
    
    S, _ = map(lambda s, i: s - TT.cumsum(i), sequences=[S_0, I])

    
    # Transmission parameter
    β = Gamma('β', 1, 0.1, testval=0.001)
    
    
    θ = Exponential('θ', 1, testval=1)
    Iw, _ = map(lambda i, θ: TT.dot(TT.exp(-θ*distance_matrix.values), i), [TT.transpose(I)],
                            non_sequences=[θ])
    
    
    λ, _ = map(lambda i, s, b: b*i[:-1]*s[:-1], [I,S], non_sequences=[β])
    
    λ_t = λ[:, -1]
    
    new_cases = Potential('new_cases', Poisson.dist(λ).logp(I[:,1:]))