Iterating over multiple regression models and data subsets in R

问题

I am trying to learn how to automate running 3 or more regression models over subsets of a dataset using the purrr and broom packages in R. I am doing this with the nest %>% mutate(map()) %>% unnest() flow in mind.

I am able to replicate examples online when there is only one regression model that is applied to several data subsets. However, I am running into problems when I have more than one regression model in my function.

What I tried to do

library(tidyverse)
library(broom)

estimate_model <- function(df) {
  model1 <- lm(mpg ~ wt, data = df)
  model2 <- lm(mpg ~ wt + gear, data = df)
  model3 <- lm(mpg ~ wt + gear + vs, data = df)
}

ols_1dep_3specs <- mtcars %>%
    nest(-cyl) %>%
    mutate(
       estimates = map(data, estimate_model), # want to run several models at once
       coef_wt = map(estimate,  ~pluck(coef(.), "wt")), # coefficient of wt only
       se_wt = map(estimate, ~pluck(tidy(.), "std.error")[[2]]), # se of wt only
       rsq = map(model, ~pluck(glance(.), "r.squared")),
       arsq = map(model, ~pluck(glance(.), "adj.r.squared"))
    ) %>%
    unnest(coef_wt, se_wt, rsq, arsq)

ols_1dep_3specs

Unfortunately, this seems to only work when the function estimate_model only contains one regression model. Any advice on how one would go about writing code when there are several specifications? Open to suggestions outside the nest() %>% mutate(map()) %>% nest() framework.

The following code sort of gets at what I am hoping to achieve but it involves a lot of repetition.

estimate_model1 <- function(df) {
  model1 <- lm(mpg ~ wt, data = df)
}
estimate_model2 <- function(df) {
  model2 <- lm(mpg ~ wt + gear, data = df)
}
estimate_model3 <- function(df) {
  model3 <- lm(mpg ~ wt + gear + vs, data = df)
}

ols_1dep_3specs <- mtcars %>%
  nest(-cyl) %>%
  mutate(model_1 = map(data, estimate_model1),
         model_2 = map(data, estimate_model2),
         model_3 = map(data, estimate_model3)) %>%
  mutate(coef_wt_1 = map_dbl(model_1, ~pluck(coef(.), "wt")),
         coef_wt_2 = map_dbl(model_2, ~pluck(coef(.), "wt")),
         coef_wt_3 = map_dbl(model_3, ~pluck(coef(.), "wt")),
         rsq_1 = map_dbl(model_1, ~pluck(glance(.), "r.squared")),
         rsq_2 = map_dbl(model_2, ~pluck(glance(.), "r.squared")),
         rsq_3 = map_dbl(model_3, ~pluck(glance(.), "r.squared"))) %>% 
  dplyr::select(starts_with("coef_wt"), starts_with("rsq"))

回答1:

In the function, there is no return call, it would be better to place all the models in a list

estimate_model <- function(df) {
        model1 <- lm(mpg ~ wt, data = df)
        model2 <- lm(mpg ~ wt + gear, data = df)
        model3 <- lm(mpg ~ wt + gear + vs, data = df)
        list(model1, model2, model3)
      }

and then apply the first piece of code by looping over each list element

mtcars %>% 
  group_by(cyl) %>%
  nest() %>% 
  mutate(estimates = map(data, estimate_model),
         coef_wt = map(estimates,  ~map_dbl(.x, ~ pluck(coef(.x), "wt"))),
         se_wt = map(estimates, ~map_dbl(.x, ~pluck(tidy(.x), "std.error")[[2]])), 
         rsq = map(estimates, ~ map_dbl(.x, ~pluck(glance(.x), "r.squared"))),
          arsq = map(estimates, ~map_dbl(.x, ~ pluck(glance(.x), "adj.r.squared")))) %>%
  unnest(c(coef_wt, se_wt, rsq, arsq))
# A tibble: 9 x 7
# Groups:   cyl [3]
#    cyl            data estimates  coef_wt se_wt   rsq  arsq
#  <dbl> <list<df[,10]>> <list>       <dbl> <dbl> <dbl> <dbl>
#1     6        [7 × 10] <list [3]>   -2.78 1.33  0.465 0.357
#2     6        [7 × 10] <list [3]>   -3.92 1.41  0.660 0.489
#3     6        [7 × 10] <list [3]>   -6.19 4.49  0.690 0.379
#4     4       [11 × 10] <list [3]>   -5.65 1.85  0.509 0.454
#5     4       [11 × 10] <list [3]>   -5.38 2.08  0.517 0.396
#6     4       [11 × 10] <list [3]>   -5.13 2.16  0.555 0.364
#7     8       [14 × 10] <list [3]>   -2.19 0.739 0.423 0.375
#8     8       [14 × 10] <list [3]>   -2.43 0.798 0.459 0.361
#9     8       [14 × 10] <list [3]>   -2.43 0.798 0.459 0.361

回答2:

Using purrr::lst automatically names that list after its elements, which helps you keep track of your models later. After applying the function to your nested data, you can pull out a column of the model names.

I chose to replace the plucking with unnesting earlier in the workflow and using 2 map calls to get values out as vectors rather than lists. It's just a preference, but I have an easier time when columns are less-deeply nested.

library(tidyverse)
library(broom)

estimate_model <- function(df) {
  model1 <- lm(mpg ~ wt, data = df)
  model2 <- lm(mpg ~ wt + gear, data = df)
  model3 <- lm(mpg ~ wt + gear + vs, data = df)
  lst(model1, model2, model3)
}

mtcars %>%
  group_by(cyl) %>%
  nest() %>%
  mutate(mods = map(data, estimate_model),
         mod_id = map(mods, names)) %>%
  unnest(c(mod_id, mods)) %>%
  mutate(coef_wt = map(mods, coef) %>% map_dbl("wt"),
         se_wt = map(mods, tidy) %>% map("std.error") %>% .[[2]],
         rsq = map(mods, glance) %>% map_dbl("r.squared"),
         arsq = map(mods, glance) %>% map_dbl("adj.r.squared"))
#> # A tibble: 9 x 8
#> # Groups:   cyl [3]
#>     cyl            data mods   mod_id coef_wt  se_wt   rsq  arsq
#>   <dbl> <list<df[,10]>> <list> <chr>    <dbl>  <dbl> <dbl> <dbl>
#> 1     6        [7 × 10] <lm>   model1   -2.78  6.36  0.465 0.357
#> 2     6        [7 × 10] <lm>   model2   -3.92  1.41  0.660 0.489
#> 3     6        [7 × 10] <lm>   model3   -6.19  0.727 0.690 0.379
#> 4     4       [11 × 10] <lm>   model1   -5.65 11.6   0.509 0.454
#> 5     4       [11 × 10] <lm>   model2   -5.38  2.08  0.517 0.396
#> 6     4       [11 × 10] <lm>   model3   -5.13  2.20  0.555 0.364
#> 7     8       [14 × 10] <lm>   model1   -2.19  4.91  0.423 0.375
#> 8     8       [14 × 10] <lm>   model2   -2.43  0.798 0.459 0.361
#> 9     8       [14 × 10] <lm>   model3   -2.43  0.835 0.459 0.361

来源：https://stackoverflow.com/questions/59569458/iterating-over-multiple-regression-models-and-data-subsets-in-r

标签

regression

tidyverse

purrr

broom