R/plot_mcmc.R

In cjcarlson/embarcadero: Species distribution models with BART

#' @title Animate the MCMC draws
#'
#' @description
#'
#' plot.mcmc() takes your BART model and your landscape, and shows each individual sum-of-trees model's predictions on the landscape, next to the accumulating posterior mean. If you start with a low number of trees (n = 5 to 20) and drop the BART burn-in (include the first few samples), you can watch the model learn in realtime, which is often very interesting. This can be a little slow, and you might want to adjust some of the settings for that reason. 
#' 
#' @param object A BART model object generated by the dbarts package 
#' @param inputstack An object of class RasterStack
#' @param iter How many of the first 1:n draws do you want to visualize?
#' @param wait Adds a Sys.sleep after the plots; because the plots are two-panel they can be a little delayed sometimes, which isn't super visually smooth. However, if you want something that's particularly crisp, I'd suggest using the animation package (see below example).
#' @param quiet Turns off progress bars if TRUE
#' 
#' @export

plot.mcmc <- function(object, inputstack, iter=100, wait=0.1, quiet=FALSE) {
  
  xnames <- attr(object$fit$data@x, "term.labels")
  if(all(xnames %in% names(inputstack))) {
    inputstack <- inputstack[[xnames]]
  } else {
    stop("Variable names of RasterStack don't match the requested names")
  }
  input.matrix <- as.matrix(getValues(inputstack))
  
  whichvals <- which(complete.cases(input.matrix))
  input.matrix <- input.matrix[complete.cases(input.matrix),]
  
  cat('\n Generating the posterior predictions (takes a second) \n')
  pred <- dbarts:::predict.bart(object, input.matrix)
  
  cat('\n Generating plots \n')
  if(!quiet){pb <- txtProgressBar(min = 0, max = iter, style = 3)}
  for (i in 1:iter){
    output.m <- t(matrix(pred[i,],
                               nrow = ncol(inputstack),
                               ncol = nrow(inputstack)))
    r <- raster(output.m,
                xmn=xmin(inputstack[[1]]), xmx=xmax(inputstack[[1]]),
                ymn=ymin(inputstack[[1]]), ymx=ymax(inputstack[[1]]),
                crs=inputstack[[1]]@crs)
    if(i==1){r1 <- r} else {r1 <- sum(r1, r)}
    par(mfrow=c(1,2))
    par(mar=c(2,1,2,5))
    plot(r, box=F, axes=F, zlim=c(0,1), main=paste('Iter',i))
    plot(r1/i, box=F, axes=F, zlim=c(0,1), main='Mean')
    Sys.sleep(0.1)
    if(!quiet){setTxtProgressBar(pb, i)}
  }
}