R/051_atoms_affine_upper_tri.R

In CVXR: Disciplined Convex Optimization

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

## CVXPY SOURCE: atoms/affine/upper_tri.py
## UpperTri -- strict upper triangle of a square matrix as a vector


UpperTri <- new_class("UpperTri", parent = AffAtom, package = "CVXR",
  constructor = function(x, id = NULL) {
    if (is.null(id)) id <- next_expr_id()
    x <- as_expr(x)
    n <- x@shape[1L]
    ## strict upper triangle has n*(n-1)/2 entries
    entries <- (n * (n - 1L)) %/% 2L
    shape <- c(entries, 1L)

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

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

method(shape_from_args, UpperTri) <- function(x) {
  n <- x@args[[1L]]@shape[1L]
  c((n * (n - 1L)) %/% 2L, 1L)
}

# -- log-log: affine (CVXPY upper_tri.py) ------------------------
method(is_atom_log_log_convex, UpperTri) <- function(x) TRUE
method(is_atom_log_log_concave, UpperTri) <- function(x) TRUE

method(numeric_value, UpperTri) <- function(x, values, ...) {
  m <- values[[1L]]
  ## CVXPY uses np.triu_indices(n, k=1, m) which returns row-major order:
  ##   (0,1), (0,2), (0,3), (1,2), (1,3), (2,3) for a 4x4 matrix.
  ## R's upper.tri() + m[mask] returns column-major order:
  ##   (1,2), (1,3), (2,3), (1,4), (2,4), (3,4) for a 4x4 matrix.
  ## We need to extract indices sorted by row then column.
  idx <- which(upper.tri(m), arr.ind = TRUE)
  idx <- idx[order(idx[, 1L], idx[, 2L]), , drop = FALSE]
  matrix(m[idx], ncol = 1L)
}

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

#' Extract strict upper triangle of a square matrix
#'
#' @param x An Expression (square matrix)
#' @returns An UpperTri atom (column vector)
#' @export
upper_tri <- function(x) {
  UpperTri(x)
}

#' Reshape a vector into an upper triangular matrix
#'
#' Inverts \code{\link{upper_tri}}. Takes a flat vector and returns an
#' upper triangular matrix (row-major order, matching CVXPY convention).
#'
#' @param expr An Expression (vector).
#' @param strict Logical. If TRUE, returns a strictly upper triangular matrix
#'   (diagonal is zero). If FALSE, includes the diagonal. Default is FALSE.
#' @returns An Expression representing the upper triangular matrix.
#' @export
vec_to_upper_tri <- function(expr, strict = FALSE) {
  expr <- as_expr(expr)

  ## Flatten to column vector if needed
  if (!expr_is_vector(expr)) {
    cli_abort("{.fn vec_to_upper_tri} requires a vector input.")
  }
  ## If expr is (n, 1) shape, get the length
  ell <- expr_size(expr)

  ## Compute n from the triangular number
  if (strict) {
    ## n * (n-1) / 2 == ell
    n <- as.integer((sqrt(8 * ell + 1) + 1) / 2)
    if ((n * (n - 1L)) %/% 2L != ell) {
      cli_abort("Vector length {ell} is not a strict triangular number.")
    }
  } else {
    ## n * (n+1) / 2 == ell
    n <- as.integer((sqrt(8 * ell + 1) - 1) / 2)
    if ((n * (n + 1L)) %/% 2L != ell) {
      cli_abort("Vector length {ell} is not a triangular number.")
    }
  }

  ## Flatten expr to (ell, 1) if not already
  if (expr@shape[1L] != ell || expr@shape[2L] != 1L) {
    expr <- vec(expr)
  }

  ## Build sparse coefficient matrix P
  ## P maps: (P @ expr).reshape((n, n)) is upper triangular
  ## CVXPY: row_idx = n * row + col for (row, col) in triu_indices(n, k)
  ## This places each vector entry at the correct position in the flattened nxn matrix
  k <- if (strict) 1L else 0L
  ## Get upper triangular indices (row-major order, matching CVXPY)
  idx <- which(upper.tri(matrix(0, n, n), diag = !strict), arr.ind = TRUE)
  idx <- idx[order(idx[, 1L], idx[, 2L]), , drop = FALSE]
  row_0 <- idx[, 1L] - 1L  # 0-based row
  col_0 <- idx[, 2L] - 1L  # 0-based col

  ## P_rows: flattened position in row-major nxn matrix (0-based)
  P_rows <- n * row_0 + col_0
  P_cols <- seq_len(ell) - 1L  # 0-based column indices
  P_vals <- rep(1, ell)

  ## Build sparse matrix (convert to 1-based for R)
  P <- Matrix::sparseMatrix(
    i = P_rows + 1L,
    j = P_cols + 1L,
    x = P_vals,
    dims = c(as.integer(n * n), as.integer(ell))
  )

  ## P %*% expr gives (n*n, 1) column, reshape to (n, n) in F-order then transpose
  ## CVXPY: reshape(P @ expr, (n, n), order='F').T
  ## In R: reshape_expr(..., c(n, n), order = "F") then transpose
  result <- Constant(P) %*% expr
  result <- reshape_expr(result, c(n, n), order = "F")
  Transpose(result)
}