R/107_atoms_sum_largest.R

In CVXR: Disciplined Convex Optimization

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

## CVXPY SOURCE: atoms/sum_largest.py
## At 1.9.0: sum_largest gained axis/keepdims support (PR #3172, subclasses
## AxisAtom). The 2D axis slice (axis=NULL/1/2, keepdims, fractional k) is fully
## ported. N/A for CVXR (2D-only): CVXPY's N-D tuple-axis reduction.
## SumLargest -- sum of k largest entries, optionally along an axis


SumLargest <- new_class("SumLargest", parent = AxisAtom, package = "CVXR",
  properties = list(
    k = new_property(class = class_numeric)
  ),
  constructor = function(x, k, 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)
    k <- as.numeric(k)
    if (!is.null(axis)) {
      axis <- as.integer(axis)
      .validate_axis(axis, length(.shape(x)))
    }
    keepdims <- as.logical(keepdims)
    ## Shape via axis-aware reduction (AxisAtom convention)
    shape <- .axis_shape(.shape(x), axis, keepdims)

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

# -- validate -----------------------------------------------------
## CVXPY SOURCE: sum_largest.py:31-37 (k > 0, then AxisAtom.validate_arguments)
method(validate_arguments, SumLargest) <- function(x) {
  if (x@k <= 0) {
    cli_abort("Second argument must be a positive number.")
  }
  ## AxisAtom.validate_arguments: axis bounds + reject complex args.
  if (!is.null(x@axis)) {
    ndim <- length(.arg_shape(x))
    axis <- x@axis
    if (axis < 0L) axis <- axis + ndim + 1L
    if (axis < 1L || axis > ndim) {
      .axis_out_of_bounds_error(x@axis, ndim)
    }
  }
  if (.any_args(x, is_complex)) {
    cli_abort("Arguments to {.cls SumLargest} cannot be complex.")
  }
  invisible(NULL)
}

# -- shape --------------------------------------------------------
## Inherits AxisAtom shape_from_args (axis-reduced shape).

# -- sign: same as arg --------------------------------------------
method(sign_from_args, SumLargest) <- function(x) {
  list(is_nonneg = is_nonneg(.args(x)[[1L]]),
       is_nonpos = is_nonpos(.args(x)[[1L]]))
}

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

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

# -- PWL ----------------------------------------------------------
method(is_pwl, SumLargest) <- function(x) is_pwl(.args(x)[[1L]])

# -- get_data -----------------------------------------------------
## CVXPY SOURCE: sum_largest.py get_data -> [self.k, self.axis, self.keepdims]
method(get_data, SumLargest) <- function(x) list(x@k, x@axis, x@keepdims)

# -- numeric ------------------------------------------------------
## Sum of the k largest entries of a fiber; fractional k interpolates
## linearly onto the (floor(k)+1)-th largest entry.
.sum_largest_fiber <- function(v, k) {
  v <- as.numeric(v)
  n <- length(v)
  k_floor <- as.integer(floor(k))
  k_frac  <- k - k_floor
  sorted  <- sort(v, decreasing = TRUE)
  result  <- if (k_floor > 0L) sum(sorted[seq_len(min(k_floor, n))]) else 0
  if (k_frac > 0 && k_floor < n) {
    result <- result + k_frac * sorted[k_floor + 1L]
  }
  result
}

method(numeric_value, SumLargest) <- function(x, values, ...) {
  v <- values[[1L]]
  k <- x@k
  if (is.null(x@axis)) {
    matrix(.sum_largest_fiber(as.numeric(v), k), 1L, 1L)
  } else if (x@axis == 2L) {
    res <- apply(v, 2L, .sum_largest_fiber, k = k)
    matrix(res, nrow = 1L)
  } else {
    res <- apply(v, 1L, .sum_largest_fiber, k = k)
    matrix(res, ncol = 1L)
  }
}

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

# -- .column_grad: per-fiber subgradient --------------------------
## CVXPY SOURCE: atoms/sum_largest.py:118-141 (sum_largest._column_grad).
## 1 at each of the floor(k) largest indices; (k - floor(k)) at the
## (floor(k) + 1)-th largest. AxisAtom._axis_grad assembles fibers into
## the (input_size x output_size) sparse Jacobian.
method(.column_grad, SumLargest) <- function(x, value, ...) {
  v <- as.numeric(value)
  n <- length(v)
  k <- x@k
  k_floor <- as.integer(floor(k))
  k_frac  <- k - k_floor

  D <- numeric(n)
  ## Decreasing-order indices (stable to give CVXPY-equivalent tie-break).
  ord <- order(-v)
  if (k_floor > 0L) {
    D[ord[seq_len(min(k_floor, n))]] <- 1
  }
  if (k_frac > 0 && k_floor < n) {
    D[ord[k_floor + 1L]] <- k_frac
  }
  D
}

#' Sum of k largest entries
#'
#' @param x An Expression
#' @param k Number of largest entries to sum
#' @param axis NULL (all entries), 1 (row-wise), or 2 (column-wise)
#' @param keepdims Logical; keep the reduced dimension as size 1
#' @returns A SumLargest atom
#' @export
sum_largest <- function(x, k, axis = NULL, keepdims = FALSE) {
  SumLargest(x, k, axis = axis, keepdims = keepdims)
}