R/244_reductions_solvers_nlp_solvers_diff_engine_converters.R

In CVXR: Disciplined Convex Optimization

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

## CVXPY SOURCE: reductions/solvers/nlp_solvers/diff_engine/converters.py
## Tree-walk converting a CVXR expression to a sparsediff expression handle.
## Per-atom dispatch is the diff_engine_convert S7 generic (registry.R).
##
## The sparsediff R API (sd_*) is NOT a 1:1 mirror of CVXPY's low-level
## _diffengine.make_* calls (e.g. separate sd_left_matmul/sd_left_matmul_dense,
## a different sd_quad_form signature, no exported dimension-check), so these
## converters adapt onto sd_* rather than translating make_* verbatim. CVXR
## shapes are native 2D (r, c), so the CVXPY normalize_shape / dimension-check
## dance is unnecessary here.


## .de_convert_expr: recursive tree-walk (CVXPY convert_expr).
## ctx = list(var_dict, n_vars, param_dict).
.de_convert_expr <- function(expr, ctx) {
  ## Base cases
  if (.s7_is(expr, Variable)) {
    return(get(as.character(expr@id), envir = ctx$var_dict))
  }
  if (.s7_is(expr, Parameter)) {
    return(get(as.character(expr@id), envir = ctx$param_dict))
  }
  if (.s7_is(expr, Constant)) {
    cval <- .de_to_dense(value(expr))
    shp <- .de_shape(expr@shape)
    return(sparsediff::sd_parameter(d1 = shp[1L], d2 = shp[2L], param_id = -1L,
                                    n_vars = ctx$n_vars, values = as.numeric(cval)))
  }
  ## Recursive case
  children <- lapply(expr@args, .de_convert_expr, ctx = ctx)
  diff_engine_convert(expr, children, ctx)
}

## convert_multiply (CVXPY convert_multiply): elementwise mult; a constant
## operand becomes scalar/vector multiplication.
.de_convert_multiply <- function(expr, children, ctx) {
  la <- expr@args[[1L]]
  ra <- expr@args[[2L]]
  if (is_constant(la)) {
    if (expr_size(la) == 1L) sparsediff::sd_scalar_mult(children[[1L]], children[[2L]])
    else                     sparsediff::sd_vector_mult(children[[1L]], children[[2L]])
  } else if (is_constant(ra)) {
    if (expr_size(ra) == 1L) sparsediff::sd_scalar_mult(children[[2L]], children[[1L]])
    else                     sparsediff::sd_vector_mult(children[[2L]], children[[1L]])
  } else {
    sparsediff::sd_elementwise_mult(children[[1L]], children[[2L]])
  }
}

## convert_div (CVXPY convert_div): x / c  ->  x * (1/c). Parametrized or
## zero divisors are rejected.
.de_convert_div <- function(expr, children, ctx) {
  divisor_expr <- expr@args[[2L]]
  if (length(parameters(divisor_expr)) > 0L) {
    cli_abort("{.fn div} does not support a parametrized divisor in the NLP diff engine.")
  }
  divisor <- .de_to_dense(value(divisor_expr))
  if (any(divisor == 0)) {
    cli_abort("Division by zero encountered in divisor.")
  }
  recip <- 1 / divisor
  shp <- .de_shape(dim(recip))
  recip_node <- sparsediff::sd_parameter(d1 = shp[1L], d2 = shp[2L], param_id = -1L,
                                         n_vars = 0L, values = as.numeric(recip))
  if (length(recip) == 1L) sparsediff::sd_scalar_mult(recip_node, children[[1L]])
  else                     sparsediff::sd_vector_mult(recip_node, children[[1L]])
}

## convert_rel_entr (CVXPY convert_rel_entr): the low-level engine auto-
## dispatches on argument dimensions; the R API splits this into three
## entry points, so we pick based on argument sizes.
.de_convert_rel_entr <- function(expr, children, ctx) {
  xs <- expr_size(expr@args[[1L]])
  ys <- expr_size(expr@args[[2L]])
  if (xs > 1L && ys > 1L && xs != ys) {
    cli_abort("rel_entr requires compatible argument sizes (x = {xs}, y = {ys}).")
  }
  if (xs == 1L && ys > 1L) {
    sparsediff::sd_rel_entr_first_scalar(children[[1L]], children[[2L]])
  } else if (ys == 1L && xs > 1L) {
    sparsediff::sd_rel_entr_second_scalar(children[[1L]], children[[2L]])
  } else {
    sparsediff::sd_rel_entr(children[[1L]], children[[2L]])
  }
}

## convert_quad_form (CVXPY convert_quad_form): scalar x^T P x with constant P.
.de_convert_quad_form <- function(expr, children, ctx) {
  P <- expr@args[[2L]]
  if (length(parameters(P)) > 0L) {
    cli_abort("quad_form with a parameterized P is not supported by the NLP diff engine.")
  }
  if (!is_constant(P)) {
    cli_abort("quad_form requires P to be a constant matrix.")
  }
  Pv <- value(P)
  if (is.null(Pv)) {
    cli_abort("quad_form with a symbolic P (no value) is not supported by the NLP diff engine.")
  }
  ## Engine expects CSR (row pointers / col indices), matching scipy.csr_matrix.
  Pcsr <- as(as(Pv, "RsparseMatrix"), "generalMatrix")
  sparsediff::sd_quad_form(children[[1L]], Pcsr@p, Pcsr@j, Pcsr@x)
}

## convert_matmul (CVXPY convert_matmul): A @ f(x), f(x) @ A, or X @ Y.
## A non-parametric constant matrix uses the sparse CSR matmul path. A
## PARAMETRIZED matrix factor uses the dense matmul fed the parameter's
## diff-engine node (with empty data: the node supplies and updates the matrix
## entries), mirroring CVXPY's make_dense_left_matmul(param_node, child, A).
## X @ Y (both variable) uses sd_matmul.
.de_convert_matmul <- function(expr, children, ctx) {
  la <- expr@args[[1L]]
  ra <- expr@args[[2L]]
  if (is_constant(la)) {
    if (length(parameters(la)) > 0L) {
      ## children[[1L]] is the converted parametric factor node.
      shp <- .de_shape(la@shape)
      return(sparsediff::sd_left_matmul_dense(children[[1L]], children[[2L]],
                                              shp[1L], shp[2L], numeric(0)))
    }
    ## Engine expects CSR (row pointers / col indices), matching scipy.csr_matrix
    ## in CVXPY. R's RsparseMatrix (dgRMatrix) gives @p (row ptr), @j (col idx).
    Acsr <- as(as(value(la), "RsparseMatrix"), "generalMatrix")
    sparsediff::sd_left_matmul(children[[2L]], Acsr@p, Acsr@j, Acsr@x, ncol(Acsr))
  } else if (is_constant(ra)) {
    if (length(parameters(ra)) > 0L) {
      shp <- .de_shape(ra@shape)
      return(sparsediff::sd_right_matmul_dense(children[[2L]], children[[1L]],
                                               shp[1L], shp[2L], numeric(0)))
    }
    Acsr <- as(as(value(ra), "RsparseMatrix"), "generalMatrix")
    sparsediff::sd_right_matmul(children[[1L]], Acsr@p, Acsr@j, Acsr@x, ncol(Acsr))
  } else {
    sparsediff::sd_matmul(children[[1L]], children[[2L]])
  }
}

## convert_reshape (CVXPY convert_reshape): F-order is direct; C-order is
## transpose(reshape(transpose(x), d2, d1)).
.de_convert_reshape <- function(expr, children, ctx) {
  shp <- .de_shape(expr@shape)
  if (expr@order == "F") {
    sparsediff::sd_reshape(children[[1L]], shp[1L], shp[2L])
  } else {
    t1 <- sparsediff::sd_transpose(children[[1L]])
    r1 <- sparsediff::sd_reshape(t1, shp[2L], shp[1L])
    sparsediff::sd_transpose(r1)
  }
}

## convert_transpose (CVXPY convert_transpose): a vector transpose is a reshape;
## a matrix transpose uses sd_transpose. (CVXR has no true 1D, so the CVXPY
## 1D no-op branch does not apply.)
.de_convert_transpose <- function(expr, children, ctx) {
  cs <- .de_shape(expr@args[[1L]]@shape)
  if (cs[1L] == 1L || cs[2L] == 1L) {
    sparsediff::sd_reshape(children[[1L]], cs[2L], cs[1L])
  } else {
    sparsediff::sd_transpose(children[[1L]])
  }
}

## convert_convolve: discrete convolution is commutative, so pass the constant
## (parameter) operand first as sd_convolve expects (param, child).
.de_convert_convolve <- function(expr, children, ctx) {
  if (is_constant(expr@args[[1L]])) {
    sparsediff::sd_convolve(children[[1L]], children[[2L]])
  } else {
    sparsediff::sd_convolve(children[[2L]], children[[1L]])
  }
}

## Flatten an Index key to 0-based column-major (F-order) indices.
## CVXR Index@key = list(row_idx, col_idx), both 1-based integer vectors.
.de_index_flat <- function(expr) {
  parent_nrow <- as.integer(expr@args[[1L]]@shape[1L])
  r0 <- as.integer(expr@key[[1L]] - 1L)
  c0 <- as.integer(expr@key[[2L]] - 1L)
  ## M[i, j] = r0[i] + c0[j]*nrow; column-major flatten == F-order
  as.integer(outer(r0, c0 * parent_nrow, "+"))
}

## SpecialIndex@select_vec holds 1-based column-major linear indices.
.de_special_index_flat <- function(expr) {
  as.integer(expr@select_vec - 1L)
}