quantnodes 3.0.0__py3-none-any.whl

QuantNodes/research/quant_alpha/mcts/tree.py

@@ -0,0 +1,201 @@
+# coding=utf-8
+"""
+mcts/tree.py - MCTS 搜索树（带谱系追踪）
+
+vs 旧 mcts_search.py:53-62 的 MCTSNode：
+- 旧：仅 parent: MCTSNode 引用，无 entry_id
+- 新：entry_id (UUID) + parent_id (UUID) + 完整谱系
+
+每节点的 entry_id 可映射到 core/trajectory.TrajectoryEntry
+"""
+
+from __future__ import annotations
+
+import math
+import uuid
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Dict, List, Optional
+
+
+class NodeStatus(str, Enum):
+    """MCTS 节点状态"""
+    PENDING = "pending"          # 未评估
+    EVALUATED = "evaluated"      # 已评估（5 通道反馈完成）
+    PRUNED = "pruned"            # 被质量门拦截
+    REJECTED = "rejected"        # 被反馈拒绝
+
+
+@dataclass
+class MCTSNode:
+    """MCTS 搜索树节点（带谱系追踪）
+
+    关键字段：
+    - entry_id: UUID，唯一标识（可映射到 TrajectoryEntry）
+    - parent_id: 父节点 entry_id（None = 根节点）
+    - formula: 当前节点公式
+    - depth: 树深度（根 = 0）
+    - children: 子节点列表
+    - visits: 访问次数
+    - dimension_scores: 5 通道反馈评分
+    - overall_score: 综合评分
+    - status: 节点状态
+    """
+    formula: str
+    parent_id: Optional[str] = None
+    entry_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    depth: int = 0
+    children: List["MCTSNode"] = field(default_factory=list)
+    visits: int = 0
+    dimension_scores: Dict[str, float] = field(default_factory=dict)
+    overall_score: float = 0.0
+    status: NodeStatus = NodeStatus.PENDING
+    is_expanded: bool = False
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+    # 暂存：UCB1 计算时用的父节点引用（避免 O(n) parent_id 查找）
+    _parent_ref: Optional["MCTSNode"] = field(default=None, repr=False, compare=False)
+
+    def add_child(self, child: "MCTSNode") -> None:
+        """添加子节点（同时设置 parent_id 和 _parent_ref）"""
+        child.parent_id = self.entry_id
+        child._parent_ref = self
+        child.depth = self.depth + 1
+        self.children.append(child)
+        self.is_expanded = True
+
+    def is_leaf(self) -> bool:
+        return len(self.children) == 0
+
+    def is_root(self) -> bool:
+        return self.parent_id is None
+
+    def ancestors(self) -> List["MCTSNode"]:
+        """返回从根到当前节点的祖先链（不含当前节点）"""
+        chain = []
+        node = self._parent_ref
+        while node is not None:
+            chain.append(node)
+            node = node._parent_ref
+        return list(reversed(chain))
+
+    def lineage_depth(self) -> int:
+        """返回谱系深度（从根到当前节点的边数）"""
+        return self.depth
+
+    def ucb1(self, exploration_weight: float = 1.414) -> float:
+        """UCB1 评分
+
+        UCB1 = Q(s) + c * sqrt(ln(N_parent) / N_s)
+
+        Returns:
+            float('inf') 当 visits = 0（优先扩展）
+        """
+        if self.visits == 0:
+            return float("inf")
+
+        parent = self._parent_ref
+        if parent is None or parent.visits == 0:
+            return float("inf")
+
+        exploit = self.overall_score
+        explore = exploration_weight * math.sqrt(
+            math.log(parent.visits) / self.visits
+        )
+        return exploit + explore
+
+    def __repr__(self) -> str:
+        preview = self.formula[:50] + "..." if len(self.formula) > 50 else self.formula
+        return (
+            f"MCTSNode(entry={self.entry_id[:8]}, "
+            f"depth={self.depth}, visits={self.visits}, "
+            f"score={self.overall_score:.3f}, "
+            f"status={self.status.value}, "
+            f"formula={preview!r})"
+        )
+
+
+@dataclass
+class MCTSTree:
+    """MCTS 搜索树容器
+
+    vs 旧 mcts_search.py：树结构本身相似，但增加：
+    - 完整 entry_id 谱系
+    - 节点状态追踪
+    - 公式缓存去重
+    - entry_id 索引（O(1) 父节点查找）
+    """
+    root: MCTSNode = field(default_factory=lambda: MCTSNode(
+        formula="__ROOT__", depth=-1,
+    ))
+    formula_cache: Dict[str, MCTSNode] = field(default_factory=dict)
+    _entry_index: Dict[str, MCTSNode] = field(default_factory=dict)
+    total_iterations: int = 0
+
+    def add_node(self, node: MCTSNode, parent: Optional[MCTSNode] = None) -> None:
+        """添加节点到树
+
+        Args:
+            node: 节点
+            parent: 父节点（None = 根节点）
+        """
+        if parent is None:
+            parent = self.root
+        parent.add_child(node)
+        # 加入公式缓存
+        self.formula_cache[node.formula] = node
+        # 加入 entry_id 索引
+        self._entry_index[node.entry_id] = node
+
+    def get_by_entry_id(self, entry_id: str) -> Optional[MCTSNode]:
+        """按 entry_id 查找节点（O(1) via index）"""
+        return self._entry_index.get(entry_id)
+
+    def get_by_formula(self, formula: str) -> Optional[MCTSNode]:
+        """按公式查找节点（O(1) via cache）"""
+        return self.formula_cache.get(formula)
+
+    def all_nodes(self) -> List[MCTSNode]:
+        """返回所有节点（DFS）"""
+        result = []
+        stack = [self.root]
+        while stack:
+            node = stack.pop()
+            if node is not self.root:
+                result.append(node)
+            stack.extend(node.children)
+        return result
+
+    def leaves(self) -> List[MCTSNode]:
+        """返回所有叶子节点"""
+        return [n for n in self.all_nodes() if n.is_leaf()]
+
+    def best_k(self, k: int = 10, metric: str = "overall_score") -> List[MCTSNode]:
+        """返回综合评分最高的 k 个节点（按 entry_id 去重）"""
+        seen = set()
+        unique = []
+        nodes = self.all_nodes()
+        # 按 metric 降序
+        nodes.sort(key=lambda n: getattr(n, metric, 0.0), reverse=True)
+        for n in nodes:
+            if n.entry_id in seen:
+                continue
+            seen.add(n.entry_id)
+            unique.append(n)
+            if len(unique) >= k:
+                break
+        return unique
+
+    def stats(self) -> Dict[str, Any]:
+        """树统计"""
+        nodes = self.all_nodes()
+        by_status: Dict[str, int] = {}
+        for n in nodes:
+            s = n.status.value
+            by_status[s] = by_status.get(s, 0) + 1
+        return {
+            "total_nodes": len(nodes),
+            "by_status": by_status,
+            "max_depth": max((n.depth for n in nodes), default=0),
+            "total_iterations": self.total_iterations,
+        }

QuantNodes/research/quant_alpha/operator_vocab/__init__.py

@@ -0,0 +1,50 @@
+# coding=utf-8
+"""
+operator_vocab - 算子词表层
+
+把分散在 3 处的 285 个算子整合成统一接口：
+- L0 Built-in: factor_node/factor_functions/ (157 算子)
+- L0 TA-Lib: factor_node/factor_functions/talib_ops.py (109 算子，需显式 import)
+- L1 Composite: operators/composite_dag.py (20 算子)
+
+主要功能：
+- 统一查询：list_operators / get_operator / get_metadata
+- 算子元数据：12 字段（含 7 个 LLM 友好）
+- Namespace 构建：build_namespace（per-date over() 修复）
+- 端到端评估：evaluate（formula + data → result）
+
+修复 3 个 latent bug：
+1. ts_corr/ts_cov 的 Series.rolling_corr 不存在
+2. rank/zscore 全局而非 per-date
+3. 异常被静默吞掉
+"""
+
+from __future__ import annotations
+
+from QuantNodes.research.quant_alpha.operator_vocab.metadata import (
+    OperatorMetadata,
+    OperatorCategory,
+)
+from QuantNodes.research.quant_alpha.operator_vocab.config import (
+    OperatorVocabConfig,
+)
+from QuantNodes.research.quant_alpha.operator_vocab.vocabulary import (
+    OperatorVocab,
+    build_namespace,
+    list_vocab_operators,
+    get_vocab_operator,
+    get_vocab_metadata,
+)
+
+__all__ = [
+    # 主类
+    "OperatorVocab",
+    "OperatorVocabConfig",
+    "OperatorMetadata",
+    "OperatorCategory",
+    # 便捷函数
+    "build_namespace",
+    "list_vocab_operators",
+    "get_vocab_operator",
+    "get_vocab_metadata",
+]

QuantNodes/research/quant_alpha/operator_vocab/config.py

@@ -0,0 +1,54 @@
+# coding=utf-8
+"""
+OperatorVocab 配置 - 控制 namespace 构建与算子查询行为
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+
+@dataclass
+class OperatorVocabConfig:
+    """OperatorVocab 配置
+
+    控制：
+    - per-date over() 默认行为（cross_sectional）
+    - 启用的算子类别
+    - TA-Lib 是否自动加载
+    - L0/L1/L2 注册表路径
+    """
+
+    # 截面算子默认行为：True=per-date over(date)，False=全局
+    # 修复旧 12-lambda namespace 的 BUG 2（rank/zscore 全局错误）
+    cross_sectional: bool = True
+
+    # 启用的算子类别（默认全部）
+    enabled_categories: List[str] = field(default_factory=lambda: [
+        "point", "time", "section", "multi_section", "talib"
+    ])
+
+    # 是否自动加载 talib_ops（109 算子）
+    # 默认 True；如 talib 未安装会自动降级为 False
+    talib_enabled: bool = True
+
+    # eval 沙箱的超时（秒），None 表示不限制
+    eval_timeout_seconds: Optional[float] = None
+
+    # 公式长度上限（字符），None 表示不限制
+    # 防止超长公式 DoS
+    max_formula_length: Optional[int] = 5000
+
+    # 公式嵌套深度上限
+    # 防止递归嵌套爆炸
+    max_formula_depth: int = 20
+
+    def __post_init__(self):
+        # 验证 enabled_categories
+        valid = {"point", "time", "section", "multi_section", "talib"}
+        for cat in self.enabled_categories:
+            if cat not in valid:
+                raise ValueError(
+                    f"Invalid category '{cat}'. Must be one of: {valid}"
+                )

QuantNodes/research/quant_alpha/operator_vocab/metadata.py

@@ -0,0 +1,263 @@
+# coding=utf-8
+"""
+算子元数据 schema - 12 字段（含 7 个 LLM 友好）
+
+扩展自 factor_node/factor_functions/_helpers.py::register_operator 的 5 字段，
+新增 7 个对 LLM 关键的字段。
+
+字段定义：
+- 基础字段（5 个，对应现有 _OPERATOR_REGISTRY）：
+    name, category, func, doc, signature, parameters
+- LLM 友好字段（7 个，本文件新增）：
+    difficulty, category_tags, default_window,
+    requires_group_by, output_dtype, examples, composes_with
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, List, Optional
+
+
+# 与 factor_node.factor_functions._helpers.OperatorCategory 保持一致
+class OperatorCategory:
+    """算子类别常量"""
+    POINT = "point"
+    TIME = "time"
+    SECTION = "section"
+    MULTI_SECTION = "multi_section"
+    TALIB = "talib"
+
+    @classmethod
+    def all(cls) -> List[str]:
+        return [cls.POINT, cls.TIME, cls.SECTION, cls.MULTI_SECTION, cls.TALIB]
+
+    @classmethod
+    def is_valid(cls, category: str) -> bool:
+        return category in cls.all()
+
+
+@dataclass
+class OperatorMetadata:
+    """算子元数据 — 12 字段 schema
+
+    用于：
+    - LLM prompt 注入（Alpha-GPT 路线 6）
+    - 算子文档自动生成
+    - 算子查询/筛选
+    - 算子推荐（composes_with）
+
+    字段分组：
+    - 基础（5）：name, category, func, doc, signature, parameters
+    - LLM 友好（7）：difficulty, category_tags, default_window,
+                requires_group_by, output_dtype, examples, composes_with
+    """
+
+    # === 基础字段（5/6） ===
+    name: str
+    category: str
+    func: Optional[Callable] = None
+    doc: str = ""
+    signature: str = ""
+    parameters: List[str] = field(default_factory=list)
+
+    # === LLM 友好字段（7） ===
+    # 难度等级：1=简单（无参/单参），2=中等（2-3 参），3=复杂（4+ 参或嵌套）
+    difficulty: int = 1
+    # 语义角色标签：central_tendency / dispersion / quantile / polarity /
+    #              position / transformation / aggregation / groupby
+    category_tags: List[str] = field(default_factory=list)
+    # 默认窗口建议：时间序列算子的推荐窗口
+    default_window: List[int] = field(default_factory=list)
+    # 是否需要 group_by(date) 上下文（用于 per-date over() 修复判断）
+    requires_group_by: bool = False
+    # 输出数据类型：float64 / int64 / bool / object
+    output_dtype: str = "float64"
+    # few-shot 示例：每项 {"input": ..., "output": ...}
+    examples: List[Dict[str, Any]] = field(default_factory=list)
+    # 推荐组合的算子名列表
+    composes_with: List[str] = field(default_factory=list)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """转为 dict（用于持久化 / LLM 注入）"""
+        return {
+            "name": self.name,
+            "category": self.category,
+            "doc": self.doc,
+            "signature": self.signature,
+            "parameters": list(self.parameters),
+            "difficulty": self.difficulty,
+            "category_tags": list(self.category_tags),
+            "default_window": list(self.default_window),
+            "requires_group_by": self.requires_group_by,
+            "output_dtype": self.output_dtype,
+            "examples": list(self.examples),
+            "composes_with": list(self.composes_with),
+        }
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "OperatorMetadata":
+        """从 dict 构造"""
+        return cls(
+            name=data["name"],
+            category=data["category"],
+            doc=data.get("doc", ""),
+            signature=data.get("signature", ""),
+            parameters=data.get("parameters", []),
+            difficulty=data.get("difficulty", 1),
+            category_tags=data.get("category_tags", []),
+            default_window=data.get("default_window", []),
+            requires_group_by=data.get("requires_group_by", False),
+            output_dtype=data.get("output_dtype", "float64"),
+            examples=data.get("examples", []),
+            composes_with=data.get("composes_with", []),
+        )
+
+    @classmethod
+    def from_registry_entry(cls, entry: Dict[str, Any]) -> "OperatorMetadata":
+        """从 factor_node 的 _OPERATOR_REGISTRY 条目构造（自动推断 7 字段）
+
+        Args:
+            entry: _OPERATOR_REGISTRY[category][name] 的值，含
+                name, category, func, doc, signature, parameters
+
+        Returns:
+            自动推断 7 字段的 OperatorMetadata
+        """
+        params = entry.get("parameters", [])
+        sig = entry.get("signature", "")
+        name = entry["name"]
+        category = entry["category"]
+
+        # 推断 difficulty
+        difficulty = _infer_difficulty(params, sig)
+
+        # 推断 category_tags
+        category_tags = _infer_category_tags(name, category, params)
+
+        # 推断 default_window
+        default_window = _infer_default_window(name, params)
+
+        # 推断 requires_group_by
+        requires_group_by = _infer_requires_group_by(name, category)
+
+        # 推断 output_dtype
+        output_dtype = _infer_output_dtype(name, category, params)
+
+        return cls(
+            name=name,
+            category=category,
+            func=entry.get("func"),
+            doc=entry.get("doc", ""),
+            signature=sig,
+            parameters=params,
+            difficulty=difficulty,
+            category_tags=category_tags,
+            default_window=default_window,
+            requires_group_by=requires_group_by,
+            output_dtype=output_dtype,
+            examples=[],
+            composes_with=[],
+        )
+
+
+def _infer_difficulty(params: List[str], signature: str) -> int:
+    """推断难度等级
+
+    - 1：0-1 参数
+    - 2：2-3 参数
+    - 3：4+ 参数 或嵌套调用
+    """
+    n = len([p for p in params if p not in ("self", "kwargs", "args")])
+    if n <= 1:
+        return 1
+    if n <= 3:
+        return 2
+    return 3
+
+
+def _infer_category_tags(name: str, category: str, params: List[str]) -> List[str]:
+    """推断语义角色标签
+
+    基于算子名前缀和类别：
+    - mean/median/sum → central_tendency
+    - std/var/skew/kurt → dispersion
+    - quantile/percentile → quantile
+    - sign/abs → polarity
+    - argmax/argmin → position
+    - rank/zscore/winsorize → transformation
+    """
+    tags = []
+    name_lower = name.lower()
+
+    if any(k in name_lower for k in ("mean", "median", "sum", "sma", "ema")):
+        tags.append("central_tendency")
+    if any(k in name_lower for k in ("std", "var", "skew", "kurt", "vol")):
+        tags.append("dispersion")
+    if any(k in name_lower for k in ("quantile", "percentile", "qtl")):
+        tags.append("quantile")
+    if any(k in name_lower for k in ("sign", "abs", "signedpower")):
+        tags.append("polarity")
+    if any(k in name_lower for k in ("argmax", "argmin", "idxmax", "idxmin")):
+        tags.append("position")
+    if any(k in name_lower for k in ("rank", "zscore", "winsorize", "scale", "normalize")):
+        tags.append("transformation")
+    if any(k in name_lower for k in ("corr", "cov", "regress", "beta")):
+        tags.append("correlation")
+    if any(k in name_lower for k in ("delta", "diff", "pct_change", "return", "roc")):
+        tags.append("momentum")
+    if any(k in name_lower for k in ("lag", "lead", "shift", "delay", "ref")):
+        tags.append("shift")
+
+    if category == OperatorCategory.SECTION:
+        tags.append("cross_sectional")
+    elif category == OperatorCategory.MULTI_SECTION:
+        tags.append("groupby")
+    elif category == OperatorCategory.TIME:
+        tags.append("time_series")
+    elif category == OperatorCategory.POINT:
+        tags.append("pointwise")
+
+    return list(set(tags))
+
+
+def _infer_default_window(name: str, params: List[str]) -> List[int]:
+    """推断默认窗口建议
+
+    时序算子默认 [5, 10, 20, 60]
+    """
+    name_lower = name.lower()
+    if any(k in name_lower for k in (
+        "ts_", "rolling_", "ewm_", "expanding_", "decay_", "lag", "shift", "delta"
+    )):
+        return [5, 10, 20, 60]
+    return []
+
+
+def _infer_requires_group_by(name: str, category: str) -> bool:
+    """推断是否需要 group_by(date) 上下文
+
+    截面算子（rank/zscore/winsorize/normalize/IndNeutralize/aggregate）需要。
+    """
+    name_lower = name.lower()
+    if category == OperatorCategory.SECTION:
+        return True
+    if category == OperatorCategory.MULTI_SECTION:
+        return True
+    if any(k in name_lower for k in ("rank", "zscore", "winsorize", "normalize", "neutralize")):
+        return True
+    return False
+
+
+def _infer_output_dtype(name: str, category: str, params: List[str]) -> str:
+    """推断输出数据类型
+
+    - argmax/argmin/sign/condition → int64 或 bool
+    - 其它 → float64
+    """
+    name_lower = name.lower()
+    if any(k in name_lower for k in ("argmax", "argmin", "idxmax", "idxmin", "row_number")):
+        return "int64"
+    if any(k in name_lower for k in ("sign", "is_", "where", "cond", "isnull", "notnull")):
+        return "bool"
+    return "float64"