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

Regression imputation with dplyr in R

我想用 dplyr 在 R 中高效地 进行 回归插补 。这是我的问题:我有一个数据集,其中一列有许多 缺失值 - 我们称它为 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 是价格,而 var1var2 是一些人口统计变量(如“教育”和“年龄”)。

到目前为止我做了什么:

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)                                     

由于我在实际问题中使用了超过 var1var2(大约 30 个变量),使用 dplyr 进行回归插补的更好方法是什么? (不过,我也对非 dplyr 解决方案持开放态度。)

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 助手对数据框进行子集化以 select 您的因变量(此处使用 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

基准测试

如评论中所述,对于大型数据框,按行操作比其他一些选项花费的时间要长得多:

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