R/094_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 a vector
## GeoMeanApprox -- SOC-based rational approximation of GeoMean
##
## The factory function geo_mean() dispatches to GeoMeanApprox (approx=TRUE)
## or GeoMean (approx=FALSE).


GeoMean <- new_class("GeoMean", parent = Atom, package = "CVXR",
  properties = list(
    p            = new_property(class = class_any),       # original weight vector (after filtering)
    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)
  ),
  constructor = function(x, p = NULL, max_denom = 1024L, id = NULL) {
    if (is.null(id)) id <- next_expr_id()
    x <- as_expr(x)

    ## Filter zero-weight elements (CVXPY geo_mean.py lines 211-220)
    if (!is.null(p)) {
      if (S7_inherits(p, Expression)) {
        cli_abort("The {.arg p} (weights) argument should not be an {.cls Expression}.")
      }
      p <- as.numeric(p)
      idxs <- which(p > 0)
      if (length(idxs) < length(p)) {
        x <- x[idxs]
        p <- p[idxs]
      }
    }

    ## x must be a vector (CVXPY geo_mean.py lines 224-227)
    if (!expr_is_vector(x)) {
      cli_abort("{.arg x} must be a row or column vector.")
    }
    n <- if (expr_is_scalar(x)) 1L else max(x@shape)

    ## Default weights: uniform
    if (is.null(p)) p <- rep(1, n)

    if (length(p) != n) {
      cli_abort("{.arg x} and {.arg p} must have the same number of elements.")
    }
    if (any(p < 0) || sum(p) <= 0) {
      cli_abort("Powers must be nonnegative and not all zero.")
    }

    ## Normalized weights (float, CVXPY geo_mean.py line 241)
    p_arr <- as.numeric(p)
    w <- p_arr / sum(p_arr)

    ## Always scalar output (CVXPY: shape = tuple())
    shape <- c(1L, 1L)

    obj <- new_object(S7_object(),
      id           = as.integer(id),
      .cache       = new.env(parent = emptyenv()),
      args         = list(x),
      shape        = shape,
      p            = p,
      w            = w,
      max_denom    = as.integer(max_denom),
      approx_error = 0.0
    )
    validate_arguments(obj)
    obj
  }
)

# -- 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 302-312
method(is_atom_convex, GeoMean) <- function(x) length(x@w) == 1L
method(is_atom_concave, 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: x[w > 0] >= 0 -----------------------------------------
method(atom_domain, GeoMean) <- function(x) {
  selection <- which(x@w > 0)
  list(x@args[[1L]][selection] >= 0)
}

# -- get_data -------------------------------------------------------
method(get_data, GeoMean) <- function(x) {
  list(x@p, x@max_denom)
}

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

# -- numeric: prod(x^w) --------------------------------------------
method(numeric_value, GeoMean) <- function(x, values, ...) {
  v <- as.numeric(values[[1L]])
  w <- as.numeric(x@w)
  val <- prod(v^w)
  matrix(val, 1L, 1L)
}

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

# ===================================================================
# GeoMeanApprox -- SOC-based rational approximation of GeoMean
# ===================================================================
## CVXPY SOURCE: atoms/geo_mean.py lines 363-436
## 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, id = NULL) {
    if (is.null(id)) id <- next_expr_id()
    x <- as_expr(x)

    ## Filter zero-weight elements (same as GeoMean)
    if (!is.null(p)) {
      if (S7_inherits(p, Expression)) {
        cli_abort("The {.arg p} (weights) argument should not be an {.cls Expression}.")
      }
      p <- as.numeric(p)
      idxs <- which(p > 0)
      if (length(idxs) < length(p)) {
        x <- x[idxs]
        p <- p[idxs]
      }
    }

    ## x must be a vector
    if (!expr_is_vector(x)) {
      cli_abort("{.arg x} must be a row or column vector.")
    }
    n <- if (expr_is_scalar(x)) 1L else max(x@shape)

    ## Default weights: uniform
    if (is.null(p)) p <- rep(1, n)

    if (length(p) != n) {
      cli_abort("{.arg x} and {.arg p} must have the same number of elements.")
    }
    if (any(p < 0) || sum(p) <= 0) {
      cli_abort("Powers must be nonnegative and not all zero.")
    }

    ## Rational approximation via fracify (CVXPY geo_mean.py line 401)
    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 404)
    tree_result <- decompose(w_dyad)

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

    ## Always scalar output
    shape <- c(1L, 1L)

    obj <- new_object(S7_object(),
      id            = as.integer(id),
      .cache        = new.env(parent = emptyenv()),
      args          = list(x),
      shape         = shape,
      p             = p,
      w             = w,
      max_denom     = as.integer(max_denom),
      approx_error  = ae,
      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 of a vector
#'
#' @param x An Expression (vector)
#' @param p Numeric weight vector (default: uniform)
#' @param max_denom Maximum denominator for rational approximation
#' @param approx If TRUE (default), use SOC approximation. If FALSE, use exact power cone.
#' @returns A GeoMean or GeoMeanApprox atom
#' @export
geo_mean <- function(x, p = NULL, max_denom = 1024L, approx = TRUE) {
  if (approx) {
    GeoMeanApprox(x, p = p, max_denom = max_denom)
  } else {
    GeoMean(x, p = p, max_denom = max_denom)
  }
}