plot_headmap.py

"""渲染 A/B Test 的热力图（Context Length × Depth Percent → Score）。

读取字段（`needlehaystack/stores/jsonl.py` 写出的 schema_version 2 行）：

- ``context_length``      —— sweep 的 context 长度（token 数）
- ``target_depth_percent`` —— sweep 的目标深度百分比
- ``score.value``         —— 命中得分，1.0 全绿 / 0.0 全红
- ``status``              —— ``"ok"`` / ``"error"``；error 行的 score 为 null

渲染逻辑要点：

- 以 ``(context_length, target_depth_percent)`` 为格子 pivot；重复行取均值，
  因此 A/B 两组的 JSONL 混在同一个文件里也能各自成图（先按 ``run_name`` 过滤）。
- ``status == "error"`` 的格子记为 ``NaN`` 而不是 0.0 —— 把"请求失败"误画成
  "检索失败"会让人读出完全错误的结论。
- 深度轴反转：0% 在上、100% 在下，与原始仓库的可视化方向一致。

用法::

    python plot_headmap.py                       # 画默认的两张图
    python plot_headmap.py --run a               # 只画 A 组（results/deepseek-a-baseline.jsonl）
    python plot_headmap.py --run b               # 只画 B 组（results/deepseek-b-phaseshift.jsonl）
    python plot_headmap.py --a X.jsonl --b Y.jsonl --out-dir out
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import matplotlib

# 无显示环境（CI / 远程）下用非交互后端，避免 "no display name" 报错。
if not sys.stdout.isatty():
    matplotlib.use("Agg")

import matplotlib.pyplot as plt  # noqa: E402
import pandas as pd  # noqa: E402
import seaborn as sns  # noqa: E402

# 默认的两组结果文件与标题。键是 `--run` 的取值。
# 与 configs/runs/deepseek-a-test.yaml / deepseek-b-test.yaml 的
# store.path 保持一致。
DEFAULT_RUNS: dict[str, tuple[str, str]] = {
    "a": ("results/deepseek-a-baseline.jsonl", "A Baseline (Random Phase)"),
    "b": ("results/deepseek-b-phaseshift.jsonl", "B Aligned (Phase 0 Protected)"),
}


def load_rows(path: Path, run_name: str | None = None) -> pd.DataFrame:
    """读 JSONL，返回含 ``context_length`` / ``depth_percent`` / ``score`` 的 DataFrame。

    ``run_name`` 非空时只保留该 ``run_name`` 的行，便于两组结果写到同一个
    文件后再分开出图。

    同一 ``(context_length, depth_percent)`` 出现多行时打印告警：正常情况下
    一个格子只有一行（单 seed）。多行意味着同一组结果被重复追加过——
    ``resume: true`` 只会跳过已存在的 key，但如果两次运行的 ``run_name``
    不同、store 路径却相同，就会静默产生重复格，让热力图"看起来有数据"。
    """
    rows: list[dict[str, float | int | None]] = []
    with open(path, encoding="utf-8") as f:
        for line_no, raw in enumerate(f, start=1):
            stripped = raw.strip()
            if not stripped:
                continue
            try:
                obj = json.loads(stripped)
            except json.JSONDecodeError:
                print(f"warning: {path}:{line_no} 不是合法 JSON，已跳过", file=sys.stderr)
                continue
            if run_name is not None and obj.get("run_name") != run_name:
                continue

            score = obj.get("score")
            # error 行的 score 是 null。把它和"真的 0 分"区分开：
            # NaN 在热力图上是灰色，不会被误读成检索失败。
            value = None if score is None else score.get("value")
            if obj.get("status") == "error":
                value = None
            rows.append(
                {
                    "context_length": int(obj["context_length"]),
                    "depth_percent": float(obj["target_depth_percent"]),
                    "score": value,
                    "run_name": obj.get("run_name"),
                }
            )

    if not rows:
        raise SystemExit(f"{path} 里没有可用的结果行（run_name={run_name!r}）")

    df = pd.DataFrame(rows)
    dupes = df.groupby(["context_length", "depth_percent"]).size()
    dupes = dupes[dupes > 1]
    if len(dupes):
        names = sorted({str(n) for n in df["run_name"].dropna().unique()})
        print(
            f"warning: {path} 有 {len(dupes)} 个格子出现多行"
            f"（涉及的 run_name: {names or '?'}）；"
            "同一格多行通常意味着两次运行写到了同一个文件，"
            "请检查 run 配置的 run_name 与 store.path 是否配套。",
            file=sys.stderr,
        )
    return df.drop(columns=["run_name"])


def plot_jsonl(file_path: str | Path, title: str, out: Path | None = None) -> None:
    df = load_rows(Path(file_path))

    # 同一格跑多次（多个 seed）时取均值，保证 pivot 不会因重复键报错。
    pivot = df.pivot_table(
        index="depth_percent",
        columns="context_length",
        values="score",
        aggfunc="mean",
    )

    plt.figure(figsize=(12, 8))
    # vmin/vmax 固定 0/1，NaN 用单独的颜色，避免不同图的色标不可比。
    sns.heatmap(
        pivot,
        cmap="RdYlGn",
        vmin=0,
        vmax=1,
        annot=False,
        cbar_kws={"label": "score (1.0 = hit, NaN = request error)"},
    )
    plt.title(title)
    plt.xlabel("Context Length (tokens)")
    plt.ylabel("Depth Percent")
    plt.gca().invert_yaxis()  # 深度 0 在上，100 在下

    if out is None:
        plt.show()
    else:
        out.parent.mkdir(parents=True, exist_ok=True)
        plt.savefig(out, dpi=150, bbox_inches="tight")
        print(f"wrote {out}")
    plt.close()


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--run",
        choices=sorted(DEFAULT_RUNS),
        action="append",
        help="只画指定组（可重复）。默认两组都画。",
    )
    parser.add_argument("--a", help="覆盖 A 组结果文件路径")
    parser.add_argument("--b", help="覆盖 B 组结果文件路径")
    parser.add_argument(
        "--out-dir",
        help="输出目录。给了就存 PNG 不弹窗；不给则 plt.show()。",
    )
    args = parser.parse_args(argv)

    chosen = args.run or list(DEFAULT_RUNS)
    overrides = {"a": args.a, "b": args.b}
    for name in chosen:
        default_path, title = DEFAULT_RUNS[name]
        path = overrides.get(name) or default_path
        out = None
        if args.out_dir:
            out = Path(args.out_dir) / f"heatmap-{name}.png"
        try:
            plot_jsonl(path, title, out=out)
        except FileNotFoundError:
            print(f"warning: 找不到 {path}，先跑对应组的 niah run", file=sys.stderr)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())