模型扩展、板块与绘图

本章介绍 QMT 内置 Python 的模型扩展、自定义指标和图表绘制功能。

自定义指标

创建自定义指标

def init(ContextInfo):
    ContextInfo.set_universe(['000001.SZ'])
    # 注册自定义指标
    ContextInfo.bindBindbar('MY_MA', 'my_ma')

def my_ma(ContextInfo, index_info):
    """自定义均线指标"""
    closes = ContextInfo.get_history_data('close', 20)
    if closes and len(closes) >= 20:
        ma = sum(closes[-20:]) / 20
        index_info.set_value(ma)

指标参数

def init(ContextInfo):
    # 注册带参数的指标
    ContextInfo.bindBindbar('MY_MACD', 'my_macd', 
        params={'short': 12, 'long': 26, 'signal': 9})

def my_macd(ContextInfo, index_info):
    short = index_info.get_param('short')
    long_period = index_info.get_param('long')
    signal = index_info.get_param('signal')

    closes = ContextInfo.get_history_data('close', long_period + signal)
    if not closes or len(closes) < long_period + signal:
        return

    # 计算 MACD
    ema_short = calc_ema(closes, short)
    ema_long = calc_ema(closes, long_period)
    dif = ema_short - ema_long
    index_info.set_value(dif)

def calc_ema(data, period):
    """计算指数移动平均"""
    multiplier = 2 / (period + 1)
    ema = data[0]
    for price in data[1:]:
        ema = price * multiplier + ema * (1 - multiplier)
    return ema

模型扩展

使用内置 talib 库

import talib

def handlebar(ContextInfo):
    if not ContextInfo.is_last_bar():
        return

    closes = ContextInfo.get_history_data('close', 50)
    if not closes or len(closes) < 50:
        return

    closes_array = np.array(closes)

    # 计算技术指标
    ma20 = talib.SMA(closes_array, timeperiod=20)
    rsi = talib.RSI(closes_array, timeperiod=14)
    macd, signal, hist = talib.MACD(closes_array)

    print(f'MA20: {ma20[-1]}, RSI: {rsi[-1]}')

可用的 talib 指标

指标	函数	说明
SMA	`talib.SMA`	简单移动平均
EMA	`talib.EMA`	指数移动平均
RSI	`talib.RSI`	相对强弱指标
MACD	`talib.MACD`	MACD 指标
KDJ	`talib.STOCH`	随机指标
BOLL	`talib.BBANDS`	布林带
ATR	`talib.ATR`	真实波幅

自定义计算函数

import numpy as np

def calc_bollinger(closes, period=20, num_std=2):
    """计算布林带"""
    closes = np.array(closes)
    ma = np.mean(closes[-period:])
    std = np.std(closes[-period:])

    upper = ma + num_std * std
    lower = ma - num_std * std

    return upper, ma, lower

def calc_kdj(highs, lows, closes, period=9):
    """计算 KDJ 指标"""
    highs = np.array(highs[-period:])
    lows = np.array(lows[-period:])
    closes = np.array(closes[-period:])

    highest = np.max(highs)
    lowest = np.min(lows)

    rsv = (closes[-1] - lowest) / (highest - lowest) * 100
    return rsv

板块管理

获取板块列表

# 获取所有板块
sectors = ContextInfo.get_sector_list()
for s in sectors:
    print(s)

获取板块成分

# 获取板块成分股
stocks = ContextInfo.get_stock_list_in_sector('沪深300')
print(f'成分股数量: {len(stocks)}')

自定义板块

# 创建自定义板块
ContextInfo.add_sector('MY_STOCKS', ['000001.SZ', '600000.SH'])

# 获取自定义板块成分
stocks = ContextInfo.get_stock_list_in_sector('MY_STOCKS')

图表绘制

基本绘图

import matplotlib
matplotlib.use('Agg')  # 使用非交互式后端
import matplotlib.pyplot as plt

def draw_chart(closes, title='价格走势'):
    """绘制价格走势图"""
    fig, ax = plt.subplots(figsize=(10, 5))
    ax.plot(closes, label='收盘价')
    ax.set_title(title)
    ax.set_xlabel('日期')
    ax.set_ylabel('价格')
    ax.legend()
    ax.grid(True)

    # 保存图片
    fig.savefig('chart.png', dpi=150, bbox_inches='tight')
    plt.close(fig)
    print('图表已保存到 chart.png')

在策略中使用绘图

def handlebar(ContextInfo):
    if not ContextInfo.is_last_bar():
        return

    closes = ContextInfo.get_history_data('close', 60)
    if not closes or len(closes) < 60:
        return

    # 每 20 天画一次图
    if ContextInfo.barpos % 20 == 0:
        draw_chart(closes, f'{ContextInfo.stockcode} 价格走势')

绘制 K 线图

from matplotlib.patches import Rectangle

def draw_candlestick(opens, highs, lows, closes, title='K线图'):
    """绘制 K 线图"""
    fig, ax = plt.subplots(figsize=(12, 6))

    for i in range(len(closes)):
        # 判断涨跌
        color = 'red' if closes[i] >= opens[i] else 'green'

        # 画实体
        rect = Rectangle(
            (i, min(opens[i], closes[i])),
            0.6,
            abs(closes[i] - opens[i]),
            facecolor=color,
            edgecolor=color,
        )
        ax.add_patch(rect)

        # 画影线
        ax.plot([i + 0.3, i + 0.3], [lows[i], highs[i]], color=color, linewidth=0.5)

    ax.set_title(title)
    ax.set_xlabel('交易日')
    ax.set_ylabel('价格')
    fig.savefig('candlestick.png', dpi=150, bbox_inches='tight')
    plt.close(fig)

数据导出

导出 CSV

import csv

def export_to_csv(ContextInfo):
    """导出策略数据到 CSV"""
    closes = ContextInfo.get_history_data('close', 20)

    with open('output.csv', 'w', newline='', encoding='utf-8') as f:
        writer = csv.writer(f)
        writer.writerow(['日期', '收盘价'])
        for i, close in enumerate(closes):
            writer.writerow([i, close])

    print('数据已导出到 output.csv')

提示

QMT 内置的 matplotlib 使用 Agg 后端，无法实时显示图表，只能保存为图片文件。