quantnodes 3.0.0__py3-none-any.whl

QuantNodes/core/knowledge/metrics/evaluator.py

@@ -0,0 +1,176 @@
+"""RAGEvaluator — 汇总 5 个指标, 输出统一报告。"""
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass, field
+from datetime import datetime
+from pathlib import Path
+
+import pandas as pd
+
+from QuantNodes.core.path_utils import ensure_parent
+
+from .metrics import (
+    hit_rate_at_k,
+    intra_list_diversity,
+    jaccard_similarity,
+    lineage_coverage,
+    mean_hit_rate_at_k,
+    mean_lineage_coverage,
+    mean_ndcg_at_k,
+    mean_reciprocal_rank,
+    ndcg_at_k,
+    reciprocal_rank,
+)
+
+
+@dataclass
+class QueryResult:
+    """单个 query 的评估结果。"""
+    query: str
+    retrieved_ids: list[str]
+    relevant_ids: list[str]
+    relevance_scores: dict[str, float] = field(default_factory=dict)
+    lineage_ids: list[str] = field(default_factory=list)
+    hit_at_5: float = 0.0
+    hit_at_10: float = 0.0
+    ndcg_at_5: float = 0.0
+    ndcg_at_10: float = 0.0
+    mrr: float = 0.0
+    lineage_cov: float = 0.0
+    diversity: float = 0.0
+
+
+@dataclass
+class EvalReport:
+    """评估汇总报告。"""
+    n_queries: int
+    hit_at_5: float
+    hit_at_10: float
+    ndcg_at_5: float
+    ndcg_at_10: float
+    mrr: float
+    lineage_coverage: float
+    diversity: float
+    per_query: list[QueryResult] = field(default_factory=list)
+    timestamp: datetime = field(default_factory=datetime.now)
+
+    def to_dict(self) -> dict:
+        return {
+            "n_queries": self.n_queries,
+            "hit_at_5": self.hit_at_5,
+            "hit_at_10": self.hit_at_10,
+            "ndcg_at_5": self.ndcg_at_5,
+            "ndcg_at_10": self.ndcg_at_10,
+            "mrr": self.mrr,
+            "lineage_coverage": self.lineage_coverage,
+            "diversity": self.diversity,
+            "timestamp": self.timestamp.isoformat(),
+            "per_query": [asdict(q) for q in self.per_query],
+        }
+
+
+class RAGEvaluator:
+    """RAG 评估器 — 接受多 query 结果, 汇总 5 个指标。
+
+    Args:
+        k_values: HitRate/NDCG 评估的 K 列表 (默认 [5, 10])
+    """
+
+    def __init__(self, k_values: list[int] | None = None):
+        self.k_values = k_values or [5, 10]
+
+    def evaluate(
+        self,
+        queries: list[str],
+        retrieved: list[list[str]],
+        relevant: list[list[str]],
+        relevance_scores: list[dict[str, float]] | None = None,
+        lineage_ids: list[list[str]] | None = None,
+        token_lists: list[list[list[str]]] | None = None,
+    ) -> EvalReport:
+        """评估多 query 结果。
+
+        Args:
+            queries: 查询文本列表
+            retrieved: 每 query 的检索结果 (entry_id 列表, 有序)
+            relevant: 每 query 的相关 entry_id 集合
+            relevance_scores: 每 query 的 doc_id -> relevance (NDCG 用, 可选)
+            lineage_ids: 每 query 的 ground truth 谱系 entry_id (可选)
+            token_lists: 每 query 的检索条目 token 列表 (diversity 用, 可选)
+
+        Returns:
+            EvalReport
+        """
+        n = len(queries)
+        relevance_scores = relevance_scores or [{} for _ in range(n)]
+        lineage_ids = lineage_ids or [[] for _ in range(n)]
+        token_lists = token_lists or [[] for _ in range(n)]
+
+        # Per-query
+        per_query: list[QueryResult] = []
+        for q, ret, rel, scores, lin, toks in zip(
+            queries, retrieved, relevant, relevance_scores, lineage_ids, token_lists
+        ):
+            k5 = self.k_values[0] if len(self.k_values) > 0 else 5
+            k10 = self.k_values[1] if len(self.k_values) > 1 else 10
+            h5 = hit_rate_at_k(ret, rel, k=k5)
+            h10 = hit_rate_at_k(ret, rel, k=k10)
+            n5 = ndcg_at_k(ret, scores, k=k5)
+            n10 = ndcg_at_k(ret, scores, k=k10)
+            mrr = reciprocal_rank(ret, rel)
+            lin_cov = lineage_coverage(ret, lin)
+            # diversity: 1 - avg pairwise jaccard of token_lists
+            if len(toks) >= 2:
+                sims = [
+                    jaccard_similarity(toks[i], toks[j])
+                    for i in range(len(toks))
+                    for j in range(i + 1, len(toks))
+                ]
+                div = 1.0 - (sum(sims) / len(sims))
+            else:
+                div = 1.0
+            per_query.append(QueryResult(
+                query=q, retrieved_ids=list(ret), relevant_ids=list(rel),
+                relevance_scores=dict(scores), lineage_ids=list(lin),
+                hit_at_5=h5, hit_at_10=h10, ndcg_at_5=n5, ndcg_at_10=n10,
+                mrr=mrr, lineage_cov=lin_cov, diversity=div,
+            ))
+
+        # 汇总
+        h5_mean = mean_hit_rate_at_k(retrieved, relevant, k=self.k_values[0])
+        h10_mean = mean_hit_rate_at_k(
+            retrieved, relevant, k=self.k_values[1] if len(self.k_values) > 1 else 10
+        )
+        n5_mean = mean_ndcg_at_k(retrieved, relevance_scores, k=self.k_values[0])
+        n10_mean = mean_ndcg_at_k(
+            retrieved, relevance_scores, k=self.k_values[1] if len(self.k_values) > 1 else 10
+        )
+        mrr_mean = mean_reciprocal_rank(retrieved, relevant)
+        lin_mean = mean_lineage_coverage(retrieved, lineage_ids)
+        div_mean = intra_list_diversity(token_lists)
+
+        return EvalReport(
+            n_queries=n,
+            hit_at_5=h5_mean, hit_at_10=h10_mean,
+            ndcg_at_5=n5_mean, ndcg_at_10=n10_mean,
+            mrr=mrr_mean,
+            lineage_coverage=lin_mean,
+            diversity=div_mean,
+            per_query=per_query,
+        )
+
+    def save(self, report: EvalReport, path: Path | str) -> None:
+        """保存为 JSON。"""
+        path = Path(path)
+        ensure_parent(path)
+        with path.open("w", encoding="utf-8") as f:
+            json.dump(report.to_dict(), f, ensure_ascii=False, indent=2)
+
+    def save_csv(self, report: EvalReport, path: Path | str) -> None:
+        """保存 per-query 为 CSV。"""
+        path = Path(path)
+        ensure_parent(path)
+        rows = [asdict(q) for q in report.per_query]
+        df = pd.DataFrame(rows)
+        df.to_csv(path, index=False)

QuantNodes/core/knowledge/metrics/metrics.py

@@ -0,0 +1,220 @@
+"""RAG 评估指标 — 5 个核心指标。
+
+- HitRate@K:    检索 Top-K 是否命中相关 entry
+- NDCG@K:       Normalized Discounted Cumulative Gain (位置权重)
+- MRR:          Mean Reciprocal Rank (首个相关 entry 的倒数排名)
+- LineageCoverage: 检索结果覆盖的谱系比例 (vs. ground truth 谱系)
+- IntraListDiversity: 检索结果内部多样性 (1 - 平均 pairwise similarity)
+"""
+from __future__ import annotations
+
+import math
+from typing import Iterable, Mapping, Sequence
+
+
+# ============================================================================
+# 1. HitRate@K
+# ============================================================================
+
+def hit_rate_at_k(
+    retrieved_ids: Sequence[str],
+    relevant_ids: Iterable[str],
+    k: int = 5,
+) -> float:
+    """检索 Top-K 中是否含至少 1 个 relevant entry。
+
+    Args:
+        retrieved_ids: 检索器返回的 entry_id 列表 (有序)
+        relevant_ids: 真实相关 entry_id 集合
+        k: 截断 K
+
+    Returns:
+        float: 1.0 (命中) / 0.0 (未命中)
+    """
+    rel_set = set(relevant_ids)
+    top_k = retrieved_ids[:k]
+    return 1.0 if any(rid in rel_set for rid in top_k) else 0.0
+
+
+def mean_hit_rate_at_k(
+    queries_retrieved: Sequence[Sequence[str]],
+    queries_relevant: Sequence[Iterable[str]],
+    k: int = 5,
+) -> float:
+    """多 query 平均 HitRate@K。"""
+    if not queries_retrieved:
+        return 0.0
+    scores = [
+        hit_rate_at_k(ret, rel, k=k)
+        for ret, rel in zip(queries_retrieved, queries_relevant)
+    ]
+    return sum(scores) / len(scores)
+
+
+# ============================================================================
+# 2. NDCG@K
+# ============================================================================
+
+def dcg_at_k(
+    retrieved_ids: Sequence[str],
+    relevance_scores: Mapping[str, float],
+    k: int = 5,
+) -> float:
+    """Discounted Cumulative Gain。
+
+    relevance_scores: doc_id -> relevance (通常 0/1, 可更高)
+    """
+    dcg = 0.0
+    for i, rid in enumerate(retrieved_ids[:k]):
+        rel = relevance_scores.get(rid, 0.0)
+        # 标准公式: rel_i / log2(i+2)  (i=0 -> log2(2)=1)
+        dcg += rel / math.log2(i + 2)
+    return dcg
+
+
+def ndcg_at_k(
+    retrieved_ids: Sequence[str],
+    relevance_scores: Mapping[str, float],
+    k: int = 5,
+) -> float:
+    """Normalized DCG (0-1)。"""
+    actual = dcg_at_k(retrieved_ids, relevance_scores, k=k)
+    # ideal: 按 relevance 降序排列的 DCG
+    ideal_rels = sorted(relevance_scores.values(), reverse=True)[:k]
+    ideal = sum(rel / math.log2(i + 2) for i, rel in enumerate(ideal_rels))
+    if ideal == 0:
+        return 0.0
+    return actual / ideal
+
+
+def mean_ndcg_at_k(
+    queries_retrieved: Sequence[Sequence[str]],
+    queries_relevance: Sequence[Mapping[str, float]],
+    k: int = 5,
+) -> float:
+    """多 query 平均 NDCG@K。"""
+    if not queries_retrieved:
+        return 0.0
+    scores = [
+        ndcg_at_k(ret, rel, k=k)
+        for ret, rel in zip(queries_retrieved, queries_relevance)
+    ]
+    return sum(scores) / len(scores)
+
+
+# ============================================================================
+# 3. MRR (Mean Reciprocal Rank)
+# ============================================================================
+
+def reciprocal_rank(
+    retrieved_ids: Sequence[str],
+    relevant_ids: Iterable[str],
+) -> float:
+    """首个 relevant entry 的倒数排名 (0 表示未命中)。"""
+    rel_set = set(relevant_ids)
+    for i, rid in enumerate(retrieved_ids, 1):
+        if rid in rel_set:
+            return 1.0 / i
+    return 0.0
+
+
+def mean_reciprocal_rank(
+    queries_retrieved: Sequence[Sequence[str]],
+    queries_relevant: Sequence[Iterable[str]],
+) -> float:
+    """多 query 平均 MRR。"""
+    if not queries_retrieved:
+        return 0.0
+    return sum(
+        reciprocal_rank(ret, rel)
+        for ret, rel in zip(queries_retrieved, queries_relevant)
+    ) / len(queries_retrieved)
+
+
+# ============================================================================
+# 4. Lineage Coverage (谱系覆盖率)
+# ============================================================================
+
+def lineage_coverage(
+    retrieved_ids: Iterable[str],
+    lineage_ids: Iterable[str],
+) -> float:
+    """检索结果中包含的 ground truth 谱系 entry 比例。
+
+    Args:
+        retrieved_ids: 检索器返回的 entry_id 列表
+        lineage_ids: ground truth 谱系 (e.g. expand_lineage 的 ancestors + descendants)
+
+    Returns:
+        float: 0-1, 1.0 = 谱系完全覆盖
+    """
+    ret_set = set(retrieved_ids)
+    lineage_set = set(lineage_ids)
+    if not lineage_set:
+        return 0.0
+    covered = sum(1 for lid in lineage_set if lid in ret_set)
+    return covered / len(lineage_set)
+
+
+def mean_lineage_coverage(
+    queries_retrieved: Sequence[Iterable[str]],
+    queries_lineage: Sequence[Iterable[str]],
+) -> float:
+    """多 query 平均 lineage coverage。"""
+    if not queries_retrieved:
+        return 0.0
+    return sum(
+        lineage_coverage(ret, lin)
+        for ret, lin in zip(queries_retrieved, queries_lineage)
+    ) / len(queries_retrieved)
+
+
+# ============================================================================
+# 5. Intra-List Diversity
+# ============================================================================
+
+def intra_list_diversity(
+    items: Sequence[Sequence[Sequence[str]]],
+    similarity_fn=None,
+) -> float:
+    """检索结果内部多样性: 1 - 平均 pairwise jaccard 相似度。
+
+    Args:
+        items: 多 query 的 token 化结果。结构 = Sequence[Sequence[Sequence[str]]]
+            - 外层: 每个 query 一份
+            - 中层: 每 query 的 Top-K 个条目
+            - 内层: 每个条目的 token 列表
+            例: [[['momentum', 'close'], ['reversal', 'open']], ...]
+        similarity_fn: 自定义相似度函数 (默认 jaccard_similarity)
+
+    Returns:
+        float: 0-1, 1.0 = 完全多样
+    """
+    if not items:
+        return 0.0
+    if similarity_fn is None:
+        similarity_fn = jaccard_similarity
+
+    diversities: list[float] = []
+    for tokens_list in items:
+        if len(tokens_list) < 2:
+            diversities.append(1.0)  # 单元素算最大多样
+            continue
+        sims: list[float] = []
+        for i in range(len(tokens_list)):
+            for j in range(i + 1, len(tokens_list)):
+                sims.append(similarity_fn(tokens_list[i], tokens_list[j]))
+        avg_sim = sum(sims) / len(sims) if sims else 0.0
+        diversities.append(1.0 - avg_sim)
+    return sum(diversities) / len(diversities)
+
+
+def jaccard_similarity(a: Sequence[str], b: Sequence[str]) -> float:
+    """Jaccard 相似度 = |A ∩ B| / |A ∪ B|。"""
+    set_a, set_b = set(a), set(b)
+    if not set_a and not set_b:
+        return 0.0
+    union = set_a | set_b
+    if not union:
+        return 0.0
+    return len(set_a & set_b) / len(union)

QuantNodes/core/knowledge/rag_prompt.py

@@ -0,0 +1,196 @@
+"""RAG Prompt — 从检索结果构造带上下文的 prompt。
+
+用途: Hypothesizer / Mutator 在生成新因子前, 检索 top-k 历史因子作为 in-context 示例。
+
+Week 8 升级: 谱系 RAG — 每个示例附 ancestors/descendants 上下文, 树状展开。
+Week 9 升级: 谱系压缩 — use_compress=True 时, 祖先/后裔段先用 LLM/启发式
+            总结为 1 段简短描述, 减少 token 消耗。
+"""
+from __future__ import annotations
+
+from typing import Optional
+
+from ..trajectory import TrajectoryEntry, TrajectoryPool
+from .knowledge_base import KnowledgeBase
+from .lineage_compress import Compressor
+from .lineage_expand import expand_lineage
+
+
+_RAG_HEADER = """你是一个量化研究员, 负责基于历史经验和当前研究假设生成新 alpha 因子。
+下方是历史表现良好的 {n_examples} 个因子作为参考, 含完整演化谱系 (ancestors / descendants):
+"""
+
+_RAG_EXAMPLE_TEMPLATE = """---
+示例 {idx}: {name}
+表达式: {expression}
+描述: {description}
+指标: sharpe={sharpe}, arr={arr}, ic_mean={ic_mean}
+"""
+
+_LINEAGE_RELATION_TEMPLATE = (
+    "{relation} (depth={depth}): {name} | sharpe={sharpe} | {expression}"
+)
+
+_LINEAGE_COMPRESSED_TEMPLATE = """{relation} ({n} entries): {summary}"""
+
+_RAG_TASK_TEMPLATE = """
+现在, 请基于以下研究假设生成新因子:
+研究假设: {hypothesis}
+补充描述: {description}
+
+请综合参考示例的设计思路, 生成一个与历史不同但经济意义清晰的因子。
+返回 JSON: {{"name": "因子名", "expression": "代码表达式", "description": "因子描述"}}
+"""
+
+
+def build_rag_prompt(
+    direction: str,
+    description: str,
+    kb: Optional[KnowledgeBase] = None,
+    top_k: int = 3,
+    min_score: float = 0.01,
+    include_lineage: bool = True,
+    max_ancestor_depth: int = 2,
+    max_descendant_depth: int = 2,
+    use_compress: bool = False,
+    compressor: Optional[Compressor] = None,
+) -> str:
+    """构造带 RAG 上下文的 prompt (Week 9 升级版: 含谱系 + 压缩)。
+
+    Args:
+        direction: 研究方向 (hypothesis)
+        description: 补充描述
+        kb: KnowledgeBase (None=不附 RAG 上下文)
+        top_k: 检索 top-k 数量
+        min_score: 最小相似度阈值
+        include_lineage: 是否附加祖先/后裔
+        max_ancestor_depth: 祖先展开深度
+        max_descendant_depth: 后裔展开深度
+        use_compress: 是否压缩谱系段 (默认 False, 沿用 Week 8 多行格式)
+        compressor: Compressor 实例 (use_compress=True 时必填,
+                    内部未传时自动构造 mock Compressor)
+
+    Returns:
+        str: 完整 prompt
+    """
+    rag_section = ""
+    if kb is not None and len(kb) > 0:
+        query = f"{direction} {description}"
+        results = kb.query(query, top_k=top_k, min_score=min_score)
+        if results:
+            example_parts = []
+            for i, (entry, score) in enumerate(results, 1):
+                example_parts.append(_format_example(i, entry, score))
+                if include_lineage and kb.pool is not None and entry is not None:
+                    lineage_str = _format_lineage(
+                        kb.pool, entry.entry_id,
+                        max_ancestor_depth=max_ancestor_depth,
+                        max_descendant_depth=max_descendant_depth,
+                        use_compress=use_compress,
+                        compressor=compressor,
+                    )
+                    if lineage_str:
+                        example_parts.append(lineage_str)
+            rag_section = (
+                _RAG_HEADER.format(n_examples=len(results))
+                + "\n".join(example_parts)
+            )
+    return rag_section + _RAG_TASK_TEMPLATE.format(
+        hypothesis=direction, description=description,
+    )
+
+
+def _format_example(idx: int, entry: TrajectoryEntry | None, score: float) -> str:
+    """格式化单个示例。"""
+    if entry is None:
+        return _RAG_EXAMPLE_TEMPLATE.format(
+            idx=idx, name="<unknown>", expression="", description="",
+            sharpe="?", arr="?", ic_mean="?",
+        )
+    cfg = entry.config_snapshot or {}
+    factor_cfg = cfg.get("factor", {}) if isinstance(cfg, dict) else {}
+    metrics = entry.metrics or {}
+    return _RAG_EXAMPLE_TEMPLATE.format(
+        idx=idx,
+        name=factor_cfg.get("name", entry.entry_id[:8]),
+        expression=factor_cfg.get("expression", ""),
+        description=factor_cfg.get("description", ""),
+        sharpe=metrics.get("sharpe", "?"),
+        arr=metrics.get("arr", "?"),
+        ic_mean=metrics.get("ic_mean", "?"),
+    )
+
+
+def _format_lineage(
+    pool: TrajectoryPool,
+    root_id: str,
+    max_ancestor_depth: int = 2,
+    max_descendant_depth: int = 2,
+    use_compress: bool = False,
+    compressor: Optional[Compressor] = None,
+) -> str:
+    """格式化谱系上下文, 嵌入示例后。"""
+    expanded = expand_lineage(
+        pool, root_id,
+        max_ancestor_depth=max_ancestor_depth,
+        max_descendant_depth=max_descendant_depth,
+    )
+    ancestors = expanded["ancestors"]
+    descendants = expanded["descendants"]
+
+    if not ancestors and not descendants:
+        return ""
+
+    if use_compress:
+        # Week 9: 压缩模式
+        comp = compressor or Compressor(model="mock")
+        return _format_lineage_compressed(ancestors, descendants, comp)
+
+    # Week 8: 多行展开模式
+    return _format_lineage_expanded(ancestors, descendants)
+
+
+def _format_lineage_expanded(ancestors, descendants) -> str:
+    """Week 8 风格: 多行展开。"""
+    parts: list[str] = ["  谱系上下文:"]
+    for depth, entry in ancestors:
+        cfg = entry.config_snapshot or {}
+        factor_cfg = cfg.get("factor", {}) if isinstance(cfg, dict) else {}
+        parts.append("  " + _LINEAGE_RELATION_TEMPLATE.format(
+            relation="↑ ancestor",
+            depth=depth,
+            name=factor_cfg.get("name", entry.entry_id[:8]),
+            sharpe=(entry.metrics or {}).get("sharpe", "?"),
+            expression=factor_cfg.get("expression", "")[:40],
+        ))
+    for depth, entry in descendants:
+        cfg = entry.config_snapshot or {}
+        factor_cfg = cfg.get("factor", {}) if isinstance(cfg, dict) else {}
+        parts.append("  " + _LINEAGE_RELATION_TEMPLATE.format(
+            relation="↓ descendant",
+            depth=depth,
+            name=factor_cfg.get("name", entry.entry_id[:8]),
+            sharpe=(entry.metrics or {}).get("sharpe", "?"),
+            expression=factor_cfg.get("expression", "")[:40],
+        ))
+    return "\n".join(parts)
+
+
+def _format_lineage_compressed(ancestors, descendants, compressor: Compressor) -> str:
+    """Week 9 风格: 1 行总结 ancestors + 1 行总结 descendants。"""
+    parts: list[str] = ["  谱系上下文 (压缩):"]
+    if ancestors:
+        c_anc = compressor.compress(ancestors, relation="ancestors")
+        parts.append("  " + _LINEAGE_COMPRESSED_TEMPLATE.format(
+            relation="↑ ancestors",
+            n=c_anc.original_count,
+            summary=c_anc.summary,
+        ))
+    if descendants:
+        c_desc = compressor.compress(descendants, relation="descendants")
+        parts.append("  " + _LINEAGE_COMPRESSED_TEMPLATE.format(
+            relation="↓ descendants",
+            n=c_desc.original_count,
+            summary=c_desc.summary,
+        ))
+    return "\n".join(parts)