R/033_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 91-97
## 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 56-63

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 115-125

method(is_nonneg, Atom) <- function(x) {
  cached <- cache_get(x, "is_nonneg")
  if (!cache_miss(cached)) return(cached)
  result <- sign_from_args(x)[["is_nonneg"]]
  cache_set(x, "is_nonneg", result)
  result
}

method(is_nonpos, Atom) <- function(x) {
  cached <- cache_get(x, "is_nonpos")
  if (!cache_miss(cached)) return(cached)
  result <- sign_from_args(x)[["is_nonpos"]]
  cache_set(x, "is_nonpos", result)
  result
}

# -- 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) {
  cached <- cache_get(x, "get_bounds")
  if (!cache_miss(cached)) return(cached)
  ba <- bounds_from_args(x)
  result <- refine_bounds_from_sign(ba[[1L]], ba[[2L]], is_nonneg(x), is_nonpos(x))
  cache_set(x, "get_bounds", result)
  result
}

# -- Complex: default FALSE ------------------------------------------
## CVXPY SOURCE: atom.py lines 127-139

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

# -- is_atom_affine: plain function ----------------------------------
## CVXPY SOURCE: atom.py lines 153-156

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 178-181)
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 195-227
## idx for is_incr/is_decr: 1-based (R convention)

method(is_convex, Atom) <- function(x) {
  key <- .dpp_key("is_convex")
  cached <- cache_get(x, key)
  if (!cache_miss(cached)) return(cached)
  result <- .atom_is_convex(x)
  cache_set(x, key, result)
  result
}

.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) {
  key <- .dpp_key("is_concave")
  cached <- cache_get(x, key)
  if (!cache_miss(cached)) return(cached)
  result <- .atom_is_concave(x)
  cache_set(x, key, result)
  result
}

.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) {
  key <- .dpp_key("is_linearizable_convex")
  cached <- cache_get(x, key)
  if (!cache_miss(cached)) return(cached)
  result <- .atom_is_linearizable_convex(x)
  cache_set(x, key, result)
  result
}

.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) {
  key <- .dpp_key("is_linearizable_concave")
  cached <- cache_get(x, key)
  if (!cache_miss(cached)) return(cached)
  result <- .atom_is_linearizable_concave(x)
  cache_set(x, key, result)
  result
}

.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 273-335

## 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 278-285
.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 159-167
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 288-310)
method(is_quasiconvex, Atom) <- function(x) {
  cached <- cache_get(x, "is_quasiconvex")
  if (!cache_miss(cached)) return(cached)
  result <- .atom_is_quasiconvex(x)
  cache_set(x, "is_quasiconvex", result)
  result
}

.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 312-335)
method(is_quasiconcave, Atom) <- function(x) {
  cached <- cache_get(x, "is_quasiconcave")
  if (!cache_miss(cached)) return(cached)
  result <- .atom_is_quasiconcave(x)
  cache_set(x, "is_quasiconcave", result)
  result
}

.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 240-271

method(is_log_log_convex, Atom) <- function(x) {
  cached <- cache_get(x, "is_log_log_convex")
  if (!cache_miss(cached)) return(cached)
  result <- .atom_is_log_log_convex(x)
  cache_set(x, "is_log_log_convex", result)
  result
}

.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 247-253)
  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) {
  cached <- cache_get(x, "is_log_log_concave")
  if (!cache_miss(cached)) return(cached)
  result <- .atom_is_log_log_concave(x)
  cache_set(x, "is_log_log_concave", result)
  result
}

.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 263-269)
  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 337-356

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 379-403

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 385-403
.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:405-453 (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) && length(z) == 1L

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:430-451 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_scalar(ga) || .is_scalar(gs)) {
          D <- ga * gs
        } else {
          D <- ga %*% gs
        }
        ## Collapse 1x1 result to a scalar to match CVXPY's
        ## np.isscalar/sp.shape == (1, 1) handling.
        if (!.is_scalar(D)) {
          dD <- dim(D)
          if (!is.null(dD) && all(dD == c(1L, 1L))) {
            D <- as.numeric(D[1L, 1L])
          }
        }
        if (key %in% names(result)) {
          result[[key]] <- result[[key]] + D
        } else {
          result[[key]] <- D
        }
      }
    }
  }
  result
}

# -- .grad: per-atom hook (abstract default) --------------------------
## CVXPY SOURCE: atom.py:455-466 (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 469-474
## 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 496-502
## 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)
}