R/106_atoms_log_sum_exp.R

In CVXR: Disciplined Convex Optimization

#####
## DO NOT EDIT THIS FILE!! EDIT THE SOURCE INSTEAD: rsrc_tree/atoms/log_sum_exp.R
#####

## CVXPY SOURCE: atoms/log_sum_exp.py
## LogSumExp -- log(sum(exp(x))), axis-aware


LogSumExp <- new_class("LogSumExp", parent = AxisAtom, package = "CVXR",
  constructor = function(x, axis = NULL, keepdims = FALSE, id = NULL) {
    if (FALSE) new_object(S7_object())  ## S7 static-check guard
    if (is.null(id)) id <- next_expr_id()
    x <- as_expr(x)
    if (!is.null(axis)) .validate_axis(axis, length(.shape(x)))
    shape <- .axis_shape(.shape(x), axis, keepdims)

    obj <- .fast_new(LogSumExp, S7_object(),
      id       = as.integer(id),
      .cache   = new.env(parent = emptyenv()),
      args     = list(x),
      shape    = shape,
      axis     = axis,
      keepdims = keepdims
    )
    validate_arguments(obj)
    obj
  }
)

# -- sign: nonneg when arg is nonneg ------------------------------
method(sign_from_args, LogSumExp) <- function(x) {
  list(is_nonneg = is_nonneg(.args(x)[[1L]]),
       is_nonpos = FALSE)
}

# -- curvature: convex --------------------------------------------
method(is_atom_convex, LogSumExp) <- function(x) TRUE
method(is_atom_concave, LogSumExp) <- function(x) FALSE

## CVXPY log_sum_exp.py: log_sum_exp is smooth.
method(is_atom_smooth, LogSumExp) <- function(x) TRUE

# -- monotonicity: always increasing ------------------------------
method(is_incr, LogSumExp) <- function(x, idx, ...) TRUE
method(is_decr, LogSumExp) <- function(x, idx, ...) FALSE

# -- numeric: log(sum(exp(x))) ------------------------------------
method(numeric_value, LogSumExp) <- function(x, values, ...) {
  v <- values[[1L]]
  axis <- x@axis
  keepdims <- x@keepdims

  if (is.null(axis)) {
    ## Reduce all: scalar
    ## Numerically stable: max + log(sum(exp(v - max)))
    mx <- max(v)
    result <- mx + log(sum(exp(v - mx)))
    matrix(result, 1L, 1L)
  } else if (axis == 2L) {
    ## axis=2: Reduce along rows (column-wise) -> one value per column
    mx <- apply(v, 2L, max)
    result <- mx + log(colSums(exp(sweep(v, 2L, mx))))
    if (keepdims) matrix(result, nrow = 1L) else matrix(result, nrow = 1L)
  } else {
    ## axis == 1: Reduce along columns -> one value per row
    mx <- apply(v, 1L, max)
    result <- mx + log(rowSums(exp(sweep(v, 1L, mx))))
    if (keepdims) matrix(result, ncol = 1L) else matrix(result, ncol = 1L)
  }
}

# -- graph_implementation: stub -----------------------------------
method(graph_implementation, LogSumExp) <- function(x, arg_objs, shape, data = NULL, ...) {
  cli_abort("graph_implementation for {.cls LogSumExp} not yet implemented.")
}

# -- .column_grad: softmax (numerically stable) -------------------
## CVXPY SOURCE: atoms/log_sum_exp.py:64-79 (log_sum_exp._column_grad).
## d/dx_i log(sum(exp(x))) = exp(x_i) / sum(exp(x_k)) = softmax.
## CVXPY uses scipy.special.logsumexp; we use a manual shift-and-sum
## for the same numerical stability.
method(.column_grad, LogSumExp) <- function(x, value, ...) {
  v <- as.numeric(value)
  m <- max(v)
  e <- exp(v - m)
  e / sum(e)
}

#' Log-sum-exp: log(sum(exp(x)))
#'
#' @param x An Expression
#' @param axis NULL (all), 1 (row-wise), or 2 (column-wise)
#' @param keepdims Logical: keep reduced dimensions?
#' @returns A LogSumExp atom
#' @export
log_sum_exp <- function(x, axis = NULL, keepdims = FALSE) {
  LogSumExp(x, axis, keepdims)
}