How to map a nested dataframe, and store multiple columns as output

user_1819186 0
Viewed 419

I have a data structure as follows:

test <- data.frame(
   id= rep(1:3, each=20),
   count  = rnorm(60, mean=5, sd=1),
   covar1 = rnorm(60, mean=10, sd=3),
   covar2 = rnorm(60, mean=95, sd=5),
   covar3 = rnorm(60, mean=30, sd=5)
   )

Then I nest it by id:

test <- test %>% nest(-id)

I want to apply a model to each data covar column, for a given id. Then I want to store the result in a separate column. I can do this as follows:

test <- test %>% mutate(covar1_lm = map(data, ~lm(count ~ covar1, data=.x)),
                        covar2_lm = map(data, ~lm(count ~ covar2, data=.x)),
                        covar3_lm = map(data, ~lm(count ~ covar3, data=.x)))

Which gives the output I want:

> test
# A tibble: 3 x 5
     id data              covar1_lm covar2_lm covar3_lm
  <int> <list>            <list>    <list>    <list>   
1     1 <tibble [20 × 4]> <lm>      <lm>      <lm>     
2     2 <tibble [20 × 4]> <lm>      <lm>      <lm>     
3     3 <tibble [20 × 4]> <lm>      <lm>      <lm>

The problem is my real data has a large number of covar columns, and so I'd like to reduce the boilerplate code. So I'm guessing I need some concept of dynamic variable names, but I cant figure out how to map over a dynamic set of column names??

purrr r
1 Answers

You can pivot_longer() the dataset first, so that there is one observation (row) for each covariate for each dataset. Then you perform the model within each covariate.

test %>%
  pivot_longer(starts_with("covar"),
               names_to = "covariate") %>%
  group_by(id, covariate) %>%
  summarize(model = list(lm(count ~ value)))

You now have one observation for each combination of ID and covariate.

# A tibble: 9 x 3
# Groups:   id [3]
     id covariate model 
  <int> <chr>     <list>
1     1 covar1    <lm>  
2     1 covar2    <lm>  
3     1 covar3    <lm>  
4     2 covar1    <lm>  
5     2 covar2    <lm>  
6     2 covar3    <lm>  
7     3 covar1    <lm>  
8     3 covar2    <lm>  
9     3 covar3    <lm>

If you want to turn that into the same kind of result, you could pipe this to pivot_wider(names_from = covariate, values_from = model). (But note that having one row for each model could make it easier to explore and visualize the models, especially if you tidy each with broom::tidy() and unnested them).

An alternative to the group_by()/summarize() above would be to nest them :

test %>%
  pivot_longer(starts_with("covar"),
               names_to = "covariate") %>%
  group_by(id, covariate) %>%
  nest() %>%
  mutate(model = map(data, ~ lm(count ~ value, data = .x)))
user_7126030
Related