Slow computing speed in fitting AR1xAR1 model

Modeling
,
9

I probably have resolved the problem. With the property of Kronecker product:
A\otimes B = (L_A L_A^\top)\otimes (L_B L_B^\top) = (L_A\otimes L_B)(L_A \otimes L_B)^\top, I found that in stan doing Kronecker product on vectors is faster than doing it on matrices.

This post transforms the Kolesky decomposition matrix L to a vector by L*z. However, it is still slow.

My idea is using the formula (L_A\otimes L_B)(z_1\otimes z_2) = (L_Az_1)\otimes (L_B z_2), with which we do Kronecker product on two vectors. It proves that this solution is much faster. I computed three matrices for 1000 times as an example and here are the results (suppose we have already got the three L matrices with sizes n1,n2,n3 and M=1000:

system.time({
   for(i in 1:M) r1 = as_matrix(chol_kronecker_product_three(L1,L2,L3,z5),n1*n2,n3)
 })
   user  system elapsed 
  28.37    0.02   28.52 
 system.time({
   for(i in 1:M) r2 = kroVecToMat(chol_kronecker_product_two(L1,L2,z4),L3%*%z3)
 })
   user  system elapsed 
   2.77    0.00    2.76 
 system.time({
   for(i in 1:M) r3 = kroVecToMat(kroVec(L1%*%z1,L2%*%z2),L3%*%z3)
 })
   user  system elapsed 
   0.16    0.00    0.16

R and rstan functions

z1 <- rnorm(n1)
z2 <- rnorm(n2)
z3 <- rnorm(n3)
z4 <- kronecker(z1,z2)
z5 <- kronecker(z4,z3)

stanfun <- 
  "
functions{
  matrix kroVecToMat(vector A, vector B){
    vector[rows(A)*rows(B)] C;
    for (j in 1:rows(A)) {
      C[((j - 1)*rows(B) + 1):(j*rows(B))] = B*A[j];
    }
    return to_matrix(C,rows(A),rows(B),0);
  }
  
  vector kroVec(vector A, vector B){
    vector[rows(A)*rows(B)] C;
    for (j in 1:rows(A)) {
      C[((j - 1)*rows(B) + 1):(j*rows(B))] = B*A[j];
    }
    return C;
  } 
  
  matrix as_matrix(vector X, int N, int K) { 
    matrix[N, K] Y; 
      for (i in 1:N) {
        Y[i] = to_row_vector(X[((i - 1) * K + 1):(i * K)]); 
      }
    return Y; 
  }

  vector chol_kronecker_product_three(matrix LA,matrix LB,matrix LC, vector d) {
    vector[num_elements(d)] new_d;
    new_d = rep_vector(0, num_elements(d));
    for(iA in 1:cols(LA)){
      for(jA in 1:iA){
        for(iB in 1:cols(LB)){
          for(jB in 1:iB){
            for(iC in 1:cols(LC)){
              for(jC in 1:iC){
                new_d[cols(LC)*(cols(LB)*(iA-1)+iB-1)+iC] = new_d[cols(LC)*(cols(LB)*(iA-1)+iB-1)+iC] + LA[iA,jA]*LB[iB,jB]*LC[iC,jC]*d[cols(LC)*(cols(LB)*(jA-1)+jB-1)+jC];
              }
            }
          }
        }
      }
    }
    return new_d;
  }
  
  vector chol_kronecker_product_two(matrix LA,matrix LB,vector d) {
    vector[num_elements(d)] new_d;
    new_d = rep_vector(0, num_elements(d));
    for(iA in 1:cols(LA)){
      for(jA in 1:iA){
        for(iB in 1:cols(LB)){
          for(jB in 1:iB){
            new_d[cols(LB)*(iA-1)+iB] = new_d[cols(LB)*(iA-1)+iB] + LA[iA,jA]*LB[iB,jB]*d[cols(LB)*(jA-1)+jB];
          }
        }
      }
    }
    return new_d;
  }
  
  matrix chol_three(matrix LA,matrix LB,matrix LC, vector d) {
    vector[num_elements(d)] new_d;
    new_d = rep_vector(0, num_elements(d));
    for(iA in 1:cols(LA)){
      for(jA in 1:iA){
        for(iB in 1:cols(LB)){
          for(jB in 1:iB){
            for(iC in 1:cols(LC)){
              for(jC in 1:iC){
                new_d[cols(LC)*(cols(LB)*(iA-1)+iB-1)+iC] = new_d[cols(LC)*(cols(LB)*(iA-1)+iB-1)+iC] + LA[iA,jA]*LB[iB,jB]*LC[iC,jC]*d[cols(LC)*(cols(LB)*(jA-1)+jB-1)+jC];
              }
            }
          }
        }
      }
    }
    return to_matrix(new_d,cols(LA)*cols(LB),cols(LC),0);
  }
  
  
}
"
expose_stan_functions(stanc(model_code = stanfun))

Comment (on 23/09/2020): Perhaps I am overthinking, but I do have a concern. The original method uses (L_1\otimes L_2)z, where z\sim N(0,1), to calculate the Kronecker product of two matrices and convert it to a vector. Then I used (L_1z_1)\otimes(L_2z_2) to calculate the Kronecker product of two vectors, where z_1\sim N(0,1), z_2\sim N(0,1).
My concern is: are these two methods identical? Numerically, I would say yes, but theoretically, I am not sure. If we expand the formula to (L_1z_1)\otimes(L_2z_2)=(L_1\otimes L_2)(z_1\otimes z_2), on the right side, z_1\otimes z_2 is supposed to be the same as z. But it is a tensor space, not a vector. I also have realized that if two variables a,b are Gaussian, the product ab is NOT Gaussian.
Please feel free to put comments below and let me know whether it is correct. Thank you.

4 Likes