R/262_reductions_cone2cone_exact.R

In CVXR: Disciplined Convex Optimization

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

## CVXPY SOURCE: reductions/cone2cone/exact.py
## PARTIAL PORT: in: PowNDConversion (+ apply_hook), PSDToSvecPSD,
##   ExactCone2Cone (apply/invert, recovery chains, apply hooks).
##   out: NonPosConversion -- implemented at reductions/utilities.R
##   (`nonpos2nonneg`); SOCConversion -- not needed (no CVXR solver is
##   PSD-only). Sequencing: notes/dual_value_shape_port_plan.md.
##
## Exact cone-to-cone conversions: rewrites that replace one cone with a set of
## simpler cones the target solver actually accepts, preserving equivalence.
## New at CVXPY 1.9.0 (PR #3079/#3223/#3268); ported here at 1.9.2 under ADR
## D_19.6, which reopened the D_19.1 cluster-2 defer.
##
## Each conversion is a list with
##   source        -- the cone class converted
##   targets       -- the cone classes it maps to
##   canonicalize(con, args, solver_context) -> list(canon_constr, aux_constrs)
##   apply_hook(constraint, canon_constr, aux_constr, inverse_data) [optional]
##       -- store metadata on inverse_data right after canonicalization, read
##          later by recover_dual.
##   recover_dual(cons, dual_var, inverse_data, dvars) [optional]
##       -- map a solver dual on the converted cone back to the original's.
##
## PORTED: PowConeND -> PowCone3D (1.9.3 cycle; before it, PowConeND reached
## only solvers that accept it natively -- CLARABEL -- and exact geo_mean
## failed on SCS and MOSEK) and PSD -> SvecPSD. Not routed through this
## reduction:
##   - NonPos -> NonNeg lives in reductions/utilities.R (`nonpos2nonneg`);
##   - SOC -> PSD is unneeded: every CVXR conic solver that lacks PSD also
##     lacks nothing SOC-shaped, and no solver here is PSD-only.


# -- EXACT_CONE_CONVERSIONS ----------------------------------------
## CVXPY SOURCE: exact.py lines 61-66 -- {source_cone: {target_cones}}, which
## must form a DAG.  Consulted by `.expand_cones` (reductions/solvers/solver.R)
## when deciding whether a solver that does not accept a cone directly can
## still take it after conversion.
##
## Represented as a list of (source, targets) pairs because S7 class objects
## cannot be list names.

EXACT_CONE_CONVERSIONS <- list(
  list(source = PowConeND, targets = list(PowCone3D)),
  list(source = PSD, targets = list(SvecPSD))
)


# -- PowNDConversion helpers ---------------------------------------
## CVXPY SOURCE: exact.py lines 69-144 (_build_pow_tree)
##
## Recursively build a balanced binary tree for column `j` of W. Each 3D cone
## (x, y, z, r) is appended to the accumulator environment `acc` (upstream
## mutates four lists). The tree is recorded for column 1 only: every column
## has the same structure. Returns list(expr, total, tree).
.build_pow_tree <- function(W, z, j, indices, alphas, acc, is_root = FALSE) {
  n <- length(indices)

  if (n == 1L) {
    ## Base case: single variable, no cone needed
    tree <- if (j == 1L) SingleVarNode(var_index = indices[1L]) else NULL
    return(list(expr = W[indices[1L], j], total = alphas[1L], tree = tree))
  }

  if (n == 2L) {
    ## Base case: two variables, create one 3D cone
    total <- alphas[1L] + alphas[2L]
    out_expr <- if (is_root) z[j] else Variable(c(1L, 1L))
    .pow_acc_push(acc, W[indices[1L], j], W[indices[2L], j], out_expr,
                  alphas[1L] / total)
    tree <- if (j == 1L) LeafNode(cone_id = length(acc$x),
                                  var_indices = indices) else NULL
    return(list(expr = out_expr, total = total, tree = tree))
  }

  ## Recursive case: split into two halves
  mid <- n %/% 2L
  left <- .build_pow_tree(W, z, j, indices[seq_len(mid)],
                          alphas[seq_len(mid)], acc, is_root = FALSE)
  right <- .build_pow_tree(W, z, j, indices[(mid + 1L):n],
                           alphas[(mid + 1L):n], acc, is_root = FALSE)

  total <- left$total + right$total
  out_expr <- if (is_root) z[j] else Variable(c(1L, 1L))
  .pow_acc_push(acc, left$expr, right$expr, out_expr, left$total / total)

  tree <- if (j == 1L) {
    SplitNode(cone_id = length(acc$x), left = left$tree, right = right$tree)
  } else NULL
  list(expr = out_expr, total = total, tree = tree)
}

.pow_acc_push <- function(acc, x, y, z, r) {
  i <- length(acc$x) + 1L
  acc$x[[i]] <- x
  acc$y[[i]] <- y
  acc$z[[i]] <- z
  acc$alpha[[i]] <- r
  invisible(NULL)
}

## CVXPY SOURCE: exact.py lines 147-179 (_extract_pow_duals_recursive)
##
## Walk the tree for column `j`, filling W duals from the raw 3 x m PowCone3D
## duals (rows x, y, z). A SingleVarNode takes its dual from the parent cone's
## x (left child) or y (right child) row; a LeafNode from its own cone. Returns
## the updated w_duals (upstream mutates it in place).
.extract_pow_duals_recursive <- function(node, parent_cone_id, is_left_child,
                                         raw_duals, w_duals, j,
                                         num_cones_per_col) {
  if (inherits(node, "SingleVarNode")) {
    col <- (j - 1L) * num_cones_per_col + parent_cone_id
    row <- if (is_left_child) 1L else 2L
    w_duals[node$var_index, j] <- raw_duals[row, col]
  } else if (inherits(node, "LeafNode")) {
    col <- (j - 1L) * num_cones_per_col + node$cone_id
    w_duals[node$var_indices[1L], j] <- raw_duals[1L, col]
    w_duals[node$var_indices[2L], j] <- raw_duals[2L, col]
  } else if (inherits(node, "SplitNode")) {
    w_duals <- .extract_pow_duals_recursive(node$left, node$cone_id, TRUE,
                                            raw_duals, w_duals, j,
                                            num_cones_per_col)
    w_duals <- .extract_pow_duals_recursive(node$right, node$cone_id, FALSE,
                                            raw_duals, w_duals, j,
                                            num_cones_per_col)
  }
  w_duals
}

## CVXPY SOURCE: exact.py lines 182-199 (_extract_pow_duals)
## W and z duals from the raw 3D cone duals, using the tree structure.
.extract_pow_duals <- function(tree, raw_duals, n, k, num_cones_per_col) {
  w_duals <- matrix(0, n, k)
  z_duals <- numeric(k)
  root_idx <- get_root_cone_id(tree)
  for (j in seq_len(k)) {
    w_duals <- .extract_pow_duals_recursive(tree, NULL, TRUE, raw_duals,
                                            w_duals, j, num_cones_per_col)
    if (!is.null(root_idx)) {
      z_duals[j] <- raw_duals[3L, (j - 1L) * num_cones_per_col + root_idx]
    }
  }
  list(w_duals = w_duals, z_duals = z_duals)
}


# -- PowNDConversion -----------------------------------------------
## CVXPY SOURCE: exact.py lines 202-305
##
## PowConeND -> PowCone3D via balanced binary tree decomposition: O(log n)
## depth instead of a linear chain, for numerical stability. R axis 2L (one
## cone per column) is CVXPY axis 0; axis 1L is transposed first, as upstream.

PowNDConversion <- list(
  source  = PowConeND,
  targets = list(PowCone3D),

  ## CVXPY SOURCE: exact.py lines 212-267
  canonicalize = function(con, args, solver_context = NULL) {
    alpha <- as.matrix(value(con@alpha))
    W <- args[[1L]]
    z <- args[[2L]]
    if (con@axis == 1L) {
      W <- t(W)
      alpha <- t(alpha)
    }
    n <- .shape(W)[1L]
    k <- .shape(W)[2L]

    if (n == 2L) {
      ## One 3D cone per column; z laid out like W's rows (1 x k).
      z_row <- if (identical(.shape(z), c(1L, k))) z else reshape_expr(z, c(1L, k))
      can_con <- PowCone3D(W[1L, ], W[2L, ], z_row,
                           matrix(alpha[1L, ], nrow = 1L))
      return(list(can_con, list()))
    }

    ## Balanced tree decomposition for n > 2
    acc <- new.env(parent = emptyenv())
    acc$x <- list(); acc$y <- list(); acc$z <- list(); acc$alpha <- list()
    tree <- NULL
    for (j in seq_len(k)) {
      res <- .build_pow_tree(W, z, j, seq_len(n), alpha[, j], acc,
                             is_root = TRUE)
      if (j == 1L) tree <- res$tree
    }

    can_con <- PowCone3D(do.call(hstack, acc$x), do.call(hstack, acc$y),
                         do.call(hstack, acc$z),
                         matrix(unlist(acc$alpha), nrow = 1L))
    ## Upstream sets `can_con._pow_tree`; the constraint's cache environment
    ## plays that role here, read back by apply_hook.
    assign("pow_tree", tree, envir = can_con@.cache)
    list(can_con, list())
  },

  ## CVXPY SOURCE: exact.py lines 270-282
  ## Store the tree on inverse_data for recover_dual, keyed by the ORIGINAL
  ## constraint's id.
  apply_hook = function(constraint, canon_constr, aux_constr, inverse_data) {
    tree <- get0("pow_tree", envir = canon_constr@.cache, inherits = FALSE)
    if (!is.null(tree)) {
      maps <- get0("pow_tree_mappings", envir = inverse_data@.extra,
                   ifnotfound = list())
      maps[[as.character(.id(constraint))]] <- tree
      assign("pow_tree_mappings", maps, envir = inverse_data@.extra)
    }
    invisible(NULL)
  },

  ## CVXPY SOURCE: exact.py lines 285-305
  ## Returns the (n + 1) x k layout save_dual_value(PowConeND) reads: W duals
  ## stacked above the z duals.
  recover_dual = function(cons, dual_var, inverse_data, dvars) {
    alpha_shape <- .shape(cons@alpha)
    n <- if (cons@axis == 2L) alpha_shape[1L] else alpha_shape[2L]
    k <- expr_size(cons@.z)
    ## The raw PowCone3D dual: rows x, y, z, one column per cone (column-major,
    ## so a flat vector in solver order gives the same matrix).
    raw <- matrix(as.numeric(dual_var), nrow = 3L)

    maps <- get0("pow_tree_mappings", envir = inverse_data@.extra,
                 ifnotfound = NULL)
    tree <- if (is.null(maps)) NULL else maps[[as.character(.id(cons))]]
    if (!is.null(tree)) {
      num_cones_per_col <- ncol(raw) %/% k
      d <- .extract_pow_duals(tree, raw, n, k, num_cones_per_col)
      w_duals <- d$w_duals
      z_duals <- d$z_duals
    } else {
      ## n == 2: direct mapping, no tree decomposition
      w_duals <- raw[1:2, , drop = FALSE]
      z_duals <- raw[3L, ]
    }
    rbind(w_duals, z_duals, deparse.level = 0L)
  }
)


# -- PSDToSvecPSD --------------------------------------------------
## CVXPY SOURCE: exact.py lines 420-446
##
## PSD -> SvecPSD via scaled vectorization.  Which triangle, and whether the
## off-diagonal is sqrt(2)-scaled, come from the solver context -- the same two
## values the solver declares as PSD_TRIANGLE_KIND / PSD_SQRT2_SCALING.

PSDToSvecPSD <- list(
  source  = PSD,
  targets = list(SvecPSD),

  ## CVXPY SOURCE: exact.py lines 430-437
  canonicalize = function(con, args, solver_context = NULL) {
    X <- args[[1L]]
    n <- .shape(X)[2L]
    M <- psd_format_mat(con, solver_context@psd_triangle_kind,
                        solver_context@psd_sqrt2_scaling)
    ## `vec(X)` is CVXPY's `vec(X, order='F')`: R is column-major throughout.
    svec_expr <- Constant(M) %*% vec(X)
    list(SvecPSD(svec_expr, n = n), list())
  },

  ## CVXPY SOURCE: exact.py lines 439-446
  recover_dual = function(cons, dual_var, inverse_data, dvars) {
    ctx <- get0("solver_context", envir = inverse_data@.extra, ifnotfound = NULL)
    n <- .arg_shape(cons)[2L]
    ## Upstream reshapes `tri_to_full`'s output to the argument's shape; CVXR's
    ## `tri_to_full` already returns the (n, n) matrix, and a PSD constraint's
    ## argument is square by construction (psd.R constructor).
    tri_to_full(dual_var, n, ctx@psd_triangle_kind, ctx@psd_sqrt2_scaling)
  }
)

## All conversions this reduction knows how to perform.
## CVXPY SOURCE: exact.py line 451 (upstream order, minus the two not routed
## here: NonPosConversion, SOCConversion)
.EXACT_CONVERSIONS <- list(PowNDConversion, PSDToSvecPSD)


# -- ExactCone2Cone ------------------------------------------------
## CVXPY SOURCE: exact.py lines 449-476
##
## `target_cones` restricts the reduction to the conversions the chain actually
## asked for (computed by `.expand_cones`); NULL means "all of them".
##
## Upstream subclasses `Canonicalization` to reuse `canonicalize_tree`.  CVXR's
## `Canonicalization` walks EXPRESSION trees through S7 generics and has no
## canon_methods dict, and by this point in the chain every constraint argument
## is already affine -- so the expression walk would be a no-op.  This is a
## `Reduction` that converts the constraints directly, which is what
## `canonicalize_tree` reduces to for affine arguments.

ExactCone2Cone <- new_class("ExactCone2Cone", parent = Reduction,
  package = "CVXR",
  properties = list(
    target_cones   = class_list,
    solver_context = class_any
  ),
  constructor = function(target_cones = list(), solver_context = NULL) {
    if (FALSE) new_object(S7_object())  ## S7 static-check guard
    .fast_new(ExactCone2Cone, S7_object(),
      .cache         = new.env(parent = emptyenv()),
      target_cones   = target_cones,
      solver_context = solver_context
    )
  }
)

## The conversion applicable to `con`, or NULL.
## CVXPY SOURCE: exact.py lines 457-460 (the `target_cones` filter) and the
## `type(current) in self.canon_methods` test at line 486.
## Membership by class NAME (`.cone_has`), as everywhere else in the chain:
## upstream's `c.source in target_cones` is type identity on a Python class,
## which is a singleton; an R class object is a closure, and `identical()` on
## it is a structural walk that fails on an instrumented copy (covr, trace())
## and costs 10-120 us per miss (sweep item 4, pass J).
.exact_conversion_for <- function(x, con) {
  for (conv in .EXACT_CONVERSIONS) {
    if (!.s7_is(con, conv$source)) next
    if (length(x@target_cones) > 0L && !.cone_has(x@target_cones, conv$source)) next
    return(conv)
  }
  NULL
}

## CVXPY SOURCE: exact.py lines 478-503
## Follow the conversion chain transitively, recording the constraints passed
## through so `invert` can unwind them in reverse.
.convert_constraint <- function(x, constraint, inverse_data, canon_constraints) {
  chain <- list()
  current <- constraint

  repeat {
    conv <- .exact_conversion_for(x, current)
    if (is.null(conv)) break
    chain[[length(chain) + 1L]] <- current
    res <- conv$canonicalize(current, .args(current), x@solver_context)
    canon_constr <- res[[1L]]
    aux_constr <- res[[2L]]

    ## CVXPY SOURCE: exact.py lines 490-492 -- per-conversion apply hook
    if (!is.null(conv$apply_hook)) {
      conv$apply_hook(current, canon_constr, aux_constr, inverse_data)
    }

    ## Auxiliary constraints convert recursively too.
    for (aux in aux_constr) {
      canon_constraints <- .convert_constraint(x, aux, inverse_data,
                                               canon_constraints)
    }
    current <- canon_constr
  }

  canon_constraints[[length(canon_constraints) + 1L]] <- current
  assign(as.character(.id(constraint)), .id(current),
         envir = inverse_data@cons_id_map)
  if (length(chain) > 0L) {
    chains <- get0("recovery_chains", envir = inverse_data@.extra,
                   ifnotfound = list())
    chains[[as.character(constraint@id)]] <- chain
    assign("recovery_chains", chains, envir = inverse_data@.extra)
  }
  canon_constraints
}

## CVXPY SOURCE: exact.py lines 505-520
method(reduction_apply, ExactCone2Cone) <- function(x, problem, ...) {
  inverse_data <- InverseData(problem)
  assign("solver_context", x@solver_context, envir = inverse_data@.extra)

  canon_constraints <- list()
  for (constraint in problem@constraints) {
    canon_constraints <- .convert_constraint(x, constraint, inverse_data,
                                             canon_constraints)
  }

  list(Problem(problem@objective, canon_constraints), inverse_data)
}

## CVXPY SOURCE: exact.py lines 522-551
method(reduction_invert, ExactCone2Cone) <- function(x, solution, inverse_data, ...) {
  ## Re-key the duals from converted-constraint ids onto original ids.
  ## O(n) via one vectorized match; see `.remap_by_id_map`.
  dvars <- .remap_by_id_map(solution@dual_vars, inverse_data@cons_id_map)

  if (length(dvars) == 0L) {
    return(Solution(status      = solution@status,
                    opt_val     = solution@opt_val,
                    primal_vars = solution@primal_vars,
                    dual_vars   = dvars,
                    attr        = solution@attr))
  }

  ## Unwind each recovery chain in reverse, so transitive conversions come
  ## apart in the order they were applied.
  chains <- get0("recovery_chains", envir = inverse_data@.extra,
                 ifnotfound = list())
  for (orig_id in names(chains)) {
    if (is.null(dvars[[orig_id]])) next
    dual <- dvars[[orig_id]]
    chain <- chains[[orig_id]]
    for (i in rev(seq_along(chain))) {
      cons <- chain[[i]]
      conv <- .exact_conversion_for(x, cons)
      if (!is.null(conv) && !is.null(conv$recover_dual)) {
        dual <- conv$recover_dual(cons, dual, inverse_data, dvars)
      }
    }
    dvars[[orig_id]] <- dual
  }

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