__CONFIG_H5_STK_WEIGHT_DIR__ = os.path.join(__CONFIG_H5_STK_DIR__, 'indexconstituent') __CONFIG_H5_STK_DIR__ = r'D:\DATA_STK' __CONFIG_TDX_STK_DIR__ = r'D:\new_hbzq' __CONFIG_H5_STK_DIVIDEND_DIR__ = os.path.join(__CONFIG_H5_STK_DIR__, 'dividend') def merge_weight(rule, wind_code, dataset_name):
    """
    合并一级文件夹
    :param rule:
    :param sector_name:
    :param dataset_name:
    :return:
    """
    path = os.path.join(__CONFIG_H5_STK_DIR__, rule, 'Symbol.csv')
    symbols = all_instruments(path)

    path = os.path.join(__CONFIG_H5_STK_DIR__, rule, 'DateTime.csv')
    DateTime = get_datetime(path)

    df = merge_weight_internal(symbols, DateTime, wind_code)

    path = os.path.join(__CONFIG_H5_STK_DIR__, rule, "%s.h5" % dataset_name)
    write_dataframe_set_dtype_remove_head(path, df, None, dataset_name)

    toc() def load_index_weight(path):
    """
    加载指定目录的数据
    通过unstack将数字设置成值
    :param path:
    :param value:
    :return:
    """
    df = read_indexconstituent_from_dir(path)
    df = df.set_index(df['_datetime_'])
    df = df.set_index([df.index, df['wind_code']], drop=False)
    df = df['i_weight'].unstack()

    return df def merge_weight_internal(symbols, DateTime, wind_code):
    """
    合并一级文件夹
    :param rule:
    :param sector_name:
    :param dataset_name:
    :return:
    """
    tic()
    path = os.path.join(__CONFIG_H5_STK_WEIGHT_DIR__, wind_code)
    df = load_index_weight(path)
    print("数据加载完成")
    # 与行业不同，行业是全部有数据，它是有一部分有数据，所以直接用fillna会出错，需要先填充
    df.fillna(-1, inplace=True)
    toc()

    # 原始数据比较简单，但与行业板块数据又不一样
    # 1.每年的约定时间会调整成份股
    # 2.每天的值都不一样
    # 约定nan表示不属于成份，0表示属于成份，但权重为0
    df = expand_dataframe_according_to(df, DateTime.index, symbols['wind_code'])
    # -1表示特殊数据，处理下
    df.replace(-1, np.nan, inplace=True)
    print("数据加载完成")
    toc()

    return df def get_symbols_from_wind_code_df(wind_codes):
    """

    :param wind_codes:
    :return:
    """
    symbols = wind_codes['wind_code'].str.split('.', expand=True)
    symbols.columns = ['code', 'market']
    symbols['wind_code'] = wind_codes['wind_code']
    symbols['market'] = symbols['market'].str.lower()
    symbols['local_symbol'] = symbols['market'] + symbols['code']
    symbols['symbol'] = symbols['wind_code'].str.replace('.SH', '.SSE').replace('.SZ', '.SZSE')

    return symbols def get_datetime(path):
    dt = pd.read_csv(path, index_col=0, parse_dates=True)
    dt['date'] = dt.index
    return dt def filter_dataframe(df, index_name=None, start_date=None, end_date=None, fields=None):
    if index_name is not None:
        df['index_datetime'] = df[index_name].apply(yyyyMMddHHmm_2_datetime)
        df = df.set_index('index_datetime')
    # 过滤时间
    if start_date is not None or end_date is not None:
        df = df[start_date:end_date]
    # 过滤字段
    if fields is not None:
        df = df[fields]
    return df def split_into_group(arr, n):
    out = [arr[i:i + n] for i in range(0, len(arr), n)]
    return out class MergeBar(object):
    def __init__(self, folder):
        self.prefix = 'tmp'
        self.folder = folder
        self.datetime = None
        self.instruments = None
        self.instruments_group = None
        self.fields = None
        self.group_len = 300
        # datetime与bar_size是相关联的
        self.bar_size = 86400
        self.init_datetime()
        self.init_symbols()
        self.init_fields()

    def init_datetime(self):
        path = os.path.join(self.folder, 'DateTime.csv')
        self.datetime.to_csv(path)

    def init_symbols(self):
        # 不再从导出列表中取，而是从文件夹中推算
        path = os.path.join(self.folder, 'sh')
        df_sh = get_symbols_from_path(path, "SSE")
        path = os.path.join(self.folder, 'sz')
        df_sz = get_symbols_from_path(path, "SZSE")
        df = pd.concat([df_sh, df_sz])

        self.instruments = df

        path = os.path.join(self.folder, 'Symbol.csv')
        self.instruments.to_csv(path, index=False)
        self.instruments_group = split_into_group(self.instruments, self.group_len)

    def init_fields(self):
        pass

    def read_data(self, market, code, bar_size):
        return None

    def _save_data(self, folder, raw_data, field):
        data = raw_data.astype(np.float64).as_matrix()
        path = os.path.join(folder, field + '.h5')
        file = h5py.File(path, 'w')
        file.create_dataset(field, data=data, compression="gzip", compression_opts=6)
        file.close()
        return None

    def _merge_data(self, datetime, instruments, i):
        t = instruments.iloc[i]
        print("%d %s" % (i, t['local_symbol']))
        df = self.read_data(t['market'], t['code'], self.bar_size)
        if df is None:
            return None
        del df['DateTime']
        df = pd.merge(df, datetime, left_index=True, right_index=True, how='right', copy=False)
        df = filter_dataframe(df, None, None, None, self.fields)

        return tuple(df.T.values)

    def _merge_branch(self, folder, datetime, instruments):
        multi = False
        if multi:
            pool_size = multiprocessing.cpu_count() - 1
            pool = multiprocessing.Pool(processes=pool_size)
            func = partial(self._merge_data, datetime, instruments)
            pool_outputs = pool.map(func, range(len(instruments)))
        else:
            pool_outputs = []
            for i in range(len(instruments)):
                x = self._merge_data(datetime, instruments, i)
                #if x is not None:
                pool_outputs.append(x)

        print("数据已经全部读取完成")
        toc()
        print("回收一下内存:%d" % gc.collect())
        # 其中可能有Nono的，需要处理成nan，不能丢弃，否则可能导致Symbol.csv对不上
        # pool_outputs
        pool_outputs = pd.Panel(pool_outputs)  # 内存不够，可能崩溃
        pool_outputs = pool_outputs.transpose(1, 2, 0)
        print("数据转置完成")
        toc()

        for i in range(len(self.fields)):
            print(self.fields[i])

            self._save_data(folder, pool_outputs.loc[i, :, :], self.fields[i])

        toc()
        print("回收一下内存:%d" % gc.collect())

    def merge(self):
        # 数据处理
        tic()

        # 想法对合约进行分组，分组后在对应目录下创建小文件，最后将小文件合并
        for index, item in enumerate(self.instruments_group):
            # 先创建目录
            sub_folder = os.path.join(self.folder, "%s_%d" % (self.prefix, index))
            os.makedirs(sub_folder, exist_ok=True)

            self._merge_branch(sub_folder, self.datetime, item)

        print("分批生成数据完成")
        toc()

    def hmerge(self):
        # 合并数据
        for i in range(len(self.fields)):
            field = self.fields[i]
            print(self.fields[i])
            data = None
            for index, item in enumerate(self.instruments_group):
                # 先创建目录
                sub_folder = os.path.join(self.folder, "%s_%d" % (self.prefix, index))
                sub_file = os.path.join(sub_folder, "%s.h5" % field)
                df = read_h5(sub_file, field)
                if data is None:
                    data = df
                else:
                    data = np.hstack((data, df))

            path = os.path.join(self.folder, field + '.h5')
            file = h5py.File(path, 'w')
            file.create_dataset(field, data=data, compression="gzip", compression_opts=6)
            file.close()

    def clear(self):
        instruments_group = split_into_group(self.instruments, self.group_len)

        # 想法对合约进行分组，分组后在对应目录下创建小文件，最后将小文件合并
        for index, item in enumerate(instruments_group):
            # 先创建目录
            sub_folder = os.path.join(self.folder, "%s_%d" % (self.prefix, index))
            shutil.rmtree(sub_folder)
            print(sub_folder)
        print("目录已删除") def read_h5_tdx(market, code, bar_size, h5_path, tdx_path, div_path):
    """
    指定时间，市场，代码，返回数据
    :param market:
    :param code:
    :param bar_size:
    :return:
    """
    if bar_size == 86400:
        # 日线优先从h5格式中取
        _h5_path = os.path.join(h5_path, get_h5_86400_path(market, code, bar_size))
        try:
            df = pd.read_hdf(_h5_path)
            # print(df.dtypes)
            df = filter_dataframe(df, 'DateTime', None, None, None)
            return df
        except Exception as e:
            # 数据没有取出来
            print(e)
            pass
    else:
        # 其它数据就直接取原始数据，这样省事
        pass

    _tdx_path = os.path.join(tdx_path, get_tdx_path(market, code, bar_size))
    try:
        df = read_file(_tdx_path)
    except FileNotFoundError:
        # 没有原始的数据文件
        return None

    # 有可能没有除权文件
    _div_path = os.path.join(div_path, "%s%s.h5" % (market, code))
    try:
        div = pd.read_hdf(_div_path)
        df = merge_adjust_factor(df, div)
    except:
        # 这里一般是文件没找到，表示没有除权信息
        df['backward_factor'] = 1
        df['forward_factor'] = 1

    return df
import os
import pandas as pd
from kquant_data.config import __CONFIG_H5_STK_WEIGHT_DIR__, __CONFIG_H5_STK_DIR__, __CONFIG_TDX_STK_DIR__, \
    __CONFIG_H5_STK_DIVIDEND_DIR__
from kquant_data.processing.merge import merge_weight, load_index_weight, merge_weight_internal
from kquant_data.stock.symbol import get_symbols_from_wind_code_df
from kquant_data.api import get_datetime
from kquant_data.processing.utils import filter_dataframe, split_into_group
from kquant_data.processing.MergeBar import MergeBar
from kquant_data.stock.stock import read_h5_tdx
﻿#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
指定数据目录，生成对应的合约行业数据

分为两种
1. 全市场数据，将部分标记上权重
2. 只对历史上成为成份股的，进行处理，由于前面已经转换了数据，这里只要跳选数据并处理即可

"""


class MergeDataStock_5min_000300(MergeBar):
    def __init__(self, folder):
        super(MergeDataStock_5min_000300, self).__init__(folder)
        self.bar_size = 300

    def init_symbols(self):
        # 不再从导出列表中取，而是从文件夹中推算
        wind_code = '000300.SH'
        path = os.path.join(__CONFIG_H5_STK_WEIGHT_DIR__, wind_code)
        df = load_index_weight(path)
        wind_codes = pd.DataFrame(list(df.columns), columns=['wind_code'])
        df = get_symbols_from_wind_code_df(wind_codes)

        self.instruments = df

        path = os.path.join(self.folder, 'Symbol.csv')
        self.instruments.to_csv(path, index=False)
        self.instruments_group = split_into_group(self.instruments, self.group_len)

    def init_datetime(self):

        df = read_h5_tdx("sh", "000300", 300, __CONFIG_H5_STK_DIR__, __CONFIG_TDX_STK_DIR__,