R/ggs_grb.R

In ggmcmc: Tools for Analyzing MCMC Simulations from Bayesian Inference

#' Gelman-Rubin-Brooks plot (Rhat shrinkage)
#'
#' Generate a Figure with the Rhat shrinkage evolution over bins of simulations, known as the Gelman-Rubin-Brooks plot, or the Gelman plot. For the Potential Scale Reduction Factor (Rhat), proposed by Gelman and Rubin (1992), the version from the second edition of Bayesian Data Analysis (Gelman, Carlin, Stern and Rubin) is used, but the version used in the package "coda" can also be used (Brooks & Gelman 1998).
#'
#' Notice that at least two chains are required.
#'
#' @references Fernández-i-Marín, Xavier (2016) ggmcmc: Analysis of MCMC Samples and Bayesian Inference. Journal of Statistical Software, 70(9), 1-20. doi:10.18637/jss.v070.i09
#' @references Gelman, Carlin, Stern and Rubin (2003) Bayesian Data Analysis. 2nd edition. Chapman & Hall/CRC, Boca Raton.
#' @references Gelman, A and Rubin, DB (1992) Inference from iterative simulation using multiple sequences, _Statistical Science_, *7*, 457-511.
#' @references Brooks, S. P., and Gelman, A. (1998). General methods for monitoring convergence of iterative simulations. _Journal of computational and graphical statistics_, 7(4), 434-455.
#' @param D Data frame whith the simulations
#' @param family Name of the family of parameters to plot, as given by a character vector or a regular expression. A family of parameters is considered to be any group of parameters with the same name but different numerical value between square brackets (as beta[1], beta[2], etc).
#' @param scaling Value of the upper limit for the x-axis. By default, it is 1.5, to help contextualization of the convergence. When 0 or NA, the axis are not scaled.
#' @param greek Logical value indicating whether parameter labels have to be parsed to get Greek letters. Defaults to false.
#' @param version_rhat Character variable with the name of the version of the potential scale reduction factor to use. Defaults to "BDA2", which refers to the second version of _Bayesian Data Analysis_ (Gelman, Carlin, Stern and Rubin). The other available version is "BG98", which refers to Brooks & Gelman (1998) and is the one used in the "coda" package.
#' @param bins Numerical value with the number of bins requested. Defaults to 50.
#' @param plot Logical value indicating whether the plot must be returned (the default) or a tidy dataframe with the results of the Rhat diagnostics per Parameter.
#' @return A \code{ggplot} object, or a \code{tidy} data frame.
#' @export
#' @examples
#' data(linear)
#' ggs_grb(ggs(s))
ggs_grb <- function(D, family = NA, scaling = 1.5, greek = FALSE, version_rhat = "BDA2", bins = 50, plot = TRUE) {
  if (attributes(D)$nChains < 2) {
    stop("At least two chains are required")
  }
  # Manage subsetting a family of parameters
  if (!is.na(family)) {
    D <- get_family(D, family = family)
  }
  if (attributes(D)$nIter <= bins) {
    stop("Insufficient number of iterations provided or too many bins required.")
  }
  # Get the data generated by the ggs_Rhat without plotting
  d <- NULL
  for (b in 1:bins) {
    d <- dplyr::filter(D, Iteration <= floor(attributes(D)$nIter / b)) %>%
      ggs_Rhat(scaling = scaling, greek = greek, version_rhat = version_rhat, plot = FALSE) %>%
      mutate(`Last iteration` = attributes(D)$nIter / b) %>%
      bind_rows(d)
  }
  # Plot or return calculations
  if (plot) {
    f <- ggplot(d, aes(x = `Last iteration`, y = Rhat)) + geom_line() +
      ggtitle("Shrinkage of Potential Scale Reduction Factors")
    if (greek) {
      f <- f + facet_wrap(~ Parameter, labeller = label_parsed, scales = "free")
    } else {
      f <- f + facet_wrap(~ Parameter, scales = "free")
    }
    # If scaling, add the scale
    if (!is.na(scaling)) {
      # Use the maximum of Rhat if it is larger than the prespecified value
      scaling <- ifelse(scaling > max(d$Rhat, na.rm=TRUE), scaling, max(d$Rhat, na.rm=TRUE))
      f <- f + ylim(min(d$Rhat, na.rm=TRUE), scaling)
    }
    return(f)
  } else {
    return(d)
  }
}