🚀 AI 一键生成 qmt 策略代码

拒绝循环:如何使用 NumPy 加速策略中的指标计算?

qmt | 发布时间: 2025/12/21
以下内容由 EasyQuant 生成。

问题描述

策略回测太慢,怎么用 NumPy 优化计算速度?

解决方案

在 QMT 中,策略回测速度慢的主要原因通常是handlebar 函数(逐 K 线执行)中进行了重复的数据获取和计算

要使用 NumPy/Pandas 优化速度,核心思想是**“向量化计算”**:

  1. 移出循环:将历史数据的获取和指标计算从 handlebar 移到 init 中。
  2. 批量获取:使用 get_market_data_ex 一次性获取全量历史数据。
  3. 向量化运算:利用 Pandas/NumPy 的矩阵运算特性,一次性计算出所有时间点的指标(如均线、RSI)和交易信号。
  4. 查表执行:在 handlebar 中仅根据当前时间查找预先计算好的信号进行下单,不做繁重的计算。

下面是一个使用 NumPy/Pandas 优化的双均线策略示例代码。

优化后的 QMT 策略代码

# -*- coding: gbk -*-
import pandas as pd
import numpy as np
import time

def init(ContextInfo):
    # 1. 设置基本参数
    ContextInfo.stock = '600000.SH'
    ContextInfo.set_account('6000000248') # 请替换为您的资金账号
    ContextInfo.set_universe([ContextInfo.stock])
    
    # 2. 设定策略参数
    ContextInfo.short_period = 5
    ContextInfo.long_period = 10
    
    # ============================================================
    # 核心优化部分:在 init 中预先计算所有信号
    # ============================================================
    
    # A. 一次性获取全量历史数据 (注意:count=-1 表示获取尽可能多的数据,或根据回测区间设置)
    # 注意:这里为了演示方便,获取了所有数据。实际回测中,QMT会自动根据回测时间段提供数据,
    # 但在init中获取全量数据需要注意内存和时间范围。
    # 更好的方式是利用 download_history_data 提前下载好数据
    print("开始预计算数据...")
    
    # 获取当前主图品种的所有历史收盘价
    # period='1d' 表示日线,根据实际需求调整
    raw_data = ContextInfo.get_market_data_ex(
        ['close'], 
        [ContextInfo.stock], 
        period='1d', 
        start_time='', 
        end_time='', 
        count=-1, 
        dividend_type='front'
    )
    
    if ContextInfo.stock not in raw_data:
        print(f"未获取到 {ContextInfo.stock} 的数据")
        ContextInfo.signals = pd.Series()
        return

    # B. 数据清洗与格式转换
    df = raw_data[ContextInfo.stock]
    
    # C. 使用 Pandas/NumPy 进行向量化计算 (比逐行计算快几十倍)
    # 计算均线
    df['ma_short'] = df['close'].rolling(window=ContextInfo.short_period).mean()
    df['ma_long'] = df['close'].rolling(window=ContextInfo.long_period).mean()
    
    # D. 生成信号 (1: 买入, -1: 卖出, 0: 持仓/空仓)
    # 逻辑:短线上穿长线做多,短线下穿长线平仓
    df['signal'] = 0
    
    # 产生金叉条件 (当前短线>长线 且 上一时刻短线<=长线)
    # 注意:这里使用 shift(1) 来比较前一刻,避免未来函数
    condition_buy = (df['ma_short'] > df['ma_long']) & (df['ma_short'].shift(1) <= df['ma_long'].shift(1))
    condition_sell = (df['ma_short'] < df['ma_long']) & (df['ma_short'].shift(1) >= df['ma_long'].shift(1))
    
    # 标记信号
    df.loc[condition_buy, 'signal'] = 1
    df.loc[condition_sell, 'signal'] = -1
    
    # 将计算好的信号表存储在 ContextInfo 中,供 handlebar 使用
    # 索引是时间字符串 (如 '20230101')
    ContextInfo.signal_df = df
    
    print("预计算完成,策略开始运行...")

def handlebar(ContextInfo):
    # ============================================================
    # 极速执行部分:仅查表,不计算
    # ============================================================
    
    # 1. 获取当前 K 线的时间
    index = ContextInfo.barpos
    timetag = ContextInfo.get_bar_timetag(index)
    # 将时间戳转换为字符串格式,与 DataFrame 的索引格式保持一致 (YYYYMMDD)
    current_date = timetag_to_datetime(timetag, '%Y%m%d')
    
    # 2. 检查预计算的数据中是否存在当前时间
    if not hasattr(ContextInfo, 'signal_df') or current_date not in ContextInfo.signal_df.index:
        return

    # 3. 直接读取预计算好的信号
    # 注意:回测时,我们通常是在当前 bar 走完后,或者当前 bar 开盘时利用上一根 bar 的信号
    # 这里演示的是:读取当前 bar 对应的计算结果。
    # 如果是日线策略,通常是在当日收盘后计算出信号,次日开盘交易。
    # 为了防止未来函数,建议读取上一日的信号,或者确保 signal_df 的计算逻辑中已经 shift 过了。
    # 在本例 init 中,我们计算的是当天的均线状态。如果要在当天收盘前下单,用 current_date。
    # 如果是次日下单,应获取前一交易日的信号。
    
    signal = ContextInfo.signal_df.loc[current_date, 'signal']
    
    # 4. 执行交易
    account = ContextInfo.get_account(0) # 获取第一个账号
    
    if signal == 1:
        # 金叉买入
        print(f"{current_date} 触发买入信号")
        order_target_percent(ContextInfo.stock, 1.0, "FIX", ContextInfo.signal_df.loc[current_date, 'close'], ContextInfo, account)
        
    elif signal == -1:
        # 死叉卖出
        print(f"{current_date} 触发卖出信号")
        order_target_percent(ContextInfo.stock, 0.0, "FIX", ContextInfo.signal_df.loc[current_date, 'close'], ContextInfo, account)

# 辅助函数:时间戳转字符串
def timetag_to_datetime(timetag, format_str):
    import time
    return time.strftime(format_str, time.localtime(timetag / 1000))

代码优化原理解析

  1. get_market_data_ex 替代 get_market_data

    • get_market_data 在循环中调用时,每次只取少量数据,IO 开销巨大。
    • get_market_data_ex 返回的是 Pandas DataFrame 格式,非常适合批量处理。
  2. init 中的预计算

    • 我们将均线计算(rolling().mean())和逻辑判断(& 运算)全部放在 init 中。
    • Pandas 底层基于 C 语言和 NumPy 优化,计算 10 年的日线数据通常只需要几毫秒,而如果在 handlebar 中用 Python for 循环计算,可能需要几秒甚至更久。
  3. handlebar 变为查表机

    • handlebar 中,我们不再做任何数学运算。
    • 我们只是把当前的 timetag 转换成日期字符串,然后去 ContextInfo.signal_df 这个大表中“查找”这一天应该做什么。
    • 这种 O(1) 复杂度的操作是极快的。

注意事项

  1. 内存占用:如果回测品种非常多(例如全市场 5000 只股票)且是分钟线,一次性加载所有数据到内存可能会导致内存溢出。这种情况下,建议分批次处理或仅在 init 中加载必要的指标数据。
  2. 未来函数:在 init 中进行向量化计算时,非常容易引入未来函数(例如使用了 shift(-1) 或者直接用当天的收盘价在当天开盘时成交)。务必确保逻辑上是“利用过去的数据预测未来”。
  3. 实盘与回测的区别
    • 上述代码主要针对回测优化。
    • 实盘模式下,init 只运行一次,后续新的行情到来时,signal_df 不会自动更新。因此,实盘代码通常需要混合写法:历史数据在 init 算好,最新的一根 K 线数据在 handlebar 中动态追加到 DataFrame 中并计算最新信号。

Q&A

Q: 为什么我的 get_market_data_ex 返回的数据是空的?
A: 请确保您已经在 QMT 客户端的“数据管理”中下载了对应品种和周期的历史数据。API 不会自动从服务器下载历史数据,只能读取本地已下载的数据。

Q: 这种写法支持多只股票吗?
A: 支持。您可以在 init 中遍历 stock_list,对每一只股票分别计算 DataFrame,然后存储在一个字典中(例如 ContextInfo.signal_dict = {code: df})。在 handlebar 中根据 ContextInfo.stockcode 获取当前正在运行的股票代码,去字典里查表即可。

Q: 如何处理分钟线回测?
A: 逻辑完全一样。只需将 period 参数改为 '1m''5m',并将时间格式化 timetag_to_datetime 的格式改为 '%Y%m%d%H%M%S' 以匹配分钟级索引即可。