gp_binance_demo/binance_order_flow/data_processing.py

import json
import time
import numpy as np
import pandas as pd
import queue
import threading
from collections import deque
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report
from sklearn.utils import resample
from logger_config import logger

# 假设我们有一个数据流，订单簿和成交数据
order_book_snapshots = deque(maxlen=100)  # 存储过去100ms的订单簿快照
trade_data = deque(maxlen=100)  # 存储过去100ms的成交数据

# 数据积累的阈值
DATA_THRESHOLD = 20
model = RandomForestClassifier(n_estimators=100)  # 使用随机森林模型
model_trained = False  # 标记模型是否已经被训练
messages = queue.Queue()  # 创建一个线程安全队列

stop_event = threading.Event()


def on_message_trade(_ws, message):
    global trade_data
    json_message = json.loads(message)
    trade = {
        'price': float(json_message['data']['p']),
        'qty': float(json_message['data']['q']),
        'timestamp': pd.to_datetime(json_message['data']['T'], unit='ms'),
        'side': 'sell' if json_message['data']['m'] else 'buy'
    }
    trade_data.append(trade)
    check_and_train_model()
    predict_market_direction()


def on_message_depth(_ws, message):
    global order_book_snapshots
    json_message = json.loads(message)
    bids = json_message['data']['b'][:10]  # Top 10 bids
    asks = json_message['data']['a'][:10]  # Top 10 asks
    timestamp = pd.to_datetime(json_message['data']['E'], unit='ms')
    order_book_snapshots.append({
        'bids': bids,
        'asks': asks,
        'timestamp': timestamp
    })
    check_and_train_model()
    predict_market_direction()


def extract_features(order_book, trade):
    # 计算买卖盘差距（spread）
    best_bid = float(order_book['bids'][0][0])
    best_ask = float(order_book['asks'][0][0])
    spread = best_ask - best_bid

    # 计算买卖盘深度
    bid_depth = sum(float(bid[1]) for bid in order_book['bids'])
    ask_depth = sum(float(ask[1]) for ask in order_book['asks'])

    # 计算成交量和方向
    trade_volume = trade['qty']
    trade_side = 1 if trade['side'] == 'buy' else -1

    # 计算买卖盘数量
    bid_count = len(order_book['bids'])
    ask_count = len(order_book['asks'])

    # 计算时间特征
    timestamp = trade['timestamp'].timestamp()

    features = {
        'spread': spread,
        'bid_depth': bid_depth,
        'ask_depth': ask_depth,
        'trade_volume': trade_volume,
        'trade_side': trade_side,
        'bid_count': bid_count,
        'ask_count': ask_count,
        'timestamp': timestamp
    }

    return features


def prepare_training_data():
    # 提取特征和标签
    X = []
    y = []

    for i in range(len(order_book_snapshots) - 1):
        if i + 1 >= len(order_book_snapshots) or i >= len(trade_data):
            break
        current_order_book = order_book_snapshots[i]
        current_trade = trade_data[i]
        future_order_book = order_book_snapshots[i + 1]

        # 提取当前的特征
        features = extract_features(current_order_book, current_trade)
        X.append(list(features.values()))

        # 生成标签
        current_price = float(current_order_book['bids'][0][0])
        future_price = float(future_order_book['bids'][0][0])
        label = generate_label(current_price, future_price)
        y.append(label)

    # 将特征和标签转换为NumPy数组
    X_train = np.array(X)
    y_train = np.array(y)

    return X_train, y_train


def generate_label(current_price, future_price):
    return 1 if future_price > current_price else 0


def balance_data(X, y):
    # 将数据转换为 DataFrame，便于处理
    df = pd.DataFrame(X, columns=['spread', 'bid_depth', 'ask_depth', 'trade_volume', 'trade_side', 'bid_count', 'ask_count', 'timestamp'])
    df['label'] = y

    # 分别获取两类数据
    df_majority = df[df.label == 0]
    df_minority = df[df.label == 1]

    if len(df_minority) == 0 or len(df_majority) == 0:
        return X, y  # 如果某一类数据为空，返回原数据

    # 使用下采样或上采样来平衡数据
    if len(df_majority) > len(df_minority):
        df_majority_downsampled = resample(df_majority, replace=False, n_samples=len(df_minority), random_state=42)
        df_balanced = pd.concat([df_majority_downsampled, df_minority])
    else:
        df_minority_upsampled = resample(df_minority, replace=True, n_samples=len(df_majority), random_state=42)
        df_balanced = pd.concat([df_majority, df_minority_upsampled])

    # 提取平衡后的特征和标签
    X_balanced = df_balanced.drop('label', axis=1).values
    y_balanced = df_balanced['label'].values

    return X_balanced, y_balanced


def check_and_train_model():
    global model_trained
    if len(order_book_snapshots) >= DATA_THRESHOLD and len(trade_data) >= DATA_THRESHOLD:
        X_train, y_train = prepare_training_data()
        if len(X_train) > 0 and len(y_train) > 0:
            X_train, y_train = balance_data(X_train, y_train)  # 平衡数据
            X_train, X_test, y_train, y_test = train_test_split(X_train, y_train, test_size=0.2, random_state=42)
            model.fit(X_train, y_train)
            model_trained = True  # 标记模型已经被训练
            logger.info("Model trained with %d samples", len(X_train))
            y_pred = model.predict(X_test)
            logger.info("\n" + classification_report(y_test, y_pred, zero_division=0))
        else:
            logger.info("Insufficient data to train the model")
    else:
        logger.info("Not enough data to train the model: %d order book snapshots, %d trades", len(order_book_snapshots), len(trade_data))


def predict_market_direction():
    global model_trained
    if len(order_book_snapshots) == 0 or len(trade_data) == 0:
        return

    if not model_trained:
        return

    features = extract_features(order_book_snapshots[-1], trade_data[-1])
    if features is not None:
        feature_vector = np.array([list(features.values())])
        prediction = model.predict(feature_vector)
        logger.info("Predicted Market Direction: %s", "Up" if prediction[0] == 1 else "Down")
        show_message(prediction[0])


def show_message(market_direction):
    global order_book_snapshots, trade_data

    if len(order_book_snapshots) > 0 and len(trade_data) > 0:
        # 获取最新的订单簿数据和成交数据
        latest_order_book = order_book_snapshots[-1]
        latest_trade = trade_data[-1]

        # 提取asks和bids数据
        asks = [[float(price), 1 if market_direction == 1 else 0] for price, qty in latest_order_book['asks']]
        bids = [[float(price), 1 if market_direction == 0 else 0] for price, qty in latest_order_book['bids']]

        # 排序asks和bids数据
        asks_sorted = sorted(asks, key=lambda x: x[0])
        bids_sorted = sorted(bids, key=lambda x: x[0], reverse=True)

        last_price = latest_trade['price']
        # last_qty = latest_trade['qty']
        last_qty = 0
        side = latest_trade['side']

        # 生成数据字典
        data = {
            "asks": asks_sorted,
            "bids": bids_sorted,
            "last_price": last_price,
            "last_qty": last_qty,
            "side": side,
            "time": int(time.time() * 1000)
        }

        # 将数据放入消息队列
        messages.put(data)