liuyu/qlib
1
0
Fork 0
mirror of https://github.com/microsoft/qlib.git synced 2026-07-15 00:36:55 +08:00
Code Activity

Add files via upload

Browse Source
This commit is contained in:
Wendi Li
2020-12-09 12:07:01 +08:00
committed by you-n-g
parent a5c098de92
commit 1bbd026195
1 changed files with 291 additions and 249 deletions
Download Patch File Download Diff File Expand all files Collapse all files

72
examples/benchmarks/TFT/tft.py
View File

@@ -19,14 +19,60 @@ from qlib.data.dataset.handler import DataHandlerLP
# To register new datasets, please add them here. # To register new datasets, please add them here.
ALLOW_DATASET = ["Alpha158"] ALLOW_DATASET = ["Alpha158", "Alpha360"]
# To register new datasets, please add their configurations here.
DATASET_SETTING = { DATASET_SETTING = {
"Alpha158": { "Alpha158": {
"feature_col": ["RESI5", "WVMA5", "RSQR5", "KLEN", "RSQR10", "CORR5", "CORD5", "CORR10", "ROC60", "RESI10"], "feature_col": [
"label_col": ["LABEL0"], "RESI5",
"WVMA5",
"RSQR5",
"KLEN",
"RSQR10",
"CORR5",
"CORD5",
"CORR10",
"ROC60",
"RESI10",
"VSTD5",
"RSQR60",
"CORR60",
"WVMA60",
"STD5",
"RSQR20",
"CORD60",
"CORD10",
"CORR20",
"KLOW",
],
"label_col": "LABEL0",
},
"Alpha360": {
"feature_col": [
"HIGH0",
"LOW0",
"OPEN0",
"CLOSE1",
"HIGH1",
"VOLUME1",
"LOW1",
"VOLUME3",
"OPEN1",
"VOLUME4",
"CLOSE2",
"CLOSE4",
"VOLUME5",
"LOW2",
"CLOSE3",
"VOLUME2",
"HIGH2",
"LOW4",
"VOLUME8",
"VOLUME11",
],
"label_col": "LABEL0",
}, },
} }
# To register new datasets, please add their configurations here.
def get_shifted_label(data_df, shifts=5, col_shift="LABEL0"): def get_shifted_label(data_df, shifts=5, col_shift="LABEL0"):
@@ -54,7 +100,7 @@ def process_qlib_data(df, dataset, fillna=False):
""" """
# Several features selected manually # Several features selected manually
feature_col = DATASET_SETTING[dataset]["feature_col"] feature_col = DATASET_SETTING[dataset]["feature_col"]
label_col = DATASET_SETTING[dataset]["label_col"] label_col = [DATASET_SETTING[dataset]["label_col"]]
temp_df = df.loc[:, feature_col + label_col] temp_df = df.loc[:, feature_col + label_col]
if fillna: if fillna:
temp_df = fill_test_na(temp_df) temp_df = fill_test_na(temp_df)
@@ -106,6 +152,8 @@ class TFTModel(ModelFT):
def __init__(self, **kwargs): def __init__(self, **kwargs):
self.model = None self.model = None
self.params = {"DATASET": "Alpha158", "label_shift": 5}
self.params.update(kwargs)
def _prepare_data(self, dataset: DatasetH): def _prepare_data(self, dataset: DatasetH):
df_train, df_valid = dataset.prepare( df_train, df_valid = dataset.prepare(
@@ -113,16 +161,10 @@ class TFTModel(ModelFT):
) )
return transform_df(df_train), transform_df(df_valid) return transform_df(df_train), transform_df(df_valid)
def fit( def fit(self, dataset: DatasetH, MODEL_FOLDER="qlib_tft_model", USE_GPU_ID=0, **kwargs):
self, DATASET = self.params["DATASET"]
dataset: DatasetH, LABEL_SHIFT = self.params["label_shift"]
DATASET="Alpha158", LABEL_COL = DATASET_SETTING[DATASET]["label_col"]
MODEL_FOLDER="qlib_alpha158_model",
LABEL_COL="LABEL0",
LABEL_SHIFT=5,
USE_GPU_ID=0,
**kwargs
):
if DATASET not in ALLOW_DATASET: if DATASET not in ALLOW_DATASET:
raise AssertionError("The dataset is not supported, please make a new formatter to fit this dataset") raise AssertionError("The dataset is not supported, please make a new formatter to fit this dataset")