Where to find stick breaking transform and Jacobian code?

Modeling
2

I found the stick breaking transform code here and here. But I could use help importing these stan library functions into my stan program.

I’m attaching two stan programs that fit the following simple model:

Z = \sum_{i=1}^k X_i
log(X_i) \sim Normal(0,1)

One version (simplex_k.stan) uses the built-in simplex type. Another verison (manual_k.stan) implements the simplex manually but does so using the less efficient transform that @aaronjg proposed in this thread.

You can test the models like so:

simplex_k_model = cmdstanr::cmdstan_model("simplex_k.stan")

simplex_k_fit = simplex_k_model$sample(data=list(z=5,k=100),
                                  seed=123,
                                  chains=4,
                                  parallel_chains=4,
                                  iter_warmup=1000,
                                  iter_sampling=1000)

And similarly for manual_k. The inference is mathematically equivalent but I find that the simplex model (which uses the stick breaking transform under the hood) has about 5x more effective samples per second.

What I would like: another manual version that uses the stick breaking transform code in the stan codebase, so I can adapt it to situations where the built-in simplex type doesn’t work (e.g. ragged array of simplexes). Thank you!
simplex_k.stan (253 Bytes)
manual_k.stan (401 Bytes)