R/034_atoms_affine_affine_atom.R

In CVXR: Disciplined Convex Optimization

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

## CVXPY SOURCE: atoms/affine/affine_atom.py
## AffAtom -- abstract base class for affine atoms
##
## Affine atoms are both convex and concave, and allow complex arguments.
## Sign propagation follows sum_signs logic by default.


AffAtom <- new_class("AffAtom", parent = Atom, package = "CVXR",
  constructor = function(args, shape, id = NULL) {
    if (FALSE) new_object(S7_object())  ## S7 static-check guard
    if (is.null(id)) id <- next_expr_id()
    args <- lapply(args, as_expr)
    if (length(args) == 0L) {
      cli_abort("No arguments given to {.cls AffAtom}.")
    }
    shape <- validate_shape(shape)
    obj <- .fast_new(AffAtom, S7_object(),
      id    = as.integer(id),
      .cache = new.env(parent = emptyenv()),
      args  = args,
      shape = shape
    )
    validate_arguments(obj)
    obj
  }
)

# -- validate_arguments: AffAtom allows complex ----------------------
## CVXPY SOURCE: affine_atom.py inherits from Atom with _allow_complex = True

method(validate_arguments, AffAtom) <- function(x) {
  ## AffAtom allows complex arguments -- no validation needed
  invisible(NULL)
}

# -- sign: sum_signs of args -----------------------------------------
## CVXPY SOURCE: affine_atom.py lines 33-36

method(sign_from_args, AffAtom) <- function(x) {
  sum_signs(.args(x))
}

# -- Complex propagation ---------------------------------------------
## CVXPY SOURCE: affine_atom.py lines 38-48

method(is_imag, AffAtom) <- function(x) {
  .all_args(x, is_imag)
}

method(is_complex, AffAtom) <- function(x) {
  .any_args(x, is_complex)
}

# -- Convexity/concavity: affine is both -----------------------------
## CVXPY SOURCE: affine_atom.py lines 50-58

method(is_atom_convex, AffAtom) <- function(x) TRUE
method(is_atom_concave, AffAtom) <- function(x) TRUE

## CVXPY affine_atom.py: affine atoms are smooth.
method(is_atom_smooth, AffAtom) <- function(x) TRUE

# -- Monotonicity: default increasing --------------------------------
## CVXPY SOURCE: affine_atom.py lines 60-70

method(is_incr, AffAtom) <- function(x, idx, ...) TRUE
method(is_decr, AffAtom) <- function(x, idx, ...) FALSE

# -- Quadratic / PWL analysis ---------------------------------------
## CVXPY SOURCE: affine_atom.py lines 72-88

method(is_quadratic, AffAtom) <- function(x) {
  .all_args(x, is_quadratic)
}

method(has_quadratic_term, AffAtom) <- function(x) {
  .any_args(x, has_quadratic_term)
}

method(is_qpwa, AffAtom) <- function(x) {
  .all_args(x, is_qpwa)
}

method(is_pwl, AffAtom) <- function(x) {
  .all_args(x, is_pwl)
}

# -- PSD/NSD propagation ---------------------------------------------
## CVXPY SOURCE: affine_atom.py lines 91-109
## For affine atoms, PSD/NSD propagates through monotonicity:
## is_psd: all args satisfy (incr(idx) and arg.is_psd()) or (decr(idx) and arg.is_nsd())
## is_nsd: mirror

method(is_psd, AffAtom) <- function(x) {
  cached <- cache_get(x, "is_psd")
  if (!cache_miss(cached)) return(cached)
  result <- .affatom_is_psd(x)
  cache_set(x, "is_psd", result)
  result
}

.affatom_is_psd <- function(x) {
  for (idx in seq_along(.args(x))) {
    arg <- .args(x)[[idx]]
    if (!((is_incr(x, idx) && is_psd(arg)) ||
          (is_decr(x, idx) && is_nsd(arg)))) {
      return(FALSE)
    }
  }
  TRUE
}

method(is_nsd, AffAtom) <- function(x) {
  cached <- cache_get(x, "is_nsd")
  if (!cache_miss(cached)) return(cached)
  result <- .affatom_is_nsd(x)
  cache_set(x, "is_nsd", result)
  result
}

.affatom_is_nsd <- function(x) {
  for (idx in seq_along(.args(x))) {
    arg <- .args(x)[[idx]]
    if (!((is_decr(x, idx) && is_psd(arg)) ||
          (is_incr(x, idx) && is_nsd(arg)))) {
      return(FALSE)
    }
  }
  TRUE
}

# -- .grad: per-atom hook for affine atoms ----------------------------
## CVXPY SOURCE: affine_atom.py:111-166 (AffAtom._grad).
##
## Strategy (line-for-line port):
##   1. Build a *fake* LinOp tree by re-running graph_implementation()
##      with placeholder Variable LinOps for each non-constant arg
##      (and the canonical-form constant LinOp for each constant arg).
##   2. Call canonInterface.get_problem_matrix() on the fake LinOp,
##      passing the per-arg column offsets via id_to_col. This returns
##      a sparse triplet (V, I, J) representing the affine map's
##      coefficient matrix at the LinOp level.
##   3. Reconstruct the (var_length, self.size) sparse matrix by
##      treating the V/I/J as transposed (var-row, output-col) entries.
##   4. Slice into per-arg Jacobians.
##
## With this in place, every affine atom (Transpose, Sum, Reshape,
## Vec, CumSum, DivExpression, Diag, Trace, ...) inherits a working
## .grad from this method. The chain rule walker handles composition.

method(.grad, AffAtom) <- function(x, values, ...) {
  ## Step 1 -- build fake_args + var_offsets.
  n_args <- length(.args(x))
  fake_args   <- vector("list", n_args)
  var_offsets <- integer(0)   # named integer vector keyed by var-id
  offset      <- 0L

  for (idx in seq_len(n_args)) {
    arg <- .args(x)[[idx]]
    if (is_constant(arg)) {
      ## CVXPY: Constant(arg.value).canonical_form[0]
      arg_val <- value(arg)
      const_lin <- canonical_form(Constant(arg_val))[[1L]]
      fake_args[[idx]] <- const_lin
    } else {
      ## CVXPY: lu.create_var(arg.shape, idx)
      ## We use idx as the var_id (per-call unique).
      fake_args[[idx]] <- create_var(.shape(arg), as.integer(idx))
      var_offsets[as.character(idx)] <- offset
      offset <- offset + as.integer(prod(.shape(arg)))
    }
  }
  var_length <- offset

  ## Step 2 -- run graph_implementation on the fakes; second-element
  ## constraint list is intentionally discarded (mirrors CVXPY).
  fake_pair <- graph_implementation(x, fake_args, .shape(x), get_data(x))
  fake_expr <- fake_pair[[1L]]

  ## Step 3 -- extract the coefficient matrix.
  pm <- get_problem_matrix(
    list(fake_expr),
    id_to_col = var_offsets,
    var_length = var_length
  )

  ## CVXPY's reshape(canon_mat, (var_length+1, self.size))[:-1, :]
  ## is equivalent to building the sparse (var_length, self.size)
  ## matrix directly from V/I/J, keeping only var-column entries.
  self_size <- as.integer(prod(.shape(x)))
  if (var_length > 0L && length(pm$V) > 0L) {
    var_mask <- pm$J < var_length    # 0-based; var cols are J in [0, var_length)
    stacked_grad <- Matrix::sparseMatrix(
      i = pm$J[var_mask] + 1L,        # var rank -> R 1-based row
      j = pm$I[var_mask] + 1L,        # output entry -> R 1-based col
      x = pm$V[var_mask],
      dims = c(var_length, self_size),
      repr = "C"
    )
  } else {
    stacked_grad <- Matrix::sparseMatrix(
      i = integer(0), j = integer(0), x = numeric(0),
      dims = c(max(var_length, 1L), self_size), repr = "C"
    )
  }

  ## Step 4 -- slice into per-arg blocks.
  grad_list <- vector("list", n_args)
  for (idx in seq_len(n_args)) {
    arg <- .args(x)[[idx]]
    arg_size <- as.integer(prod(.shape(arg)))
    if (is_constant(arg)) {
      ## CVXPY: zero placeholder (scalar 0 if 1x1; sparse zero otherwise).
      ## CVXR's chain-rule walker ignores Constant args' grads (their
      ## variables() is empty), so the value here is irrelevant; we
      ## return a sparse zero of the matching shape for consistency.
      if (arg_size == 1L && self_size == 1L) {
        grad_list[[idx]] <- 0
      } else {
        grad_list[[idx]] <- Matrix::sparseMatrix(
          i = integer(0), j = integer(0), x = numeric(0),
          dims = c(arg_size, self_size), repr = "C"
        )
      }
    } else {
      start_row <- var_offsets[as.character(idx)] + 1L
      end_row   <- start_row + arg_size - 1L
      grad_list[[idx]] <- stacked_grad[start_row:end_row, , drop = FALSE]
    }
  }
  grad_list
}