如何用modelsummary自动标注回归模型中控制变量组的纳入情况

核心思路 #

通过提取模型的自变量列表，结合预定义的控制变量组，自动判断每个模型是否纳入对应控制组，生成标注行后用add_rows添加到表格中，无需手动逐行编写。

完整示例代码 #

import pandas as pd
import statsmodels.api as sm
from modelsummary import summary

# 加载对应Stata示例的数据集
url = "https://www.stata-press.com/data/r17/nlsw88b.dta"
df = pd.read_stata(url)

# 拟合目标模型
model1 = sm.OLS.from_formula("wage ~ occupation", data=df).fit()
model2 = sm.OLS.from_formula("wage ~ occupation + age + race + married", data=df).fit()
model3 = sm.OLS.from_formula("wage ~ occupation + age + race + married + industry + union", data=df).fit()
models = [model1, model2, model3]

# 定义控制变量组（对应Stata的indicate参数逻辑）
control_groups = {
    "Individual controls": ["age", "race", "married"],
    "Industry controls": ["industry", "union"]
}

# 自动生成标注行
add_rows_list = []
for group_name, vars_list in control_groups.items():
    # 检查每个模型是否包含该组所有变量，可根据需求替换all()为any()
    row_values = ["Yes" if all(var in model.model.exog_names for var in vars_list) else "No" for model in models]
    add_rows_list.append([group_name] + row_values)

# 生成带标注的表格
summary(
    models,
    add_rows=add_rows_list,
    notes=["Control variable inclusion status: Yes = all variables in group are included"],
    title="Wage Regression Models"
)

关键调整说明 #

• 控制组可自由扩展：新增控制组只需在control_groups字典中添加键值对
• 判断逻辑灵活切换：若需组内有任意变量纳入即标Yes，将代码中的all()替换为any()即可
• 适配多模型场景：新增模型时无需修改标注行生成代码，会自动匹配模型数量

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