Question

我有一个Pandas DataFrame，结构如下：

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

我还有一个包含以下内容的DataFrame：

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)

我想按location code进行分组，然后对amount的值执行线性回归并存储系数。我尝试了以下代码：

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），而且有点复杂，这需要非常长的时间。我想知道是否有一个矢量化操作可以更有效地做到这一点？我知道在Pandas DataFrame上使用循环是不好的。

解

经过一些修修补补，我写了一个函数，可以在apply上使用groupby方法。

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

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

Answer 1

如果不了解有关数据，记录数量等的更多信息，此代码应该运行得更快：

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)

通过DataFrame在Python Pandas组上执行数学运算

1 个答案: