Объединить две модели предсказания регрессии (с подмножествами кадра данных) обратно в кадр данных (один столбец)

Question

Я строил поверх аналогичного вопроса, заданного и ответившего на SO год назад. Это относится к этому сообщению: how to merge two linear regression prediction models (each per data frame's subset) into one column of the data frame

Я буду использовать те же данные, что и там, но с новым столбцом. Я создаю данные:

dat = read.table(text = " cats birds wolfs  snakes  trees 
0  3  8   7  2 
1  3  8   7  3 
1  1  2   3  2 
0  1  2   3  1 
0  1  2   3  2 
1  6  1   1  3 
0  6  1   1  1 
1  6  1   1  1 " ,header = TRUE) 

Модель количество волков, используя два подмножества данных, чтобы различать условия. Уравнения для каждого подмножества различны.

f0 = lm(wolfs~snakes,data = dat,subset=dat$cats==0) 
f1 = lm(wolfs~snakes + trees,data = dat,subset=dat$cats==1) 

Предскажите количество волков для каждого подмножества.

f0_predict = predict(f0,data = dat,subset=dat$cats==1,type='response') 
f1_predict = predict(f1,data = dat,subset=dat$cats==0,type='response') 

Затем (опять же, за сообщение 2015 года) Я разделил данные на переменную кошек.

dat.l = split(dat, dat$cats) 
dat.l 

... Вот где это становится немного сложнее. Повестка дня 2015 года предложила использовать lapply для присоединения двух наборов предсказаний к набору данных. Но здесь функция респондента не сработала, так как предполагалось, что оба уравнения регрессии были в основном одинаковыми. Вот моя попытка (это близко к оригиналу, просто переделаны):

dat.l = lapply(dat.l, function(x){ 
mod = 

ifelse(dat$cats==0,lm(wolfs~snakes,data=x),lm(wolfs~snakes+trees,data=x)) 
       x$full_prediction = predict(mod,data=x,type='response') 
       return(x) 
    }) 
    unsplit(dat.l, dat$cats) 

Любые идеи относительно последних нескольких шагов? Я до сих пор относительно новичок в S.O. и являюсь промежуточным звеном с R, поэтому, пожалуйста, пойдите осторожно, если я не разместил точно, как предпочитает сообщество.

Answer 1

Вот dplyr решение, здание от предыдущего поста вы цитируемой:

library(dplyr) 

# create a new column defining the lm formula for each level of cats 
dat <- dat %>% mutate(formula = ifelse(cats==0, "wolfs ~ snakes", 
     "wolfs ~ snakes + trees")) 

# build model and find predicted values for each value of cats 
dat <- dat %>% group_by(cats) %>% 
    do({ 
     mod <- lm(as.formula(.$formula[1]), data = .) 
     pred <- predict(mod) 
     data.frame(., pred) 
    }) 

> dat 
Source: local data frame [8 x 7] 
Groups: cats [2] 
    cats birds wolfs snakes trees    formula  pred 
    (int) (int) (int) (int) (int)     (chr)  (dbl) 
1  0  3  8  7  2   wolfs ~ snakes 7.5789474 
2  0  1  2  3  1   wolfs ~ snakes 2.6315789 
3  0  1  2  3  2   wolfs ~ snakes 2.6315789 
4  0  6  1  1  1   wolfs ~ snakes 0.1578947 
5  1  3  8  7  3 wolfs ~ snakes + trees 7.6800000 
6  1  1  2  3  2 wolfs ~ snakes + trees 2.9600000 
7  1  6  1  1  3 wolfs ~ snakes + trees 0.8400000 
8  1  6  1  1  1 wolfs ~ snakes + trees 0.5200000