R/winsor.R

In jmbarbone/qpm: Quantitative Psychological Methods

#' Winsor
#'
#' @param x A numeric vector or data.frame.
#' @param method A character value of either "mean", "var", or "sd" to denote the type of analysis.
#' @param trim The proportion of the data by which to trim from the ends
#' @param na.rm Logica.  If TRUE, NA values are removed from the data.
#'
#' @export

winsor <- function(x, method = c("mean", "var", "sd"), trim = 0.2, na.rm = FALSE) {
  UseMethod("winsor")
}

#' @export
winsor.default <- function(x, .method = c("mean", "var", "sd"), .trim = 0.2, .na_rm = FALSE, ...) {
  method <- match.arg(method)
  switch(method,
         "mean" = win_fun(x, mean, .trim, .na_rm),
         "var" = win_fun(x, var, .trim, .na_rm),
         "sd" = sqrt(win_fun(x, var, .trim, .na_rm)))
}

#' @export
winsor.data.frame <- function(x, .method = c("mean", "var", "sd"), .trim = 0.2, .na_rm = FALSE, ...) {
  sapply(x, winsor.default, .method = method, .trim = trim, .na_rm = .na_rm)
}

win_fun <- function(x, .f, .trim, .na_rm, ...) {
  if (trim < .5) {
    .f(winsorize(x, trim = .trim, na.rm = .na_rm), ...)
  } else {
    median(x, na.rm = TRUE)
  }
}


#' @export
#' @rdname winsor
winsorize <- function(x, trim = 0.2, na.rm = FALSE) {
  UseMethod("winsorize")
}


#' @export
#' @rdname winsor
winsorise <- function(x, trim = 0.2, na.rm = FALSE) {
  UseMethod("winsorize", x)
}

#' @export
#' @rdname winsor
winsorize.default <- function(x, trim = 0.2, na.rm = FALSE) {
  if ((trim < 0) | (trim > 0.5) ) warning("Consider less extreme values.", call. = FALSE)
  qtrim <- quantile(x, probs = c(trim, 0.5, 1 - trim), na.rm = na.rm, names = FALSE)
  lower <- qtrim[1]
  upper <- qtrim[3]
  if(trim < .5) {
    x[x < lower] <- lower
    x[x > upper] <- upper
  } else {
    x[!is.na(x)] <- qtrim[2]
  }
  x
}

#' @export
#' @rdname winsor
winsorize.data.frame <- function(x, trim = 0.2, na.rm = FALSE) {
  x[] <- sapply(x, winsorize, trim = trim, na.rm = na.rm)
  x
}