代码传递

from jqdata import *
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import datetime
import warnings
import numpy as np
from tqdm import tqdm
import gc
plt.rcParams['font.sans-serif'] = ['SimHei']  # 中文字体设置-黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决保存图像是负号'-'显示为方块的问题
# 过滤科创版股票
def filter_gem_stock(stock_list):
    """过滤掉科创板股票"""
    return [stock for stock in stock_list if not stock.startswith(('68', '4', '8'))]

# 过滤新股次新股
def filter_new_stock(end_date, stock_list):
    """过滤上市小于240天的股票"""
    return [stock for stock in stock_list if (end_date - datetime.timedelta(days=240)) > get_security_info(stock).start_date]

# 过滤ST股
def filter_st_stock(end_date, stocks):
    """过滤掉ST股票"""
    df_stocks = get_extras('is_st', stocks, end_date=end_date, df=True, count=1).T
    df_stocks = df_stocks[df_stocks == False].dropna()
    return df_stocks.index.tolist()

# 计算微盘股指数
def cal_smallcap_index(start_date, end_date, n=400):
    """计算微盘股指数及成分股数据"""
    trade_days = get_trade_days(start_date=start_date, end_date=end_date)
    df_index = []
    stock_data = {}
    
    for day in tqdm(trade_days, desc="计算微盘股指数"):
        # 获取当前交易日所有股票
        stock_list = get_all_securities(date=day).index.tolist()
        
        # 过滤股票
        stock_list = filter_gem_stock(stock_list)
        stock_list = filter_new_stock(day, stock_list)
        stock_list = filter_st_stock(day, stock_list)
        
        # 按市值排序，选取市值最小的n只股票
        tmp = get_valuation(stock_list, count=1, end_date=day, fields=['market_cap']).sort_values('market_cap', ascending=True)
        microcap_stocks = tmp.iloc[:n]['code'].tolist()
        
        # 获取当日价格数据
        daily_prices = get_price(microcap_stocks, end_date=day, frequency='1d', fields=['close'], count=1, panel=False)
        
        # 计算等权指数
        if not daily_prices.empty:
            # 计算非零价格的股票数量
            valid_stocks = daily_prices[daily_prices['close'] > 0]
            if len(valid_stocks) == 0:
                continue
                
            # 计算等权重平均值
            avg_close = valid_stocks['close'].mean()
            
            df_index.append({
                'time': day,
                'index_close': avg_close,
                'index_num_stocks': len(microcap_stocks),
                'stocks': microcap_stocks
            })
            
            # 存储成分股数据
            stock_data[day] = microcap_stocks
            
        gc.collect()  # 手动垃圾回收
    
    # 创建微盘股指数DataFrame
    if df_index:
        df_index = pd.DataFrame(df_index).set_index('time')
        # 添加百分比变化
        df_index['index_pct_change'] = df_index['index_close'].pct_change() * 100
    else:
        df_index = pd.DataFrame()
    
    return df_index, stock_data

# 计算微盘股市场宽度
def calculate_microcap_breadth(stock_data, days=20):
    """计算微盘股市场宽度指标"""
    if not stock_data:
        return pd.DataFrame()
    
    # 获取交易日列表（按日期排序）
    dates = sorted(stock_data.keys())
    df_breadth = pd.DataFrame(index=dates, columns=['breadth_num_stocks', 'breadth_above_ma20'])
    
    for i, date in enumerate(tqdm(dates, desc="计算市场宽度")):
        if i < days:  # 需要足够的天数来计算MA20
            continue
            
        # 获取当前微盘股成分股
        stocks = stock_data[date]
        
        # 获取这些股票的收盘价数据（包含MA计算所需天数）
        prices = get_price(stocks, end_date=date, frequency='1d', fields=['close'], count=days+1, panel=False)
        
        if prices.empty:
            df_breadth.loc[date, 'breadth_num_stocks'] = len(stocks)
            df_breadth.loc[date, 'breadth_above_ma20'] = np.nan
            continue
        
        # 过滤掉无效价格
        prices = prices[prices['close'] > 0]
        
        # 计算每只股票的20日移动平均线
        # 创建宽格式数据框
        df_prices = prices.pivot(index='time', columns='code', values='close')
        
        # 检查是否有足够的数据点
        if len(df_prices) < days:
            df_breadth.loc[date, 'breadth_num_stocks'] = len(stocks)
            df_breadth.loc[date, 'breadth_above_ma20'] = np.nan
            continue
        
        # 计算移动平均
        ma20 = df_prices.rolling(window=days).mean()
        
        # 获取最新收盘价
        latest_prices = df_prices.iloc[-1]
        
        # 计算股价高于20日均线的股票数量
        above_ma20 = (latest_prices > ma20.iloc[-1]).sum()
        
        # 存储结果
        df_breadth.loc[date, 'breadth_num_stocks'] = len(stocks)
        df_breadth.loc[date, 'breadth_above_ma20'] = above_ma20
    
    # 清理数据
    df_breadth = df_breadth.dropna()
    if df_breadth.empty:
        return pd.DataFrame()
    
    # 计算宽度百分比
    df_breadth['breadth_pct'] = (df_breadth['breadth_above_ma20'] / df_breadth['breadth_num_stocks']) * 100
    
    # 计算宽度变化
    df_breadth['breadth_change'] = df_breadth['breadth_pct'].diff()
    
    return df_breadth

# 可视化市场宽度
def visualize_microcap_breadth(microcap_index, breadth_data):
    """可视化微盘股指数和市场宽度"""
    if microcap_index.empty or breadth_data.empty:
        print("没有足够的数据进行可视化")
        return
        
    # 确保索引类型一致
    if not microcap_index.index.equals(breadth_data.index):
        # 只保留共同的日期
        common_dates = microcap_index.index.intersection(breadth_data.index)
        microcap_index = microcap_index.loc[common_dates]
        breadth_data = breadth_data.loc[common_dates]
    
    plt.figure(figsize=(16, 12))
    
    # 微盘股指数走势
    plt.subplot(3, 1, 1)
    plt.plot(microcap_index.index, microcap_index['index_close'], label='微盘股指数')
    plt.title('微盘股指数走势')
    plt.ylabel('指数值')
    plt.grid(True)
    
    # 添加百分比变化 - 红涨绿跌
    ax2 = plt.twinx()
    colors = np.where(microcap_index['index_pct_change'] >= 0, 'red', 'green')
    ax2.bar(microcap_index.index, microcap_index['index_pct_change'], 
            color=colors, alpha=0.3, label='涨跌幅')
    plt.ylabel('涨跌幅(%)')
    
    # 合并图例
    lines, labels = plt.gca().get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    plt.legend(lines + lines2, labels + labels2, loc='upper left')
    
    # 市场宽度 - 使用红色表示强势，绿色表示弱势
    plt.subplot(3, 1, 2)
    # 确保数据为数值类型且有限
    breadth_pct = pd.to_numeric(breadth_data['breadth_pct'], errors='coerce').replace([np.inf, -np.inf], np.nan).dropna()
    
    # 创建颜色数组：宽度大于80%用红色，小于20%用绿色，其他用蓝色
    line_colors = []
    for pct in breadth_pct.values:
        if pct > 80:
            line_colors.append('red')  # 强势 - 红色
        elif pct < 20:
            line_colors.append('green')  # 弱势 - 绿色
        else:
            line_colors.append('blue')  # 中性 - 蓝色
    
    # 绘制线段连接的点
    for i in range(len(breadth_pct) - 1):
        plt.plot(breadth_pct.index[i:i+2], breadth_pct.values[i:i+2], 
                 color=line_colors[i], zorder=5)
    
    # 填充区域 - 使用相同的颜色逻辑
    for i in range(len(breadth_pct)):
        if i < len(breadth_pct) - 1:
            start_idx = breadth_pct.index[i]
            end_idx = breadth_pct.index[i+1]
            avg_val = (breadth_pct.loc[start_idx] + breadth_pct.loc[end_idx]) / 2
            
            # 确定填充颜色
            fill_color = 'red' if avg_val > 80 else 'green' if avg_val < 20 else 'blue'
            fill_alpha = 0.1
            
            # 填充区域
            plt.fill_between(
                [start_idx, end_idx],
                [0, 0],
                [breadth_pct.loc[start_idx], breadth_pct.loc[end_idx]],
                color=fill_color,
                alpha=fill_alpha
            )
    
    # 添加中性区域的标记线
    plt.axhline(40, color='green', linestyle='-.', alpha=0.5, zorder=3)
    plt.axhline(60, color='red', linestyle='-.', alpha=0.5, zorder=3)
    plt.axhline(50, color='gray', linestyle='--', alpha=0.7, zorder=2)
    
    # 添加图例
    plt.plot([], [], color='red', label='强势 (>80%)')
    plt.plot([], [], color='blue', label='中性 (20%-80%)')
    plt.plot([], [], color='green', label='弱势 (<20%)')
    
    plt.title('微盘股市场宽度')
    plt.ylabel('百分比(%)')
    plt.grid(True)
    plt.legend(loc='upper left')
    
    # 宽度变化
    ax2 = plt.twinx()
    breadth_change = pd.to_numeric(breadth_data['breadth_change'], errors='coerce').replace([np.inf, -np.inf], np.nan).dropna()
    
    if not breadth_change.empty:
        ax2.bar(breadth_change.index, breadth_change, 
                color=np.where(breadth_change > 0, 'blue', 'red'),  # 正变蓝，负变红
                alpha=0.3, label='宽度变化', zorder=3)
        plt.ylabel('变化')
    
    # 热力图 - 
    plt.subplot(3, 1, 3)
    if len(breadth_pct) >= 10:
        # 创建热度图的数据
        recent_dates = sorted(breadth_pct.index, reverse=False)[-30:]
        recent_data = breadth_pct.loc[recent_dates]
        
        # 创建2D数据
        grid_data = np.tile(recent_data.values, (5, 1))
        
        # 自定义热力图颜色映射 - 
        from matplotlib.colors import LinearSegmentedColormap
        cmap_custom = LinearSegmentedColormap.from_list("rg", ["green", "yellow", "red"], N=256)
        
        # 绘制热度图
        plt.imshow(grid_data, cmap=cmap_custom, aspect='auto', vmin=0, vmax=100)
        cbar = plt.colorbar(label='市场宽度(%)', shrink=0.8)
        cbar.set_label('市场宽度(%)', fontsize=12)
        
        # 添加数值标注
        for i, (date, val) in enumerate(zip(recent_dates, recent_data)):
            color = 'white' if val < 30 or val > 70 else 'black' 
            plt.text(i, 2, f'{val:.1f}%', ha='center', va='center', 
                     color=color, fontsize=10)
        
        # 设置坐标轴
        plt.xticks(range(len(recent_dates)), [d.strftime('%m-%d') for d in recent_dates], rotation=45)
        plt.yticks([])
        plt.title('近期市场宽度动态')
    else:
        plt.text(0.5, 0.5, '没有足够的数据生成热力图', 
                 ha='center', va='center', fontsize=12)
    
    plt.tight_layout()
    plt.show()


# 设置基本参数
today = datetime.date.today()
start_date = today - datetime.timedelta(days=count_)  
ma_days = 20  # 移动平均天数

print(f"开始微盘股市场宽度分析 ({start_date} 至 {today})")

# 计算微盘股指数
print("计算微盘股指数...")
microcap_index, stock_data = cal_smallcap_index(start_date, today,n)

if microcap_index.empty:
    print("警告：微盘股指数数据为空，请检查股票筛选条件")
    exit()

print(f"获得 {len(microcap_index)} 天的微盘股指数数据")

# 计算市场宽度
print("计算市场宽度指标...")
if stock_data:
    print(f"共有 {len(stock_data)} 天成分股数据可用于宽度计算")
    breadth_data = calculate_microcap_breadth(stock_data, days=ma_days)
else:
    print("警告：没有足够的微盘股成分股数据计算市场宽度")
    breadth_data = pd.DataFrame()