R/046_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) {
    if (FALSE) new_object(S7_object())  ## S7 static-check guard
    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]], .shape(expr)[1L])
    col_idx <- .validate_index_key(key[[2L]], .shape(expr)[2L])
    validated_key <- list(row_idx, col_idx)

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

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

# -- bounds_from_args -------------------------------------------------
## CVXPY SOURCE: index.py:75-78 (#3080). Select the same submatrix from the
## arg's bounds. CVXR's `key` is list(row_idx, col_idx) of 1-based R indices,
## so `a[rows, cols]` is the direct selection (no order translation needed).
method(bounds_from_args, Index) <- function(x) {
  b <- get_bounds(.args(x)[[1L]])
  key <- x@key
  kf <- function(a) a[key[[1L]], key[[2L]], drop = FALSE]
  index_bounds(b[[1L]], b[[2L]], kf)
}

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

method(shape_from_args, Index) <- function(x) {
  .index_shape(x@key, .arg_shape(x))
}

# -- 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]], .arg_shape(x)[1L])
  col_str <- .key_to_str(x@key[[2L]], .arg_shape(x)[2L])
  sprintf("%s[%s, %s]", expr_name(.args(x)[[1L]]), row_str, col_str)
}

## CVXPY SOURCE: atoms/affine/index.py Index.format_labeled.
method(format_labeled, Index) <- function(x) {
  lbl <- label(x); if (!is.null(lbl)) return(lbl)
  .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]], .arg_shape(x)[1L])
  col_str <- .key_to_str(x@key[[2L]], .arg_shape(x)[2L])
  sprintf("%s[%s, %s]", format_labeled(.args(x)[[1L]]), row_str, col_str)
}

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

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

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

# ========================================================================
# SpecialIndex -- element-wise indexing (2-column matrix, logical, linear)
# ========================================================================
## CVXPY SOURCE: atoms/affine/index.py class special_index
## Supports R's single-subscript matrix indexing forms:
##   x[cbind(rows, cols)]   -- 2-column matrix (element-wise)
##   x[logical_matrix]      -- logical matrix
##   x[integer_vector]      -- linear column-major indexing

SpecialIndex <- new_class("SpecialIndex", parent = AffAtom, package = "CVXR",
  properties = list(
    key        = new_property(class = class_any),       # original key
    select_vec = new_property(class = class_integer)    # 1-based linear indices (column-major)
  ),
  constructor = function(expr, key) {
    if (FALSE) new_object(S7_object())  ## S7 static-check guard
    expr <- as_expr(expr)
    eshape <- .shape(expr)

    ## -- Validate key and compute select_vec (1-based linear indices) --
    if (is.matrix(key) && ncol(key) == 2L && (is.integer(key) || is.double(key))) {
      ## 2-column matrix: each row is a (row, col) pair
      if (is.double(key)) storage.mode(key) <- "integer"
      if (nrow(key) == 0L) {
        select_vec <- integer(0L)
      } else {
        if (anyNA(key))
          cli_abort("Index key contains NA values.")
        if (any(key[, 1L] < 1L | key[, 1L] > eshape[1L]))
          cli_abort("Row index out of bounds: must be in [1, {eshape[1L]}].")
        if (any(key[, 2L] < 1L | key[, 2L] > eshape[2L]))
          cli_abort("Column index out of bounds: must be in [1, {eshape[2L]}].")
        ## Fast path: compute linear indices directly (no idx_mat allocation)
        select_vec <- (key[, 2L] - 1L) * eshape[1L] + key[, 1L]
      }
    } else if (is.matrix(key) && is.logical(key)) {
      ## Logical matrix: select where TRUE
      if (nrow(key) != eshape[1L] || ncol(key) != eshape[2L])
        cli_abort("Logical matrix dimensions ({nrow(key)} x {ncol(key)}) must match expression shape ({eshape[1L]} x {eshape[2L]}).")
      idx_mat <- matrix(seq_len(prod(eshape)), nrow = eshape[1L], ncol = eshape[2L])
      select_vec <- idx_mat[key]
    } else if (is.integer(key) && !is.matrix(key)) {
      ## Integer vector: linear column-major indexing
      n <- prod(eshape)
      if (anyNA(key))
        cli_abort("Index key contains NA values.")
      if (any(key < 1L))
        cli_abort("Negative and zero indices are not supported for element-wise indexing.")
      if (any(key > n))
        cli_abort("Linear index out of bounds: must be in [1, {n}].")
      select_vec <- key
    } else {
      cli_abort("Unsupported key type for {.cls SpecialIndex}.")
    }

    shape <- c(length(select_vec), 1L)

    .fast_new(SpecialIndex, S7_object(),
      id    = next_expr_id(),
      .cache = new.env(parent = emptyenv()),
      args  = list(expr),
      shape = shape,
      key   = key,
      select_vec = select_vec
    )
  }
)

# -- bounds_from_args -------------------------------------------------
## CVXPY SOURCE: index.py (special_index):134-137 (#3080). Select the same
## entries from the arg's bounds. `select_vec` is 1-based column-major linear
## indices, so index the column-major flattened bounds and reshape to (k, 1).
method(bounds_from_args, SpecialIndex) <- function(x) {
  b <- get_bounds(.args(x)[[1L]])
  sel <- x@select_vec
  kf <- function(a) matrix(as.numeric(a)[sel], ncol = 1L)
  index_bounds(b[[1L]], b[[2L]], kf)
}

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

method(shape_from_args, SpecialIndex) <- function(x) {
  c(length(x@select_vec), 1L)
}

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

# -- log-log curvature ------------------------------------------------

method(is_atom_log_log_convex, SpecialIndex) <- function(x) TRUE
method(is_atom_log_log_concave, SpecialIndex) <- function(x) TRUE

# -- numeric_value ----------------------------------------------------
## Use select_vec (linear indices) for consistency with graph_implementation

method(numeric_value, SpecialIndex) <- function(x, values, ...) {
  val <- values[[1L]]
  if (inherits(val, "sparseMatrix")) val <- as.matrix(val)
  if (!is.matrix(val)) val <- as.matrix(val)
  ## Flatten column-major, select by linear index
  matrix(as.vector(val)[x@select_vec], ncol = 1L)
}

# -- get_data ---------------------------------------------------------
## Must match constructor signature: SpecialIndex(expr, key)
## expr_copy calls do.call(SpecialIndex, c(new_args, get_data(x)))

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

# -- graph_implementation ---------------------------------------------
## CVXPY SOURCE: index.py special_index.graph_implementation (lines 189-214)
## Flatten → sparse selection matrix → multiply

method(graph_implementation, SpecialIndex) <- function(x, arg_objs, shape, data = NULL, ...) {
  select_vec <- x@select_vec       # 1-based linear indices
  n <- expr_size(.args(x)[[1L]])     # total elements in original expr
  k <- length(select_vec)

  ## 1. Flatten arg to column vector (column-major)
  vec_arg <- reshape_linop(arg_objs[[1L]], c(n, 1L))

  ## 2. Sparse selection matrix: k × n, entry (i, select_vec[i]) = 1
  sel_mat <- Matrix::sparseMatrix(
    i = seq_len(k), j = select_vec,
    x = rep(1.0, k), dims = c(k, n)
  )
  ## Ensure dgCMatrix for C++ bridge
  if (!inherits(sel_mat, "dgCMatrix")) {
    sel_mat <- methods::as(sel_mat, "dgCMatrix")
  }
  mul_const <- create_const(sel_mat, c(k, n), sparse = TRUE)

  ## 3. Multiply: sel_mat %*% vec(expr) → c(k, 1)
  result <- mul_expr_linop(mul_const, vec_arg, shape)

  list(result, list())
}

# -- expr_name --------------------------------------------------------

method(expr_name, SpecialIndex) <- function(x) {
  key <- x@key
  if (is.matrix(key) && ncol(key) == 2L && is.integer(key)) {
    n <- nrow(key)
    if (n <= 3L) {
      pairs <- vapply(seq_len(n), function(i)
        sprintf("(%d,%d)", key[i, 1L], key[i, 2L]), character(1L))
      idx_str <- paste(pairs, collapse = ",")
    } else {
      idx_str <- sprintf("%d elements", n)
    }
  } else if (is.matrix(key) && is.logical(key)) {
    idx_str <- sprintf("%d TRUE of %dx%d", sum(key), nrow(key), ncol(key))
  } else {
    idx_str <- sprintf("%d indices", length(x@select_vec))
  }
  sprintf("%s[%s]", expr_name(.args(x)[[1L]]), idx_str)
}

## CVXPY SOURCE: atoms/affine/index.py SpecialIndex shares Index format.
method(format_labeled, SpecialIndex) <- function(x) {
  lbl <- label(x); if (!is.null(lbl)) return(lbl)
  key <- x@key
  if (is.matrix(key) && ncol(key) == 2L && is.integer(key)) {
    n <- nrow(key)
    if (n <= 3L) {
      pairs <- vapply(seq_len(n), function(i)
        sprintf("(%d,%d)", key[i, 1L], key[i, 2L]), character(1L))
      idx_str <- paste(pairs, collapse = ",")
    } else {
      idx_str <- sprintf("%d elements", n)
    }
  } else if (is.matrix(key) && is.logical(key)) {
    idx_str <- sprintf("%d TRUE of %dx%d", sum(key), nrow(key), ncol(key))
  } else {
    idx_str <- sprintf("%d indices", length(x@select_vec))
  }
  sprintf("%s[%s]", format_labeled(.args(x)[[1L]]), idx_str)
}

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

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

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