Hierarchical model with Autoregressive residuals

#1


I am trying to duplicate the hierarchical model shown above in Pymc3.
I have no problem to build this model in pymc3 with the residuals for each OS and PD follow by N~(0,sigma^2). After reading this Analysis of AR(1) in Pymc3, I still do not know how to incorporate the first order Autoregressive residuals AR(1) into my model in Pymc3.
Therefore, I have no choice but to post this question here.
Below is the model I built without the AR(1) as the residuals and feel free to give me advice or feedback on my model as well. Thank you for your precious time! Appreciate it!

mu_vector_prior = np.array([0.9,0.5])
cov_matrix_prior = np.array([[0.1,0],[0,0.1]])
mu_vector = np.array([5,0.5])
cov_matrix= np.array([[0.5,0],[0,0.1]])


with pm.Model() as hierarchical_model:
    
    #Prior for HyperParameters
    prior_mu = pm.MvNormal('prior', mu = mu_vector_prior , cov = cov_matrix_prior, shape = 2) 
    prior_cov = pm.LKJCholeskyCov('packed_L', n=2, eta=2., sd_dist=pm.HalfCauchy.dist(2.5))
    L = pm.expand_packed_triangular(2, prior_cov)
    sig = pm.Deterministic('sig', L.dot(L.T))
    
    #Model's parameter
    
    #random-effect
    model_param_rnd = pm.MvNormal("rnd_var", mu = prior_mu, chol = L, shape = (10,2))

    #fixed-effect
    model_param_fix = pm.MvNormal("fix_var", mu = mu_vector, cov = cov_matrix, shape = 2)
    
    #Model error
    epsOS = pm.Uniform('epsOS', upper = 10000, lower = 0)
    epsPD = pm.Uniform('epsPD', upper = 5000, lower = 0)
    
    #Expected value
    RLR = model_param_rnd[npcohort,0]
    RRF = model_param_rnd[npcohort,1]
    Ker = model_param_fix[0]
    Kp = model_param_fix[1]
    
    A = (P*Ker*RLR)/(Ker-Kp)
    
    B = (P*Ker*RLR*RRF)/(Ker-Kp)
    
    OS = A*((np.exp(-Kp*nptime))-np.exp(-(Ker)*nptime))
    
    PD = B*((Ker*(1-np.exp(-Kp*nptime))) - (Kp*(1-np.exp(-Ker*nptime))))
    
    #likelihood
    y_OS = pm.Normal('y_OS', mu = OS, sd = epsOS, observed = dataOS.Loss.values)
    y_PD = pm.Normal('y_PD', mu = PD, sd = epsPD, observed = dataPD.Loss.values)
#2

Something along the line of:

beta = ...
epsOS = pm.AR('y', beta, sd=1.0, shape=len(dataOS.Loss.values))