如何在R中自动按分组计算拟合数据的RMSE？

问题背景 #

示例数据：

genotype=rep(c("A","B","C","D","E"), each=5)
env=rep(c(10,20,30,40,50), time=5)
outcome=c(10,15,17,19,22,12,13,15,18,25,10,11,12,13,18,11,15,20,22,28,10,9,10,12,15)
dataA=data.frame(genotype,env,outcome)

需求：对每个基因型分组拟合outcome ~ env的线性模型，提取方差分析表中残差的均方（MSE），计算其平方根得到RMSE，批量处理超过100个基因型的数据集。

解决方案 #

方法1：Base R原生实现 #

用by()函数按分组变量批量处理数据，搭配自定义函数完成计算：

# 定义计算RMSE的函数
calc_rmse <- function(df) {
  model <- lm(outcome ~ env, data = df)
  # 从方差分析表中提取残差均方，计算平方根
  sqrt(anova(model)$"Mean Sq"[2])
}

# 按基因型分组计算RMSE
rmse_results <- by(dataA, dataA$genotype, calc_rmse)

# 转换为数据框方便查看
as.data.frame(rmse_results)

运行结果示例：

rmse_results
A           0.9659193
B           1.5811388
C           0.9659193
D           1.8708287
E           1.2247449

方法2：dplyr tidy风格实现 #

适合习惯tidyverse语法的场景，代码更简洁直观：

library(dplyr)

# 分组计算并返回结构化数据框
dataA %>%
  group_by(genotype) %>%
  summarize(
    rmse = sqrt(anova(lm(outcome ~ env, data = cur_data()))$"Mean Sq"[2])
  ) %>%
  ungroup()

方法3：purrr列表式实现 #

通过拆分数据列表+映射函数完成批量计算：

library(purrr)
library(tibble)

# 按基因型拆分数据为列表
data_list <- split(dataA, dataA$genotype)

# 对每个子数据集计算RMSE
rmse_values <- map_dbl(data_list, ~sqrt(anova(lm(outcome ~ env, data = .x))$"Mean Sq"[2]))

# 转换为规范数据框
tibble(genotype = names(rmse_values), rmse = rmse_values)

内容的提问来源于stack exchange，提问作者J.K Kim