基于EWMA的滚动波动率计算性能优化方案咨询

你的问题核心是rolling.apply逐窗口调用Python自定义函数带来的巨大开销，以下是几种保证结果一致的高性能优化方法：

方法一：全量向量化计算（最优方案） #

利用numpy的滑动窗口视图一次性处理所有窗口，完全避开Python级别的循环，速度提升最明显。

步骤与代码： #

1. 预计算固定权重数组（因为窗口大小和加权参数固定，无需每个窗口重复生成）
2. 用滑动窗口提取所有窗口数据，批量完成均值、平方差、加权求和等计算

import numpy as np
import pandas as pd

# 固定参数
WINDOW_SIZE = 255
PARAM = 0.99

# 预计算权重数组及权重和，与原函数逻辑完全一致
coefs = PARAM ** np.arange(WINDOW_SIZE)[::-1]
coefs_sum = coefs.sum()

def fast_ewm_std(df):
    # 将DataFrame转为numpy数组
    data = df.values
    # 生成滑动窗口，shape: (有效窗口数, 列数, 窗口大小)
    windows = np.lib.stride_tricks.sliding_window_view(data, WINDOW_SIZE, axis=0)
    
    # 批量计算每个窗口的算术均值（与原函数逻辑一致，非加权均值）
    window_means = windows.mean(axis=2)
    # 扩展维度以便广播计算平方差
    window_means_expanded = window_means[:, :, np.newaxis]
    
    # 计算平方差并加权求和
    squared_diffs = (windows - window_means_expanded) ** 2
    weighted_sums = (squared_diffs * coefs).sum(axis=2)
    
    # 计算波动率并包装为DataFrame
    std_values = np.sqrt(weighted_sums / coefs_sum)
    result = pd.DataFrame(
        data=std_values,
        index=df.index[WINDOW_SIZE-1:],
        columns=df.columns
    )
    
    # 补全前WINDOW_SIZE-1行的NaN，与rolling结果格式对齐
    result = pd.concat([pd.DataFrame(index=df.index[:WINDOW_SIZE-1], columns=df.columns), result])
    return result

调用方式：fast_ewm_std(your_df)

方法二：用Numba编译自定义函数（代码改动最小） #

如果不想大幅修改原有逻辑，可借助Numba将Python函数编译为机器码，减少单窗口计算的开销：

from numba import jit
import numpy as np
import pandas as pd

@jit(nopython=True)
def ewm_std_numba(x, param=0.99):
    n = len(x)
    coefs = param ** np.arange(n)[::-1]
    mean_x = np.mean(x)
    squared_diff = (x - mean_x) ** 2
    res = np.sqrt(np.dot(squared_diff, coefs) / np.sum(coefs))
    return res

# 调用方式与原代码几乎一致
result = your_df.rolling(window=255).apply(ewm_std_numba, raw=True)

方法三：预计算权重（小幅提升原方法速度） #

即使继续使用rolling.apply，也可通过预计算权重避免每个窗口重复生成，小幅降低开销：

import numpy as np
import pandas as pd

WINDOW_SIZE = 255
PARAM = 0.99
coefs = PARAM ** np.arange(WINDOW_SIZE)[::-1]
coefs_sum = coefs.sum()

def ewm_std_optimized(x):
    mean_x = np.mean(x)
    squared_diff = (x - mean_x) ** 2
    res = np.sqrt(np.dot(squared_diff, coefs) / coefs_sum)
    return res

result = your_df.rolling(window=255).apply(ewm_std_optimized, raw=True)

验证结果一致性 #

可通过以下代码验证优化后的结果与原方法是否一致（浮点误差范围内）：

# 生成测试数据
np.random.seed(42)
test_df = pd.DataFrame(np.random.randn(300, 2), columns=['A', 'B'])

# 原方法结果
original = test_df.rolling(window=255).apply(ewm_std, raw=True)
# 优化方法结果
optimized = fast_ewm_std(test_df)

# 对比非NaN部分
print(np.allclose(original.dropna(), optimized.dropna()))  # 应返回True

内容的提问来源于stack exchange，提问作者NCall