需要 pandas 个具有 100 万个股票数据的优化代码

Question

目前我的密码是

self.df['sma'] = self.df['Close'].rolling(window=30).mean()

self.df['cma'] = self.df.apply(lambda x: self.get_cma(x), axis=1)

def get_cma(self, candle):
    if np.isnan(candle['sma']):
        return np.nan
    secma = (candle['sma'] - self.previous_cma if self.previous_cma is not None else 0) ** 2 
    ka = 1 - (candle['var']/secma) if candle['var'] < secma else 0
    cma = ((ka * candle['sma']) + ((1 - ka) * self.previous_cma)) if self.previous_cma is not None else candle[self.src]
    self.previous_cma = cma
    return cma

上面的优化可以让它更快吗？

Answer 1

您可能已经知道，Pandas 的性能秘诀在于以矢量化形式执行此操作。这意味着没有 apply。以下是通过将 get_cma() 函数的部分提取为其向量化等价物来加速代码需要采取的前几个步骤。

if np.isnan(candle['sma']):
    return np.nan

get_cma()不需要提前退出，我们可以这样做：

self.df['cma'] = np.nan
valid = self.df['sma'].notnull()
# this comment is a placeholder for step 2
self.df.loc[valid, 'cma'] = self.df[valid].apply(self.get_cma, axis=1)

这不仅矢量化了 get_cma() 的前两行，这意味着 get_cma() 现在只在 not-null 行而不是每一行上调用。根据您的数据，仅此一项就可以提供明显的加速。

如果这还不够，我们需要更大的锤子。根本问题是 get_cma() 的每次迭代都依赖于前一次迭代，因此向量化并不容易。因此，让我们使用 Numba 来 JIT 编译代码。首先，我们需要通过在各个列上使用一个很好的旧 for 循环来摆脱 apply，这是等效的（并且仍然很慢）。请注意，这是一个免费（全局）函数，而不是成员函数，它采用 NumPy 数组而不是 Pandas 类型，因为这些是 Numba 所理解的：

def get_cma(sma, var, src):
    cma = np.empty_like(sma)
    # take care of the initial value first, to avoid unnecessary branches later
    cma[0] = src[0]

    # now do all remaining rows, cma[ii-1] is previous_cma and is never None
    for ii in range(1, len(sma)):
        secma = (sma[ii] - cma[ii-1]) ** 2 
        ka = 1 - (var[ii] / secma) if var[ii] < secma else 0
        cma[ii] = (ka * sma[ii]]) + ((1 - ka) * cma[ii-1])

    return cma

像这样调用它，将所需的列作为 NumPy 数组传递：

valid_rows = self.df[valid]
self.df.loc[valid, 'cma'] = get_cma(
    valid_rows['sma'].to_numpy(),
    valid_rows['var'].to_numpy(),
    valid_rows[self.src].to_numpy())

最后，确认代码可以正常工作后，装饰get_cma()使其自动使用Numba编译，如下所示：

import numba

@numba.njit
def get_cma(sma, var, src):
    ...

就是这样。请让我们知道这对您的真实数据的运行速度有多快。我希望它会足够快。

需要 pandas 个具有 100 万个股票数据的优化代码

Need pandas optimized code with 1 million stock data

vectorization

apply

pandas