R/utils.R

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

#' Calculate Replicated Lugsail Batch Means Estimator
#'
#' Computes the replicated lugsail batch means variance estimator (univariate)
#' or time-average covariance matrix (multivariate) for MCMC chains as described
#' by Vats and Knudson (2021).
#'
#' @param arr 3D numeric array of dimension \code{(n, p, m)} representing \code{n} steps,
#'   \code{p} parameters, and \code{m} chains.
#' @param b Batch size. If \code{NULL}, defaults to \code{floor(sqrt(n))}.
#'
#' @return A list containing:
#'   \item{tau_L}{Scalar variance estimate if \code{p = 1}, or \code{p x p} covariance matrix if \code{p > 1}.}
#'   \item{b}{Batch size used.}
#'   \item{a}{Number of batches.}
#' @keywords internal
calc_lugsail_bm <- function(arr, b = NULL) {
  dim_arr <- dim(arr)
  n <- dim_arr[1]
  p <- dim_arr[2]
  m <- dim_arr[3]

  if (is.null(b)) {
    b <- max(3, floor(sqrt(n)))
  } else {
    b <- max(3, floor(b))
  }

  if (b >= n / 2) {
    b <- max(3, floor(sqrt(n)))
  }

  # Helper for single batch size calculation
  compute_Tb <- function(batch_sz) {
    a <- floor(n / batch_sz)
    if (a < 1) return(matrix(0, p, p))
    
    # Batch means array: (a, p, m)
    Y_ik <- array(0, dim = c(a, p, m))
    for (i in seq_len(m)) {
      for (k in seq_len(a)) {
        idx <- ((k - 1) * batch_sz + 1):(k * batch_sz)
        if (p == 1) {
          Y_ik[k, 1, i] <- mean(arr[idx, 1, i])
        } else {
          Y_ik[k, , i] <- colMeans(arr[idx, , i, drop = FALSE])
        }
      }
    }

    # Overall mean vector across all chains and steps
    mu_hat <- numeric(p)
    for (j in seq_len(p)) {
      mu_hat[j] <- mean(arr[1:(a * batch_sz), j, ])
    }

    # Sum of outer products
    denom <- a * m - 1
    if (denom <= 0) denom <- 1

    Tb <- matrix(0, p, p)
    for (i in seq_len(m)) {
      for (k in seq_len(a)) {
        diff_vec <- Y_ik[k, , i] - mu_hat
        Tb <- Tb + outer(diff_vec, diff_vec)
      }
    }
    Tb <- (batch_sz / denom) * Tb
    return(Tb)
  }

  # Batch size b and b/3
  Tb <- compute_Tb(b)
  b3 <- max(1, floor(b / 3))
  Tb3 <- compute_Tb(b3)

  # Lugsail estimator: 2 * Tb - Tb3
  TL <- 2 * Tb - Tb3

  # Ensure positive definiteness / symmetry
  TL <- (TL + t(TL)) / 2
  if (p == 1) {
    TL <- max(1e-10, as.numeric(TL))
  } else {
    eig <- eigen(TL, symmetric = TRUE, only.values = TRUE)$values
    if (any(eig <= 1e-10)) {
      TL <- TL + diag(max(1e-6, abs(min(eig)) + 1e-6), p)
    }
  }

  list(tau_L = TL, b = b, a = floor(n / b))
}

#' Calculate Minimum Required Effective Sample Size (M_alpha_eps_p)
#'
#' Computes the theoretical minimum effective sample size required to obtain a 
#' confidence region with relative volume \code{epsilon} and coverage \code{1 - alpha},
#' based on Vats, Flegal and Jones (2019) and Vats and Knudson (2021).
#'
#' @param p Integer, number of parameters.
#' @param alpha Numeric, significance level (default 0.05 for 95\% confidence).
#' @param epsilon Numeric, relative volume tolerance (default 0.10).
#'
#' @return Numeric scalar, required minimum ESS.
#' @export
#'
#' @examples
#' calc_ess_bound(p = 1, alpha = 0.05, epsilon = 0.10)
#' calc_ess_bound(p = 5, alpha = 0.05, epsilon = 0.05)
calc_ess_bound <- function(p = 1, alpha = 0.05, epsilon = 0.10) {
  if (p < 1) stop("p must be a positive integer.")
  if (alpha <= 0 || alpha >= 1) stop("alpha must be between 0 and 1.")
  if (epsilon <= 0) stop("epsilon must be positive.")

  chi2_val <- stats::qchisq(1 - alpha, df = p)
  term1 <- (2^(2 / p) * pi) / ((p * gamma(p / 2))^(2 / p))
  M_val <- (term1 * chi2_val) / (epsilon^2)
  return(M_val)
}

#' Calculate Gelman-Rubin Termination Threshold (delta_epsilon)
#'
#' Computes the principled Gelman-Rubin diagnostic termination threshold 
#' \code{delta_epsilon} based on the target ESS bound \code{M_alpha_eps_p} and 
#' number of chains \code{m}.
#'
#' @param m Integer, number of chains.
#' @param M_val Numeric, minimum required ESS obtained from \code{\link{calc_ess_bound}}.
#'
#' @return Numeric scalar, target threshold \code{delta_epsilon}.
#' @export
#'
#' @examples
#' M_bound <- calc_ess_bound(p = 1, alpha = 0.05, epsilon = 0.10)
#' calc_psrf_cutoff(m = 3, M_val = M_bound)
calc_psrf_cutoff <- function(m = 1, M_val = 1537) {
  if (m < 1) stop("m must be at least 1.")
  if (M_val <= 0) stop("M_val must be positive.")

  delta_eps <- sqrt(1 + m / M_val)
  return(delta_eps)
}

#' Format Input Chains into a 3D Array
#'
#' @param x Input object: vector, matrix, 3D array, or list of matrices/data.frames.
#'
#' @return 3D numeric array of dimension \code{(n, p, m)}.
#' @keywords internal
format_chains <- function(x) {
  if (is.array(x) && length(dim(x)) == 3) {
    # Already 3D array (n, p, m)
    arr <- x
  } else if (is.list(x) && !is.data.frame(x)) {
    # List of matrices/data.frames per chain
    m <- length(x)
    mats <- lapply(x, function(item) {
      if (is.data.frame(item) || is.vector(item)) as.matrix(item) else item
    })
    n <- nrow(mats[[1]])
    p <- ncol(mats[[1]])
    arr <- array(0, dim = c(n, p, m))
    for (i in seq_len(m)) {
      arr[, , i] <- mats[[i]]
    }
  } else if (is.matrix(x) || is.data.frame(x)) {
    # Single chain matrix (n, p) -> 3D array (n, p, 1)
    mat <- as.matrix(x)
    n <- nrow(mat)
    p <- ncol(mat)
    arr <- array(mat, dim = c(n, p, 1))
  } else if (is.vector(x) && is.numeric(x)) {
    # Single chain vector n -> 3D array (n, 1, 1)
    n <- length(x)
    arr <- array(x, dim = c(n, 1, 1))
  } else {
    stop("Input 'x' must be a vector, matrix, 3D array, or list of matrices/data.frames.")
  }
  return(arr)
}