plot_results.py

#!/usr/bin/env python3
# plot_results.py
import os
import argparse
import pandas as pd
import matplotlib.pyplot as plt

def ensure_dir(p):
    os.makedirs(p, exist_ok=True)
    return p

def load_data(csv_path: str) -> pd.DataFrame:
    if not os.path.exists(csv_path):
        raise SystemExit(f"[error] CSV not found at: {csv_path}")

    # Try normal read
    df = pd.read_csv(csv_path)

    needed = {"model","device","precision","batch","p50_ms","throughput_sps"}
    if not needed.issubset(df.columns):
        # Fallback: assume headerless CSV in the expected order and assign names
        # Expected order:
        # model,device,precision,batch,seq_len,p50_ms,p90_ms,p99_ms,throughput_sps,iters,warmup,torch,python,machine
        try:
            df2 = pd.read_csv(csv_path, header=None)
            expected_cols = [
                "model","device","precision","batch","seq_len",
                "p50_ms","p90_ms","p99_ms","throughput_sps","iters","warmup",
                "torch","python","machine"
            ]
            if df2.shape[1] == len(expected_cols):
                df2.columns = expected_cols
                df = df2
            else:
                print("[error] CSV seems malformed and column count doesn't match expected schema.")
                print(open(csv_path).read())
                raise SystemExit(f"[error] CSV has {df2.shape[1]} columns, expected {len(expected_cols)}.")
        except Exception:
            print("[error] CSV seems malformed. Contents:\n")
            try:
                print(open(csv_path).read())
            except Exception:
                pass
            missing = needed - set(df.columns)
            raise SystemExit(f"[error] CSV is missing columns: {missing}")

    # Normalize types
    df["batch"] = df["batch"].astype(int)
    for col in ["p50_ms","throughput_sps"]:
        df[col] = pd.to_numeric(df[col], errors="coerce")
    if "seq_len" not in df.columns:
        df["seq_len"] = None
    return df


def plot_latency(df: pd.DataFrame, out_dir: str):
    for model in sorted(df["model"].unique()):
        sub = df[df["model"] == model].copy()
        sub = sub.sort_values(by=["device","precision","batch"])
        fig = plt.figure(figsize=(8,5))
        for (device, precision), grp in sub.groupby(["device","precision"]):
            x = grp["batch"].astype(int)
            y = grp["p50_ms"].astype(float)
            plt.plot(x, y, marker="o", label=f"{device}-{precision}")
        plt.title(f"{model} — Latency (P50)")
        plt.xlabel("Batch size"); plt.ylabel("Latency (ms)")
        plt.grid(True, alpha=0.3); plt.legend()
        out_path = os.path.join(out_dir, f"{model}_latency_p50.png")
        plt.tight_layout(); plt.savefig(out_path, dpi=160); plt.close(fig)

def plot_throughput(df: pd.DataFrame, out_dir: str):
    for model in sorted(df["model"].unique()):
        sub = df[df["model"] == model].copy()
        sub = sub.sort_values(by=["device","precision","batch"])
        fig = plt.figure(figsize=(8,5))
        for (device, precision), grp in sub.groupby(["device","precision"]):
            x = grp["batch"].astype(int)
            y = grp["throughput_sps"].astype(float)
            plt.plot(x, y, marker="o", label=f"{device}-{precision}")
        plt.title(f"{model} — Throughput (samples/sec)")
        plt.xlabel("Batch size"); plt.ylabel("Throughput (samples/sec)")
        plt.grid(True, alpha=0.3); plt.legend()
        out_path = os.path.join(out_dir, f"{model}_throughput.png")
        plt.tight_layout(); plt.savefig(out_path, dpi=160); plt.close(fig)

def print_speedup_table(df: pd.DataFrame):
    rows = []
    for model in sorted(df["model"].unique()):
        for batch in sorted(df["batch"].unique()):
            cpu = df[(df.model==model)&(df.device=="cpu")&(df.precision=="fp32")&(df.batch==batch)]
            mps16 = df[(df.model==model)&(df.device=="mps")&(df.precision=="fp16")&(df.batch==batch)]
            mps32 = df[(df.model==model)&(df.device=="mps")&(df.precision=="fp32")&(df.batch==batch)]
            if cpu.empty: continue
            cpu_ms = cpu["p50_ms"].mean()
            if not mps16.empty:
                mps_ms = mps16["p50_ms"].mean(); pair="cpu-fp32 vs mps-fp16"
            elif not mps32.empty:
                mps_ms = mps32["p50_ms"].mean(); pair="cpu-fp32 vs mps-fp32"
            else:
                continue
            if cpu_ms>0 and mps_ms>0:
                rows.append({
                    "model": model, "batch": int(batch),
                    "cpu_p50_ms": round(cpu_ms,2),
                    "mps_p50_ms": round(mps_ms,2),
                    "speedup_x": round(cpu_ms/mps_ms,2),
                    "pair": pair
                })
    if rows:
        table = pd.DataFrame(rows).sort_values(by=["model","batch"])
        print("\n=== CPU → MPS Latency Speedup (P50) ===")
        print(table.to_string(index=False)); print()
    else:
        print("\n[info] Not enough data to compute speedup.\n")

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--csv", default="results/bench.csv")
    ap.add_argument("--out", default="results/plots")
    args = ap.parse_args()

    df = load_data(args.csv)
    os.makedirs(args.out, exist_ok=True)
    print_speedup_table(df)
    plot_latency(df, args.out)
    plot_throughput(df, args.out)
    print(f"[ok] Charts saved to: {args.out}")

if __name__ == "__main__":
    main()