在Python中为多数据集拟合相同模型（对应R嵌套流程）

报错根源 #

核心问题：Polars嵌套列data的类型为List(Struct)，statsmodels依赖的patsy公式解析器无法直接从该结构中提取字段（如recipe），会将整个嵌套结构视为单个Series元素，导致无法通过字段名访问数据。

分步解决方案 #

1. 编写自定义拟合函数 #

定义函数将嵌套子数据集转换为可被statsmodels识别的格式，拟合二项式GLM并提取recipeB的系数：

import polars as pl
import statsmodels.api as sm
import statsmodels.formula.api as smf

def extract_recipeB_coef(sub_data: list[dict]) -> float | None:
    # 将List(Struct)转换为Polars DataFrame
    sub_df = pl.DataFrame(sub_data)
    try:
        # 拟合二项式GLM，转换为Pandas DataFrame确保兼容性
        model = smf.glm(
            formula="events ~ recipe",
            data=sub_df.to_pandas(),
            family=sm.families.Binomial()
        ).fit()
        # 提取recipe[T.B]的系数，无对应项则返回None
        return model.params.get("recipe[T.B]", None)
    except Exception:
        # 处理子数据集样本不足、类别缺失等拟合失败情况
        return None

2. 在Polars中批量应用函数 #

使用map_elements对嵌套列data批量执行拟合逻辑：

# 假设你的嵌套数据集为df_nested，结构包含test_id和data列
df_result = df_nested.with_columns(
    pl.col("data").map_elements(extract_recipeB_coef).alias("recipeB_coef")
)

3. 性能优化替代方案（展开分组拟合） #

若嵌套数据量较大，可先展开嵌套结构再分组拟合，避免map_elements的单元素处理开销：

# 展开嵌套的data列并解析Struct字段
df_exploded = df_nested.explode("data").unnest("data")

# 按test_id分组拟合模型的Pandas兼容写法
def group_fit(group):
    model = smf.glm(
        "events ~ recipe",
        data=group,
        family=sm.families.Binomial()
    ).fit()
    return model.params.get("recipe[T.B]", None)

# 分组拟合后转回Polars
df_result = pl.from_pandas(
    df_exploded.to_pandas().groupby("test_id").apply(group_fit).reset_index(name="recipeB_coef")
)

关键注意事项 #

• 必须转换嵌套结构：List(Struct)无法直接被patsy解析，转换为DataFrame后才能让公式识别events、recipe等字段。
• 兼容性处理：部分statsmodels版本对Polars DataFrame支持有限，转换为Pandas DataFrame可避免兼容性问题。
• 异常捕获：子数据集可能存在样本量不足、类别缺失等问题，加入异常捕获可保证整个流程不中断。

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