R/loess.R

In twolodzko/misster: Simple tools for dealing with missing values

#' Replace missing values using local polynomial regression
#'
#' @param data       a data.frame.
#' @param formula    an object of class "\code{\link[stats]{formula}}":
#'                   a symbolic description of the model to be fitted.
#'                   Alternatively, if \code{formula} is a numeric vector,
#'                   a bivariate model \code{x ~ seq_along(x)} is used.
#'
#' @param \dots      further arguments passed to \code{\link[stats]{loess}}.
#'
#' @examples
#'
#' na_loess( c(NA, 1, 2, NA, 4, NA, NA, NA, 8))
#'
#' dat <- mtcars
#' dat$disp[sample.int(nrow(dat), 10)] <- NA
#'
#' na_loess(dat, disp ~ mpg + cyl)
#'
#' set.seed(123)
#'
#' x <- sin(seq(-4, 4, length.out = 120))
#' x <- rnorm(120, x, 0.5)
#' x[sample.int(length(x), 20)] <- NA
#'
#' plot(x)
#' points(which(is.na(x)), na_loess(x)[is.na(x)], col = "red", pch = 16)
#'
#' @importFrom stats loess
#'
#' @export

na_loess <- function(data, formula, ...) UseMethod("na_loess")

#' @export

na_loess.data.frame <- function(data, formula, ...) {
  na_predict(data, formula, learnFun = loess, ...)
}

#' @export

na_loess.numeric <- function(data, ...) {

  if ( !is.null(formula) )
    message("formula argument was ignored")

  x <- as_imputed(data)
  nas <- is.na(x)
  x_pos <- seq_along(x)

  data <- data.frame(
    y = x[!nas],
    x = x_pos[!nas]
  )

  model <- loess(y ~ x, data = data, ...)
  pred <- predict(model, newdata = data.frame(x = x_pos[nas]))

  if ( !is_simple_vector(pred) )
    stop("invalid format of predictions")
  if ( length(pred) != sum(nas) )
    stop("model failed predict all the missing values")

  x[nas] <- pred
  x
}