python patsy中的分组因子

Question

这是this older SO question的一个扩展，但不是R，而是python。我也认为这个解决方案不是最好的。

假设我有这样的数据...

State   Y
AL      5
AK      10
AZ      8

我想写一个patsy公式来将State转换为Region，然后使用statsmodel来使用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()