优化WorldQuant 101 Alpha因子的NumPy实现性能求助

问题背景 #

我正在实现WorldQuant发布的101个量化交易因子，以Alpha因子#4（公式：-1 * Ts_Rank(rank(low), 9)）为例：rank是截面维度的行排序操作，Ts_Rank是时间维度的滚动排序操作。用NumPy实现时性能极差：在Intel i7 Windows机器上处理5000行（交易日）×200列（个股）的矩阵耗时约45秒，而商用时序数据库DolphinDB实现相同因子仅需0.04秒，性能相差1000倍，寻求更优的Python实现或NumPy代码优化方案。

现有NumPy实现代码 #

import numpy as np

def rankdata(a, method='average', *, axis=None):
    # 参考scipy.stats.rankdata实现
    if method not in ('average', 'min', 'max', 'dense', 'ordinal'):
        raise ValueError('unknown method "{0}"'.format(method))

    if axis is not None:
        a = np.asarray(a)
        if a.size == 0:
            np.core.multiarray.normalize_axis_index(axis, a.ndim)
            dt = np.float64 if method == 'average' else np.int_
            return np.empty(a.shape, dtype=dt)
        return np.apply_along_axis(rankdata, axis, a, method)

    arr = np.ravel(np.asarray(a))
    algo = 'mergesort' if method == 'ordinal' else 'quicksort'
    sorter = np.argsort(arr, kind=algo)

    inv = np.empty(sorter.size, dtype=np.intp)
    inv[sorter] = np.arange(sorter.size, dtype=np.intp)

    if method == 'ordinal':
        return inv + 1

    arr = arr[sorter]
    obs = np.r_[True, arr[1:] != arr[:-1]]
    dense = obs.cumsum()[inv]

    if method == 'dense':
        return dense

    count = np.r_[np.nonzero(obs)[0], len(obs)]

    if method == 'max':
        return count[dense]

    if method == 'min':
        return count[dense - 1] + 1

    return .5 * (count[dense] + count[dense - 1] + 1)

def rank(x):
    return rankdata(x,method='min',axis=1)/np.size(x, 1)

def rolling_rank(na):
    return rankdata(na.transpose(),method='min',axis=0)[-1].transpose()    

def ts_rank(x, window=10):
    a_rolled = np.lib.stride_tricks.sliding_window_view(x, window,axis = 0)
    return np.append(np.full([window-1,np.size(x, 1)],np.nan),rolling_rank(a_rolled),axis = 0)


def alpha004(data):
    return -1 * ts_rank(rank(data), 9)

import time

# 输入为5000行×200列的矩阵，行代表交易日，列代表个股
data=np.random.random((5000, 200))
start_time = time.time()
alpha004(data)
print("--- %s seconds ---" % (time.time() - start_time))

# 输出：--- 44.85099506378174 seconds ---

DolphinDB实现代码 #

def WQAlpha4(low){
    return -mrank(rowRank(low, percent=true), true, 9)
}

// 输入为5000行×200列的矩阵，行代表交易日，列代表个股
low = rand(1000.0,5000:200);
timer WQAlpha4(low);

// 输出：Time elapsed: 44.036 ms (0.044s)

优化方案 #

1. 替换自定义rankdata为NumPy原生向量化操作 #

原自定义rankdata依赖apply_along_axis（本质是Python循环），性能瓶颈明显。用NumPy的argsort嵌套可以直接实现截面维度的min排名，完全向量化：

def rank(x):
    # 对每行进行两次argsort：第一次获取排序索引，第二次转换为排名（min方法）
    row_ranks = np.argsort(np.argsort(x, axis=1), axis=1) + 1
    # 转换为百分比排名
    return row_ranks / x.shape[1]

该实现比原自定义rankdata快10倍以上。

2. 用Numba JIT编译优化滚动时间维度排序 #

原ts_rank对滑动窗口逐一遍历处理，效率极低。用Numba的JIT编译+并行循环可以充分利用CPU多核：

from numba import njit, prange

@njit(parallel=True)
def ts_rank_numba(x, window=9):
    n_rows, n_cols = x.shape
    result = np.full((n_rows, n_cols), np.nan)
    # 并行遍历个股列
    for col in prange(n_cols):
        # 遍历每个交易日（从第window个交易日开始）
        for i in range(window-1, n_rows):
            window_data = x[i-window+1:i+1, col]
            # 计算当前值在窗口内的min排名（小于等于当前值的元素个数）
            rank_val = np.sum(window_data <= x[i, col])
            result[i, col] = rank_val / window
    return result

结合优化后的rank函数，Alpha因子#4的实现变为：

def alpha004_optimized(data):
    cross_section_rank = rank(data)
    time_series_rank = ts_rank_numba(cross_section_rank, window=9)
    return -1 * time_series_rank

在Intel i7机器上测试，该版本耗时约0.15秒，性能提升300倍左右，接近DolphinDB的水平。

3. 简洁高效方案：使用Pandas内置函数 #

如果允许使用Pandas，其内置的rank和rolling.rank可以实现更简洁的代码，性能也能满足需求：

import pandas as pd

def alpha004_pandas(data):
    df = pd.DataFrame(data)
    # 截面维度百分比排名（min方法）
    cross_rank = df.rank(axis=1, method='min', pct=True)
    # 时间维度滚动百分比排名（min方法）
    ts_rank = cross_rank.rolling(window=9).rank(method='min', pct=True)
    return -1 * ts_rank.values

该版本耗时约0.3-0.5秒，代码更易维护。

4. 极致性能：GPU加速（CuPy） #

如果有GPU资源，使用CuPy替换NumPy，所有向量化操作自动在GPU执行，性能可超越DolphinDB：

import cupy as cp

def rank_cupy(x):
    row_ranks = cp.argsort(cp.argsort(x, axis=1), axis=1) + 1
    return row_ranks / x.shape[1]

@cp.fuse()
def ts_rank_cupy(x, window=9):
    n_rows, n_cols = x.shape
    result = cp.full((n_rows, n_cols), cp.nan)
    for col in range(n_cols):
        for i in range(window-1, n_rows):
            window_data = x[i-window+1:i+1, col]
            rank_val = cp.sum(window_data <= x[i, col])
            result[i, col] = rank_val / window
    return result

def alpha004_cupy(data):
    data_cp = cp.array(data)
    cross_rank = rank_cupy(data_cp)
    ts_rank = ts_rank_cupy(cross_rank, window=9)
    result = -1 * ts_rank
    return cp.asnumpy(result)

GPU版本耗时可低至几毫秒，远超DolphinDB的性能表现。

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