plot_efficiency.py

"""Plot the SparseKAN efficiency figure: accuracy vs cost, three regimes.

Overlays the three efficiency regimes on one accuracy-vs-cost-ratio plot so the
paper's central trade-off is visible at a glance:

  1. Unstructured pruning frontier (near-lossless, best accuracy/cost, but no
     dense-hardware speedup): from ablate.py CSVs
     (importance_prune_ablation.csv, topk_harden_ablation.csv).
  2. Structured 'neuron' pruning + compaction (physically smaller & faster,
     modest accuracy cost): from a small CSV you assemble during the neuron
     keep-ratio sweep (see --structured-csv format below).
  3. Optional quantization annotation: a text callout showing the
     multiplicative bit-cost saving (cost_ratio x bits/32).

Inputs
------
--importance-csv  path to importance_prune_ablation.csv (cols: value,
                  test_metric, cost_ratio, ...)
--topk-csv        path to topk_harden_ablation.csv (same schema)
--structured-csv  optional CSV with columns: label,cost_ratio,accuracy,
                  params_reduction[,latency_ratio]. One row per neuron
                  keep-ratio (after compaction). Example rows:
                      label,cost_ratio,accuracy,params_reduction,latency_ratio
                      neuron0.75,0.75,0.966,0.25,0.85
                      neuron0.5,0.50,0.964,0.50,0.66
                      neuron0.25,0.25,0.951,0.75,0.55
--quant-bits      if set, annotate the joint bit-cost of the best unstructured
                  point at these bits (e.g. 8).
--dense-acc       dense (unpruned) accuracy, drawn as a reference line.
--out             output image path (default sparsekan_efficiency.png).

Usage
-----
python plot_efficiency.py \
    --importance-csv runs_unified/mnist_efficientkan/my_first_run/importance_prune_ablation.csv \
    --topk-csv       runs_unified/mnist_efficientkan/my_first_run/topk_harden_ablation.csv \
    --structured-csv structured_points.csv \
    --quant-bits 8 --dense-acc 0.9750 --out sparsekan_efficiency.png
"""

import argparse
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import pandas as pd  # noqa: E402


def _load(path):
    if path and Path(path).exists():
        return pd.read_csv(path)
    return None


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--importance-csv", type=str, default=None)
    p.add_argument("--topk-csv", type=str, default=None)
    p.add_argument("--structured-csv", type=str, default=None)
    p.add_argument("--quant-bits", type=int, default=None)
    p.add_argument("--dense-acc", type=float, default=None)
    p.add_argument("--acc-percent", action="store_true",
                   help="Multiply accuracies by 100 for the y-axis.")
    p.add_argument("--title", type=str, default="SparseKAN: accuracy vs cost")
    p.add_argument("--out", type=str, default="sparsekan_efficiency.png")
    args = p.parse_args()

    scale = 100.0 if args.acc_percent else 1.0
    ylab = "Test accuracy (%)" if args.acc_percent else "Test accuracy"

    fig, ax = plt.subplots(figsize=(7.5, 5.2))

    # --- Regime 1: unstructured frontier -------------------------------
    imp = _load(args.importance_csv)
    if imp is not None:
        imp = imp.sort_values("cost_ratio")
        ax.plot(imp["cost_ratio"], imp["test_metric"] * scale,
                "o-", color="#2c7fb8", label="Unstructured: importance prune",
                zorder=3)
    topk = _load(args.topk_csv)
    if topk is not None:
        topk = topk.sort_values("cost_ratio")
        ax.plot(topk["cost_ratio"], topk["test_metric"] * scale,
                "s--", color="#41b6c4", label="Unstructured: top-k harden",
                zorder=3)

    # --- Regime 2: structured + compaction -----------------------------
    st = _load(args.structured_csv)
    if st is not None:
        st = st.sort_values("cost_ratio")
        ax.plot(st["cost_ratio"], st["accuracy"] * scale,
                "^-", color="#d95f0e", label="Structured neuron + compaction",
                zorder=4, markersize=9)
        # annotate each structured point with its physical param reduction.
        for _, r in st.iterrows():
            txt = f"-{r['params_reduction']*100:.0f}% params"
            if "latency_ratio" in st.columns and pd.notna(r.get("latency_ratio")):
                txt += f"\n{r['latency_ratio']:.2f}x lat"
            ax.annotate(txt, (r["cost_ratio"], r["accuracy"] * scale),
                        textcoords="offset points", xytext=(8, -18),
                        fontsize=8, color="#d95f0e")

    # --- dense reference line ------------------------------------------
    if args.dense_acc is not None:
        ax.axhline(args.dense_acc * scale, ls=":", color="grey", lw=1,
                   label="Dense (unpruned)")

    # --- quantization annotation ---------------------------------------
    if args.quant_bits and imp is not None:
        # best unstructured point = highest accuracy at lowest cost.
        best = imp.iloc[imp["test_metric"].idxmax()]
        joint = best["cost_ratio"] * args.quant_bits / 32.0
        ax.annotate(
            f"+{args.quant_bits}-bit QAT →\njoint bit-cost {joint:.3f}\n"
            f"(~{1/joint:.0f}x fewer bit-ops)",
            (best["cost_ratio"], best["test_metric"] * scale),
            textcoords="offset points", xytext=(10, 12), fontsize=8,
            color="#2c7fb8",
            arrowprops=dict(arrowstyle="->", color="#2c7fb8", lw=0.8),
        )

    ax.set_xlabel("Active cost ratio (active / dense)")
    ax.set_ylabel(ylab)
    ax.set_title(args.title)
    ax.grid(alpha=0.3)
    ax.invert_xaxis()  # cheaper (more pruned) to the right
    ax.legend(loc="lower left", fontsize=9)
    fig.tight_layout()
    fig.savefig(args.out, dpi=200)
    print(f"[plot] wrote {args.out}")

    # Also emit a compact LaTeX-ready summary table to stdout.
    print("\n% --- efficiency summary (paste-ready rows) ---")
    if imp is not None:
        for _, r in imp.iterrows():
            print(f"Unstructured (importance) & {r['cost_ratio']:.3f} & "
                  f"{r['test_metric']*100:.2f} \\\\")
    if st is not None:
        for _, r in st.iterrows():
            lat = f" & {r['latency_ratio']:.2f}" if "latency_ratio" in st.columns else ""
            print(f"Structured+compact ({r['label']}) & {r['cost_ratio']:.3f} & "
                  f"{r['accuracy']*100:.2f} & -{r['params_reduction']*100:.0f}\\%{lat} \\\\")


if __name__ == "__main__":
    main()