Fit a Mixed effect model without loop for likelihood

Modeling
1

Hello Everyone,

I am trying to fit a mixed effect model in which there is a correlation within each id. Here is my data structure.
image

I fitted stan model with loop, but it takes longer time. I would really appreciate if anyone could given suggestion to make it fast. I have given the part of stan code for your review.

Thanks
Marimuthu.
N - Total number of observations
K - Number of id
Ni[K] - Number of observations in each id
P - Number of covariates
S (N by N) - within study covariance matrix
L (P by P)- between study covariance matrix

 int pos;
   pos = 1;
 // Likelihood
   for (k in 1:K) {
    int nk;
    
    nk=Ni[k];
    
    matrix[nk,P] Xk;
    vector[nk] Yk;
    matrix[nk, nk] Sk ;
       	
    Xk = block(X, pos, 1, nk, P);
    Yk = segment(y, pos, nk);
    Sk=block(S, pos, pos, nk, nk);	

      
    Yk ~ multi_normal(Xk*beta+ Xk*u_beta[k,]', Sk+Xk*L*Xk');

    // Update position
    pos = pos + nk;
  }
2

Hi

I do not think it would be possible to do this without a loop. If you however have access to multiple cores you can make use of the reduce_sum function in stan to segment the data and run the loop in parallel segments, see the user guide: Parallelization.

Sample code would be to define a new function

real partial_sum(array[] int index, int start, int end, array[] Ni, matrix X, vector y, matrix S, int P, array[] u_beta) {
    int K = num_elements(index); //the number of elements in the partitioned array
    real partial_target = 0; // target for the given data partition
    for (k in 1:K) {
        int pos = sum(Ni[:index[k]-1]) + 1; // position of first element
        int nk = Ni[index[k]];
        matrix[nk,P] Xk = block(X, pos, 1, nk, P);;
        vector[nk] Yk = segment(y, pos, nk);
        matrix[nk, nk] Sk = block(S, pos, pos, nk, nk);
        matrix[nk, nk] L_k = cholesky_decompose(Sk+Xk*L*Xk'); // cholesky decomposition of covariance matrix
        partial_target += multi_normal_cholesky_lpdf( Yk | Xk*beta+ Xk*u_beta[index[k],]', L_k);
        }
        return partial_target;
}

In the transformed data block you can define index to be an array of integers with values ranging from 1:K

transformed data {
   array[K] int index;
   for (k in 1:K) {
       index[k] = k;
   }
}

And in the model block use the reduce_sum function:

model {
  int grainsize=1;
  // other code
  target += reduce_sum(partial_sum, index, grainsize, Ni, X, ... // all other inputs);
}
1 Like
3

Hi Hermanus,

I appreciate your response. Hope this will sort out the problem. I will try and get back. Have a good day!

Marimuthu

4

In Stan, for loops usually don’t induce substantial penalty. Most commonly, sampling speed is limited by a problematic geometry of the model (e.g. non-identifiability, weak identifiability). Are your rhats and ess diagnostic good? Do you get divergent transitions? Do your pairs plots look nice?

If the model had good geometry a potential performance bottleneck is the multi_normal which will need to compute the inverse of the covariance matrix. If you can reformulate your model such that you parametrize the model with the precision matrix and use multi_normal_prec or cholesky decomposition of the covariance matrix and multi_normal_cholesky, you are likely to gain noticeable speedup. You will however likely not gain much speedup if you need to run the decomposition yourself, only if you can avoid computing it at all.

If you cannot avoid the matrix addition, I think (not sure if I remember correctly) that some people have seen a speedup using the QR approach to get the Cholesky of the sum from the Cholesky of the components (linear algebra - How do you find the Cholesky decomposition of the sum of two positive definite matrices without adding the matrices directly? - MathOverflow)

Best of luck with the model!

5

Hi Martinmodrak,

Thank you so much for your suggestions. My diagnostics are good only.
Regarding the Cholesky decomposition, I split the covariance matrix into standard deviation and correlation matrix, can I still use Cholesky decomposition for correlation matrix? If I use, does it help to reduce the running time? I really appreciate your hints / suggestions.

Marimuthu

6

Yes, Cholesky decomposition of correlation matrix can be useful.

7

Thank you, Martin. Let me try this.

Marimuthu.