diff --git a/R/type_loess.R b/R/type_loess.R index 3f751d31..b7d51af4 100644 --- a/R/type_loess.R +++ b/R/type_loess.R @@ -6,6 +6,9 @@ #' @inheritParams stats::loess #' @param se logical. If `TRUE` (the default), confidence intervals are drawn. #' @param level the confidence level required if `se = TRUE`. Default is 0.95. +#' @param locfit local. If `TRUE` the LOCFIT library is used to estimate the +#' local polynomial regression which is much faster. Note that `span` is reused +#' for locfits `nn` option. #' @importFrom stats loess loess.control predict #' @examples #' # "loess" type convenience string @@ -20,11 +23,21 @@ type_loess = function( family = "gaussian", control = loess.control(), se = TRUE, - level = 0.95 + level = 0.95, + locfit = FALSE ) { + if (locfit == TRUE) { + pkg = "locfit" + if (!requireNamespace(pkg, quietly = TRUE)) { + msg = sprintf("Package '%s' is required by type_locfit(). Install with install.packages('%s').", + pkg, pkg) + warning(msg, call. = FALSE) + stop(msg, call. = FALSE) + } + } out = list( draw = draw_ribbon(), - data = data_loess(span = span, degree = degree, family = family, control = control, se = se, level = level), + data = data_loess(span = span, degree = degree, family = family, control = control, se = se, level = level, locfit = locfit), name = if (isTRUE(se)) "ribbon" else "l" ) class(out) = "tinyplot_type" @@ -32,18 +45,29 @@ type_loess = function( } -data_loess = function(span, degree, family, control, se, level, ...) { +data_loess = function(span, degree, family, control, se, level, locfit, ...) { fun = function(datapoints, ...) { datapoints = split(datapoints, list(datapoints$facet, datapoints$by)) datapoints = Filter(function(k) nrow(k) > 0, datapoints) datapoints = lapply(datapoints, function(dat) { - fit = loess(y ~ x, data = dat, span = span, degree = degree, family = family, control = control) + if (locfit == FALSE) { + fit = loess(y ~ x, data = dat, span = span, degree = degree, family = family, control = control) + } + else { + fit = locfit::locfit(y ~ locfit::lp(x, nn = span, deg = degree), data = dat, family = family) + } if (se == TRUE) { + if (locfit == FALSE) { p = predict(fit, newdata = dat, se = TRUE) - p = ci(p$fit, p$se.fit, conf.level = level, p$df) - dat$y = p$estimate - dat$ymax = p$conf.high - dat$ymin = p$conf.low + } + else { + p = predict(fit, newdata = dat, se.fit = TRUE) + p$df = Inf + } + p = ci(p$fit, p$se.fit, conf.level = level, p$df) + dat$y = p$estimate + dat$ymax = p$conf.high + dat$ymin = p$conf.low } else { dat$y = predict(fit, dat) }