AI写量化代码【1】-alpha#1

Alpha#1 本质是一个结合 “涨跌状态、波动 / 价格、短期极值时机” 的复合因子，核心逻辑是：
通过捕捉股票在 “下跌时的波动特征” 和 “上涨时的价格特征”，结合近期极端值出现的时机，筛选出在短期趋势节奏上具有优势的股票（因子值越大，选股优先级越高）。

在实际应用中，该因子通常用于多因子模型，与其他因子（如动量、价值、质量等）结合，提升选股效果。

import warnings
import tushare as ts
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib.colors import LinearSegmentedColormap
from mplfinance.original_flavor import candlestick_ohlc
import mplfinance as mpf

# 基础配置：中文字体+图表风格（专业金融风格）
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
plt.rcParams['axes.grid'] = True
plt.rcParams['grid.alpha'] = 0.3
plt.rcParams['figure.figsize'] = (16, 12)


# 1. 初始化Tushare接口
def init_tushare(token):
    ts.set_token(token)
    return ts.pro_api()


# 2. 获取并筛选股票池（仅60/0开头）
def get_filtered_stock_pool(pro, is_hs300=False):
    if is_hs300:
        hs300 = pro.index_weight(index_code='000300.SH', trade_date='')
        stock_pool = hs300['con_code'].unique().tolist()
    else:
        stocks = pro.stock_basic(list_status='L', fields='ts_code')
        stock_pool = stocks['ts_code'].unique().tolist()
    # 筛选60/0开头股票
    filtered_pool = [code for code in stock_pool if code.split('.')[0][:2] in ['60', '00']]
    return filtered_pool


# 3. 获取日线数据（扩展为OHLC格式，支持mplfinance绘图）
def get_daily_data(pro, stock_codes, start_date, end_date):
    start_str = start_date.strftime('%Y%m%d')
    end_str = end_date.strftime('%Y%m%d')
    all_data = []
    batch_size = 300
    for i in range(0, len(stock_codes), batch_size):
        batch_codes = stock_codes[i:i + batch_size]
        # 获取OHLC数据用于金融图表
        df = pro.daily(
            ts_code=','.join(batch_codes),
            start_date=start_str,
            end_date=end_str,
            fields='ts_code, trade_date, open, high, low, close, vol'
        )
        if not df.empty:
            all_data.append(df)

    if not all_data:
        raise ValueError("未获取到有效数据，检查日期或Token")

    daily_data = pd.concat(all_data, ignore_index=True)
    daily_data['trade_date'] = pd.to_datetime(daily_data['trade_date'])
    daily_data = daily_data.sort_values(['ts_code', 'trade_date'])

    # 计算日收益率
    daily_data['daily_return'] = daily_data.groupby('ts_code')['close'].pct_change()
    return daily_data.dropna(subset=['daily_return'])


# 4. 计算Alpha#1因子
def calculate_alpha1(daily_data):
    df = daily_data.copy()

    # 步骤1：条件判断（收益率<0用20日标准差，否则用收盘价）
    df['roll_std_20d'] = df.groupby('ts_code')['daily_return'].rolling(20, min_periods=10).std().reset_index(level=0,
                                                                                                             drop=True)
    df['condition_val'] = np.where(df['daily_return'] < 0, df['roll_std_20d'], df['close'])

    # 步骤2：带符号幂运算
    df['signed_power'] = np.sign(df['condition_val']) * (np.abs(df['condition_val']) ** 2)

    # 步骤3：过去5天最大值位置
    def argmax_window(window):
        return window.argmax()

    df['ts_argmax'] = df.groupby('ts_code')['signed_power'].rolling(5, min_periods=3).apply(argmax_window).reset_index(
        level=0, drop=True)

    # 步骤4：横截面排名+因子计算
    df['rank_pct'] = df.groupby('trade_date')['ts_argmax'].rank(pct=True)
    df['alpha1'] = df['rank_pct'] - 0.5

    return df[['ts_code', 'trade_date', 'alpha1', 'open', 'high', 'low', 'close', 'vol', 'daily_return']]


# 5. 筛选Top50股票
def select_top50_stocks(pro, alpha_data, select_date):
    select_dt = pd.to_datetime(select_date)
    date_alpha = alpha_data[alpha_data['trade_date'] == select_dt].copy()

    if date_alpha.empty:
        raise ValueError(f"{select_date}无数据，可能为非交易日")

    # 选Alpha1最高的50只股票
    top50 = date_alpha.nlargest(50, 'alpha1')[['ts_code', 'alpha1']]
    # 补充股票名称
    stock_names = pro.stock_basic(ts_code=','.join(top50['ts_code']), fields='ts_code, name')

    result = top50.merge(stock_names, on='ts_code', how='left')
    result = result.sort_values('alpha1', ascending=False).reset_index(drop=True)
    result['rank'] = range(1, 51)  # 1-50排名
    return result[['rank', 'ts_code', 'name', 'alpha1']]


# 6. 优化绘图函数：使用mplfinance增强金融图表
def plot_top50_charts(top50_stocks, alpha_data, daily_data, select_date, start_date):
    """
    绘制4个核心图表：
    1. 平均价格K线图（mplfinance专业金融风格）
    2. Top5股票K线图对比（示例）
    3. Top50股票累计收益率曲线
    4. Alpha1因子分布直方图
    """
    top50_codes = top50_stocks['ts_code'].tolist()
    top5_codes = top50_codes[:5]  # 取前5名做K线对比
    date_filter = (daily_data['trade_date'] >= start_date) & (daily_data['trade_date'] <= select_date)
    plot_data = daily_data[date_filter & daily_data['ts_code'].isin(top50_codes)].copy()
    factor_data = alpha_data[alpha_data['ts_code'].isin(top50_codes) & (alpha_data['trade_date'] == select_date)]

    # 创建4子图布局
    fig = plt.figure(figsize=(20, 18))
    gs = fig.add_gridspec(3, 2, height_ratios=[1.2, 1.2, 1], width_ratios=[1, 1])

    ax1 = fig.add_subplot(gs[0, :])  # 平均价格K线图（全宽）
    ax2 = fig.add_subplot(gs[1, 0])  # Top1股票K线
    ax3 = fig.add_subplot(gs[1, 1])  # Top2股票K线
    ax4 = fig.add_subplot(gs[2, 0])  # 累计收益率曲线
    ax5 = fig.add_subplot(gs[2, 1])  # 因子分布直方图

    fig.suptitle(f'Top50股票分析图表 ({start_date.strftime("%Y-%m-%d")}至{select_date.strftime("%Y-%m-%d")})',
                 fontsize=18, fontweight='bold', y=0.98)

    # ---------------------- 子图1：平均价格K线图（mplfinance） ----------------------
    # 计算平均OHLC
    mean_ohlc = plot_data.groupby('trade_date').agg({
        'open': 'mean',
        'high': 'mean',
        'low': 'mean',
        'close': 'mean',
        'vol': 'sum'
    }).reset_index()

    # 转换为mplfinance所需格式
    mean_ohlc['trade_date'] = mean_ohlc['trade_date'].apply(mdates.date2num)
    ohlc_data = mean_ohlc[['trade_date', 'open', 'high', 'low', 'close', 'vol']].values

    ax1.set_title('Top50股票平均价格K线图', fontsize=14, fontweight='bold', pad=15)
    candlestick_ohlc(ax1, ohlc_data, width=0.6, colorup='red', colordown='green', alpha=0.8)
    ax1.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
    ax1.xaxis.set_major_locator(mdates.WeekdayLocator(byweekday=mdates.MO))
    plt.setp(ax1.xaxis.get_majorticklabels(), rotation=45)
    ax1.set_ylabel('价格（元）', fontsize=12)
    ax1.grid(True, linestyle='--')
    ax1.annotate('数据来源: Tushare', xy=(0.02, 0.02), xycoords='axes fraction', fontsize=10,
                 bbox=dict(facecolor='white', alpha=0.8))

    # ---------------------- 子图2-3：Top2股票K线图（mplfinance） ----------------------
    for i, (code, ax) in enumerate(zip(top5_codes[:2], [ax2, ax3])):
        stock_data = plot_data[plot_data['ts_code'] == code].copy()
        stock_name = top50_stocks[top50_stocks['ts_code'] == code]['name'].values[0]

        stock_data['trade_date'] = stock_data['trade_date'].apply(mdates.date2num)
        stock_ohlc = stock_data[['trade_date', 'open', 'high', 'low', 'close', 'vol']].values

        ax.set_title(f'Top{i + 1}: {stock_name}({code}) K线图', fontsize=12, fontweight='bold', pad=15)
        candlestick_ohlc(ax, stock_ohlc, width=0.6, colorup='red', colordown='green', alpha=0.8)
        ax.xaxis.set_major_formatter(mdates.DateFormatter('%m-%d'))
        ax.set_ylabel('价格（元）', fontsize=10)
        ax.grid(True, linestyle='--')

    # ---------------------- 子图4：Top50股票累计收益率曲线 ----------------------
    ax4.set_title('Top50股票累计收益率曲线', fontsize=14, fontweight='bold', pad=15)

    # 计算累计收益率
    cum_returns = []
    for code in top50_codes:
        stock_return = plot_data[plot_data['ts_code'] == code].copy()
        stock_return['cum_return'] = (1 + stock_return['daily_return']).cumprod() - 1
        cum_returns.append(stock_return[['trade_date', 'cum_return']].rename(columns={'cum_return': code}))

    return_df = cum_returns[0]
    for df in cum_returns[1:]:
        return_df = return_df.merge(df, on='trade_date', how='outer')

    # 绘制所有股票曲线
    colors = plt.cm.viridis(np.linspace(0.3, 0.9, 50))
    for i, code in enumerate(top50_codes):
        ax4.plot(return_df['trade_date'], return_df[code] * 100, color=colors[i], linewidth=1, alpha=0.6)

    # 平均收益率曲线
    return_df['mean_cum'] = return_df[top50_codes].mean(axis=1)
    ax4.plot(return_df['trade_date'], return_df['mean_cum'] * 100, color='#E74C3C',
             linewidth=2.5, label='Top50平均收益率', alpha=0.9)

    # 基准线和标注
    ax4.axhline(y=0, color='black', linewidth=1, linestyle='--', alpha=0.5)
    ax4.xaxis.set_major_formatter(mdates.DateFormatter('%m-%d'))
    ax4.set_xlabel('日期', fontsize=12)
    ax4.set_ylabel('累计收益率（%）', fontsize=12)
    ax4.legend(loc='upper right', fontsize=10)

    # 添加最终收益率标注
    final_date = return_df['trade_date'].iloc[-1]
    final_mean = return_df['mean_cum'].iloc[-1] * 100
    ax4.annotate(f'最终平均收益率: {final_mean:.2f}%',
                 xy=(final_date, final_mean),
                 xytext=(10, 0), textcoords='offset points',
                 bbox=dict(facecolor='yellow', alpha=0.5),
                 arrowprops=dict(arrowstyle='->'))

    # ---------------------- 子图5：Alpha1因子分布直方图 ----------------------
    ax5.set_title('Top50股票Alpha1因子分布', fontsize=14, fontweight='bold', pad=15)
    alpha_values = factor_data['alpha1'].dropna()

    # 统计指标
    mean_alpha = alpha_values.mean()
    median_alpha = alpha_values.median()
    std_alpha = alpha_values.std()

    # 绘制带渐变的直方图
    n, bins, patches = ax5.hist(alpha_values, bins=15, alpha=0.8, edgecolor='white', linewidth=0.8)
    cmap = LinearSegmentedColormap.from_list('alpha_cmap', ['#64B5F6', '#1565C0'])
    for i, patch in enumerate(patches):
        patch.set_facecolor(cmap(i / len(patches)))

    # 统计信息标注
    stats_text = f'均值: {mean_alpha:.4f}\n中位数: {median_alpha:.4f}\n标准差: {std_alpha:.4f}'
    ax5.text(0.05, 0.95, stats_text, transform=ax5.transAxes, fontsize=11,
             verticalalignment='top', bbox=dict(boxstyle='round', facecolor='lightgrey', alpha=0.8))

    # 均值和中位数线
    ax5.axvline(mean_alpha, color='#D32F2F', linewidth=2, linestyle='-', label=f'均值')
    ax5.axvline(median_alpha, color='#1976D2', linewidth=2, linestyle='--', label=f'中位数')
    ax5.set_xlabel('Alpha1因子值', fontsize=12)
    ax5.set_ylabel('股票数量', fontsize=12)
    ax5.legend(loc='upper right', fontsize=10)

    # 调整布局
    plt.tight_layout()
    plt.subplots_adjust(top=0.94, hspace=0.3, wspace=0.25)

    # 保存图表
    plot_file = f"top50_stock_analysis_{select_date.strftime('%Y%m%d')}.png"
    plt.savefig(plot_file, dpi=300, bbox_inches='tight', facecolor='white')
    plt.close()
    print(f"\n✅ 图表已保存至：{plot_file}")


# 主函数
def main():
    # 配置参数
    TUSHARE_TOKEN = ""  # 替换为个人Token
    SELECT_DATE = "2024-05-31"  # 选股日期（交易日）
    USE_HS300 = False  # 不使用沪深300，全市场筛选60/0开头
    HISTORY_DAYS = 60  # 历史数据范围（60天）

    try:
        # 1. 初始化API
        print("1. 初始化Tushare接口...")
        pro = init_tushare(TUSHARE_TOKEN)
        print("✅ 接口初始化成功")

        # 2. 计算时间范围
        select_dt = datetime.strptime(SELECT_DATE, '%Y-%m-%d')
        start_dt = select_dt - timedelta(days=HISTORY_DAYS)
        print(f"\n2. 数据范围：{start_dt.strftime('%Y-%m-%d')} 至 {SELECT_DATE}")

        # 3. 获取筛选股票池
        print("\n3. 获取股票池（仅60/0开头）...")
        stock_pool = get_filtered_stock_pool(pro, is_hs300=USE_HS300)
        print(f"✅ 筛选后股票池：{len(stock_pool)}只")

        # 4. 获取日线数据
        print("\n4. 获取日线数据...")
        daily_data = get_daily_data(pro, stock_pool, start_dt, select_dt)
        covered_stocks = daily_data['ts_code'].nunique()
        print(f"✅ 成功获取{covered_stocks}只股票数据")

        # 5. 计算Alpha1因子
        print("\n5. 计算Alpha1因子...")
        alpha_data = calculate_alpha1(daily_data)
        print(f"✅ 因子计算完成，覆盖{alpha_data['trade_date'].nunique()}个交易日")

        # 6. 筛选Top50股票
        print("\n6. 筛选Top50股票...")
        top50_stocks = select_top50_stocks(pro, alpha_data, SELECT_DATE)
        print("\n" + "=" * 80)
        print(f"Top50选股结果（{SELECT_DATE}）")
        print("=" * 80)
        print(top50_stocks.to_string(index=False))

        # 7. 保存结果到Excel
        excel_file = f"top50_stocks_{SELECT_DATE.replace('-', '')}.xlsx"
        top50_stocks.to_excel(excel_file, index=False)
        print(f"\n💾 选股结果已保存至：{excel_file}")

        # 8. 绘制分析图表
        print("\n8. 生成Top50股票分析图表...")
        plot_top50_charts(top50_stocks, alpha_data, daily_data, select_dt, start_dt)

        print("\n" + "=" * 80)
        print("🎉 全流程执行完成！")
        print("=" * 80)

    except Exception as e:
        print(f"\n❌ 运行错误：{str(e)}")
        print("🔍 建议检查：1）Token有效性；2）选股日期是否为交易日；3）网络")


if __name__ == "__main__":
    main()