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()