R/060_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",
  ## cumsum.py:75 `axis = 0` (down the columns) is the default, as for
  ## Cumprod and Cummax.
  constructor = function(x, axis = 2L, id = NULL) {
    if (FALSE) new_object(S7_object())  ## S7 static-check guard
    if (is.null(id)) id <- next_expr_id()
    x <- as_expr(x)
    ## Cumsum keeps the same shape as input along an axis; with axis = NULL
    ## it is a column of all entries (see shape_from_args).
    shape <- .cumsum_shape(.shape(x), axis)

    obj <- .fast_new(Cumsum, S7_object(),
      id       = as.integer(id),
      .cache   = new.env(parent = emptyenv()),
      args     = list(x),
      shape    = shape,
      axis     = axis,
      keepdims = FALSE
    )
    ## AxisAffAtom's axis-range check (as Cumprod does). Cumsum skipped it, so
    ## any axis was accepted and failed later in the canonicalizer (1.9.3
    ## plan defect 0.1).
    validate_arguments(obj)
    obj
  }
)

## CVXPY SOURCE: atoms/affine/cumsum.py:41-57 (_cumsum_axis_op) -- the sparse
## operator for cumsum along one axis, acting on the column-major vec of an
## array with `shape`. In R axis terms on an m x n array: axis 2 (down
## columns, CVXPY axis 0) is kron(I_n, L_m); axis 1 (across rows, CVXPY
## axis 1) is kron(L_n, I_m), with L the lower-triangular matrix of ones.
## Consumer: cumsum_canon's parameter-affine branch (CVXPY v1.9.3, #3398).
.cumsum_axis_op <- function(shape, axis) {
  m <- shape[1L]
  n <- shape[2L]
  tril <- function(k) {
    idx <- which(lower.tri(matrix(0, k, k), diag = TRUE), arr.ind = TRUE)
    Matrix::sparseMatrix(i = idx[, 1L], j = idx[, 2L], x = 1, dims = c(k, k))
  }
  if (axis == 2L) {
    if (n == 1L) return(tril(m))
    Matrix::kronecker(Matrix::Diagonal(n), tril(m))
  } else {
    if (m == 1L) return(tril(n))
    Matrix::kronecker(tril(n), Matrix::Diagonal(m))
  }
}

## CVXPY SOURCE: cumsum.py shape_from_args -- the input shape along an axis,
## `(size,)` for axis None. R deviation (1.9.3 plan decision D2, defect 0.2):
## axis = NULL gives the column c(size, 1) in COLUMN-MAJOR order, as base R's
## cumsum() does; CVXPY flattens in C order. Pre-fix the atom kept the input
## shape while the canonicalizer flattened to c(size, 1) in C order, so the
## solver and value() disagreed on matrices.
.cumsum_shape <- function(arg_shape, axis) {
  if (is.null(axis)) c(as.integer(prod(arg_shape)), 1L) else arg_shape
}

method(shape_from_args, Cumsum) <- function(x) {
  .cumsum_shape(.arg_shape(x), x@axis)
}

## No log-log hooks: cumsum.py defines none, so Atom's FALSE applies
## (atom.py:195-202). A TRUE/TRUE pair here made cumsum(x) log-log affine;
## `min sum(cumsum(x)) s.t. x1 x2 >= 1` solved under gp = TRUE to 3e8 against
## a true 2.83, where CVXPY raises DGPError (sweep item 5, finding 5.1).

method(numeric_value, Cumsum) <- function(x, values, ...) {
  v <- values[[1L]]
  ax <- x@axis
  if (is.null(ax)) {
    ## Cumsum over all entries, column-major, as a column (see .cumsum_shape)
    matrix(cumsum(as.vector(v)), ncol = 1L)
  } 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 2 (down columns, the default, as CVXPY's \code{axis = 0}),
#'   1 (across rows), or NULL (all entries, in column-major order, as a
#'   column)
#' @returns A Cumsum atom
#' @export
cumsum_axis <- function(x, axis = 2L) {
  Cumsum(x, axis)
}