pandas - python patsy中的分组因子

Question

这是这个较旧的 SO 问题的扩展，但适用于 python 而不是 R。我也认为这个解决方案不是最好的。

假设我有看起来像这样的数据......

State   Y
AL      5
AK      10
AZ      8

我想编写一个 patsy 公式将 State 转换为 Region，然后使用 statsmodels 使用 Region 进行预测。所以桌子看起来像......

State   Region    Y
AL      Southeast 5
AK      Northwest 10
AZ      Southwest 8

我想有一个功能

model = sm.OLS('Y ~ C(State, StateToRegionGrouping)').fit()

我认为有两种方法。首先，在原始数据上添加一个查找列或编写一个分类转换器函数供 patsy 处理。

哪种方式更好，如果 patsy 分类转换器更好，那么编程它的好方法是什么？

Answer 1

保持简单，只需使用字典映射：

import statsmodels.formula.api as smf

mapping = {'AL': 'Southeast',
           'AK': 'Northwest',
           'AZ': 'Southwest'}

df = pd.DataFrame({'State': ['AL', 'AK', 'AZ'], 'Y': [5, 10, 8]})
df['Region'] = df.State.map(mapping)

>>> df
  State   Y     Region
0    AL   5  Southeast
1    AK  10  Northwest
2    AZ   8  Southwest

model = smf.ols('Y ~ Region', data=df).fit()