Search code examples
rregressionlinear-regressionpurrrsample

Tag regression samples for many models using purrr


I want to tag the samples used in multiple regression models with a purrr function. Drawing on this Q&A, I can accomplish this goal on an ad hoc basis as follows:

library(dplyr)

df <- mtcars %>% 
  mutate(disp = replace(hp, c(2, 3), NA)) %>% 
  mutate(wt = replace(wt, c(3, 4, 5), NA))

s1 <- lm(mpg ~ disp, data = df)
df$samp1 <- TRUE
df$samp1[na.action(s1)] <- FALSE         

s2 <- lm(mpg ~ wt, data = df)
df$samp2 <- TRUE
df$samp2[na.action(s2)] <- FALSE

How can I add samp1 and samp2 to df using purrr?


Solution

  • I'm not quite there yet but here's a tidy approach using a custom function:

    flag_use <- function(df, model, name) {
      mutate(df, {{name}} := !row_number() %in% na.action( {{model}} ))
    }
    
    df %>%
      flag_use(s1, "samp1") %>%
      flag_use(s2, "samp2")