R/228_utilities_coeff_extractor.R

In CVXR: Disciplined Convex Optimization

#####
## DO NOT EDIT THIS FILE!! EDIT THE SOURCE INSTEAD: rsrc_tree/utilities/coeff_extractor.R
#####

## CVXPY SOURCE: utilities/coeff_extractor.py
## CoeffExtractor -- non-parametric coefficient extraction
##
## Extracts A matrix and b vector from affine expressions via C++ canonInterface.
## Non-parametric simplification: no parameter tensor, just direct A and b.


# -- CoeffExtractor class ------------------------------------------
## CVXPY SOURCE: coeff_extractor.py lines 36-78

CoeffExtractor <- new_class("CoeffExtractor", package = "CVXR",
  properties = list(
    id_to_col     = class_any,      # named integer vector: var_id_str -> offset
    x_length      = class_integer,
    param_to_size = class_list,     # param.id -> size (includes CONSTANT_ID)
    param_id_map  = class_list      # param.id -> column offset in tensor
  ),
  constructor = function(inverse_data) {
    ## Build id_to_col from InverseData@var_offsets
    ## var_offsets is a named list: "var_id" -> offset
    id_to_col <- as.integer(unlist(inverse_data@var_offsets))
    names(id_to_col) <- names(inverse_data@var_offsets)
    new_object(S7_object(),
      id_to_col     = id_to_col,
      x_length      = inverse_data@x_length,
      param_to_size = inverse_data@param_to_size,
      param_id_map  = inverse_data@param_id_map
    )
  }
)

# -- coeff_affine: extract A, b from a list of affine expressions --
## Returns list(A, b) where:
##   A is a sparse Matrix (num_rows x x_length)
##   b is a numeric vector (length num_rows)
## Sign convention from C++: A*x + b = expression_value
##   For Zero constraints: A*x + b = 0
##   For NonNeg constraints: A*x + b >= 0

coeff_affine <- function(extractor, expr_list) {
  if (!is.list(expr_list)) expr_list <- list(expr_list)

  ## Total rows = sum of expression sizes
  num_rows <- sum(vapply(expr_list, expr_size, integer(1L)))

  ## Get LinOp trees from canonical_form
  linop_list <- lapply(expr_list, function(e) canonical_form(e)[[1L]])

  ## Call C++ via canonInterface
  ## get_problem_matrix returns list(V, I, J, const_vec)
  ## I and J are 0-based from C++
  result <- get_problem_matrix(linop_list, extractor@id_to_col,
                                var_length = extractor@x_length)

  V <- result$V
  I <- result$I
  J <- result$J
  const_vec <- as.vector(result$const_vec)

  ## Build sparse matrix A (num_rows x x_length)
  ## I, J are 0-based from C++; add 1 for R's 1-based indexing
  if (length(V) > 0L && extractor@x_length > 0L) {
    ## Filter to variable columns only (J < x_length)
    ## The C++ may return entries for the constant column too,
    ## but const_vec already captures those.
    var_mask <- J < extractor@x_length
    A <- Matrix::sparseMatrix(
      i = I[var_mask] + 1L,
      j = J[var_mask] + 1L,
      x = V[var_mask],
      dims = c(num_rows, extractor@x_length)
    )
  } else {
    A <- Matrix::sparseMatrix(
      i = integer(0), j = integer(0), x = numeric(0),
      dims = c(num_rows, max(extractor@x_length, 1L))
    )
  }

  list(A = A, b = const_vec)
}

# -- coeff_affine_tensor: extract tensor from affine expressions --
## CVXPY SOURCE: coeff_extractor.py lines 47-78
## Returns a sparse tensor matrix (constr*(var+1), param_size_plus_one).
## Used by DPP path: parameters are NOT substituted, so the tensor
## captures how each parameter element contributes to each entry.

coeff_affine_tensor <- function(extractor, expr_list) {
  if (!is.list(expr_list)) expr_list <- list(expr_list)

  ## Get LinOp trees from canonical_form
  linop_list <- lapply(expr_list, function(e) canonical_form(e)[[1L]])

  ## Call C++ via canonInterface -- returns sparse tensor matrix
  get_problem_matrix_tensor(linop_list, extractor@id_to_col,
                             var_length = extractor@x_length,
                             param_to_size = extractor@param_to_size,
                             param_id_map = extractor@param_id_map)
}

# -- coeff_quad_form: extract P and q from a quadratic expression --
## CVXPY SOURCE: coeff_extractor.py lines 284-339
## Non-parametric simplification of extract_quadratic_coeffs.
##
## Returns list(P, q) where:
##   P is a sparse Matrix (x_length x x_length)  -- already 2x scaled
##   q is a numeric vector (length x_length)       -- linear term
## Solver minimizes 0.5*x'*P*x + q'*x, so P = 2 * (extracted coefficients).

coeff_quad_form <- function(extractor, expr) {
  ## Step 0: If root IS a SymbolicQuadForm, wrap it so replace_quad_forms
  ## can find it as a child (it only replaces children, not the root).
  ## CVXPY wraps in LinOp(NO_OP,...); we wrap in 0 + expr.
  if (S7_inherits(expr, SymbolicQuadForm) || S7_inherits(expr, QuadForm)) {
    expr <- Constant(matrix(0, 1L, 1L)) + expr
  }

  ## Step 1: Replace SymbolicQuadForm nodes with dummy Variables
  rqf <- replace_quad_forms(expr, list())
  modified_expr <- rqf$expr
  quad_forms <- rqf$quad_forms

  ## Step 2: Build LOCAL variable offsets for the modified expression
  ## The modified expr has real variables + dummy placeholder variables
  local_vars <- variables(modified_expr)
  local_offsets <- list()
  local_id_map <- list()
  local_var_shapes <- list()
  offset <- 0L
  for (v in local_vars) {
    vid <- as.character(v@id)
    sz <- expr_size(v)
    local_offsets[[vid]] <- offset
    local_id_map[[vid]] <- c(offset, sz)
    local_var_shapes[[vid]] <- v@shape
    offset <- offset + sz
  }
  local_x_length <- offset

  ## Step 3: Build local id_to_col for get_problem_matrix
  local_id_to_col <- as.integer(unlist(local_offsets))
  names(local_id_to_col) <- names(local_offsets)

  ## Step 4: Call C++ to get affine coefficients of modified (affine) expression
  linop <- canonical_form(modified_expr)[[1L]]
  result <- get_problem_matrix(list(linop), local_id_to_col,
                                var_length = local_x_length)
  V <- result$V; I <- result$I; J <- result$J
  const_offset <- as.numeric(result$const_vec)

  ## Build local coefficient vector (1 x local_x_length)
  ## For a scalar expression, I should all be 0
  if (length(V) > 0L && local_x_length > 0L) {
    var_mask <- J < local_x_length
    local_c <- numeric(local_x_length)
    for (k in which(var_mask)) {
      local_c[J[k] + 1L] <- local_c[J[k] + 1L] + V[k]
    }
  } else {
    local_c <- numeric(local_x_length)
  }

  ## Step 5: Separate quad form coefficients from linear coefficients
  ## Build global P (x_length x x_length) and q (x_length)
  global_x <- extractor@x_length
  n_local <- length(local_vars)
  triplet_i_chunks <- vector("list", n_local)
  triplet_j_chunks <- vector("list", n_local)
  triplet_x_chunks <- vector("list", n_local)
  q_vec <- numeric(global_x)

  for (vi in seq_len(n_local)) {
    v <- local_vars[[vi]]
    vid <- as.character(v@id)
    lo <- local_offsets[[vid]]
    sz <- expr_size(v)

    if (!is.null(quad_forms[[vid]])) {
      ## This is a dummy variable replacing a SymbolicQuadForm
      sqf <- quad_forms[[vid]]$quad_form
      ## Coefficient of this dummy in the affine expression (scalar for scalar objectives)
      c_part <- local_c[(lo + 1L):(lo + sz)]

      ## The SymbolicQuadForm's inner variable (the x in quad_form(x, P))
      orig_var <- sqf@args[[1L]]
      orig_vid <- as.character(orig_var@id)
      orig_offset <- extractor@id_to_col[[orig_vid]]
      orig_size <- expr_size(orig_var)

      ## Get P matrix from the SymbolicQuadForm
      P_val <- value(sqf@args[[2L]])
      if (inherits(P_val, "sparseMatrix")) {
        ## Coerce diagonalMatrix/symmetricMatrix etc. to dgTMatrix for
        ## reliable triplet summary (ddiMatrix summary returns diagSummary)
        if (!inherits(P_val, "dgTMatrix"))
          P_val <- methods::as(P_val, "TsparseMatrix")
        P_coo <- Matrix::summary(P_val)  # i, j, x triplets (1-based)
      } else {
        P_mat <- as.matrix(P_val)
        nz <- which(P_mat != 0, arr.ind = TRUE)
        P_coo <- if (nrow(nz) > 0L) {
          data.frame(i = nz[, 1L], j = nz[, 2L], x = P_mat[nz])
        } else {
          data.frame(i = integer(0), j = integer(0), x = numeric(0))
        }
      }

      if (nrow(P_coo) > 0L) {
        ## Scalar coefficient (most common case: sum of scalar quad forms)
        scalar_coeff <- if (sz == 1L) c_part[1L] else 1.0
        ## Shift to global indices -- collect chunks
        triplet_i_chunks[[vi]] <- P_coo$i + orig_offset  # already 1-based
        triplet_j_chunks[[vi]] <- P_coo$j + orig_offset
        triplet_x_chunks[[vi]] <- P_coo$x * scalar_coeff
      }
    } else {
      ## Real variable: its coefficient goes into q
      global_offset <- extractor@id_to_col[[vid]]
      q_vec[(global_offset + 1L):(global_offset + sz)] <- local_c[(lo + 1L):(lo + sz)]
    }
  }

  P_triplets_i <- unlist(triplet_i_chunks)
  P_triplets_j <- unlist(triplet_j_chunks)
  P_triplets_x <- unlist(triplet_x_chunks)
  if (is.null(P_triplets_i)) {
    P_triplets_i <- integer(0)
    P_triplets_j <- integer(0)
    P_triplets_x <- numeric(0)
  }

  ## Step 6: Build sparse P matrix
  if (length(P_triplets_i) > 0L) {
    P_mat <- Matrix::sparseMatrix(
      i = P_triplets_i, j = P_triplets_j, x = P_triplets_x,
      dims = c(global_x, global_x)
    )
  } else {
    P_mat <- Matrix::sparseMatrix(
      i = integer(0), j = integer(0), x = numeric(0),
      dims = c(global_x, global_x)
    )
  }

  ## Step 7: Factor of 2 -- solver minimizes 0.5*x'Px + q'x
  ## Our extracted P represents x'Px (no 0.5), so multiply by 2.
  P_mat <- 2 * P_mat

  ## Step 8: Restore original expression (for caching correctness)
  ## Not strictly needed since we don't store the modified expr,
  ## but matches CVXPY's restore_quad_forms call.

  list(P = P_mat, q = q_vec, offset = const_offset)
}