Matrix operation in Stan

I would like to create a particular matrix in Stan but I am not sure how to do it.

We have defined an identity matrix:

diagonalSigma = diag(N)

and within the stan program I have defined as input an array of bins (two bins in this case)

array[N] int<lower=1,upper=2> bins;

and I would like now to define a matrix to be a diagonal matrix consisting of the two sampled sigmas on the diagonal distributed according to the bins, which I am doing like this:

parameters {
  real<lower=0> sigmaD[2];
}


transformed parameters {
  matrix[N, N] cov = diagonalSigma*(sigmaD[bins]^2)
}
...

Which is wrong (Stan gives an error). What would be the correct syntax for this? I think it’s because I probably can’t do sigmaD[bins] but then how do I define the vector? To put it more concretely, for example if I have a list of bins:

\{1, 1,1,1,1,1,2,2,2,2,1,1,2,2,\ldots\}

and

\sigma_D=\{a,b\}

How do I create my vector in Stan to now be:

\{a, a,a,a,a,a,b,b,b,b,a,a,b,b,\ldots\}

So that I can then multiply with the identity matrix (which I am assuming should work) and hence create a new diagonal matrix. I can of course use a for loop:

cov = diagonalSigma;
for (i in 1:N) {
    cov[i,i] = sigmaD[bins[i]]^2
}

and do it one by one, but I am wondering if there is a vectorized, more efficient, way.

Thank you!

to do it elementwise?

Of course, made a clarification in the message, I was looking for a vectorized way that would be more efficient.

I don’t think your error is likely a problem with how you index. Can help more if you can provide the error, but first guess is that you intended to compute the power function elementwise, in which case you need .^.

Maybe, but if I remove that I get (I will define diagonalMatrix and sigmaD for context):

data {
  matrix[N,N] diagonalSigma;
  ...
}

parameters {
  real<lower=0> sigmaD[2];
}
   -------------------------------------------------
    14:  transformed parameters {
    15:    cov_matrix[N] choleskyCov;
    16:    matrix[N, N] cov = diagonalSigma*(sigmaD[bins]);
                              ^
    17:    choleskyCov = cholesky_decompose(cov);
    18:  }
   -------------------------------------------------

Ill-typed arguments supplied to infix operator *.

You probably need to coerce to vector.

Got it:
diag_matrix(vector x)
!!