Multinomial regression model

multinomial_regression.stan

/*

This model is based on the multinomial model in Bayesian Data Analysis 3, p426.
I used the multinomial regression example from the Stan website
https://mc-stan.org/docs/2_27/stan-users-guide/multi-logit-section.html and
then modified it to match our model.

For this model to work, you need an up to date version of `rstan`, at least
2.26.2.  Follow the instructions here: https://mc-stan.org/r-packages/ to
upgrade your `rstan` if you have an old version. You can check which version
you have by running:

library(rstan)
stan_version()

in R.

*/
data {
  /* Hyperparameters */
  real A;
  real B;

  /* Dimensions */
  int<lower = 0> K;
  int<lower = 0> N;
  int<lower = 0> D;

  /* Data */
  int<lower = 0> y[N, K];
  matrix[N, D] x;
}
parameters {
  matrix[D, K] beta;
  real mu[D];
  real<lower = 0> sigma2[K];
}
model {
  matrix[N, K] x_beta = x * beta;

  for (i in 1:D) {
    mu[i] ~ normal(0, 10);
    sigma2[i] ~ gamma(A, B);
    for (j in 1:K) {
      beta[i, j] ~ normal(mu[i], sigma2[j]);
      /* beta[, j] ~ multi_normal(mu, sigma2[j] * identity_matrix(D)); */
      /* beta[, j] ~ normal(0, sigma2[j]); */
    }
  }

  for (i in 1:N) {
    y[i,] ~ multinomial(softmax(x_beta[i,]'));
  }
}