将自定义函数应用于数据框列

Question

有没有办法优化下面的代码？ 我正在尝试使用前一行列值和自定义函数中指定的时间段以及当前行列中的价格来计算当前行列的值。

import pandas as pd

class EMA_Period:
   fast = 8
   slow = 17

def calculate_ema(prev_ema, price, period):
    return prev_ema + (2.0 / (1.0 + period)) * (price - prev_ema)

times = [1578614400, 1578614700, 1578615000, 1578615300, 1578615600]
closes = [10278.6, 10276.0, 10275.6, 10274.8, 10277.0]
fast_ema = [10278.6, 0, 0, 0, 0]

df = pd.DataFrame(data={'time': times, 'close': closes, 'fast_ema': fast_ema})

df.set_index('time', inplace=True)

for i in range(1, df.shape[0]):
    df.iloc[i]['fast_ema'] = calculate_ema(df.iloc[i-1]['fast_ema'], df.iloc[i]['close'], EMA_Period.fast)

Answer 1

谢谢@火星

def calc_ema(df, period=8, col_name='fast'):
    prev_value = df.iloc[0][col_name]
    def func2(row):
        # non local variable ==> will use pre_value from the new_fun function
        nonlocal prev_value
        prev_value = prev_value + (2.0 / (1.0 + period)) * (row['close'] - prev_value)
        return prev_value
    # This line might throw a SettingWithCopyWarning warning
    df.iloc[1:][col_name] = df.iloc[1:].apply(func2, axis=1)
    return df

df = calc_ema(df)

Answer 2

如果您关心速度，您真的应该使用矢量化方法。 遍历行始终是最慢的选择（尽管有时不可避免）

您甚至不需要更改您的函数来使其矢量化！

def calculate_ema(prev_ema, price, period):
    return prev_ema + (2.0 / (1.0 + period)) * (price - prev_ema)

# though we will make your dataframe longer: 500 rows instead of 5 rows
df = pd.concat([df] * 100)

print(df)
              close  fast_ema
time                         
1578614400  10278.6   10278.6
1578614700  10276.0       0.0
1578615000  10275.6       0.0
1578615300  10274.8       0.0
1578615600  10277.0       0.0
...             ...       ...
1578614400  10278.6   10278.6
1578614700  10276.0       0.0
1578615000  10275.6       0.0
1578615300  10274.8       0.0
1578615600  10277.0       0.0

[500 rows x 2 columns]

请注意，这些测试是计时 2 件重要的事情：

• 计算本身的性能
• 将值分配回数据帧的性能

行循环解决方案

%%timeit
for i in range(1, df.shape[0]):
    df.iloc[i]['fast_ema'] = calculate_ema(df.iloc[i-1]['fast_ema'], df.iloc[i]['close'], EMA_Period.fast)

10 个循环，最好的 5 个：每个循环 86.1 毫秒

对于这么小的数据集，86.1 毫秒是相当慢的。 让我们看看矢量化方法的比较：

---

矢量化解决方案

• 通过在“fast_ema”列上使用.shift()我们可以改变这些向量的对齐方式，使得“close”中的每个值都与之前的“fast_ema”对齐。
• 处理对齐后，我们可以将这些向量直接输入到calculate_ema函数中，而无需进行任何更改

%%timeit 
df["fast_ema"].iloc[1:] = calculate_ema(df["fast_ema"].shift(), df["close"], EMA_Period.fast).iloc[1:]

1000 个循环，最好的 5 个：每个循环 569 µs

时间对比：

方法	时间
行循环	86.1 毫秒
矢量化	569 微秒