multilogit.R

library(rstan)

model <- "                                                                                                                                                                             
data {                                                                                                                                                                                 
       int<lower=2> K;                                                                                                                                                                 
       int<lower=0> N;                                                                                                                                                                 
       int<lower=1> D;                                                                                                                                                                 
       int<lower=1,upper=K> y[N];                                                                                                                                                      
       vector[D] x[N];                                                                                                                                                                 
     }                                                                                                                                                                                 
     parameters {                                                                                                                                                                      
       matrix[K,D] beta;                                                                                                                                                               
     }                                                                                                                                                                                 
     model {                                                                                                                                                                           
       for (k in 1:K)                                                                                                                                                                  
         for (d in 1:D)                                                                                                                                                                
           beta[k,d] ~ normal(0,5);                                                                                                                                                    
       for (n in 1:N)                                                                                                                                                                  
         y[n] ~ categorical(softmax(beta * x[n]));                                                                                                                                     
}                                                                                                                                                                                      
"
#                                                                                                                                                                                      
N <- 9
D <- 2
x <- seq(1,N)+rnorm(N,0)
x <- cbind(x,x)
y <- floor(seq(0,N-1)/3)+1
K <- 3
fit <- stan(model_code = model,
            pars=c("beta"),
            data = list(N=N,
              D=D,K=K,
              y=y,x=x),
           chains = 1, iter = 100)

print(fit)
png("multilogit_cate.png")
traceplot(fit)
dev.off()