R/tune.R

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

#' Hyperparameter Tuning for Moreau-Yosida Regularization Parameter
#'
#' Evaluates the relative Hamiltonian error metric across a grid of
#' candidate \code{lambda_g} values to select the optimal regularization
#' parameter balancing potential smoothness and Hamiltonian conservation,
#' as detailed in Section V of Shukla, Vats, and Chi (2025).
#'
#' @param fn Function. Smooth component \eqn{f(x)} of the potential.
#'   Must accept parameter vector as first argument.
#' @param grad_f Function or \code{NULL}. Analytical gradient of
#'   \eqn{f(x)}. If \code{NULL}, finite differences are used.
#' @param g Function or \code{NULL}. Non-smooth component \eqn{g(x)}.
#' @param prox_fn Function or character. Proximal operator or
#'   built-in name (\code{"l1"}, \code{"l2"}, \code{"nuclear_norm"},
#'   \code{"elastic_net"}, \code{"none"}).
#' @param start Numeric vector. Starting parameter value.
#' @param lambda_grid Numeric vector. Grid of candidate
#'   \code{lambda_g} values to test.
#' @param epsilon Numeric scalar > 0. Leapfrog step size.
#' @param L Integer >= 1. Number of leapfrog steps.
#' @param M Mass matrix or \code{NULL} (defaults to identity matrix).
#' @param target_rel_err Numeric scalar. Maximum acceptable relative
#'   Hamiltonian error (default 1e-4).
#' @param data Optional dataset passed as second argument to \code{fn}
#'   and \code{grad_f}.
#' @param seed Optional integer or \code{NULL}. Random seed to set conditionally
#'   for reproducible initial momentum sampling. Default is \code{NULL}.
#' @param ... Additional arguments forwarded to \code{fn},
#'   \code{grad_f}, or \code{prox_fn}.
#'
#' @return An object of class \code{"phmc_tune"}, which is a list containing:
#'   \item{optimal_lambda_g}{A numeric scalar of class \code{"numeric"} specifying the selected optimal Moreau-Yosida regularization parameter \code{lambda_g} that satisfies the relative Hamiltonian error threshold.}
#'   \item{grid_results}{A data frame of class \code{"data.frame"} containing tested candidate \code{lambda_g} values, computed relative Hamiltonian errors (\code{R_lambda_g}), and absolute Hamiltonian difference metrics (\code{H_diff}).}
#'   \item{target_rel_err}{A numeric scalar of class \code{"numeric"} giving the target relative Hamiltonian error threshold used during grid search tuning.}
#'
#' @details The relative Hamiltonian error metric is defined as
#'   \deqn{R_{\lambda_g} = \left| \frac{H(x_0, p_0) -
#'     H(\tilde{T}_{\epsilon, L}^{\lambda_g}(x_0, p_0))}{H(x_0, p_0)}
#'     \right|.}
#'
#' @references
#' Shukla A, Vats D, Chi EC (2025).
#' \dQuote{Proximal Hamiltonian Monte Carlo.}
#' \emph{arXiv preprint}, \doi{10.48550/arXiv.2510.22252}.
#'
#' @export
phmc_tune <- function(fn,
                      grad_f = NULL,
                      g = NULL,
                      prox_fn = "l1",
                      start,
                      lambda_grid = 10^seq(-5, 0, length.out = 15),
                      epsilon = 0.001,
                      L = 10,
                      M = NULL,
                      target_rel_err = 1e-4,
                      data = NULL,
                      seed = NULL,
                      ...) {

  d <- length(start)
  if (is.null(M)) {
    M_inv <- diag(d)
    M_mat <- diag(d)
  } else if (is.vector(M) && !is.matrix(M)) {
    M_inv <- diag(1 / M, nrow = d)
    M_mat <- diag(M, nrow = d)
  } else {
    M_mat <- M
    M_inv <- solve(M)
  }

  ## Define smooth potential evaluation
  eval_f <- function(x) {
    res <- tryCatch(
      if (!is.null(data)) fn(x, data, ...) else fn(x, ...),
      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 evaluation
  eval_g <- function(x) {
    if (!is.null(g)) {
      res <- tryCatch(
        if (!is.null(data)) g(x, data, ...) else g(x, ...),
        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") {
        if (!is.matrix(x)) {
          n_side <- as.integer(round(sqrt(length(x))))
          x <- matrix(x, nrow = n_side)
        }
        return(sum(svd(x)$d))
      }
    }
    return(0)
  }

  ## Define grad_f evaluation
  eval_grad_f <- function(x) {
    if (is.function(grad_f)) {
      if (!is.null(data)) return(grad_f(x, data, ...))
      return(grad_f(x, ...))
    }
    ## Numerical gradient via central finite difference
    eps_diff <- 1e-6
    grad_vec <- numeric(length(x))
    for (i in seq_along(x)) {
      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)
  }

  if (!is.null(seed)) {
    set.seed(seed)
  }
  p0 <- stats::rnorm(d)
  if (!is.null(M)) p0 <- as.vector(t(chol(M_mat)) %*% p0)

  kin_energy0 <- 0.5 * sum(p0 * (M_inv %*% p0))
  pot_energy0 <- eval_f(start) + eval_g(start)
  H0 <- pot_energy0 + kin_energy0

  results <- data.frame(
    lambda_g   = lambda_grid,
    R_lambda_g = numeric(length(lambda_grid)),
    H_diff     = numeric(length(lambda_grid))
  )

  for (k in seq_along(lambda_grid)) {
    lam <- lambda_grid[k]

    ## Perform leapfrog integration
    curr_x <- start
    curr_p <- p0

    grad_total <- eval_grad_f(curr_x) +
      grad_my_envelope(curr_x, prox_fn = prox_fn, lambda_g = lam, ...)
    curr_p <- curr_p - 0.5 * epsilon * grad_total

    for (step in seq_len(L)) {
      curr_x <- curr_x + epsilon * as.vector(M_inv %*% curr_p)
      if (step < L) {
        grad_total <- eval_grad_f(curr_x) +
          grad_my_envelope(curr_x, prox_fn = prox_fn, lambda_g = lam, ...)
        curr_p <- curr_p - epsilon * grad_total
      }
    }
    grad_total <- eval_grad_f(curr_x) +
      grad_my_envelope(curr_x, prox_fn = prox_fn, lambda_g = lam, ...)
    curr_p <- curr_p - 0.5 * epsilon * grad_total

    kin_energy_L <- 0.5 * sum(curr_p * (M_inv %*% curr_p))
    pot_energy_L <- eval_f(curr_x) + eval_g(curr_x)
    HL <- pot_energy_L + kin_energy_L

    H_diff <- abs(H0 - HL)
    R_lam <- H_diff / max(abs(H0), 1e-10)

    results$R_lambda_g[k] <- R_lam
    results$H_diff[k] <- H_diff
  }

  valid_idx <- which(results$R_lambda_g <= target_rel_err)
  if (length(valid_idx) > 0L) {
    opt_lambda <- max(results$lambda_g[valid_idx])
  } else {
    opt_lambda <- results$lambda_g[which.min(results$R_lambda_g)]
  }

  res <- list(
    optimal_lambda_g = opt_lambda,
    grid_results     = results,
    target_rel_err   = target_rel_err
  )
  class(res) <- "phmc_tune"
  return(res)
}