R/lugsail_gr.R

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

#' Upgraded Univariate and Multivariate Gelman-Rubin Diagnostic
#'
#' Computes the upgraded Gelman-Rubin convergence statistics (\eqn{\hat{R}_L} and \eqn{\hat{R}^p_L}), 
#' Effective Sample Size (\eqn{\widehat{\mathrm{ESS}}}), and principled termination thresholds (\eqn{\delta_{\epsilon}}) 
#' proposed by Vats and Knudson (2021) using replicated lugsail batch means variance estimators.
#'
#' @param x Input MCMC chain output. Can be a numeric vector (single univariate chain), 
#'   matrix of dimension \code{(n, p)} (single multivariate chain), 3D array of dimension \code{(n, p, m)} 
#'   (multiple multivariate chains), or a list of length \code{m} containing matrices or data.frames.
#' @param alpha Numeric, significance level for the confidence region (default \code{0.05} for 95\% confidence).
#' @param epsilon Numeric, relative volume tolerance for target precision (default \code{0.10}).
#' @param b Optional batch size for lugsail batch means. If \code{NULL}, defaults to \code{floor(sqrt(n))}.
#' @param multivariate Logical, if \code{TRUE} (default), computes multivariate Gelman-Rubin diagnostic 
#'   \eqn{\hat{R}^p_L} and multivariate ESS when \code{p > 1}.
#'
#' @return An object of class \code{"lugsail_gr"} containing:
#'   \item{psrf}{Vector of univariate lugsail potential scale reduction factors \eqn{\hat{R}_L} for each parameter.}
#'   \item{mpsrf}{Multivariate lugsail potential scale reduction factor \eqn{\hat{R}^p_L}.}
#'   \item{ess}{Vector of univariate effective sample sizes for each parameter.}
#'   \item{mess}{Multivariate effective sample size \eqn{\widehat{\mathrm{ESS}}_p}.}
#'   \item{delta_eps}{Principled target threshold \eqn{\delta_{\epsilon}}.}
#'   \item{M_alpha_eps_p}{Minimum required effective sample size bound \eqn{M_{\alpha, \epsilon, p}}.}
#'   \item{converged}{Logical, indicating if diagnostic has converged (\code{mpsrf <= delta_eps}).}
#'   \item{means}{Posterior sample mean vector across chains.}
#'   \item{sd}{Posterior standard deviation vector.}
#'   \item{s2}{Sample variance (scalar or matrix \code{S}).}
#'   \item{tau_L}{Replicated lugsail variance/covariance matrix estimate \eqn{\hat{T}_L}.}
#'   \item{n}{Chain length (number of iterations per chain).}
#'   \item{m}{Number of parallel chains.}
#'   \item{p}{Number of parameter dimensions.}
#'   \item{alpha}{Significance level.}
#'   \item{epsilon}{Relative volume tolerance.}
#'
#' @references
#' Vats, D. and Knudson, C. (2021). Revisiting the Gelman–Rubin Diagnostic. 
#' \emph{Statistical Science}, 36(4), 518–529. \doi{10.1214/20-STS812}.
#'
#' Gelman, A. and Rubin, D. B. (1992). Inference from iterative simulation using multiple sequences. 
#' \emph{Statistical Science}, 7(4), 457–472. \doi{10.1214/ss/1177011136}.
#'
#' Brooks, S. P. and Gelman, A. (1998). General methods for monitoring convergence of iterative simulations. 
#' \emph{Journal of Computational and Graphical Statistics}, 7(4), 434–455. \doi{10.1080/10618600.1998.10474787}.
#'
#' @export
#'
#' @examples
#' # Example 1: Univariate MCMC output (3 chains of 1000 iterations)
#' set.seed(123)
#' chain1 <- rnorm(1000, mean = 0, sd = 1)
#' chain2 <- rnorm(1000, mean = 0.05, sd = 1)
#' chain3 <- rnorm(1000, mean = -0.05, sd = 1)
#' res1 <- lugsail_gr(list(chain1, chain2, chain3))
#' print(res1)
#'
#' # Example 2: Multivariate MCMC output (2 chains of 500 iterations, 2 parameters)
#' mat1 <- matrix(rnorm(1000), ncol = 2)
#' mat2 <- matrix(rnorm(1000), ncol = 2)
#' res2 <- lugsail_gr(list(mat1, mat2))
#' summary(res2)
lugsail_gr <- function(x, alpha = 0.05, epsilon = 0.10, b = NULL, multivariate = TRUE) {
  arr <- format_chains(x)
  dim_arr <- dim(arr)
  n <- dim_arr[1]
  p <- dim_arr[2]
  m <- dim_arr[3]

  if (n < 10) {
    stop("Chain length 'n' must be at least 10 for batch means estimation.")
  }

  # Variable names if present
  param_names <- NULL
  if (is.list(x) && !is.data.frame(x) && is.matrix(x[[1]])) {
    param_names <- colnames(x[[1]])
  } else if (is.matrix(x) || is.data.frame(x)) {
    param_names <- colnames(x)
  }
  if (is.null(param_names) || length(param_names) != p) {
    param_names <- paste0("param_", seq_len(p))
  }

  # Overall sample means and per-chain statistics
  means <- numeric(p)
  names(means) <- param_names

  # Calculate per-chain mean & sample covariance matrix Si
  Si_list <- list()
  chain_means <- matrix(0, nrow = m, ncol = p)

  for (i in seq_len(m)) {
    mat_i <- arr[, , i, drop = FALSE]
    dim(mat_i) <- c(n, p)
    chain_means[i, ] <- colMeans(mat_i)
    if (n > 1) {
      Si_list[[i]] <- stats::cov(mat_i)
    } else {
      Si_list[[i]] <- matrix(0, p, p)
    }
  }

  means <- colMeans(chain_means)
  names(means) <- param_names

  # Average sample covariance S
  S <- matrix(0, p, p)
  for (i in seq_len(m)) {
    S <- S + Si_list[[i]]
  }
  S <- S / m
  S <- (S + t(S)) / 2

  # Calculate lugsail batch means estimator TL
  bm_res <- calc_lugsail_bm(arr, b = b)
  TL <- as.matrix(bm_res$tau_L)
  S <- as.matrix(S)

  # Univariate calculation for each parameter component
  psrf_uni <- numeric(p)
  ess_uni <- numeric(p)
  sd_vec <- numeric(p)
  names(psrf_uni) <- param_names
  names(ess_uni) <- param_names
  names(sd_vec) <- param_names

  for (j in seq_len(p)) {
    s2_j <- S[j, j]
    sd_vec[j] <- sqrt(s2_j)
    tau_L_j <- TL[j, j]

    # Biased-from-above target variance sigma^2_L
    sigma2_L_j <- ((n - 1) / n) * s2_j + (tau_L_j / n)

    if (s2_j <= 1e-12) {
      psrf_uni[j] <- 1.0
      ess_uni[j] <- m * n
    } else {
      psrf_uni[j] <- sqrt(max(1.0, sigma2_L_j / s2_j))
      ess_uni[j] <- (m * n * s2_j) / tau_L_j
    }
  }

  # Multivariate calculation
  mpsrf_val <- NA_real_
  mess_val <- NA_real_

  if (multivariate && p > 1) {
    det_S <- det(S)
    det_TL <- det(TL)

    if (det_S <= 1e-12 || det_TL <= 1e-12 || is.nan(det_S) || is.nan(det_TL)) {
      # Regularization for determinant calculation
      S_reg <- S + diag(1e-6, p)
      TL_reg <- TL + diag(1e-6, p)
      det_ratio <- max(1e-10, det(TL_reg) / det(S_reg))
    } else {
      det_ratio <- det_TL / det_S
    }

    ratio_root <- max(0, det_ratio)^(1 / p)

    # Multivariate R_L^p
    mpsrf_val <- sqrt(((n - 1) / n) + (ratio_root / n))
    
    # Multivariate ESS
    mess_val <- (m * n) / ratio_root
  } else {
    mpsrf_val <- psrf_uni[1]
    mess_val <- ess_uni[1]
  }

  # Threshold calculations
  M_bound <- calc_ess_bound(p = p, alpha = alpha, epsilon = epsilon)
  delta_eps <- calc_psrf_cutoff(m = m, M_val = M_bound)

  # Check convergence
  target_psrf <- if (!is.na(mpsrf_val)) mpsrf_val else max(psrf_uni)
  is_converged <- (target_psrf <= delta_eps)

  res <- list(
    psrf = psrf_uni,
    mpsrf = mpsrf_val,
    ess = ess_uni,
    mess = mess_val,
    delta_eps = delta_eps,
    M_alpha_eps_p = M_bound,
    converged = is_converged,
    means = means,
    sd = sd_vec,
    s2 = S,
    tau_L = TL,
    n = n,
    m = m,
    p = p,
    alpha = alpha,
    epsilon = epsilon,
    batch_size = bm_res$b,
    arr = arr
  )
  class(res) <- "lugsail_gr"
  return(res)
}