---
title: "Using a LKJ prior in Stan"
date: "2016-03-07"
output: html_document
---
```{r setup, include=FALSE}
knitr::opts_chunk$set(echo = TRUE, message=FALSE)
```
There are two ways to use a LKJ prior distribution for a correlation matrix in STAN. The first one assigns the distribution on the correlation matrix, whereas the second one assigns the distribution on the lower Cholesky factor of the correlation matrix.
I am going to show an example for a trivariate normal sample with a fixed mean:
$$
y_i \sim_{\text{iid}} {\cal N}\left(
\begin{pmatrix}
0 \\ 0 \\ 0
\end{pmatrix}, \Sigma\right).
$$
Recall the relation between the covariance matrix and the correlation matrix:
$$
\begin{pmatrix}
\sigma_{1}^2 & \sigma_{12} & \sigma_{13} \\
\sigma_{12} & \sigma_2^2 & \sigma_{23} \\
\sigma_{13} & \sigma_{23} & \sigma_3^2
\end{pmatrix}
=
\begin{pmatrix}
\sigma_{1}^2 & 0 & 0 \\
0 & \sigma_2^2 & 0 \\
0 & 0 & \sigma_3^2
\end{pmatrix}
\Omega
\begin{pmatrix}
\sigma_{1}^2 & 0 & 0 \\
0 & \sigma_2^2 & 0 \\
0 & 0 & \sigma_3^2
\end{pmatrix}
$$
This operation is performed in Stan by the function `quad_form_diag`.
I do not assume the reader familiar with Stan or the `rstan` package, so I will comment each step.
First I load the `rstan` package with the usual options:
```{r}
library(rstan)
rstan_options(auto_write = TRUE)
options(mc.cores = parallel::detectCores())
```
I will run Stan on these simulated data:
```{r}
set.seed(666)
Omega <- rbind(
c(1, 0.3, 0.2),
c(0.3, 1, 0.1),
c(0.2, 0.1, 1)
)
sigma <- c(1, 2, 3)
Sigma <- diag(sigma) %*% Omega %*% diag(sigma)
N <- 100
y <- mvtnorm::rmvnorm(N, c(0,0,0), Sigma)
```
## Prior on the correlation matrix
Below is the Stan code for the Bayesian model assigning a LKJ prior on the correlation matrix $\Omega$.
I use the LKJ distribution with shape parameter $1$, which is the uniform distribution on the space of correlation matrices.
```{r}
stancode1 <- 'data {
int<lower=1> N; // number of observations
int<lower=1> J; // dimension of observations
vector[J] y[N]; // observations
vector[J] Zero; // a vector of Zeros (fixed means of observations)
}
parameters {
corr_matrix[J] Omega;
vector<lower=0>[J] sigma;
}
transformed parameters {
cov_matrix[J] Sigma;
Sigma <- quad_form_diag(Omega, sigma);
}
model {
y ~ multi_normal(Zero,Sigma); // sampling distribution of the observations
sigma ~ cauchy(0, 5); // prior on the standard deviations
Omega ~ lkj_corr(1); // LKJ prior on the correlation matrix
}'
```
The `stan_model` function of the `rstan` package runs the Stan compilation of the model.
```{r compilation1, eval=TRUE, cache=TRUE}
stanmodel1 <- stan_model(model_code = stancode1, model_name="stanmodel1")
```
Note that this function only takes the Stan code as input. Once the model is compiled, the `stanmodel` object can be used to run the model on different datasets with the `sampling` function.
The data must be passed to the `sampling` function into a list:
```{r}
standata <- list(J = ncol(y), N=N, y = y, Zero=rep(0, ncol(y)))
```
The algorithms used by Stan to generate the posterior distributions require initial values of the parameters. One can let the `sampling` function generate random initial values, or pass them in its `init` argument. I prefer to give my initial values. More precisely, one must pass to the `init` argument a function that returns the initial values in a list (see `?sampling`).
Here is the way I use to generate initial values. Firstly, I write a function that returns some estimates of the parameters, taking the observations as input and allowing to randomly perturb these observations:
```{r}
estimates <- function(y, perturb=FALSE){
if(perturb) y <- y + rnorm(length(y), 0, 1)
sigma <- sqrt(diag(var(y)))
Omega <- cor(y)
return(list(sigma=sigma, Omega=Omega))
}
```
I run Stan with several chains, for instance four chains. The function passed to the `init` argument of the `sampling` function takes an optional argument `chain_id`. For the first chain, I use the estimates calculated from the original data as initial values, and for the other chains I use the estimates calculated from the disturbed original data:
```{r}
inits <- function(chain_id){
values <- estimates(standata$y, perturb = chain_id > 1)
return(values)
}
```
Now we are ready to run Stan:
```{r sampling1}
samples1 <- sampling(stanmodel1, data = standata, init=inits,
iter = 10000, warmup = 1000, chains = 4)
```
Some numerical problems occur but they are benign. It is not abnormal to get some messages such as
```
validate transformed params: Sigma is not positive definite
validate transformed params: Sigma is not symmetric. Sigma[1,2] = -nan, but Sigma[2,1] = -nan
Exception thrown at line 23: lkj_corr_log: y is not positive definite
validate transformed params: Sigma is not symmetric. Sigma[1,2] = inf, but Sigma[2,1] = inf
```
These problems will not occur with the LKJ prior on the Cholesky factor.
I like to use the `coda` package for output analysis.
This is the way I use to store the samples in a `coda` object:
```{r}
library(coda)
codasamples1 <- do.call(mcmc.list,
plyr::alply(rstan::extract(samples1,
pars=c("sigma", "Omega[1,2]", "Omega[1,3]", "Omega[2,3]"),
permuted=FALSE), 2, mcmc))
```
## Prior on the Cholesky factor
The correlation matrix $\Omega$ has a Cholesky factorization $\Omega = LL'$ where $L$ is a lower triangular matrix.
Instead of assigning a prior distribution on $\Omega$, on can assign a prior dsitribution on $L$. By this way, the numerical problems encountered with the previous way are overcome, and this way is also better for a speed perspective.
```{r}
stancode2 <- 'data {
int<lower=1> N; // number of observations
int<lower=1> J; // dimension of observations
vector[J] y[N]; // observations
vector[J] Zero; // a vector of Zeros (fixed means of observations)
}
parameters {
cholesky_factor_corr[J] Lcorr;
vector<lower=0>[J] sigma;
}
model {
y ~ multi_normal_cholesky(Zero, diag_pre_multiply(sigma, Lcorr));
sigma ~ cauchy(0, 5);
Lcorr ~ lkj_corr_cholesky(1);
}
generated quantities {
matrix[J,J] Omega;
matrix[J,J] Sigma;
Omega <- multiply_lower_tri_self_transpose(Lcorr);
Sigma <- quad_form_diag(Omega, sigma);
}'
```
Note the `generated quantities` block as compared to the `transformed quantities` block in the first code. The objects declared in the `transformed parameters` block are intended to be used in the likelihood of the data, whereas the objects declared in the `generated quantities` block are not.
Now we only have to change the `estimates` function:
```{r}
estimates <- function(y, perturb=FALSE){
if(perturb) y <- y + rnorm(length(y), 0, 1)
Lcorr <- t(chol(cor(y)))
sigma <- sqrt(diag(var(y)))
return(list(Lcorr=Lcorr, sigma=sigma))
}
```
Then compile and run as before:
```{r compilation2, eval=TRUE, cache=TRUE}
stanmodel2 <- stan_model(model_code = stancode2, model_name="stanmodel2")
samples2 <- sampling(stanmodel2, data = standata, init=inits,
iter = 10000, warmup = 1000, chains = 4)
```
```{r}
library(coda)
codasamples2 <- do.call(mcmc.list,
plyr::alply(rstan::extract(samples2,
pars=c("sigma", "Omega[1,2]", "Omega[1,3]", "Omega[2,3]"),
permuted=FALSE), 2, mcmc))
```
## Comparison of results
Results are almost identical:
```{r}
summary(codasamples1)
summary(codasamples2)
```