在 R 中使用 dplyr 进行回归插补

Question

我想高效地在R中使用dplyr进行回归插补。这是我的问题：我有一个数据集，其中一列有许多缺失值 - 我们称之为p。现在我想用回归插补方法估计p 的缺失值。为此，我使用未经审查的数据（p 没有缺失值的数据集的子集）使用 OLS 对一组变量进行回归 p。然后我用估计的系数来计算p的缺失值。

我的数据集是这样的：

df = data.frame(
  id = c(1, 1, 1, 2, 2, 2),
  group = c(1, 1, 2, 1, 1, 2),
  sub_group = c(1, 2, 3, 1, 2, 3),
  p = c(4.3, 5.7, NA, NA, NA, 10),
  var1 = c(0.3, 0.1, 0.4, 0.9, 0.1, 0.2),
  var2 = c(0, 0, 0, 1, 1, 1)
)

其中id 代表个人，他们从group（例如“食物”）和subgroups（例如“面包”）购买商品。 p 是价格，而 var1 和 var2 是一些人口统计变量（如“教育”和“年龄”）。

到目前为止我做了什么：

library(dplyr)

df <- as_tibble(df)

# Create uncensored data
uncensored_df <- df %>%
filter(!is.na(p))

# Run regression on uncensored data
imp_model <- lm(p ~ var1 + var2, data = uncensored_df)

# Get the coefficients of the fitted model
coefs <- unname(imp_model$coefficients)

# Use coefficients to compute missing values of p
censored_df <-df %>%
filter(is.na(p)) %>%
group_by(id, group, sub_group) %>%
  mutate(p = coefs[1] + coefs[2] * var1 + coefs[3] * var2)  

# And finally combine the two subsets                                 
bind_rows(uncensored_df, censored_df) %>% arrange(id, group, sub_group)                                     

由于我在实际问题中使用的不仅仅是 var1 和 var2（大约 30 个变量），使用 dplyr 进行回归插补的更好方法是什么？ （不过，我也对非 dplyr 解决方案持开放态度。）

Answer 1

library(dplyr)

fit <- lm(p ~ ., data = select(df, p, starts_with("var")))


df %>% 
  rowwise() %>% 
  mutate(p = ifelse(is.na(p), predict(fit, newdata = across()), p)) %>% 
  ungroup()

工作原理

• 对于初学者，在拟合模型时，您可以使用select 和任何tidyselect 助手来选择您的因变量（此处使用starts_with("var")）来对数据框进行子集化。然后，此子集数据框允许您使用 ~ . 表示法，这意味着在子集数据框中的所有内容上回归 p。
• 接下来，您将创建一个逐行数据框，并使用您的模型来预测p 的缺失位置。在这种情况下，across 将每一行转换为一个 1x6 的 tibble，您可以将其传递给 newdata 参数。 predict 然后使用模型拟合和这个新数据来预测 p 的值。

输出

     id group sub_group     p  var1  var2
  <dbl> <dbl>     <dbl> <dbl> <dbl> <dbl>
1     1     1         1  4.3    0.3     0
2     1     1         2  5.7    0.1     0
3     1     2         3  3.60   0.4     0
4     2     1         1  5.10   0.9     1
5     2     1         2 10.7    0.1     1
6     2     2         3 10      0.2     1

---

基准测试

如 cmets 中所述，对于大型数据帧，按行操作所需的时间明显长于其他一些选项：

library(microbenchmark)

set.seed(1)
df1 <- df %>%
  slice_sample(n = 1E5, replace = T)

fit <- lm(p ~ ., data = select(df1, p, starts_with("var")))


dplyr_rowwise <- function(){
  df1 %>% 
    rowwise() %>% 
    mutate(p = ifelse(is.na(p), predict(fit, newdata = across()), p)) %>% 
    ungroup()
}

dplyr_coalesce <- function(){
  df1 %>%
    mutate(p = coalesce(p, predict(fit, newdata = df1)))
}

base_index <- function(){
  isna <- is.na(df1$p)
  df1$p[isna] <- predict(fit, newdata = subset(df1, isna))
}

microbenchmark(
  dplyr_rowwise(),
  dplyr_coalesce(),
  base_index(),
  times = 10L
)

Unit: milliseconds
             expr        min         lq        mean      median         uq  
  dplyr_rowwise() 63739.9512 64441.0800 66926.46041 65513.51785 66923.0241
 dplyr_coalesce()     6.5901     6.9037     8.55971     7.21125     7.7157
     base_index()    13.0368    13.1790    15.73682    13.53310    19.3004