R/r-hat.R

In bbsBayes: Hierarchical Bayesian Analysis of North American BBS Data

#' Generate Gelman-Rubin's R-Hat statistic
#'
#' \code{r_hat} returns a dataframe of Gelman-Rubin's R-hat statistics for each
#'   parameter tracked in the model.
#'
#'   R-hat, also known as the potential scale reduction factor (PSRF) was described by
#'   Gelman & Rubin (1992) as a way of calculating convergence of parameters given
#'   2 or more chains. See citation below for details.
#'
#'   Gelman, Andrew; Rubin, Donald B. Inference from Iterative Simulation Using Multiple
#'   Sequences. Statist. Sci. 7 (1992), no. 4, 457--472. doi:10.1214/ss/1177011136.
#'   https://projecteuclid.org/euclid.ss/1177011136
#'
#' @param jags_mod JAGS list generated by \code{run_model}
#' @param parameter_list Optional list of parameters to subset
#' @param threshold Return only r-hat values greater than OR equal to this threshold
#'   (floating point value)
#'
#' @return Dataframe consisting of r-hat values per parameter.
#'
#' @export
#'
#' @examples
#' # Toy example with Pacific Wren sample data
#' # First, stratify the sample data
#'
#' strat_data <- stratify(by = "bbs_cws", sample_data = TRUE)
#'
#' # Prepare the stratified data for use in a JAGS model.
#' jags_data <- prepare_jags_data(strat_data = strat_data,
#'                                species_to_run = "Pacific Wren",
#'                                model = "firstdiff",
#'                                min_year = 2009,
#'                                max_year = 2018)
#'
#' # Now run a JAGS model.
#'
#' jags_mod <- run_model(jags_data = jags_data,
#'                       n_adapt = 0,
#'                       n_burnin = 0,
#'                       n_iter = 10,
#'                       n_thin = 1,
#'                       parameters_to_track = c("n","strata"))
#'
#' # Check convergence for all parameters
#' convergence <- r_hat(jags_mod = jags_mod)
#'
#' # Check convergence for only subset of parameters
#' convergence <- r_hat(jags_mod = jags_mod, parameter_list = "strata")
#'
#' # Only return R Hat values greater than or equal to specified value
#' convergence <- r_hat(jags_mod = jags_mod, threshold = 1.1)
#'
#'

r_hat <- function(jags_mod = NULL,
                  parameter_list = NULL,
                  threshold = NULL)
{
  rhat_list <- jags_mod$Rhat
  if (is.null(parameter_list))
  {
    parameter_list <- names(rhat_list)
  }
  to_return <- data.frame(Parameter = character(),
                          R_hat = double(),
                          stringsAsFactors = FALSE)

  for (p in parameter_list)
  {
    if (isFALSE(p %in% names(rhat_list)))
    {
      warning(paste("Parameter", p, "not found in supplied model."))
      next
    }
    current <- rhat_list[[p]]
    l <- length(dim(current))
    if (l == 0)
    {
      to_return <- rbind(to_return,
                         data.frame(Parameter = p,
                                    R_hat = current))
    }else if (l == 1)
    {
      for (i in 1:dim(current))
      {
        to_return <- rbind(to_return,
                           data.frame(Parameter = paste0(p, "[", i, "]"),
                                      R_hat = current[i]))
      }
    }else if (l == 2)
    {
      for (i in 1:dim(current)[1])
      {
        for (j in 1:dim(current)[2])
        {
          to_return <- rbind(to_return,
                             data.frame(Parameter = paste0(p, "[", i, "][", j, "]"),
                                        R_hat = current[i,j]))
        }
      }
    }else if (l == 3)
    {
      for (i in 1:dim(current)[1])
      {
        for (j in 1:dim(current)[2])
        {
          for (k in 1:dim(current)[3])
          {
            to_return <- rbind(to_return,
                               data.frame(Parameter = paste0(p, "[", i, "][", j, "][", k, "]"),
                                          R_hat = current[i,j,k]))
          }
        }
      }
    }
  }

  if (!is.null(threshold))
  {
    to_return <- to_return[which(to_return$R_hat >= threshold), ]
  }

  return(to_return)

}