Stanでマルチレベル（マニュアルコピペ）

参考：
https://mc-stan.org/docs/2_18/stan-users-guide/multivariate-hierarchical-priors-section.html
1.13 Multivariate Priors for Hierarchical Modelsより
高速なコードにはなっていないが，これで基本的にはOK。
添字が慣れたものではないが，読めばわかるような感じになっている。
cross classifiedモデルへの拡張や，適当なlink functionを仮定して従属変数が連続量以外でもOKな状況に持っていけるはず。

#
# Stanでマルチレベル lkj分布を使ったりなんだり
#
library("rstan")

# データ
library(haven)
mlmdata <- read_dta("https://stats.idre.ucla.edu/stat/examples/imm/imm10.dta")
write.csv(mlmdata,file="imm.csv",row.names = F)
dat <- read.csv(file = "imm.csv")
head(dat)

# データの用意
N <- nrow(dat)
x <- dat[,c("ses","homework","sex","white")]
K <- ncol(x)
J <- length(unique(dat$schid))
u <- dat[,c("meanses","public" )]
u <- u[!duplicated(u[,1]),]
L <- ncol(u)
y <- dat$math

uniq_sch_id <- unique(dat$schid)
jj <- rep(NA, N)
for (j in 1:J){
  jj[dat$schid==uniq_sch_id[j]] <- j
}
dat_stan <- list(N = N,
            K = K,
            J = J,
            L = L,
            jj = jj,
            x = x,
            u = u,
            y = y)
# int<lower=0> N;              // num individuals
# int<lower=1> K;              // num ind predictors
# int<lower=1> J;              // num groups
# int<lower=1> L;              // num group predictors
# int<lower=1,upper=J> jj[N];  // group for individual
# matrix[N, K] x;               // individual predictors
# row_vector[L] u[J];          // group predictors; uはJ×Lの行列でデータを指定する。
# vector[N] y;    
help(package="rstan")

cat("data {
  int<lower=0> N;              // num individuals
  int<lower=1> K;              // num ind predictors
  int<lower=1> J;              // num groups
  int<lower=1> L;              // num group predictors
  int<lower=1,upper=J> jj[N];  // group for individual
  matrix[N, K] x;               // individual predictors
  row_vector[L] u[J];          // group predictors
  vector[N] y;                 // outcomes
}
parameters {
  corr_matrix[K] Omega;        // prior correlation
  vector<lower=0>[K] tau;      // prior scale
  matrix[L, K] gamma;           // group coeffs
  vector[K] beta[J];           // indiv coeffs by group
  real<lower=0> sigma;         // prediction error scale
}
model {
  tau ~ cauchy(0, 2.5);
  Omega ~ lkj_corr(2);
  to_vector(gamma) ~ normal(0, 5);
  {
    row_vector[K] u_gamma[J];
    for (j in 1:J)
      u_gamma[j] = u[j] * gamma;
    beta ~ multi_normal(u_gamma, quad_form_diag(Omega, tau));
  }

//  for (n in 1:N)
//    y[n] ~ normal(x[n] * beta[jj[n]], sigma);
// Speed UP!!!
 {
    vector[N] x_beta_jj;
    for (n in 1:N)
      x_beta_jj[n] = x[n] * beta[jj[n]];
    y ~ normal(x_beta_jj, sigma);
  }
}
",file="multilevel.stan")

fit1 <- stan(
  file = "multilevel.stan",  # Stan program
  data = dat_stan,    # named list of data
  chains = 2,             # number of Markov chains
  warmup = 1000,          # number of warmup iterations per chain
  iter = 2000,            # total number of iterations per chain
  cores = 2,              # number of cores (using 2 just for the vignette)
  refresh = 1000         # show progress every 'refresh' iterations
)

stan_trace(fit1, pars=c("gamma"))
stan_trace(fit1, pars=c("Omega"))
stan_trace(fit1, pars=c("tau"))
stan_ac(fit1, pars=c("gamma"))
stan_ac(fit1, pars=c("Omega"))
stan_ac(fit1, pars=c("tau"))
stan_hist(fit1, pars=c("gamma"))
stan_hist(fit1, pars=c("Omega"))
stan_hist(fit1, pars=c("tau"))

print(fit1, pars=c("gamma", "Omega", "tau"), probs=c(.1,.5,.9))
plot(fit1, pars=c("gamma", "Omega"))

以下のチュートリアルも非常に参考になる
https://web.stanford.edu/class/psych201s/psych201s/papers/SorensenEtAl.pdf