regression_functions.R

globalVariables(c(
  "term", ".resid", "AIC", "BIC", "deviance", "df.residual", "logLik", "ID",
  "mse", "rmse", "residual", "r_squared", "adj_r_squared", "conf_low",
  "conf_high"
))

#' Get regression table
#'
#' Output regression table for an `lm()` regression in "tidy" format. This function
#' is a wrapper function for `broom::tidy()` and includes confidence
#' intervals in the output table by default.
#'
#' @param model an `lm()` model object
#' @inheritParams broom::tidy.lm
#' @param digits number of digits precision in output table
#' @param print If TRUE, return in print format suitable for R Markdown
#' @param default_categorical_levels If TRUE, do not change the non-baseline
#'  categorical variables in the term column. Otherwise non-baseline
#'  categorical variables will be displayed in the format
#'  "categorical_variable_name-level_name"
#'
#' @return A tibble-formatted regression table along with lower and upper end
#' points of all confidence intervals for all parameters `lower_ci` and
#' `upper_ci`; the confidence levels default to 95\%.
#' @importFrom stats lm
#' @importFrom stats predict
#' @importFrom formula.tools lhs
#' @importFrom formula.tools rhs
#' @importFrom broom tidy
#' @importFrom tibble as_tibble
#' @importFrom janitor clean_names
#' @importFrom knitr kable
#' @importFrom purrr imap
#' @importFrom stringr str_replace_all
#' @importFrom stringr coll
#' @export
#' @seealso [`tidy()`][broom::reexports], [get_regression_points()], [get_regression_summaries()]
#'
#' @examples
#' library(moderndive)
#'
#' # Fit lm() regression:
#' mpg_model <- lm(mpg ~ cyl, data = mtcars)
#'
#' # Get regression table:
#' get_regression_table(mpg_model)
#'
#' # Vary confidence level of confidence intervals
#' get_regression_table(mpg_model, conf.level = 0.99)
get_regression_table <- function(model, conf.level = 0.95, digits = 3, 
                                 print = FALSE, 
                                 default_categorical_levels = FALSE) {
  # Check inputs
  input_checks(model, digits, print)
  
  if (!default_categorical_levels && length(model[["xlevels"]]) > 0) {
    # Add delimiter in dummy-coded categorical variables, as in "var-level"
    delim <- "-"
    old_names <- unlist(imap(model[["xlevels"]], ~ paste0(.y, .x)))
    new_names <- unlist(imap(model[["xlevels"]], ~ paste0(.y, delim, .x)))
    names(new_names) <- old_names
    names(model[["coefficients"]]) <- names(model[["coefficients"]]) %>%
      str_replace_all(coll(new_names))
  }
  
  # Create output tibble
  regression_table <- model %>%
    tidy(conf.int = TRUE, conf.level) %>%
    mutate_if(is.numeric, round, digits = digits) %>%
    mutate(term = ifelse(term == "(Intercept)", "intercept", term)) %>%
    as_tibble() %>%
    clean_names() %>%
    rename(
      lower_ci = conf_low,
      upper_ci = conf_high
    )
  
  # Transform to markdown
  if (print) {
    regression_table <- regression_table %>%
      kable()
  }
  
  return(regression_table)
}


#' Get regression points
#'
#' Output information on each point/observation used in an `lm()` regression in
#' "tidy" format. This function is a wrapper function for `broom::augment()`
#' and renames the variables to have more intuitive names.
#'
#' @inheritParams get_regression_table
#' @param newdata A new data frame of points/observations to apply `model` to
#' obtain new fitted values and/or predicted values y-hat. Note the format of
#' `newdata` must match the format of the original `data` used to fit
#' `model`.
#' @param ID A string indicating which variable in either the original data used
#'  to fit `model` or `newdata` should be used as
#' an identification variable to distinguish the observational units
#' in each row. This variable will be the left-most variable in the output data
#' frame. If `ID` is unspecified, a column `ID` with values 1 through the number 
#' of rows is returned as the identification variable.
#'
#' @return A tibble-formatted regression table of outcome/response variable,
#' all explanatory/predictor variables, the fitted/predicted value, and residual.
#' @importFrom dplyr select
#' @importFrom dplyr rename_at
#' @importFrom dplyr vars
#' @importFrom dplyr rename
#' @importFrom dplyr mutate
#' @importFrom dplyr pull
#' @importFrom dplyr everything
#' @importFrom dplyr mutate_if
#' @importFrom dplyr summarise
#' @importFrom stats formula
#' @importFrom formula.tools lhs
#' @importFrom formula.tools rhs
#' @importFrom broom augment
#' @importFrom tibble as_tibble
#' @importFrom janitor clean_names
#' @importFrom stringr str_c
#' @importFrom knitr kable
#' @importFrom rlang sym
#' @importFrom rlang ":="
#' @importFrom stats na.omit
#' @importFrom stats terms
#' @export
#' @seealso [`augment()`][broom::reexports], [get_regression_table()], [get_regression_summaries()]
#'
#' @examples
#' library(dplyr)
#' library(tibble)
#'
#' # Convert rownames to column
#' mtcars <- mtcars %>%
#'   rownames_to_column(var = "automobile")
#'
#' # Fit lm() regression:
#' mpg_model <- lm(mpg ~ cyl, data = mtcars)
#'
#' # Get information on all points in regression:
#' get_regression_points(mpg_model, ID = "automobile")
#'
#' # Create training and test set based on mtcars:
#' training_set <- mtcars %>%
#'   sample_frac(0.5)
#' test_set <- mtcars %>%
#'   anti_join(training_set, by = "automobile")
#'
#' # Fit model to training set:
#' mpg_model_train <- lm(mpg ~ cyl, data = training_set)
#'
#' # Make predictions on test set:
#' get_regression_points(mpg_model_train, newdata = test_set, ID = "automobile")
get_regression_points <-
  function(model, digits = 3, print = FALSE, newdata = NULL, ID = NULL) {
    # Check inputs
    input_checks(model, digits, print)
    if (!is.null(ID)) {
      check_character(ID)
    }
    if (!is.null(newdata)) {
      check_data_frame(newdata)
    }
    
    # Define outcome variable
    outcome_variable <- formula(model) %>%
      lhs() %>%
      all.vars()
    outcome_variable_hat <- str_c(outcome_variable, "_hat")
    
    # Predictor names in the *data*
    explanatory_vars_data <- formula(model) %>%
      rhs() %>%
      all.vars()
    
    # Term labels used in the model matrix / potentially in broom::augment()
    term_labels <- attr(terms(model), "term.labels")
    explanatory_vars_aug <- make.names(term_labels)
    
    # Compute all fitted/predicted values and residuals for three possible
    # cases/scenarios
    if (is.null(newdata)) {
      # Case 1: same data set used to fit model
      
      aug <- broom::augment(model)
      
      # columns that broom uses for meta information
      broom_meta_cols <- c(
        ".fitted", ".se.fit", ".resid",
        ".hat", ".sigma", ".cooksd", ".std.resid"
      )
      
      # treat everything that is not the outcome and not meta as a predictor
      predictor_cols <- setdiff(
        names(aug),
        c(outcome_variable, broom_meta_cols)
      )
      
      # final column order: outcome, predictors, fitted, residual
      cols_to_keep <- c(outcome_variable, predictor_cols, ".fitted", ".resid")
      
      regression_points <- aug %>%
        dplyr::select(dplyr::all_of(cols_to_keep)) %>%
        dplyr::rename_with(~ outcome_variable_hat, ".fitted") %>%
        dplyr::rename(residual = .resid)
    } else {
      # Two cases when we wanted to return point information on a new data set,
      # newdata, different than the one used to fit the model with:
      if (outcome_variable %in% names(newdata)) {
        # Case 2.a) If outcome variable is included, we can compute both fitted
        # values and residuals.
        regression_points <- newdata %>%
          select(!!c(outcome_variable, explanatory_vars_data)) %>%
          # Compute fitted values
          mutate(y_hat = predict(model, newdata = newdata)) %>%
          rename_at(vars("y_hat"), ~outcome_variable_hat) %>%
          # Compute residuals
          mutate(
            residual := !!sym(outcome_variable) - !!sym(outcome_variable_hat)
          )
      } else {
        # Case 2.b) If outcome variable is not included, we can only return
        # predicted values and not the residuals. This corresponds to typical
        # prediction scenario.
        regression_points <- model %>%
          # Compute fitted values:
          augment(newdata = newdata) %>%
          select(!!c(explanatory_vars_aug, ".fitted")) %>%
          rename_at(vars(".fitted"), ~ str_c(outcome_variable, "_hat"))
      }
    }
    
    # Set identification variable for three possible cases/scenarios
    if (is.null(ID)) {
      # Case 1: If ID argument is not specified, set as ID variable as 1 through
      # number of rows
      regression_points <- regression_points %>%
        na.omit() %>%
        mutate(ID = 1:n()) %>%
        select(ID, everything())
    } else {
      # Two cases when ID argument is specified:
      if (is.null(newdata)) {
        # Case 2.a) When computing fitted values and residuals for the same data
        # used to fit the model, extract ID variable from original model fit.
        identification_variable <- eval(model$call$data, 
                                        environment(formula(model))) %>%
          pull(!!ID)
      } else {
        # Case 2.b) When computing predicted values for a new dataset newdata 
        # than the one used to fit the model, extract ID variable from newdata.
        identification_variable <- newdata %>%
          pull(!!ID)
      }
      # Set ID variable
      regression_points <- regression_points %>%
        na.omit() %>%
        mutate(ID = identification_variable) %>%
        select(ID, everything()) %>%
        rename_at(vars("ID"), ~ID)
    }
    
    # Final clean-up
    regression_points <- regression_points %>%
      mutate_if(is.double, round, digits = digits) %>%
      as_tibble()
    
    # Transform to markdown
    if (print) {
      regression_points <- regression_points %>%
        kable()
    }
    
    return(regression_points)
  }


#' Get regression summary values
#'
#' Output scalar summary statistics for an `lm()` regression in "tidy"
#' format. This function is a wrapper function for `broom::glance()`.
#'
#' @inheritParams get_regression_table
#'
#' @return A single-row tibble with regression summaries. Ex: `r_squared` and `mse`.
#' @importFrom dplyr select
#' @importFrom dplyr rename_at
#' @importFrom dplyr vars
#' @importFrom dplyr rename
#' @importFrom dplyr mutate
#' @importFrom dplyr everything
#' @importFrom dplyr mutate_if
#' @importFrom dplyr summarise
#' @importFrom dplyr bind_cols
#' @importFrom dplyr n
#' @importFrom stats formula
#' @importFrom formula.tools lhs
#' @importFrom formula.tools rhs
#' @importFrom broom glance
#' @importFrom broom augment
#' @importFrom tibble as_tibble
#' @importFrom janitor clean_names
#' @importFrom knitr kable
#' @export
#' @seealso [`glance()`][broom::reexports], [get_regression_table()], [get_regression_points()]
#'
#' @examples
#' library(moderndive)
#'
#' # Fit lm() regression:
#' mpg_model <- lm(mpg ~ cyl, data = mtcars)
#'
#' # Get regression summaries:
#' get_regression_summaries(mpg_model)
get_regression_summaries <-
  function(model, digits = 3, print = FALSE) {
    # Check inputs
    input_checks(model, digits, print)
    
    # Define outcome variable
    outcome_variable <- formula(model) %>%
      lhs() %>%
      all.vars()
    
    # Use term labels for augment() columns
    term_labels <- attr(terms(model), "term.labels")
    explanatory_vars_aug <- make.names(term_labels)
    
    # Compute mean-squared error and root mean-squared error
    mse_and_rmse <- model %>%
      augment() %>%
      # Note: no explanatory vars here at all
      select(!!c(outcome_variable, ".fitted", ".resid")) %>%
      rename_at(vars(".fitted"), ~ str_c(outcome_variable, "_hat")) %>%
      rename(residual = .resid) %>%
      summarise(
        mse  = mean(residual^2, na.rm = TRUE),
        rmse = sqrt(mse)
      )
    
    # Create output tibble
    regression_summaries <- model %>%
      glance() %>%
      mutate_if(is.numeric, round, digits = digits) %>%
      select(-c(AIC, BIC, deviance, df.residual, logLik)) %>%
      as_tibble() %>%
      clean_names() %>%
      bind_cols(mse_and_rmse) %>%
      select(r_squared, adj_r_squared, mse, rmse, everything())
    
    # Transform to markdown
    if (print) {
      regression_summaries <- regression_summaries %>%
        kable()
    }
    
    return(regression_summaries)
  }


# Check input functions ----
input_checks <- function(model, digits = 3, print = FALSE, 
                         default_categorical_levels = FALSE) {
  # Since the `"glm"` class also contains the `"lm"` class
  if (length(class(model)) != 1 | !("lm" %in% class(model))) {
    stop(paste(
      "Only simple linear regression",
      "models are supported. Try again using only `lm()`",
      "models as appropriate."
    ))
  }
  check_numeric(digits)
  check_logical(print)
  check_logical(default_categorical_levels)
}

check_numeric <- function(input) {
  if (!is.numeric(input)) {
    stop("The input entry must be numeric.")
  }
}

check_logical <- function(input) {
  if (!is.logical(input)) {
    stop("The input must be logical.")
  }
}

check_character <- function(input) {
  if (!is.character(input)) {
    stop("The input must be a character.")
  }
}

check_data_frame <- function(input) {
  if (!is.data.frame(input)) {
    stop("The input must be a data frame.")
  }
}