R/phmc.R

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

#' Generalized Proximal Hamiltonian Monte Carlo Sampler and Estimator
#'
#' Fits Bayesian models and estimates parameters using the Proximal
#' Hamiltonian Monte Carlo (p-HMC) algorithm for non-differentiable
#' target densities as proposed by Shukla, Vats, and Chi (2025).
#'
#' @param fn Function. The smooth potential component \eqn{f(x)} or
#'   complete potential \eqn{U(x) = -\log \pi(x)}.
#'   Must accept parameter vector \code{x} as first argument.
#' @param grad_f Function or \code{NULL}. Gradient of the smooth
#'   potential component \eqn{f(x)}. If \code{NULL},
#'   finite difference approximation is automatically computed.
#' @param g Function or \code{NULL}. The non-smooth penalty
#'   component \eqn{g(x)}.
#' @param prox_fn Function or character. Proximal mapping operator for
#'   \eqn{g(x)}, or name of built-in operator (\code{"l1"},
#'   \code{"l2"}, \code{"elastic_net"}, \code{"nuclear_norm"},
#'   \code{"none"}). Default is \code{"l1"}.
#' @param start Numeric vector or matrix. Initial parameter values.
#' @param data Optional dataset passed as second argument to
#'   \code{fn}, \code{grad_f}, and \code{g}.
#' @param n_draws Integer > 0. Total number of MCMC iterations to
#'   run (default 2000).
#' @param burnin Integer >= 0. Number of initial draws to discard
#'   as burn-in (default \code{floor(n_draws / 2)}).
#' @param thin Integer >= 1. Thinning interval (default 1).
#' @param epsilon Numeric scalar > 0. Leapfrog step size
#'   (default 0.01).
#' @param L Integer >= 1. Number of leapfrog steps per proposal
#'   (default 10).
#' @param lambda_g Numeric scalar > 0. Moreau-Yosida regularization
#'   parameter (default 0.01).
#' @param M Mass matrix or \code{NULL} (defaults to identity matrix).
#' @param tune_lambda Logical. If \code{TRUE}, uses
#'   \code{\link{phmc_tune}} to automatically select optimal
#'   \code{lambda_g}.
#' @param verbose Logical. If \code{TRUE}, prints sampling progress.
#' @param ... Additional arguments passed to \code{fn}, \code{grad_f},
#'   \code{g}, and \code{prox_fn}.
#'
#' @return An object of class \code{"phmc"}, which is a list containing:
#'   \item{draws}{A numeric matrix of class \code{"matrix"} containing MCMC parameter samples after burn-in and thinning.}
#'   \item{estimates}{A summary matrix of class \code{"matrix"} with rows corresponding to parameters and columns containing posterior Mean, MAP, Median, StdErr, 2.5\% and 97.5\% credible interval bounds, ESS, and ESS per second.}
#'   \item{accept_rate}{A numeric scalar of class \code{"numeric"} giving the overall Metropolis-Hastings acceptance rate (between 0 and 1).}
#'   \item{ess}{A named numeric vector of class \code{"numeric"} containing Effective Sample Size estimates for each parameter.}
#'   \item{ess_per_sec}{A named numeric vector of class \code{"numeric"} containing Effective Sample Size per second for each parameter.}
#'   \item{logLik}{A numeric scalar of class \code{"numeric"} giving the log-likelihood value evaluated at the Maximum A Posteriori (MAP) parameter estimate.}
#'   \item{AIC}{A numeric scalar of class \code{"numeric"} giving the Akaike Information Criterion value for model assessment.}
#'   \item{BIC}{A numeric scalar of class \code{"numeric"} giving the Bayesian Information Criterion value for model assessment.}
#'   \item{DIC}{A numeric scalar of class \code{"numeric"} giving the Deviance Information Criterion value for model assessment.}
#'   \item{elapsed_time}{A numeric scalar of class \code{"numeric"} giving total sampler execution time in seconds.}
#'   \item{lambda_g}{A numeric scalar of class \code{"numeric"} specifying the Moreau-Yosida regularization parameter used during sampling.}
#'   \item{call}{An object of class \code{"call"} recording the matched function call.}
#'
#' @references
#' Shukla A, Vats D, Chi EC (2025).
#' \dQuote{Proximal Hamiltonian Monte Carlo.}
#' \emph{arXiv preprint}, \doi{10.48550/arXiv.2510.22252}.
#'
#' @examples
#' \donttest{
#' set.seed(42)
#' y_data <- rnorm(100, mean = 1, sd = 0.5)
#' f_smooth <- function(x, y) 0.5 * sum((y - x)^2)
#' grad_f_smooth <- function(x, y) -sum(y - x)
#' fit <- phmc(fn = f_smooth, grad_f = grad_f_smooth,
#'             prox_fn = "l1", start = 0.5,
#'             data = y_data, lambda_g = 0.01,
#'             n_draws = 500)
#' summary(fit)
#' }
#'
#' @export
phmc <- function(fn,
                 grad_f = NULL,
                 g = NULL,
                 prox_fn = "l1",
                 start,
                 data = NULL,
                 n_draws = 2000,
                 burnin = floor(n_draws / 2),
                 thin = 1,
                 epsilon = 0.01,
                 L = 10,
                 lambda_g = 0.01,
                 M = NULL,
                 tune_lambda = FALSE,
                 verbose = FALSE,
                 ...) {

  cl <- match.call()
  start_time <- Sys.time()

  ## Flatten start if matrix or vector
  is_matrix_param <- is.matrix(start)
  orig_dim <- if (is_matrix_param) dim(start) else NULL
  start_vec <- as.vector(start)
  d <- length(start_vec)

  param_names <- names(start_vec)
  if (is.null(param_names)) {
    if (is_matrix_param) {
      grid <- expand.grid(seq_len(orig_dim[1]), seq_len(orig_dim[2]))
      param_names <- paste0("X[", grid[, 1], ",", grid[, 2], "]")
    } else {
      param_names <- paste0("param_", seq_len(d))
    }
  }

  ## Automated lambda tuning if requested
  if (tune_lambda) {
    if (verbose) message("Running automated lambda_g tuning...")
    tune_res <- phmc_tune(fn = fn, grad_f = grad_f, g = g,
                          prox_fn = prox_fn, start = start_vec,
                          epsilon = epsilon, L = L, M = M,
                          data = data, ...)
    lambda_g <- tune_res$optimal_lambda_g
    if (verbose) message(sprintf("Selected optimal lambda_g = %g", lambda_g))
  }

  ## Mass matrix setup
  if (is.null(M)) {
    M_mat <- diag(d)
    M_inv <- diag(d)
    M_chol <- diag(d)
  } else if (is.vector(M) && !is.matrix(M)) {
    M_mat <- diag(M, nrow = d)
    M_inv <- diag(1 / M, nrow = d)
    M_chol <- diag(sqrt(M), nrow = d)
  } else {
    M_mat <- M
    M_inv <- solve(M)
    M_chol <- t(chol(M_mat))
  }

  ## Define smooth potential evaluation f(x)
  eval_f <- function(x) {
    x_val <- if (is_matrix_param) {
      matrix(x, nrow = orig_dim[1], ncol = orig_dim[2])
    } else {
      x
    }
    res <- tryCatch(
      if (!is.null(data)) fn(x_val, data, ...) else fn(x_val, ...),
      error = function(e) NA_real_
    )
    res <- as.numeric(res)
    if (is.na(res) || is.nan(res) || is.infinite(res)) return(1e10)
    return(res)
  }

  ## Define non-smooth g(x)
  eval_g <- function(x) {
    if (!is.null(g)) {
      x_val <- if (is_matrix_param) {
        matrix(x, nrow = orig_dim[1], ncol = orig_dim[2])
      } else {
        x
      }
      res <- tryCatch(
        if (!is.null(data)) g(x_val, data, ...) else g(x_val, ...),
        error = function(e) 0
      )
      return(as.numeric(res))
    }
    if (is.character(prox_fn)) {
      p_type <- match.arg(tolower(prox_fn),
        c("l1", "l2", "elastic_net", "nuclear_norm", "none"))
      if (p_type == "l1") return(sum(abs(x)))
      if (p_type == "l2") return(sqrt(sum(x^2)))
      if (p_type == "nuclear_norm") {
        x_m <- if (is_matrix_param) {
          matrix(x, nrow = orig_dim[1], ncol = orig_dim[2])
        } else {
          n_side <- as.integer(round(sqrt(length(x))))
          matrix(x, nrow = n_side)
        }
        return(sum(svd(x_m)$d))
      }
    }
    return(0)
  }

  ## Define grad_f evaluation
  eval_grad_f <- function(x) {
    if (is.function(grad_f)) {
      x_val <- if (is_matrix_param) {
        matrix(x, nrow = orig_dim[1], ncol = orig_dim[2])
      } else {
        x
      }
      res <- if (!is.null(data)) grad_f(x_val, data, ...) else grad_f(x_val, ...)
      return(as.vector(res))
    } else {
      ## Central finite-difference gradient
      eps_diff <- 1e-5
      grad_vec <- numeric(d)
      for (i in seq_len(d)) {
        x_plus <- x
        x_plus[i] <- x[i] + eps_diff
        x_minus <- x
        x_minus[i] <- x[i] - eps_diff
        grad_vec[i] <- (eval_f(x_plus) - eval_f(x_minus)) / (2 * eps_diff)
      }
      return(grad_vec)
    }
  }

  ## Define total gradient of approximate potential
  eval_grad_total <- function(x) {
    gf <- eval_grad_f(x)
    x_val <- if (is_matrix_param) {
      matrix(x, nrow = orig_dim[1], ncol = orig_dim[2])
    } else {
      x
    }
    g_my <- grad_my_envelope(x_val, prox_fn = prox_fn,
                             lambda_g = lambda_g, ...)
    gf + as.vector(g_my)
  }

  ## Kinetic energy and Hamiltonian
  eval_kinetic <- function(p) 0.5 * sum(p * as.vector(M_inv %*% p))
  eval_potential <- function(x) eval_f(x) + eval_g(x)
  eval_hamiltonian <- function(x, p) eval_potential(x) + eval_kinetic(p)

  ## Storage for MCMC
  raw_draws <- matrix(0, nrow = n_draws, ncol = d)
  colnames(raw_draws) <- param_names
  accepted <- logical(n_draws)
  energies <- numeric(n_draws)

  current_x <- start_vec
  raw_draws[1, ] <- current_x
  energies[1] <- eval_potential(current_x)
  accepted[1] <- TRUE

  n_accepted <- 0L

  for (t in 2:n_draws) {
    ## Sample momentum p ~ N(0, M)
    p0 <- as.vector(M_chol %*% stats::rnorm(d))

    ## Hamiltonian at initial state
    H0 <- eval_hamiltonian(current_x, p0)

    ## Leapfrog integrator
    q <- current_x
    p <- p0

    ## Half step for momentum
    p <- p - 0.5 * epsilon * eval_grad_total(q)

    ## Full steps for position and momentum
    for (l in seq_len(L)) {
      q <- q + epsilon * as.vector(M_inv %*% p)
      if (l < L) {
        p <- p - epsilon * eval_grad_total(q)
      }
    }
    ## Final half step for momentum
    p <- p - 0.5 * epsilon * eval_grad_total(q)

    ## Calculate proposal energy
    HL <- eval_hamiltonian(q, p)

    ## Metropolis-Hastings acceptance ratio
    log_alpha <- -HL + H0
    if (is.na(log_alpha) || is.nan(log_alpha)) {
      log_alpha <- -Inf
    }
    alpha_val <- min(1, exp(log_alpha))

    if (stats::runif(1) <= alpha_val) {
      current_x <- q
      n_accepted <- n_accepted + 1L
      accepted[t] <- TRUE
    } else {
      accepted[t] <- FALSE
    }

    raw_draws[t, ] <- current_x
    energies[t] <- eval_potential(current_x)

    if (verbose && (t %% max(1L, floor(n_draws / 10)) == 0L)) {
      message(sprintf("Iter %d / %d | Accept rate: %.2f%%",
                      t, n_draws, 100 * n_accepted / t))
    }
  }

  end_time <- Sys.time()
  elapsed <- as.numeric(difftime(end_time, start_time, units = "secs"))

  ## Post-processing: Burnin and Thinning
  keep_idx <- seq(from = burnin + 1L, to = n_draws, by = thin)
  draws <- raw_draws[keep_idx, , drop = FALSE]
  kept_energies <- energies[keep_idx]

  ## Effective Sample Size calculation
  calc_ess <- function(x_vec) {
    n <- length(x_vec)
    if (n < 4L || stats::var(x_vec) == 0) return(as.numeric(n))
    acf_vals <- stats::acf(x_vec, plot = FALSE,
                           lag.max = min(n - 1L, 100L))$acf[, 1, 1]
    neg_idx <- which(acf_vals < 0)
    max_lag <- if (length(neg_idx) > 0L) neg_idx[1L] - 1L else length(acf_vals)
    if (max_lag <= 1L) return(as.numeric(n))
    rho_sum <- sum(acf_vals[2:max_lag])
    ess_val <- n / (1 + 2 * rho_sum)
    max(1, min(n, ess_val))
  }

  ess_vec <- apply(draws, 2, calc_ess)
  ess_per_sec <- ess_vec / max(elapsed, 1e-4)

  ## Parameter Summary Statistics
  means <- colMeans(draws)
  medians <- apply(draws, 2, stats::median)
  sds <- apply(draws, 2, stats::sd)
  map_idx <- keep_idx[which.min(kept_energies)]
  map_est <- raw_draws[map_idx, ]

  ci_lower <- apply(draws, 2, function(z) {
    stats::quantile(z, probs = 0.025, names = FALSE)
  })
  ci_upper <- apply(draws, 2, function(z) {
    stats::quantile(z, probs = 0.975, names = FALSE)
  })

  estimates <- cbind(
    Mean       = means,
    MAP        = map_est,
    Median     = medians,
    StdErr     = sds,
    `2.5%`     = ci_lower,
    `97.5%`    = ci_upper,
    ESS        = ess_vec,
    `ESS/sec`  = ess_per_sec
  )

  ## Model selection metrics (AIC, BIC, DIC, logLik)
  best_pot <- eval_potential(map_est)
  logLik_val <- -best_pot

  n_obs <- if (!is.null(data)) {
    if (is.matrix(data) || is.data.frame(data)) nrow(data) else length(data)
  } else {
    1L
  }

  aic_val <- 2 * d - 2 * logLik_val
  bic_val <- d * log(max(n_obs, 1L)) - 2 * logLik_val

  ## DIC calculation
  mean_pot <- mean(kept_energies)
  pot_at_mean <- eval_potential(means)
  pD <- mean_pot - pot_at_mean
  dic_val <- mean_pot + pD

  res <- list(
    draws        = draws,
    raw_draws    = raw_draws,
    estimates    = estimates,
    accept_rate  = mean(accepted),
    ess          = ess_vec,
    ess_per_sec  = ess_per_sec,
    logLik       = logLik_val,
    AIC          = aic_val,
    BIC          = bic_val,
    DIC          = dic_val,
    elapsed_time = elapsed,
    lambda_g     = lambda_g,
    epsilon      = epsilon,
    L            = L,
    call         = cl
  )
  class(res) <- "phmc"
  return(res)
}