R/ops.R

In tf: S3 Classes and Methods for Tidy Functional Data

#' Math, Summary and Ops Methods for `tf`
#'
#' These methods and operators mostly work `arg`-value-wise on `tf` objects, see
#' [vctrs::vec_arith()] etc. for implementation details.
#'
#' - Operations on `tfd`-objects do not extrapolate functions on a common grid first,
#' they operate on the function at argument values that both objects have in
#' common.
#' - With the exception of addition and
#' multiplication, operations on `tfb`-objects first evaluate the data on their
#' `arg`, perform computations on these evaluations and then convert back to an
#' `tfb`- object, so a loss of precision should be expected -- especially so for
#' small spline bases and/or very wiggly data.
#' - Equality checks of functional objects are even more iffy
#' than usual for computer math and not very reliable.
#' - Note that `max` and `min` are not guaranteed to be maximal/minimal over the
#'  entire domain, only at the argument values used for computation.
#'
#' See examples below, many more are in `tests/testthat/test-ops.R`.
#'
#' @param x a `tf` or `numeric` object.
#' @param y a `tf` or `numeric` object.
#' @param op An arithmetic operator as a string.
#' @param ... `tf`-objects (not used for `Math` group generic).
#' @param e1 an `tf` or a numeric vector.
#' @param e2 an `tf` or a numeric vector.
#' @returns a `tf`- or `logical` vector with the computed result.
#' @seealso [tf_fwise()] for scalar summaries of each function in a `tf`-vector
#' @name tfgroupgenerics
#' @family tidyfun compute functions
#' @examples
#' set.seed(1859)
#' f <- tf_rgp(4)
#' 2 * f == f + f
#' sum(f) == f[1] + f[2] + f[3] + f[4]
#' log(exp(f)) == f
#' plot(f, points = FALSE)
#' lines(range(f), col = 2, lty = 2)
#'
#' f2 <- tf_rgp(5) |> exp() |> tfb(k = 25)
#' layout(t(1:3))
#' plot(f2, col = gray.colors(5))
#' plot(cummin(f2), col = gray.colors(5))
#' plot(cumsum(f2), col = gray.colors(5))
#'
#' # ?tf_integrate for integrals, ?tf_fwise for scalar summaries of each function
NULL

#' @rdname tfgroupgenerics
#' @export
`==.tfd` <- function(e1, e2) {
  assert_compatible_size("==", e1, e2)
  # not comparing names, as per convention...
  same <- all(compare_tf_attribs(e1, e2))
  if (!same) {
    rep(FALSE, vec_size(e1))
  } else {
    map2_lgl(e1, e2, \(x, y) isTRUE(all.equal(x, y)))
  }
}

#' @rdname tfgroupgenerics
#' @export
`!=.tfd` <- function(e1, e2) !(e1 == e2)

# need to copy instead of defining tf-method s.t. dispatch in Ops works
#' @rdname tfgroupgenerics
#' @export
`==.tfb` <- eval(`==.tfd`)

#' @rdname tfgroupgenerics
#' @export
`!=.tfb` <- eval(`!=.tfd`)

#' @rdname tfgroupgenerics
#' @export
#' @method vec_arith tfd
vec_arith.tfd <- function(op, x, y, ...) {
  UseMethod("vec_arith.tfd", y)
}

#' @export
#' @method vec_arith.tfd default
vec_arith.tfd.default <- function(op, x, y, ...) {
  stop_incompatible_op(op, x, y)
}

#' @export
#' @method vec_arith.tfd tfd
vec_arith.tfd.tfd <- function(op, x, y, ...) {
  switch(
    op,
    `+` = ,
    `-` = ,
    `*` = ,
    `/` = ,
    `^` = ,
    `%%` = ,
    `%/%` = tfd_op_tfd(op, x, y),
    stop_incompatible_op(op, x, y)
  )
}

#' @export
#' @method vec_arith.tfd numeric
vec_arith.tfd.numeric <- function(op, x, y, ...) {
  switch(
    op,
    `+` = ,
    `-` = ,
    `*` = ,
    `/` = ,
    `^` = ,
    `%%` = ,
    `%/%` = tfd_op_numeric(op, x, y),
    stop_incompatible_op(op, x, y)
  )
}

#' @export
#' @method vec_arith.numeric tfd
vec_arith.numeric.tfd <- function(op, x, y, ...) {
  switch(
    op,
    `+` = ,
    `-` = ,
    `*` = ,
    `/` = ,
    `^` = ,
    `%%` = ,
    `%/%` = numeric_op_tfd(op, x, y),
    stop_incompatible_op(op, x, y)
  )
}

#' @export
#' @method vec_arith.tfd MISSING
vec_arith.tfd.MISSING <- function(op, x, y, ...) {
  tf_op_missing(op, x, y, ...)
}

#' @rdname tfgroupgenerics
#' @export
#' @method vec_arith tfb
vec_arith.tfb <- function(op, x, y, ...) {
  UseMethod("vec_arith.tfb", y)
}

#' @export
#' @method vec_arith.tfb default
vec_arith.tfb.default <- function(op, x, y, ...) {
  stop_incompatible_op(op, x, y)
}

#' @export
#' @method vec_arith.tfb tfb
vec_arith.tfb.tfb <- function(op, x, y, ...) {
  switch(
    op,
    `+` = ,
    `-` = tfb_plusminus_tfb(op, x, y),
    `*` = ,
    `/` = ,
    `^` = ,
    `%%` = ,
    `%/%` = tfb_op_tfb(op, x, y),
    stop_incompatible_op(op, x, y)
  )
}

#' @export
#' @method vec_arith.tfb numeric
vec_arith.tfb.numeric <- function(op, x, y, ...) {
  switch(
    op,
    `/` = ,
    `*` = tfb_multdiv_numeric(op, x, y),
    `+` = ,
    `-` = ,
    `^` = ,
    `%%` = ,
    `%/%` = tfb_op_numeric(op, x, y),
    stop_incompatible_op(op, x, y)
  )
}

#' @export
#' @method vec_arith.numeric tfb
vec_arith.numeric.tfb <- function(op, x, y, ...) {
  switch(
    op,
    `*` = tfb_multdiv_numeric(op, y, x),
    `+` = ,
    `-` = ,
    `/` = ,
    `^` = ,
    `%%` = ,
    `%/%` = numeric_op_tfb(op, x, y),
    stop_incompatible_op(op, x, y)
  )
}

#' @export
#' @method vec_arith.tfb MISSING
vec_arith.tfb.MISSING <- function(op, x, y, ...) {
  tf_op_missing(op, x, y, ...)
}

#-------------------------------------------------------------------------------

tfd_op_tfd <- function(op, x, y) {
  assert_compatible_size(op, x, y)

  if (!domains_overlap(x, y)) {
    message <- cli::format_inline(
      "<{vec_ptype_full(x)}> {op} <{vec_ptype_full(y)}> not permitted for non-overlapping domains"
    )
    stop_incompatible_op(op, x, y, message = message)
  }

  same_domain <- all.equal(
    tf_domain(x),
    tf_domain(y),
    check.attributes = FALSE
  ) |>
    isTRUE()
  same_arg <- all.equal(tf_arg(x), tf_arg(y), check.attributes = FALSE) |>
    isTRUE()

  if (same_arg) {
    arg_ret <- tf_arg(x) |> ensure_list()
    x_ <- tf_evaluations(x)
    y_ <- tf_evaluations(y)
  } else {
    arg_ret <- common_args(x, y) |> ensure_list()
    if (all(lengths(arg_ret) == 0)) {
      message <- cli::format_inline(
        "<{vec_ptype_full(x)}> {op} <{vec_ptype_full(y)}> not possible: no common argument values."
      )
      cli::cli_abort(c(
        message,
        "i" = "Use {.fn tf_rebase} or {.fn tf_interpolate} to change argument values."
      ))
    }
    use_x <- map2(
      ensure_list(tf_arg(x)),
      arg_ret,
      \(arg, grid) which(arg %in% grid)
    )
    use_y <- map2(
      ensure_list(tf_arg(y)),
      arg_ret,
      \(arg, grid) which(arg %in% grid)
    )
    x_ <- map2(tf_evaluations(x), use_x, `[`)
    y_ <- map2(tf_evaluations(y), use_y, `[`)
  }

  domain_ret <- c(
    min(tf_domain(x), tf_domain(y)),
    max(tf_domain(x), tf_domain(y))
  )

  if (!same_domain || !same_arg) {
    message <- cli::format_inline(
      "Attempting <{vec_ptype_full(x)}> {op} <{vec_ptype_full(y)}> for different",
      "{ifelse(same_domain, '', ' domains')}",
      "{ifelse(!same_domain & !same_arg, ' and', '')}",
      "{ifelse(same_arg, '', ' argument values')}"
    )
    cli::cli_warn(c(
      message,
      "i" = "Use {.fn tf_rebase} or {.fn tf_interpolate} to adjust argument values first if necessary."
    ))
  }

  evals_ret <- map2(x_, y_, \(x, y) {
    if (is.null(x) || is.null(y)) return(NULL)
    do.call(op, list(x, y))
  })
  if (vec_size(x) >= vec_size(y)) {
    names(evals_ret) <- names(x)
  } else {
    names(evals_ret) <- names(y)
  }

  evaluator_ret <- ifelse(
    vec_size(x) >= vec_size(y),
    attr(x, "evaluator_name"),
    attr(y, "evaluator_name")
  ) |>
    sym()
  if (attr(x, "evaluator_name") != attr(y, "evaluator_name")) {
    cli::cli_inform(c(
      "!" = "Inputs have different evaluators, result has {.val {evaluator_ret}}."
    ))
  }

  do.call(
    tfd,
    list(
      data = evals_ret,
      arg = arg_ret,
      domain = domain_ret,
      evaluator = evaluator_ret
    )
  )
}

tfd_op_numeric <- function(op, x, y, ...) {
  assert_compatible_size(op, x, y)
  ret <- map2(tf_evaluations(x), y, \(x, y) {
    if (is.null(x)) return(NULL)
    result <- do.call(op, list(x, y))
    if (allMissing(result)) NULL else result
  })
  if (is_irreg(x)) {
    ret <- map2(tf_arg(x), ret, \(.arg, .ret) {
      if (is.null(.ret)) return(NULL)
      list(arg = .arg, value = .ret)
    })
  }
  attributes(ret) <- attributes(x)
  if (vec_size(y) > 1) {
    names(ret) <- NULL
  }
  ret
}

# some code-duplication here, this makes non-commutative ops work for tfd and numeric
numeric_op_tfd <- function(op, x, y) {
  assert_compatible_size(op, x, y)
  ret <- map2(x, tf_evaluations(y), \(x, y) {
    if (is.null(y)) return(NULL)
    result <- do.call(op, list(x, y))
    if (allMissing(result)) NULL else result
  })
  if (is_irreg(y)) {
    ret <- map2(tf_arg(y), ret, \(.arg, .ret) {
      if (is.null(.ret)) return(NULL)
      list(arg = .arg, value = .ret)
    })
  }
  attributes(ret) <- attributes(y)
  if (vec_size(x) > 1) {
    names(ret) <- NULL
  }
  ret
}

# construct a tfb shell with NULL entries, using ref_tfb for attributes
tfb_na_result <- function(eval, ref_tfb) {
  result <- vector("list", length(eval))
  result[] <- list(NULL)
  names(result) <- names(eval)
  attributes(result) <- attributes(ref_tfb)
  names(result) <- names(eval)
  result
}

#-------------------------------------------------------------------------------

tfb_multdiv_numeric <- function(op, x, y) {
  # dispatch to general operator implementation (i.e. cast to tfd and back)
  # if link functions are involved:
  if (is_tfb_spline(x) && attr(x, "family")$link != "identity") {
    return(tfb_op_numeric(op, x, y))
  }
  # if not, * and / can simply be applied to the basis coefficients:
  ret <- map2(vec_data(x), y, \(x, y) {
    if (is.null(x)) return(NULL)
    result <- do.call(op, list(e1 = x, e2 = y))
    if (allMissing(result)) NULL else result
  })
  attributes(ret) <- attributes(x)
  ret
}

tfb_op_numeric <- function(op, x, y) {
  cli::cli_warn(
    "Potentially lossy cast to {.cls tfd} and back in {.cls {vec_ptype_full(x)}} {op} {.cls {vec_ptype_full(y)}}."
  )
  eval <- tfd_op_numeric(op, tfd(x), y)
  na_entries <- is.na(eval)
  if (all(na_entries)) return(tfb_na_result(eval, x))
  if (any(na_entries)) {
    rebased <- tf_rebase(
      eval[!na_entries],
      x[!na_entries],
      penalized = FALSE,
      verbose = FALSE
    )
    result <- tfb_na_result(eval, rebased)
    result[!na_entries] <- unclass(rebased)
    return(result)
  }
  tf_rebase(eval, x, penalized = FALSE, verbose = FALSE)
  #TODO: restore sp afterwards so all properties are preserved?
}

numeric_op_tfb <- function(op, x, y) {
  cli::cli_warn(
    "Potentially lossy cast to {.cls tfd} and back in {.cls {vec_ptype_full(x)}} {op} {.cls {vec_ptype_full(y)}}."
  )
  eval <- numeric_op_tfd(op, x, tfd(y))
  na_entries <- is.na(eval)
  if (all(na_entries)) return(tfb_na_result(eval, y))
  if (any(na_entries)) {
    rebased <- tf_rebase(
      eval[!na_entries],
      y[!na_entries],
      penalized = FALSE,
      verbose = FALSE
    )
    result <- tfb_na_result(eval, rebased)
    result[!na_entries] <- unclass(rebased)
    return(result)
  }
  tf_rebase(eval, y, penalized = FALSE, verbose = FALSE) #TODO: see tfb_op_numeric
}

tfb_op_tfb <- function(op, x, y) {
  cli::cli_warn(
    "Potentially lossy casts to {.cls tfd} and back for {.cls {vec_ptype_full(x)}} {op} {.cls {vec_ptype_full(y)}}."
  )
  eval <- tfd_op_tfd(op, tfd(x), tfd(y))
  ret_ptype <- if (vec_size(x) >= vec_size(y)) vec_ptype(x) else vec_ptype(y)
  na_entries <- is.na(eval)
  if (all(na_entries)) return(tfb_na_result(eval, ret_ptype))
  if (any(na_entries)) {
    rebased <- tf_rebase(
      eval[!na_entries],
      ret_ptype,
      penalized = FALSE,
      verbose = FALSE
    )
    result <- tfb_na_result(eval, rebased)
    result[!na_entries] <- unclass(rebased)
    return(result)
  }
  tf_rebase(eval, ret_ptype, penalized = FALSE, verbose = FALSE) #TODO: see tfb_op_numeric
}

tfb_plusminus_tfb <- function(op, x, y) {
  assert_compatible_size(op, x, y)
  # rebase to basis of shorter or first argument if bases are not compatible
  # less computation to tf_rebase shorter input.
  # would be potentially more accurate to decide based on basis dimension and
  # tf_rebase the one with the smaller basis...
  if (!same_basis(x, y)) {
    if (vec_size(x) >= vec_size(y)) {
      cli::cli_warn(
        "Bases unequal -- potentially lossy {.fn tf_rebase} for 2nd argument in {.cls {vec_ptype_full(x)}} {op} {.cls {vec_ptype_full(y)}}."
      )
      y <- tf_rebase(y, x)
    } else {
      cli::cli_warn(
        "Bases unequal -- potentially lossy {.fn tf_rebase} for 1st argument in {.cls {vec_ptype_full(x)}} {op} {.cls {vec_ptype_full(y)}}."
      )
      x <- tf_rebase(x, y)
    }
  }
  # dispatch to general operator implementation (i.e. cast to tfd and back)
  # if link functions are involved
  if (!all(c(attr(x, "family")$link, attr(y, "family")$link) == "identity")) {
    return(tfb_op_tfb(op, x, y))
  }
  ret <- map2(
    vec_data(x),
    vec_data(y),
    \(x, y) {
      if (is.null(x) || is.null(y)) return(NULL)
      result <- do.call(op, list(e1 = x, e2 = y))
      if (allMissing(result)) NULL else result
    }
  )
  attributes(ret) <- if (vec_size(x) >= vec_size(y)) {
    attributes(x)
  } else {
    attributes(y)
  }
  ret
}

#-------------------------------------------------------------------------------

tf_op_missing <- function(op, x, y, ...) {
  switch(op, `-` = x * -1, `+` = x, stop_incompatible_op(op, x, y))
}