58 lines
1.5 KiB
Python
58 lines
1.5 KiB
Python
|
"""函数装饰器 - 用于标记因子函数类型并自动注入 over 处理。
|
|||
|
|
|
|||
|
|
提供三种装饰器：
|
|||
|
|
- @time_series: 时序因子，自动添加 .over("ts_code")
|
|||
|
|
- @cross_section: 截面因子，自动添加 .over("trade_date")
|
|||
|
|
- @element_wise: 元素级运算，不添加 over
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
from functools import wraps
|
|||
|
|
from typing import Callable
|
|||
|
|
|
|||
|
|
|
|||
|
|
def time_series(func: Callable) -> Callable:
|
|||
|
|
"""标记为时序因子，自动添加 .over('ts_code')。
|
|||
|
|
|
|||
|
|
用于 ts_* 函数，如 ts_mean, ts_std, ts_corr 等。
|
|||
|
|
每个时序计算都按股票代码分组，防止跨股票串数据。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
@wraps(func)
|
|||
|
|
def wrapper(self, node):
|
|||
|
|
expr = func(self, node)
|
|||
|
|
return expr.over("ts_code")
|
|||
|
|
|
|||
|
|
wrapper._factor_type = "ts"
|
|||
|
|
return wrapper
|
|||
|
|
|
|||
|
|
|
|||
|
|
def cross_section(func: Callable) -> Callable:
|
|||
|
|
"""标记为截面因子，自动添加 .over('trade_date')。
|
|||
|
|
|
|||
|
|
用于 cs_* 函数，如 cs_rank, cs_zscore 等。
|
|||
|
|
每个截面计算都按交易日分组，在同一天的所有股票间计算。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
@wraps(func)
|
|||
|
|
def wrapper(self, node):
|
|||
|
|
expr = func(self, node)
|
|||
|
|
return expr.over("trade_date")
|
|||
|
|
|
|||
|
|
wrapper._factor_type = "cs"
|
|||
|
|
return wrapper
|
|||
|
|
|
|||
|
|
|
|||
|
|
def element_wise(func: Callable) -> Callable:
|
|||
|
|
"""标记为元素级运算，不添加 over。
|
|||
|
|
|
|||
|
|
用于数学函数，如 log, exp, sqrt, cos, sin 等。
|
|||
|
|
这些函数对每个元素独立计算，不需要分组。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
@wraps(func)
|
|||
|
|
def wrapper(self, node):
|
|||
|
|
return func(self, node)
|
|||
|
|
|
|||
|
|
wrapper._factor_type = "element"
|
|||
|
|
return wrapper
|