量化第一课：backtrader实现双均线策略

# -*- coding: utf-8 -*-
"""数据加载模块"""


import pandas as pd
import pymysql
from typing import Optional

# 数据库配置
DB_CONFIG = {
    "host": "223.109.142.60",
    "port": 3306,
    "user": "huayang",
    "password": "huayang555",
    "database": "etf",
    "charset": "utf8mb4"
}

#从数据库拿数据
def read_etf_data(code: str, start: Optional[str] = None, end: Optional[str] = None) -> pd.DataFrame:
    sql = "SELECT trade_date as date, open, high, low, close, volume FROM etf_hist_data WHERE code = %s"
    params = [code]
    
    if start:
        sql += " AND trade_date >= %s"
        params.append(start)
    if end:
        sql += " AND trade_date <= %s"
        params.append(end)
    
    sql += " ORDER BY trade_date ASC"

    try:
        conn = pymysql.connect(**DB_CONFIG)
        df = pd.read_sql(sql, conn, params=params)
        conn.close()
        return df
    except Exception as e:
        print(f"⚠️ 数据库读取失败 [{code}]: {e}")
        return pd.DataFrame()

#转化成bt格式
def load_etf(code: str, start: str, end: str):
    df = read_etf_data(code, start, end)
    if df.empty or len(df) < 10:
        print(f"⚠️ {code} 数据不足")
        return None
    
    df['date'] = pd.to_datetime(df['date'])
    df = df.dropna(subset=['open', 'high', 'low', 'close', 'volume'])
    df.set_index('date', inplace=True)
    
    return df

# -*- coding: utf-8 -*-
"""可视化模块"""

import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
import data_loader

matplotlib.use('TkAgg')
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

def plot_return_curve(strategy, etf_code: str, start: str, end: str, bench_code: str = None):
    """
    绘制收益率曲线
    """
    # 策略数据
    df = pd.DataFrame({
        'date': pd.to_datetime(strategy.date_list),
        'nav': strategy.nav_list
    }).drop_duplicates(subset='date', keep='last')
    
    initial = df['nav'].iloc[0]
    df['strat_return'] = (df['nav'] - initial) / initial
    
    # 加载成交量
    vol_df = data_loader.read_etf_data(etf_code, start, end)
    if not vol_df.empty:
        vol_df['date'] = pd.to_datetime(vol_df['date'])
        df = pd.merge(df, vol_df[['date', 'volume']], on='date', how='left')
    
    # 加载基准数据
    if bench_code:
        bench_df = data_loader.read_etf_data(bench_code, start, end)
        if not bench_df.empty:
            bench_df['date'] = pd.to_datetime(bench_df['date'])
            bench_df = bench_df.dropna(subset=['close']).sort_values('date')
            bench_df['bench_return'] = bench_df['close'] / bench_df['close'].iloc[0] - 1
            df = pd.merge(df, bench_df[['date', 'bench_return']], on='date', how='left')
            df['bench_return'] = df['bench_return'] - df['bench_return'].iloc[0]
    
    # 绘图
    has_vol = 'volume' in df.columns
    has_bench = bench_code and 'bench_return' in df.columns
    
    if has_vol:
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8),  gridspec_kw={'height_ratios': [3, 1]}, sharex=True)
    else:
        fig, ax1 = plt.subplots(figsize=(12, 6))
    
    # 收益曲线
    ax1.plot(df['date'], df['strat_return'], label='策略收益', c='red', lw=2)
    if has_bench:
        ax1.plot(df['date'], df['bench_return'], label='基准收益', c='blue', lw=1.5, alpha=0.7)
    
    ax1.set_title('策略收益率曲线', fontsize=14)
    ax1.set_ylabel('累计收益率')
    ax1.legend()
    ax1.yaxis.set_major_formatter(plt.FuncFormatter(lambda x, _: f'{x*100:.0f}%'))
    ax1.grid(True, alpha=0.3)
    ax1.axhline(y=0, color='gray', lw=0.5)
    
    # 成交量
    if has_vol:
        ax2.bar(df['date'], df['volume'], color='gray', alpha=0.7)
        ax2.set_ylabel('成交量')
        ax2.yaxis.set_major_formatter(plt.FuncFormatter(lambda x, _: f'{x/1e8:.0f}亿'))
    
    fig.autofmt_xdate(rotation=45)
    plt.tight_layout()
    plt.show()