MCMC_ggmcmc.R

library(ggmcmc) # ver. 1.1
library(data.table)
ggs2 <- function(result, pattern=NULL, invert=F, ...){
  if(class(result)[1] =="stanfit" ){
    nThin       <- result@sim$thin
    nChains     <- result@sim$chains
    nBurnin     <- result@sim$warmup2[1]
    nIterations  <- result@sim$n_save[1] - nBurnin
    result <- as.data.frame(result) %>% 
      mutate(Chain=rep(1:nChains, nIterations), Iteration=rep(1:nIterations, nChains)) %>% 
      gather(key=Parameter, value=value, -Chain, -Iteration)
    if(!is.null(pattern)){
      if(invert){result <- filter(result, !grepl(pattern, Parameter, ...))}
      else{result <- filter(result, grepl(pattern, Parameter, ...))}
    }
    result <- arrange(result, Parameter) %>% mutate(Parameter = factor(Parameter, levels= unique(Parameter)))
    attr(result, "nChains")     <- nChains
    attr(result, "nIterations") <- nIterations
  }
  else if(class(result) %in% c("runjags", "mcmc.list")){
    if( class(result) == "runjags"){
      nThin   <- result$thin
      nBurnin <- result$burnin/nThin
      result <- result$mcmc
    }
    else{
      nBurnin <- 0
      nThin  <- 1
    }
    result <- rbindlist(lapply(result, as.data.frame), idcol="Chain") %>% as.data.frame %>%
      group_by(Chain) %>% mutate(Iteration=row_number(Chain)) %>%
      gather(key=Parameter, value=value, -Chain, -Iteration) %>% ungroup 
    if(!is.null(pattern)){
      if(invert){result <- filter(result, !grepl(pattern, Parameter, ...))}
      else{result <- filter(result, grepl(pattern, Parameter, ...))}
    }
    result <- arrange(result, Parameter) %>% mutate(Parameter = factor(Parameter, levels = unique(Parameter)))
    attr(result, "nChains")     <- length(unique(result$Chain))
    attr(result, "nIterations") <- length(unique(result$Iteration))
  }
  attr(result, "nBurnin")       <- nBurnin
  attr(result, "nParameters")   <- length(unique(result$Parameter))
  attr(result, "nThin")         <- nThin
  return(result)
}