2-深度学习挖短线股-2-训练数据计算

2  训练数据计算

        选择2017年12月31日及之前的数据作为训练数据，计算训练数据的输入特征及分类标签。将10个交易内上涨幅度超过10%的股票标记为类别1，其他股票标记为类别0。为了给深度学习模型提供更多的输入维度，可以基于日线数据的开盘价、收盘价、最高价、最低价、成交量等进行扩展。这里只扩展5、10、20、30、60、120、240均线及均量线数据，后续可根据需要计算MACD、KDJ、RSI等指标数据，进一步扩展输入维度。

         程序的主要处理流程是：首先定义了买入信号的计算逻辑，即未来 N 个交易日内涨幅达到 INC_PER 则标记为买入信号；然后读取股票代码列表，逐个处理每只股票的数据，计算多种移动平均线指标；最后根据预设条件生成买入信号，并将所有处理后的数据保存到新文件中。注释部分还列出了其他可扩展的技术指标计算函数，方便后续扩展使用。

# -*- coding: utf-8 -*-
"""
Created on Thu Jun  5 08:54:22 2025

@author: Administrator
"""
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)
import datetime  # 导入datetime模块用于日期时间操作
import os.path  # 导入os.path模块用于路径处理
import sys  # 导入sys模块用于获取脚本路径等系统信息
import pandas as pd  # 导入pandas模块用于数据处理
import numpy as np  # 导入numpy模块用于数值计算

# 引入topq_talib，计算技术指标
script_path = os.path.dirname(os.path.abspath(sys.argv[0]))  # 获取当前脚本路径
talib_path = os.path.join(script_path, '../../TQDat/TQDown2020v1/topqt/')  # 计算topq_talib模块路径
sys.path.append(talib_path)  # 将topq_talib模块路径添加到系统路径
import topq_talib1 as tt  # 导入topq_talib1模块并命名为tt

# 类别标准：N个交易日内上涨INC_PER
N = 10  # 定义预测周期为10个交易日
INC_PER = 0.1  # 定义上涨阈值为10%

# n 天内上涨inc_per
def calc_buy_signal(df, n, inc_per):
    df = df[df.notnull().T.all()]  # 移除包含空值的列
    buy = np.zeros(df.shape[0])  # 创建全0数组用于存储买入信号
    df.insert(df.shape[1], 'buy', buy)  # 在DataFrame中插入buy列，初始值为0
    df = df.copy()  # 复制DataFrame防止SettingWithCopyWarning
    
    row_i = 0
    for row_i in range(df.shape[0]):  # 遍历DataFrame的每一行
        for n_i in range(n):  # 遍历未来n个交易日
            # 检查是否超出数据范围，并且是否满足上涨条件
            if row_i + n_i + 1 < df.shape[0] and \
                df.iloc[row_i + n_i + 1].at['close'] - df.iloc[row_i].at['close'] > inc_per * df.iloc[row_i].at['close']:
                df.loc[df.index[row_i], 'buy'] = 1  # 满足条件则标记为买入信号
                break  # 找到第一个满足条件的交易日即退出内层循环
    return df  # 返回添加了买入信号的DataFrame

# 读入股票代码
stk_code_file = './stk_data/dp_stock_list.csv'  # 定义股票代码文件路径
stk_pools = pd.read_csv(stk_code_file, encoding='gbk')  # 读取股票代码文件

# 对每个股票添加衍生数据
for stk_code in stk_pools['code']:  # 遍历每只股票
    # 读入数据
    input_file = './stk_data/d/' + stk_code + '.csv'  # 定义股票数据文件路径
    if not os.path.exists(input_file):  # 检查文件是否存在
        continue  # 文件不存在则跳过当前循环
    
    df = pd.read_csv(input_file, index_col=0)  # 读取股票数据，设置index列为索引
    df = df.sort_index(ascending=True)  # 按索引升序排序
    
    # MA
    ma_list = [5, 10, 20, 30, 60, 120, 240]  # 定义移动平均线周期列表
    for i in ma_list:  # 遍历每个周期
        df = tt.MA_n(df, i)  # 计算i周期的移动平均线并添加到DataFrame
    
    # vol_MA
    vol_ma_list = [5, 10, 20, 30, 60, 120, 240]  # 定义成交量移动平均线周期列表
    for i in vol_ma_list:  # 遍历每个周期
        df = tt.vol_MA_n(df, i)  # 计算i周期的成交量移动平均线并添加到DataFrame
    
    '''可扩展指标
    #BBANDS
    df = tt.BBANDS_UpLow(df, 20)
    #MACD
    df = tt.MACD020(df, 12, 26)
    #KDJ
    df = tt.KDJ(df, 9, 3)
    #RSI
    df = tt.RSI(df, 14)
    # A/D
    ad_list = [5, 10, 20, 30, 60, 120, 240]
    for i in ad_list:
        df = tt.ACCDIST(df, i)
    # ATRcd
    atr_list = [5, 10, 20, 30, 60, 120, 240]
    for i in atr_list:
        df = tt.ATR(df, i)
    '''
    
    df = calc_buy_signal(df, N, INC_PER)  # 计算买入信号
    
    # 写出文件
    output_file = './baostock/data_ext/' + stk_code + '.csv'  # 定义输出文件路径
    df.to_csv(output_file)  # 将处理后的数据保存为CSV文件
    print(stk_code + ' done!')  # 打印完成提示