在purrr嵌套回归中动态创建treatment_var变量的实现方法

场景说明 #

现有两个数据框：df是回归分析的数据集，regression_models是待执行的回归任务列表，需要根据regression_models每行的arm1和arm2，动态生成treatment_var变量，再运行对应的逻辑回归分析。

原始数据 #

library(tidyr)

set.seed(1)
df <- data.frame(x1 = rnorm(100, 0, 1),
                 x2 = rnorm(100, 0, 1),
                 y = sample(0:1, 100, replace = TRUE), 
                 group = sample(c("control", "treatment1", "treatment2", "treatment3"), 100, replace = TRUE),
                 var = 1) |>
  pivot_wider(names_from = "group", values_from = "var")

regression_models <- data.frame(outcome = c("y", "y", "y", "y", "y", "y"),
                                arm1 = c("treatment1", "treatment2", "treatment3", "treatment1", "treatment1", "treatment2"),
                                arm2 = c("control", "control", "control", "treatment2", "treatment3", "treatment3"),
                                covariates = c("x1 + x2", "x1 + x2", "x1 + x2", "x1 + x2", "x1 + x2", "x1 + x2"))

解决方案 #

使用purrr::pmap_dfr遍历每个回归任务，动态生成分组变量并执行回归，最后合并所有结果：

library(dplyr)
library(purrr)
library(broom)

regression_output <- regression_models |>
  pmap_dfr(function(outcome, arm1, arm2, covariates) {
    # 动态生成treatment_var，仅保留当前对比组的观测
    analysis_df <- df |>
      mutate(treatment_var = case_when(
        .data[[arm1]] == 1 ~ TRUE,
        .data[[arm2]] == 1 ~ FALSE,
        TRUE ~ NA
      )) |>
      filter(!is.na(treatment_var))
    
    # 构建回归公式
    formula <- as.formula(paste(outcome, "~ treatment_var +", covariates))
    
    # 运行逻辑回归并整理结果，添加任务标识
    glm(formula, family = binomial, data = analysis_df) |>
      tidy() |>
      mutate(outcome = outcome, arm1 = arm1, arm2 = arm2, .before = 1)
  })

# 查看结果
head(regression_output)

关键细节 #

• pmap_dfr会遍历regression_models的每一行，自动将每行的列值作为参数传递给匿名函数，最终把所有回归结果合并成一个数据框
• .data[[arm1]]是动态引用列名的方式，能根据arm1和arm2的字符串值提取df中的对应列，实现分组变量的动态生成
• 过滤掉treatment_var为NA的行，确保回归仅针对当前任务指定的两组观测
• 在整理结果时添加outcome、arm1、arm2列，方便后续区分不同的回归任务

内容的提问来源于stack exchange，提问作者Rob F