R/proximal.R

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

#' Proximal Mapping Operators and Moreau-Yosida Envelope Gradients
#'
#' Evaluates proximal operators for common non-smooth penalties (L1 norm,
#' nuclear norm, L2 norm, elastic net) and computes the gradient of the
#' Moreau-Yosida envelope as described in Shukla, Vats, and Chi (2025).
#'
#' @param x Numeric vector or matrix. Parameter value at which to evaluate
#'   the operator.
#' @param tau Numeric scalar. Thresholding parameter (typically
#'   \code{lambda_g * scale}).
#' @param alpha Numeric scalar in [0, 1]. Mixing parameter for elastic net
#'   penalty.
#' @param prox_fn Function or character string. Proximal mapping function or
#'   name of built-in proximal operator (\code{"l1"}, \code{"l2"},
#'   \code{"elastic_net"}, \code{"nuclear_norm"}, \code{"none"}).
#' @param lambda_g Numeric scalar > 0. Moreau-Yosida regularization
#'   parameter.
#' @param ... Additional arguments passed to the proximal operator.
#'
#' @return For \code{prox_l1}, \code{prox_l2}, \code{prox_elastic_net}, and
#'   \code{prox_nuclear}, returns a numeric vector or matrix of class \code{"numeric"}
#'   or \code{"matrix"} (matching the shape and dimensions of input \code{x})
#'   representing the evaluated proximal point operator. For \code{grad_my_envelope},
#'   returns a numeric vector or matrix of class \code{"numeric"} or \code{"matrix"}
#'   containing the computed gradient of the Moreau-Yosida envelope at \code{x}.
#'
#' @details The Moreau-Yosida envelope of a proper, lower-semicontinuous,
#'   convex function \eqn{g} with scaling parameter
#'   \eqn{\lambda_g > 0} is defined as
#'   \deqn{g_{\lambda_g}(x) = \inf_y \{g(y) +
#'     \frac{1}{2\lambda_g} \|y - x\|^2\}.}
#'   Its gradient is given by
#'   \deqn{\nabla g_{\lambda_g}(x) =
#'     \frac{1}{\lambda_g}(x - \mathrm{prox}_{\lambda_g}^{g}(x)).}
#'
#' @references
#' Shukla A, Vats D, Chi EC (2025).
#' \dQuote{Proximal Hamiltonian Monte Carlo.}
#' \emph{arXiv preprint}, \doi{10.48550/arXiv.2510.22252}.
#'
#' @importFrom Matrix Matrix
#' @name proximal_operators
#' @export
prox_l1 <- function(x, tau) {
  sign(x) * pmax(abs(x) - tau, 0)
}

#' @rdname proximal_operators
#' @export
prox_l2 <- function(x, tau) {
  norm_x <- sqrt(sum(x^2))
  if (norm_x <= tau || norm_x == 0) {
    return(numeric(length(x)))
  }
  (1 - tau / norm_x) * x
}

#' @rdname proximal_operators
#' @export
prox_elastic_net <- function(x, tau, alpha = 0.5) {
  tau_l1 <- tau * alpha
  tau_l2 <- tau * (1 - alpha)
  x_soft <- prox_l1(x, tau_l1)
  x_soft / (1 + tau_l2)
}

#' @rdname proximal_operators
#' @export
prox_nuclear <- function(x, tau) {
  if (!is.matrix(x)) {
    n_side <- as.integer(round(sqrt(length(x))))
    x <- matrix(x, nrow = n_side, ncol = n_side)
  }
  mat_obj <- as.matrix(Matrix::Matrix(x))
  svd_decomp <- svd(mat_obj)
  d_soft <- pmax(svd_decomp$d - tau, 0)
  if (all(d_soft == 0)) {
    res <- matrix(0, nrow = nrow(x), ncol = ncol(x))
  } else {
    pos_idx <- which(d_soft > 0)
    if (length(pos_idx) == 1L) {
      res <- svd_decomp$u[, 1L, drop = FALSE] %*%
        (d_soft[1L] * t(svd_decomp$v[, 1L, drop = FALSE]))
    } else {
      res <- svd_decomp$u[, pos_idx, drop = FALSE] %*%
        diag(d_soft[pos_idx], nrow = length(pos_idx)) %*%
        t(svd_decomp$v[, pos_idx, drop = FALSE])
    }
  }
  return(res)
}

#' @rdname proximal_operators
#' @export
grad_my_envelope <- function(x, prox_fn = "l1", lambda_g = 0.01, ...) {
  if (is.character(prox_fn)) {
    prox_type <- match.arg(tolower(prox_fn),
      c("l1", "l2", "elastic_net", "nuclear_norm", "none"))
    if (prox_type == "none") {
      if (is.matrix(x)) return(matrix(0, nrow = nrow(x), ncol = ncol(x)))
      return(numeric(length(x)))
    }
    fn <- switch(prox_type,
      "l1"           = function(p) prox_l1(p, tau = lambda_g),
      "l2"           = function(p) prox_l2(p, tau = lambda_g),
      "elastic_net"  = function(p) prox_elastic_net(p, tau = lambda_g, ...),
      "nuclear_norm" = function(p) prox_nuclear(p, tau = lambda_g))
  } else if (is.function(prox_fn)) {
    fn <- function(p) prox_fn(p, lambda_g = lambda_g, ...)
  } else {
    stop("'prox_fn' must be a character string or a function.")
  }

  prox_val <- fn(x)
  grad <- (x - prox_val) / lambda_g
  return(grad)
}