如何优化Pandas实现的交易盘整区间识别代码？

问题描述 #

我有一个从TradingView Pine Script转译到Python的函数，用于接收OHLCV蜡烛图数据，识别价格在区间内震荡的盘整时段。函数逻辑正确，但存在严重性能瓶颈：处理30K根蜡烛时，循环前的代码仅需约5秒，而末尾的for循环耗时超过2分钟。暂时无法将这段循环改为非迭代实现，寻求该段代码的优化方案。

原代码实现 #

import numpy as np
import pandas as pd
from pandas import (DataFrame, Series)

def _find_zz(row: Series):
    if pd.notnull(row['hb']) and pd.notnull(row['lb']):
        if row['dir'] == 1:
            return row['hb']
        else:
            return row['lb']
    else:
        return row['hb'] if pd.notnull(row['hb']) else row['lb'] if pd.notnull(row['lb']) else np.NaN

def consolidation_zones(dataframe: DataFrame, timeperiod: int = 100,
        minlength: int = 20) -> DataFrame:
    rolling = dataframe.rolling(timeperiod, min_periods=1)
    idxmax = rolling['high'].apply(lambda x: x.idxmax()).astype(int)
    idxmin = rolling['low'].apply(lambda x: x.idxmin()).astype(int)
    highest = pd.concat({'value': dataframe['high'], 'offset': dataframe.index - idxmax}, axis=1)
    lowest = pd.concat({'value': dataframe['low'], 'offset': dataframe.index - idxmin}, axis=1)
    hb = highest.apply(lambda x: x['value'] if x['offset'] == 0 else np.NaN, axis=1)
    lb = lowest.apply(lambda x: x['value'] if x['offset'] == 0 else np.NaN, axis=1)
    direction = pd.concat({'hb': hb, 'lb': lb}, axis=1).apply(lambda x: 1 if pd.notnull(x['hb']) and pd.isnull(x['lb']) else -1 if pd.isnull(x['hb']) and pd.notnull(x['lb']) else np.NaN, axis=1).fillna(method='ffill').fillna(0).astype(int)
    zz = pd.concat({'hb': hb, 'lb': lb, 'dir': direction}, axis=1).apply(_find_zz, axis=1)

    group = direction.ne(direction.shift()).cumsum()
    zzdir = pd.concat({'zz': zz, 'dir': direction, 'group': group}, axis=1)
    zzdir['min'] = zzdir.groupby('group')['zz'].cummin().fillna(method='ffill')
    zzdir['max'] = zzdir.groupby('group')['zz'].cummax().fillna(method='ffill')
    zzdir['pp'] = np.NaN
    pp = Series(np.where(zzdir['dir'] == 1, zzdir['max'], np.where(zzdir['dir'] == -1, zzdir['min'], zzdir['pp'])))

    H = dataframe.rolling(minlength, min_periods=1)['high'].max()
    L = dataframe.rolling(minlength, min_periods=1)['low'].min()
    prevpp = np.NaN
    conscnt = 0
    condhigh = np.NaN
    condlow = np.NaN
    zones = DataFrame(index=dataframe.index, columns=['upper_bound', 'lower_bound'])
    indexes = [] # will keep indexes of candles that are part of the consolidation
#----------------
    for index, value in pp.items():
        # pp is a value computed before: when it changes, it *may* be the end of a consolidation zone
        if value != prevpp:
            if conscnt > 0 and value <= condhigh and value >= condlow:
                # if condlow <= pp <= condhigh, we are still in consolidation
                conscnt = conscnt + 1
                indexes.append(index)
            else: # end of consolidation
                conscnt = 0
                indexes = []
        else:
            conscnt = conscnt + 1
            indexes.append(index)

        if conscnt >= minlength:
            if conscnt == minlength:
                # initially, condhigh/low is equal to the highest/lowest value in last minlength candles
                condhigh = H.get(index)
                condlow = L.get(index)
            else:
                # update condhigh/low with new high/low
                condhigh = max(condhigh, dataframe.loc[index, 'high'])
                condlow = min(condlow, dataframe.loc[index, 'low'])
            zones.loc[zones.index.isin(indexes), 'upper_bound'] = condhigh
            zones.loc[zones.index.isin(indexes), 'lower_bound'] = condlow
        prevpp = value
#----------------

    return zones

核心优化思路 #

原循环的性能瓶颈主要在于频繁修改DataFrame、低效的索引查找和重复的属性访问。以下是针对性的优化方案：

1. 批量处理盘整区间，避免循环内修改DataFrame #

原循环中每次满足条件时都执行zones.loc[zones.index.isin(indexes), ...]，这会反复触发DataFrame的索引匹配和赋值，是性能杀手。改为先记录所有盘整区间的起止索引和对应上下界，最后一次性批量赋值。

2. 用numpy数组替代列表和DataFrame属性访问 #

将pp、high、low等序列转为numpy数组，大幅提升循环内的访问速度；用起止索引替代存储所有索引的列表，减少内存占用和查找时间。

3. 简化条件判断逻辑 #

提前处理pp的变化点，减少循环内的重复判断，同时优化盘整区间的起止记录逻辑。

优化后的核心循环代码 #

替换原代码中被标记的for循环部分：

# 提前转为numpy数组，提升循环内访问速度
pp_vals = pp.values
high_vals = dataframe['high'].values
low_vals = dataframe['low'].values
H_vals = H.values
L_vals = L.values

# 初始化变量
prevpp = np.nan
conscnt = 0
condhigh = np.nan
condlow = np.nan
# 存储所有盘整区间：(起始索引, 结束索引, 上界, 下界)
consolidation_intervals = []
current_interval_start = None

for i in range(len(pp_vals)):
    value = pp_vals[i]
    
    if value != prevpp:
        if conscnt > 0 and condlow <= value <= condhigh:
            # 当前仍处于盘整状态，计数累加
            conscnt += 1
        else:
            # 结束当前盘整区间（如果已满足最小长度）
            if conscnt >= minlength:
                consolidation_intervals.append((current_interval_start, i-1, condhigh, condlow))
            # 重置计数，开始新的潜在盘整
            conscnt = 1
            current_interval_start = i
    else:
        conscnt += 1

    # 更新盘整区间的上下界
    if conscnt >= minlength:
        if conscnt == minlength:
            # 首次满足最小长度，初始化上下界
            condhigh = H_vals[i]
            condlow = L_vals[i]
        else:
            # 后续更新上下界
            condhigh = max(condhigh, high_vals[i])
            condlow = min(condlow, low_vals[i])
    
    prevpp = value

# 处理最后一个未闭合的盘整区间
if conscnt >= minlength:
    consolidation_intervals.append((current_interval_start, len(pp_vals)-1, condhigh, condlow))

# 批量赋值到zones，替代循环内的多次修改
zones = pd.DataFrame(index=dataframe.index, columns=['upper_bound', 'lower_bound'])
for start_idx, end_idx, upper, lower in consolidation_intervals:
    zones.loc[zones.index[start_idx:end_idx+1], 'upper_bound'] = upper
    zones.loc[zones.index[start_idx:end_idx+1], 'lower_bound'] = lower

额外优化建议（非循环部分） #

除了循环，前面的apply调用也可以进一步优化：

• 将_find_zz函数改为numpy向量化实现，替代apply：

# 替代原zz的计算逻辑
hb_vals = hb.values
lb_vals = lb.values
dir_vals = direction.values

zz = np.where(
    np.logical_and(pd.notnull(hb_vals), pd.notnull(lb_vals)),
    np.where(dir_vals == 1, hb_vals, lb_vals),
    np.where(pd.notnull(hb_vals), hb_vals, lb_vals)
)
zz = pd.Series(zz, index=dataframe.index)

这能大幅减少apply带来的性能开销。

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