R/methods.R

In pHMC: Proximal Hamiltonian Monte Carlo for Non-Smooth Bayesian Inference

#' S3 Methods for Proximal Hamiltonian Monte Carlo Objects
#'
#' Provides summary, printing, extraction, and diagnostic plotting
#' methods for objects returned by \code{\link{phmc}}.
#'
#' @param x An object of class \code{"phmc"}.
#' @param object An object of class \code{"phmc"}.
#' @param type Character string specifying plot type: \code{"trace"}
#'   for trace plots, \code{"acf"} for autocorrelation,
#'   \code{"density"} for posterior density, or \code{"all"}
#'   (default).
#' @param par_indices Optional integer vector specifying parameter
#'   indices to plot.
#' @param ... Additional arguments passed to generic methods.
#'
#' @return Depending on the S3 method invoked, returns the following:
#' \describe{
#'   \item{\code{print.phmc}}{Invisibly returns the input object \code{x} of class \code{"phmc"} (called for its side effect of printing summary metrics to the console).}
#'   \item{\code{summary.phmc}}{Returns an object of class \code{"summary.phmc"}, which is a list containing model parameter estimates, acceptance rate, log-likelihood, information criteria (AIC, BIC, DIC), elapsed time, and regularization parameter \code{lambda_g}.}
#'   \item{\code{print.summary.phmc}}{Invisibly returns the input object \code{x} of class \code{"summary.phmc"} (called for its side effect of printing detailed summary results to the console).}
#'   \item{\code{coef.phmc}}{Returns a named numeric vector of class \code{"numeric"} containing posterior mean parameter estimates.}
#'   \item{\code{vcov.phmc}}{Returns a numeric matrix of class \code{"matrix"} containing the empirical posterior variance-covariance matrix of the MCMC parameter draws.}
#'   \item{\code{logLik.phmc}}{Returns an object of class \code{"logLik"} representing the log-likelihood value evaluated at the MAP estimate, with attributes \code{"df"} (number of estimated parameters) and \code{"nobs"} (number of retained MCMC draws).}
#'   \item{\code{plot.phmc}}{Invisibly returns the input object \code{x} of class \code{"phmc"} (called for its side effect of generating diagnostic MCMC trace plots, autocorrelation functions, and posterior density curves).}
#' }
#'
#' @importFrom grDevices col2rgb rgb
#' @importFrom utils head tail
#' @importFrom stats acf coef cov density median quantile rnorm
#'   runif sd var
#'
#' @name phmc_methods
#' @export
print.phmc <- function(x, ...) {
  cat("\nProximal Hamiltonian Monte Carlo (p-HMC) Estimation\n")
  cat("---------------------------------------------------\n")
  cat(sprintf("Acceptance rate: %.2f%%\n", x$accept_rate * 100))
  cat(sprintf("Draws retained:  %d\n", nrow(x$draws)))
  cat(sprintf("Elapsed time:    %.3f seconds\n", x$elapsed_time))
  cat(sprintf("lambda_g:        %g\n\n", x$lambda_g))

  cat("Posterior Estimates (Mean & 95%% Credible Intervals):\n")
  show_cols <- c("Mean", "StdErr", "2.5%", "97.5%", "ESS")
  show_mat <- x$estimates[, show_cols, drop = FALSE]
  print(round(show_mat, 4))
  invisible(x)
}

#' @rdname phmc_methods
#' @export
summary.phmc <- function(object, ...) {
  res <- list(
    call         = object$call,
    estimates    = object$estimates,
    accept_rate  = object$accept_rate,
    logLik       = object$logLik,
    AIC          = object$AIC,
    BIC          = object$BIC,
    DIC          = object$DIC,
    elapsed_time = object$elapsed_time,
    lambda_g     = object$lambda_g
  )
  class(res) <- "summary.phmc"
  return(res)
}

#' @rdname phmc_methods
#' @export
print.summary.phmc <- function(x, ...) {
  cat("\n===================================================\n")
  cat("Summary of Proximal Hamiltonian Monte Carlo (p-HMC)\n")
  cat("===================================================\n\n")

  cat("Parameter Summary Table:\n")
  print(round(x$estimates, 4))
  cat("\n---------------------------------------------------\n")
  cat(sprintf("Acceptance Rate:  %.2f%%\n", x$accept_rate * 100))
  cat(sprintf("Log-Likelihood:   %.4f\n", x$logLik))
  cat(sprintf("AIC:              %.4f\n", x$AIC))
  cat(sprintf("BIC:              %.4f\n", x$BIC))
  cat(sprintf("DIC:              %.4f\n", x$DIC))
  cat(sprintf("Computation Time: %.3f seconds\n", x$elapsed_time))
  cat("===================================================\n")
  invisible(x)
}

#' @rdname phmc_methods
#' @export
coef.phmc <- function(object, ...) {
  object$estimates[, "Mean"]
}

#' @rdname phmc_methods
#' @export
vcov.phmc <- function(object, ...) {
  stats::cov(object$draws)
}

#' @rdname phmc_methods
#' @export
logLik.phmc <- function(object, ...) {
  val <- object$logLik
  attr(val, "df") <- nrow(object$estimates)
  attr(val, "nobs") <- nrow(object$draws)
  class(val) <- "logLik"
  return(val)
}

#' @rdname phmc_methods
#' @export
plot.phmc <- function(x,
                      type = c("all", "trace", "acf", "density"),
                      par_indices = NULL, ...) {
  type <- match.arg(type)
  draws <- x$draws
  p <- ncol(draws)

  if (is.null(par_indices)) {
    par_indices <- seq_len(min(p, 4L))
  }

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

  if (type == "trace" || type == "all") {
    graphics::par(mfrow = c(length(par_indices), 1),
                  mar = c(3, 4, 2, 1))
    for (idx in par_indices) {
      graphics::plot(draws[, idx], type = "l", col = "steelblue",
        main = paste("Traceplot:", colnames(draws)[idx]),
        xlab = "Iteration", ylab = "Value")
    }
  }

  if (type == "acf" || type == "all") {
    graphics::par(mfrow = c(length(par_indices), 1),
                  mar = c(3, 4, 2, 1))
    for (idx in par_indices) {
      stats::acf(draws[, idx],
        main = paste("ACF:", colnames(draws)[idx]),
        col = "darkred")
    }
  }

  if (type == "density" || type == "all") {
    graphics::par(mfrow = c(length(par_indices), 1),
                  mar = c(3, 4, 2, 1))
    for (idx in par_indices) {
      dens <- stats::density(draws[, idx])
      graphics::plot(dens,
        main = paste("Posterior Density:", colnames(draws)[idx]),
        col = "darkgreen", lwd = 2, xlab = "Parameter Value")
      graphics::polygon(dens,
        col = .make_alpha("darkgreen", 0.3),
        border = NA)
    }
  }

  invisible(x)
}

## Internal helper for color transparency using grDevices
.make_alpha <- function(col, alpha = 0.3) {
  rgb_vals <- grDevices::col2rgb(col) / 255
  grDevices::rgb(rgb_vals[1L], rgb_vals[2L], rgb_vals[3L],
                 alpha = alpha)
}