From 3f8617105116b4ffc33759965dd2727a7231f276 Mon Sep 17 00:00:00 2001 From: Linlang Date: Thu, 4 Jul 2024 20:43:41 +0800 Subject: [PATCH] add data loder test --- qlib/contrib/data/handler.py | 2 - qlib/contrib/data/loader.py | 15 +++++--- qlib/data/dataset/loader.py | 20 ++++++---- tests/data_mid_layer_tests/test_dataloader.py | 37 +++++++++++++------ 4 files changed, 47 insertions(+), 27 deletions(-) diff --git a/qlib/contrib/data/handler.py b/qlib/contrib/data/handler.py index 802cc564d..01d83dd6a 100644 --- a/qlib/contrib/data/handler.py +++ b/qlib/contrib/data/handler.py @@ -90,7 +90,6 @@ class Alpha360(DataHandlerLP): return ["Ref($close, -2)/Ref($close, -1) - 1"], ["LABEL0"] - class Alpha360vwap(Alpha360): def get_label_config(self): return ["Ref($vwap, -2)/Ref($vwap, -1) - 1"], ["LABEL0"] @@ -154,7 +153,6 @@ class Alpha158(DataHandlerLP): return ["Ref($close, -2)/Ref($close, -1) - 1"], ["LABEL0"] - class Alpha158vwap(Alpha158): def get_label_config(self): return ["Ref($vwap, -2)/Ref($vwap, -1) - 1"], ["LABEL0"] diff --git a/qlib/contrib/data/loader.py b/qlib/contrib/data/loader.py index 12f3fb41e..432caedf1 100644 --- a/qlib/contrib/data/loader.py +++ b/qlib/contrib/data/loader.py @@ -1,11 +1,12 @@ - from qlib.data.dataset.loader import DataLoader, QlibDataLoader + class Alpha360DL(QlibDataLoader): """Dataloader to get Alpha360""" + def __init__(self, config=None, **kwargs): _config = { - "feature":self.get_feature_config(), + "feature": self.get_feature_config(), } if config is not None: _config.update(config) @@ -59,23 +60,26 @@ class Alpha360DL(QlibDataLoader): class Alpha158DL(QlibDataLoader): """Dataloader to get Alpha158""" + def __init__(self, config=None, **kwargs): _config = { - "feature":self.get_feature_config(), + "feature": self.get_feature_config(), } if config is not None: _config.update(config) super().__init__(config=_config, **kwargs) @staticmethod - def get_feature_config(config={ + def get_feature_config( + config={ "kbar": {}, "price": { "windows": [0], "feature": ["OPEN", "HIGH", "LOW", "VWAP"], }, "rolling": {}, - }): + } + ): """create factors from config config = { @@ -304,4 +308,3 @@ class Alpha158DL(QlibDataLoader): names += ["VSUMD%d" % d for d in windows] return fields, names - diff --git a/qlib/data/dataset/loader.py b/qlib/data/dataset/loader.py index b64a169c9..f2921124b 100644 --- a/qlib/data/dataset/loader.py +++ b/qlib/data/dataset/loader.py @@ -7,7 +7,7 @@ from pathlib import Path import warnings import pandas as pd -from typing import Tuple, Union, List +from typing import Tuple, Union, List, Dict from qlib.data import D from qlib.utils import load_dataset, init_instance_by_config, time_to_slc_point @@ -283,7 +283,8 @@ class NestedDataLoader(DataLoader): """ We have multiple DataLoader, we can use this class to combine them. """ - def __init__(self, dataloader_l: list[dict], join="left") -> None: + + def __init__(self, dataloader_l: List[Dict], join="left") -> None: """ Parameters @@ -307,18 +308,23 @@ class NestedDataLoader(DataLoader): } ] ) - join : + join : it will pass to pd.concat when merging it. """ super().__init__() - self.data_loader_l = [(dl if isinstance(dl, DataLoader) else init_instance_by_config(dl)) for dl in dataloader_l] + self.data_loader_l = [ + (dl if isinstance(dl, DataLoader) else init_instance_by_config(dl)) for dl in dataloader_l + ] self.join = join def load(self, instruments=None, start_time=None, end_time=None) -> pd.DataFrame: - df_l = [] + df_full = None for dl in self.data_loader_l: - df_l = dl.load(instruments, start_time, end_time) - df_full = pd.concat(df_l, axis=1, join=self.join) + df_current = dl.load(instruments, start_time, end_time) + if df_full is None: + df_full = df_current + else: + df_full = pd.merge(df_full, df_current, left_index=True, right_index=True, how=self.join) return df_full.sort_index(axis=1) diff --git a/tests/data_mid_layer_tests/test_dataloader.py b/tests/data_mid_layer_tests/test_dataloader.py index 9f12180ce..15f4f6cf8 100644 --- a/tests/data_mid_layer_tests/test_dataloader.py +++ b/tests/data_mid_layer_tests/test_dataloader.py @@ -1,35 +1,48 @@ - # TODO: # dump alpha 360 to dataframe and merge it with Alpha158 - +import sys import unittest +import qlib +from pathlib import Path +sys.path.append(str(Path(__file__).resolve().parent)) from qlib.data.dataset.loader import NestedDataLoader +from qlib.contrib.data.loader import Alpha158DL, Alpha360DL class TestDataLoader(unittest.TestCase): - def test_nested_data_loader(self): + qlib.init(provider_uri="C:/Users/v-lvlinlang/.qlib/qlib_data/cn_data_simple") nd = NestedDataLoader( dataloader_l=[ { "class": "qlib.contrib.data.loader.Alpha158DL", - }, { + }, + { "class": "qlib.contrib.data.loader.Alpha360DL", - "kwargs": { - "config": { - "label": ( ["Ref($close, -2)/Ref($close, -1) - 1"], ["LABEL0"]) - } - } - } + "kwargs": {"config": {"label": (["Ref($close, -2)/Ref($close, -1) - 1"], ["LABEL0"])}}, + }, ] ) # Of course you can use StaticDataLoader - nd.load - ... + dataset = nd.load() + + assert dataset is not None + + columns = dataset.columns.tolist() + columns_list = [tup[1] for tup in columns] + + for col in Alpha158DL.get_feature_config()[1]: + assert col in columns_list + + for col in Alpha360DL.get_feature_config()[1]: + assert col in columns_list + + # 断言标签也包含在数据中 + assert "LABEL0" in columns_list # Then you can use it wth DataHandler;