Functions for calculating WAIC from a linear model

waic_lm.R

#from rethinking library for numerically stable log sums
log_sum_exp <- function (x) {
  xmax <- max(x)
  xsum <- sum(exp(x - xmax))
  xmax + log(xsum)
}
  

#function for WAIC from an LM
waic.lm <- function(mod, n.sims=1e3){
  mod_sims <- sim(mod, n.sims=n.sims)
  mod_X <- model.matrix(mod)
  mod_Y <- mod$fitted.values+mod$residuals
  
  #generate distribution of observations
  pred_sims <- apply(mod_sims@coef, 1, function(b) mod_X %*% b )
 # pred_sims_err <- sapply(1:nrow(pred_sims), 
  #                     function(i) rnorm(ncol(pred_sims), pred_sims[i,], mod_sims@sigma[i]))
  
  #I hate nested loops, but it's expedient
  ll <- sapply(1:ncol(pred_sims), 
               function(j){
                 sapply(1:nrow(pred_sims),
                        function(i) dnorm(pred_sims[i,j], mod_Y[i], mod_sims@sigma[i], log=TRUE))
  })
  
  #get the things we'll need...
  lppd <- apply(ll, 1, function(arow) log_sum_exp(arow) - log(length(arow)))
  pWAIC <- apply(ll, 1, var)
  
  -2*sum(lppd) + 2*sum(pWAIC)             
}