R/customPredict.R

In twolodzko/twextras: twextras

#' Create a custom predict function
#' 
#' \code{customPredict} creates a custom predict function
#' that transforms the output of usual \code{\link[stats]{predict}}.
#' 
#' @param model     a model (e.g. of class \code{lm}, \code{glm},
#'                  \code{rpart} etc.), that has associated
#'                  \code{\link[stats]{predict}} function.
#' @param transform a function used to transform output of
#'                  \code{predict(model, ...)}.
#' @param \dots     futher arguments passed to \code{\link[stats]{predict}}.
#' @param x         an \code{customPredict}-class object.
#'        
#' @examples 
#' 
#' model1 <- lm(log(mpg) ~ 1, data = mtcars)
#' model1_cp <- customPredict(model1, transform = exp)
#' 
#' predict(model1)
#' exp(predict(model1))
#' predict(model1_cp)
#'
#' @importFrom stats predict
#' @export

customPredict <- function(model, transform = identity, ...) {
  if (!is.function(transform))
    stop("transform is not a function")
  class(model) <- c("customPredict", attr(model, "class"))
  attr(model, "customPredict_transform") <- transform
  attr(model, "customPredict_param") <- list(...)
  invisible(model)
}

#' @rdname customPredict
#' @export

is_customPredict <- function(x) {
  inherits(x, "customPredict") &&
    "customPredict_transform" %in% names(attributes(x))
}

#' @rdname customPredict
#' @export

predict.customPredict <- function(model, ...) {
  transform <- attr(model, "customPredict_transform")
  param <- attr(model, "customPredict_param")
  attr(model, "customPredict_transform") <- NULL
  attr(model, "customPredict_param") <- NULL
  class(model) <- class(model)[-inherits(model, "customPredict", TRUE)]
  transform(do.call(predict, list(model, param, ...)))
}