R/s3_methods.R

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

#' Print Method for lugsail_gr Objects
#'
#' @param x An object of class \code{"lugsail_gr"}.
#' @param ... Further arguments passed to print.
#'
#' @return Invisibly returns \code{x}.
#' @export
print.lugsail_gr <- function(x, ...) {
  cat("\n=========================================================\n")
  cat("  Upgraded Lugsail Gelman-Rubin Diagnostic (Vats & Knudson, 2021)\n")
  cat("=========================================================\n")
  cat(sprintf("Iterations per chain (n) : %d\n", x$n))
  cat(sprintf("Number of chains (m)     : %d\n", x$m))
  cat(sprintf("Number of parameters (p) : %d\n", x$p))
  cat(sprintf("Batch size (b)           : %d\n", x$batch_size))
  cat("---------------------------------------------------------\n")

  # Table of univariate results
  df_res <- data.frame(
    Estimate = round(x$means, 5),
    `Std.Error` = round(x$sd / sqrt(x$ess), 5),
    `PSRF (R_L)` = round(x$psrf, 6),
    `ESS` = round(x$ess, 1),
    check.names = FALSE
  )
  print(df_res)

  cat("---------------------------------------------------------\n")
  if (x$p > 1 && !is.na(x$mpsrf)) {
    cat(sprintf("Multivariate PSRF (R^p_L)  : %.6f\n", x$mpsrf))
    cat(sprintf("Multivariate ESS           : %.1f\n", x$mess))
  }
  cat(sprintf("Required ESS (M_alpha,eps) : %.1f (alpha = %.2f, eps = %.2f)\n", 
              x$M_alpha_eps_p, x$alpha, x$epsilon))
  cat(sprintf("Target Threshold (delta_eps): %.6f\n", x$delta_eps))
  cat("---------------------------------------------------------\n")
  if (x$converged) {
    cat("Status: CONVERGED (PSRF <= Target Threshold delta_eps)\n")
  } else {
    cat("Status: NOT CONVERGED (PSRF > Target Threshold delta_eps)\n")
  }
  cat("=========================================================\n\n")
  invisible(x)
}

#' Summary Method for lugsail_gr Objects
#'
#' @param object An object of class \code{"lugsail_gr"}.
#' @param ... Further arguments passed to summary.
#'
#' @return A summary table object.
#' @export
summary.lugsail_gr <- function(object, ...) {
  print(object, ...)
  invisible(object)
}

#' Plot Method for lugsail_gr Objects
#'
#' Produces diagnostic plots for MCMC chains, including trace plots, 
#' running Gelman-Rubin statistic \eqn{\hat{R}_L} plots, and density estimates.
#'
#' @param x An object of class \code{"lugsail_gr"}.
#' @param type Character string specifying plot type: \code{"all"} (default), 
#'   \code{"trace"}, \code{"running_gr"}, or \code{"density"}.
#' @param ... Additional graphical parameters.
#'
#' @return No return value, called for side effects (generating plots).
#' @export
#'
#' @examples
#' set.seed(123)
#' chain1 <- rnorm(500)
#' chain2 <- rnorm(500)
#' fit <- lugsail_gr(list(chain1, chain2))
#' plot(fit, type = "running_gr")
plot.lugsail_gr <- function(x, type = c("all", "trace", "running_gr", "density"), ...) {
  type <- match.arg(type)

  arr <- x$arr
  n <- x$n
  p <- x$p
  m <- x$m
  param_names <- names(x$means)

  old_par <- graphics::par(no.readonly = TRUE)
  on.exit(graphics::par(old_par))

  if (type == "trace" || type == "all") {
    graphics::par(mfrow = c(min(p, 3), 1), mar = c(4, 4, 2, 1))
    for (j in seq_len(min(p, 3))) {
      ymin <- min(arr[, j, ])
      ymax <- max(arr[, j, ])
      graphics::plot(1:n, arr[, j, 1], type = "l", col = 1,
                     ylim = c(ymin, ymax),
                     xlab = "Iteration", ylab = param_names[j],
                     main = paste("Trace Plot:", param_names[j]))
      if (m > 1) {
        for (i in 2:m) {
          graphics::lines(1:n, arr[, j, i], col = i)
        }
      }
    }
  }

  if (type == "running_gr" || type == "all") {
    graphics::par(mfrow = c(min(p, 3), 1), mar = c(4, 4, 2, 1))
    step_size <- max(10, floor(n / 20))
    eval_seq <- seq(from = max(30, step_size), to = n, by = step_size)

    for (j in seq_len(min(p, 3))) {
      gr_seq <- numeric(length(eval_seq))
      for (k in seq_along(eval_seq)) {
        sub_arr <- arr[1:eval_seq[k], j, , drop = FALSE]
        sub_fit <- lugsail_gr(sub_arr, alpha = x$alpha, epsilon = x$epsilon, multivariate = FALSE)
        gr_seq[k] <- sub_fit$psrf[1]
      }

      graphics::plot(eval_seq, gr_seq, type = "b", pch = 19, col = "blue",
                     xlab = "Iteration (n)", ylab = expression(hat(R)[L]),
                     main = paste("Running Gelman-Rubin Statistic:", param_names[j]),
                     ylim = c(min(0.99, min(gr_seq, na.rm = TRUE)), max(1.15, max(gr_seq, na.rm = TRUE))))
      graphics::abline(h = x$delta_eps, col = "red", lty = 2, lwd = 2)
      graphics::legend("topright", legend = c(expression(hat(R)[L]), expression(delta[epsilon])),
                       col = c("blue", "red"), lty = c(1, 2), lwd = c(1, 2))
    }
  }

  if (type == "density" || type == "all") {
    graphics::par(mfrow = c(min(p, 3), 1), mar = c(4, 4, 2, 1))
    for (j in seq_len(min(p, 3))) {
      d_all <- stats::density(arr[, j, ])
      graphics::plot(d_all, main = paste("Density Estimate:", param_names[j]),
                     xlab = param_names[j], col = "black", lwd = 2)
      if (m > 1) {
        for (i in seq_len(m)) {
          d_i <- stats::density(arr[, j, i])
          graphics::lines(d_i, col = i + 1, lty = 2)
        }
      }
    }
  }

  invisible(NULL)
}