如何向量化改造SmartMoney类函数以实现DataFrame全量运算提速

问题背景 #

开发Smart Money Concept交易策略的Python SmartMoney 类时，最初将is_uptrend()和has_bull_choch()设计为仅处理DataFrame最后一行/单根K线的逻辑，导致回测半年数据时逐行循环效率极低。需将这两个方法通过向量化改造，实现批量处理整个DataFrame，返回包含新增对应列（列值为1或0）的完整DataFrame。

测试数据 #

import yfinance as yf
ticker_symbol = "BTC-USD"

start_date = "2023-06-01"
end_date = "2023-12-31"

bitcoin_data = yf.download(ticker_symbol, start=start_date, end=end_date)

# 将日期转为常规列并重置索引
df = bitcoin_data.reset_index()

# 重命名列以适配策略逻辑
df.rename(columns={
    'Date': 'time', 
    'Open': 'open', 
    'High': 'high', 
    'Low': 'low', 
    'Close': 'close', 
    'Adj Close': 'adj_close', 
    'Volume': 'volume'
}, inplace=True)

期望调用方式 #

from smart_money import SmartMoney
sm = SmartMoney()

# 获取支撑阻力位（假设该方法已实现向量化）
df = sm.get_supports_and_resistances(df=df)
# 新增is_uptrend列（1=处于上升趋势，0=非上升趋势）
df = sm.is_uptrend(df=df)
# 新增has_bull_choch列（1=出现牛势CHOCH信号，0=无信号）
df = sm.has_bull_choch(df=df)

向量化改造实现 #

利用Pandas的矢量化操作（如shift()、rolling()、逻辑运算符）替代逐行循环，直接批量计算整列数据。以下是改造后的SmartMoney类核心代码：

import pandas as pd

class SmartMoney:
    def get_supports_and_resistances(self, df):
        # 示例：用20周期滚动窗口计算支撑阻力（可根据策略替换为实际逻辑）
        df['support'] = df['low'].rolling(window=20).min()
        df['resistance'] = df['high'].rolling(window=20).max()
        return df
    
    def is_uptrend(self, df):
        # 上升趋势判断逻辑（可根据策略调整）：当前高点>前一高点 且 当前低点>前一低点
        df['is_uptrend'] = ((df['high'] > df['high'].shift(1)) & 
                            (df['low'] > df['low'].shift(1))).astype(int)
        return df
    
    def has_bull_choch(self, df):
        # 牛势CHOCH信号判断逻辑（示例，需匹配你的策略定义）：
        # 前一根K线跌破前低，当前K线收盘价收回前低上方且高点突破跌破前的高点
        prev_low = df['low'].shift(1)
        break_prev_low = (df['low'].shift(1) < df['low'].shift(2)).astype(int)
        reclaim = (df['close'] > prev_low).astype(int)
        higher_high = (df['high'] > df['high'].shift(2)).astype(int)
        
        df['has_bull_choch'] = ((break_prev_low == 1) & (reclaim == 1) & (higher_high == 1)).astype(int)
        return df

关键说明 #

• 趋势和CHOCH的判断逻辑是通用示例，需根据你实际的Smart Money策略规则调整（比如修改高低点的确认周期、添加成交量过滤等）。
• 矢量化操作完全避免了逐行循环，处理半年K线数据的效率会提升数倍甚至数十倍。
• 新增列is_uptrend和has_bull_choch的值严格为1（满足条件）或0（不满足条件），符合需求。

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