1 Introduction

Disclaimer: I’m not a statistician (if you spot any errors please leave a message here or contact me). I’m not endorsing or advocating the use of \(R^2\) to evaluate model fit for linear mixed models. This is just a blog post to show how it can be done with JAGS. For further info on the topic please refer to the papers cited in this post.

Nakagawa and Schielzeth (2013) proposed a method to quantify the proportion of variance explained in a linear mixed model and provided code to do so with the results of models fitted with lme4. This post shows how to do the same for a random intercept mixed model fitted with JAGS. A similar blog post that helped me understand how to do this is available here. It should be noted that the method proposed by Nakagawa and Schielzeth (2013) that we’ll cover here is limited to random intercept models. Johnson (2014) proposed an extension for models that include random slopes (see also Nakagawa, Schielzeth, and Johnson 2017).

Nakagawa and Schielzeth (2013) defined two \(R^2\) measures for linear mixed models. The first one, marginal \(R^2\) (\(R_{m}^2\)), represents the proportion of variance explained by the fixed effects alone over the overall variance. The second one, conditional \(R^2\) (\(R_{c}^2\)), represents the proportion of variance explained by both fixed and random effects over the overall variance. Suppose you have a model with one or more fixed effects, and a single random effect, such as defined by Equations (1.1), (1.2), and (1.3).

\[\begin{equation} y_{ij} = \beta_0 + \sum_{h=1}^p \beta_{h}x_{hij} + \alpha_j + e_{ij} \tag{1.1} \end{equation}\] \[\begin{equation} a_j \sim \mathcal{N}(0, \sigma_\alpha^2) \tag{1.2} \end{equation}\] \[\begin{equation} e_{ij} \sim \mathcal{N}(0, \sigma_e^2) \tag{1.3} \end{equation}\] In this model the \(\beta_h\)s are the slopes of the fixed effects, and \(\alpha_j\) represents the \(j^{th}\) random intercept. Equations (1.4), and (1.5) give the marginal and conditional \(R^2\) for this model, respectively. \[\begin{equation} R^2_{(m)} = \frac{\sigma_f^2}{\sigma_f^2 + \sigma_\alpha^2 + \sigma_e^2} \tag{1.4} \end{equation}\] \[\begin{equation} R^2_{(c)} = \frac{\sigma_f^2 + \sigma_\alpha^2}{\sigma_f^2 + \sigma_\alpha^2 + \sigma_e^2} \tag{1.5} \end{equation}\] In these equations \(\sigma_f^2\) is the variance of the fixed effects. This is given in Equation (1.6). Another way to think about this is as the variance of the values that would be predicted by the fixed effects alone. \[\begin{equation} \sigma_f^2 = var\left(\sum_{h=1}^p \beta_{h}x_{hij}\right) \tag{1.6} \end{equation}\]

1.1 Calculating \(R^2\) from lmer output

We’ll first have a look at how marginal and conditional \(R^2\) are calculated for a model fitted with the lmer function from the lme4 package. There is actually a function in the MuMIn package to automatically calculate \(R_{m}^2\) and \(R_{c}^2\) from lmer output, but it is instructive to look at the step by step calculations. We first load some packages and source some code that we’ll need. You can download the DBDA2E-utilities_SC.R file here. It is a customized version of a file from Krushcke’s code accompanying the DBDA2E book (Kruschke 2014).

library(lme4)
library(MuMIn)
library(ggplot2)
library(dplyr)
library(runjags)
library(coda)
source("DBDA2E-utilities_SC.R")

We’ll use the “Orthodont” data from the nlme package to run our examples. The choice was based mainly on convenience and I don’t claim the analysis shown below is the best one for this dataset. This dataset contains an orthodontic measurement (distance from the pituitary to the pterygomaxillary fissure in mm) on a number of young subject. For each subject there are four measurements taken each at a different age. The code below loads the data and plots them:

data(Orthodont, package = "nlme")
dat = Orthodont
p = ggplot(dat, aes(age, distance))+geom_point()+facet_wrap(~Subject)
p = p + xlab("Age (Years)") + ylab("Distance (mm)") + theme_bw()
p

As a matter of personal preference the code below sets the contrasts to contr.sum. We’ll also center the variables, and for consistency with the JAGS code that will be shown in the next section, we’ll create and use indicator variables for the Sex and Sex by age interaction factors.

options(contrasts=c('contr.sum','contr.poly'))
dat$age = dat$age - mean(dat$age)
dat$sexNum = ifelse(dat$Sex=="Female", -1, 1)
dat$sexXage = dat$sexNum*dat$age
dat$distance = dat$distance - mean(dat$distance)

Below we fit the model including a random intercept for Subject. The function r.squaredGLMM comes from the MuMIn package and gives us the marginal and conditional \(R^2\) values:

lmeFit = lmer(distance ~ sexNum + age + sexXage + (1 | Subject), data = dat)
r.squaredGLMM(lmeFit)
##       R2m       R2c 
## 0.4098416 0.7827266

Although the r.squaredGLMM function provides an easy way to compute marginal and conditional \(R^2\), to understand how the computations are done we’ll look at some code (adapted from the supplementary material in Nakagawa and Schielzeth (2013)).

lmeFixedEff = fixef(lmeFit)[2] * lmeFit@pp$X[, 2] +
    fixef(lmeFit)[3] * lmeFit@pp$X[, 3] +
    fixef(lmeFit)[4] * lmeFit@pp$X[, 4]
lmeVarFixed = var(lmeFixedEff)
lmeVarRandom = VarCorr(lmeFit)$Subject[1]
lmeVarResidual = attr(VarCorr(lmeFit), "sc")^2
lmeMarginalR2 = lmeVarFixed/(lmeVarFixed + lmeVarRandom + lmeVarResidual)
lmeConditionalR2 = (lmeVarFixed + lmeVarRandom)/(lmeVarFixed +
                                                 lmeVarRandom + lmeVarResidual)

in the first line the function fixef extracts the estimates of the fixed effects, that is the \(\beta\)s. Each \(\beta\) is multiplied with its corresponding vector of predictors (obtained via lmeFit@pp$X[, h]). The variance of these values is then computed in the second line. The third line extracts the random effect variance for the Subject effect, while the fourth extracts the residual error variance. The last two lines compute the marginal and conditional \(R^2\).

To check the consistency of the results between the lmer model shown in this section and the JAGS model that will be shown in the next section we can also print some summaries for the current model:

summary(lmeFit)   
## Linear mixed model fit by REML ['lmerMod']
## Formula: distance ~ sexNum + age + sexXage + (1 | Subject)
##    Data: dat
## 
## REML criterion at convergence: 436.5
## 
## Scaled residuals: 
##     Min      1Q  Median      3Q     Max 
## -3.5980 -0.4546  0.0158  0.5024  3.6862 
## 
## Random effects:
##  Groups   Name        Variance Std.Dev.
##  Subject  (Intercept) 3.299    1.816   
##  Residual             1.922    1.386   
## Number of obs: 108, groups:  Subject, 27
## 
## Fixed effects:
##             Estimate Std. Error t value
## (Intercept) -0.21491    0.38071  -0.564
## sexNum       1.16051    0.38071   3.048
## age          0.63196    0.06071  10.409
## sexXage      0.15241    0.06071   2.511
## 
## Correlation of Fixed Effects:
##         (Intr) sexNum age   
## sexNum  -0.185              
## age      0.000  0.000       
## sexXage  0.000  0.000 -0.185
print(lmeVarFixed)
## [1] 3.625561
print(lmeVarRandom)
## [1] 3.298634
print(lmeVarResidual)
## [1] 1.922055

1.2 Calculating \(R^2\) from JAGS

The JAGS model code is given below:

modelString = "
# Standardize the data:
data {
  ym <- mean(y)
  ysd <- sd(y)
  for ( i in 1:Ntotal ){
    zy[i] <- (y[i] - ym) / ysd
  }
  for (j in 1:Nx) {
    xm[j]  <- mean(x[,j])
    xsd[j] <-   sd(x[,j])
    for ( i in 1:Ntotal ) {
      zx[i,j] <- (x[i,j] - xm[j]) / xsd[j]
    }
  }
}
# Specify the model for standardized data:
model{
  for (i in 1:Ntotal){
    zy[i] ~ dnorm(zbeta0 + sum(zbeta[1:Nx] * zx[i,1:Nx]) + subjEff[subjID[i]], 1/zsigma^2)
  }
  # Priors vague on standardized scale:
  zbeta0 ~ dnorm(0 , 1/10^2)  
  for (j in 1:Nx){
    zbeta[j] ~ dnorm(0 , 1/10^2)
  }
  zsigma ~ dunif(0.00001, 10)
  # Transform to original scale:
  beta[1:Nx] <- ( zbeta[1:Nx] / xsd[1:Nx] )*ysd
  beta0 <- zbeta0*ysd  + ym - sum( zbeta[1:Nx] * xm[1:Nx] / xsd[1:Nx] )*ysd
  sigma <- zsigma*ysd

  for (s in 2:nSubj){
    subjEff[s] ~ dnorm(0, 1/zSubjEffSigma^2)
  }
  subjEff[1] <- -sum(subjEff[2:nSubj]) #sum-to-zero constraint
  zSubjEffSigma ~ dunif(0.00001, 10)
  subjEffSigma <- zSubjEffSigma*ysd

  for (j in 1:Ntotal){
     yPredFixed[j] <-  sum(beta[1:Nx] * x[j,1:Nx])
  }

  varFixed <- (sd(yPredFixed))^2
  varResidual <- sigma^2 # get the variance of residuals
  varRandom <- subjEffSigma^2  # get the variance of random plot effect
  # calculate marginal R^2
  marginalR2 <- varFixed / (varFixed + varRandom + varResidual)
  # calculate conditional R^2
  conditionalR2 <- (varRandom + varFixed) / (varFixed + varRandom + varResidual) 
}
"

The variance of the random effects and the residual variance are directly estimated by the model. The only quantity that we need to compute is the variance of the fixed effect components. The values predicted from the fixed effects alone yPredFixed are calculated in the last for loop. The following line computes their variance. The R code needed to run the JAGS model is given below:

yName = "distance"
xName = c("sexNum", "age", "sexXage")
subjIDName = "Subject"
subjID = as.numeric(as.factor(dat[,subjIDName]))
x = as.matrix(dat[,xName],ncol=length(xName))
y = dat[, yName]

# Specify the data in a list, for later shipment to JAGS:
dataList = list(
  x = x,
  y = y ,
  subjID=subjID,
  nSubj = length(unique(subjID)),
  Nx = dim(x)[2] ,
  Ntotal = dim(x)[1])
  
parameters = c( "beta0" ,  "beta" ,  "sigma", 
               "zbeta0" , "zbeta" , "zsigma",
               "subjEff", "subjEffSigma", "zSubjEffSigma",
               "varFixed", "varResidual", "varRandom",
               "marginalR2", "conditionalR2")

#function to set random seed for each chain
initsFunction = function(chain){
    .RNG.seed = c(11, 22, 33, 44)[chain]
    .RNG.name = c("base::Super-Duper",
                  "base::Wichmann-Hill",
                  "base::Marsaglia-Multicarry",
                  "base::Mersenne-Twister")[chain]
    return(list(.RNG.seed=.RNG.seed,
                .RNG.name=.RNG.name))
}

adaptSteps = 500  # Number of steps to "tune" the samplers
burnInSteps = 1000
numSavedSteps = 30000
thinSteps = 10
nChains = 4
runJagsOut = run.jags(method="parallel",
                      model=modelString, 
                      monitor=parameters , 
                      data=dataList ,  
                      inits=initsFunction, 
                      n.chains=nChains ,
                      adapt=adaptSteps ,
                      burnin=burnInSteps , 
                      sample=ceiling(numSavedSteps/nChains) ,
                      thin=thinSteps ,
                      summarise=FALSE ,
                      plots=FALSE )
mcmcCoda = as.mcmc.list(runJagsOut)
mcmcMat = as.matrix(mcmcCoda, chains=T)

For brevity we’ll skip on the MCMC diagnostics and the other parameters of the model and we’ll plot directly the posterior distributions for the marginal and conditional \(R^2\). The plotPost function comes from Krushcke’s code accompanying the DBDA2E book (Kruschke 2014).

par(mfrow=c(1,2))
plotPost(mcmcMat[,"varFixed"]/
         (mcmcMat[,"varFixed"]+mcmcMat[,"varRandom"]+mcmcMat[,"varResidual"]),
         main=expression(paste("Marginal ", R^2)), col="gray80")
plotPost((mcmcMat[,"varFixed"]+mcmcMat[,"varRandom"])/
         (mcmcMat[,"varFixed"]+mcmcMat[,"varRandom"]+mcmcMat[,"varResidual"]),
         main=expression(paste("Conditional ", R^2)), col="gray80")

The values are very close to those computed from the lmer output. The code below can be used to plot diagnostics for the parameters of the model (these will be saved as pdf files in a folder called diagnostics).

dir.create("diagnostics/")    
parameterNames = varnames(mcmcCoda) # get all parameter names
for (parName in parameterNames){
    if (substr(parName, 1, 5) == "beta["){
        thisIdx = as.numeric(strsplit(strsplit(parName, "\\[")[[1]][2], "]")[[1]][1])
        pdf(paste0("diagnostics/", parName, "_", xName[thisIdx], ".pdf"))
    } else {
        pdf(paste0("diagnostics/", parName, ".pdf"))
    }
    diagMCMC(codaObject=mcmcCoda , parName=parName)
    dev.off()
}

We can also plot some other model parameters to check consistencey with lmer results.

par(mfrow=c(2,2))
plotPost(mcmcMat[,"beta0"], main="Intercept", col="gray80")
plotPost(mcmcMat[,"beta[1]"], main="Sex", col="gray80")
plotPost(mcmcMat[,"beta[2]"], main="Age", col="gray80")
plotPost(mcmcMat[,"beta[3]"], main="Sex X Age", col="gray80")

par(mfrow=c(1,3))
plotPost(mcmcMat[,"varFixed"], main="Fixed effects variance", col="gray80")
plotPost(mcmcMat[,"varRandom"], main="Random effects variance", col="gray80")
plotPost(mcmcMat[,"varResidual"], main="Residual error variance", col="gray80")

References

Johnson, Paul C. D. 2014. “Extension of Nakagawa & Schielzeth’s \(R^2_{GLMM}\) to random slopes models.” Methods in Ecology and Evolution, no. 5: 944–46. doi:10.1111/2041-210X.12225.

Kruschke, John K. 2014. Doing Bayesian Data Analysis: A Tutorial with R, JAGS, and Stan. 2nd ed. London: Academid Press / Elsevier.

Nakagawa, Shinichi, and Holger Schielzeth. 2013. “A general and simple method for obtaining \(R^2\) from generalized linear mixed-effects models.” Methods in Ecology and Evolution 4 (2): 133–42. doi:10.1111/j.2041-210x.2012.00261.x.

Nakagawa, Shinichi, Holger Schielzeth, and Paul C. D. Johnson. 2017. “The coefficient of determination R 2 and intra-class correlation coefficient from generalized linear mixed-effects models revisited and expanded.” Journal of the Royal Society Interface 14 (134). doi:10.1098/rsif.2017.0213.