kernel-design-agents/plot_experiment_results.py

#!/usr/bin/env python3
"""Plot comparable latency progress from the kernel ablation benchmark ledgers.

Only complete, all-PASS canonical official-suite runs are comparable:
  * GDN prefill: 100 workloads
  * FP8 MoE:     19 workloads
  * DSA:         23 workloads

The time plots include clean and dirty-worktree development runs.  The token
plots select the lowest-mean clean run for each model/workflow when available.
"""

from __future__ import annotations

import argparse
import bisect
import csv
import json
import math
import statistics
from collections import Counter, defaultdict
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path

import matplotlib.pyplot as plt
from matplotlib.lines import Line2D


ROOT = Path(__file__).resolve().parent

KERNELS = {
    "gdn-prefill": {
        "title": "GDN Prefill",
        "expected": 100,
        "token_key": "total_seq_len",
        "token_label": "Total sequence tokens",
        "xscale": "log",
    },
    "fp8-moe": {
        "title": "FP8 MoE",
        "expected": 19,
        "token_key": "seq_len",
        "token_label": "Sequence tokens",
        "xscale": "log",
    },
    "dsa": {
        "title": "DeepSeek Sparse Attention",
        "expected": 23,
        "token_key": "num_tokens",
        "token_label": "Query tokens",
        "xscale": "linear",
    },
}

MODELS = {
    "claude-opus-4.8": ("Opus 4.8", "#1f77b4"),
    "fable-5": ("Fable 5", "#ff7f0e"),
    "gpt-5.5": ("GPT-5.5", "#2ca02c"),
    "gpt-5.6-sol": ("GPT-5.6-sol", "#d62728"),
}

WORKFLOWS = {
    "goal": ("goal", "-"),
    "goal-arar": ("goal-arar", "-"),
    "kda-humanize": ("kda-humanize", "--"),
}


@dataclass(frozen=True)
class Series:
    model_dir: str
    workflow_dir: str
    kernel: str
    path: Path

    @property
    def model(self) -> str:
        return MODELS[self.model_dir][0]

    @property
    def workflow(self) -> str:
        return WORKFLOWS[self.workflow_dir][0]

    @property
    def label(self) -> str:
        return f"{self.model} | {self.workflow}"

    @property
    def color(self) -> str:
        return MODELS[self.model_dir][1]

    @property
    def linestyle(self) -> str:
        return WORKFLOWS[self.workflow_dir][1]


@dataclass
class Run:
    series: Series
    timestamp: datetime
    timestamp_text: str
    commit: str
    dirty: bool
    rows: list[dict[str, str]]
    mean_kernel_ms: float
    mean_baseline_ms: float
    arithmetic_mean_speedup: float


def parse_bool(value: str) -> bool:
    return value.strip().lower() == "true"


def discover_series() -> list[Series]:
    found: list[Series] = []
    for model_dir in MODELS:
        for workflow_dir in WORKFLOWS:
            for kernel in KERNELS:
                path = ROOT / model_dir / workflow_dir / kernel / "benchmark.csv"
                if path.exists():
                    found.append(Series(model_dir, workflow_dir, kernel, path))
    return found


def read_valid_runs(series: Series) -> list[Run]:
    expected = KERNELS[series.kernel]["expected"]
    groups: dict[tuple[str, str, str, str, str], list[dict[str, str]]] = defaultdict(list)
    with series.path.open(newline="") as handle:
        for row in csv.DictReader(handle):
            if row.get("suite") != "official":
                continue
            if int(row.get("workload_count", "0")) != expected:
                continue
            key = (
                row["timestamp_utc"],
                row["git_commit"],
                row["git_dirty"],
                row["op"],
                row["workload_count"],
            )
            groups[key].append(row)

    runs: list[Run] = []
    for key, rows in groups.items():
        if len(rows) != expected:
            continue
        if len({row["workload_id"] for row in rows}) != expected:
            continue
        if not all(parse_bool(row["passed"]) for row in rows):
            continue
        try:
            kernel_ms = [float(row["kernel_ms"]) for row in rows]
            baseline_ms = [float(row["baseline_ms"]) for row in rows]
            speedups = [float(row["speedup"]) for row in rows]
        except (KeyError, ValueError):
            continue
        if not all(math.isfinite(x) and x > 0 for x in kernel_ms + baseline_ms):
            continue
        timestamp_text, commit, dirty_text, _, _ = key
        runs.append(
            Run(
                series=series,
                timestamp=datetime.fromisoformat(timestamp_text.replace("Z", "+00:00")),
                timestamp_text=timestamp_text,
                commit=commit,
                dirty=parse_bool(dirty_text),
                rows=rows,
                mean_kernel_ms=statistics.fmean(kernel_ms),
                mean_baseline_ms=statistics.fmean(baseline_ms),
                arithmetic_mean_speedup=statistics.fmean(speedups),
            )
        )
    return sorted(runs, key=lambda run: run.timestamp)


def cumulative_timeline(events: list[tuple[datetime, int]]) -> list[tuple[datetime, int]]:
    events.sort(key=lambda item: item[0])
    total = 0
    timeline = []
    for timestamp, count in events:
        if count <= 0:
            continue
        total += count
        timeline.append((timestamp, total))
    return timeline


def parse_timestamp(value: str | None) -> datetime | None:
    if not value:
        return None
    try:
        return datetime.fromisoformat(value.replace("Z", "+00:00"))
    except ValueError:
        return None


def claude_token_timeline(series: Series) -> tuple[list[tuple[datetime, int]], str]:
    encoded = str(series.path.parent.resolve()).replace("/", "-").replace(".", "-")
    project_dir = Path.home() / ".claude" / "projects" / encoded
    events: list[tuple[datetime, int]] = []
    seen_responses: set[str] = set()
    for path in project_dir.rglob("*.jsonl") if project_dir.exists() else []:
        with path.open(errors="replace") as handle:
            for line in handle:
                try:
                    record = json.loads(line)
                except json.JSONDecodeError:
                    continue
                message = record.get("message") or {}
                usage = message.get("usage")
                if record.get("type") != "assistant" or not isinstance(usage, dict):
                    continue
                response_id = record.get("requestId") or message.get("id") or record.get("uuid")
                if not response_id or response_id in seen_responses:
                    continue
                timestamp = parse_timestamp(record.get("timestamp"))
                if timestamp is None:
                    continue
                seen_responses.add(response_id)
                count = sum(
                    int(usage.get(field, 0) or 0)
                    for field in (
                        "input_tokens",
                        "cache_creation_input_tokens",
                        "cache_read_input_tokens",
                        "output_tokens",
                    )
                )
                events.append((timestamp, count))
    return cumulative_timeline(events), "Claude session JSONL (main + subagents)"


def codex_session_files(cwd: Path) -> list[Path]:
    result = []
    session_root = Path.home() / ".codex" / "sessions"
    for path in session_root.rglob("*.jsonl") if session_root.exists() else []:
        try:
            with path.open() as handle:
                first = json.loads(handle.readline())
            if first.get("type") != "session_meta":
                continue
            recorded_cwd = Path(first["payload"]["cwd"]).resolve()
        except (OSError, KeyError, json.JSONDecodeError):
            continue
        if recorded_cwd == cwd.resolve():
            result.append(path)
    return result


def codex_token_timeline(series: Series) -> tuple[list[tuple[datetime, int]], str]:
    events: list[tuple[datetime, int]] = []
    files = codex_session_files(series.path.parent)
    for path in files:
        previous_total = 0
        with path.open(errors="replace") as handle:
            for line in handle:
                try:
                    record = json.loads(line)
                except json.JSONDecodeError:
                    continue
                payload = record.get("payload") or {}
                if record.get("type") != "event_msg" or payload.get("type") != "token_count":
                    continue
                total = (
                    ((payload.get("info") or {}).get("total_token_usage") or {}).get("total_tokens")
                )
                timestamp = parse_timestamp(record.get("timestamp"))
                if total is None or timestamp is None:
                    continue
                total = int(total)
                delta = total - previous_total if total >= previous_total else total
                previous_total = total
                if delta > 0:
                    events.append((timestamp, delta))
    return cumulative_timeline(events), f"Codex token_count events ({len(files)} sessions)"


def omh_token_timeline(series: Series) -> tuple[list[tuple[datetime, int]], str]:
    artifact_dir = series.path.parent / "workflow-output" / "omh-runtime" / "artifacts"
    events: list[tuple[datetime, int]] = []
    seen_responses: set[str] = set()
    files = list(artifact_dir.rglob("*.jsonl")) if artifact_dir.exists() else []
    for path in files:
        with path.open(errors="replace") as handle:
            for line in handle:
                try:
                    record = json.loads(line)
                except json.JSONDecodeError:
                    continue
                message = record.get("message") or {}
                usage = message.get("usage")
                if (
                    record.get("type") != "message"
                    or message.get("role") != "assistant"
                    or not isinstance(usage, dict)
                ):
                    continue
                response_id = message.get("responseId") or record.get("id")
                if not response_id or response_id in seen_responses:
                    continue
                timestamp = parse_timestamp(record.get("timestamp") or message.get("timestamp"))
                if timestamp is None:
                    continue
                seen_responses.add(response_id)
                count = usage.get("totalTokens")
                if count is None:
                    count = sum(
                        int(usage.get(field, 0) or 0)
                        for field in ("input", "output", "cacheRead", "cacheWrite")
                    )
                events.append((timestamp, int(count or 0)))
    return cumulative_timeline(events), f"OMH artifact usage ({len(files)} transcripts)"


def token_timeline(series: Series) -> tuple[list[tuple[datetime, int]], str]:
    if series.workflow_dir == "kda-humanize":
        return omh_token_timeline(series)
    if series.model_dir in {"claude-opus-4.8", "fable-5"}:
        return claude_token_timeline(series)
    return codex_token_timeline(series)


def tokens_at(timeline: list[tuple[datetime, int]], timestamp: datetime) -> int | None:
    times = [item[0] for item in timeline]
    index = bisect.bisect_right(times, timestamp) - 1
    return timeline[index][1] if index >= 0 else None


def set_plot_style() -> None:
    plt.rcParams.update(
        {
            "figure.dpi": 140,
            "savefig.dpi": 180,
            "font.size": 10,
            "axes.titlesize": 14,
            "axes.labelsize": 11,
            "axes.grid": True,
            "grid.alpha": 0.24,
            "grid.linestyle": ":",
            "legend.fontsize": 8.5,
        }
    )


def plot_time(kernel: str, runs_by_series: dict[Series, list[Run]], out_dir: Path) -> None:
    config = KERNELS[kernel]
    fig, ax = plt.subplots(figsize=(11.8, 6.8))
    for series in sorted(runs_by_series, key=lambda item: item.label):
        runs = runs_by_series[series]
        if not runs:
            continue
        start = runs[0].timestamp
        hours = [(run.timestamp - start).total_seconds() / 3600 for run in runs]
        values = [run.mean_kernel_ms for run in runs]
        best = []
        current = math.inf
        for value in values:
            current = min(current, value)
            best.append(current)

        if len(hours) == 1:
            ax.scatter(hours, best, color=series.color, s=38, label=series.label, zorder=3)
        else:
            ax.plot(
                hours,
                best,
                color=series.color,
                linestyle=series.linestyle,
                linewidth=2.2,
                label=series.label,
            )

    ax.set_yscale("log")
    ax.set_xlabel("Hours since first complete all-PASS official run (per experiment)")
    ax.set_ylabel("Arithmetic mean kernel latency across official suite (ms, log scale)")
    ax.set_title(f"{config['title']}: latency progress over experiment time")
    ax.text(
        0.01,
        0.01,
        "Continuous lines connect observed best-so-far full-suite checkpoints; isolated dots denote one-result experiments",
        transform=ax.transAxes,
        fontsize=8.5,
        color="#555555",
    )
    ax.legend(loc="upper left", bbox_to_anchor=(1.01, 1.0), frameon=False)
    fig.tight_layout()
    save_figure(fig, out_dir / f"{kernel}_latency_vs_time")


def choose_best_run(runs: list[Run]) -> Run:
    clean = [run for run in runs if not run.dirty]
    return min(clean or runs, key=lambda run: run.mean_kernel_ms)


def token_curve(
    runs: list[Run], timeline: list[tuple[datetime, int]]
) -> list[tuple[int, float, float, Run]]:
    mapped = []
    for run in runs:
        token_count = tokens_at(timeline, run.timestamp)
        if token_count is not None and token_count > 0:
            mapped.append((token_count, run.mean_kernel_ms, run))
    mapped.sort(key=lambda item: (item[0], item[2].timestamp))

    # Multiple benchmark runs can occur before another model response updates
    # the session counter. Keep the lowest latency at each cumulative count.
    collapsed: dict[int, tuple[float, Run]] = {}
    for token_count, latency, run in mapped:
        previous = collapsed.get(token_count)
        if previous is None or latency < previous[0]:
            collapsed[token_count] = (latency, run)

    result = []
    best = math.inf
    for token_count, (latency, run) in sorted(collapsed.items()):
        best = min(best, latency)
        result.append((token_count, latency, best, run))
    return result


def plot_tokens(
    kernel: str,
    runs_by_series: dict[Series, list[Run]],
    timelines: dict[Series, tuple[list[tuple[datetime, int]], str]],
    out_dir: Path,
) -> None:
    config = KERNELS[kernel]
    fig, ax = plt.subplots(figsize=(11.8, 6.8))
    for series in sorted(runs_by_series, key=lambda item: item.label):
        runs = runs_by_series[series]
        if not runs:
            continue
        points = token_curve(runs, timelines[series][0])
        if not points:
            continue
        xs = [point[0] / 1_000_000 for point in points]
        ys = [point[2] for point in points]
        if len(xs) == 1:
            ax.scatter(xs, ys, color=series.color, s=38, label=series.label, zorder=3)
        else:
            ax.plot(
                xs,
                ys,
                color=series.color,
                linestyle=series.linestyle,
                linewidth=2.2,
                label=series.label,
            )

    ax.set_xscale("log")
    ax.set_yscale("log")
    ax.set_xlabel("Cumulative agent-session tokens processed (millions, log scale)")
    ax.set_ylabel("Best-so-far mean official-suite kernel latency (ms, log scale)")
    ax.set_title(f"{config['title']}: latency versus cumulative agent tokens")
    ax.text(
        0.01,
        0.01,
        "Lines connect observed monotonic best-so-far checkpoints; usage includes cached input/output and workflow subagents/reviewers",
        transform=ax.transAxes,
        fontsize=8.5,
        color="#555555",
    )
    ax.legend(loc="upper left", bbox_to_anchor=(1.01, 1.0), frameon=False)
    fig.tight_layout()
    save_figure(fig, out_dir / f"{kernel}_latency_vs_tokens")


def save_figure(fig: plt.Figure, stem: Path) -> None:
    fig.savefig(stem.with_suffix(".png"), bbox_inches="tight")
    fig.savefig(stem.with_suffix(".svg"), bbox_inches="tight")
    plt.close(fig)


def write_summary(
    all_series: list[Series],
    runs_by_kernel: dict[str, dict[Series, list[Run]]],
    timelines: dict[Series, tuple[list[tuple[datetime, int]], str]],
    out_dir: Path,
) -> None:
    fields = [
        "kernel",
        "model",
        "workflow",
        "benchmark_csv",
        "valid_full_runs",
        "token_mapped_runs",
        "token_source",
        "tokens_at_first_mapped_run",
        "tokens_at_last_mapped_run",
        "first_valid_utc",
        "last_valid_utc",
        "observed_hours",
        "best_commit",
        "best_git_dirty",
        "best_mean_kernel_ms",
        "best_mean_baseline_ms",
        "best_arithmetic_mean_speedup",
    ]
    series_lookup = {(s.kernel, s.model_dir, s.workflow_dir): s for s in all_series}
    rows = []
    for kernel in KERNELS:
        for model_dir in MODELS:
            workflows = ("goal", "kda-humanize") if model_dir in {"claude-opus-4.8", "fable-5"} else ("goal-arar", "kda-humanize")
            for workflow_dir in workflows:
                series = series_lookup.get((kernel, model_dir, workflow_dir))
                runs = runs_by_kernel[kernel].get(series, []) if series else []
                if runs:
                    best = choose_best_run(runs)
                    curve = token_curve(runs, timelines[series][0])
                    observed_hours = (runs[-1].timestamp - runs[0].timestamp).total_seconds() / 3600
                    rows.append(
                        {
                            "kernel": kernel,
                            "model": MODELS[model_dir][0],
                            "workflow": WORKFLOWS[workflow_dir][0],
                            "benchmark_csv": str(series.path.relative_to(ROOT)),
                            "valid_full_runs": len(runs),
                            "token_mapped_runs": len(curve),
                            "token_source": timelines[series][1],
                            "tokens_at_first_mapped_run": curve[0][0] if curve else "",
                            "tokens_at_last_mapped_run": curve[-1][0] if curve else "",
                            "first_valid_utc": runs[0].timestamp_text,
                            "last_valid_utc": runs[-1].timestamp_text,
                            "observed_hours": f"{observed_hours:.4f}",
                            "best_commit": best.commit,
                            "best_git_dirty": best.dirty,
                            "best_mean_kernel_ms": f"{best.mean_kernel_ms:.9g}",
                            "best_mean_baseline_ms": f"{best.mean_baseline_ms:.9g}",
                            "best_arithmetic_mean_speedup": f"{best.arithmetic_mean_speedup:.9g}",
                        }
                    )
                else:
                    rows.append(
                        {
                            "kernel": kernel,
                            "model": MODELS[model_dir][0],
                            "workflow": WORKFLOWS[workflow_dir][0],
                            "benchmark_csv": str(series.path.relative_to(ROOT)) if series else "missing",
                            "valid_full_runs": 0,
                            "token_mapped_runs": 0,
                            "token_source": timelines[series][1] if series else "missing benchmark.csv",
                            "tokens_at_first_mapped_run": "",
                            "tokens_at_last_mapped_run": "",
                            "first_valid_utc": "",
                            "last_valid_utc": "",
                            "observed_hours": "",
                            "best_commit": "",
                            "best_git_dirty": "",
                            "best_mean_kernel_ms": "",
                            "best_mean_baseline_ms": "",
                            "best_arithmetic_mean_speedup": "",
                        }
                    )
    with (out_dir / "coverage_summary.csv").open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields)
        writer.writeheader()
        writer.writerows(rows)

    token_fields = [
        "kernel",
        "model",
        "workflow",
        "benchmark_timestamp_utc",
        "git_commit",
        "git_dirty",
        "cumulative_agent_tokens",
        "mean_kernel_ms",
        "best_so_far_mean_kernel_ms",
        "token_source",
    ]
    token_rows = []
    for kernel, by_series in runs_by_kernel.items():
        for series, runs in by_series.items():
            source = timelines[series][1]
            for token_count, latency, best, run in token_curve(runs, timelines[series][0]):
                token_rows.append(
                    {
                        "kernel": kernel,
                        "model": series.model,
                        "workflow": series.workflow,
                        "benchmark_timestamp_utc": run.timestamp_text,
                        "git_commit": run.commit,
                        "git_dirty": run.dirty,
                        "cumulative_agent_tokens": token_count,
                        "mean_kernel_ms": f"{latency:.9g}",
                        "best_so_far_mean_kernel_ms": f"{best:.9g}",
                        "token_source": source,
                    }
                )
    token_rows.sort(key=lambda row: (row["kernel"], row["model"], row["workflow"], int(row["cumulative_agent_tokens"])))
    with (out_dir / "token_curve_points.csv").open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=token_fields)
        writer.writeheader()
        writer.writerows(token_rows)

    missing = [row for row in rows if row["valid_full_runs"] == 0]
    with (out_dir / "README.md").open("w") as handle:
        handle.write("# Kernel ablation latency plots\n\n")
        generated_at = datetime.now(timezone.utc).isoformat(timespec="seconds")
        handle.write(f"Snapshot generated at `{generated_at}` from the workspace `benchmark.csv` ledgers.\n\n")
        handle.write("## Comparison rules\n\n")
        handle.write("- Only canonical `official` runs with the full expected workload count and all rows passing are plotted.\n")
        handle.write("- Expected suite sizes: GDN prefill 100, FP8 MoE 19, DSA 23.\n")
        handle.write("- Time starts at each experiment's first valid full-suite result; curves are not extended to 12 hours.\n")
        handle.write("- Time charts connect observed best-so-far mean kernel latency checkpoints with continuous lines.\n")
        handle.write("- Token charts use cumulative agent-session usage from local session logs, not workload sequence length.\n")
        handle.write("- Token curves connect best-so-far checkpoints and are therefore monotonically non-increasing.\n")
        handle.write("- Isolated dots are retained only for experiments with exactly one comparable result.\n")
        handle.write("- Claude goal runs include main/subagent usage; Codex goal runs include all exact-cwd sessions; KDA includes OMH builder/reviewer/judge artifacts.\n")
        handle.write("- Usage is provider-native total processed tokens, including cached input and output.\n")
        handle.write("- Latency axes use milliseconds and logarithmic scaling.\n\n")
        handle.write("## Missing comparable results\n\n")
        if missing:
            for row in missing:
                handle.write(f"- {row['kernel']}: {row['model']} | {row['workflow']} (no complete all-PASS official run)\n")
        else:
            handle.write("None.\n")
        handle.write("\nSee `coverage_summary.csv` for coverage and `token_curve_points.csv` for every plotted token/latency point.\n")
        handle.write("\nRegenerate with `uv run --with matplotlib python plot_experiment_results.py`.\n")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=ROOT / "ablation-plots",
        help="Directory for PNG, SVG, and coverage metadata",
    )
    args = parser.parse_args()
    out_dir = args.output_dir.resolve()
    out_dir.mkdir(parents=True, exist_ok=True)

    set_plot_style()
    all_series = discover_series()
    runs_by_kernel: dict[str, dict[Series, list[Run]]] = {kernel: {} for kernel in KERNELS}
    for series in all_series:
        runs_by_kernel[series.kernel][series] = read_valid_runs(series)

    timelines = {series: token_timeline(series) for series in all_series}

    for kernel in KERNELS:
        plot_time(kernel, runs_by_kernel[kernel], out_dir)
        plot_tokens(kernel, runs_by_kernel[kernel], timelines, out_dir)
    write_summary(all_series, runs_by_kernel, timelines, out_dir)

    print(f"Wrote plots to {out_dir}")
    for kernel in KERNELS:
        count = sum(bool(runs) for runs in runs_by_kernel[kernel].values())
        points = sum(len(runs) for runs in runs_by_kernel[kernel].values())
        print(f"  {kernel}: {count} series, {points} complete all-PASS runs")


if __name__ == "__main__":
    main()