R/053_atoms_affine_cumsum.R

In CVXR: Disciplined Convex Optimization

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

## CVXPY SOURCE: atoms/affine/cumsum.py
## Cumsum -- cumulative sum along an axis (MI: AxisAtom + AffAtom -> AxisAffAtom)


Cumsum <- new_class("Cumsum", parent = AxisAffAtom, package = "CVXR",
  constructor = function(x, axis = NULL, id = NULL) {
    if (is.null(id)) id <- next_expr_id()
    x <- as_expr(x)
    ## Cumsum keeps the same shape as input (no reduction)
    shape <- x@shape

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

## Override shape_from_args: cumsum preserves shape
method(shape_from_args, Cumsum) <- function(x) {
  x@args[[1L]]@shape
}

# -- log-log: affine (CVXPY cumsum.py) ---------------------------
method(is_atom_log_log_convex, Cumsum) <- function(x) TRUE
method(is_atom_log_log_concave, Cumsum) <- function(x) TRUE

method(numeric_value, Cumsum) <- function(x, values, ...) {
  v <- values[[1L]]
  ax <- x@axis
  if (is.null(ax)) {
    ## Cumsum over all entries (column-major order)
    matrix(cumsum(as.vector(v)), nrow = nrow(v), ncol = ncol(v))
  } else if (ax == 2L) {
    ## Cumsum down columns (column-wise)
    apply(v, 2L, cumsum)
  } else {
    ## Cumsum across rows
    t(apply(v, 1L, cumsum))
  }
}

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

#' Cumulative sum along an axis
#'
#' @param x An Expression
#' @param axis NULL (all), 1 (across rows), or 2 (down columns)
#' @returns A Cumsum atom
#' @export
cumsum_axis <- function(x, axis = NULL) {
  Cumsum(x, axis)
}