From 73669de3923cc1a43baa6ec2741ed512c107ad39 Mon Sep 17 00:00:00 2001 From: lwwang1995 Date: Wed, 9 Dec 2020 18:11:35 +0800 Subject: [PATCH] Fix code. --- qlib/data/dataset/processor.py | 6 +----- qlib/workflow/record_temp.py | 13 +++++++------ 2 files changed, 8 insertions(+), 11 deletions(-) diff --git a/qlib/data/dataset/processor.py b/qlib/data/dataset/processor.py index 0bdb192a5..5a06f66be 100755 --- a/qlib/data/dataset/processor.py +++ b/qlib/data/dataset/processor.py @@ -113,11 +113,7 @@ class FilterCol(Processor): all_cols = df.columns diff_cols = np.setdiff1d(all_cols.get_level_values(-1), cols.get_level_values(-1)) self.col_list = np.union1d(diff_cols, self.col_list) - - if isinstance(df.columns, pd.MultiIndex): - mask = df.columns.get_level_values(-1).isin(self.col_list) - else: - mask = df.columns.isin(self.col_list) + mask = df.columns.get_level_values(-1).isin(self.col_list) return df.loc[:, mask] diff --git a/qlib/workflow/record_temp.py b/qlib/workflow/record_temp.py index 1cbd4bcab..be4ccdb77 100644 --- a/qlib/workflow/record_temp.py +++ b/qlib/workflow/record_temp.py @@ -131,7 +131,13 @@ class SignalRecord(RecordTemp): pprint(pred.head(5)) # save according label - if isinstance(self.dataset, DatasetH): + if isinstance(self.dataset, TSDatasetH): + index = raw_label.get_index() + raw_label = raw_label.data.loc[index] + raw_label = raw_label[:, -1:] + self.recorder.save_objects(**{"label.pkl": raw_label}) + + elif isinstance(self.dataset, DatasetH): params = dict(self=self.dataset, segments="test", col_set="label", data_key=DataHandlerLP.DK_R) try: # Assume the backend handler is DataHandlerLP @@ -142,11 +148,6 @@ class SignalRecord(RecordTemp): # The backend handler should be DataHandler raw_label = DatasetH.prepare(**params) - if not isinstance(raw_label, pd.DataFrame): - index = raw_label.get_index() - raw_label = raw_label.data.loc[index] - raw_label = raw_label.iloc[:, -1:] - self.recorder.save_objects(**{"label.pkl": raw_label}) def list(self):