R/103_atoms_geo_mean.R

In CVXR: Disciplined Convex Optimization

#####
## DO NOT EDIT THIS FILE!! EDIT THE SOURCE INSTEAD: rsrc_tree/atoms/geo_mean.R
#####

## CVXPY SOURCE: atoms/geo_mean.py
## GeoMean -- (weighted) geometric mean of an expression, optionally along an axis
## GeoMeanApprox -- SOC-based rational approximation of GeoMean
##
## The factory function geo_mean() dispatches to GeoMeanApprox (approx=TRUE)
## or GeoMean (approx=FALSE).
##
## CVXPY 1.9.3 (#3493): GeoMean is an AxisAtom. Zero-weight entries stay in the
## argument and are dropped through the `.keep` mask in the aligned view
## (.gm_aligned_arg), so the argument keeps the shape the axis bookkeeping is
## based on. CVXR is 2-D: axis is NULL, 1L or 2L (no axis tuples).


## Shared constructor logic for GeoMean and GeoMeanApprox.
## CVXPY SOURCE: atoms/geo_mean.py GeoMean.__init__ (lines 219-257)
.gm_init <- function(x, p, axis, keepdims) {
  if (!is.null(p) && .s7_is(p, Expression)) {
    cli_abort("The {.arg p} (weights) argument should not be an {.cls Expression}.")
  }
  x <- as_expr(x)
  if (!is.null(axis)) axis <- as.integer(axis)
  keepdims <- as.logical(keepdims)
  ## AxisAtom.__init__: validates the axis and gives the reduced shape.
  shape <- .axis_shape(.shape(x), axis, keepdims)

  ## _reduced_size(): how many entries each geometric mean is taken over.
  x_shape <- .shape(x)
  n <- if (is.null(axis)) as.integer(prod(x_shape)) else x_shape[3L - axis]

  if (!is.null(p)) {
    p <- as.numeric(p)
    if (length(p) != n) {
      cli_abort("{.arg x} and {.arg p} must have the same number of elements.")
    }
    ## Entries with zero weight take no part. They are dropped from the
    ## canonicalization rather than from the argument.
    keep <- p > 0
    p <- p[keep]
  } else {
    keep <- rep(TRUE, n)
    p <- rep(1, n)
  }
  if (any(p < 0) || sum(p) <= 0) {
    cli_abort("Powers must be nonnegative and not all zero.")
  }
  list(x = x, p = p, keep = keep, axis = axis, keepdims = keepdims, shape = shape)
}

GeoMean <- new_class("GeoMean", parent = AxisAtom, package = "CVXR",
  properties = list(
    p            = new_property(class = class_any),       # weights of the kept entries
    w            = new_property(class = class_any),       # normalized weights (numeric or bigq)
    max_denom    = new_property(class = class_integer),
    approx_error = new_property(class = class_numeric),
    .keep        = new_property(class = class_logical)    # entries with positive weight
  ),
  constructor = function(x, p = NULL, max_denom = 1024L, axis = NULL,
                         keepdims = FALSE, id = NULL) {
    if (FALSE) new_object(S7_object())  ## S7 static-check guard
    if (is.null(id)) id <- next_expr_id()
    init <- .gm_init(x, p, axis, keepdims)

    ## Normalized weights (float, CVXPY geo_mean.py line 256)
    w <- init$p / sum(init$p)

    obj <- .fast_new(GeoMean, S7_object(),
      id           = as.integer(id),
      .cache       = new.env(parent = emptyenv()),
      args         = list(init$x),
      shape        = init$shape,
      axis         = init$axis,
      keepdims     = init$keepdims,
      p            = init$p,
      w            = w,
      max_denom    = as.integer(max_denom),
      approx_error = 0.0,
      .keep        = init$keep
    )
    validate_arguments(obj)
    obj
  }
)

# -- aligned view ----------------------------------------------------
## CVXPY SOURCE: atoms/geo_mean.py GeoMean._aligned_arg (lines 267-285)
## The argument with the reduced entries along the rows (one column per
## output), zero-weight rows dropped. The domain and all canonicalizers work
## on this view. `arg` is a canonicalizer's own (same-shaped) argument.
## 2-D only, and built from reshape_expr / t() alone: no moveaxis analog,
## whose C++ lowering CVXPY 1.9.3 gets wrong (plan notes/cvxpy_1.9.3_plan.md
## §4 U1).
.gm_aligned_arg <- function(expr, arg = NULL) {
  x <- if (is.null(arg)) .args(expr)[[1L]] else arg
  axis <- expr@axis
  if (is.null(axis)) {
    shp <- .shape(x)
    aligned <- if (shp[2L] == 1L) x else reshape_expr(x, c(as.integer(prod(shp)), 1L))
  } else if (axis == 2L) {
    aligned <- x
  } else {
    aligned <- t(x)
  }
  keep <- expr@.keep
  if (!all(keep)) aligned <- aligned[which(keep), , drop = FALSE]
  aligned
}

## The same view of a numeric value (for numeric_value).
.gm_aligned_value <- function(expr, v) {
  shp <- .arg_shape(expr)
  M <- matrix(as.vector(as.matrix(v)), shp[1L], shp[2L])
  axis <- expr@axis
  A <- if (is.null(axis)) matrix(as.vector(M), ncol = 1L) else if (axis == 2L) M else t(M)
  A[expr@.keep, , drop = FALSE]
}

## Map the k x 1 or 1 x k result of a canonicalizer onto the atom's shape.
.gm_to_shape <- function(out, shape) {
  if (identical(.shape(out), as.integer(shape))) out else reshape_expr(out, shape)
}

# -- sign: always nonneg --------------------------------------------
method(sign_from_args, GeoMean) <- function(x) {
  list(is_nonneg = TRUE, is_nonpos = FALSE)
}

# -- curvature: concave (affine when single weight) -----------------
## CVXPY geo_mean.py lines 352-362
method(is_atom_convex, GeoMean) <- function(x) length(x@w) == 1L
method(is_atom_concave, GeoMean) <- function(x) TRUE

## CVXPY geo_mean.py: geo_mean is smooth.
method(is_atom_smooth, GeoMean) <- function(x) TRUE

# -- log-log --------------------------------------------------------
method(is_atom_log_log_convex, GeoMean) <- function(x) TRUE
method(is_atom_log_log_concave, GeoMean) <- function(x) TRUE

# -- monotonicity: always increasing --------------------------------
method(is_incr, GeoMean) <- function(x, idx, ...) TRUE
method(is_decr, GeoMean) <- function(x, idx, ...) FALSE

# -- domain: aligned >= 0 -------------------------------------------
## CVXPY SOURCE: atoms/geo_mean.py GeoMean._domain (lines 303-308)
## Zero-weight entries are already gone from the aligned view, so they carry
## no implicit nonnegativity.
method(atom_domain, GeoMean) <- function(x) {
  list(.gm_aligned_arg(x) >= 0)
}

# -- get_data -------------------------------------------------------
## CVXPY SOURCE: atoms/geo_mean.py GeoMean.get_data (line 388)
## CVXPY returns the kept weights and copies `_keep` by hand in copy(). R
## rebuilds an atom through its constructor from get_data() (expr_copy), so
## the weights go back at full length, zeros included, and the constructor
## re-derives the mask. Recorded: ADR D_PORT.14 (claims audit C0203).
method(get_data, GeoMean) <- function(x) {
  p_full <- numeric(length(x@.keep))
  p_full[x@.keep] <- x@p
  list(p_full, x@max_denom, x@axis, x@keepdims)
}

# -- name -----------------------------------------------------------
method(expr_name, GeoMean) <- function(x) {
  weights <- paste(x@w, collapse = ", ")
  sprintf("%s(%s, (%s))", class(x)[[1L]], expr_name(.args(x)[[1L]]), weights)
}

## CVXPY SOURCE: atoms/geo_mean.py GeoMean.format_labeled (lines 339-343).
method(format_labeled, GeoMean) <- function(x) {
  lbl <- label(x); if (!is.null(lbl)) return(lbl)
  weights <- paste(x@w, collapse = ", ")
  sprintf("%s(%s, (%s))", class(x)[[1L]], format_labeled(.args(x)[[1L]]), weights)
}

# -- numeric: prod(x^w) per slice -----------------------------------
## CVXPY SOURCE: atoms/geo_mean.py GeoMean.numeric (lines 287-301)
method(numeric_value, GeoMean) <- function(x, values, ...) {
  A <- .gm_aligned_value(x, values[[1L]])
  w <- as.numeric(x@w)
  out <- apply(A^w, 2L, prod)
  shp <- .shape(x)
  matrix(out, shp[1L], shp[2L])
}

# -- graph_implementation: stub -------------------------------------
method(graph_implementation, GeoMean) <- function(x, arg_objs, shape, data = NULL, ...) {
  cli_abort("graph_implementation for {.cls GeoMean} not yet implemented.")
}

# -- .column_grad: per-fiber subgradient ------------------------------
## CVXPY SOURCE: atoms/geo_mean.py GeoMean._column_grad (lines 310-320);
## _grad (lines 322-333) is AxisAtom._axis_grad, which walks the fibers
## (method(.grad, AxisAtom)).
## D[i] = w[i] / x[i] * geo_mean(x; w) on kept entries, 0 elsewhere.
## Domain: NULL if any kept entry is <= 0.
## w may be bigq; coerce via as.numeric (CVXPY's float(w_i)).
method(.column_grad, GeoMean) <- function(x, value, ...) {
  column <- as.numeric(value)
  kept <- column[x@.keep]
  w <- as.numeric(x@w)
  if (any(kept <= 0)) return(NULL)
  gm_val <- prod(kept^w)
  D <- numeric(length(column))
  D[x@.keep] <- w / kept * gm_val
  D
}

# ===================================================================
# GeoMeanApprox -- SOC-based rational approximation of GeoMean
# ===================================================================
## CVXPY SOURCE: atoms/geo_mean.py lines 418-481
## Subclass of GeoMean. Overrides w with rational approximation
## and adds w_dyad (dyadic completion), tree (decomposition),
## cone_lb, cone_num_over, cone_num.

GeoMeanApprox <- new_class("GeoMeanApprox", parent = GeoMean, package = "CVXR",
  properties = list(
    w_dyad        = new_property(class = class_any),       # bigq dyadic completion
    tree          = new_property(class = class_any),       # decomposition tree
    cone_lb       = new_property(class = class_integer),
    cone_num_over = new_property(class = class_integer),
    cone_num      = new_property(class = class_integer)
  ),
  constructor = function(x, p = NULL, max_denom = 1024L, axis = NULL,
                         keepdims = FALSE, id = NULL) {
    if (FALSE) new_object(S7_object())  ## S7 static-check guard
    if (is.null(id)) id <- next_expr_id()
    init <- .gm_init(x, p, axis, keepdims)
    p <- init$p

    ## Rational approximation via fracify (CVXPY geo_mean.py line 456)
    frac_result <- fracify(p, max_denom)
    w <- frac_result$w
    w_dyad <- frac_result$w_dyad
    ae <- approx_error(p, w)

    ## Decompose for SOC construction (CVXPY geo_mean.py line 459)
    tree_result <- decompose(w_dyad)

    ## Cone bounds (CVXPY geo_mean.py lines 462-468)
    clb <- lower_bound(w_dyad)
    cno <- over_bound(w_dyad, tree_result)
    cn <- clb + cno

    obj <- .fast_new(GeoMeanApprox, S7_object(),
      id            = as.integer(id),
      .cache        = new.env(parent = emptyenv()),
      args          = list(init$x),
      shape         = init$shape,
      axis          = init$axis,
      keepdims      = init$keepdims,
      p             = p,
      w             = w,
      max_denom     = as.integer(max_denom),
      approx_error  = ae,
      .keep         = init$keep,
      w_dyad        = w_dyad,
      tree          = tree_result,
      cone_lb       = as.integer(clb),
      cone_num_over = as.integer(cno),
      cone_num      = as.integer(cn)
    )
    validate_arguments(obj)
    obj
  }
)

# -- Factory function ---------------------------------------------
#' (Weighted) geometric mean
#'
#' The (weighted) geometric mean of the entries of `x`, of all of them or of
#' each row or column.
#'
#' @param x An Expression of any shape. With `axis = NULL` its entries are
#'   taken in column-major order, so a row vector, a column vector and a
#'   matrix holding the same entries give the same result.
#' @param p Numeric weight vector (default: uniform), one weight per entry
#'   being reduced: `length(x)` entries with `axis = NULL`, `ncol(x)` with
#'   `axis = 1`, `nrow(x)` with `axis = 2`. Entries with zero weight take no
#'   part.
#' @param max_denom Maximum denominator for rational approximation
#' @param approx If TRUE (default), use SOC approximation. If FALSE, use exact power cone.
#' @param axis `NULL` (default) for one geometric mean of all entries, `1` for
#'   one per row, `2` for one per column.
#' @param keepdims Keep the reduced dimension with length one.
#' @returns A GeoMean or GeoMeanApprox atom: a scalar when `axis = NULL`,
#'   otherwise one geometric mean per row or column.
#' @export
geo_mean <- function(x, p = NULL, max_denom = 1024L, approx = TRUE,
                     axis = NULL, keepdims = FALSE) {
  if (approx) {
    GeoMeanApprox(x, p = p, max_denom = max_denom, axis = axis, keepdims = keepdims)
  } else {
    GeoMean(x, p = p, max_denom = max_denom, axis = axis, keepdims = keepdims)
  }
}