69 lines
2.2 KiB
Python
69 lines
2.2 KiB
Python
"""挂载 vnpy 自带 Alpha101/Alpha158 表达式因子到注册表.
|
|
|
|
vnpy 的 AlphaDataset 子类在 __init__ 里只做 add_feature 字符串注册(不触发计算),
|
|
用空 schema df 实例化即可零成本提取全部表达式.
|
|
"""
|
|
import sys
|
|
import os
|
|
import polars as pl
|
|
|
|
_VNPY_SRC = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "vnpy_v4.4.0"))
|
|
if _VNPY_SRC not in sys.path:
|
|
sys.path.insert(0, _VNPY_SRC)
|
|
|
|
from .registry import register_factor, list_factors, _REGISTRY
|
|
|
|
_ALPHA_SCHEMA = {
|
|
"vt_symbol": pl.Utf8,
|
|
"datetime": pl.Datetime,
|
|
"open": pl.Float64,
|
|
"high": pl.Float64,
|
|
"low": pl.Float64,
|
|
"close": pl.Float64,
|
|
"volume": pl.Float64,
|
|
"turnover": pl.Float64,
|
|
"open_interest": pl.Float64,
|
|
"vwap": pl.Float64,
|
|
}
|
|
|
|
_DUMMY_PERIOD = ("2018-01-01", "2018-01-02")
|
|
|
|
|
|
def _extract_expressions(dataset_cls) -> dict[str, str]:
|
|
"""实例化数据集类,提取其注册的全部 {name: expression}."""
|
|
df = pl.DataFrame(schema=_ALPHA_SCHEMA)
|
|
ds = dataset_cls(df, _DUMMY_PERIOD, _DUMMY_PERIOD, _DUMMY_PERIOD)
|
|
return {name: str(expr) for name, expr in ds.feature_expressions.items()}
|
|
|
|
|
|
def _mount(dataset_cls, category: str) -> int:
|
|
"""把一个数据集类的表达式挂进注册表(已存在同名跳过,幂等)."""
|
|
mounted = 0
|
|
for name, expression in _extract_expressions(dataset_cls).items():
|
|
if name in _REGISTRY:
|
|
continue
|
|
register_factor(name, expression, category)
|
|
mounted += 1
|
|
return mounted
|
|
|
|
|
|
def mount_alpha101() -> int:
|
|
"""挂载 WorldQuant Alpha101(82 个),category=alpha101."""
|
|
from vnpy.alpha.dataset.datasets.alpha_101 import Alpha101
|
|
return _mount(Alpha101, "alpha101")
|
|
|
|
|
|
def mount_alpha158() -> int:
|
|
"""挂载 Qlib Alpha158(158 个),category=alpha158."""
|
|
from vnpy.alpha.dataset.datasets.alpha_158 import Alpha158
|
|
return _mount(Alpha158, "alpha158")
|
|
|
|
|
|
def mount_all() -> dict[str, int]:
|
|
"""挂载全部内置数据集,返回各类新挂数量."""
|
|
return {"alpha101": mount_alpha101(), "alpha158": mount_alpha158()}
|
|
|
|
|
|
# 模块导入即挂载(与 library.py 内置因子同一模式,下游 import sanguo_factor 即全量可见)
|
|
mount_all()
|