R/relabel_predictors.R

In dotwhisker: Dot-and-Whisker Plots of Regression Results

#' Relabel the Predictors in a Tidy Data Frame of Regression Results
#'
#' \code{relabel_predictors} is a convenience function for relabeling the predictors in a tidy data frame to be passed to \code{\link[dotwhisker]{dwplot}} or a plot generated by \code{\link[dotwhisker]{dwplot}}
#'
#' @param x Either a tidy data frame to be passed to \code{\link[dotwhisker]{dwplot}} or a plot generated by \code{\link[dotwhisker]{dwplot}}.
#' @param ... Named replacements, as in \code{\link[dplyr]{recode}}. The argument names should be the current values to be replaced, and the argument values should be the new (replacement) values.  For backwards compatibility, a named character vector, with new values as values, and old values as names may also be used.  The order of the named replacements will be preserved, so this function also serves the purpose of reordering variables.
#'
#' @return The function returns an object of the same type as it is passed: a tidy data frame or a plot generated by \code{\link[dotwhisker]{dwplot}}.
#'
#' @examples
#' library(broom)
#' library(dplyr)
#'
#' data(mtcars)
#' m1 <- lm(mpg ~ wt + cyl + disp, data = mtcars)
#' m1_df <- broom::tidy(m1) %>%
#'          relabel_predictors("(Intercept)" = "Intercept",
#'                               wt = "Weight",
#'                               disp = "Displacement",
#'                               cyl = "Cylinder")
#' dwplot(m1_df)
#'
#' dwplot(m1, show_intercept = TRUE) %>%
#'     relabel_predictors("(Intercept)" = "Intercept",
#'                               wt = "Weight",
#'                               disp = "Displacement",
#'                               cyl = "Cylinder")
#'
#'
#' @importFrom dplyr select one_of arrange
#' @importFrom rlang exprs
#' @importFrom stats setNames
#'
#' @export

relabel_predictors <- function(x, ...) {
    model <- term <- y_ind <- NULL # Set to NULL to make R CMD check happy

    dots <- rlang::exprs(...)
    if (is.language(dots[[1]])) {                   # if a named vector is passed,
        dots <- stats::setNames(as.list(...), names(...))      # convert it to a list
    }

    if (is.data.frame(x)) {
        if ("model" %in% names(x)) {
            x <- arrange(x, model, match(term, names(dots)))
            x$term <- factor(x$term, levels = unique(x$term))
        } else {
            x <- arrange(x, match(term, names(dots)))
            x$term <- factor(x$term, levels = unique(x$term))
        }
        x$term <- dplyr::recode(x$term, !!! dots)
        return(x)
    }
    else if (is.ggplot(x)) {
        if ("model" %in% names(x$data)) {
            x$data <- arrange(x$data, model, match(term, names(dots))) %>%
                select(-y_ind)
            x$data$term <- factor(x$data$term, levels = unique(x$data$term))
        } else {
            x$data <- arrange(x$data, match(term, names(dots))) %>%
                select(-y_ind)
            x$data$term <- factor(x$data$term, levels = unique(x$data$term))
        }

        x$data$term <- dplyr::recode(x$data$term, !!! dots)

        args_list <- list(x = x$data, x$args)

        args_list <- x$args
        args_list$x <- x$data
        args_list$vars_order <- NULL # removing the vars_order since the var names are already changed

        p <- do.call(dwplot, args_list)
        return(p)
    } else stop("x should be either a tidy data frame or a plot generated by dwplot")
}