寻求优化Pandas加权移动平均（WMA）性能的高效方案

核心思路 #

WMA的本质是滑动窗口内的加权求和再除以权重总和，要避开apply的逐行计算，优先用向量化操作或编译型函数，利用numpy/numba的底层优化提升速度。

方案1：Numpy卷积（最快方案之一） #

利用numpy.convolve实现滑动加权求和，这是numpy底层C实现的操作，速度远高于Python循环。

假设窗口周期为period，权重采用最近元素权重最高的线性权重（如窗口3，权重为3,2,1）：

import numpy as np
import pandas as pd

# 定义参数
period = 20
weight_sum = period * (period + 1) // 2  # 线性权重的总和
weights = np.arange(1, period + 1)  # 权重数组：1,2,...,period

# 获取目标序列的numpy数组
series_vals = df['target_column'].values

# 计算卷积（mode='valid'只保留完整窗口的结果）
conv_result = np.convolve(series_vals, weights[::-1], mode='valid') / weight_sum

# 将结果对齐到原DataFrame，前面补NaN
df['wma'] = np.concatenate([np.full(period-1, np.nan), conv_result])

• 优势：速度极快，414k行数据的计算耗时可控制在0.1-0.5秒级别，和EMA的性能接近。
• 注意：如果需要最远元素权重最高，直接使用weights而非weights[::-1]。

方案2：Numba编译加速 #

用numba将自定义WMA函数编译为机器码，避免Python解释器的开销，适合需要灵活调整权重逻辑的场景。

from numba import jit
import numpy as np

@jit(nopython=True)  # 编译为原生机器码
def compute_wma(arr, period):
    weight_sum = period * (period + 1) // 2
    result = np.empty(arr.shape[0])
    result[:period-1] = np.nan  # 前period-1个值为NaN
    
    for i in range(period-1, arr.shape[0]):
        # 提取当前窗口，计算加权和
        window = arr[i - period + 1 : i + 1]
        weighted_sum = np.sum(window * np.arange(1, period + 1))
        result[i] = weighted_sum / weight_sum
    
    return result

# 应用到DataFrame
df['wma'] = compute_wma(df['target_column'].values, period)

• 优势：编译后循环速度接近C语言，比纯Python apply快10-20倍，且逻辑灵活易调整。
• 注意：首次调用会触发编译，后续调用速度极快。

方案3：滑动窗口矩阵乘法 #

利用numpy.lib.stride_tricks.as_strided生成滑动窗口矩阵，再通过矩阵乘法完成加权求和，适合需要批量处理多个序列的场景。

from numpy.lib.stride_tricks import as_strided
import numpy as np

def wma_strided(arr, period):
    weight_sum = period * (period + 1) // 2
    weights = np.arange(1, period + 1).reshape(-1, 1)  # 转为列向量
    
    # 生成滑动窗口矩阵（内存复用，无额外拷贝）
    n = arr.shape[0]
    window_shape = (n - period + 1, period)
    strides = (arr.strides[0], arr.strides[0])
    window_matrix = as_strided(arr, shape=window_shape, strides=strides)
    
    # 矩阵乘法计算加权和，再归一化
    wma_vals = (window_matrix @ weights).flatten() / weight_sum
    
    # 补NaN对齐原序列
    result = np.full(n, np.nan)
    result[period-1:] = wma_vals
    return result

# 应用到DataFrame
df['wma'] = wma_strided(df['target_column'].values, period)

• 优势：向量化操作，速度接近卷积方案，适合复杂加权逻辑的批量计算。
• 注意：大窗口（如period>1000）可能会占用较多内存，需根据实际情况调整。

方案4：Pandas Rolling + Raw参数优化 #

如果必须使用Pandas的rolling接口，开启raw=True传递numpy数组而非Series，可大幅降低apply的开销。

import numpy as np

weight_sum = period * (period + 1) // 2
weights = np.arange(1, period + 1)

df['wma'] = df['target_column'].rolling(window=period).apply(
    lambda x: np.sum(x * weights) / weight_sum,
    raw=True  # 关键：传递numpy数组，避免Series的开销
)

• 优势：保留Pandas的rolling接口语义，比普通apply快5-10倍，但速度仍不如前三种方案。

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