R/230_reductions_dcp2cone_cone_matrix_stuffing.R

In CVXR: Disciplined Convex Optimization

#####
## DO NOT EDIT THIS FILE!! EDIT THE SOURCE INSTEAD: rsrc_tree/reductions/dcp2cone/cone_matrix_stuffing.R
#####

## CVXPY SOURCE: reductions/dcp2cone/cone_matrix_stuffing.py
## ConeMatrixStuffing -- convert affine expressions to sparse matrices
## ConeDims -- summary of cone dimensions


# ==================================================================
# ConeDims -- summary of cone dimensions present in constraints
# ==================================================================
## CVXPY SOURCE: cone_matrix_stuffing.py lines 57-141

ConeDims <- new_class("ConeDims", package = "CVXR",
  properties = list(
    zero   = class_integer,
    nonneg = class_integer,
    exp    = class_integer,
    soc    = class_integer,     # vector of SOC cone dims
    psd    = class_integer,     # vector of PSD sizes (n for n x n)
    p3d    = class_numeric,     # vector of PowCone3D alphas
    pnd    = class_list         # list of numeric vectors for PowConeND
  ),
  constructor = function(constr_map) {
    ## Zero cone dimension
    zero_constrs <- constr_map[["Zero"]]
    zero <- if (length(zero_constrs) == 0L) 0L
            else as.integer(sum(vapply(zero_constrs, constr_size, numeric(1L))))

    ## NonNeg cone dimension
    nn_constrs <- constr_map[["NonNeg"]]
    nonneg <- if (length(nn_constrs) == 0L) 0L
              else as.integer(sum(vapply(nn_constrs, constr_size, numeric(1L))))

    ## Exponential cone count
    exp_constrs <- constr_map[["ExpCone"]]
    exp_count <- if (length(exp_constrs) == 0L) 0L
                 else as.integer(sum(vapply(exp_constrs, num_cones, numeric(1L))))

    ## SOC dimensions (one entry per individual cone)
    soc_constrs <- constr_map[["SOC"]]
    soc <- if (length(soc_constrs) == 0L) integer(0)
           else as.integer(unlist(lapply(soc_constrs, cone_sizes)))

    ## PSD dimensions (n for each n x n PSD constraint)
    psd_constrs <- constr_map[["PSD"]]
    psd <- if (length(psd_constrs) == 0L) integer(0)
           else vapply(psd_constrs, function(c) c@shape[1L], integer(1L))

    ## PowCone3D alphas
    p3d_constrs <- constr_map[["PowCone3D"]]
    p3d <- if (length(p3d_constrs) == 0L) numeric(0)
           else unlist(lapply(p3d_constrs, function(c) as.numeric(value(c@alpha))))

    ## PowConeND alphas
    pnd_constrs <- constr_map[["PowConeND"]]
    pnd <- if (length(pnd_constrs) == 0L) list()
           else {
             alpha_chunks <- vector("list", length(pnd_constrs))
             for (ci in seq_along(pnd_constrs)) {
               a <- value(pnd_constrs[[ci]]@alpha)
               if (is.matrix(a)) {
                 ## Each column is one cone's alpha vector
                 alpha_chunks[[ci]] <- lapply(seq_len(ncol(a)), function(j) a[, j])
               } else {
                 alpha_chunks[[ci]] <- list(as.numeric(a))
               }
             }
             unlist(alpha_chunks, recursive = FALSE)
           }

    new_object(S7_object(),
      zero = zero, nonneg = nonneg, exp = exp_count,
      soc = soc, psd = psd, p3d = p3d, pnd = pnd)
  }
)

method(print, ConeDims) <- function(x, ...) {
  cat(sprintf("%d equalities, %d inequalities, %d exponential cones\n",
    x@zero, x@nonneg, x@exp))
  cat(sprintf("SOC: %s, PSD: %s\n",
    paste(x@soc, collapse = ","),
    paste(x@psd, collapse = ",")))
  invisible(x)
}

# -- dims_to_solver_dict -------------------------------------------
## CVXPY SOURCE: conic_solver.py lines 87-97

dims_to_solver_dict <- function(cone_dims) {
  list(
    f   = cone_dims@zero,
    l   = cone_dims@nonneg,
    q   = cone_dims@soc,
    ep  = cone_dims@exp,
    s   = cone_dims@psd,
    p   = cone_dims@p3d,
    pnd = cone_dims@pnd
  )
}


# -- extract_mip_idx -----------------------------------------------
## CVXPY SOURCE: matrix_stuffing.py lines 79-96
## Maps per-variable boolean/integer flags to global flattened variable indices.

.extract_mip_idx <- function(problem, inverse_data) {
  vars <- variables(problem)
  bool_chunks <- vector("list", length(vars))
  int_chunks  <- vector("list", length(vars))

  for (i in seq_along(vars)) {
    var <- vars[[i]]
    vid <- as.character(var@id)
    offset <- inverse_data@var_offsets[[vid]]
    sz <- expr_size(var)
    global_indices <- seq(offset + 1L, offset + sz)  # 1-based R indices

    if (isTRUE(var@attributes$boolean)) {
      bool_chunks[[i]] <- global_indices
    } else if (isTRUE(var@attributes$integer)) {
      int_chunks[[i]] <- global_indices
    }
  }

  bool_idx <- unlist(bool_chunks)
  int_idx  <- unlist(int_chunks)
  list(bool_idx = if (is.null(bool_idx)) integer(0) else bool_idx,
       int_idx  = if (is.null(int_idx))  integer(0) else int_idx)
}

# ==================================================================
# ConeMatrixStuffing -- reduction from affine expressions to matrices
# ==================================================================
## CVXPY SOURCE: cone_matrix_stuffing.py lines 321-470

ConeMatrixStuffing <- new_class("ConeMatrixStuffing", parent = Reduction,
  package = "CVXR",
  properties = list(
    quad_obj = class_logical
  ),
  constructor = function(quad_obj = FALSE) {
    new_object(S7_object(),
      .cache = new.env(parent = emptyenv()),
      quad_obj = quad_obj
    )
  }
)

## accepts: affine (or quadratic when quad_obj) objective (Minimize),
## no convex attributes, all constraint args affine
## CVXPY SOURCE: cone_matrix_stuffing.py lines 335-342
method(reduction_accepts, ConeMatrixStuffing) <- function(x, problem, ...) {
  is_min <- S7_inherits(problem@objective, Minimize)
  obj_expr <- problem@objective@args[[1L]]
  valid_obj <- if (x@quad_obj) {
    is_affine(obj_expr) || is_quadratic(obj_expr)
  } else {
    is_affine(obj_expr)
  }
  no_cvx_attr <- length(convex_attributes(variables(problem))) == 0L
  aff_con <- are_args_affine(problem@constraints)
  is_min && valid_obj && no_cvx_attr && aff_con
}

## apply: lower constraints and extract A, b, c matrices
## CVXPY SOURCE: cone_matrix_stuffing.py lines 360-430
method(reduction_apply, ConeMatrixStuffing) <- function(x, problem, ...) {
  inverse_data <- InverseData(problem)

  ## Step 1: Lower constraints -- pre-allocate
  n_cons_raw <- length(problem@constraints)
  cons <- vector("list", n_cons_raw)
  for (i in seq_len(n_cons_raw)) {
    con <- problem@constraints[[i]]
    if (S7_inherits(con, Equality)) {
      con <- lower_equality(con)
    } else if (S7_inherits(con, Inequality)) {
      con <- lower_ineq_to_nonneg(con)
    } else if (S7_inherits(con, NonPos)) {
      con <- nonpos2nonneg(con)
    } else if (S7_inherits(con, SOC) && con@axis == 1L) {
      ## Transpose X for axis=1 -> axis=2
      con <- SOC(con@args[[1L]], t(con@args[[2L]]), axis = 2L,
                 constr_id = con@id)
    } else if (S7_inherits(con, PowConeND) && con@axis == 1L) {
      ## CVXPY SOURCE: cone_matrix_stuffing.py lines 382-388
      ## Transpose W and alpha for axis=1 -> axis=2
      con <- PowConeND(t(con@.W), con@.z,
                        t(con@alpha), axis = 2L,
                        constr_id = con@id)
    }
    ## ExpCone/PowCone3D flattening: only needed if multidimensional
    ## For Phase 5b, our canonicalizers produce 1D args; skip flattening for now
    cons[[i]] <- con
  }

  ## Step 2: Group and order constraints
  ## CVXPY order: Zero, NonNeg, SOC, PSD, ExpCone, PowCone3D, PowConeND
  constr_map <- group_constraints(cons)
  ordered_cons <- c(constr_map[["Zero"]], constr_map[["NonNeg"]],
                    constr_map[["SOC"]], constr_map[["PSD"]],
                    constr_map[["ExpCone"]], constr_map[["PowCone3D"]],
                    constr_map[["PowConeND"]])

  ## Step 3: Store constraint ID mapping (identity -- already lowered)
  ## CVXPY SOURCE: cone_matrix_stuffing.py line 406
  id2cons <- new.env(hash = TRUE, parent = emptyenv())
  for (con in ordered_cons) {
    assign(as.character(con@id), con@id, envir = inverse_data@cons_id_map)
    assign(as.character(con@id), con, envir = id2cons)
  }

  ## Step 4: Create CoeffExtractor and extract objective
  extractor <- CoeffExtractor(inverse_data)

  ## Create the flattened variable
  x_var <- Variable(c(extractor@x_length, 1L))

  ## Detect DPP path: parameters exist and no EvalParams was applied
  ## (i.e., parameters are still present in the expressions).
  has_params <- length(parameters(problem)) > 0L

  ## Extract objective
  P_mat <- NULL
  c_tensor <- NULL
  P_tensor <- NULL
  pv <- NULL  ## parameter vector, lazily computed when has_params
  if (x@quad_obj) {
    ## QP path: extract 0.5*x'*P*x + q'*x + d
    quad_result <- coeff_quad_form(extractor, problem@objective@args[[1L]])
    P_mat <- quad_result$P          # already 2x scaled for solver
    c_vec <- quad_result$q          # linear term
    d_offset <- quad_result$offset  # constant
    ## TODO: QP tensor path for DPP (P_tensor, c_tensor)
  } else {
    if (has_params) {
      ## DPP tensor path: extract c_tensor (x_length+1, param_size)
      c_tensor <- coeff_affine_tensor(extractor, list(problem@objective@args[[1L]]))
      ## Also get current numeric values via apply
      pv <- get_parameter_vector(extractor@param_to_size,
                                  extractor@param_id_map,
                                  parameters(problem))
      c_flat <- as.numeric(c_tensor %*% pv)
      c_vec <- c_flat[seq_len(extractor@x_length)]
      d_offset <- c_flat[extractor@x_length + 1L]
    } else {
      ## LP/conic path: extract c'*x + d
      obj_result <- coeff_affine(extractor, list(problem@objective@args[[1L]]))
      c_vec <- as.numeric(obj_result$A)  # x_length vector
      d_offset <- as.numeric(obj_result$b)  # scalar offset
    }
  }

  ## Step 5: Batch constraint args and extract A, b
  expr_chunks <- vector("list", length(ordered_cons))
  for (i in seq_along(ordered_cons)) {
    expr_chunks[[i]] <- ordered_cons[[i]]@args
  }
  expr_list <- unlist(expr_chunks, recursive = FALSE)
  if (is.null(expr_list)) expr_list <- list()

  A_tensor <- NULL
  if (length(expr_list) > 0L) {
    if (has_params) {
      ## DPP tensor path: extract A_tensor
      A_tensor <- coeff_affine_tensor(extractor, expr_list)
      ## Also get current numeric values via apply
      if (is.null(pv)) {
        pv <- get_parameter_vector(extractor@param_to_size,
                                    extractor@param_id_map,
                                    parameters(problem))
      }
      Ab_flat <- as.numeric(A_tensor %*% pv)
      n_cols <- extractor@x_length + 1L
      n_rows <- length(Ab_flat) %/% n_cols
      Ab_mat <- matrix(Ab_flat, nrow = n_rows, ncol = n_cols)
      ## Force dgCMatrix (not triangular/symmetric auto-detected types)
      A <- as(Matrix::Matrix(Ab_mat[, seq_len(extractor@x_length), drop = FALSE],
                              sparse = TRUE), "generalMatrix")
      b <- Ab_mat[, n_cols]
    } else {
      con_result <- coeff_affine(extractor, expr_list)
      A <- con_result$A
      b <- con_result$b
    }
  } else {
    A <- Matrix::sparseMatrix(i = integer(0), j = integer(0), x = numeric(0),
                              dims = c(0L, extractor@x_length))
    b <- numeric(0)
  }

  ## Step 6: Build ConeDims
  cone_dims <- ConeDims(constr_map)

  ## Step 7: Store extra data for inversion
  inverse_data@.extra <- new.env(hash = TRUE, parent = emptyenv())
  inverse_data@.extra$minimize <- TRUE
  inverse_data@.extra$constraints <- ordered_cons
  inverse_data@.extra$id2cons <- id2cons

  ## Return data as named list + inverse data
  data <- list()
  data[[SD_C]]      <- c_vec
  data[[SD_OFFSET]]  <- d_offset
  data[[SD_A]]      <- A
  data[[SD_B]]      <- b
  data[[SD_DIMS]]   <- cone_dims
  data[["x_id"]]    <- x_var@id
  data[["constraints"]] <- ordered_cons
  if (!is.null(P_mat)) data[[SD_P]] <- P_mat

  ## Step 8: Build ParamConeProg for DPP caching
  if (has_params && !is.null(c_tensor) && !is.null(A_tensor)) {
    data[[SD_PARAM_PROB]] <- ParamConeProg(
      c_tensor = c_tensor,
      A_tensor = A_tensor,
      x_length = extractor@x_length,
      x_id = x_var@id,
      parameters = parameters(problem),
      param_id_to_col = extractor@param_id_map,
      param_to_size = extractor@param_to_size,
      variables = variables(problem),
      var_id_to_col = inverse_data@var_offsets,
      constraints = ordered_cons,
      cone_dims = cone_dims,
      P_tensor = P_tensor
    )
  }

  ## Extract MIP indices for boolean/integer variables
  mip_idx <- .extract_mip_idx(problem, inverse_data)
  data[[SD_BOOL_IDX]] <- mip_idx$bool_idx
  data[[SD_INT_IDX]]  <- mip_idx$int_idx

  list(data, inverse_data)
}

## invert: split solution vector back into per-variable values
## CVXPY SOURCE: cone_matrix_stuffing.py lines 432-470
method(reduction_invert, ConeMatrixStuffing) <- function(x, solution, inverse_data, ...) {
  var_map <- inverse_data@var_offsets

  ## Flip sign of opt val if maximize
  opt_val <- solution@opt_val
  if (!(solution@status %in% ERROR_STATUS) &&
      !is.null(inverse_data@.extra$minimize) &&
      !inverse_data@.extra$minimize) {
    opt_val <- -opt_val
  }

  primal_vars <- list()
  dual_vars <- list()

  if (!(solution@status %in% SOLUTION_PRESENT)) {
    return(Solution(solution@status, opt_val, primal_vars, dual_vars,
                    solution@attr))
  }

  ## Split vectorized variable into components
  x_opt <- solution@primal_vars[[1L]]
  if (is.null(x_opt)) {
    ## Try to get the single primal var
    if (length(solution@primal_vars) > 0L) {
      x_opt <- solution@primal_vars[[1L]]
    }
  }

  if (!is.null(x_opt)) {
    for (var_id in names(var_map)) {
      offset <- var_map[[var_id]]
      shape <- inverse_data@var_shapes[[var_id]]
      sz <- prod(shape)
      val <- x_opt[(offset + 1L):(offset + sz)]
      primal_vars[[var_id]] <- matrix(val, nrow = shape[1L], ncol = shape[2L])
    }
  }

  ## Remap dual variables
  if (length(solution@dual_vars) > 0L) {
    con_map_env <- inverse_data@cons_id_map
    id2cons_env <- inverse_data@.extra$id2cons
    old_ids <- ls(con_map_env, all.names = TRUE)
    for (old_id in old_ids) {
      new_id <- as.character(get(old_id, envir = con_map_env))
      if (!is.null(solution@dual_vars[[new_id]])) {
        dual_vars[[old_id]] <- solution@dual_vars[[new_id]]
      }
    }
  }

  Solution(solution@status, opt_val, primal_vars, dual_vars, solution@attr)
}