plot_final.py

import os
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import pandas as pd
import numpy as np

os.makedirs("plots", exist_ok=True)

df = pd.read_csv("results/combined_results.csv")
baseline = df[df["method"] == "baseline"].copy()
spec = df[df["method"] == "speculative"].copy()

# ── Best K summary ────────────────────────────────────────────────────────
best_rows = []
for prompt in sorted(spec["prompt"].unique()):
    sub = spec[spec["prompt"] == prompt]
    best = sub.loc[sub["speedup_vs_baseline"].idxmax()]
    base = baseline[baseline["prompt"] == prompt].iloc[0]
    best_rows.append({
        "prompt": prompt,
        "prompt_len": int(base["prompt_len"]),
        "baseline_tput": base["throughput_tok_s"],
        "max_new_tokens": int(base["max_new_tokens"]),
        "best_k": int(best["k"]),
        "best_tput": best["throughput_tok_s"],
        "best_speedup": best["speedup_vs_baseline"],
        "best_accept": best["mean_accept_rate"],
        "achieved_speedup": best["speedup_vs_baseline"] > 1.0,
    })

best_df = pd.DataFrame(best_rows).sort_values("best_speedup", ascending=False)
best_df.to_csv("results/best_k_summary.csv", index=False)
print("=== Best K per prompt ===")
print(best_df.to_string(index=False))

# ── Plot 1: speedup vs K ─────────────────────────────────────────────────
fig, ax = plt.subplots(figsize=(12, 7))
for prompt in sorted(spec["prompt"].unique()):
    sub = spec[spec["prompt"] == prompt].sort_values("k")
    ax.plot(sub["k"], sub["speedup_vs_baseline"], marker="o", linewidth=2, label=prompt)
ax.axhline(1.0, color="gray", linestyle="--", linewidth=1.5, label="baseline")
ax.set_xlabel("Draft tokens K", fontsize=12)
ax.set_ylabel("Speedup vs baseline (×)", fontsize=12)
ax.set_title("Speculative Decoding Speedup vs K\n(12 prompt types)", fontsize=13)
ax.legend(fontsize=8, ncol=2)
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig("plots/speedup_vs_k_all.png", dpi=180, bbox_inches="tight")
plt.close()
print("plot: plots/speedup_vs_k_all.png")

# ── Plot 2: acceptance vs speedup scatter ─────────────────────────────────
fig, ax = plt.subplots(figsize=(12, 7))
colors = plt.cm.tab20(np.linspace(0, 1, len(spec["prompt"].unique())))
for (prompt, sub), color in zip(spec.groupby("prompt"), colors):
    ax.scatter(sub["mean_accept_rate"], sub["speedup_vs_baseline"],
               c=[color], s=80, label=prompt)
    for _, r in sub.iterrows():
        ax.annotate(f"K={int(r['k'])}", (r["mean_accept_rate"], r["speedup_vs_baseline"]),
                    textcoords="offset points", xytext=(4, 4), fontsize=6, alpha=0.7)
ax.axhline(1.0, color="gray", linestyle="--", linewidth=1)
ax.axvline(0.80, color="red", linestyle=":", alpha=0.5, label="~80% threshold")
ax.set_xlabel("Mean acceptance rate", fontsize=12)
ax.set_ylabel("Speedup vs baseline (×)", fontsize=12)
ax.set_title("Acceptance Rate vs Speedup\n(threshold ~80% for positive speedup)", fontsize=13)
ax.legend(fontsize=7, ncol=2)
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig("plots/acceptance_vs_speedup_all.png", dpi=180, bbox_inches="tight")
plt.close()
print("plot: plots/acceptance_vs_speedup_all.png")

# ── Plot 3: best speedup per prompt bar chart ─────────────────────────────
fig, ax = plt.subplots(figsize=(14, 6))
colors = ["#2ECC71" if s > 1 else "#E74C3C" for s in best_df["best_speedup"]]
bars = ax.barh(best_df["prompt"], best_df["best_speedup"], color=colors)
ax.axvline(1.0, color="gray", linestyle="--", linewidth=1.5)
for i, (_, r) in enumerate(best_df.iterrows()):
    ax.text(
        r["best_speedup"] + 0.02,
        i,
        f"K={int(r['best_k'])} acc={r['best_accept']:.2f}",
        va="center",
        fontsize=9,
    )
ax.set_xlabel("Best speedup vs baseline (×)", fontsize=12)
ax.set_title("Best Speculative Decoding Speedup per Prompt\n(green > 1×, red < 1×)", fontsize=13)
ax.grid(True, alpha=0.3, axis="x")
patches = [
    mpatches.Patch(color="#2ECC71", label="Speedup > 1×"),
    mpatches.Patch(color="#E74C3C", label="Speedup < 1×"),
]
ax.legend(handles=patches)
plt.tight_layout()
plt.savefig("plots/best_speedup_per_prompt.png", dpi=180, bbox_inches="tight")
plt.close()
print("plot: plots/best_speedup_per_prompt.png")

# ── Plot 4: throughput vs K by category ───────────────────────────────────
fig, axes = plt.subplots(2, 2, figsize=(16, 12))
axes = axes.flatten()

categories = {
    "High agreement": ["repetitive_pattern", "list_items", "wikipedia_style"],
    "Moderate agreement": ["technical_ml", "code_json", "dialogue"],
    "Code": ["code_python", "code_json"],
    "Creative": ["creative_fiction", "dialogue"],
}

for ax, (cat_name, prompts) in zip(axes, categories.items()):
    for prompt in prompts:
        sub_s = spec[spec["prompt"] == prompt].sort_values("k")
        if len(sub_s) == 0:
            continue
        base_tput = baseline[baseline["prompt"] == prompt]["throughput_tok_s"].iloc[0]
        ax.plot(sub_s["k"], sub_s["throughput_tok_s"], marker="o", linewidth=2, label=prompt)
        ax.axhline(base_tput, linestyle="--", alpha=0.3)
    ax.set_title(cat_name)
    ax.set_xlabel("Draft tokens K")
    ax.set_ylabel("Throughput (tok/s)")
    ax.grid(True, alpha=0.3)
    ax.legend(fontsize=8)

plt.suptitle("Throughput vs K by prompt category", y=1.02, fontsize=14)
plt.tight_layout()
plt.savefig("plots/throughput_by_category.png", dpi=180, bbox_inches="tight")
plt.close()
print("plot: plots/throughput_by_category.png")

# ── Summary stats ─────────────────────────────────────────────────────────
print(f"\n=== Overall stats ===")
print(f"Prompts with speedup > 1×: {best_df['achieved_speedup'].sum()} / {len(best_df)}")
print(f"Best overall speedup: {best_df['best_speedup'].max():.3f}× ({best_df.iloc[0]['prompt']})")
print(f"Mean best speedup: {best_df['best_speedup'].mean():.3f}×")
print(f"Mean acceptance at best K: {best_df['best_accept'].mean():.3f}")