R/models.R

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

#' Specialized High-Level Estimators Using Proximal HMC
#'
#' Fits sparse logistic regression, nuclear-norm low-rank matrix
#' recovery, and Bayesian Lasso linear regression models using the
#' p-HMC algorithm, matching the case studies in Section VI of
#' Shukla, Vats, and Chi (2025).
#'
#' @param X Numeric design matrix (for regression models) or observed
#'   noisy matrix (for matrix recovery).
#' @param y Numeric binary response vector for logistic regression or
#'   continuous response vector for Lasso.
#' @param alpha Numeric scalar > 0. Regularization penalty parameter
#'   (default 1.0).
#' @param sigma_sq Numeric scalar > 0. Error variance parameter
#'   (default 1.0).
#' @param lambda_g Numeric scalar > 0. Moreau-Yosida envelope
#'   regularization parameter.
#' @param n_draws Integer > 0. Total number of MCMC iterations.
#' @param epsilon Numeric scalar > 0. Step size parameter.
#' @param L Integer >= 1. Number of leapfrog steps.
#' @param ... Additional arguments passed to \code{\link{phmc}}.
#'
#' @return An object of class \code{"phmc"}, which is a list containing posterior MCMC draws, parameter summary estimates matrix, log-likelihood, information criteria (AIC, BIC, DIC), acceptance rate, and execution metadata. See \code{\link{phmc}} for detailed descriptions of the list elements and output meaning.
#'
#' @references
#' Shukla A, Vats D, Chi EC (2025).
#' \dQuote{Proximal Hamiltonian Monte Carlo.}
#' \emph{arXiv preprint}, \doi{10.48550/arXiv.2510.22252}.
#'
#' @name specialized_models
#' @export
phmc_logistic <- function(X, y, alpha = 1.0, lambda_g = 0.01,
                          n_draws = 2000, epsilon = 0.002,
                          L = 10, ...) {

  X <- as.matrix(X)
  y <- as.vector(y)
  d <- ncol(X)

  ## Negative log-likelihood f(beta)
  f_logistic <- function(beta, data) {
    eta <- data$X %*% beta
    ## Numerically stable log(1 + exp(eta))
    sum(pmax(eta, 0) + log1p(exp(-abs(eta))) - data$y * eta)
  }

  ## Gradient of f(beta)
  grad_f_logistic <- function(beta, data) {
    eta <- data$X %*% beta
    prob <- 1 / (1 + exp(-eta))
    as.vector(t(data$X) %*% (prob - data$y))
  }

  ## Proximal mapping for alpha * ||beta||_1
  prox_l1_alpha <- function(beta, lambda_g) {
    prox_l1(beta, tau = alpha * lambda_g)
  }

  data_list <- list(X = X, y = y)
  start_beta <- numeric(d)
  names(start_beta) <- if (!is.null(colnames(X))) {
    colnames(X)
  } else {
    paste0("beta_", seq_len(d))
  }

  phmc(
    fn       = f_logistic,
    grad_f   = grad_f_logistic,
    prox_fn  = prox_l1_alpha,
    start    = start_beta,
    data     = data_list,
    lambda_g = lambda_g,
    n_draws  = n_draws,
    epsilon  = epsilon,
    L        = L,
    ...
  )
}

#' @rdname specialized_models
#' @export
phmc_matrix <- function(X, alpha = 1.0, sigma_sq = 1.0,
                        lambda_g = 0.0001, n_draws = 1000,
                        epsilon = 0.001, L = 10, ...) {

  Y <- as.matrix(X)
  m <- nrow(Y)
  k <- ncol(Y)

  ## Smooth potential f(X) = ||Y - X||_F^2 / (2 * sigma_sq)
  f_matrix <- function(X_mat, data) {
    0.5 * sum((data$Y - X_mat)^2) / sigma_sq
  }

  ## Gradient of f(X)
  grad_f_matrix <- function(X_mat, data) {
    (X_mat - data$Y) / sigma_sq
  }

  ## Proximal mapping for nuclear norm alpha * ||X||_*
  prox_nuclear_alpha <- function(X_mat, lambda_g) {
    prox_nuclear(X_mat, tau = alpha * lambda_g)
  }

  data_list <- list(Y = Y)
  start_mat <- matrix(0, nrow = m, ncol = k)

  phmc(
    fn       = f_matrix,
    grad_f   = grad_f_matrix,
    prox_fn  = prox_nuclear_alpha,
    start    = start_mat,
    data     = data_list,
    lambda_g = lambda_g,
    n_draws  = n_draws,
    epsilon  = epsilon,
    L        = L,
    ...
  )
}

#' @rdname specialized_models
#' @export
phmc_lasso <- function(X, y, alpha = 1.0, sigma_sq = 1.0,
                       lambda_g = 0.01, n_draws = 2000,
                       epsilon = 0.005, L = 10, ...) {

  X <- as.matrix(X)
  y <- as.vector(y)
  d <- ncol(X)

  f_lasso <- function(beta, data) {
    0.5 * sum((data$y - data$X %*% beta)^2) / sigma_sq
  }

  grad_f_lasso <- function(beta, data) {
    -as.vector(t(data$X) %*% (data$y - data$X %*% beta)) / sigma_sq
  }

  prox_l1_alpha <- function(beta, lambda_g) {
    prox_l1(beta, tau = alpha * lambda_g)
  }

  data_list <- list(X = X, y = y)
  start_beta <- numeric(d)
  names(start_beta) <- if (!is.null(colnames(X))) {
    colnames(X)
  } else {
    paste0("beta_", seq_len(d))
  }

  phmc(
    fn       = f_lasso,
    grad_f   = grad_f_lasso,
    prox_fn  = prox_l1_alpha,
    start    = start_beta,
    data     = data_list,
    lambda_g = lambda_g,
    n_draws  = n_draws,
    epsilon  = epsilon,
    L        = L,
    ...
  )
}