R/256_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) {
    if (FALSE) new_object(S7_object())  ## S7 static-check guard
    ## 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)
    .fast_new(CoeffExtractor, 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_is(expr, SymbolicQuadForm) || .s7_is(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(.id(v))
    sz <- expr_size(v)
    local_offsets[[vid]] <- offset
    local_id_map[[vid]] <- c(offset, sz)
    local_var_shapes[[vid]] <- .shape(v)
    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(.id(v))
    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 <- .args(sqf)[[1L]]
      orig_vid <- as.character(.id(orig_var))
      orig_offset <- extractor@id_to_col[[orig_vid]]
      orig_size <- expr_size(orig_var)

      ## Get P matrix from the SymbolicQuadForm
      P_val <- value(.args(sqf)[[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) {
        ## CVXPY SOURCE: coeff_extractor.py:227-262 -- extract_quadratic_coeffs
        ## branches on the SHAPE of the placeholder, because a non-scalar
        ## placeholder carries ONE COEFFICIENT PER COMPONENT, not one overall.
        if (sz == 1L) {
          ## SCALAR PATH (coeff_extractor.py:227-239). One quadratic term, so
          ## the whole of P is scaled by the placeholder's affine coefficient.
          ## Reached by quad_form() and quad_over_lin(), both of which are
          ## scalar-valued.
          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 * c_part[1L]
        } else {
          ## DIAGONAL PATH (coeff_extractor.py:245-262).
          ##
          ## Reached from power_canon, which builds
          ## SymbolicQuadForm(x, Diagonal(n), square(x)). `square(x)` is
          ## ELEMENTWISE, so the placeholder is a VECTOR: component i stands for
          ## P[i,i] * x_i^2 and gets its own affine coefficient c_part[i]. The
          ## contribution is therefore diag(P[i,i] * c_part[i]) -- NOT P scaled
          ## by a single number.
          ##
          ## Until 1.9.1.9049 this branch did `scalar_coeff <- 1.0`, silently
          ## DISCARDING every coefficient: `lambda * sum(square(w))` solved as
          ## though lambda were 1, so ridge regression returned the same fit for
          ## every lambda. No error, no warning, wrong minimizer. It survived
          ## because sum_squares()/quad_form() take the scalar path above, and
          ## because a PURE quadratic has the same argmin under any positive
          ## scale -- it takes a linear term to expose it (see
          ## test_tree/test-quad-coeff-scaling.R).
          ##
          ## The three aborts below are cases upstream handles differently or
          ## asserts against; each means an assumption here has broken.
          if (!is.null(sqf@block_indices)) {
            cli_abort(c(
              "Cannot extract the quadratic form: block-structured quadratic terms are not supported.",
              "i" = "CVXR reaches this only if a canonicalizer set {.code block_indices}; none does today.",
              "i" = "This is a bug in CVXR, not in your problem. Please report it, with a reproducible example, at {.url https://github.com/cvxgrp/CVXR/issues}."
            ), class = "CVXR_internal_error")
          }
          if (sz != orig_size) {
            cli_abort(c(
              "Cannot extract the quadratic form: placeholder size ({sz}) does not match the variable size ({orig_size}).",
              "i" = "The component-wise mapping between the quadratic term and its variable is what makes this path valid.",
              "i" = "This is a bug in CVXR, not in your problem. Please report it, with a reproducible example, at {.url https://github.com/cvxgrp/CVXR/issues}."
            ), class = "CVXR_internal_error")
          }
          if (!all(P_coo$i == P_coo$j)) {
            cli_abort(c(
              "Cannot extract the quadratic form: a vector-valued quadratic term needs a diagonal {.var P}.",
              "i" = "Upstream CVXPY asserts the same condition (coeff_extractor.py:250-254).",
              "i" = "This is a bug in CVXR, not in your problem. Please report it, with a reproducible example, at {.url https://github.com/cvxgrp/CVXR/issues}."
            ), class = "CVXR_internal_error")
          }
          idx  <- P_coo$i
          vals <- P_coo$x * c_part[idx]
          keep <- vals != 0
          triplet_i_chunks[[vi]] <- idx[keep] + orig_offset
          triplet_j_chunks[[vi]] <- idx[keep] + orig_offset
          triplet_x_chunks[[vi]] <- vals[keep]
        }
      }
    } 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)
}

# -- coeff_quad_form_tensor: DPP P_tensor + c_tensor extraction ----
## CVXPY SOURCE: coeff_extractor.py extract_quadratic_coeffs() +
## quad_form() (the parametric-P path added in v1.9.0 #3142).
##
## Tensor analogue of coeff_quad_form for the DPP path: instead of baking
## P's numeric value, it tracks P's affine dependence on parameters so the
## QP fast path can rebuild P from new parameter values on re-solve (no
## staleness). Returns list(P_tensor, c_tensor):
##   P_tensor : (x_length*x_length, param_size_plus_one) sparse -- column-major
##              vec(P) contracted with the parameter vector pv gives the
##              solver P (already 2x scaled for the 0.5 x'Px convention).
##   c_tensor : (x_length+1, param_size_plus_one) sparse -- linear term over
##              x (rows 1..x_length) plus the constant offset (last row).
## apply_parameters() (param_prob.R) contracts both with pv on every solve.
##
## Requires each quad_form to have a SCALAR output and a param-free affine
## multiplier (enforced upstream by QuadForm._check_dpp_args in the
## quad_form_dpp_scope). A param-affine P (bare Parameter, 2*P, P1+P2, ...)
## is handled because get_problem_matrix_tensor extracts vec(P)'s parameter
## coefficients directly.
coeff_quad_form_tensor <- function(extractor, expr) {
  ## Step 0-1: wrap a bare quad form so replace_quad_forms can reach it, then
  ## swap SymbolicQuadForm/QuadForm nodes for dummy Variables (affine head).
  if (.s7_is(expr, SymbolicQuadForm) || .s7_is(expr, QuadForm)) {
    expr <- Constant(matrix(0, 1L, 1L)) + expr
  }
  rqf <- replace_quad_forms(expr, list())
  modified_expr <- rqf$expr
  quad_forms <- rqf$quad_forms

  ## Step 2: local variable offsets for the modified (affine) expression.
  local_vars <- variables(modified_expr)
  local_offsets <- list()
  local_id_to_col <- integer(0)
  offset <- 0L
  for (v in local_vars) {
    vid <- as.character(.id(v))
    sz <- expr_size(v)
    local_offsets[[vid]] <- offset
    local_id_to_col[vid] <- offset
    offset <- offset + sz
  }
  local_x_length <- offset

  ## Parameter-tensor column geometry (shared with c/A tensors and pv).
  param_size_plus_one <- as.integer(sum(unlist(extractor@param_to_size)))
  const_col0 <- extractor@param_id_map[["-1"]]   # 0-based constant-slice column
  global_x <- extractor@x_length

  ## Step 3: affine head as a parameter tensor.
  ## T_head: (local_x_length + 1, param_size_plus_one); row j (1-based) is the
  ## coefficient of local entry j (offset row last); columns are param slices
  ## then the constant slice.
  linop <- canonical_form(modified_expr)[[1L]]
  T_head <- get_problem_matrix_tensor(list(linop), local_id_to_col,
                                      var_length = local_x_length,
                                      param_to_size = extractor@param_to_size,
                                      param_id_map = extractor@param_id_map)
  T_head <- methods::as(T_head, "TsparseMatrix")

  ## Step 4: split into c_tensor (real vars + offset) and P_tensor (quad forms).
  c_i <- integer(0); c_j <- integer(0); c_x <- numeric(0)
  P_i <- integer(0); P_j <- integer(0); P_x <- numeric(0)

  ## Helper: copy a contiguous T_head row block [lo+1 .. lo+sz] to global rows
  ## starting at g0 (0-based) in c_tensor.
  Th <- Matrix::summary(T_head)   # i, j, x (1-based)

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

    if (!is.null(quad_forms[[vid]])) {
      ## Quad-form dummy. The objective term is m(theta) * x' P(theta) x, where
      ## m is the dummy's affine-head coefficient and P the quad-form matrix.
      ## For DPP the product m * P must be parameter-AFFINE, i.e. AT MOST ONE of
      ## {m, P} is parametric (both parametric => quadratic-in-params, non-DPP).
      ## Two valid shapes:
      ##   (1) m param-free, P param-affine  -- e.g. 0.5*quad_form(x, P_param)
      ##   (2) m param-affine, P constant    -- e.g. p * x^2  (P = I, m = p)
      sqf <- quad_forms[[vid]]$quad_form
      if (sz != 1L) {
        cli_abort("DPP quad_form with parametric P requires a scalar quad_form output.")
      }
      ## m_row: the dummy's coefficient across all parameter slices (+const).
      m_row <- as.numeric(T_head[lo + 1L, , drop = TRUE])
      m_is_const <- all(m_row[-(const_col0 + 1L)] == 0)

      orig_var    <- .args(sqf)[[1L]]
      orig_offset <- extractor@id_to_col[[as.character(.id(orig_var))]]
      n <- as.integer(expr_size(orig_var))   # P is n x n

      ## vec(P_expr) parameter coefficients: (n*n, param_size_plus_one).
      P_expr  <- .args(sqf)[[2L]]
      Plinop  <- canonical_form(P_expr)[[1L]]
      Ptens   <- methods::as(
        get_problem_matrix_tensor(
          list(Plinop), integer(0), var_length = 0L,
          param_to_size = extractor@param_to_size,
          param_id_map  = extractor@param_id_map),
        "TsparseMatrix")
      Pcoo <- Matrix::summary(Ptens)
      P_is_const <- all(Pcoo$j == (const_col0 + 1L) | Pcoo$x == 0)

      ## contrib triplets: (local vec index k 1-based, param column j, value)
      ck <- integer(0); cj <- integer(0); cx <- numeric(0)
      if (m_is_const) {
        ## (1) scale every P parameter slice by the constant multiplier.
        m_c <- m_row[const_col0 + 1L]
        if (m_c != 0 && nrow(Pcoo) > 0L) {
          ck <- Pcoo$i; cj <- Pcoo$j; cx <- m_c * Pcoo$x
        }
      } else if (P_is_const) {
        ## (2) outer product of constant vec(P) with the parametric multiplier:
        ## contrib[k, p] = P_const[k] * m_row[p].
        Pc <- Pcoo[Pcoo$j == (const_col0 + 1L) & Pcoo$x != 0, , drop = FALSE]
        m_nz <- which(m_row != 0)
        if (nrow(Pc) > 0L && length(m_nz) > 0L) {
          ck <- rep(Pc$i, times = length(m_nz))
          cj <- rep(m_nz, each = nrow(Pc))
          cx <- rep(Pc$x, times = length(m_nz)) * rep(m_row[m_nz], each = nrow(Pc))
        }
      } else {
        cli_abort(c(
          "Quadratic objective is quadratic in parameters (not DPP).",
          "i" = "Both the quad_form multiplier and its matrix depend on parameters."
        ))
      }

      if (length(ck) > 0L) {
        ## Local column-major vec index k -> (row, col) in n x n -> global flat.
        k  <- ck - 1L
        lr <- k %% n
        lc <- k %/% n
        g_flat0 <- (orig_offset + lc) * global_x + (orig_offset + lr)
        P_i <- c(P_i, g_flat0 + 1L)              # 1-based global flat row
        P_j <- c(P_j, cj)                        # param column (1-based)
        ## Factor of 2: solver minimizes 0.5 x'Px; quad_form is x'Px.
        P_x <- c(P_x, 2 * cx)
      }
    } else {
      ## Real variable: its affine-head rows become c_tensor rows.
      g0 <- extractor@id_to_col[[vid]]
      blk <- Th[Th$i >= (lo + 1L) & Th$i <= (lo + sz), , drop = FALSE]
      if (nrow(blk) > 0L) {
        c_i <- c(c_i, (blk$i - lo) + g0)         # 1-based global row
        c_j <- c(c_j, blk$j)
        c_x <- c(c_x, blk$x)
      }
    }
  }

  ## Offset row (last row of T_head) -> last row of c_tensor.
  off_blk <- Th[Th$i == (local_x_length + 1L), , drop = FALSE]
  if (nrow(off_blk) > 0L) {
    c_i <- c(c_i, rep(global_x + 1L, nrow(off_blk)))
    c_j <- c(c_j, off_blk$j)
    c_x <- c(c_x, off_blk$x)
  }

  c_tensor <- Matrix::sparseMatrix(
    i = c_i, j = c_j, x = c_x,
    dims = c(global_x + 1L, param_size_plus_one))
  P_tensor <- Matrix::sparseMatrix(
    i = P_i, j = P_j, x = P_x,
    dims = c(global_x * global_x, param_size_plus_one))

  list(P_tensor = P_tensor, c_tensor = c_tensor)
}