Ragged array of simplexes

Modeling
10

One option (for the left hand side) is to use a softmax transform for pi_slice. The last post in this thread has a function that does the transformation and automatically performs the Jacobian adjustment.

In theory, you could then do something like:

functions{
vector simplex_constrain_softmax_lp(vector v) {
     int K = size(v) + 1;
     vector[K] v0 = append_row(0, v);
     // Jacobian
     target += sum(v) - K * log_sum_exp(v0);
     return softmax(v0);
  }
}
data {
   int<lower=0> N; // number of states
   int<lower=0> M; // number of groups
   array[M] int sizes; // size of groups
   vector[N] alpha0; // hyperp for state probabilities
}
parameters {
  vector[N-M] pi_unconstrained; 
}
model {
 int  pos_pi = 1;
  int pos_alpha = 1;
 for(k in 1:M){
  simplex_constrain_softmax_lp(segment(pi_unconstrained, pos_pi, sizes[k] - 1)) 
   ~ dirichlet(segment(alpha0,pos_alpha, sizes[k]));
  pos_alpha += sizes[k];
  pos_pi += sizes[k]-1;
 }

}

It may require some slight adjustments if you want to actually use your adjusted simplex parameters for something (perhaps defining them in the transformed parameters and re-using them elsewhere).

1 Like