I would like to fit the following model using a multivariate normal, it is a Hidden Markov Model where we know the latent states. We have the following observations

where x_{i,s(i)} corresponds to observation i assigned to latent state s(i) \in \{1, \ldots,K\}.

We will model the x_{i,s(i)} as a multivariate normal centered on zero defined by the following (I don’t write out the entire expression so as to not detract from the main issue):

and

and f and g are functions of \sigma_{s(i)} which I don’t write down on purpose. In other words, consecutive measurements of the same state are correlated. We can also express this as

where \sum_{i,j}=\langle x_{i,s(i)} x_{j,s(j)} \rangle

I have been trying to implement this but it’s extremely inefficient and I would like to see if it could be improved. Here is what my current model looks like simplified (note that the functions f and g are placeholders for a more complicated expression in the stan model):

```
data {
int<lower=0> N;
int<lower=0> K;
int<lower=1,upper=K> state[N];
vector[N] x;
}
parameters {
real sigma[K];
}
transformed parameters {
cov_matrix[N] cov;
for (i in 1:N) {
for (j in 1:N) {
if (i==j) {
cov[i,j] = f(sigma[state[i]]);
} else if (abs(i-j)==1 && state[i]==state[j] ) {
cov[i,j] = g(sigma[state[i]]);
} else {
cov[i,j] = 0;
}
}
}
}
model {
sigma ~ gamma(3,0.2);
x ~ multi_normal(rep_vector(0,N),cov);
}
```

Any tips appreciated!