Recovering class index in mixture model using generated quantities

[nhuurre]

@emanudelfava

When generating ZZ using the categorical_rng function, I obtain that all individuals are assigned to the component 2, as its mixture weight is higher than 0.5.

That doesn’t sound right. categorical_rng does not just select the maximum probability category.
If \theta is (0.25,0.75) then one-in-four individuals are assigned to the first category.

Anyway, mixture weight theta is essentially a prior for category assignment and categorical_rng(theta) samples a new individual.
You can recover the categories for observed individuals by applying Bayes’ theorem and sampling the posterior assignment probability. This is straightforward; posterior is equal to normalized likelihood.

generated quantities{
  vector[N] log_lik; // log-likelihood for calculation of LOO
  int<lower=1> Z[N];          // group index

  for (n in 1:N) {
    vector[K] ll;
    for(k in 1:K) {
      ll[k] = log(theta[k]) + normal_lpdf(logod[n] | mu[k], sigma[k]);      
    }
    log_lik[n] = log_sum_exp(ll);
    Z[n] = categorical_rng(exp(ll - log_lik[n]));
  }
}