R/051_atoms_affine_trace.R

In CVXR: Disciplined Convex Optimization

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

## CVXPY SOURCE: atoms/affine/trace.py
## Trace -- sum of diagonal entries of a square matrix


Trace <- new_class("Trace", parent = AffAtom, package = "CVXR",
  constructor = function(x, id = NULL) {
    if (FALSE) new_object(S7_object())  ## S7 static-check guard
    if (is.null(id)) id <- next_expr_id()
    x <- as_expr(x)
    shape <- c(1L, 1L)  # scalar

    obj <- .fast_new(Trace, S7_object(),
      id    = as.integer(id),
      .cache = new.env(parent = emptyenv()),
      args  = list(x),
      shape = shape
    )
    validate_arguments(obj)
    obj
  }
)

method(validate_arguments, Trace) <- function(x) {
  arg <- .args(x)[[1L]]
  if (.shape(arg)[1L] != .shape(arg)[2L]) {
    cli_abort("{.cls Trace} requires a square matrix, got shape ({arg@shape[1L]}, {arg@shape[2L]}).")
  }
  invisible(NULL)
}

method(shape_from_args, Trace) <- function(x) c(1L, 1L)

# -- sign: PSD/NSD propagation (CVXPY trace.py lines 54-61) ------
method(sign_from_args, Trace) <- function(x) {
  list(
    is_nonneg = is_nonneg(.args(x)[[1L]]) || is_psd(.args(x)[[1L]]),
    is_nonpos = is_nonpos(.args(x)[[1L]]) || is_nsd(.args(x)[[1L]])
  )
}

# -- is_real / is_complex (CVXPY trace.py lines 100-104) ---------
## The trace of a HERMITIAN matrix is real: its diagonal entries satisfy
## H[i,i] == Conj(H[i,i]), so the sum has zero imaginary part.
##
## `matrix_trace()` already handled the Hermitian-PRODUCT case by wrapping the
## result in Real_() (see its body above), but a bare
## `Variable(c(n, n), hermitian = TRUE)` builds a plain Trace, which inherited a
## complex verdict.  Complex2Real then split it and emitted a redundant
## `Im(.) == 0` row -- correct answers, extra work, and a `value()` a user has
## to call Re() on.
method(is_real, Trace) <- function(x) {
  is_real(.args(x)[[1L]]) || is_hermitian(.args(x)[[1L]])
}
method(is_complex, Trace) <- function(x) !is_real(x)

# -- log-log: convex, not concave (CVXPY trace.py lines 87-92) --
method(is_atom_log_log_convex, Trace) <- function(x) TRUE
method(is_atom_log_log_concave, Trace) <- function(x) FALSE

method(numeric_value, Trace) <- function(x, values, ...) {
  matrix(sum(diag(values[[1L]])), 1L, 1L)
}

method(graph_implementation, Trace) <- function(x, arg_objs, shape, data = NULL, ...) {
  list(trace_linop(arg_objs[[1L]]), list())
}

#' Trace of a square matrix expression
#'
#' For \code{matrix_trace(A \%*\% B)}, uses the O(n^2) identity
#' \code{trace(A \%*\% B) = sum(A * t(B))} instead of forming the
#' full matrix product.
#'
#' @param x An Expression (square matrix)
#' @returns A Trace atom or equivalent expression (scalar)
#' @export
matrix_trace <- function(x) {
  ## CVXPY v1.8.2 fix: trace(A@B) = sum(A * B.T) avoids O(n^3) matmul.
  ## Also detects Hermitian products and wraps with Real_() so
  ## is_real() propagates correctly for complex problems.
  if (.s7_is(x, MulExpression)) {
    result <- SumEntries(Multiply(.args(x)[[1L]], t(.args(x)[[2L]])))
    if (is_hermitian(x)) {
      return(Real_(result))
    }
    return(result)
  }
  Trace(x)
}