如何通过pandas数据框中的记录迭代计算指数移动平均线（EMA）

Question

我创建了一个 pandas 数据框，如下所示：

import pandas as pd
import numpy as np
    
ds = { 'trend' : [1,1,1,1,2,2,3,3,3,3,3,3,4,4,4,4,4], 'price' : [23,43,56,21,43,55,54,32,9,12,11,12,23,3,2,1,1]}

df = pd.DataFrame(data=ds)

数据框如下所示：

显示(df)

   trend    price
0   1        23
1   1        43
2   1        56
3   1        21
4   2        43
5   2        55
6   3        54
7   3        32 
8   3         9
9   3        12
10  3        11
11  3        12
12  4        23
13  4         3
14  4         2
15  4         1
16  4         1

我已将数据帧保存到名为 df.csv 的 .csv 文件中：

df.to_csv("df.csv", index = False)

我需要创建一个名为

ema2

的新字段，其中：

1. 迭代数据帧的每条记录

2. 通过考虑每次迭代中观察到的价格以及之前趋势中观察到的价格（本例中 EMA 长度为 2）来计算指数移动平均线 (EMA)。 例如：

3. 我在记录 0 处迭代，EMA 为 NaN（缺失）。

4. 我在记录 1 处迭代，EMA 仍然为 NaN（缺失）

5. I 在记录 12 处迭代，EMA 为 24.20（它考虑记录 3 处的价格、记录 5 处的价格和记录 12 处的价格

6. I 在记录 13 处迭代，EMA 为 13.53（它考虑记录 3 处的价格、记录 5 处的价格和记录 13 处的价格

7. I 在记录 15 处迭代，EMA 为 12.46（它考虑记录 3 处的价格、记录 5 处的价格和记录 15 处的价格 等等......

我写了以下代码：

time_window = 2

ema= []

for i in range(len(df)):

    ds = pd.read_csv("df.csv", nrows=i+1)
    d = ds.groupby(['trend'], as_index=False).agg(
                                                    {'price':'last'})

    d['ema2'] = d['price'].ewm(com=time_window - 1, min_periods=time_window).mean()    
    ema.append(d['ema2'].iloc[-1])

df['ema2'] = ema

这会产生正确的数据帧：

print(df)

    trend  price       ema2
0       1     23        NaN
1       1     43        NaN
2       1     56        NaN
3       1     21        NaN
4       2     43  35.666667
5       2     55  43.666667
6       3     54  49.571429
7       3     32  37.000000
8       3      9  23.857143
9       3     12  25.571429
10      3     11  25.000000
11      3     12  25.571429
12      4     23  24.200000
13      4      3  13.533333
14      4      2  13.000000
15      4      1  12.466667
16      4      1  12.466667

问题在于，当数据帧有数百万条记录时：运行需要很长时间。

有人知道如何快速有效地获得相同的结果吗？

Answer 1

你可以试试这个版本的代码，应该会更快。如果可能的话，尽量避免使用 for 循环。在您的代码中，您在每次迭代中从文件“df.csv”中读取数据，这使得您的代码非常非常慢。

import pandas as pd
import numpy as np

# Data setup
ds = {
    'trend': [1,1,1,1,2,2,3,3,3,3,3,3,4,4,4,4,4],
    'price': [23,43,56,21,43,55,54,32,9,12,11,12,23,3,2,1,1]
}
df = pd.DataFrame(data=ds)


time_window = 2

df['ema2'] = df.groupby('trend')['price'].transform(lambda x: x.ewm(com=time_window - 1, min_periods=time_window).mean())

print(df)