R/245_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. An argument the converter reads by value() is not
  ## converted (its slot is NULL): converting a constant densifies it, O(m n)
  ## even for a sparse matrix, only for the converter to discard the node.
  ## DELIBERATE R ADAPTATION: CVXPY's convert_expr converts every argument
  ## (converters.py:112-119); results are unchanged.
  by_value <- .de_value_read_args(expr)
  children <- lapply(seq_along(expr@args), function(i) {
    if (i %in% by_value) NULL else .de_convert_expr(expr@args[[i]], ctx)
  })
  diff_engine_convert(expr, children, ctx)
}

## Indices of the arguments that the converter dispatched for expr reads by
## value() rather than as engine nodes: the non-parametric constant factor of
## a matrix product (.de_convert_matmul; Multiply has its own converter) and a
## non-parametric constant P of quad_form (.de_convert_quad_form).
.de_value_read_args <- function(expr) {
  const_np <- function(a) is_constant(a) && length(parameters(a)) == 0L
  if (.s7_is(expr, QuadForm)) {
    return(if (const_np(expr@args[[2L]])) 2L else integer(0))
  }
  if (.s7_is(expr, MulExpression) && !.s7_is(expr, Multiply)) {
    la <- expr@args[[1L]]
    if (is_constant(la)) return(if (const_np(la)) 1L else integer(0))
    if (const_np(expr@args[[2L]])) return(2L)
  }
  integer(0)
}

## 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.
## A PARAMETRIC P (Parameter, diag(p), ...) goes to the dense quad form fed the
## converted P node children[[2L]] (empty data: the engine re-evaluates the node
## each solve), mirroring CVXPY's make_quad_form(children[1], children[0],
## "dense", None, n). That path is dense-only: the Hessian block is a full n x n
## whatever P's structure (quad_form(x, diag(p)), n = 40: 820 nnz vs 40 for the
## same P as a constant). A constant P given as a sparse Matrix uses CSR; a dense
## one uses CSR below SPARSE_DENSITY_THRESHOLD and the dense binding at or above it
## (CVXPY v1.9.3 #3366, registry.py:136-144; ADR D_193.5).
.de_convert_quad_form <- function(expr, children, ctx) {
  P <- expr@args[[2L]]
  if (length(parameters(P)) > 0L) {
    ## CVXPY v1.9.3 #3365: parametric P via the dense quad form (registry.py:121-122).
    return(sparsediff::sd_quad_form_dense(children[[2L]], children[[1L]], numeric(0)))
  }
  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.")
  }
  if (!.de_is_sparse(Pv)) {
    Pv <- .de_to_dense(Pv)
    if (.de_density(Pv) >= SPARSE_DENSITY_THRESHOLD) {
      return(sparsediff::sd_quad_form_dense(NULL, children[[1L]], as.vector(t(Pv))))
    }
  }
  ## 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 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).
## A constant factor given as a sparse Matrix (scipy sparse upstream) uses CSR.
## A dense constant on the left uses CSR below SPARSE_DENSITY_THRESHOLD and the
## dense binding at or above it; on the right it always uses the dense binding
## (CVXPY v1.9.3 #3366, converters.py:47-65; ADR D_193.5). Dense data is
## row-major, as the engine requires. 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)))
    }
    A <- value(la)
    if (!.de_is_sparse(A)) {
      A <- .de_to_dense(A)
      if (.de_density(A) >= SPARSE_DENSITY_THRESHOLD) {
        return(sparsediff::sd_left_matmul_dense(NULL, children[[2L]],
                                                nrow(A), ncol(A), as.vector(t(A))))
      }
    }
    ## 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(A, "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)))
    }
    A <- value(ra)
    if (!.de_is_sparse(A)) {
      A <- .de_to_dense(A)
      return(sparsediff::sd_right_matmul_dense(NULL, children[[1L]],
                                               nrow(A), ncol(A), as.vector(t(A))))
    }
    Acsr <- as(as(A, "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)
}