ubiquant量化竞赛

import pickle   
def load_pickle(path):
    with open(path, 'rb') as f:
        return pickle.load(f)

path = '/home/mw/input/debug2469/debug'

x_A = load_pickle(f'{path}/A.pickle')
x_H = load_pickle(f'{path}/H.pickle')

all_dict = {}
all_dict['A'] = x_A
all_dict['H'] = x_H
INT_dates_all = list(sorted(set(all_dict['A']['INT_dates']).union(set(all_dict['H']['INT_dates']))))
all_dict['INT_dates_all'] = INT_dates_all

x_A['close'].shape, x_H['close'].shape #

list(x_A.keys()) # 所有字段

 

x_A['INT_tickers'] # 股票id

 

list(x_A['INT_dates']) # 日期

 

x_A['INT_date_2_di'] # 日期对应数据的di, 因为存在A股和H股不同时开市的情况

 

import numpy as np
import pandas as pd

class DepthInfo:
    def __init__(self):
        self.prices = []
        self.volumes = []
        
class BasicStockInfo():
    def __init__(self):
        self.buy_infos = DepthInfo()
        self.sell_infos = DepthInfo()
    def show_all(self):
        print(f'[+] stock_id = {self.stock_id}')
        print(f'    market_id = {self.market_id}')
        print(f'    open = {self.open}')
        print(f'    close = {self.close}')
        print(f'    close_prev = {self.close_prev}')
        print(f'    high = {self.high}')
        print(f'    low = {self.low}')
        print(f'    volume = {self.volume}')
        print(f'    tvr = {self.tvr}')
        print(f'    is_tradable = {self.is_tradable}')
        print(f'    buy_infos.prices = {self.buy_infos.prices}')
        print(f'    buy_infos.volumes = {self.buy_infos.volumes}')
        print(f'    sell_infos.prices = {self.sell_infos.prices}')
        print(f'    sell_infos.volumes = {self.sell_infos.volumes}')

class StockManager:
    def __init__(self, all_dict):
        self.all_dict = all_dict

    def get_stock_info_dict(self, x_dict, di, ti, market_id):
        stock_info_dict = {}
        if di == -1: 
            return stock_info_dict
        if market_id == 0 and ti >= 25 + 3:
            return stock_info_dict
        if market_id == 1 and ti >= 34:
            return stock_info_dict

        # A: 
        if market_id == 0:
            # 11:30-12:00, not open
            if ti >= 13 and ti < 13 + 3: 
                return stock_info_dict
            
            # A: 13:00-15:00: [13+3, 24+3]
            if ti >= 16:
                ti -= 3

        # not valid di
        if di >= x_dict['open'].shape[0]:
            return {}

        # eacho stock
        for si, stock_id in enumerate(x_dict['INT_tickers']):
            stock_info = BasicStockInfo()
            stock_info.stock_id = stock_id
            stock_info.market_id = market_id

            stock_info.open = x_dict['open'][di, ti, si]
            stock_info.close = x_dict['close'][di, ti, si]
            stock_info.high = x_dict['high'][di, ti, si]
            stock_info.low = x_dict['low'][di, ti, si]
            stock_info.volume = x_dict['volume'][di, ti, si]
            stock_info.tvr = x_dict['tvr'][di, ti, si]
            stock_info.is_tradable = x_dict['mask'][di, ti, si]
            stock_info.close_prev = x_dict['close_prev'][di, si] if ti == 0 else np.nan
            for a in list(x_dict['bidprice'][:, di, ti, si]):
                stock_info.buy_infos.prices.append(a)
            for a in list(x_dict['bidvolume'][:, di, ti, si]):
                stock_info.buy_infos.volumes.append(a)
            for a in list(x_dict['askprice'][:, di, ti, si]):
                stock_info.sell_infos.prices.append(a)
            for a in list(x_dict['askvolume'][:, di, ti, si]):
                stock_info.sell_infos.volumes.append(a)
            stock_info_dict[stock_id] = stock_info
        return stock_info_dict

    def get_sequence_stocks_info(self, date_i, ti):
        print(f'date_i = {date_i}, ti = {ti}')
        # A: 09:30-11:30, 13:00-15:00, total = 4   hours, 4 * 6 + 1 = 25
        # H: 09:30-12:00, 13:00-16:00, total = 5.5 hours, 5 * 6 + 3 + 1 = 34

        # A: 
        # 1) 09:30-11:30: [0, 12]
        # 2) 13:00-15:00: [13+3, 24+3]

        # H: 
        # 1) 09:30-12:30: [0, 15]
        # 2) 13:00-16:00: [16, 33]

        all_dict = self.all_dict
        date_i += 20 # 20 days for history data
        date = all_dict['INT_dates_all'][date_i]
        
        res = []
        di = all_dict['A']['INT_date_2_di'].get(date, -1)
        stock_info_dict = self.get_stock_info_dict(all_dict['A'], di, ti, market_id = 0)
        res.append(stock_info_dict)   

        di = all_dict['H']['INT_date_2_di'].get(date, -1)
        stock_info_dict = self.get_stock_info_dict(all_dict['H'], di, ti, market_id = 1)
        res.append(stock_info_dict) 
        return res    

模拟拿到数据, 方便debug. date = , ti =