R/034_atoms_atom.R

In CVXR: Disciplined Convex Optimization

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

## CVXPY SOURCE: atoms/atom.py
## Atom -- abstract base class for all atoms (operations on expressions)
##
## An Atom represents a mathematical function applied to one or more
## Expression arguments. The DCP composition rules are implemented here.
##
## CVXPY 1.9 parity notes:
##  - DNLP atom predicates mirror atom.py:191-298: `is_atom_smooth` defaults to
##    FALSE, and the linearizable-convex / linearizable-concave composition
##    rules accept smooth children or monotone-compatible linearizable children.


Atom <- new_class("Atom", parent = Expression, 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 Atom}.")
    }
    shape <- validate_shape(shape)
    obj <- .fast_new(Atom, S7_object(),
      id    = as.integer(id),
      .cache = new.env(parent = emptyenv()),
      args  = args,
      shape = shape
    )
    validate_arguments(obj)
    obj
  }
)

# -- validate_arguments ----------------------------------------------
## CVXPY SOURCE: atom.py lines 92-98
## Default: rejects complex arguments unless _allow_complex is TRUE.
## Subclasses override via validate_arguments generic.

method(validate_arguments, Atom) <- function(x) {
  if (.any_args(x, is_complex)) {
    cli_abort("Arguments to {.cls {class(x)[[1L]]}} cannot be complex.")
  }
  invisible(NULL)
}

# -- name -------------------------------------------------------------
## CVXPY SOURCE: atom.py lines 57-64

method(expr_name, Atom) <- function(x) {
  data <- get_data(x)
  data_str <- if (is.null(data)) character(0) else vapply(data, as.character, character(1))
  arg_strs <- vapply(.args(x), expr_name, character(1))
  sprintf("%s(%s)", class(x)[[1L]], paste(c(arg_strs, data_str), collapse = ", "))
}

## CVXPY SOURCE: atom.py Atom.format_labeled (recursive default).
## Mirrors expr_name(Atom) but recurses with format_labeled() on each
## argument so a labelled sub-expression renders as its label.
method(format_labeled, Atom) <- function(x) {
  lbl <- label(x)
  if (!is.null(lbl)) return(lbl)
  data <- get_data(x)
  data_str <- if (is.null(data)) character(0) else vapply(data, as.character, character(1))
  arg_strs <- vapply(.args(x), format_labeled, character(1))
  sprintf("%s(%s)", class(x)[[1L]], paste(c(arg_strs, data_str), collapse = ", "))
}

# -- Sign: cached, from sign_from_args --------------------------------
## CVXPY SOURCE: atom.py lines 148-158

## Every cached member goes through compute_once (utilities/performance_utils.R),
## the one mirror of CVXPY's @perf.compute_once; see that file for the key rule.
method(is_nonneg, Atom) <- function(x) compute_once(x, "is_nonneg", .atom_is_nonneg)
.atom_is_nonneg <- function(x) sign_from_args(x)[["is_nonneg"]]

method(is_nonpos, Atom) <- function(x) compute_once(x, "is_nonpos", .atom_is_nonpos)
.atom_is_nonpos <- function(x) sign_from_args(x)[["is_nonpos"]]

# -- Bounds propagation (interval arithmetic) ------------------------
## CVXPY SOURCE: atom.py:116-146 (#3080).
## bounds_from_args: an atom's interval bounds from its arguments' bounds.
## Default is unbounded; atoms override with the relevant utilities/bounds.R
## helper (Batches 3-4). get_bounds combines that with sign info and is cached
## (CVXPY @compute_once).

method(bounds_from_args, Atom) <- function(x) {
  unbounded(.shape(x))
}

method(get_bounds, Atom) <- function(x) compute_once(x, "get_bounds", .atom_get_bounds)

.atom_get_bounds <- function(x) {
  ba <- bounds_from_args(x)
  refine_bounds_from_sign(ba[[1L]], ba[[2L]], is_nonneg(x), is_nonpos(x))
}

# -- Complex: default FALSE ------------------------------------------
## CVXPY SOURCE: atom.py lines 160-172

method(is_imag, Atom) <- function(x) FALSE
method(is_complex, Atom) <- function(x) FALSE

# -- is_atom_affine: plain function ----------------------------------
## CVXPY SOURCE: atom.py lines 186-189

is_atom_affine <- function(x) is_atom_convex(x) && is_atom_concave(x)

# -- Log-log DGP atom hooks: default FALSE -------------------------
## CVXPY SOURCE: atom.py -- default is False
method(is_atom_log_log_convex, Atom) <- function(x) FALSE
method(is_atom_log_log_concave, Atom) <- function(x) FALSE

## is_atom_log_log_affine: convenience (CVXPY atom.py lines 215-218)
is_atom_log_log_affine <- function(x) {
  is_atom_log_log_concave(x) && is_atom_log_log_convex(x)
}

# -- DCP composition: is_convex / is_concave -------------------------
## CVXPY SOURCE: atom.py lines 232-264
## idx for is_incr/is_decr: 1-based (R convention)

method(is_convex, Atom) <- function(x) compute_once(x, "is_convex", .atom_is_convex)

.atom_is_convex <- function(x) {
  if (is_constant(x)) return(TRUE)
  if (!is_atom_convex(x)) return(FALSE)
  for (idx in seq_along(.args(x))) {
    arg <- .args(x)[[idx]]
    if (!(is_affine(arg) ||
          (is_convex(arg) && is_incr(x, idx)) ||
          (is_concave(arg) && is_decr(x, idx)))) {
      return(FALSE)
    }
  }
  TRUE
}

method(is_concave, Atom) <- function(x) compute_once(x, "is_concave", .atom_is_concave)

.atom_is_concave <- function(x) {
  if (is_constant(x)) return(TRUE)
  if (!is_atom_concave(x)) return(FALSE)
  for (idx in seq_along(.args(x))) {
    arg <- .args(x)[[idx]]
    if (!(is_affine(arg) ||
          (is_concave(arg) && is_incr(x, idx)) ||
          (is_convex(arg) && is_decr(x, idx)))) {
      return(FALSE)
    }
  }
  TRUE
}

# -- DNLP: linearizable-convex / linearizable-concave composition --
## CVXPY SOURCE: atom.py lines 191-193, 266-298
## idx for is_incr/is_decr: 1-based (R convention)

## is_atom_smooth: default FALSE. Smooth atoms (trig/hyperbolic/normcdf,
## a later 1.9 slice) override this to TRUE. CVXPY SOURCE: atom.py:191-193
method(is_atom_smooth, Atom) <- function(x) FALSE

method(is_linearizable_convex, Atom) <- function(x) {
  compute_once(x, "is_linearizable_convex", .atom_is_linearizable_convex)
}

.atom_is_linearizable_convex <- function(x) {
  if (is_constant(x)) return(TRUE)
  if (!(is_atom_smooth(x) || is_atom_convex(x))) return(FALSE)
  for (idx in seq_along(.args(x))) {
    arg <- .args(x)[[idx]]
    if (!(is_smooth(arg) ||
          (is_linearizable_convex(arg) && is_incr(x, idx)) ||
          (is_linearizable_concave(arg) && is_decr(x, idx)))) {
      return(FALSE)
    }
  }
  TRUE
}

method(is_linearizable_concave, Atom) <- function(x) {
  compute_once(x, "is_linearizable_concave", .atom_is_linearizable_concave)
}

.atom_is_linearizable_concave <- function(x) {
  if (is_constant(x)) return(TRUE)
  if (!(is_atom_smooth(x) || is_atom_concave(x))) return(FALSE)
  for (idx in seq_along(.args(x))) {
    arg <- .args(x)[[idx]]
    if (!(is_smooth(arg) ||
          (is_linearizable_concave(arg) && is_incr(x, idx)) ||
          (is_linearizable_convex(arg) && is_decr(x, idx)))) {
      return(FALSE)
    }
  }
  TRUE
}

# -- DQCP: quasiconvex / quasiconcave composition ------------------
## CVXPY SOURCE: atom.py lines 345-406

## Helper: indices of non-constant arguments (1-based, R convention)
.non_const_idx <- function(x) {
  idx <- integer(0)
  for (i in seq_along(.args(x))) {
    if (!is_constant(.args(x)[[i]])) idx <- c(idx, i)
  }
  idx
}

## Helper: is this a "real" function? (scalar in, scalar out, one non-const arg)
## CVXPY SOURCE: atom.py lines 349-356
.is_real_fn <- function(x) {
  non_const <- .non_const_idx(x)
  expr_is_scalar(x) &&
    length(non_const) == 1L &&
    expr_is_scalar(.args(x)[[non_const[1L]]])
}

## Atom-level quasiconvex/quasiconcave hooks (defaults: fall back to DCP)
## CVXPY SOURCE: atom.py lines 205-213
method(is_atom_quasiconvex, Atom) <- function(x) is_atom_convex(x)
method(is_atom_quasiconcave, Atom) <- function(x) is_atom_concave(x)

## Quasiconvex composition (CVXPY atom.py lines 358-381)
method(is_quasiconvex, Atom) <- function(x) {
  compute_once(x, "is_quasiconvex", .atom_is_quasiconvex)
}

.atom_is_quasiconvex <- function(x) {
  if (is_convex(x)) return(TRUE)
  ## Maximum / MaxEntries: all args quasiconvex
  if (.s7_is(x, Maximum) || .s7_is(x, MaxEntries)) {
    return(.all_args(x, is_quasiconvex))
  }
  ## Real function with monotone argument
  non_const <- .non_const_idx(x)
  if (.is_real_fn(x) && is_incr(x, non_const[1L])) {
    return(is_quasiconvex(.args(x)[[non_const[1L]]]))
  }
  if (.is_real_fn(x) && is_decr(x, non_const[1L])) {
    return(is_quasiconcave(.args(x)[[non_const[1L]]]))
  }
  ## Atom-level quasiconvex with DCP-like arg checks
  if (is_atom_quasiconvex(x)) {
    for (idx in seq_along(.args(x))) {
      arg <- .args(x)[[idx]]
      if (!(is_affine(arg) ||
            (is_convex(arg) && is_incr(x, idx)) ||
            (is_concave(arg) && is_decr(x, idx)))) {
        return(FALSE)
      }
    }
    return(TRUE)
  }
  FALSE
}

## Quasiconcave composition (CVXPY atom.py lines 383-406)
method(is_quasiconcave, Atom) <- function(x) {
  compute_once(x, "is_quasiconcave", .atom_is_quasiconcave)
}

.atom_is_quasiconcave <- function(x) {
  if (is_concave(x)) return(TRUE)
  ## Minimum / MinEntries: all args quasiconcave
  if (.s7_is(x, Minimum) || .s7_is(x, MinEntries)) {
    return(.all_args(x, is_quasiconcave))
  }
  ## Real function with monotone argument
  non_const <- .non_const_idx(x)
  if (.is_real_fn(x) && is_incr(x, non_const[1L])) {
    return(is_quasiconcave(.args(x)[[non_const[1L]]]))
  }
  if (.is_real_fn(x) && is_decr(x, non_const[1L])) {
    return(is_quasiconvex(.args(x)[[non_const[1L]]]))
  }
  ## Atom-level quasiconcave with DCP-like arg checks
  if (is_atom_quasiconcave(x)) {
    for (idx in seq_along(.args(x))) {
      arg <- .args(x)[[idx]]
      if (!(is_affine(arg) ||
            (is_concave(arg) && is_incr(x, idx)) ||
            (is_convex(arg) && is_decr(x, idx)))) {
        return(FALSE)
      }
    }
    return(TRUE)
  }
  FALSE
}

# -- Log-log DCP -----------------------------------------------------
## CVXPY SOURCE: atom.py lines 311-342

## is_log_log_constant -> is_constant and the per-atom DPP rules (e.g.
## gmatmul) answer differently in scope; a plain key returned whichever
## answer was computed first (fixed 1.9.2.9125, generalized in 9126).
method(is_log_log_convex, Atom) <- function(x) {
  compute_once(x, "is_log_log_convex", .atom_is_log_log_convex)
}

.atom_is_log_log_convex <- function(x) {
  if (is_log_log_constant(x)) return(TRUE)
  if (!is_atom_log_log_convex(x)) return(FALSE)
  ## Check args satisfy DGP composition (CVXPY atom.py lines 318-323)
  for (i in seq_along(.args(x))) {
    arg <- .args(x)[[i]]
    if (!(is_log_log_affine(arg) ||
          (is_log_log_convex(arg) && is_incr(x, i)) ||
          (is_log_log_concave(arg) && is_decr(x, i)))) {
      return(FALSE)
    }
  }
  TRUE
}

method(is_log_log_concave, Atom) <- function(x) {
  compute_once(x, "is_log_log_concave", .atom_is_log_log_concave)
}

.atom_is_log_log_concave <- function(x) {
  if (is_log_log_constant(x)) return(TRUE)
  if (!is_atom_log_log_concave(x)) return(FALSE)
  ## Check args satisfy DGP composition (CVXPY atom.py lines 335-340)
  for (i in seq_along(.args(x))) {
    arg <- .args(x)[[i]]
    if (!(is_log_log_affine(arg) ||
          (is_log_log_concave(arg) && is_incr(x, i)) ||
          (is_log_log_convex(arg) && is_decr(x, i)))) {
      return(FALSE)
    }
  }
  TRUE
}

# -- canonicalize ----------------------------------------------------
## CVXPY SOURCE: atom.py lines 408-427

method(canonicalize, Atom) <- function(x) {
  ## Constant atoms (with no parameters) -> wrap as Constant
  if (is_constant(x) && length(parameters(x)) == 0L) {
    return(canonical_form(Constant(value(x))))
  }
  ## Recursively canonicalize args
  arg_objs <- vector("list", length(.args(x)))
  constraints <- list()
  for (i in seq_along(.args(x))) {
    cf <- canonical_form(.args(x)[[i]])
    arg_objs[[i]] <- cf[[1L]]
    constraints <- c(constraints, cf[[2L]])
  }
  ## Graph implementation
  data <- get_data(x)
  gi <- graph_implementation(x, arg_objs, .shape(x), data)
  list(gi[[1L]], c(constraints, gi[[2L]]))
}

# -- value -----------------------------------------------------------
## CVXPY SOURCE: atom.py lines 450-474

method(value, Atom) <- function(x) {
  ## Check if any parameter lacks a value
  params <- parameters(x)
  if (length(params) > 0L) {
    for (p in params) {
      if (is.null(value(p))) return(NULL)
    }
  }
  ## Evaluate
  .atom_value_impl(x)
}

## Internal value implementation (recursive)
## CVXPY SOURCE: atom.py lines 456-474
.atom_value_impl <- function(x) {
  if (0L %in% .shape(x)) return(numeric(0))
  arg_values <- vector("list", length(.args(x)))
  for (i in seq_along(.args(x))) {
    arg_val <- value(.args(x)[[i]])
    ## A NULL arg makes higher-level values NULL, unless atom is constant
    if (is.null(arg_val) && !is_constant(x)) return(NULL)
    arg_values[[i]] <- arg_val
  }
  numeric_value(x, arg_values)
}

# -- grad: chain-rule walker ------------------------------------------
## CVXPY SOURCE: atom.py:477-524 (Atom.grad property).
## Walks the expression DAG, calling .grad(x, values) at each atom to get
## per-arg Jacobians, and composes via the chain rule against each child
## Expression's grad() (which is recursive). Returns a list keyed by
## variable id (as.character) of sparse Jacobians of shape
## (prod(var@shape), prod(x@shape)).

.is_scalar <- function(z) (is.numeric(z) || is.complex(z)) && length(z) == 1L

## Dense base copy of a Jacobian for the complex branches of the walker below.
## A 1x1 becomes a true scalar, so that `.is_scalar` and `*` treat it as one
## (a 1x1 base matrix times an n x 1 matrix is "non-conformable" in base R,
## where the same product of Matrix objects is a scalar multiple).
.grad_dense <- function(m) {
  m <- as.matrix(m)
  if (length(m) == 1L) m[1L, 1L] else m
}

method(grad, Atom) <- function(x) {
  ## Constants short-circuit: zero Jacobians for every variable in scope.
  if (is_constant(x)) {
    return(.constant_grad(x))
  }

  ## Pull values for each arg; if any is missing, propagate NULL up.
  arg_values <- vector("list", length(.args(x)))
  for (i in seq_along(.args(x))) {
    v <- value(.args(x)[[i]])
    if (is.null(v)) {
      return(.error_grad(x))
    }
    arg_values[[i]] <- v
  }

  ## Per-arg Jacobians from the .grad hook.
  grad_self <- .grad(x, arg_values)

  ## Compose against each child's chain via matrix multiplication (or
  ## scalar multiplication when either operand is a scalar). Mirrors
  ## CVXPY's atom.py:501-522 line-for-line.
  result <- list()
  for (idx in seq_along(.args(x))) {
    grad_arg <- grad(.args(x)[[idx]])
    for (key in names(grad_arg)) {
      ga <- grad_arg[[key]]
      gs <- grad_self[[idx]]
      if (is.null(ga) || is.null(gs)) {
        ## Propagate NULL using single-bracket assignment so the entry
        ## is set to NULL rather than removed.
        result[key] <- list(NULL)
      } else {
        if (is.complex(ga) || is.complex(gs)) {
          ## A complex Jacobian is a dense base complex matrix (see
          ## .dense_to_csc_vector, axis_atom.R): Matrix has no complex sparse
          ## class and cannot multiply or add a dgCMatrix with complex data,
          ## so the real operand is densified. Real Jacobians never enter
          ## this branch. Claims-audit decision 15, CVXR 1.9.2.9153.
          ga <- .grad_dense(ga)
          gs <- .grad_dense(gs)
        }
        if (.is_scalar(ga) || .is_scalar(gs)) {
          D <- ga * gs
        } else {
          D <- ga %*% gs
        }
        ## Collapse a 1x1 result to a scalar to match CVXPY's
        ## np.isscalar/sp.shape == (1, 1) handling; a complex 1x1 stays a
        ## complex scalar, as `D[0, 0]` does.
        dD <- dim(D)
        if (!is.null(dD) && all(dD == c(1L, 1L))) {
          D <- if (is.complex(D)) D[1L, 1L] else as.numeric(D[1L, 1L])
        }
        if (key %in% names(result)) {
          prev <- result[[key]]
          if (!is.null(prev) && (is.complex(prev) || is.complex(D))) {
            prev <- .grad_dense(prev)
            D <- .grad_dense(D)
          }
          result[[key]] <- prev + D
        } else {
          result[[key]] <- D
        }
      }
    }
  }
  result
}

# -- .grad: per-atom hook (abstract default) --------------------------
## CVXPY SOURCE: atom.py:526-538 (Atom._grad abstract method).
## Each non-affine non-leaf atom must register a method(.grad, MyAtom)
## returning a list of one sparse Jacobian per argument, shape
## (prod(arg.shape), prod(self.shape)).
##
## AffAtom has its own .grad method (built from the LinOp coefficient);
## see atoms/affine/affine_atom.R.

method(.grad, Atom) <- function(x, values, ...) {
  cli_abort(c(
    ".grad() not yet implemented for {.cls {S7::S7_class(x)@name}}.",
    "i" = "Each non-affine atom must register a {.fn method(.grad, MyAtom)} hook.",
    "i" = "See {.help dev_handbook} for the porting convention."
  ))
}

# -- .column_grad: per-fiber hook (abstract default) --------------
## Used by axis-aware atoms that inherit `.grad(AxisAtom)` from
## atoms/axis_atom.R. Each axis atom registers `method(.column_grad,
## MyAtom)` returning the per-fiber gradient (a numeric vector or
## NULL-on-domain-violation). Atoms that don't go through `.column_grad`
## register their own `method(.grad, MyAtom)` and never reach this
## default.

method(.column_grad, Atom) <- function(x, value, ...) {
  cli_abort(c(
    ".column_grad() not yet implemented for {.cls {S7::S7_class(x)@name}}.",
    "i" = "Axis atoms must register either {.fn method(.column_grad, MyAtom)} (for fiber-level grad) or {.fn method(.grad, MyAtom)} (to bypass the axis machinery).",
    "i" = "See {.help dev_handbook} for the porting convention."
  ))
}

# -- domain ----------------------------------------------------------
## CVXPY SOURCE: atom.py lines 541-551
## self._domain() + [con for arg in self.args for con in arg.domain]

## Default atom_domain: no extra constraints
## CVXPY SOURCE: atom.py::_domain -- default returns []
method(atom_domain, Atom) <- function(x) list()

method(domain, Atom) <- function(x) {
  ## CVXPY: self._domain() + [con for arg in self.args for con in arg.domain]
  my_domain <- atom_domain(x)
  arg_domains <- unlist(lapply(.args(x), domain), recursive = FALSE)
  if (is.null(arg_domains)) arg_domains <- list()
  c(my_domain, arg_domains)
}

# -- atoms -----------------------------------------------------------
## CVXPY SOURCE: atom.py lines 567-573
## Returns atom types present in the expression tree.
## In R, we collect S7 class objects, deduplicated by class name.

method(atoms, Atom) <- function(x) {
  nargs <- length(.args(x))
  parts <- vector("list", nargs + 1L)
  for (i in seq_len(nargs)) {
    parts[[i]] <- atoms(.args(x)[[i]])
  }
  parts[[nargs + 1L]] <- list(S7_class(x))
  .unique_atom_types(unlist(parts, recursive = FALSE))
}

## Helper: deduplicate a list of S7 class objects by class name
.unique_atom_types <- function(class_list) {
  n <- length(class_list)
  if (n == 0L) return(list())
  seen <- new.env(hash = TRUE, parent = emptyenv())
  result <- vector("list", n)
  ri <- 0L
  for (cls in class_list) {
    name <- if (inherits(cls, "S7_class")) cls@name else as.character(cls)
    if (!exists(name, envir = seen, inherits = FALSE)) {
      assign(name, TRUE, envir = seen)
      ri <- ri + 1L
      result[[ri]] <- cls
    }
  }
  result[seq_len(ri)]
}

# -- print -----------------------------------------------------------

method(print, Atom) <- function(x, ...) {
  cat(sprintf("%s\n", expr_name(x)))
  invisible(x)
}