ProStock/src/factors/api.py

"""DSL API 层 - 提供常用的符号和函数。

该模块提供量化因子表达式中常用的符号（如 close, open 等）
和函数（如 ts_mean, cs_rank 等），用户可以直接导入使用。

示例:
    >>> from src.factors.api import close, open, ts_mean, cs_rank
    >>> expr = ts_mean(close - open, 20) / close
    >>> print(expr)
    ts_mean(((close - open), 20)) / close
"""

from src.factors.dsl import Symbol, FunctionNode, Node, _ensure_node
from typing import Union

# ==================== 常用价格符号 ====================

#: 收盘价
close = Symbol("close")

#: 开盘价
open = Symbol("open")

#: 最高价
high = Symbol("high")

#: 最低价
low = Symbol("low")

#: 成交量
volume = Symbol("volume")

#: 成交额
amount = Symbol("amount")

#: 前收盘价
pre_close = Symbol("pre_close")

#: 涨跌额
change = Symbol("change")

#: 涨跌幅
pct_change = Symbol("pct_change")


# ==================== 时间序列函数 (ts_*) ====================


def ts_mean(x: Union[Node, str], window: int) -> FunctionNode:
    """时间序列均值。

    计算给定因子在滚动窗口内的平均值。

    Args:
        x: 输入因子表达式或字段名字符串
        window: 滚动窗口大小

    Returns:
        FunctionNode: 函数调用节点

    Example:
        >>> from src.factors.api import close, ts_mean
        >>> expr = ts_mean(close, 20)  # 20日收盘价均值
        >>> expr = ts_mean("close", 20)  # 使用字符串
        >>> print(expr)
        ts_mean(close, 20)
    """
    return FunctionNode("ts_mean", x, window)


def ts_std(x: Union[Node, str], window: int) -> FunctionNode:
    """时间序列标准差。

    计算给定因子在滚动窗口内的标准差。

    Args:
        x: 输入因子表达式或字段名字符串
        window: 滚动窗口大小

    Returns:
        FunctionNode: 函数调用节点
    """
    return FunctionNode("ts_std", x, window)


def ts_max(x: Union[Node, str], window: int) -> FunctionNode:
    """时间序列最大值。

    计算给定因子在滚动窗口内的最大值。

    Args:
        x: 输入因子表达式或字段名字符串
        window: 滚动窗口大小

    Returns:
        FunctionNode: 函数调用节点
    """
    return FunctionNode("ts_max", x, window)


def ts_min(x: Union[Node, str], window: int) -> FunctionNode:
    """时间序列最小值。

    计算给定因子在滚动窗口内的最小值。

    Args:
        x: 输入因子表达式或字段名字符串
        window: 滚动窗口大小

    Returns:
        FunctionNode: 函数调用节点
    """
    return FunctionNode("ts_min", x, window)


def ts_sum(x: Union[Node, str], window: int) -> FunctionNode:
    """时间序列求和。

    计算给定因子在滚动窗口内的求和。

    Args:
        x: 输入因子表达式或字段名字符串
        window: 滚动窗口大小

    Returns:
        FunctionNode: 函数调用节点
    """
    return FunctionNode("ts_sum", x, window)


def ts_delay(x: Union[Node, str], periods: int) -> FunctionNode:
    """时间序列滞后。

    获取给定因子在 N 个周期前的值。

    Args:
        x: 输入因子表达式或字段名字符串
        periods: 滞后期数

    Returns:
        FunctionNode: 函数调用节点
    """
    return FunctionNode("ts_delay", x, periods)


def ts_delta(x: Union[Node, str], periods: int) -> FunctionNode:
    """时间序列差分。

    计算给定因子与 N 个周期前的差值。

    Args:
        x: 输入因子表达式或字段名字符串
        periods: 差分期数

    Returns:
        FunctionNode: 函数调用节点
    """
    return FunctionNode("ts_delta", x, periods)


def ts_corr(x: Union[Node, str], y: Union[Node, str], window: int) -> FunctionNode:
    """时间序列相关系数。

    计算两个因子在滚动窗口内的相关系数。

    Args:
        x: 第一个因子表达式或字段名字符串
        y: 第二个因子表达式或字段名字符串
        window: 滚动窗口大小

    Returns:
        FunctionNode: 函数调用节点
    """
    return FunctionNode("ts_corr", x, y, window)


def ts_cov(x: Union[Node, str], y: Union[Node, str], window: int) -> FunctionNode:
    """时间序列协方差。

    计算两个因子在滚动窗口内的协方差。

    Args:
        x: 第一个因子表达式或字段名字符串
        y: 第二个因子表达式或字段名字符串
        window: 滚动窗口大小

    Returns:
        FunctionNode: 函数调用节点
    """
    return FunctionNode("ts_cov", x, y, window)


def ts_rank(x: Union[Node, str], window: int) -> FunctionNode:
    """时间序列排名。

    计算当前值在过去窗口内的分位排名。

    Args:
        x: 输入因子表达式或字段名字符串
        window: 滚动窗口大小

    Returns:
        FunctionNode: 函数调用节点
    """
    return FunctionNode("ts_rank", x, window)


# ==================== 截面函数 (cs_*) ====================


def cs_rank(x: Union[Node, str]) -> FunctionNode:
    """截面排名。

    计算因子在横截面上的排名（分位数）。

    Args:
        x: 输入因子表达式或字段名字符串

    Returns:
        FunctionNode: 函数调用节点

    Example:
        >>> from src.factors.api import close, cs_rank
        >>> expr = cs_rank(close)  # 收盘价截面排名
        >>> expr = cs_rank("close")  # 使用字符串
        >>> print(expr)
        cs_rank(close)
    """
    return FunctionNode("cs_rank", x)


def cs_zscore(x: Union[Node, str]) -> FunctionNode:
    """截面标准化 (Z-Score)。

    计算因子在横截面上的 Z-Score 标准化值。

    Args:
        x: 输入因子表达式或字段名字符串

    Returns:
        FunctionNode: 函数调用节点
    """
    return FunctionNode("cs_zscore", x)


def cs_neutralize(
    x: Union[Node, str], group: Union[Symbol, str, None] = None
) -> FunctionNode:
    """截面中性化。

    对因子进行行业/市值中性化处理。

    Args:
        x: 输入因子表达式或字段名字符串
        group: 分组变量（如行业分类），可以为字符串或 Symbol，默认为 None

    Returns:
        FunctionNode: 函数调用节点
    """
    if group is not None:
        return FunctionNode("cs_neutralize", x, group)
    return FunctionNode("cs_neutralize", x)


def cs_winsorize(
    x: Union[Node, str], lower: float = 0.01, upper: float = 0.99
) -> FunctionNode:
    """截面缩尾处理。

    对因子进行截面缩尾处理，去除极端值。

    Args:
        x: 输入因子表达式或字段名字符串
        lower: 下尾分位数，默认 0.01
        upper: 上尾分位数，默认 0.99

    Returns:
        FunctionNode: 函数调用节点
    """
    return FunctionNode("cs_winsorize", x, lower, upper)


def cs_demean(x: Union[Node, str]) -> FunctionNode:
    """截面去均值。

    计算因子在横截面上减去均值。

    Args:
        x: 输入因子表达式或字段名字符串

    Returns:
        FunctionNode: 函数调用节点
    """
    return FunctionNode("cs_demean", x)


# ==================== 数学函数 ====================


def log(x: Union[Node, str]) -> FunctionNode:
    """自然对数。

    Args:
        x: 输入因子表达式或字段名字符串

    Returns:
        FunctionNode: 函数调用节点
    """
    return FunctionNode("log", x)


def exp(x: Union[Node, str]) -> FunctionNode:
    """指数函数。

    Args:
        x: 输入因子表达式或字段名字符串

    Returns:
        FunctionNode: 函数调用节点
    """
    return FunctionNode("exp", x)


def sqrt(x: Union[Node, str]) -> FunctionNode:
    """平方根。

    Args:
        x: 输入因子表达式或字段名字符串

    Returns:
        FunctionNode: 函数调用节点
    """
    return FunctionNode("sqrt", x)


def sign(x: Union[Node, str]) -> FunctionNode:
    """符号函数。

    返回 -1, 0, 1 表示输入值的符号。

    Args:
        x: 输入因子表达式或字段名字符串

    Returns:
        FunctionNode: 函数调用节点
    """
    return FunctionNode("sign", x)


def abs(x: Union[Node, str]) -> FunctionNode:
    """绝对值。

    Args:
        x: 输入因子表达式或字段名字符串

    Returns:
        FunctionNode: 函数调用节点
    """
    return FunctionNode("abs", x)


def max_(x: Union[Node, str], y: Union[Node, str, int, float]) -> FunctionNode:
    """逐元素最大值。

    Args:
        x: 第一个因子表达式或字段名字符串
        y: 第二个因子表达式、字段名字符串或数值

    Returns:
        FunctionNode: 函数调用节点
    """
    return FunctionNode("max", x, _ensure_node(y))


def min_(x: Union[Node, str], y: Union[Node, str, int, float]) -> FunctionNode:
    """逐元素最小值。

    Args:
        x: 第一个因子表达式或字段名字符串
        y: 第二个因子表达式、字段名字符串或数值

    Returns:
        FunctionNode: 函数调用节点
    """
    return FunctionNode("min", x, _ensure_node(y))


def clip(
    x: Union[Node, str],
    lower: Union[Node, str, int, float],
    upper: Union[Node, str, int, float],
) -> FunctionNode:
    """数值裁剪。

    将因子值限制在 [lower, upper] 范围内。

    Args:
        x: 输入因子表达式或字段名字符串
        lower: 下限（因子表达式、字段名字符串或数值）
        upper: 上限（因子表达式、字段名字符串或数值）

    Returns:
        FunctionNode: 函数调用节点
    """
    return FunctionNode("clip", x, _ensure_node(lower), _ensure_node(upper))


# ==================== 条件函数 ====================


def if_(
    condition: Union[Node, str],
    true_val: Union[Node, str, int, float],
    false_val: Union[Node, str, int, float],
) -> FunctionNode:
    """条件选择。

    根据条件选择值。

    Args:
        condition: 条件表达式或字段名字符串
        true_val: 条件为真时的值（因子表达式、字段名字符串或数值）
        false_val: 条件为假时的值（因子表达式、字段名字符串或数值）

    Returns:
        FunctionNode: 函数调用节点
    """
    return FunctionNode(
        "if", condition, _ensure_node(true_val), _ensure_node(false_val)
    )


def where(
    condition: Union[Node, str],
    true_val: Union[Node, str, int, float],
    false_val: Union[Node, str, int, float],
) -> FunctionNode:
    """条件选择（if_ 的别名）。

    Args:
        condition: 条件表达式或字段名字符串
        true_val: 条件为真时的值（因子表达式、字段名字符串或数值）
        false_val: 条件为假时的值（因子表达式、字段名字符串或数值）

    Returns:
        FunctionNode: 函数调用节点
    """
    return if_(condition, true_val, false_val)