## model "matrix-normal distribution"
## https://en.wikipedia.org/wiki/Matrix_normal_distribution
## Y ~ MN(M, U, V)
## Y is NxT; M is NxT; U is NxN; V is TxT
## here, M is assumed fixed and full of 0
## U is assumed equal to the identity
## V is "unconstrained", i.e. full of unknown parameters to estimate

## trick: vec(Y) ~ MVN(vec(M), V kx U) where "kx" is the Kronecker product

library(MASS)
library(rstan)

## simulate data
set.seed(1)
N <- 500
T <- 2
M <- matrix(data=0, nrow=N, ncol=T)
U <- diag(N)
V <- matrix(data=NA, nrow=T, ncol=T)
diag(V) <- c(4,8)
cor.1.2 <- 0.7
V[upper.tri(V)] <- c(cor.1.2 * sqrt(V[1,1] * V[2,2]))
V[lower.tri(V)] <- V[upper.tri(V)]
V
Y <- matrix(data=mvrnorm(n=1, mu=c(M), Sigma=V %x% U), nrow=N, ncol=T)
cov(Y[,1], Y[,2])
cor(Y[,1], Y[,2])

## check stan code below to compute the Kronecker product
Sigma <- matrix(data=NA, nrow=T*N, ncol=T*N)
for(ti in 1:T)
  for(tj in 1:T)
    for(ni in 1:N)
      for(nj in 1:N)
        Sigma[N*(ti-1)+ni, N*(tj-1)+nj] <- V[ti,tj] * U[ni,nj];
all(Sigma == V %x% U)

## write and compile model in Stan
model <- "data {
  int<lower=0> N;
  int<lower=0> T;
  vector[N*T] vecY;
  vector[N*T] vecM;
  matrix[N,N] U;
  matrix[T,T] S;
}
parameters {
  matrix[T,T] V;
}
model {
  matrix[T*N,T*N] Sigma;
  V ~ inv_wishart(10, S);
  for(ti in 1:T)
    for(tj in 1:T)
      for(ni in 1:N)
        for(nj in 1:N)
          Sigma[N*(ti-1)+ni, N*(tj-1)+nj] <- V[ti,tj] * U[ni,nj];
  vecY ~ multi_normal(vecM, Sigma);
}
"
model.file <- "model_matrix-normal.stan"
write(x=model, file=model.file)
rt <- stanc(file=model.file, model_name="Matrix-Normal")
sm <- stan_model(stanc_ret=rt, verbose=FALSE)

## fit
nb.iters <- 1000
warmup <- 100
thin <- 2
nb.chains <- 2
fit <- sampling(object=sm,
                data=list(N=N, T=T, vecY=c(Y), vecM=c(M), U=U, S=diag(T)),
                iter=nb.iters + warmup,
                warmup=warmup,
                thin=thin,
                chains=nb.chains)

## results
fit


library(MCMCpack)
riwish(v=T+1, S=diag(T))