data.py

"""
data.py

此模块定义了用于处理序列数据的Dataset类和相关函数。它包括了从文件中读取数据、构建特征序列以及将数据转换为适合模型训练的格式的功能。根据模型模式的不同，数据可以以不同的方式进行处理和组织，以满足不同类型模型的需求。
"""

import json
import math
import traceback

import torch
from torch.utils.data.dataloader import DataLoader
from torch.utils.data.dataset import Dataset

import constants


def get_data_meta(meta_path):
    """
    获取元数据
    :param meta_path: str - 元数据路径

    :return: dict - 包含元数据的字典对象
    """
    with open(meta_path, 'r', encoding=constants.DEFAULT_CHARSET) as file:
        return json.load(file)


class FMDataset(Dataset):
    """
    FM 数据集。

    构造函数参数：
        problem_type: bool - 为分类任务时为 True，为回归任务时为 False
        device: torch.device - 数据加载到的设备
        feat_index: list - 特征索引列表
        feat_value: list - 特征值列表
        label: list (可选) - 标签列表，默认为 None
    """

    def __init__(self, problem_type, device, feat_index, feat_value, label=None):
        super(FMDataset, self).__init__()
        self.is_train = False
        self.problem_type = problem_type
        self.device = device
        self.feat_index = feat_index
        self.feat_value = feat_value
        if label is not None:
            self.is_train = True
            self.label = label

    def __len__(self):
        """
        获取数据集大小。

        :return: 数据集大小
        """
        return len(self.feat_index)

    def __getitem__(self, index):
        """
        获取当前索引位置数据。

        :param index: int - 当前索引位置

        :return: tuple - 包含特征索引、特征值和标签（如果是训练模式）的元组
        """
        try:
            if self.is_train:
                return self.feat_index[index], self.feat_value[index], self.label[index]
            else:
                return self.feat_index[index], self.feat_value[index]
        except BaseException as e:
            traceback.print_exc()

    def collate_fn(self, batch_data):
        """
        数据批处理函数，将一个批次的数据转换为模型输入格式。

        :param batch_data: 一个批次的数据

        :return: tuple - 包含特征索引 Ternsor 、特征值 Ternsor 和标签 Ternsor（如果是训练模式）的元组
        """
        labels = [] if self.is_train else None
        feat_indices = []
        feat_values = []
        for instance in batch_data:
            j = 0
            feat_indices.append(instance[j])
            j += 1
            feat_values.append(instance[j])
            if self.is_train:
                j += 1
                labels.append(instance[j])
        if self.is_train:
            return torch.tensor(feat_indices, dtype=torch.int).to(self.device), torch.tensor(feat_values,
                                                                                             dtype=torch.float).to(
                self.device), torch.tensor(labels, dtype=(torch.long if self.problem_type else torch.float)).to(
                self.device)
        else:
            return torch.tensor(feat_indices, dtype=torch.int).to(self.device), torch.tensor(feat_values,
                                                                                             dtype=torch.float).to(
                self.device)


class NonFMDataset(Dataset):
    """
    非 FM 数据集。

    构造函数参数：
        problem_type: bool - 为分类任务时为 True，为回归任务时为 False
        device: torch.device - 数据加载到的设备
        feat_value: list - 特征值列表
        label: list (可选) - 标签列表，默认为 None
    """

    def __init__(self, problem_type, device, feat_value, label=None):
        super(NonFMDataset, self).__init__()
        self.is_train = False
        self.problem_type = problem_type
        self.device = device
        self.feat_value = feat_value
        if label is not None:
            self.is_train = True
            self.label = label

    def __len__(self):
        """
        获取数据集大小。

        :return: 数据集大小
        """
        return len(self.feat_value)

    def __getitem__(self, index):
        """
        获取当前索引位置数据。

        :param index: int - 当前索引位置

        :return: 如果是训练模式是一个包含特征值和标签的元组，否则是特征值
        """
        try:
            if self.is_train:
                return self.feat_value[index], self.label[index]
            else:
                return self.feat_value[index]
        except BaseException as e:
            traceback.print_exc()

    def collate_fn(self, batch_data):
        """
        数据批处理函数，将一个批次的数据转换为模型输入格式。

        :param batch_data: 一个批次的数据

        :return: 如果是训练模式是一个包含特征值 Ternsor 和标签 Ternsor 的元组，否则是特征值 Ternsor
        """
        labels = [] if self.is_train else None
        feat_values = []
        for instance in batch_data:
            j = 0
            feat_values.append(instance[j])
            if self.is_train:
                j += 1
                labels.append(instance[j])
        if self.is_train:
            return torch.tensor(feat_values, dtype=torch.float).to(self.device), torch.tensor(labels, dtype=(
                torch.long if self.problem_type else torch.float)).to(self.device)
        else:
            return torch.tensor(feat_values, dtype=torch.float).to(self.device)


def build_dataset(problem_type, model_mode, code, meta, sensitive_rate, data_path, feature_size, seq_size, buffer_size,
                  batch_size, num_workers, is_multiprocess, device, is_train=True):
    """
    构建用于训练或评测模型的数据集。

    :param problem_type: bool - 为分类任务时为 True，为回归任务时为 False
    :param model_mode: int - 为 FMLSTMAttentionModel 时为 0，为 FMLlamaModel 时为 1
    :param code: str - 标识
    :param meta: dict - 包含数据特征范围的元数据字典
    :param sensitive_rate: float - 敏感率
    :param data_path: str - 数据文件路径
    :param feature_size: int - 输入特征的维度大小
    :param seq_size: int - 序列长度（时间步数）
    :param buffer_size: int - 数据加载时的缓冲区大小，通常用于控制数据预处理和加载的效率
    :param batch_size: int - 每个训练或评测批次的数据量大小
    :param num_workers: int - 工作进程数
    :param is_multiprocess: bool - 如果设置为 True ，则数据加载器在数据集被使用一次后不会关闭工作进程。这使得能够保持工作进程中的数据集实例处于活跃状态。（默认值：False）
    :param device: torch.device - 数据加载到的设备
    :param is_train:bool - 如果设置为 True ，则构建用于训练的数据集；如果设置为 False ，则构建用于评测的数据集（默认值：True）

    :return: tuple - 包含标识和 DataLoader 的元组，DataLoader 用于迭代数据集以进行训练或评测
    """
    label = [] if is_train else None
    feat_index = None if model_mode == 2 else []
    feat_value = []
    price_meta = meta['price']
    price_min = price_meta['min']
    price_max = price_meta['max']
    turnover_meta = meta['turnover']
    turnover_min = turnover_meta['min']
    turnover_max = turnover_meta['max']
    macd_meta = meta['macd']
    macd_min = macd_meta['min']
    macd_max = macd_meta['max']
    expect_change_meta = meta['expect_change']
    expect_change_min = expect_change_meta['min']
    expect_change_max = expect_change_meta['max']
    amount_meta = meta['amount']
    amount_min = amount_meta['min']
    amount_max = amount_meta['max']
    amplitude_meta = meta['amplitude']
    amplitude_min = amplitude_meta['min']
    amplitude_max = amplitude_meta['max']
    change_meta = meta['change']
    change_min = change_meta['min']
    change_max = change_meta['max']
    price_change_meta = meta['price_change']
    price_change_min = price_change_meta['min']
    price_change_max = price_change_meta['max']
    turnover_rate_meta = meta['turnover_rate']
    turnover_rate_min = turnover_rate_meta['min']
    turnover_rate_max = turnover_rate_meta['max']
    with open(data_path, 'r', encoding=constants.DEFAULT_CHARSET) as file:
        line = file.readline().strip()
        while line:
            instance = json.loads(line)
            if is_train:
                expect_change = instance['expect_change']
                label.append(
                    expect_change if problem_type else sigmoid(sensitive_rate, expect_change_min, expect_change_max,
                                                               expect_change))
            if model_mode == 2:
                feat_value.append(
                    build_seq(model_mode, sensitive_rate, price_min, price_max, turnover_min, turnover_max, macd_min,
                              macd_max, amount_min, amount_max, amplitude_min, amplitude_max, change_min, change_max,
                              price_change_min, price_change_max, turnover_rate_min, turnover_rate_max, feature_size,
                              seq_size, instance[
                                  'seq']))  # , turnover_rsi_min, turnover_rsi_max, turnover_kdj_min, turnover_kdj_max, turnover_macd_min, turnover_macd_max, turnover_wr_min, turnover_wr_max, turnover_dmi_min, turnover_dmi_max, turnover_bias_min, turnover_bias_max, turnover_cci_min, turnover_cci_max, turnover_roc_min, turnover_roc_max, turnover_cr_min, turnover_cr_max, feature_size, seq_size, instance['seq']))
            else:
                final_seq_indices, final_seq_values = build_seq(model_mode, sensitive_rate, price_min, price_max,
                                                                turnover_min, turnover_max, macd_min, macd_max,
                                                                amount_min, amount_max, amplitude_min, amplitude_max,
                                                                change_min, change_max, price_change_min,
                                                                price_change_max, turnover_rate_min, turnover_rate_max,
                                                                feature_size, seq_size, instance[
                                                                    'seq'])  # , turnover_rsi_min, turnover_rsi_max, turnover_kdj_min, turnover_kdj_max, turnover_macd_min, turnover_macd_max, turnover_wr_min, turnover_wr_max, turnover_dmi_min, turnover_dmi_max, turnover_bias_min, turnover_bias_max, turnover_cci_min, turnover_cci_max, turnover_roc_min, turnover_roc_max, turnover_cr_min, turnover_cr_max, feature_size, seq_size, instance['seq'])
                feat_index.append(final_seq_indices)
                feat_value.append(final_seq_values)
            line = file.readline().strip()
    if model_mode == 2:
        non_fm_dataset = NonFMDataset(problem_type, device, feat_value, label) if is_train else NonFMDataset(
            problem_type, device, feat_value)
        return code, DataLoader(non_fm_dataset, batch_size, shuffle=buffer_size > batch_size,
                                collate_fn=non_fm_dataset.collate_fn, num_workers=num_workers,
                                persistent_workers=is_multiprocess, pin_memory=False, prefetch_factor=2)
    else:
        fm_dataset = FMDataset(problem_type, device, feat_index, feat_value, label) if is_train else FMDataset(
            problem_type, device, feat_index, feat_value)
        return code, DataLoader(fm_dataset, batch_size, shuffle=buffer_size > batch_size,
                                collate_fn=fm_dataset.collate_fn, num_workers=num_workers,
                                persistent_workers=is_multiprocess, pin_memory=False, prefetch_factor=2)


def build_seq(model_mode, sensitive_rate, price_min, price_max, turnover_min, turnover_max, macd_min, macd_max,
              amount_min, amount_max, amplitude_min, amplitude_max, change_min, change_max, price_change_min,
              price_change_max, turnover_rate_min, turnover_rate_max, feature_size, seq_size, seq):
    """
    构建用于训练或评测模型的序列数据。

    :param model_mode: int - 为 FMLSTMAttentionModel 时为 0，为 FMLlamaModel 时为 1
    :param sensitive_rate: float - 敏感率
    :param price_min: float - 价格最小值
    :param price_max: float - 价格最大值
    :param turnover_min: float - 成交量最小值
    :param turnover_max: float - 成交量最大值
    :param macd_min: float - MACD最小值
    :param macd_max: float - MACD最大值
    :param amount_min: float - 成交金额最小值
    :param amount_max: float - 成交金额最大值
    :param amplitude_min: float - 振幅最小值
    :param amplitude_max: float - 振幅最大值
    :param change_min: float - 涨跌幅最小值
    :param change_max: float - 涨跌幅最大值
    :param price_change_min: float - 涨跌额最小值
    :param price_change_max: float - 涨跌额最大值
    :param turnover_rate_min: float - 成交量比率最小值
    :param turnover_rate_max: float - 成交量比率最大值
    :param feature_size: int - 特征数量
    :param seq_size: int - 序列长度（时间步数）
    :param seq: list - 输入序列

    :return: tuple - 包含最终序列索引和值的元组
    """
    i = 0
    final_seq_indices = None if model_mode == 2 else []
    final_seq_values = []
    while i < seq_size:
        j = 0
        final_feature_indices = None if model_mode == 2 else []
        final_feature_values = []
        features = seq[i]
        i += 1
        while j < feature_size:
            if j <= 27:
                if model_mode != 2:
                    final_feature_indices.append(j)
                if j == 4 or j == 15 or j == 16 or j == 26 or j == 27:
                    final_feature_values.append(sigmoid(sensitive_rate, turnover_min, turnover_max, features[j]))
                elif j == 5:
                    final_feature_values.append(sigmoid(sensitive_rate, amount_min, amount_max, features[j]))
                elif j == 6:
                    final_feature_values.append(sigmoid(sensitive_rate, amplitude_min, amplitude_max, features[j]))
                elif j == 7:
                    final_feature_values.append(sigmoid(sensitive_rate, change_min, change_max, features[j]))
                elif j == 8:
                    final_feature_values.append(
                        sigmoid(sensitive_rate, price_change_min, price_change_max, features[j]))
                elif j == 9:
                    final_feature_values.append(
                        sigmoid(sensitive_rate, turnover_rate_min, turnover_rate_max, features[j]))
                elif j >= 23 and j <= 25:
                    final_feature_values.append(sigmoid(sensitive_rate, macd_min, macd_max, features[j]))
                else:
                    final_feature_values.append(sigmoid(sensitive_rate, price_min, price_max, features[j]))
            else:
                if model_mode == 2:
                    final_feature_values.append(features[j])
                elif features[j] == 1:
                    final_feature_indices.append(j)
                    final_feature_values.append(1)
            j += 1
        if model_mode != 2:
            final_seq_indices.append(final_feature_indices)
        final_seq_values.append(final_feature_values)
    if model_mode == 2:
        return final_seq_values
    else:
        return final_seq_indices, final_seq_values


def sigmoid(sensitive_rate, value_min, value_max, value):
    """
    将输入值按照敏感率和范围使用 Sigmod 进行缩放。

    :param sensitive_rate: float - 敏感率
    :param value_min: float - 值的最小值
    :param value_max: float - 值的最大值
    :param value: float - 输入值

    :return: float - 缩放后的值
    """
    try:
        return 1.0 / (1.0 + math.exp(
            (((value_min + value_max) * 3.0) - (value * 6.0)) / ((value_max - value_min) * sensitive_rate)))
    except OverflowError:
        if value > value_max:
            return 1.0
        elif value < value_min:
            return 0.0
        else:
            return 0.5