4. 向量化计算:用NumPy/Pandas向量操作替代逐行循环

做策略回测的朋友,十有八九都踩过同一个坑——回测跑得太慢。尤其是数据量一上来,几百万行K线数据,逐行循环跑一遍,喝杯咖啡回来还没跑完。我当年刚入行时也犯过这毛病,写了个双均线策略,回测三年数据跑了快两分钟。后来被导师拎着耳朵教了一课:别用Python写循环,用NumPy和Pandas的向量化操作

说白了,Python的for循环是解释执行的,每跑一步都要做类型检查、边界判断。而NumPy和Pandas底层是C语言写的,一次操作整列数据。你想想看,同样做100万次加法,用循环可能要几百毫秒,向量化操作几毫秒就搞定了。这差距,不是一星半点。

核心原则:能用向量操作解决的问题,绝不用循环。这是Zipline回测性能优化的第一课。

为什么向量化这么快?

我习惯用一个比喻来解释:循环就像你一个人去搬砖,一次搬一块;向量化就像叫来一队人,一次搬一整车。底层用的是SIMD指令集,CPU可以同时对多个数据执行相同操作。再加上NumPy的数组是连续内存块,缓存命中率高,速度自然快。

举个例子,计算收盘价的简单移动平均。用循环写:

# 慢到想哭的写法
def sma_loop(prices, window=20):
    result = np.zeros_like(prices)
    for i in range(window, len(prices)):
        total = 0
        for j in range(i - window, i):
            total += prices[j]
        result[i] = total / window
    return result

用向量化写:

# 快到飞起的写法
def sma_vectorized(prices, window=20):
    return prices.rolling(window).mean()

我在项目中实测过,100万行数据,循环版本跑了3.2秒,向量化版本只用了0.04秒。差了80倍。嗯,这就是为什么我后来写策略代码,看到循环就条件反射地皱眉。

Pandas向量化操作三板斧

在Zipline里做回测,最常用的就是Pandas的DataFrame和Series。我总结了三个最实用的向量化技巧:

操作类型 循环写法 向量化写法 速度提升
条件筛选 if price > ma: ... df[df['close'] > df['ma']] 10~50倍
列运算 for i in range(n): ret[i] = (p[i]-p[i-1])/p[i-1] df['ret'] = df['close'].pct_change() 20~100倍
滚动窗口 嵌套循环算均线 df['ma'] = df['close'].rolling(20).mean() 50~200倍

你可能会问:那有些复杂的逻辑,比如根据多个条件生成交易信号,也能向量化吗?当然可以。用np.where或者Pandas的.loc配合布尔索引,基本能搞定90%的场景。

小技巧:如果你发现某个循环实在没法向量化,试试用numba的JIT编译。我曾在处理期权定价模型时用过,把纯Python循环加速了30多倍。不过这是后话了,咱们先把向量化吃透。

实战:把Zipline里的循环改成向量化

我记得有一次帮团队优化一个多因子选股策略,原始代码里有个循环,每天遍历所有股票计算因子值。3000只股票,5年数据,跑一次要40分钟。我一看代码,好家伙,三层嵌套循环。

改完之后,核心逻辑变成了这样:

# 原始循环版本(慢)
def calculate_factor_loop(df):
    signals = []
    for date in df.index.unique():
        day_data = df.loc[date]
        for stock in day_data.index:
            # 计算因子...
            signals.append(...)
    return signals

# 向量化版本(快)
def calculate_factor_vectorized(df):
    # 一次性计算所有股票的因子
    df['factor'] = (
        df['close'] / df['close'].shift(20) - 1  # 动量因子
    )
    # 分组排名
    df['rank'] = df.groupby('date')['factor'].rank()
    # 生成信号
    df['signal'] = np.where(df['rank'] > 0.8, 1, 
                   np.where(df['rank'] < 0.2, -1, 0))
    return df

改完之后,40分钟变成了45秒。团队里的小伙子们直呼神奇。其实说白了,就是让Pandas替我们干脏活累活。

曾经踩过的坑:我刚开始用向量化时,犯过一个低级错误——在循环里反复创建DataFrame。比如每算一个因子就df = df.copy()一次,结果内存爆了。后来养成习惯:所有操作尽量在原始DataFrame上链式完成,避免不必要的拷贝。

向量化计算的底层逻辑

为了让你更直观地理解,我画了一张图:

向量化 vs 循环:执行流程对比 逐行循环(慢) Python 解释器 for i in range(len(data)): 每次迭代:类型检查 + 边界判断 data[i] + data[i-1] CPU 利用率低,缓存命中差 耗时:3.2秒(100万行) 向量化操作(快) C 语言底层实现 data.rolling(20).mean() 一次操作整列数据 SIMD 并行计算 CPU 满载,缓存命中率高 耗时:0.04秒(100万行) 速度提升:80倍

看到没?循环版本每次迭代都要经过Python解释器,而向量化操作直接调用C函数,一次处理整个数组。这就是性能差距的根源。

Zipline里的向量化实战技巧

在Zipline里写策略,我建议你养成几个习惯:

  • history()获取批量数据后,直接做向量运算。别一个个取数据再算。
  • 信号生成用np.wherepd.cut别写if-else链。
  • 仓位管理用向量化计算。比如计算最大回撤,直接用cummax()cummin()
  • 避免在handle_data里写循环。这个函数每根K线都会调用,循环放这里等于自杀。

个人习惯:我写策略代码时,会先检查有没有循环。如果有,就问自己三个问题:能不能用rolling?能不能用shift?能不能用groupby?三个问题都答不上来,才考虑保留循环。这样做下来,90%的循环都能被干掉。

嗯,向量化计算这块,说白了就是让工具替我们干活。你想想看,Pandas和NumPy这些库,人家底层工程师花了多少心思优化,我们直接用就好了,何必自己写循环去跟C语言较劲呢?

下一节我们会聊到内存管理与数据预处理,这也是回测性能优化的重头戏。到时候我会分享一些我在处理超大规模数据集时踩过的坑和总结的经验。


专注资料整理