R/252_transforms_partial_optimize.R

In CVXR: Disciplined Convex Optimization

#####
## DO NOT EDIT THIS FILE!! EDIT THE SOURCE INSTEAD: rsrc_tree/transforms/partial_optimize.R
#####

## CVXPY SOURCE: transforms/partial_optimize.py
##
## CVXPY 1.9 parity notes:
##  - DNLP linearizable predicates mirror partial_optimize.py:145-148:
##    PartialProblem delegates linearizable-convex / linearizable-concave to
##    its convex / concave predicates.
##
## partial_optimize takes a Problem and turns it into an Expression
## (PartialProblem) representing the optimal value of that problem
## as a function of a subset of its variables (`dont_opt_vars`).
##
## Typical use:
##   x <- Variable(n); t <- Variable(n)
##   abs_x <- partial_optimize(
##     Problem(Minimize(sum(t)), list(-t <= x, x <= t)),
##     opt_vars = list(t))
##
## The resulting PartialProblem can be embedded in a larger problem.
## When that larger problem is canonicalized, the inner objective +
## inner constraints flow through into the outer cone-form data,
## because PartialProblem's `variables()` / `parameters()` / `constants()`
## methods delegate to the inner problem (so fresh inner variables are
## visible to the outer ConeMatrixStuffing) and the `canonicalize()`
## method emits the inner objective's canonical form plus every inner
## constraint's canonical form (mirrors CVXPY:308-311).

# ==================================================================
# .partial_optimize_subs -- helper: tree_copy with id substitution
# ==================================================================
## CVXPY SOURCE: partial_optimize.py:94-99 (tree_copy with id_to_new_var).
##
## CVXR's tree_copy(Canonical) (canonical.R:89-99) memoises by
## `as.character(@id)`.  Pre-seeding the id_objects env with
## { as.character(old@id) -> new_var } makes the walk return the
## replacement whenever the original is encountered.

.partial_optimize_subs <- function(expr, id_to_new_var) {
  env <- new.env(hash = TRUE, parent = emptyenv())
  for (key in names(id_to_new_var)) {
    assign(key, id_to_new_var[[key]], envir = env)
  }
  tree_copy(expr, env)
}

# ==================================================================
# PartialProblem class
# ==================================================================
## CVXPY SOURCE: partial_optimize.py:104-311
##
## PartialProblem inherits from Expression (NOT Atom).  Like Problem
## itself, it's a special compound Expression: every method needs to
## be explicitly defined -- it does not get the Atom defaults for
## expr_name / grad / canonicalize.

#' Partial optimization of a Problem
#'
#' A `PartialProblem` is an Expression that represents the optimal
#' value of an inner Problem as a function of the variables you choose
#' NOT to optimise over.  Build one with [partial_optimize()] rather
#' than constructing the class directly.
#'
#' @keywords internal
#' @export
## NOTE: CVXPY stores the inner Problem inside `self.args = [prob]`,
## but Problem in CVXR is NOT a Canonical/Expression subclass and has
## no `@args` slot.  The outer Dcp2Cone tree walker recurses into
## every arg's `@args`, which would fail.  We keep `@args = list()`
## (PartialProblem is a leaf to the tree walker) and stash the inner
## Problem in a dedicated `inner_problem` property.

PartialProblem <- new_class("PartialProblem", parent = Expression,
  package = "CVXR",
  properties = list(
    inner_problem = class_any,    # Problem
    opt_vars      = class_list,
    dont_opt_vars = class_list,
    solver        = class_any,    # NULL or character
    solve_kwargs  = class_list
  ),
  constructor = function(prob, opt_vars, dont_opt_vars,
                          solver = NULL, solve_kwargs = list(),
                          id = NULL) {
    if (FALSE) new_object(S7_object())  ## S7 static-check guard
    if (is.null(id)) id <- next_expr_id()
    if (!.s7_is(prob, Problem)) {
      cli_abort("{.cls PartialProblem} requires a {.cls Problem} object.")
    }
    .fast_new(PartialProblem, S7_object(),
      id            = as.integer(id),
      .cache        = new.env(parent = emptyenv()),
      args          = list(),       # leaf to the tree walker
      shape         = c(1L, 1L),
      inner_problem = prob,
      opt_vars      = opt_vars,
      dont_opt_vars = dont_opt_vars,
      solver        = solver,
      solve_kwargs  = solve_kwargs
    )
  }
)

# ==================================================================
# partial_optimize -- the public factory function
# ==================================================================
## CVXPY SOURCE: partial_optimize.py:29-101.

#' Partial optimization transform
#'
#' Builds an Expression representing the optimal value of `prob` as
#' a function of the variables you choose NOT to optimise over.  Useful
#' for two-stage / hierarchical optimisation, custom atom definitions,
#' and embedding sub-problems inside larger problems.
#'
#' Exactly one of `opt_vars` or `dont_opt_vars` may be `NULL`; the
#' missing list is taken to be the complement (relative to the full
#' list of variables in `prob`).  If both are supplied, they must
#' together cover every variable in `prob`.
#'
#' The returned `PartialProblem` is an Expression with scalar shape:
#' it is convex when `prob` is DCP with a `Minimize` objective and
#' concave when DCP with `Maximize`.  Embed it like any other
#' expression in a larger Problem; the larger problem's canonicalizer
#' will pull the inner objective and constraints into the outer cone
#' form so a single solve handles both layers.
#'
#' @param prob A [Problem] to partially optimise.
#' @param opt_vars Optional list of [Variable]s to optimise over.
#' @param dont_opt_vars Optional list of [Variable]s to keep as free
#'   arguments of the resulting expression.
#' @param solver Optional solver name (passed to [psolve()] when the
#'   PartialProblem is evaluated via [value()] or [grad()]).
#' @param ... Additional named arguments forwarded to [psolve()] when
#'   `value()` / `grad()` are called.
#' @returns A `PartialProblem` expression.
#'
#' @examples
#' \dontrun{
#' x <- Variable(3)
#' t <- Variable(3)
#' abs_x <- partial_optimize(
#'   Problem(Minimize(sum_entries(t)), list(-t <= x, x <= t)),
#'   opt_vars = list(t)
#' )
#' ## abs_x is now an expression of x alone, equivalent to sum(abs(x)).
#' }
#'
#' @seealso [Problem], [psolve()]
#' @export
partial_optimize <- function(prob, opt_vars = NULL, dont_opt_vars = NULL,
                              solver = NULL, ...) {
  if (!.s7_is(prob, Problem)) {
    cli_abort("{.fn partial_optimize} requires a {.cls Problem} object.")
  }
  if (is.null(opt_vars) && is.null(dont_opt_vars)) {
    cli_abort(c(
      "{.fn partial_optimize} called with neither {.arg opt_vars} nor {.arg dont_opt_vars}.",
      "i" = "Specify one (the other is filled in as the complement)."
    ))
  }
  all_vars <- variables(prob)
  all_ids  <- vapply(all_vars, function(v) as.character(v@id), character(1L))

  if (is.null(opt_vars)) {
    do_not_ids <- vapply(dont_opt_vars,
                          function(v) as.character(v@id), character(1L))
    opt_vars <- all_vars[!all_ids %in% do_not_ids]
  } else if (is.null(dont_opt_vars)) {
    opt_ids <- vapply(opt_vars,
                       function(v) as.character(v@id), character(1L))
    dont_opt_vars <- all_vars[!all_ids %in% opt_ids]
  } else {
    given_ids <- c(
      vapply(opt_vars,      function(v) as.character(v@id), character(1L)),
      vapply(dont_opt_vars, function(v) as.character(v@id), character(1L))
    )
    missing_ids <- setdiff(all_ids, given_ids)
    if (length(missing_ids) > 0L) {
      cli_abort(c(
        paste0("If both {.arg opt_vars} and {.arg dont_opt_vars} are ",
               "specified, they must together cover every variable in {.arg prob}."),
        "x" = "Missing {length(missing_ids)} variable{?s}."
      ))
    }
  }

  ## CVXPY SOURCE: partial_optimize.py:94-100 -- replace each opt_var
  ## with a fresh Variable of the same shape and attributes.  The
  ## tree_copy walker then substitutes whenever the original is
  ## encountered, so the inner problem references only fresh
  ## variables for opt_vars (isolated from outer scope).
  id_to_new_var <- list()
  for (v in opt_vars) {
    fresh <- do.call(Variable, c(list(shape = v@shape), v@attributes))
    id_to_new_var[[as.character(v@id)]] <- fresh
  }
  new_obj_arg <- .partial_optimize_subs(prob@objective@args[[1L]],
                                         id_to_new_var)
  ## Reconstruct objective of same class with substituted argument.
  obj_cls <- S7_class(prob@objective)
  new_obj <- obj_cls(new_obj_arg)
  new_constrs <- lapply(prob@constraints, function(c)
                          .partial_optimize_subs(c, id_to_new_var))
  new_prob <- Problem(new_obj, new_constrs)

  PartialProblem(new_prob, opt_vars = opt_vars,
                 dont_opt_vars = dont_opt_vars,
                 solver = solver,
                 solve_kwargs = list(...))
}

# ==================================================================
# DCP / curvature methods
# ==================================================================
## CVXPY SOURCE: partial_optimize.py:130-183

## is_constant: TRUE when the inner problem has no variables.
method(is_constant, PartialProblem) <- function(x) {
  length(variables(x@inner_problem)) == 0L
}

## is_convex: inner is DCP + Minimize.
method(is_convex, PartialProblem) <- function(x) {
  is_dcp(x@inner_problem) && .s7_is(x@inner_problem@objective, Minimize)
}

## is_concave: inner is DCP + Maximize.
method(is_concave, PartialProblem) <- function(x) {
  is_dcp(x@inner_problem) && .s7_is(x@inner_problem@objective, Maximize)
}

## DNLP curvature: delegate to convex/concave.
## CVXPY SOURCE: partial_optimize.py lines 145-148
method(is_linearizable_convex, PartialProblem) <- function(x) is_convex(x)
method(is_linearizable_concave, PartialProblem) <- function(x) is_concave(x)

## is_dpp: delegate to inner.  CVXPY takes a context arg; CVXR's
## generic is parameterless and uses with_dpp_scope() out-of-band.
method(is_dpp, PartialProblem) <- function(x, context = "dcp") {
  is_dpp(x@inner_problem, context)
}

method(is_log_log_convex, PartialProblem) <- function(x) {
  is_dgp(x@inner_problem) && .s7_is(x@inner_problem@objective, Minimize)
}

method(is_log_log_concave, PartialProblem) <- function(x) {
  is_dgp(x@inner_problem) && .s7_is(x@inner_problem@objective, Maximize)
}

method(is_nonneg, PartialProblem) <- function(x) {
  is_nonneg(x@inner_problem@objective@args[[1L]])
}

method(is_nonpos, PartialProblem) <- function(x) {
  is_nonpos(x@inner_problem@objective@args[[1L]])
}

method(is_imag, PartialProblem)    <- function(x) FALSE
method(is_complex, PartialProblem) <- function(x) FALSE

# ==================================================================
# Identity methods
# ==================================================================

method(expr_name, PartialProblem) <- function(x) {
  sprintf("PartialProblem(%s)", expr_name(x@inner_problem@objective@args[[1L]]))
}

## variables / parameters / constants delegate to the inner Problem
## so the outer pipeline sees the inner's fresh opt_vars + the
## dont_opt_vars (the substituted tree shares the original
## dont_opt_var instances).
method(variables, PartialProblem) <- function(x) {
  variables(x@inner_problem)
}

method(parameters, PartialProblem) <- function(x) {
  parameters(x@inner_problem)
}

method(constants, PartialProblem) <- function(x) {
  constants(x@inner_problem)
}

## domain: inner constraints + inner objective expression's domain.
## CVXPY SOURCE: partial_optimize.py:265-272.
method(domain, PartialProblem) <- function(x) {
  inner <- x@inner_problem
  obj_expr <- inner@objective@args[[1L]]
  c(inner@constraints, domain(obj_expr))
}

# ==================================================================
# value -- numeric evaluation by inner solve
# ==================================================================
## CVXPY SOURCE: partial_optimize.py:274-295.

method(value, PartialProblem) <- function(x, ...) {
  inner <- x@inner_problem
  vars  <- variables(inner)

  ## Snapshot original values so we can restore after solving.
  old_vals <- lapply(vars, function(v) value(v))
  names(old_vals) <- vapply(vars, function(v) as.character(v@id), character(1L))

  fix_vars <- list()
  for (v in x@dont_opt_vars) {
    if (is.null(value(v))) {
      ## Restore (no-op here, but keep symmetric with grad).
      return(NULL)
    }
    fix_vars <- c(fix_vars, list(v == value(v)))
  }

  sub_prob <- Problem(inner@objective,
                      c(fix_vars, inner@constraints))
  result <- do.call(psolve,
                    c(list(problem = sub_prob, solver = x@solver),
                      x@solve_kwargs))

  ## Restore the original values.
  for (v in vars) {
    key <- as.character(v@id)
    value(v) <- old_vals[[key]]
  }
  result
}

# ==================================================================
# dcp_canonicalize -- embed inner objective + constraints into outer
# ==================================================================
## CVXPY SOURCE: partial_optimize.py:297-311 (PartialProblem.canonicalize).
##
## CVXPY's `canonical_form` is the legacy LinOp-based path; CVXR's
## modern Dcp2Cone reduction drives canonicalization via the
## `dcp_canonicalize` S7 generic.  We hook into that generic so the
## outer chain can resolve PartialProblem like any other expression.
##
## Strategy: re-enter the chain's tree walker (`.dcp2cone_tree`) for
## the inner objective's expression and for each inner constraint;
## collect their replacement expressions / new constraints; return the
## inner objective's canonicalized expression as the replacement and
## the concatenation of all collected constraints + the original
## inner constraints' canonicalized forms.
##
## `quad_obj = FALSE` for the inner re-walk: the inner sub-problem
## takes the DCP cone path even if the outer chose the QP path.
## Slightly suboptimal for embedded least-squares but always correct;
## fixing this would require threading `quad_obj` through the
## `dcp_canonicalize` generic's signature.

method(dcp_canonicalize, PartialProblem) <- function(expr, args,
                                                       solver_context = NULL) {
  inner    <- expr@inner_problem
  obj_expr <- inner@objective@args[[1L]]
  obj_result <- .dcp2cone_tree(quad_obj = FALSE, obj_expr,
                                affine_above = TRUE,
                                solver_context = solver_context)
  collected <- obj_result[[2L]]
  for (con in inner@constraints) {
    con_result <- .dcp2cone_tree(quad_obj = FALSE, con,
                                  affine_above = FALSE,
                                  solver_context = solver_context)
    ## Per .dcp2cone_tree contract, the canonicalized constraint comes
    ## back as element 1; new sub-constraints come back in element 2.
    collected <- c(collected,
                   list(con_result[[1L]]),
                   con_result[[2L]])
  }
  list(obj_result[[1L]], collected)
}

method(has_dcp_canon, PartialProblem) <- function(expr) TRUE

# ==================================================================
# grad -- Lagrangian gradient w.r.t. dont_opt_vars
# ==================================================================
## CVXPY SOURCE: partial_optimize.py:211-263.
##
## subgrad of g(y) = min_x f_0(x, y) s.t. f_i(x, y) <= 0, h_j(x, y) == 0
## is given by  D f_0(x*, y) + sum_i lambda*_i D f_i(x*, y)
##                            + sum_j  nu*_j  D h_j(x*, y)
## evaluated at the optimal (primal, dual) point.  In code: pin every
## dont_opt_var to its current value via equality constraints, solve,
## then assemble the Lagrangian as a CVXR expression and call .grad().

## Generic signature is .grad(x, values, ...) -- match formals exactly
## so S7's method-formals check passes.  PartialProblem doesn't use the
## values argument (CVXPY's PartialProblem._grad ignores arg values too,
## reading variable values directly via inner_problem.dont_opt_vars).
method(.grad, PartialProblem) <- function(x, values, ...) {
  ## Short-circuit for constants.
  if (is_constant(x)) {
    return(.constant_grad(x))
  }

  inner <- x@inner_problem
  vars  <- variables(inner)
  old_vals <- lapply(vars, function(v) value(v))
  names(old_vals) <- vapply(vars, function(v) as.character(v@id), character(1L))

  fix_vars <- list()
  for (v in x@dont_opt_vars) {
    if (is.null(value(v))) {
      return(.error_grad(x))
    }
    fix_vars <- c(fix_vars, list(v == value(v)))
  }

  sub_prob <- Problem(inner@objective, c(fix_vars, inner@constraints))
  do.call(psolve,
          c(list(problem = sub_prob, solver = x@solver), x@solve_kwargs))

  if (!(status(sub_prob) %in% c("optimal", "optimal_inaccurate"))) {
    for (v in vars) {
      value(v) <- old_vals[[as.character(v@id)]]
    }
    return(.error_grad(x))
  }

  ## Build the Lagrangian: f_0 + sign * sum_i (lam_i . constr_i),
  ## where sign = +1 for convex (Minimize) and -1 for concave (Maximize).
  ## CVXPY SOURCE: partial_optimize.py:243-256.
  sign <- as.integer(is_convex(x)) - as.integer(is_concave(x))
  lagr <- inner@objective@args[[1L]]
  for (con in inner@constraints) {
    dv <- dual_value(con)
    if (is.null(dv)) next
    lam <- as_expr(sign * dv)
    prod <- t(lam) %*% con@args[[1L]]
    if (expr_size(prod) == 1L) {
      lagr <- lagr + sum_entries(prod)
    } else {
      lagr <- lagr + matrix_trace(prod)
    }
  }

  grad_map <- .grad(lagr)
  result <- list()
  for (v in x@dont_opt_vars) {
    key <- as.character(v@id)
    result[[key]] <- grad_map[[key]]
  }

  for (v in vars) {
    value(v) <- old_vals[[as.character(v@id)]]
  }
  result
}