Compute wAIC with Jags

compute-wAIC-in-Jags.R

##################################################################################
# How to calculate wAIC with JAGS: Linear regression example                     #
# https://sourceforge.net/p/mcmc-jags/discussion/610036/thread/8211df61/#ea5c    #
##################################################################################

#--- Simulate data
set.seed(2020) # set seed
n <- 100 # sample size
x <- sort(rnorm(n)) # covariate values
int <- 30 # true intercept
slope <- 10 # true slope
mu <- int + slope * x # mean
sigma <- 5 # standard deviation
y <- rnorm(n, mean = mu, sd = sigma) # response
dat <- data.frame(y = y, x = x)
head(dat)

#--- Fit linear regression using lm()

# model with covariate
m1.lm <- lm(y ~ x, data = dat) 
# model with intercept only
m0.lm <- lm(y ~ 1, data = dat)

# As expected (true slope is != 0), the difference in AIC tells us that the covariate has some importance
AIC(m0.lm, m1.lm)
  
#--- Fit linear regression using Jags

# load R2jags package
library(R2jags)

# build models
m1.jags <- function(){ # model with covariate
  # Likelihood
  for (i in 1:n){
    y[i] ~ dnorm(mu[i], tau)
    mu[i] <- int + slope * x[i]
  }
  # Priors:
  int ~ dnorm(0, 0.1) # intercept
  slope ~ dnorm(0, 0.1) # slope
  sigma ~ dunif(0, 100) # standard deviation
  tau <- 1 / (sigma * sigma) # precision
}

m0.jags <- function(){ # model with intercept only
  # Likelihood
  for (i in 1:n){
    y[i] ~ dnorm(mu[i], tau)
    mu[i] <- int 
  }
  # Priors:
  int ~ dnorm(0, 0.1) # intercept
  sigma ~ dunif(0, 100) # standard deviation
  tau <- 1 / (sigma * sigma) # precision
}

# initial values for model with covariate
inits.m1 <- function(){ 
  list(int = rnorm(1), 
       slope = rnorm(1), 
       sigma = runif(1))
}

# initial values for model with intercept only
inits.m0 <- function(){ 
  list(int = rnorm(1), 
       sigma = runif(1))
}

# parameters to be monitored
params.m1 <- c("int", "slope", "sigma") # for model with covariate
params.m0 <- c("int", "sigma") # for model with intercept only

# data
jags.data.m1 <- list(y = dat$y, # for model with covariate
                     x = dat$x,
                     n = length(dat$y))
jags.data.m0 <- list(y = dat$y, # for model with intercept only
                     n = length(dat$y))

# run jags and fit model with covariate
m1.jags <- jags(data = jags.data.m1, 
                inits = inits.m1, 
                parameters.to.save = params.m1, 
                model.file = m1.jags,
                n.chains = 2, 
                n.iter = 5000, 
                n.burnin = 1000, 
                n.thin = 1)

# run jags and fit model with intercept only
m0.jags <- jags(data = jags.data.m0, 
                inits = inits.m0, 
                parameters.to.save = params.m0, 
                model.file = m0.jags,
                n.chains = 2, 
                n.iter = 5000, 
                n.burnin = 1000, 
                n.thin = 1)

# compute wAIC for model with intercept only
samples.m0 <- jags.samples(m0.jags$model, 
                           c("WAIC","deviance"), 
                           type = "mean", 
                           n.iter = 5000,
                           n.burnin = 1000,
                           n.thin = 1)
samples.m0$p_waic <- samples.m0$WAIC
samples.m0$waic <- samples.m0$deviance + samples.m0$p_waic
tmp <- sapply(samples.m0, sum)
waic.m0 <- round(c(waic = tmp[["waic"]], p_waic = tmp[["p_waic"]]),1)

# compute wAIC for model with covariate
samples.m1 <- jags.samples(m1.jags$model, 
                           c("WAIC","deviance"), 
                           type = "mean", 
                           n.iter = 5000,
                           n.burnin = 1000,
                           n.thin = 1)
samples.m1$p_waic <- samples.m1$WAIC
samples.m1$waic <- samples.m1$deviance + samples.m1$p_waic
tmp <- sapply(samples.m1, sum)
waic.m1 <- round(c(waic = tmp[["waic"]], p_waic = tmp[["p_waic"]]),1)

# The difference in wAIC tells us that the covariate has some effect
# wAIC of m1 the model with covariate << wAIC of m0 intercept only
data.frame(m0 = waic.m0, m1 = waic.m1)