1
0
mirror of https://github.com/microsoft/qlib.git synced 2026-07-11 14:56:55 +08:00

Optimize the pit collector script (#982)

* Optimize the pit collector script

* Add copyright notice to collector.py

* Remove unnecessary parameters for test_pit.py

* Update test_pit.py

* Update test_pit.py
This commit is contained in:
Chauncey
2022-03-18 21:51:36 +08:00
committed by GitHub
parent f2a5ecd98a
commit 8efc8b92ef
3 changed files with 208 additions and 300 deletions

View File

@@ -1,20 +1,28 @@
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
import pandas as pd
import qlib
from qlib.data import D
import unittest
pd.set_option("display.width", 1000)
pd.set_option("display.max_columns", None)
class TestPIT(unittest.TestCase):
"""
NOTE!!!!!!
The assert of this test assumes that users follows the cmd below and only download 2 stock.
`python collector.py download_data --source_dir ./csv_pit --start 2000-01-01 --end 2020-01-01 --interval quarterly --symbol_flt_regx "^(600519|000725).*"`
1. `python scripts/get_data.py qlib_data --target_dir ~/.qlib/qlib_data/cn_data --region cn`
2. `python scripts/data_collector/pit/collector.py download_data --source_dir ~/.qlib/stock_data/source/pit --start 2000-01-01 --end 2020-01-01 --interval quarterly --symbol_regex "^(600519|000725).*"`
3. `python scripts/data_collector/pit/collector.py normalize_data --interval quarterly --source_dir ~/.qlib/stock_data/source/pit --normalize_dir ~/.qlib/stock_data/source/pit_normalized`
4. `python scripts/dump_pit.py dump --csv_path ~/.qlib/stock_data/source/pit_normalized --qlib_dir ~/.qlib/qlib_data/cn_data --interval quarterly`
"""
def setUp(self):
# qlib.init(kernels=1) # NOTE: set kernel to 1 to make it debug easier
qlib.init() # NOTE: set kernel to 1 to make it debug easier
qlib.init()
def to_str(self, obj):
return "".join(str(obj).split())
@@ -27,10 +35,7 @@ class TestPIT(unittest.TestCase):
fields = ["P($$roewa_q)", "P($$yoyni_q)"]
# Mao Tai published 2019Q2 report at 2019-07-13 & 2019-07-18
# - http://www.cninfo.com.cn/new/commonUrl/pageOfSearch?url=disclosure/list/search&lastPage=index
data = D.features(instruments, fields, start_time="2019-01-01", end_time="20190719", freq="day")
print(data)
data = D.features(instruments, fields, start_time="2019-01-01", end_time="2019-07-19", freq="day")
res = """
P($$roewa_q) P($$yoyni_q)
count 133.000000 133.000000
@@ -57,12 +62,11 @@ class TestPIT(unittest.TestCase):
def test_no_exist_data(self):
fields = ["P($$roewa_q)", "P($$yoyni_q)", "$close"]
data = D.features(["sh600519", "sh601988"], fields, start_time="2019-01-01", end_time="20190719", freq="day")
data = D.features(["sh600519", "sh601988"], fields, start_time="2019-01-01", end_time="2019-07-19", freq="day")
data["$close"] = 1 # in case of different dataset gives different values
print(data)
expect = """
P($$roewa_q) P($$yoyni_q) $close
instrument datetime
instrument datetime
sh600519 2019-01-02 0.25522 0.243892 1
2019-01-03 0.25522 0.243892 1
2019-01-04 0.25522 0.243892 1
@@ -74,7 +78,7 @@ class TestPIT(unittest.TestCase):
2019-07-17 NaN NaN 1
2019-07-18 NaN NaN 1
2019-07-19 NaN NaN 1
[266 rows x 3 columns]
"""
self.check_same(data, expect)
@@ -115,12 +119,12 @@ class TestPIT(unittest.TestCase):
fields = ["P($$roewa_q)"]
instruments = ["sh600519"]
_ = D.features(instruments, fields, freq="day") # this should not raise error
data = D.features(instruments, fields, end_time="20200101", freq="day") # this should not raise error
data = D.features(instruments, fields, end_time="2020-01-01", freq="day") # this should not raise error
s = data.iloc[:, 0]
# You can check the expected value based on the content in `docs/advanced/PIT.rst`
expect = """
instrument datetime
sh600519 1999-11-10 NaN
sh600519 2005-01-04 NaN
2007-04-30 0.090219
2007-08-17 0.139330
2007-10-23 0.245863
@@ -156,7 +160,7 @@ class TestPIT(unittest.TestCase):
2014-10-30 0.234085
2015-04-21 0.078494
2015-08-28 0.137504
2015-10-26 0.201709
2015-10-23 0.201709
2016-03-24 0.264205
2016-04-21 0.073664
2016-08-29 0.136576
@@ -176,7 +180,6 @@ class TestPIT(unittest.TestCase):
2019-10-16 0.255819
Name: P($$roewa_q), dtype: float32
"""
self.check_same(s[~s.duplicated().values], expect)
def test_expr2(self):
@@ -186,8 +189,6 @@ class TestPIT(unittest.TestCase):
fields += ["P(Sum($$yoyni_q, 4))"]
fields += ["$close", "P($$roewa_q) * $close"]
data = D.features(instruments, fields, start_time="2019-01-01", end_time="2020-01-01", freq="day")
print(data)
print(data.describe())
if __name__ == "__main__":