Skip to content

Instantly share code, notes, and snippets.

@andrewheiss
Created July 17, 2023 23:11
Show Gist options
  • Save andrewheiss/ade01b8557dce1a67ceaa87de66c8b75 to your computer and use it in GitHub Desktop.
Save andrewheiss/ade01b8557dce1a67ceaa87de66c8b75 to your computer and use it in GitHub Desktop.
library(tidyverse)
library(broom)

# Model with a cateogorical predictor
example_model <- lm(hwy ~ displ + drv, data = mpg)

# Extract all the right-hand variables
rhs <- all.vars(stats::update(formula(example_model), 0 ~ .))
rhs
#> [1] "displ" "drv"

# Make a dataframe of all the levels of all the non-numeric things
model_variable_levels <- tibble(
  variable = rhs
) %>% 
  mutate(levels = map(variable, ~{
    x <- mpg[[.x]]
    if (is.numeric(x)) {
      ""
    } else if (is.factor(x)) {
      levels(x)
    } else {
      sort(unique(x))
    }
  })) %>% 
  unnest(levels) %>% 
  mutate(term = paste0(variable, levels))
model_variable_levels
#> # A tibble: 4 × 3
#>   variable levels term 
#>   <chr>    <chr>  <chr>
#> 1 displ    ""     displ
#> 2 drv      "4"    drv4 
#> 3 drv      "f"    drvf 
#> 4 drv      "r"    drvr

# Extract model results
model_results <- tidy(example_model, conf.int = TRUE)

# Combine full dataset of factor levels with model results
thing_to_plot <- model_variable_levels %>% 
  left_join(model_results, by = join_by(term)) %>% 
  mutate(levels = fct_inorder(levels)) %>% 
  # Make these zero
  mutate(
    across(c(estimate, conf.low, conf.high), 
           ~ifelse(is.na(.x), 0, .x)))
thing_to_plot
#> # A tibble: 4 × 9
#>   variable levels term  estimate std.error statistic   p.value conf.low
#>   <chr>    <fct>  <chr>    <dbl>     <dbl>     <dbl>     <dbl>    <dbl>
#> 1 displ    ""     displ    -2.91     0.218    -13.4   1.73e-30    -3.34
#> 2 drv      "4"    drv4      0       NA         NA    NA            0   
#> 3 drv      "f"    drvf      4.79     0.530      9.05  6.40e-17     3.75
#> 4 drv      "r"    drvr      5.26     0.734      7.17  1.03e-11     3.81
#> # ℹ 1 more variable: conf.high <dbl>

# Make coefficient plot
ggplot(thing_to_plot, aes(x = estimate, y = fct_rev(term), color = variable)) +
  geom_pointrange(aes(xmin = conf.low, xmax = conf.high)) +
  guides(color = "none")

Created on 2023-07-17 with reprex v2.0.2

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment