Как использовать функцию прогнозирования для линейных моделей, хранящихся в столбцах в виде столбцов

Question

У меня есть код, похожий на этот:

library(tidyverse)
library(broom)

df <- tribble(~group, ~set, ~x, ~y,
  0, "training", 1, 2,
  0, "training", 2, 4,
  0, "training", 4, 9,
  0, "training", 5, 9,
  0, "training", 15, 29,
  0, "training", 20, 42,
  0, "training", 21, 42,
  0, "training", 25, 49,
  0, "test", 10, NA,
)
df %>%
  filter(set == "training") %>%
  group_by(group) %>%
  do(pm = lm(y ~ x, data = .)) %>%
  left_join(df, ., by = "group")

Итак, я получаю что-то вроде этого:

  group      set    x     y     model   
  <dbl>  <chr>    <dbl> <dbl>  <list>  
1    0. training    1.    2. <S3: lm>
2    0. training    2.    4. <S3: lm>
3    0. training    4.    9. <S3: lm>
4    0. training    5.    9. <S3: lm>
5    0. training   15.   29. <S3: lm>
6    0. training   20.   42. <S3: lm>
7    0. training   21.   42. <S3: lm>
8    0. training   25.   49. <S3: lm>
9    0. test       10.   NA  <S3: lm>

Как использовать линейную модель в 9-й строке, где set равно «test» для прогнозирования y из x ? Заранее благодарю за помощь

Answer 1

Мы можем replace NA элементов в 'y' с выводом predict на основе модели lm в той же строке после группировки по 'group'

library(dplyr)
df %>%
  filter(set == "training") %>%
  group_by(group) %>%
  do(pm = lm(y ~ x, data = .)) %>%
  left_join(df, ., by = "group") %>%
  group_by(group) %>% 
  mutate(y = replace(y, is.na(y), predict(pm[[1]],  
        newdata = data.frame(x = x[set == "test"]))))
# A tibble: 9 x 5
# Groups:   group [1]
#  group set          x     y pm    
#  <dbl> <chr>    <dbl> <dbl> <list>
#1     0 training     1   2   <lm>  
#2     0 training     2   4   <lm>  
#3     0 training     4   9   <lm>  
#4     0 training     5   9   <lm>  
#5     0 training    15  29   <lm>  
#6     0 training    20  42   <lm>  
#7     0 training    21  42   <lm>  
#8     0 training    25  49   <lm>  
#9     0 test        10  20.0 <lm>