Manual ordering of two simplexes for emission probabilities in multinomial HMM

Hi, @Marty.

I encountered similar issue and I think I kind of figured out the answer.
Basically, you can re-calculate filtered state probabilities using log_omega, Gamma, rho, which are required for hmm_marginal function.

One thing to note here is that if you try to modify the hmm_gidden_state_prob function, the results produce NaNs at times, possibly because of under/overflow.

I needed to use some parts of the code that I originally used, to avoid NaN issue.

The implemented function is as below:

functions {
  
  // ref: https://github.com/luisdamiano/gsoc17-hhmm/blob/master/hmm/stan/hmm-multinom.stan
  matrix hmm_forward_prob(matrix log_omega, matrix Gamma, vector rho) {
    
    int n_state = rows(log_omega);
    int n_obs = cols(log_omega);
    matrix[n_state, n_obs] unnorm_log_alpha; 
    matrix[n_state, n_obs] alpha; // p(x_t=j|y_{1:t})
    array[n_state] real accumulator;
    
    // first alpha
    unnorm_log_alpha[,1] = log(rho) + log_omega[,1];
    
    // other alphas
    for (t in 2:n_obs) {
      for (j in 1:n_state) {
        for (i in 1:n_state) {
          accumulator[i] = unnorm_log_alpha[i, t-1] + log(Gamma[i, j]) + log_omega[j, t];
        }
        unnorm_log_alpha[j, t] = log_sum_exp(accumulator);
      }
    }

    // normalize alpha later for numerical stability
    for (t in 1:n_obs) {
      alpha[,t] = softmax(unnorm_log_alpha[,t]);
    }
    return alpha;
  } // !hmm_forward_prob
}

Also , I attached the sample code, for someone who may be interested in.

hmm-multinom-example.stan (2.1 KB)