plot.py

import csv
import argparse
import matplotlib.pyplot as plt
from pathlib import Path
import json

def main():
    p = argparse.ArgumentParser()
    p.add_argument("--results", "-r", default="results/summary.csv",
                   help="Path to summary.csv generated earlier")
    p.add_argument("--out", "-o", default="results",
                   help="Directory to write new plots")
    args = p.parse_args()

    csv_path = Path(args.results)
    outdir = Path(args.out)
    outdir.mkdir(parents=True, exist_ok=True)

    # Load CSV
    rows = []
    with csv_path.open() as f:
        reader = csv.DictReader(f)
        for row in reader:
            if row["test_accuracy_percent"] and row["test_accuracy_percent"] != "None":
                row["test_accuracy_percent"] = float(row["test_accuracy_percent"])
                row["workers"] = int(row["workers"])
                rows.append(row)

    if not rows:
        print("No accuracy data found.")
        return

    # Group: binary -> workers -> values
    data = {}
    for row in rows:
        binname = row["binary"]
        w = row["workers"]
        acc = row["test_accuracy_percent"]
        data.setdefault(binname, {}).setdefault(w, []).append(acc)

    # Find global min & max
    all_acc = [acc for binr in data.values() for lst in binr.values() for acc in lst]
    min_acc = min(all_acc)
    max_acc = max(all_acc)

    print(f"Global min accuracy: {min_acc:.4f}%")
    print(f"Global max accuracy: {max_acc:.4f}%")

    # Plot
    plt.figure(figsize=(10,6))
    binaries = sorted(data.keys())
    all_workers = sorted({w for binr in data.values() for w in binr.keys()})
    bar_width = 0.8 / max(1, len(binaries))
    x = list(range(len(all_workers)))

    for i, binname in enumerate(binaries):
        means = []
        stds = []
        for w in all_workers:
            vals = data[binname].get(w, [])
            if vals:
                mean = sum(vals) / len(vals)
                std = (sum((v-mean)**2 for v in vals)/len(vals))**0.5 if len(vals)>1 else 0
            else:
                mean = min_acc
                std = 0
            means.append(mean)
            stds.append(std)

        positions = [xi + i*bar_width for xi in x]
        plt.bar(positions, means, width=bar_width, label=binname,
                yerr=stds, capsize=3)

    plt.xticks([xi + (len(binaries)-1)*bar_width/2 for xi in x], all_workers)
    plt.ylabel("Test Accuracy (%)")
    plt.xlabel("MPI ranks (workers)")
    plt.title("Test Accuracy (zoomed to min)")

    # zoom the y-axis so min value is the lower bound
    margin = (max_acc - min_acc) * 0.2 if max_acc > min_acc else 0.1
    plt.ylim(min_acc - margin*0.2, max_acc + margin)

    plt.legend()
    plt.grid(axis='y')

    outpath = outdir / "test_accuracy_zoomed.png"
    plt.savefig(outpath, bbox_inches="tight")
    plt.close()

    print("Saved:", outpath)

    # ============================
    # Per-worker loss comparison
    # ============================

    # Reload CSV to capture losses
    loss_data = {}
    with csv_path.open() as f:
        reader = csv.DictReader(f)
        for row in reader:
            if not row["losses"] or row["losses"] == "[]":
                continue
            binname = row["binary"]
            w = int(row["workers"])
            try:
                losses = json.loads(row["losses"])
            except:
                print("Could not parse. Make sure python environment has json module.")
                continue
            loss_data.setdefault(w, {}).setdefault(binname, []).append(losses)

    # Generate plots: for each worker count, compare algos
    for w, binmap in sorted(loss_data.items()):
        plt.figure(figsize=(10,6))
        any_curve = False

        for binname, loss_lists in binmap.items():

            # compute average loss across repeats
            max_epochs = max(len(l) for l in loss_lists)
            avg_losses = []
            for e in range(max_epochs):
                vals = [l[e] for l in loss_lists if e < len(l)]
                if vals:
                    avg_losses.append(sum(vals)/len(vals))
                else:
                    avg_losses.append(float('nan'))

            xs = list(range(len(avg_losses)))
            plt.plot(xs, avg_losses, marker='o', label=binname)
            any_curve = True

        if not any_curve:
            continue

        plt.xlabel("Epoch")
        plt.ylabel("Loss (log)")
        plt.title(f"Loss vs Epoch — workers={w}")
        plt.yscale("log")
        plt.legend()
        plt.grid(True)

        outpath_loss = outdir / f"loss_vs_epoch_workers_{w}.png"
        plt.savefig(outpath_loss, bbox_inches="tight")
        plt.close()
        print("Saved:", outpath_loss)

    

if __name__ == "__main__":
    main()