update TimeSeriesDataset

tests/test_all_pipeline.py

@@ -22,6 +22,8 @@ from qlib.contrib.evaluate import (
 from qlib.utils import exists_qlib_data, init_instance_by_config, flatten_dict
 from qlib.workflow import R
 from qlib.workflow.record_temp import SignalRecord, SigAnaRecord, PortAnaRecord
+from qlib.tests.get_data import GetData
+from qlib.tests import TestAutoData
 
 
 market = "csi300"
@@ -156,26 +158,12 @@ def backtest_analysis(pred, rid):
     return analysis_df
 
 
-class TestAllFlow(unittest.TestCase):
+class TestAllFlow(TestAutoData):
     PRED_SCORE = None
     REPORT_NORMAL = None
     POSITIONS = None
     RID = None
 
-    @classmethod
-    def setUpClass(cls) -> None:
-        # use default data
-        provider_uri = "~/.qlib/qlib_data/cn_data_simple"  # target_dir
-        if not exists_qlib_data(provider_uri):
-            print(f"Qlib data is not found in {provider_uri}")
-            sys.path.append(str(Path(__file__).resolve().parent.parent.joinpath("scripts")))
-            from get_data import GetData
-
-            GetData().qlib_data(
-                name="qlib_data_simple", region="cn", version="latest", interval="1d", target_dir=provider_uri
-            )
-        qlib.init(provider_uri=provider_uri, region=REG_CN)
-
     @classmethod
     def tearDownClass(cls) -> None:
         shutil.rmtree(str(Path(C["exp_manager"]["kwargs"]["uri"].strip("file:")).resolve()))

tests/test_dataset.py

@@ -0,0 +1,53 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import unittest
+import sys
+from qlib.tests import TestAutoData
+from qlib.data.dataset import TSDatasetH
+import numpy as np
+
+
+class TestDataset(TestAutoData):
+    def testTSDataset(self):
+        tsdh = TSDatasetH(
+            handler={
+                "class": "Alpha158",
+                "module_path": "qlib.contrib.data.handler",
+                "kwargs": {
+                    "start_time": "2008-01-01",
+                    "end_time": "2020-08-01",
+                    "fit_start_time": "2008-01-01",
+                    "fit_end_time": "2014-12-31",
+                    "instruments": "csi300",
+                },
+            },
+            segments={
+                "train": ("2008-01-01", "2014-12-31"),
+                "valid": ("2015-01-01", "2016-12-31"),
+                "test": ("2017-01-01", "2020-08-01"),
+            },
+        )
+        _ = tsdh.prepare("train")  # Test the correctness
+        tsds = tsdh.prepare("valid")  # prepare a dataset with is friendly to converting tabular data to time-series
+
+        # The dimension of sample is same as tabular data, but it will return timeseries data of the sample
+
+        # We have two method to get the time-series of a sample
+
+        # 1) sample by int index directly
+        tsds[len(tsds) - 1]
+
+        # 2) sample by <datetime,instrument> index
+        data_from_ds = tsds["2016-12-31", "SZ300315"]
+
+        # Check the data
+        # Get data from DataFrame Directly
+        data_from_df = tsdh._handler.fetch().loc(axis=0)["2015-01-01":"2016-12-31", "SZ300315"].iloc[-30:].values
+
+        equal = np.isclose(data_from_df, data_from_ds)
+        self.assertTrue(equal[~np.isnan(data_from_df)].all())
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=10)