recreate the stan code from https://arxiv.org/pdf/1808.06399.pdf

dirichletreg-stan-ex.R

# stan code from https://arxiv.org/pdf/1808.06399.pdf
library("DirichletReg")
Bld <- BloodSamples
Bld <- na.omit(Bld)
Bld$Smp <- DR_data(Bld[, 1:4])

stan_code <- '
data {
  int<lower=1> N;  // total number of observations
  int<lower=2> ncolY; // number of categories
  int<lower=2> ncolX; // number of predictor levels
  matrix[N,ncolX] X; // predictor design matrix
  matrix[N,ncolY] Y; // response variable
  real sd_prior; // Prior standard deviation
}
parameters {
  matrix[ncolY-1,ncolX] beta_raw; // coefficients (raw)
  real theta;
}
transformed parameters{
  real exptheta = exp(theta);
  matrix[ncolY,ncolX] beta; // coefficients
  for (l in 1:ncolX) {
    beta[ncolY,l] = 0.0;
  }
  for (k in 1:(ncolY-1)) {
    for (l in 1:ncolX) {
      beta[k,l] = beta_raw[k,l];
    }
  } }
model {
  // prior:
  theta ~ normal(0,sd_prior);
  for (k in 1:(ncolY-1)) {
    for (l in 1:ncolX) {
      beta_raw[k,l] ~ normal(0,sd_prior);
    }
  }
  // likelihood
  for (n in 1:N) {
    vector[ncolY] logits;
    for (m in 1:ncolY){
      logits[m] = X[n,] * transpose(beta[m,]);
    }
    transpose(Y[n,]) ~ dirichlet(softmax(logits) * exptheta);
  }
}'

library(rstan)
prg <- stan_model(model_code = stan_code)
X <- as.matrix(model.matrix(lm(Albumin ~ Disease, data = Bld)))
X <- matrix(nrow = nrow(X), ncol = ncol(X), data = as.numeric(X))
Y <- Bld$Smp
D <- list(N = nrow(Y), ncolY = ncol(Y), ncolX = ncol(X), 
          X = X, Y = Y, sd_prior = 1)
fit1 <- sampling(prg, data = D, chains = 4, iter = 2000, cores = 4,
                 control = list(adapt_delta = 0.95, max_treedepth = 20),
                 refresh = 100)

B <- extract(fit1)$beta

simplex <- function(x){
  exp(x)/sum(exp(x))
}

# plot stan
my_colors <- scales::hue_pal()(4)
layout(matrix(1:2, ncol = 2))
plot(1:4, Bld[1, 1:4], ylim = c(0, 0.6), type = "n", xaxt = "n", las = 1,
     xlab = "", ylab = "Proportion", main = "Disease A", xlim = c(0.6, 4.4))
abline(h = seq(0, 0.6, by = 0.1), col = "grey", lty = 3)
axis(1, at = 1:4, labels = names(Bld)[1:4], las = 2)
aux <- t(apply(B[, , 1], MAR = 1, FUN = simplex))
apply(subset(Bld, Disease == "A")[, 1:4], MAR = 1, FUN = points, pch = 16,
      col = "grey")
lines(apply(subset(Bld, Disease == "A")[, 1:4], MAR = 2, FUN = mean),
      type = "b", pch = 16, cex = 1.2, lwd = 2)
lines(apply(aux, MAR = 2, FUN = quantile, prob = 0.975), type = "b", pch = 4,
      lty = 2, col = my_colors[1])
lines(apply(aux, MAR = 2, FUN = quantile, prob = 0.025), type = "b", pch = 4,
      lty = 2, col = my_colors[1])
lines(apply(aux, MAR = 2, FUN = mean), lwd = 2, col = my_colors[1], type = "b",
      pch = 16)

plot(1:4, Bld[1, 1:4], ylim = c(0, 0.6), type = "n", xaxt = "n", las = 1,
     xlab = "", ylab = "Proportion", main = "Disease B", xlim = c(0.6, 4.4))
abline(h = seq(0, 0.6, by = 0.1), col = "grey", lty = 3)
axis(1, at = 1:4, labels = names(Bld)[1:4], las = 2)
aux <- t(apply(B[, , 1] + B[, , 2], MAR = 1, FUN = simplex))
apply(subset(Bld, Disease == "B")[, 1:4], MAR = 1, FUN = points, pch = 16,
      col = "grey")
lines(apply(subset(Bld, Disease == "B")[, 1:4], MAR = 2, FUN = mean),
      type = "b", pch = 16, cex = 1.2, lwd = 2)
lines(apply(aux, MAR = 2, FUN = quantile, prob = 0.975), type = "b", pch = 4,
      lty = 2, col = my_colors[2])
lines(apply(aux, MAR = 2, FUN = quantile, prob = 0.025), type = "b", pch = 4,
      lty = 2, col = my_colors[2])
lines(apply(aux, MAR = 2, FUN = mean), lwd = 2, col = my_colors[2], type = "b",
      pch = 16)
layout(1)