股票强度计算错误

您好，您的代码在运行时出现了KeyError: 'ma5'的错误，这是因为在select_stock函数中，您在访问qualified_stocks列表中的元素时，字典键的层次结构有误。

具体来说，在以下代码中：

python
features = np.array([[s['strength']['ma5'], s['strength']['ma13'], s['strength']['ma20'], s['strength']['ma60']] for s in qualified_stocks])

这里的s['strength']实际上是整个stock_info字典，而不是直接的strength字典，因此需要进一步访问s['strength']['strength']。

为了解决这个问题，我们需要在将股票信息添加到qualified_stocks列表时，确保'strength'键对应的是stock_info['strength']，而不是整个stock_info。修改后的代码如下：

python
qualified_stocks.append({
    'code': stock,
    'strength': stock_info['strength'],
    'buy_signals': buy_signals,
})

以下是更新后的完整代码：

python
from jqdata import *
import numpy as np
import pandas as pd
from sklearn.ensemble import GradientBoostingClassifier

def initialize(context):
    # 设置基准
    set_benchmark('000300.XSHG')
    # 设置滑点
    set_slippage(PriceRelatedSlippage(0.002))
    # 设置手续费
    set_order_cost(OrderCost(open_tax=0, close_tax=0.001, open_commission=0.0003, close_commission=0.0003, min_commission=5), type='stock')
    # 设置运行时间
    run_daily(select_stock, time='15:30', reference_security='000300.XSHG')
    run_daily(market_open, time='9:30', reference_security='000300.XSHG')
    run_daily(market_close, time='14:58', reference_security='000300.XSHG')
    # 初始化全局变量
    g.current_stock = None
    g.selected_stock = None
    g.holding_stocks = {}
    g.max_hold_days = 2

def calculate_strength(close_data):
    """计算强度得分"""
    ma5 = np.mean(close_data[-5:])
    ma13 = np.mean(close_data[-13:])
    ma20 = np.mean(close_data[-20:])
    ma60 = np.mean(close_data[-60:])
    
    return {
        'ma5': ma5,
        'ma13': ma13,
        'ma20': ma20,
        'ma60': ma60
    }

def get_stock_strength(context):
    """获取股票强度数据"""
    stocks = list(get_all_securities(['stock']).index)
    stocks = [s for s in stocks if not (s.startswith(('688', '689')) or 'ST' in get_security_info(s).display_name)]
    strength_scores = {}

    for stock in stocks:
        try:
            # 获取过去121个交易日的收盘价
            prices = get_price(stock, count=121, end_date=context.current_dt, fields=['close'], skip_paused=False, fq='pre')
            close = prices['close'].values
            
            # 计算强度得分
            strength = calculate_strength(close)

            # 计算强度历史和日变化
            strength_history = []
            daily_changes = []
            for i in range(60):
                if i > 0:
                    day_close = close[:-i]
                else:
                    day_close = close
                day_strength = calculate_strength(day_close)
                strength_history.append({
                    'date': prices.index[-(i + 1)],
                    'strength': day_strength['ma5']  # 使用ma5作为强度得分
                })
                if i > 0:
                    daily_change = strength_history[-1]['strength'] - strength_history[-2]['strength']
                    daily_changes.append({
                        'date': prices.index[-(i + 1)],
                        'change': daily_change
                    })

            strength_scores[stock] = {
                'strength': strength,
                'strength_history': strength_history,
                'daily_changes': daily_changes
            }
        except Exception as e:
            continue

    return strength_scores

def check_buy_signals(stock_data):
    """检查股票是否满足买入信号"""
    buy_signals = []
    strength_history = [d['strength'] for d in stock_data['strength_history']]

    # 早期信号条件
    if stock_data['strength']['ma5'] > 12.5 and stock_data['strength']['ma5'] < 13:
        daily_changes = [d['change'] for d in stock_data['daily_changes'][:2]]
        if all(dc > 0 for dc in daily_changes) and all(1.0 <= dc <= 1.5 for dc in daily_changes):
            buy_signals.append("接近关键值的趋势判断")

    # 买入信号1: 强度得分突破前的增长趋势
    if len(strength_history) >= 4:
        growth_rates = [(strength_history[i] - strength_history[i + 1]) / strength_history[i + 1] * 100 for i in range(3)]
        daily_changes = [strength_history[i] - strength_history[i + 1] for i in range(3)]
        if all(0.5 <= gr <= 0.8 for gr in growth_rates) and all(0.5 <= dc <= 1.5 for dc in daily_changes):
            buy_signals.append("强度得分突破前的增长趋势")

    # 买入信号2: 强度得分即将突破关键值（13）
    if stock_data['strength']['ma5'] > 12.5:
        if len(stock_data['daily_changes']) >= 2:
            daily_changes = [d['change'] for d in stock_data['daily_changes'][:2]]
            if all(1.0 <= dc <= 1.5 for dc in daily_changes):
                buy_signals.append("强度得分即将突破关键值（13）")

    # 买入信号3: 强度连续增长趋势形成
    if len(stock_data['daily_changes']) >= 3:
        daily_changes = [d['change'] for d in stock_data['daily_changes'][:3]]
        if all(dc > 0 for dc in daily_changes):
            if daily_changes[0] < daily_changes[1] < daily_changes[2]:
                if daily_changes[0] >= 0.3 and daily_changes[2] > 1.0:
                    buy_signals.append("强度连续增长趋势形成")

    # 买入信号4: 当前强度与前一天强度对比
    if stock_data['strength']['ma5'] > stock_data['strength']['ma5']:
        if stock_data['strength']['ma5'] - stock_data['strength']['ma5'] > 0.5:
            buy_signals.append("当前强度与前一天强度对比")

    return buy_signals

def check_sell_signals(stock_data):
    """检查股票是否满足卖出信号"""
    sell_signals = []

    # 卖出信号1: 强度得分达到14.5以上后的衰退迹象
    if stock_data['strength']['ma5'] > 14.5:
        if stock_data['daily_changes'][0]['change'] <= -0.5:
            sell_signals.append("强度得分达到14.5以上后的衰退迹象")

    # 卖出信号2: 日强度变化连续2天为负，且平均值低于-0.5
    if len(stock_data['daily_changes']) >= 2:
        daily_changes = [d['change'] for d in stock_data['daily_changes'][:2]]
        if all(dc < 0 for dc in daily_changes):
            if np.mean(daily_changes) < -0.5:
                sell_signals.append("日强度变化连续2天为负，平均值低于-0.5")

    # 卖出信号3: 高点滞涨信号
    if len(stock_data['strength_history']) >= 3:
        growth_rates = [(stock_data['strength_history'][i]['strength'] - stock_data['strength_history'][i + 1]['strength']) / stock_data['strength_history'][i + 1]['strength'] * 100 for i in range(2)]
        daily_changes = [stock_data['daily_changes'][i]['change'] for i in range(2)]
        if all(gr < 0.2 for gr in growth_rates) and all(-0.5 <= dc <= 0.5 for dc in daily_changes):
            sell_signals.append("高点滞涨信号")

    # 卖出信号4: 急剧回调信号
    if stock_data['daily_changes'][0]['change'] < -1.5:
        sell_signals.append("急剧回调信号")

    # 卖出信号5: 强度得分连续下降信号
    if len(stock_data['strength_history']) >= 3:
        decreases = all(stock_data['strength_history'][i]['strength'] < stock_data['strength_history'][i + 1]['strength'] for i in range(2))
        decreases_amount = all((stock_data['strength_history'][i + 1]['strength'] - stock_data['strength_history'][i]['strength']) > 0.8 for i in range(2))
        if decreases and decreases_amount:
            sell_signals.append("强度得分连续下降信号")

    return sell_signals

def select_stock(context):
    """选股函数"""
    strength_scores = get_stock_strength(context)
    qualified_stocks = []

    for stock in strength_scores:
        stock_info = strength_scores[stock]
        buy_signals = check_buy_signals(stock_info)

        if buy_signals:
            qualified_stocks.append({
                'code': stock,
                'strength': stock_info['strength'],  # 修改这里
                'buy_signals': buy_signals,
            })

    # 使用梯度提升树进行排序
    if qualified_stocks:
        # 假设我们有特征和标签
        features = np.array([[s['strength']['ma5'], s['strength']['ma13'], s['strength']['ma20'], s['strength']['ma60']] for s in qualified_stocks])
        labels = np.array([1 if '接近关键值的趋势判断' in s['buy_signals'] else 0 for s in qualified_stocks])
        
        model = GradientBoostingClassifier()
        model.fit(features, labels)
        predictions = model.predict(features)

        # 根据预测结果排序
        qualified_stocks = [s for s, p in zip(qualified_stocks, predictions) if p == 1]
        qualified_stocks.sort(key=lambda x: x['strength']['ma5'], reverse=True)

        if qualified_stocks:
            g.selected_stock = qualified_stocks[0]['code']
            print(f"选中的股票: {g.selected_stock}")

def market_open(context):
    """开盘时执行的函数"""
    positions = context.portfolio.positions
    position_count = len(positions)

    if g.selected_stock and g.selected_stock not in positions:
        if position_count < 10:  # 每日最多持有10只股票
            cash = context.portfolio.available_cash / (10 - position_count)
            if cash > 0:
                order_value(g.selected_stock, cash)
                g.holding_stocks[g.selected_stock] = context.current_dt.date()
                print(f"购买股票: {g.selected_stock}")

def market_close(context):
    """收盘前执行的函数"""
    positions = context.portfolio.positions
    to_sell = []

    for stock in positions:
        hold_days = (context.current_dt.date() - g.holding_stocks.get(stock, context.current_dt.date())).days + 1
        stock_data = get_bars(stock, count=1, unit='1d', fields=['close', 'high_limit'], include_now=False)

        if isinstance(stock_data, pd.DataFrame) and not stock_data.empty:
            close_price = stock_data['close'].iloc[-1]
            high_limit = stock_data['high_limit'].iloc[-1]

            if hold_days >= g.max_hold_days and close_price < high_limit:
                to_sell.append(stock)

    for stock in to_sell:
        order_target(stock, 0)
        g.holding_stocks.pop(stock, None)
        print(f"卖出股票: {stock}")

主要修改：

• 在select_stock函数中，修改了qualified_stocks.append部分，将'strength'键对应的值改为stock_info['strength']，即：

  python
  'strength': stock_info['strength'],

  这样，在后续访问s['strength']['ma5']时，就不会出现KeyError错误了。

希望以上修改能解决您的问题。