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:
@@ -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__":
|
||||
|
||||
Reference in New Issue
Block a user