How to sample an integer from a range, given a simplex of probabilities

Is there a nice way to sample an integer from a range 1:K, given a simplex of the K probabilities? An example is for sampling which component of a finite mixture model resulted in an observation. Say the simplex is simplex[K] theta (for most applications this would be after reweighting by conditioning on the data). If K=2, i.e. a two-component mixture model, we can easily sample with int which_component = bernoulli_rng(theta[2]) + 1 (adding 1 for 1-based indexing instead of 0-based). However for K>2, all I can think of is this, which works but is ugly:

int which_component;
// Sample a K-vector with one element equal to 1, the rest equal to 0:
array[K] int counts_by_component = multinomial_rng(theta, 1);
// reproduce the behaviour of e.g. R's "which" function:
for (component in 1:K) {
  if (counts_by_component[component]) {
    which_component = component;
    break;
  } 
}

You’re looking for the categorical distribution.

int which_component = categorical_rng(theta);
1 Like