Bayesian multi-level model for data with structured dispersion

This post provides code to model variance components in a multi-level setting. This is useful when variances are heteroskedastic and, we suspect, heteroskedasticity is driven by some observed variable(s).

Steps:

1. Load the data. For this demonstration I will simulate data in which 600 individual (\(i\)) are nested into 200 groups (\(j\)). \(Y\) is the outcome variable and I allow for random intercepts at the group level plus structured dispersion at both group-level and individual-level. In this particular case the variance components are functions of the same variable \(z\), but this predictor has opposite effects on each of them.

More formally, the DGP can be described as: \(Y = \mathbf{1}_{k} \otimes \mathbf{1}_{n}\mu + \mathbf{I}_{k} \otimes \mathbf{1}_{n} \mathbf{a} + \mathbf{I}_{k} \otimes \mathbf{I}_{n} \mathbf{e}\), where \(\mathbf{a}\) and \(\mathbf{e}\) are random effects, such that \(\mathbf{a} \sim N(\mathbf{0}, \mathbf{\sigma_{a}})\) and \(\mathbf{e} \sim N(\mathbf{0}, \mathbf{\sigma_{e}})\). Moreover, both variances are functions of \(z\): \(\mathbf{\sigma_{a}}=f(\mathbf{z})\) and \(\mathbf{\sigma_{e}}=g(\mathbf{z})\).

n=3   # Number of individuals by group
k=200 # Number of  groups
# Design matrices for multi-levels
One.k=matrix(rep(1,k))
One.n=matrix(rep(1,n))
I.k=diag(k)
I.n=diag(n)
mu=0
# variance predictor 
z = runif(k,1,5)               # indep variable in variance function at group level
Z = (I.k%x%One.n)%*%z          # indep variable in variance function at individual level 
# An arbitrary function for variance at the group level
THETA_a0 =  -5
THETA_a1  = 1.9
var.a = exp(THETA_a0 + THETA_a1*z + rnorm(k))
sd.a  = sqrt(var.a)
# An arbitrary function for variance at individual level (residual variance)
THETA_e0 =  5
THETA_e1  = -2.7
var.e = exp(THETA_e0 + THETA_e1*Z + rnorm(k*n))
sd.e  = sqrt(var.e)
# Random intercepts
a=rnorm(k,0,sd.a)
# Random error 
e=rnorm(k*n,0,sd.e) 
# Creates Y = mu + random intercept + random error
Y = (One.k%x%One.n)*mu + (I.k%x%One.n)%*%a + (I.k%x%I.n)%*%e   
# Generate indicators. To be used later. 
group=matrix(rep(1:k))  # vector with group indicator for each group
Group=(I.k%x%One.n)%*%group # vector with Group indicator for each individual

A simple scatterplot of \(Y\) and \(z\) should reveal variance structure.

2. Load rstan to perform Bayesian estimation via Halmiltonian Monte Carlo. Put the data into a list.

library("rstan")
Data_l1 <- data.frame(Y,Group,Z) # data individual-level
Data_l2 <- data.frame(group,z)   # data group-level 
data    <- list(N=n*k,J=k,y=Data_l1$Y,G=Data_l1$Group,Z=Data_l1$Z, g=Data_l2$group,z=Data_l2$z)

3. Write the model in stan.

structured_dispersion <- '
data {
  int<lower=0> J;                   // number of groups 
  int<lower=0> N;                   // number of individuals
  int<lower=1, upper=J>  G[N];      // map individuals to groups
  vector[N] y;                      // outcome variable
  vector[J] z;                      // group level dispersion predictor
  vector[N] Z;                      // individual level dispersion predictor
}
parameters {
  vector[J] a;                           // Random Intercept           
  
  real theta_a0;                         // Intercept model for variance of random intercept
  real theta_e0;                         // Intercept model for variance of residual error
  
  real theta_a1;                         // Slope model for variance of random intercept
  real theta_e1;                         // Slope model for variance of residual error
  vector<lower=0>[J] sigma2_a;           // heterogeneous var for random intercept 
  vector<lower=0>[N] sigma2_e;           // heterogeneous var residual error
}
transformed parameters {
  
  vector[N] y_mu;                        // Random effects
  vector[J] logsigma2_a_mu;              // mean of log-variance for random intercept 
  vector[J] sigma2_a_mu;                 // mean of variance for random intercept 
  vector[J] sigma_a;                     // var for random intercept 
  vector[N] logsigma2_e_mu;              // mean of log-variance residual error
  vector[N] sigma2_e_mu;                 // mean of var residual error
  vector[N] sigma_e;                     // var residual error
  for (i in 1:J) {                       
                                         // Assigns a different sd to each random intercept
    logsigma2_a_mu[i] = theta_a0 + theta_a1*z[i];
    sigma2_a_mu[i]    = exp(logsigma2_a_mu[i]);
    sigma_a[i]        = sqrt(sigma2_a_mu[i]);
  }
  for (i in 1:N) {                    
                                          // Assigns a different sd to each residual error   
    logsigma2_e_mu[i] = theta_e0 + theta_e1*Z[i];
    sigma2_e_mu[i]    = exp(logsigma2_e_mu[i]);
    sigma_e[i]        = sqrt(sigma2_e_mu[i]);
                                          // Assigns a different predicted mean to ech individual depending of group of belonging
    y_mu[i] = a[G[i]];                  
  }
}
model {
  // Priors 
  theta_a0    ~ cauchy(0,2.5);
  theta_e0    ~ cauchy(0,2.5);
  theta_a1    ~ cauchy(0,2.5);
  theta_e1    ~ cauchy(0,2.5);
  // Variances
  sigma2_a ~ lognormal(sigma2_a_mu,0.25);
  sigma2_e ~ lognormal(sigma2_e_mu,0.25);
  // Random intercept 
  a ~ normal(0, sigma_a);
  // Outcome variable
  y ~ normal(y_mu, sigma_e);
}
'

4. Fit the model and compare the estimates to the parameters used in the DGP.

mymodel <- stan(model_code=structured_dispersion, model_name="structured_dispersion", 
iter=2000, pars=c("theta_e0", "theta_e1","theta_a0", "theta_a1"), data=data, chains=3)

print(mymodel)

Inference for Stan model: structured_dispersion.
3 chains, each with iter=2000; warmup=1000; thin=1; 
post-warmup draws per chain=1000, total post-warmup draws=3000.

           mean se_mean    sd    2.5%    25%    50%    75%  97.5% n_eff Rhat
theta_e0   5.63    0.06  0.19    5.31   5.46   5.63   5.73   6.04    10 1.17
theta_e1  -2.79    0.02  0.06   -2.91  -2.82  -2.79  -2.75  -2.68    12 1.13
theta_a0  -3.87    0.20  0.46   -4.69  -4.24  -3.83  -3.53  -3.07     5 1.80
theta_a1   1.73    0.06  0.12    1.53   1.64   1.73   1.81   1.96     4 1.96
lp__     -63.59    3.02 23.12 -110.10 -78.88 -63.47 -47.50 -20.06    59 1.04

Samples were drawn using NUTS(diag_e) at Mon Dec  3 17:15:39 2018.
For each parameter, n_eff is a crude measure of effective sample size,
and Rhat is the potential scale reduction factor on split chains (at 
convergence, Rhat=1).

5. Uset he functionlaunch_shinystan from package rstanarm to get visual and numerical summaries of model parameters and convergence diagnostics.