Efficiently populate a matrix in transformed parameters block

Modeling
,
1

For a model I’m working on, I need to populate a covariance matrix with entries dependent on the parameter lambda. The matrix QCorr is an already populated matrix of the same size as q_var and the function tukey_qd takes in a known p and the parameter lambda. Is there a faster way to fill the matrix other than using a for loop?

for (i in 1:N) {
    for (j in 1:i) {

      q_var[i, j] = (QCorr[i,j] *
                    (tukey_qd(p[i], lambda) *
                      tukey_qd(p[j], lambda)));

      q_var[j, i] = q_var[i, j];

      if (i == j) {q_var[i,j] = q_var[i,j] + .0000001;}

    }
  }
2

Loops in Stan are just compiled down to C++ and are very fast. What will slow you down is the arithmetic and autodiff. This can perhaps be sped up a bit.

First, you want to precompute the Tukey values, so you only have N calls rather than the (N * (N + 1)) / 2 calls in your code:

vector[N] tuk;
for (n in 1:N) {
  tuk[n] = tukey_qd(p[n], lambda);
}

Using tuk[i] where you have tukey_qd(p[I], lambda) and tuk[j] where you have tukey_qd(p[j], lambda) will give you a big speedup, especially if the tukey_qd function is expensive.

I’m not sure if it’ll be faster, but you could define your whole matrix in one go this way:

matrix[N, N] q_var = QCorr .* (tuk * tuk');
for (n in 1:N) {
  q_var[n, n] += 1e-7;

I would hope you could get away without conditioning the diagonal, but I left it in there.

Ideally, you’d use a Cholesky factor the whole thing:

cholesky_factor_corr[N] L_q_corr = cholesky_decompose(QCorr);
cholesky_factor[N] L_q_var = diag_pre_multiply(tuk, L_QCorr);

It is much faster working with Cholesky factors than with dense covariance matrices (quadratic solves instead of cubic and better numerical stability). Ideally, you’d declare QCorr as a Cholesky factor for a correlation matrix in order to skip the decomposition. This is what we outline in the regression chapter of the User’s Guide: Regression Models (link is to relevant section).