如何用Numpy加速DataFrame遍历的回测优化代码？

核心优化方向 #

• 抛弃逐参数循环的Pandas操作，改用Numpy向量化运算处理所有有效参数
• 用Numpy广播、cumsum、reduceat等原生函数替代groupby+lambda的组合，减少Python层级循环开销
• 全程基于Numpy数组操作，避免反复修改Pandas DataFrame带来的内存和性能损耗

假设原代码场景（基于常见回测逻辑） #

先模拟你可能的原代码结构（如果实际逻辑不同，可对应调整优化部分）：

import pandas as pd
import numpy as np

# 行情数据示例
df = pd.DataFrame({
    'close': np.random.randn(100000).cumsum(),
    'ind1': np.random.rand(100000),
    'ind2': np.random.rand(100000),
    'ind3': np.random.rand(100000)
})

# 筛选后的166375组有效参数（假设是三个指标的阈值）
valid_params = np.random.rand(166375, 3)

# 原PNL计算函数
def calc_pnl(params, df):
    th1, th2, th3 = params
    # 生成持仓信号
    df['signal'] = np.where(
        (df['ind1'] > th1) & (df['ind2'] < th2) & (df['ind3'] > th3), 1, 0
    )
    # 计算每日收益率
    df['ret'] = df['close'].pct_change()
    # 分组计算持仓时段累计收益
    pnl = df.groupby((df['signal'] != df['signal'].shift(1)).cumsum())['ret'].apply(
        lambda x: (1 + x).prod() - 1 if x.name % 2 == 1 else 0
    ).sum()
    return pnl

# 遍历所有参数计算PNL（耗时约30分钟）
results = [calc_pnl(p, df) for p in valid_params]

Numpy优化实现 #

import numpy as np

# 1. 预提取数据为Numpy数组，避免Pandas索引开销
close = df['close'].values
ind1 = df['ind1'].values
ind2 = df['ind2'].values
ind3 = df['ind3'].values

# 2. 预计算每日收益率（等价于pct_change()，长度比原数据少1）
rets = np.diff(close) / close[:-1]

# 3. 将有效参数转为Numpy数组（确保格式正确）
valid_params = np.asarray(valid_params)
n_params = valid_params.shape[0]

# 4. 向量化生成所有参数对应的信号矩阵（广播机制，shape: [参数数量, 交易日数量]）
signals = (ind1[np.newaxis, :] > valid_params[:, 0, np.newaxis]) & \
          (ind2[np.newaxis, :] < valid_params[:, 1, np.newaxis]) & \
          (ind3[np.newaxis, :] > valid_params[:, 2, np.newaxis])

# 5. 计算信号变化点，生成分组ID（shape: [参数数量, 交易日数量-1]）
signal_shifts = signals[:, 1:] != signals[:, :-1]
group_ids = np.cumsum(signal_shifts, axis=1)

# 6. 批量计算每个参数的PNL
def batch_calc_pnl(group_ids, rets, signals):
    pnl_results = np.zeros(n_params, dtype=np.float64)
    # 遍历每个参数（相比原逻辑，这里的循环只处理参数，内部是Numpy原生运算）
    for i in range(n_params):
        # 提取当前参数的信号和分组ID
        curr_signal = signals[i, 1:]  # 和rets长度匹配
        curr_groups = group_ids[i]
        # 只保留持仓时段的收益率和对应分组（信号为1时持仓）
        hold_rets = rets[curr_signal]
        hold_groups = curr_groups[curr_signal]
        
        if len(hold_rets) == 0:
            pnl_results[i] = 0.0
            continue
        
        # 用reduceat计算每个分组的(1+ret)乘积
        unique_groups, split_idx = np.unique(hold_groups, return_index=True)
        sorted_split_idx = np.sort(split_idx)
        # 用log转换避免数值溢出，同时加速乘积计算
        log_rets = np.log(1 + hold_rets)
        group_log_sum = np.add.reduceat(log_rets, sorted_split_idx)
        group_prods = np.exp(group_log_sum) - 1
        # 累加所有持仓时段的收益
        pnl_results[i] = group_prods.sum()
    return pnl_results

# 执行计算，耗时可压缩至数分钟甚至更短
results_np = batch_calc_pnl(group_ids, rets, signals)

进一步提速：Numba JIT编译 #

如果想把速度再提升一个量级，可以用Numba对批量计算函数做JIT编译，消除Python循环的开销：

from numba import jit

@jit(nopython=True)
def batch_calc_pnl_numba(group_ids, rets, signals):
    n_params = group_ids.shape[0]
    pnl_results = np.zeros(n_params, dtype=np.float64)
    
    for i in range(n_params):
        curr_signal = signals[i, 1:]
        curr_groups = group_ids[i]
        hold_rets = rets[curr_signal]
        hold_groups = curr_groups[curr_signal]
        
        if len(hold_rets) == 0:
            continue
        
        # 排序分组ID和对应收益率，方便批量计算
        sort_idx = np.argsort(hold_groups)
        sorted_groups = hold_groups[sort_idx]
        sorted_rets = hold_rets[sort_idx]
        
        # 找到分组分割点
        split_points = np.where(sorted_groups[1:] != sorted_groups[:-1])[0] + 1
        split_points = np.concatenate(([0], split_points, [len(sorted_rets)]))
        
        total_pnl = 0.0
        for j in range(len(split_points)-1):
            start = split_points[j]
            end = split_points[j+1]
            prod = 1.0
            for k in range(start, end):
                prod *= (1 + sorted_rets[k])
            total_pnl += (prod - 1)
        
        pnl_results[i] = total_pnl
    return pnl_results

# 编译后执行，速度可提升10~20倍
results_numba = batch_calc_pnl_numba(group_ids, rets, signals)

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