AI量化qlib学习笔记 一：清洗数据

最近关注了一个AI量化投资框架qlib，略做笔记，因为某些无法实现原因，不得不放弃实战。

首先这个框架仅仅是对特征有效，并没有其它关联，比如有些股是有时间规律的，比如农产品，还有一些有环境 规律 ，比如A股一般是权重先涨，稳住大盘后，换成小盘股唱戏。还有一些是政策类的，比如出台了某些可能是较长影响的政策利好某板块的。

当然这些可以自己在训练的时候当成因子写进去，这个就比较花时间和精力去做了。

首先把源码从github git下来 地址：https://github.com/microsoft/qlib

有个注意情况 ，项目主目录和源码目录不要用qlib做为目录 ，改个名，比如 qlib_src这样的。

是python环境 ，我用的是poetry 搞的虚拟 环境 ,后面所有都是在当前项目的虚拟环境中进行。

然后进入源码目录安装打包安装以及相关依赖  pip install -e .

然后开始相关准备，最开始是准备数据，它的示例中是从yahoo财经下载的，这不是开玩笑嘛，不过有个捷径，自行下载：https://github.com/chenditc/investment_data/releases

说实话，如果要是搞的不是A股，请自行准备数据，这个我花了很长时间，因为qlib对数据的要求太高了，或太多不支持了。

我从通达信获取的数据，不仅有日线，还有分时的，不得不说，通达信还是很不错的，无偿提供。

拿到数据，需要清洗，以下是清洗脚本：

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
简单的CSV数据清洗脚本 - 为qlib准备数据
"""

import pandas as pd
import os

def convert_filename(original_filename):
    """转换文件名格式：品种代码.市场代码.周期.csv -> 品种代码.交易所代码.周期.csv"""
    # 市场代码到交易所代码的映射
    market_mapping = {
        '30': 'SHFE',  # 上海期货交易所
        '28': 'CZCE',  # 郑州商品交易所
        '29': 'DCE',   # 大连商品交易所
        '31': 'INE',   # 上海国际能源交易中心
        '66': 'GFEX',   # 广州期货交易所
        '42': 'ZS'   # 指数
    }
    
    parts = original_filename.split('.')
    if len(parts) >= 4:  # 品种代码.市场代码.周期.csv
        symbol = parts[0]
        market_code = parts[1]
        period = parts[2]
        extension = parts[3]
        
        # 转换市场代码
        exchange_code = market_mapping.get(market_code, market_code)
        new_filename = f"{symbol}.{exchange_code}.{period}.{extension}"
        
        print(f"文件名转换: {original_filename} -> {new_filename}")
        return new_filename
    
    # 如果格式不匹配，返回原文件名
    return original_filename

def format_time_column(time_value):
    """格式化时间列：905 -> 09:05"""
    time_str = str(int(time_value)).zfill(4)  # 确保4位数字，前面补0
    return f"{time_str[:2]}:{time_str[2:]}"

def clean_csv_file(input_file, output_file):
    """清洗单个CSV文件"""
    print(f"正在处理: {input_file}")
    
    # 读取原始文件
    with open(input_file, 'r', encoding='gbk') as f:
        lines = f.readlines()
    
    # 去掉第一行（提示信息）和最后一行
    data_lines = lines[1:-1]
    
    # 处理标题行（第二行，现在是第一行）
    header_line = data_lines[0].strip()
    
    # 中文到英文的映射
    header_mapping = {
        '日期': 'date',
        '时间': 'time', 
        '开盘': 'open',
        '最高': 'high',
        '最低': 'low',
        '收盘': 'close',
        '成交量': 'volume',
        '持仓量': 'open_interest',
        '结算价': 'settlement'
    }
    
    # 将tab分隔的标题转为英文逗号分隔
    headers = [h.strip() for h in header_line.split('\t')]
    english_headers = [header_mapping.get(h, h) for h in headers]
    
    # 创建输出目录
    os.makedirs(os.path.dirname(output_file), exist_ok=True)
    
    # 先写入临时文件，然后格式化时间列
    temp_file = output_file + '.tmp'
    with open(temp_file, 'w', encoding='utf-8') as f:
        # 写入英文标题
        f.write(','.join(english_headers) + '\n')
        # 写入数据行
        for line in data_lines[1:]:
            f.write(line)
    
    # 读取临时文件并格式化时间列
    df = pd.read_csv(temp_file, encoding='utf-8')
    
    # 如果有时间列，格式化时间
    if 'time' in df.columns:
        df['time'] = df['time'].apply(format_time_column)
    
    # 关键修改：如果是分钟级数据（1m或5m），合并日期和时间字段为datetime
    filename = os.path.basename(output_file)
    is_minute_data = any(period in filename for period in ['1m', '5m'])
    
    if is_minute_data and 'date' in df.columns and 'time' in df.columns:
        period_type = '1分钟' if '1m' in filename else '5分钟'
        print(f"检测到{period_type}原始数据，正在合并日期和时间字段...")
        
        # 合并日期和时间
        df['datetime'] = pd.to_datetime(df['date'].astype(str) + ' ' + df['time'].astype(str))
        
        # 添加instrument列（从文件名提取）
        instrument_name = filename.split('.')[0] + '.' + filename.split('.')[1]  # 如RBL8.SHFE
        df['instrument'] = instrument_name
        
        # 重新排列列顺序，符合qlib标准格式
        qlib_columns = ['datetime', 'instrument']
        for col in ['open', 'high', 'low', 'close', 'volume']:
            if col in df.columns:
                qlib_columns.append(col)
        
        # 添加其他列
        for col in df.columns:
            if col not in qlib_columns and col not in ['date', 'time']:
                qlib_columns.append(col)
        
        df = df[qlib_columns]
        
        print(f"已合并datetime字段，数据形状: {df.shape}")
        print(f"时间范围: {df['datetime'].min()} 到 {df['datetime'].max()}")
    
    # 如果是日线数据，也需要处理datetime格式
    elif 'date' in df.columns and 'time' not in df.columns:
        print("检测到日线数据，正在处理日期格式...")
        
        # 日线数据只有日期，添加默认时间
        df['datetime'] = pd.to_datetime(df['date'].astype(str) + ' 15:00:00')
        
        # 添加instrument列
        filename = os.path.basename(output_file)
        instrument_name = filename.split('.')[0] + '.' + filename.split('.')[1]
        df['instrument'] = instrument_name
        
        # 重新排列列顺序
        qlib_columns = ['datetime', 'instrument']
        for col in ['open', 'high', 'low', 'close', 'volume']:
            if col in df.columns:
                qlib_columns.append(col)
        
        # 添加其他列
        for col in df.columns:
            if col not in qlib_columns and col not in ['date', 'time']:
                qlib_columns.append(col)
        
        df = df[qlib_columns]
        
        print(f"已处理日线数据，数据形状: {df.shape}")
        print(f"时间范围: {df['datetime'].min()} 到 {df['datetime'].max()}")
    
    # 保存最终文件
    df.to_csv(output_file, index=False, encoding='utf-8')
    
    # 删除临时文件
    os.remove(temp_file)
    
    print(f"已保存到: {output_file}")
    return output_file

def find_csv_files(directory='.'):
    """查找目录中的所有CSV文件"""
    csv_files = []
    for file in os.listdir(directory):
        if file.endswith('.csv') and os.path.isfile(file):
            csv_files.append(file)
    return sorted(csv_files)

def main():
    """主函数"""
    # 自动查找所有CSV文件
    csv_files = find_csv_files()
    
    if not csv_files:
        print("当前目录没有找到CSV文件")
        return
    
    print(f"找到 {len(csv_files)} 个CSV文件:")
    for file in csv_files:
        print(f"  - {file}")
    
    # 创建输出目录
    output_dir = 'cleaned_data'
    os.makedirs(output_dir, exist_ok=True)
    
    # 清洗所有文件
    cleaned_files = []
    for input_file in csv_files:
        try:
            # 转换文件名
            new_filename = convert_filename(input_file)
            output_file = os.path.join(output_dir, new_filename)
            cleaned_file = clean_csv_file(input_file, output_file)
            cleaned_files.append(os.path.basename(cleaned_file))
        except Exception as e:
            print(f"处理文件 {input_file} 时出错: {e}")
    
    # 从5分钟数据生成更高周期数据（15m, 30m, 60m）
    print(f"\n开始从5分钟数据生成更高周期数据...")
    all_generated_files = []
    for cleaned_file in cleaned_files:
        file_path = os.path.join(output_dir, cleaned_file)
        generated = generate_higher_timeframes(file_path, output_dir)
        all_generated_files.extend(generated)
    
    print(f"\n数据清洗完成！")
    print(f"清洗后的文件保存在: {output_dir} 目录")
    print(f"原始数据文件: {len(cleaned_files)} 个 (1m, 5m, 日线)")
    print(f"合成的高周期文件: {len(all_generated_files)} 个 (15m, 30m, 60m)")
    print(f"总计: {len(cleaned_files + all_generated_files)} 个文件")

def is_session_start_time(time_str):
    """判断是否是交易时段开始时间（09:05或21:05）"""
    if isinstance(time_str, str) and ':' in time_str:
        return time_str in ['09:05', '21:05']
    else:
        time_num = int(time_str)
        return time_num in [905, 2105]

def detect_trading_sessions(df):
    """检测交易时段边界"""
    sessions = []
    current_session = []
    
    # 检查是否有datetime列
    has_datetime_column = 'datetime' in df.columns
    
    for i, row in df.iterrows():
        if has_datetime_column:
            # 从datetime列提取时间
            datetime_obj = pd.to_datetime(row['datetime'])
            time_str = datetime_obj.strftime('%H:%M')
        else:
            # 使用原有的time列
            time_str = row['time']
        
        if i == 0:
            current_session = [i]
        else:
            # 如果当前时间是新时段开始时间，结束上一个时段
            if is_session_start_time(time_str) and len(current_session) > 0:
                sessions.append(current_session)
                current_session = [i]
            else:
                current_session.append(i)
    
    if current_session:
        sessions.append(current_session)
    
    return sessions

def resample_session_data(df_session, period_minutes, base_minutes=5, has_datetime_column=False):
    """按你的逻辑进行重采样"""
    if len(df_session) == 0:
        return []
    
    # 根据周期确定每组K线个数
    klines_per_group = period_minutes // base_minutes  # 如：15分钟÷5分钟=3根，15分钟÷1分钟=15根
    
    resampled_data = []
    i = 0
    
    while i < len(df_session):
        # 收集当前组的K线
        group_indices = []
        
        # 正常情况：收集指定数量的K线
        for j in range(klines_per_group):
            if i + j < len(df_session):
                group_indices.append(i + j)
            else:
                break
        
        # 检查是否需要提前结束当前组
        # 如果下一根K线是新时段开始（09:05或21:05），当前组就结束
        next_index = i + len(group_indices)
        if next_index < len(df_session):
            if has_datetime_column:
                # 如果有datetime列，从中提取时间
                next_datetime = pd.to_datetime(df_session.iloc[next_index]['datetime'])
                next_time = next_datetime.strftime('%H:%M')
            else:
                next_time = df_session.iloc[next_index]['time']
            
            if is_session_start_time(next_time):
                # 下一根是新时段开始，当前组结束（不管够不够数）
                pass  # 使用当前已收集的K线
        
        # 如果已经到了时段末尾，也要结束当前组
        if next_index >= len(df_session):
            # 已经是最后几根K线了
            pass
        
        # 合并当前组的K线
        if group_indices:
            group = df_session.iloc[group_indices]
            
            # 根据是否有datetime列来处理时间
            if has_datetime_column:
                # 使用最后一根K线的datetime
                last_datetime = pd.to_datetime(group.iloc[-1]['datetime'])
                new_kline = {
                    'datetime': last_datetime,
                    'instrument': group.iloc[0]['instrument'],
                    'open': group.iloc[0]['open'],
                    'high': group['high'].max(),
                    'low': group['low'].min(),
                    'close': group.iloc[-1]['close'],
                    'volume': group['volume'].sum()
                }
                
                # 添加其他列
                for col in group.columns:
                    if col not in ['datetime', 'instrument', 'open', 'high', 'low', 'close', 'volume']:
                        if col == 'open_interest':
                            new_kline[col] = group.iloc[-1][col]
                        elif col == 'settlement':
                            new_kline[col] = group.iloc[-1][col]
                
                print(f"合并K线组: {pd.to_datetime(group.iloc[0]['datetime']).strftime('%H:%M')}-{last_datetime.strftime('%H:%M')} ({len(group)}根) -> {last_datetime.strftime('%H:%M')}")
            else:
                # 原有逻辑，使用date和time列
                new_kline = {
                    'date': group.iloc[0]['date'],
                    'time': group.iloc[-1]['time'],  # 关键：使用最后一根K线的时间
                    'open': group.iloc[0]['open'],
                    'high': group['high'].max(),
                    'low': group['low'].min(),
                    'close': group.iloc[-1]['close'],
                    'volume': group['volume'].sum(),
                    'open_interest': group.iloc[-1]['open_interest']
                }
                
                if 'settlement' in group.columns:
                    new_kline['settlement'] = group.iloc[-1]['settlement']
                
                print(f"合并K线组: {group.iloc[0]['time']}-{group.iloc[-1]['time']} ({len(group)}根) -> {new_kline['time']}")
            
            resampled_data.append(new_kline)
        
        # 移动到下一组
        i += len(group_indices)
        
        # 如果下一根是新时段开始，跳出当前时段的处理
        if next_index < len(df_session):
            if has_datetime_column:
                next_datetime = pd.to_datetime(df_session.iloc[next_index]['datetime'])
                next_time = next_datetime.strftime('%H:%M')
            else:
                next_time = df_session.iloc[next_index]['time']
            
            if is_session_start_time(next_time):
                break
    
    return resampled_data

def resample_kline_data(df, period_minutes, base_minutes=5):
    """智能重采样K线数据，考虑交易时段边界"""
    print(f"智能重采样：从{base_minutes}分钟生成{period_minutes}分钟K线数据")
    
    # 检查是否有datetime列
    has_datetime_column = 'datetime' in df.columns
    
    # 检测交易时段
    sessions = detect_trading_sessions(df)
    print(f"检测到 {len(sessions)} 个交易时段")
    
    all_resampled_data = []
    
    # 对每个交易时段分别处理
    for session_idx, session_indices in enumerate(sessions):
        df_session = df.iloc[session_indices]
        session_data = resample_session_data(df_session, period_minutes, base_minutes, has_datetime_column)
        all_resampled_data.extend(session_data)
        
        if len(session_data) > 0:
            if has_datetime_column:
                start_datetime = pd.to_datetime(df_session.iloc[0]['datetime'])
                end_datetime = pd.to_datetime(df_session.iloc[-1]['datetime'])
                start_time = start_datetime.strftime('%H:%M')
                end_time = end_datetime.strftime('%H:%M')
            else:
                start_time = df_session.iloc[0]['time']
                end_time = df_session.iloc[-1]['time']
            
            print(f"时段{session_idx+1}: {start_time}-{end_time}, 原始{len(df_session)}根 -> 重采样{len(session_data)}根")
    
    return pd.DataFrame(all_resampled_data)

def generate_higher_timeframes(file_path, output_dir):
    """只从5分钟数据生成15分钟、30分钟、60分钟数据"""
    base_filename = os.path.basename(file_path)
    
    # 只处理5分钟数据
    if '5m' not in base_filename:
        return []
    
    print(f"正在从5分钟数据生成更高周期: {base_filename}")
    
    # 读取5分钟数据
    df = pd.read_csv(file_path, encoding='utf-8')
    
    # 检查是否有datetime列
    if 'datetime' not in df.columns:
        print(f"跳过 {base_filename}：没有datetime列")
        return []
    
    generated_files = []
    
    # 从5分钟数据生成：15m, 30m, 60m
    periods = {
        '15m': 15,
        '30m': 30,
        '60m': 60
    }
    
    for period_name, period_minutes in periods.items():
        try:
            # 重采样数据（基础周期是5分钟）
            resampled_df = resample_kline_data(df.copy(), period_minutes, base_minutes=5)
            
            if len(resampled_df) == 0:
                print(f"警告：{period_name} 重采样结果为空")
                continue
            
            # 生成新文件名
            new_filename = base_filename.replace('5m', period_name)
            new_file_path = os.path.join(output_dir, new_filename)
            
            # 保存文件
            resampled_df.to_csv(new_file_path, index=False, encoding='utf-8')
            generated_files.append(new_filename)
            
            print(f"已生成 {period_name} 数据: {new_filename}")
            print(f"  数据形状: {resampled_df.shape}")
            print(f"  时间范围: {resampled_df['datetime'].min()} 到 {resampled_df['datetime'].max()}")
            
        except Exception as e:
            print(f"生成 {period_name} 数据时出错: {e}")
            import traceback
            traceback.print_exc()
    
    return generated_files
if __name__ == "__main__":
    main()