R/159_reductions_dcp2cone_canonicalizers_sum_largest_canon.R

In CVXR: Disciplined Convex Optimization

#####
## DO NOT EDIT THIS FILE!! EDIT THE SOURCE INSTEAD: rsrc_tree/reductions/dcp2cone/canonicalizers/sum_largest_canon.R
#####

## CVXPY SOURCE: reductions/eliminate_pwl/canonicalizers/sum_largest_canon.py
## At 1.9.0: axis-aware canon + DNLP value-initialization (PR #3172). The 2D axis
## slice (NULL/1/2 + keepdims) and the axis-None DNLP init are fully ported.
## N/A for CVXR (2D-only): N-D tuple axis.
## sum_largest(x, k, axis) = min sum(t, axis) + k*q  s.t. x <= t + q, t >= 0


sum_largest_canon <- function(expr, args, solver_context = NULL) {
  x <- args[[1L]]
  k <- expr@k
  axis <- expr@axis

  ## min sum(t, axis) + k*q  s.t.  x <= t + q,  0 <= t
  t <- Variable(shape = .shape(x))
  q <- Variable(shape = .shape(expr))

  ## Promote q (shape = reduced expr shape) back to x's shape. CVXPY relies on
  ## numpy broadcasting; CVXR has none, so promote explicitly (cf. max_canon).
  if (is.null(axis)) {
    promoted_q <- cvxr_promote(q, .shape(x))
  } else if (axis == 2L) {
    ## axis=2: reduce rows -> q is (1, ncol); broadcast down the rows.
    ones_col <- Constant(matrix(1, nrow = .shape(x)[1L], ncol = 1L))
    q_row <- reshape_expr(q, c(1L, .shape(x)[2L]))
    promoted_q <- ones_col %*% q_row
  } else {
    ## axis=1: reduce cols -> q is (nrow, 1); broadcast across the cols.
    q_col <- reshape_expr(q, c(.shape(x)[1L], 1L))
    ones_row <- Constant(matrix(1, nrow = 1L, ncol = .shape(x)[2L]))
    promoted_q <- q_col %*% ones_row
  }

  obj <- SumEntries(t, axis = axis, keepdims = expr@keepdims) + k * q
  constraints <- list(x <= t + promoted_q, t >= 0)

  ## CVXPY sum_largest_canon: for DNLP we initialize the auxiliary t, q so every
  ## variable has a value (DNLP guarantees x.value is set). CVXPY guards this on
  ## `axis is None`; mirror that. q = max of the non-top-k entries (x_{[k+1]});
  ## t = the positive part above q on the top-k entries.
  if (is.null(axis)) {
    xval <- value(x)
    if (!is.null(xval)) {
      xv <- as.numeric(xval)
      n <- length(xv)
      k_floor <- as.integer(floor(k))
      ord <- order(xv)
      idx_largest <- if (k_floor > 0L) ord[(n - k_floor + 1L):n] else integer(0)
      qv <- if (n > k_floor) max(xv[ord[seq_len(n - k_floor)]]) else min(xv)
      tv <- numeric(n)
      tv[idx_largest] <- xv[idx_largest] - qv
      value(q) <- matrix(qv, 1L, 1L)
      value(t) <- matrix(tv, nrow = .shape(x)[1L], ncol = .shape(x)[2L])
    }
  }

  list(obj, constraints)
}

method(dcp_canonicalize, SumLargest) <- sum_largest_canon
method(has_dcp_canon, SumLargest) <- function(expr) TRUE