按数据帧在Python组上执行数学运算

我有一个具有以下结构的熊猫数据帧：

In [1]: df
Out[1]: 
    location_code    month    amount
0    1               1        10
1    1               2        11
2    1               3        12
3    1               4        13
4    1               5        14
5    1               6        15
6    2               1        23
7    2               2        25
8    2               3        27
9    2               4        29
10   2               5        31
11   2               6        33

我还有一个具有以下内容的数据帧：

In [2]: output_df
Out[2]: 
    location_code    regression_coef
0   1                None
1   2                None

我想要什么：

output_df = df.groupby('location_code')[amount].apply(linear_regression_and_return_coefficient)

我想按

位置代码

分组，然后对

数量

的值进行线性回归，并存储系数。我尝试了以下代码：

import pandas as pd
import statsmodels.api as sm
import numpy as np

gb = df.groupby('location_code')['amount']

x = []
for j in range(6): x.append(j+1)

for location_code, amount in gb:
    trans = amount.tolist()
    x = sm.add_constant(x)
    model = sm.OLS(trans, x)
    results = model.fit()
    output_df['regression_coef'][merchant_location_code] = results.params[1]/np.mean(trans)

这段代码可以工作，但是我的数据集有点大（大约5 gb），有点复杂，这需要花费很长时间。我想知道是否有一个矢量化的操作可以更有效地做到这一点？我知道在数据帧上使用循环是不好的

解决方案

经过一些修改，我编写了一个函数，可以与

groupby

上的

apply

方法一起使用

def get_lin_reg_coef(series):
    x=sm.add_constant(range(1,7))
    result = sm.OLS(series, x).fit().params[1]
    return result/series.mean()

gb = df.groupby('location_code')['amount']

output_df['lin_reg_coef'] = gb.apply(get_lin_reg_coef)

与我以前使用的迭代解决方案（输入大小不同）进行基准测试得到：

DataFrame Rows    Iterative Solution (sec)    Vectorized Solution (sec)
       370,000    81.42                       84.46
     1,850,000    448.36                      365.66
     3,700,000    1282.83                     715.89
     7,400,000    5034.62                     1407.88         

---

显然，随着数据集大小的增长，速度会快得多

在不了解更多数据、记录数等信息的情况下，此代码应该运行得更快：

import pandas as pd
import statsmodels.api as sm
import numpy as np

gb = df.groupby('location_code')['amount']

x = sm.add_constant(range(1,7))

def fit(stuff):
    return sm.OLS(stuff["amount"], x).fit().params[1] / stuff["amount"].mean()

output = gb.apply(fit)

---

交易金额

未定义。这应该是

amount

？很好的捕获，是的，修复了这个问题。

x

也是未定义的。我们是否应该假设这是一个常数（关于

1

？是的，我忘记了

x

。这就是我的代码中的方式。你不需要调用

tolist

，只要用

trans

替换

直行

——无需强制转换为列表。这将节省复制数据集的时间（如果真的很长，这需要时间）.当我调用

tolist

时，我只创建了一个groupby金额中的内容列表，即6个项目。这是否达到了您想要的速度？不，不是。根据我的测试，如果不是稍微慢一点的话，它几乎是一样的。我正在寻找一个可以应用于groupby对象的函数，这样我就不必进行迭代。