R/mcmc_gr_test.R

In LugsailGR: Generalized Gelman-Rubin Diagnostic and Effective Sample Size for MCMC

#' Evaluate Gelman-Rubin Diagnostics for User-Supplied Target Density Function
#'
#' Runs Markov chain Monte Carlo (MCMC) sampling for a user-supplied probability 
#' density function (PDF) or log-likelihood function of variable and parameters 
#' with data, and returns upgraded lugsail Gelman-Rubin convergence test statistics,
#' effective sample sizes, parameter estimates, and termination decisions.
#'
#' @param target_pdf Function taking parameter vector \code{par} as its first argument and optional \code{data}. 
#'   It should return the probability density value or log-density value for \code{par}.
#' @param data Optional data vector, matrix, data.frame, or list passed to \code{target_pdf}.
#' @param start_par Initial parameter vector. If named, parameter names will be preserved.
#' @param n_iter Integer, total number of MCMC iterations per chain (default \code{2000}).
#' @param n_chains Integer, number of parallel MCMC chains to generate (default \code{3}).
#' @param burn_in Numeric, fraction of initial burn-in samples to discard (default \code{0.50}).
#' @param scale Numeric or vector, proposal standard deviation for random walk Metropolis-Hastings sampler (default \code{0.5}).
#' @param alpha Numeric, significance level for convergence bound (default \code{0.05}).
#' @param epsilon Numeric, relative volume tolerance (default \code{0.10}).
#' @param ... Additional parameters passed to \code{target_pdf}.
#'
#' @return An object of class \code{"lugsail_gr"} containing Gelman-Rubin test statistics, 
#'   effective sample sizes, cutoff thresholds, parameter estimates, standard errors, 
#'   and MCMC chain samples.
#'
#' @references
#' Vats, D. and Knudson, C. (2021). Revisiting the Gelman–Rubin Diagnostic. 
#' \emph{Statistical Science}, 36(4), 518–529. \doi{10.1214/20-STS812}.
#'
#' @export
#'
#' @examples
#' # Example 1: Normal distribution target PDF with data
#' set.seed(42)
#' true_mean <- 3
#' my_data <- rnorm(50, mean = true_mean, sd = 1)
#'
#' # Log-posterior density function for parameter mu
#' log_post <- function(mu, data) {
#'   sum(dnorm(data, mean = mu, sd = 1, log = TRUE))
#' }
#'
#' # Run diagnostic test
#' fit <- mcmc_gr_test(target_pdf = log_post, data = my_data, start_par = c(mu = 0), 
#'                     n_iter = 1000, n_chains = 3)
#' print(fit)
mcmc_gr_test <- function(target_pdf, data = NULL, start_par, n_iter = 2000, 
                         n_chains = 3, burn_in = 0.5, scale = 0.5, 
                         alpha = 0.05, epsilon = 0.10, ...) {
  if (!is.function(target_pdf)) {
    stop("'target_pdf' must be a valid function of parameter vector 'par' and 'data'.")
  }

  p <- length(start_par)
  par_names <- names(start_par)
  if (is.null(par_names) || length(par_names) != p) {
    par_names <- paste0("param_", seq_len(p))
  }

  if (n_iter < 20) stop("'n_iter' must be at least 20.")
  if (n_chains < 1) stop("'n_chains' must be at least 1.")

  # Check if target_pdf returns log-density or density
  test_val <- suppressWarnings(tryCatch(target_pdf(start_par, data, ...), error = function(e) NA))
  if (is.na(test_val) || !is.numeric(test_val)) {
    stop("Error evaluating 'target_pdf' at initial parameter vector.")
  }

  is_log_scale <- (test_val <= 0)

  eval_log_post <- function(par) {
    val <- suppressWarnings(tryCatch(target_pdf(par, data, ...), error = function(e) -Inf))
    if (!is.numeric(val) || length(val) != 1 || is.nan(val) || is.na(val)) return(-Inf)
    if (!is_log_scale) {
      if (val <= 0) return(-Inf)
      return(log(val))
    }
    return(val)
  }

  # Generate overdispersed starting values across chains
  init_mats <- matrix(0, nrow = n_chains, ncol = p)
  init_mats[1, ] <- start_par
  if (n_chains > 1) {
    for (i in 2:n_chains) {
      init_mats[i, ] <- start_par + stats::rnorm(p, mean = 0, sd = 2 * max(0.1, abs(start_par)))
    }
  }

  # Scale for proposal random walk
  if (length(scale) == 1) {
    scale_vec <- rep(scale, p)
  } else if (length(scale) == p) {
    scale_vec <- scale
  } else {
    scale_vec <- rep(scale[1], p)
  }

  # Run Metropolis-Hastings MCMC chains
  chains_list <- list()

  for (i in seq_len(n_chains)) {
    chain_mat <- matrix(0, nrow = n_iter, ncol = p)
    colnames(chain_mat) <- par_names
    curr_par <- init_mats[i, ]
    curr_lp <- eval_log_post(curr_par)

    for (t in seq_len(n_iter)) {
      prop_par <- curr_par + stats::rnorm(p, mean = 0, sd = scale_vec)
      prop_lp <- eval_log_post(prop_par)

      log_alpha <- prop_lp - curr_lp
      if (!is.na(log_alpha) && log(stats::runif(1)) < log_alpha) {
        curr_par <- prop_par
        curr_lp <- prop_lp
      }
      chain_mat[t, ] <- curr_par
    }
    chains_list[[i]] <- chain_mat
  }

  # Discard burn-in
  burn_steps <- floor(burn_in * n_iter)
  if (burn_steps >= n_iter - 10) burn_steps <- floor(0.1 * n_iter)

  post_chains <- lapply(chains_list, function(mat) {
    mat[(burn_steps + 1):n_iter, , drop = FALSE]
  })

  # Compute Gelman-Rubin lugsail diagnostic
  res <- lugsail_gr(post_chains, alpha = alpha, epsilon = epsilon, multivariate = (p > 1))
  res$raw_chains <- chains_list
  res$burn_in_steps <- burn_steps
  return(res)
}