python 疾病传播模型，在PyMC 2中

Posted 2021-05-09

def measles_model(obs_date, confirmation=True, all_traces=False):
    
    n_districts, n_periods, n_age_groups = I_obs.shape
    
    ### Confirmation sub-model
    
    if confirmation:

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

        mu = Normal("mu", mu=0, tau=0.0001, value=[0]*len(age_classes))
        sig = HalfCauchy('sig', 0, 25, value=1)
        var = sig**2
        cor = Uniform('cor', -1, 1, value=0)

        # Build variance-covariance matrix with first-order correlation 
        # among age classes
        @deterministic
        def Sigma(var=var, cor=cor):
            I = np.eye(len(age_classes))*var
            E = np.diag(np.ones(len(age_classes)-1), k=-1)*var*cor
            return I + E + E.T

        # Age-specific probabilities of confirmation as multivariate normal 
        # random variables
        beta_age = MvNormalCov("beta_age", mu=mu, C=Sigma, 
                        value=[1]*len(age_classes))
        p_age = Lambda('p_age', lambda t=beta_age: invlogit(t))

        @deterministic(trace=False)
        def p_confirm(beta=beta_age):
            return invlogit(beta[age_index])


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


    '''
    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
    
    if confirmation:
        
        @stochastic(trace=all_traces, dtype=int)
        def I(value=(I_obs_t*0.7).astype(int), n=I_obs_t, p=p_age):
            # Binomial confirmation process: confirm by age, then re-combine
            
            return np.sum([[binomial_like(value[d,:,a], n[d,:,a], p[a]) 
                            for a in range(n_age_groups)] 
                                for d in range(n_districts)])

    else:
        
        I = I_obs_t
        
        
    # Estimage age distribution from observed distribution of infecteds to date
    age_dist = Lambda('age_dist', lambda I=I: np.sum(I, (0,1))/I.sum())

    # Weakly-informative prior on proportion susceptible being 
    # between 0 and 0.07
    p_susceptible = Beta('p_susceptible', 2, 100, value=0.08)
    
    # Estimated total susceptibles by district
    S_0 = Binomial('S_0', n=N.values.astype(int), p=p_susceptible)

    @deterministic(trace=all_traces)
    def S(I=I, S_0=S_0):
        # Calculate susceptibles from total number of infections
        return np.array([S_0[d] - np.array([I[d,:t].sum() 
                        for t in range(t_obs+1)])
                            for d in range(n_districts)])
    
    # Susceptibles at time t, by age
    S_age_init = [rmultinomial(s[-1], age_dist.value) for s in S.value]
    @stochastic(dtype=int)
    def S_age(value=S_age_init, S=S, p=age_dist):
        return sum([multinomial_like(v, s[-1], p) for v,s in zip(value, S)])
    
    # Transmission parameter
    β = Gamma('β', 1, 0.1, value=0.001) 
    
    θ = Exponential('θ', 1, value=100)
    @deterministic
    def Iw(I=I, θ=θ): 
        # Distance-weighted infecteds
        return np.transpose([np.exp(-θ*distance_matrix.values).dot(I_t) for I_t in I.sum(2).T])
    
    α = Exponential('α', 1, value=1)
    
    # Force of infection
    @deterministic
    def λ(β=β, I=Iw, S=S, α=α): 
        return np.array([β*((I_d+0.1)**α)*S_d for I_d, S_d in zip(I,S)])
    
    # FOI in observation period
    λ_t = Lambda('λ_t', lambda λ=λ: λ[:, -1])
    
    # Poisson likelihood for observed cases
    @potential
    def new_cases(I=I, λ=λ): 
        return poisson_like(I.sum(2)[:,1:], λ[:, :-1])