Ragged array of simplexes

Modeling
14

For @mjhajharia’s 2023 summer project with @Bob_Carpenter she started experiments on a bunch of different simplex parameterizations that I helped with. @Seth_Axen added a few new transforms, checked the math twice (by hand and with Jax), and put together a pipeline to evaluate the efficiency of the different transforms. The forthcoming preprint will have the derivations.

We included the log versions - the return vector is a log simplex - because you typically can stay on the log scale. In fact, the log_dirichlet should be more numerically stable and faster (assuming we write the derivatives in stan-math).

The transforms are all in _lp functions that you can test and should be able to reuse the form that @Christopher-Peterson suggested.

All the transforms live at:

If you want to stay on the log scale but use the dirichlet you can use this lpdf

real log_dirichlet_lpdf(vector log_theta, vector alpha) {
  int N = rows(log_theta);
  if (N != rows(alpha))
    reject("Input must contain same number of elements as alpha");
  return dot_product(alpha, log_theta) - log_theta[N]
         + lgamma(sum(alpha)) - sum(lgamma(alpha));
}
3 Likes