R/040_atoms_affine_index.R

In CVXR: Disciplined Convex Optimization

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

## CVXPY SOURCE: atoms/affine/index.py
## Index -- indexing/slicing into an Expression

# -- Key validation helpers ---------------------------------------------
## Convert R indexing arguments to validated integer index vectors.
## R's `[` can pass: integer, logical, missing (=NULL here), negative, etc.
## We normalize to a list of two integer vectors (1-based) or NULL (=all).

#' Validate and normalize an index key for a given dimension length
#' @param idx The index argument (integer, logical, or NULL for all)
#' @param dim_len The length of the dimension being indexed
#' @returns Integer vector of 1-based positive indices
#' @noRd
.validate_index_key <- function(idx, dim_len) {
  if (is.null(idx)) {
    ## NULL means "all" -- return full sequence
    return(seq_len(dim_len))
  }
  if (is.logical(idx)) {
    if (length(idx) != dim_len) {
      cli_abort("Logical index length ({length(idx)}) must match dimension ({dim_len}).")
    }
    return(which(idx))
  }
  idx <- as.integer(idx)
  if (any(idx < 0L)) {
    ## Negative indexing: exclude those positions
    all_idx <- seq_len(dim_len)
    return(all_idx[idx])  # R handles negative indexing
  }
  if (any(idx < 1L | idx > dim_len)) {
    cli_abort("Index out of bounds: must be in [1, {dim_len}].")
  }
  idx
}

#' Compute the shape resulting from indexing
#' @param key List of two integer index vectors
#' @param orig_shape Integer(2) original shape
#' @returns Integer(2) resulting shape
#' @noRd
.index_shape <- function(key, orig_shape) {
  nrow <- length(key[[1L]])
  ncol <- length(key[[2L]])
  c(as.integer(nrow), as.integer(ncol))
}

# -- Index class ------------------------------------------------------
## CVXPY SOURCE: atoms/affine/index.py lines 31-115

Index <- new_class("Index", parent = AffAtom, package = "CVXR",
  properties = list(
    key      = new_property(class = class_list),    # list(row_idx, col_idx)
    orig_key = new_property(class = class_list)      # list(orig_row, orig_col)
  ),
  constructor = function(expr, key, orig_key = NULL) {
    expr <- as_expr(expr)
    if (is.null(orig_key)) orig_key <- key

    ## Validate key indices against expression shape
    row_idx <- .validate_index_key(key[[1L]], expr@shape[1L])
    col_idx <- .validate_index_key(key[[2L]], expr@shape[2L])
    validated_key <- list(row_idx, col_idx)

    shape <- .index_shape(validated_key, expr@shape)

    new_object(S7_object(),
      id    = next_expr_id(),
      .cache = new.env(parent = emptyenv()),
      args  = list(expr),
      shape = shape,
      key   = validated_key,
      orig_key = orig_key
    )
  }
)

# -- shape_from_args --------------------------------------------------

method(shape_from_args, Index) <- function(x) {
  .index_shape(x@key, x@args[[1L]]@shape)
}

# -- sign_from_args ---------------------------------------------------
## Inherits from AffAtom: sum_signs(args)

# -- log-log curvature: affine (CVXPY index.py lines 68-72) ----------
method(is_atom_log_log_convex, Index) <- function(x) TRUE
method(is_atom_log_log_concave, Index) <- function(x) TRUE

# -- numeric_value ----------------------------------------------------
## CVXPY SOURCE: index.py lines 88-90

method(numeric_value, Index) <- function(x, values, ...) {
  val <- values[[1L]]
  if (!is.matrix(val)) val <- matrix(val, ncol = 1L)
  result <- val[x@key[[1L]], x@key[[2L]], drop = FALSE]
  result
}

# -- get_data ---------------------------------------------------------
## CVXPY SOURCE: index.py lines 96-98

method(get_data, Index) <- function(x) {
  list(x@key, x@orig_key)
}

# -- graph_implementation ---------------------------------------------
## CVXPY SOURCE: index.py lines 100-115

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

# -- expr_name --------------------------------------------------------
## CVXPY SOURCE: index.py lines 76-79

method(expr_name, Index) <- function(x) {
  .key_to_str <- function(idx, dim_len) {
    if (length(idx) == dim_len) return("")
    if (length(idx) == 1L) return(as.character(idx))
    paste0(idx[1L], ":", idx[length(idx)])
  }
  row_str <- .key_to_str(x@key[[1L]], x@args[[1L]]@shape[1L])
  col_str <- .key_to_str(x@key[[2L]], x@args[[1L]]@shape[2L])
  sprintf("%s[%s, %s]", expr_name(x@args[[1L]]), row_str, col_str)
}

# -- is_symmetric / is_hermitian --------------------------------------

method(is_symmetric, Index) <- function(x) {
  x@shape[1L] == x@shape[2L] && x@shape[1L] == 1L
}

method(is_hermitian, Index) <- function(x) {
  x@shape[1L] == x@shape[2L] && x@shape[1L] == 1L
}