R/245_reductions_solvers_nlp_solvers_diff_engine_registry.R

In CVXR: Disciplined Convex Optimization

#####
## DO NOT EDIT THIS FILE!! EDIT THE SOURCE INSTEAD: rsrc_tree/reductions/solvers/nlp_solvers/diff_engine/registry.R
#####

## CVXPY SOURCE: reductions/solvers/nlp_solvers/diff_engine/registry.py
##              (+ diff_engine/__init__.py, pure re-exports, not mirrored)
## Atom -> sparsediff converter registrations. CVXPY keys an ATOM_CONVERTERS
## dict by atom name; CVXR dispatches via the diff_engine_convert S7 generic
## (default: abort "unsupported"). Each method receives (expr, children, ctx)
## where children are already-converted sparsediff handles and ctx holds
## var_dict / n_vars / param_dict.
##
## SCOPE: smooth-atom core + affine + bivariate + matmul (constant, variable,
## and parametrized matrix factors) + structural atoms (reshape/transpose/
## index/special_index/hstack/vstack/diag/trace/upper_tri/broadcast/convolve) +
## pass-through wraps. Still DEFERRED (clear abort): DiagVec/DiagMat with
## k != 0. (Parametrized divisors in div and a parametrized P in quad_form are
## rejected to match CVXPY's diff engine, which rejects them too.)

#' @noRd
diff_engine_convert <- new_generic("diff_engine_convert", "expr",
  function(expr, children, ctx) {
    S7_dispatch()
  }
)

method(diff_engine_convert, S7_object) <- function(expr, children, ctx) {
  cli_abort(c(
    "Atom {.cls {class(expr)[[1L]]}} is not supported by the NLP diff engine.",
    "i" = "matmul and structural atoms are deferred to a Step-5a follow-on."
  ))
}

# -- elementwise smooth atoms ------------------------------------------------
method(diff_engine_convert, Exp)     <- function(expr, children, ctx) sparsediff::sd_exp(children[[1L]])
method(diff_engine_convert, Log)     <- function(expr, children, ctx) sparsediff::sd_log(children[[1L]])
method(diff_engine_convert, Sin)     <- function(expr, children, ctx) sparsediff::sd_sin(children[[1L]])
method(diff_engine_convert, Cos)     <- function(expr, children, ctx) sparsediff::sd_cos(children[[1L]])
method(diff_engine_convert, Tan)     <- function(expr, children, ctx) sparsediff::sd_tan(children[[1L]])
method(diff_engine_convert, Sinh)    <- function(expr, children, ctx) sparsediff::sd_sinh(children[[1L]])
method(diff_engine_convert, Tanh)    <- function(expr, children, ctx) sparsediff::sd_tanh(children[[1L]])
method(diff_engine_convert, Asinh)   <- function(expr, children, ctx) sparsediff::sd_asinh(children[[1L]])
method(diff_engine_convert, Atanh)   <- function(expr, children, ctx) sparsediff::sd_atanh(children[[1L]])
method(diff_engine_convert, Entr)    <- function(expr, children, ctx) sparsediff::sd_entr(children[[1L]])
method(diff_engine_convert, Logistic)<- function(expr, children, ctx) sparsediff::sd_logistic(children[[1L]])
method(diff_engine_convert, Xexp)    <- function(expr, children, ctx) sparsediff::sd_xexp(children[[1L]])
method(diff_engine_convert, Normcdf) <- function(expr, children, ctx) sparsediff::sd_normal_cdf(children[[1L]])

# -- power (p stored as p_used) ----------------------------------------------
method(diff_engine_convert, Power)       <- function(expr, children, ctx) sparsediff::sd_power(children[[1L]], as.numeric(expr@p_used))
method(diff_engine_convert, PowerApprox) <- function(expr, children, ctx) sparsediff::sd_power(children[[1L]], as.numeric(expr@p_used))

# -- affine ------------------------------------------------------------------
method(diff_engine_convert, NegExpression) <- function(expr, children, ctx) sparsediff::sd_neg(children[[1L]])
method(diff_engine_convert, AddExpression) <- function(expr, children, ctx) .de_chain_add(children)
method(diff_engine_convert, Promote)       <- function(expr, children, ctx) {
  shp <- .de_shape(expr@shape)
  sparsediff::sd_promote(children[[1L]], shp[1L], shp[2L])
}
method(diff_engine_convert, SumEntries) <- function(expr, children, ctx) {
  ## CVXR axis is 1-based (1 = reduce columns = CVXPY/engine axis 1; 2 = reduce
  ## rows = engine axis 0). NULL = reduce all = -1.
  axis <- expr@axis
  engine_axis <- if (is.null(axis)) -1L else (2L - as.integer(axis))
  sparsediff::sd_sum(children[[1L]], engine_axis)
}
method(diff_engine_convert, Multiply)      <- .de_convert_multiply
method(diff_engine_convert, DivExpression) <- .de_convert_div

# -- bivariate ---------------------------------------------------------------
method(diff_engine_convert, QuadForm)    <- .de_convert_quad_form
method(diff_engine_convert, QuadOverLin) <- function(expr, children, ctx) sparsediff::sd_quad_over_lin(children[[1L]], children[[2L]])
method(diff_engine_convert, RelEntr)     <- .de_convert_rel_entr

# -- reductions returning a scalar -------------------------------------------
method(diff_engine_convert, Prod) <- function(expr, children, ctx) {
  ## CVXR axis is 1-based (1 = reduce columns = CVXPY/engine axis 1; 2 = reduce
  ## rows = CVXPY/engine axis 0). NULL = product of all entries.
  axis <- expr@axis
  if (is.null(axis)) return(sparsediff::sd_prod(children[[1L]]))
  engine_axis <- 2L - as.integer(axis)
  if (engine_axis == 1L) sparsediff::sd_prod_axis_one(children[[1L]])
  else sparsediff::sd_prod_axis_zero(children[[1L]])
}

# -- matrix multiplication ---------------------------------------------------
method(diff_engine_convert, MulExpression) <- .de_convert_matmul

# -- structural / affine atoms -----------------------------------------------
method(diff_engine_convert, Reshape)   <- .de_convert_reshape
method(diff_engine_convert, Transpose) <- .de_convert_transpose
method(diff_engine_convert, Convolve)  <- .de_convert_convolve
method(diff_engine_convert, Trace)     <- function(expr, children, ctx) sparsediff::sd_trace(children[[1L]])
method(diff_engine_convert, UpperTri)  <- function(expr, children, ctx) sparsediff::sd_upper_tri(children[[1L]])
## NOTE: CVXPY's broadcast_to atom has no CVXR counterpart (broadcast_to.R is
## NOT IMPLEMENTED -- R handles broadcasting differently), so no converter.
method(diff_engine_convert, HStack) <- function(expr, children, ctx) sparsediff::sd_hstack(children, ctx$n_vars)
method(diff_engine_convert, VStack) <- function(expr, children, ctx) sparsediff::sd_vstack(children, ctx$n_vars)
method(diff_engine_convert, DiagVec) <- function(expr, children, ctx) {
  if (expr@k != 0L) cli_abort("{.cls DiagVec} with k != 0 is not supported in the NLP diff engine.")
  sparsediff::sd_diag_vec(children[[1L]])
}
method(diff_engine_convert, DiagMat) <- function(expr, children, ctx) {
  if (expr@k != 0L) cli_abort("{.cls DiagMat} with k != 0 is not supported in the NLP diff engine.")
  sparsediff::sd_diag_mat(children[[1L]])
}

# -- indexing ----------------------------------------------------------------
method(diff_engine_convert, Index) <- function(expr, children, ctx) {
  shp <- .de_shape(expr@shape)
  sparsediff::sd_index(children[[1L]], shp[1L], shp[2L], .de_index_flat(expr))
}
method(diff_engine_convert, SpecialIndex) <- function(expr, children, ctx) {
  shp <- .de_shape(expr@shape)
  sparsediff::sd_index(children[[1L]], shp[1L], shp[2L], .de_special_index_flat(expr))
}

# -- pass-through wraps (no-ops for real-valued expressions) -----------------
.de_passthrough <- function(expr, children, ctx) children[[1L]]
method(diff_engine_convert, Conj_)               <- .de_passthrough
method(diff_engine_convert, nonneg_wrap)         <- .de_passthrough
method(diff_engine_convert, nonpos_wrap)         <- .de_passthrough
method(diff_engine_convert, psd_wrap)            <- .de_passthrough
## CVXPY SOURCE: registry.py:249 ("nsd_wrap"). Registered per CLASS, not on the
## Wrap base, so a new Wrap subclass needs its own line or the diff engine
## aborts on it -- which is what would have happened to an NSD variable in a
## DNLP problem when nsd_wrap was added at 1.9.1.9042.
method(diff_engine_convert, nsd_wrap)            <- .de_passthrough
method(diff_engine_convert, hermitian_wrap)      <- .de_passthrough
method(diff_engine_convert, symmetric_wrap)      <- .de_passthrough
method(diff_engine_convert, skew_symmetric_wrap) <- .de_passthrough