R/227_utilities_replace_quad_forms.R

In CVXR: Disciplined Convex Optimization

#####
## DO NOT EDIT THIS FILE!! EDIT THE SOURCE INSTEAD: rsrc_tree/utilities/replace_quad_forms.R
#####

## CVXPY SOURCE: utilities/replace_quad_forms.py
## Replace SymbolicQuadForm nodes with dummy Variables for P extraction.
##
## S7 objects are immutable -- cannot mutate expr@args directly.
## Instead, we rebuild the tree bottom-up via expr_copy.


# -- replace_quad_forms ------------------------------------------
## Walk the expression tree. Replace any SymbolicQuadForm (or QuadForm)
## child with a dummy Variable of the same shape and ID.
## Returns list(new_expr, quad_forms) where quad_forms is a named list
## mapping dummy_id_str -> list(parent, idx, quad_form_atom).
##
## Because S7 is immutable, we reconstruct parent nodes bottom-up
## using expr_copy(parent, new_args).

replace_quad_forms <- function(expr, quad_forms = list()) {
  new_args <- list()
  changed <- FALSE

  for (idx in seq_along(expr@args)) {
    arg <- expr@args[[idx]]
    if (S7_inherits(arg, SymbolicQuadForm) || S7_inherits(arg, QuadForm)) {
      ## Replace with dummy Variable
      placeholder <- Variable(arg@shape, var_id = arg@id)
      new_args[[idx]] <- placeholder
      quad_forms[[as.character(placeholder@id)]] <-
        list(parent = expr, idx = idx, quad_form = arg)
      changed <- TRUE
    } else {
      ## Recurse into child
      child_result <- replace_quad_forms(arg, quad_forms)
      new_args[[idx]] <- child_result$expr
      quad_forms <- child_result$quad_forms
      if (!identical(child_result$expr, arg)) changed <- TRUE
    }
  }

  if (changed) {
    new_expr <- expr_copy(expr, new_args)
  } else {
    new_expr <- expr
  }
  list(expr = new_expr, quad_forms = quad_forms)
}

# -- restore_quad_forms ------------------------------------------
## Restore original SymbolicQuadForm nodes in the expression tree.
## Returns the restored expression (again, bottom-up reconstruction).

restore_quad_forms <- function(expr, quad_forms) {
  new_args <- list()
  changed <- FALSE

  for (idx in seq_along(expr@args)) {
    arg <- expr@args[[idx]]
    arg_id <- as.character(arg@id)
    if (S7_inherits(arg, Variable) && !is.null(quad_forms[[arg_id]])) {
      ## Restore the original quad form
      new_args[[idx]] <- quad_forms[[arg_id]]$quad_form
      changed <- TRUE
    } else {
      ## Recurse
      restored <- restore_quad_forms(arg, quad_forms)
      new_args[[idx]] <- restored
      if (!identical(restored, arg)) changed <- TRUE
    }
  }

  if (changed) {
    expr_copy(expr, new_args)
  } else {
    expr
  }
}