python - 提高加权移动平均线的表现

Question

我一直在玩弄一个包含 414,000 行的 pandas 数据框。

pandas 内置的是指数移动平均值，计算公式为:

series.ewm(span=period).mean()

以上执行时间< 0.3 秒。然而，我正在寻求尝试使用加权移动平均值(每个元素具有线性线性权重)。我遇到了以下函数:

def WMA(self, s, period):
    return s.rolling(period).apply(lambda x: (np.arange(period)+1*x).sum()/(np.arange(period)+1).sum(), raw=True)

上述函数执行花费了 27 秒。我注意到 arange 函数可以被缓存并生成以下内容:

def WMA(self, s, period):
    weights = np.arange(period)+1
    weights_sum = weights.sum()
    return s.rolling(period).apply(lambda x: (weights*x).sum()/weights_sum, raw=True)

上述函数花费了11秒，这是一个显着的改进。

我想弄清楚是否有某种方法可以进一步优化它(最好替换 apply 函数)，但我真的不确定如何去做。

任何想法将不胜感激!

Answer 1

可以使用np滑动窗口函数docs ，那么它看起来像这样:

import numpy as np
import pandas as pd

d1 = pd.DataFrame(np.random.randint(0, 10, size=(500_000))) # x=500_000

p = 50
w = np.arange(p)+1
w_s = w.sum()

########## for comparison purpose ##########
# 1.47 s ± 12.5 ms per loop (mean ± std. dev. of 7 runs, 2 loops each)
r = d1.rolling(p).apply(lambda x: (w*x).sum()/w_s, raw=True)

# 62.1 ms ± 4.57 ms per loop (mean ± std. dev. of 7 runs, 2 loops each)
swv = np.lib.stride_tricks.sliding_window_view(d1.values.flatten(), window_shape=p)
sw = (swv*w).sum(axis=1) / w_s

########## for comparison purpose ##########
np.array_equal(r.iloc[p - 1:].values.flatten(), sw) # True

因此，整体加速约为 ~23.67x。但是，您需要随后将形状调整为您想要的形状。由于 sw 从 0 开始，形状为 x-p。而 r 从 p 开始，形状为 x 且第一个 p 值为 -> nan .