R/evaluate.R

In tf: S3 Classes and Methods for Tidy Functional Data

#' Evaluate `tf`-vectors for given argument values
#'
#' Also used internally by the `[`-operator for `tf` data (see `?tfbrackets`) to
#' evaluate `object`, see examples.
#'
#' @param object a `tf`, or a `data.frame`-like object with `tf` columns.
#' @param arg optional evaluation grid (vector or list of vectors).
#'   Defaults to `tf_arg(object)`, implicitly.
#' @param evaluator optional. The function to use for inter/extrapolating the
#'  `tfd`. Defaults to `tf_evaluator(object)`.
#'  See e.g. [tf_approx_linear()] for details.
#' @param ... not used.
#' @returns A list of numeric vectors containing the function
#'   evaluations on `arg`.
#' @export
#' @family tidyfun inter/extrapolation functions
#' @examples
#' f <- tf_rgp(3, arg = seq(0, 1, length.out = 11))
#' tf_evaluate(f) |> str()
#' tf_evaluate(f, arg = 0.5) |> str()
#' # equivalent, as matrix:
#' f[, 0.5]
#' new_grid <- seq(0, 1, length.out = 6)
#' tf_evaluate(f, arg = new_grid) |> str()
#' # equivalent, as matrix:
#' f[, new_grid]
tf_evaluate <- function(object, arg, ...) UseMethod("tf_evaluate")

#' @export
#' @rdname tf_evaluate
tf_evaluate.default <- function(object, arg, ...) .NotYetImplemented()

#' @export
#' @rdname tf_evaluate
tf_evaluate.tfd <- function(
  object,
  arg,
  evaluator = tf_evaluator(object),
  ...
) {
  if (missing(arg) || is.null(arg)) {
    return(tf_evaluations(object))
  }
  arg <- ensure_list(arg)
  assert_arg(arg, object, check_unique = FALSE)
  ret <- pmap(
    list(arg, ensure_list(tf_arg(object)), tf_evaluations(object)),
    function(x, y, z) {
      if (!length(z)) {
        return(rep(NA_real_, length(x)))
      }
      evaluate_tfd_once(
        new_arg = x,
        arg = y,
        evaluations = z,
        evaluator = evaluator
      )
    }
  )

  setNames(ret, names(object))
}

evaluate_tfd_once <- function(
  new_arg,
  arg,
  evaluations,
  evaluator
) {
  if (isTRUE(all.equal(new_arg, arg))) return(evaluations)
  seen <- match(new_arg, arg)
  seen_index <- na.omit(seen)
  seen <- !is.na(seen)
  ret <- rep(NA_real_, length(new_arg))
  ret[seen] <- evaluations[seen_index]
  ret[!seen] <-
    evaluator(new_arg[!seen], arg = arg, evaluations = evaluations)
  ret
}

#' @export
#' @rdname tf_evaluate
tf_evaluate.tfb <- function(object, arg, ...) {
  if (missing(arg) || is.null(arg)) {
    return(tf_evaluations(object))
  }
  arg <- ensure_list(arg)
  assert_arg(arg, object, check_unique = FALSE)
  if (length(arg) == 1) {
    arg <- unlist(arg, use.names = FALSE)
    na_entries <- is.na(object)
    ret <- vector("list", length(object))
    if (any(!na_entries)) {
      evals <- evaluate_tfb_once(
        x = arg,
        arg = tf_arg(object),
        coefs = do.call(cbind, coef(object[!na_entries])),
        basis = attr(object, "basis"),
        X = attr(object, "basis_matrix")
      )
      ret_non_na <- if (length(arg) == 1) {
        # avoid cast to simple vector for point evaluations
        split(evals, seq_along(evals))
      } else {
        split(evals, col(as.matrix(evals)))
      }
      ret[!na_entries] <- ret_non_na
    }
    if (any(na_entries)) {
      ret[na_entries] <- list(rep(NA_real_, length(arg)))
    }
  } else {
    ret <- pmap(
      list(arg, ensure_list(tf_arg(object)), coef(object)),
      function(x, y, z) {
        if (!length(z) || anyNA(z)) {
          return(rep(NA_real_, length(x)))
        }
        evaluate_tfb_once(
          x = x,
          arg = y,
          coefs = z,
          basis = attr(object, "basis"),
          X = attr(object, "basis_matrix")
        )
      }
    )
  }
  if (!is_tfb_fpc(object)) {
    ret <- map(ret, attr(object, "family")$linkinv)
  }
  setNames(ret, names(object))
}

evaluate_tfb_once <- function(x, arg, coefs, basis, X) {
  seen <- match(x, arg)
  seen_index <- na.omit(seen)
  seen <- !is.na(seen)
  if (all(seen)) return(drop(X[seen_index, , drop = FALSE] %*% coefs))
  Xnew <- X[rep(1, length(x)), ]
  if (any(seen)) Xnew[seen, ] <- X[seen_index, , drop = FALSE]
  Xnew[!seen, ] <- basis(x[!seen])
  drop(Xnew %*% coefs)
}