R/079_atoms_total_variation.R

In CVXR: Disciplined Convex Optimization

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

## CVXPY SOURCE: atoms/total_variation.py
## Total variation of a vector or matrix

#' Total variation of a vector or matrix
#'
#' @description
#' Computes total variation using L1 norm of discrete gradients for vectors
#' and L2 norm of discrete gradients for matrices.
#'
#' @param value An Expression or numeric constant (vector or matrix)
#' @param ... Additional matrix expressions extending the third dimension
#' @returns An Expression representing the total variation
#' @export
total_variation <- function(value, ...) {
  value <- as_expr(value)
  s <- value@shape
  n <- s[1L]
  m <- s[2L]

  if (n == 1L && m == 1L) {
    cli_abort("total_variation cannot take a scalar argument.")
  }

  ## 1D case: vector (n,1) or (1,n) -> L1 norm of differences
  if (m == 1L) {
    ## Column vector (n, 1)
    return(norm1(value[2L:n, ] - value[1L:(n - 1L), ]))
  }
  if (n == 1L) {
    ## Row vector (1, m)
    return(norm1(value[, 2L:m] - value[, 1L:(m - 1L)]))
  }

  ## 2D case: matrix -> L2 norms of discrete gradients
  extra_args <- list(...)
  extra_args <- lapply(extra_args, as_expr)
  all_values <- c(list(value), extra_args)

  diffs <- list()
  for (mat in all_values) {
    ## Horizontal differences: mat[1:(n-1), 2:m] - mat[1:(n-1), 1:(m-1)]
    diffs <- c(diffs, list(mat[1L:(n - 1L), 2L:m] - mat[1L:(n - 1L), 1L:(m - 1L)]))
    ## Vertical differences: mat[2:n, 1:(m-1)] - mat[1:(n-1), 1:(m-1)]
    diffs <- c(diffs, list(mat[2L:n, 1L:(m - 1L)] - mat[1L:(n - 1L), 1L:(m - 1L)]))
  }

  len <- (n - 1L) * (m - 1L)
  stacked <- do.call(vstack,
    lapply(diffs, function(d) reshape_expr(d, c(1L, len)))
  )
  sum_entries(p_norm(stacked, p = 2, axis = 2L))
}