Search code examples
rworkflownon-linear-regressionpiecewise

How to predict the outcome of a regression holding regressor constant?


Hi everyone based on the wage-dataset (wage being the dependent variable) and on the workflow created below, I would like to find out the following:

  1. What is the predicted wage of a person with age equal to 30 for each piecewise model?
  2. Considering the flexible pw6_wf_fit model configuration and in particular the six breakpoints above: Exceeding which (approximate) value of age correlates strongest with wage?

I tried to use versions of extract but so far I don´t know how to apply it in R. Helpful for any comment

The code I use is the following:

if (!require("pacman")) install.packages("pacman")
# load (or install if pacman cannot find an existing installation) the relevant packages
pacman::p_load(
  tidyverse, tidymodels, ISLR, patchwork,
  rpart, rpart.plot, randomForest, gbm, kernlab, parsnip, skimr
)
data(Wage, package = "ISLR")

Wage %>% 
  tibble::as_tibble() %>% 
  skimr::skim()
lin_rec   <- recipe(wage ~ age, data = Wage)
# Specify as linear regression
lm_spec <- 
  linear_reg() %>%
  set_mode("regression") %>%
  set_engine("lm")
plot_model <- function(wf_fit, data) {
  
  predictions <- 
    tibble::tibble(age = seq(min(data$age), max(data$age))) %>% 
    dplyr::bind_cols(
      predict(wf_fit, new_data = .),
      predict(wf_fit, new_data = ., type = "conf_int")
    )
    
  p <- ggplot2::ggplot(aes(age, wage), data = data) +
    geom_point(alpha = 0.05) +
    geom_line(aes(y = .pred),
              data = predictions, color = "darkgreen") +
    geom_line(aes(y = .pred_lower),
              data = predictions, linetype = "dashed", color = "blue") +
    geom_line(aes(y = .pred_upper),
              data = predictions, linetype = "dashed", color = "blue") +
    scale_x_continuous(breaks = seq(20, 80, 5)) +
    labs(title = substitute(wf_fit)) +
    theme_classic()
  
  return(p)
  
}

pw3_rec <- lin_rec %>% step_discretize(age, num_breaks = 3, min_unique = 5) 
pw4_rec <- lin_rec %>% step_discretize(age, num_breaks = 4, min_unique = 5) 
pw5_rec <- lin_rec %>% step_discretize(age, num_breaks = 5, min_unique = 5) 
pw6_rec <- lin_rec %>% step_discretize(age, num_breaks = 6, min_unique = 5)
pw3_wf_fit <- workflow(pw3_rec, lm_spec) %>% fit(data = Wage)
pw4_wf_fit <- workflow(pw4_rec, lm_spec) %>% fit(data = Wage)
pw5_wf_fit <- workflow(pw5_rec, lm_spec) %>% fit(data = Wage)
pw6_wf_fit <- workflow(pw6_rec, lm_spec) %>% fit(data = Wage)
(plot_model(pw3_wf_fit, Wage) + plot_model(pw4_wf_fit, Wage)) /
(plot_model(pw5_wf_fit, Wage) + plot_model(pw6_wf_fit, Wage)) 

Solution

  • The answer to the first question is pretty straightforward:

    map(list(pw3_wf_fit, pw4_wf_fit, pw5_wf_fit, pw6_wf_fit), 
        ~predict(.x, new_data=tibble(age=30))) %>% 
      bind_rows()
    
    # # A tibble: 4 × 1
    #   .pred
    #   <dbl>
    # 1  99.3
    # 2  94.2
    # 3  92.3
    # 4  89.5