R/148_reductions_dcp2cone_dcp2cone.R

In CVXR: Disciplined Convex Optimization

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

## CVXPY SOURCE: reductions/dcp2cone/dcp2cone.py
## Dcp2Cone -- reduce DCP problems to conic form
##
## Uses S7 generic dispatch (dcp_canonicalize, quad_canonicalize, has_dcp_canon)
## for expression-level canonicalization. When quad_obj=TRUE, also tries
## quad_canonicalize for quadratic atoms in the objective's affine subtree (QP path).
## This is NOT the same as graph_implementation (which is LinOp-level).


## -- S7 generics for canonicalization dispatch ----------------------
## These replace the old environment-based CANON_METHODS / QUAD_CANON_METHODS
## registries with proper S7 method dispatch.
##
## INHERITANCE SAFETY INVARIANT: Every atom subclass of a canon atom MUST have
## its own explicit method() registration. Without it, S7 would silently dispatch
## to the parent's canonicalizer -- potentially producing incorrect results.
## Verified safe pairs: Log/Log1p, Pnorm/PnormApprox, Power/PowerApprox,
## GeoMean/GeoMeanApprox -- all have explicit separate registrations.

#' DCP cone canonicalization dispatch
#'
#' Replaces CANON_METHODS environment lookup. Default: identity copy.
#' Each DCP atom registers its own method returning list(canon_expr, constraints).
#' @noRd
dcp_canonicalize <- new_generic("dcp_canonicalize", "expr",
  function(expr, args, ...) {
    S7_dispatch()
  }
)

method(dcp_canonicalize, S7_object) <- function(expr, args, ...) {
  list(expr_copy(expr, args), list())
}

#' Predicate: does this expression class have a DCP cone canonicalizer?
#'
#' Used by .dcp2cone_tree() for affine_above tracking. Returns TRUE only for
#' atoms with a registered dcp_canonicalize method (NOT FiniteSet).
#' @noRd
has_dcp_canon <- new_generic("has_dcp_canon", "expr",
  function(expr) {
    S7_dispatch()
  }
)

method(has_dcp_canon, S7_object) <- function(expr) FALSE

#' Quadratic canonicalization dispatch
#'
#' Replaces QUAD_CANON_METHODS environment lookup. Default: NULL sentinel
#' (meaning "no quad canon for this class" -- fall through to dcp_canonicalize).
#' Guard logic (.quadratic_power(), is_qpwa()) is inside each method.
#' @noRd
quad_canonicalize <- new_generic("quad_canonicalize", "expr",
  function(expr, args, ...) {
    S7_dispatch()
  }
)

method(quad_canonicalize, S7_object) <- function(expr, args, ...) NULL

#' Predicate: does this expression class have a QUAD canonicalizer?
#'
#' The `has_dcp_canon` twin, needed by the CSE cache: a subtree's
#' canonicalization depends on `affine_above` only if the quad branch could fire
#' somewhere inside it (`.affine_above_relevant`). CVXPY spells this
#' `type(expr) in self.quad_canon_methods` (dcp2cone.py:291); CVXR dispatches
#' canonicalizers through S7 generics, so membership needs its own predicate.
#'
#' INHERITANCE SAFETY: same rule as `has_dcp_canon` -- every subclass of a quad
#' atom registers explicitly. PowerApprox does, beside Power.
#' @noRd
has_quad_canon <- new_generic("has_quad_canon", "expr",
  function(expr) {
    S7_dispatch()
  }
)

method(has_quad_canon, S7_object) <- function(expr) FALSE

# -- Helper: vectorize an expression (column-major) -----------------
## Equivalent to CVXPY vec(x, order='F')
.cvxr_vec <- function(x) {
  Reshape(x, c(expr_size(x), 1L), order = "F")
}

# -- Dcp2Cone class -------------------------------------------------
## CVXPY SOURCE: dcp2cone.py lines 31-145

Dcp2Cone <- new_class("Dcp2Cone", parent = Canonicalization,
  package = "CVXR",
  properties = list(
    quad_obj       = class_logical,
    solver_context = class_any  # SolverInfo or NULL
  ),
  constructor = function(quad_obj = FALSE, solver_context = NULL) {
    if (FALSE) new_object(S7_object())  ## S7 static-check guard
    .fast_new(Dcp2Cone, S7_object(),
      .cache         = new.env(parent = emptyenv()),
      quad_obj       = quad_obj,
      solver_context = solver_context
    )
  }
)

## accepts: problem must be Minimize and DCP
## CVXPY SOURCE: dcp2cone.py lines 47-50
method(reduction_accepts, Dcp2Cone) <- function(x, problem, ...) {
  .s7_is(problem@objective, Minimize) && is_dcp(problem)
}

## apply: override to pass affine_above flag
## CVXPY SOURCE: dcp2cone.py lines 52-76
method(reduction_apply, Dcp2Cone) <- function(x, problem, ...) {
  if (!reduction_accepts(x, problem)) {
    cli_abort("Cannot reduce problem to cone program: must be a minimization DCP problem.")
  }

  inverse_data <- InverseData(problem)

  ## CVXPY SOURCE: dcp2cone.py:74-77 -- the per-apply caches, created here and
  ## threaded through the whole walk so a subtree canonicalized for the
  ## objective is reused by the constraints. Upstream keeps three dicts; here
  ## they are three maps inside ONE C++ store (see subexpr_cache.R for why):
  ##   result cache     structural key -> canonicalized expression
  ##   key memo         expression id  -> interned structural key
  ##   affine_above     expression id  -> does affine_above matter here?
  ##
  ## PER-APPLY, and that is load-bearing: the affine_above memo depends on
  ## `quad_obj`, so a store shared across applies with different `quad_obj`
  ## would return answers computed under the wrong mode.
  cse <- list(store = StructuralKeyCache())

  ## Canonicalize objective (affine_above = TRUE for objective)
  obj_result <- .dcp2cone_tree(x@quad_obj, problem@objective, TRUE,
                                x@solver_context, cse)
  canon_objective <- obj_result[[1L]]

  ## Canonicalize each constraint -- collect chunks, flatten once
  n_cons <- length(problem@constraints)
  all_chunks <- vector("list", n_cons + 1L)
  all_chunks[[1L]] <- obj_result[[2L]]
  for (i in seq_len(n_cons)) {
    con <- problem@constraints[[i]]
    con_result <- .dcp2cone_tree(x@quad_obj, con, FALSE, x@solver_context, cse)
    all_chunks[[i + 1L]] <- c(con_result[[2L]], list(con_result[[1L]]))
    assign(as.character(.id(con)), .id(con_result[[1L]]),
           envir = inverse_data@cons_id_map)
  }
  canon_constraints <- unlist(all_chunks, recursive = FALSE)
  if (is.null(canon_constraints)) canon_constraints <- list()

  new_problem <- Problem(canon_objective, canon_constraints)
  list(new_problem, inverse_data)
}

# -- Dcp2Cone-specific tree walk ------------------------------------
## CVXPY SOURCE: dcp2cone.py lines 78-107
## The key difference from base Canonicalization is the affine_above tracking:
## if the path from root to current node is all affine, we may skip cone canon.

.dcp2cone_tree <- function(quad_obj, expr, affine_above, solver_context = NULL,
                            cse = NULL) {
  ## CVXPY SOURCE: dcp2cone.py:155-171
  ## Only Expression subtrees are eligible: Objectives and user Constraints are
  ## excluded so their ids flow through to inverse_data unchanged.
  cache_key <- NULL
  if (!is.null(cse) && .s7_is(expr, Expression)) {
    cache_key <- .dcp2cone_cache_key(quad_obj, expr, affine_above, cse)
    if (!is.null(cache_key)) {
      ## wrapper: .CseKeys__result_get(cse$store, cache_key)
      hit <- .Call(`_CVXR_CseKeys__result_get`, cse$store, cache_key)
      ## An empty constraint list on a hit is not an oversight: the FIRST
      ## emission already handed the generated constraints to its caller, and
      ## they are in the problem exactly once (dcp2cone.py:186-190).
      if (!is.null(hit)) return(list(hit, list()))
    }
  }

  ## Determine if this atom is affine (no DCP cone canonicalizer registered)
  affine_atom <- !has_dcp_canon(expr)

  ## Recurse into each argument -- pre-allocate, flatten once
  n_args <- length(.args(expr))
  canon_args <- vector("list", n_args)
  constr_chunks <- vector("list", n_args + 1L)
  for (i in seq_len(n_args)) {
    arg_result <- .dcp2cone_tree(quad_obj,
                                  .args(expr)[[i]],
                                  affine_atom && affine_above,
                                  solver_context,
                                  cse)
    canon_args[[i]] <- arg_result[[1L]]
    constr_chunks[[i]] <- arg_result[[2L]]
  }

  ## Canonicalize this node
  ## CVXPY passes affine_above (NOT affine_atom && affine_above) here.
  ## affine_above means "is the path ABOVE this node all affine?"
  ## This allows QuadForm (non-affine) at the top of an affine path
  ## to be dispatched to quad canon methods.
  node_result <- .dcp2cone_expr(quad_obj, expr, canon_args,
                                 affine_above = affine_above,
                                 solver_context = solver_context)
  constr_chunks[[n_args + 1L]] <- node_result[[2L]]
  constrs <- unlist(constr_chunks, recursive = FALSE)
  if (is.null(constrs)) constrs <- list()

  if (!is.null(cache_key)) {
    ## wrapper: .CseKeys__result_set(cse$store, cache_key, node_result[[1L]])
    .Call(`_CVXR_CseKeys__result_set`, cse$store, cache_key, node_result[[1L]])
  }
  list(node_result[[1L]], constrs)
}

# -- CSE cache key --------------------------------------------------
## CVXPY SOURCE: dcp2cone.py:248-273 (`_make_cache_key`)
##
## The structural key of the subtree, paired with `affine_above` ONLY when
## canonicalization could depend on it -- i.e. when the quad branch might fire
## inside this subtree. Pure cone-mode subtrees then merge across contexts,
## while a quad-eligible subtree cannot share a result between the two modes.
## Returns NULL when no safe key exists, and the caller skips the cache.

## The key is an INTEGER, not a string. `expr_key` already hands back a small
## consecutive integer, so the variant is folded in arithmetically rather than
## by pasting the integer back into text: `3k + v` is injective for
## v in {0,1,2}, which is all this ever needed.
##   v = 0  affine_above is irrelevant to this subtree
##   v = 1  relevant, affine_above TRUE
##   v = 2  relevant, affine_above FALSE

.dcp2cone_cache_key <- function(quad_obj, expr, affine_above, cse) {
  structural <- expr_key(expr, cse$store)
  if (.is_uncacheable(structural)) return(NULL)
  variant <- if (.affine_above_relevant(quad_obj, expr, cse$store)) {
    if (isTRUE(affine_above)) 1L else 2L
  } else {
    0L
  }
  3L * structural + variant
}

## CVXPY SOURCE: dcp2cone.py:275-302 (`_affine_above_relevant`)
##
## TRUE when this node or a descendant could take the quad-canon path, which
## needs `quad_obj` AND an unbroken chain of affine atoms from the root. A
## non-affine atom forces `affine_above = FALSE` for its children, so nothing
## below it can reach the quad branch through this node.

.affine_above_relevant <- function(quad_obj, expr, store) {
  if (!isTRUE(quad_obj) || !.s7_is(expr, Expression)) return(FALSE)
  expr_id <- .id(expr)
  ## wrapper: .CseKeys__aa_get(store, expr_id)
  hit <- .Call(`_CVXR_CseKeys__aa_get`, store, expr_id)
  if (!is.na(hit)) return(hit == 1L)

  relevant <- if (has_quad_canon(expr)) {
    TRUE
  } else if (has_dcp_canon(expr)) {
    FALSE
  } else {
    ## Affine atom: forwards affine_above to its children, so look below.
    any(vapply(.args(expr),
               function(a) .affine_above_relevant(quad_obj, a, store),
               logical(1L)))
  }
  ## wrapper: .CseKeys__aa_set(store, expr_id, if (relevant) 1L else 0L)
  .Call(`_CVXR_CseKeys__aa_set`, store, expr_id, if (relevant) 1L else 0L)
  relevant
}

## .dcp2cone_expr: canonicalize a single node (Dcp2Cone version)
## CVXPY SOURCE: dcp2cone.py lines 109-145
.dcp2cone_expr <- function(quad_obj, expr, args, affine_above = FALSE,
                            solver_context = NULL) {
  ## Skip constants with no parameters
  if (.s7_is(expr, Expression) &&
      is_constant(expr) && length(parameters(expr)) == 0L) {
    return(list(expr, list()))
  }

  ## QP path: try quad_canonicalize first when quad_obj=TRUE
  ## and the path above is all affine.
  ## Guard logic (.quadratic_power, is_qpwa) is inside each quad method.
  ## NULL return = guard failed or no quad method -> fall through to DCP.
  ## CVXPY SOURCE: dcp2cone.py lines 128-139
  if (quad_obj && affine_above) {
    quad_result <- quad_canonicalize(expr, args, solver_context = solver_context)
    if (!is.null(quad_result)) return(quad_result)
  }

  ## DCP cone canonicalization (S7 dispatch -- default returns identity copy).
  ## solver_context flows through `...` to canon-method implementations
  ## (each declares `solver_context = NULL`); used today by the SOC-approx
  ## warnings in power/geo_mean/pnorm approx canonicalizers.
  dcp_canonicalize(expr, args, solver_context = solver_context)
}