[tab] slurm/utils/mihir/wandb_csv_to_matplotlib_minimal.py (python)
#!/usr/bin/env python3
import argparse
import os
import re
import sys
from matplotlib.ticker import MaxNLocator

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


def find_metric_columns(cols, metric):
    # Matches "<run> - metric" and ignores "__MIN"/"__MAX"
    pat = re.compile(rf"^(.*?)\s*-\s*{re.escape(metric)}$")
    hits = []
    for c in cols:
        if c.endswith("__MIN") or c.endswith("__MAX"):
            continue
        m = pat.match(c)
        if m:
            run = m.group(1).strip()
            hits.append((run, c))
        if c.lower() == metric.lower():
            hits.append(("jafar", c))
    return hits  # list of (run_name, col_name)


def find_step_column(cols, run_name):
    # Prefer "<run> - _step", fallback to a generic "_step" if present
    if "step" in cols:
        return "step"
    exact = f"{run_name} - _step"
    if exact in cols:
        return exact
    if "_step" in cols:
        return "_step"
    # Sometimes exported as "<run> - step" (rare)
    alt = f"{run_name} - step"
    if alt in cols:
        return alt
    return None


def pick_first_numeric_metric(cols):
    # Choose a numeric-looking metric column that's not step/min/max/time
    skip_suffixes = {"__MIN", "__MAX"}
    skip_exact = {
        "_step",
        "global_step",
    }
    for c in cols:
        if c in skip_exact:
            continue
        if any(c.endswith(suf) for suf in skip_suffixes):
            continue
        # Heuristic: columns with " - _step" are steps, skip
        if c.endswith(" - _step") or c.endswith(" - step"):
            continue
        # Likely a metric column of the form "<run> - name"
        if " - " in c:
            return c.split(" - ", 1)[1]
    return None


def roll_mean(x, w):
    if w <= 1:
        return x
    # NaN-safe rolling via convolution over valid entries
    x = np.asarray(x, dtype=float)
    mask = ~np.isnan(x)
    num = np.convolve(np.nan_to_num(x), np.ones(w), mode="same")
    den = np.convolve(mask.astype(float), np.ones(w), mode="same")
    with np.errstate(invalid="ignore", divide="ignore"):
        y = num / den
    y[den == 0] = np.nan
    return y


def main():
    ap = argparse.ArgumentParser(description="Plot W&B CSV runs in Matplotlib.")
    ap.add_argument("csvs", nargs="+", help="One or more W&B CSV files")
    ap.add_argument("--metric", help="Metric name to plot (e.g., cross_entropy_loss)")
    ap.add_argument("--x", choices=["step", "time", "index"], default="step", help="X-axis")
    ap.add_argument("--smooth", type=int, default=1, help="Rolling window size")
    ap.add_argument("--title", default=None, help="Figure title")
    ap.add_argument("--save", default=None, help="Path to save the figure (e.g., out.png). If omitted, shows the plot.")
    ap.add_argument("--y_log", action="store_true", help="Enable log scale for the y axis.")
    ap.add_argument("--x_log", action="store_true", help="Enable log scale for the x axis.")
    ap.add_argument("--max_xticks", type=int, default=8, help="Maximum number of x-axis ticks")
    args = ap.parse_args()

    plt.rcParams.update({'font.size': 20})

    plt.figure(figsize=(12, 7))

    any_plotted = False
    chosen_metric = args.metric

    for path in args.csvs:
        try:
            df = pd.read_csv(path)
        except Exception as e:
            print(f"Failed to read {path}: {e}", file=sys.stderr)
            continue

        if chosen_metric is None:
            chosen_metric = pick_first_numeric_metric(df.columns)
            if chosen_metric is None:
                print(f"{path}: Could not auto-detect a metric. Use --metric.", file=sys.stderr)
                continue

        metric_hits = find_metric_columns(df.columns, chosen_metric)
        if not metric_hits:
            print(f"{path}: Metric '{chosen_metric}' not found.", file=sys.stderr)
            continue

        # Plot each run found in this CSV for the chosen metric
        for run_name, metric_col in metric_hits:
            if args.x == "time":
                x = df.get("Relative Time (Process)")
                if x is None:
                    print(f"{path}: 'Relative Time (Process)' not found; using index.", file=sys.stderr)
                    x = pd.RangeIndex(len(df))
            elif args.x == "step":
                step_col = find_step_column(df.columns, run_name)
                x = df.get(step_col) if step_col else None
                if x is None:
                    print(f"{path}: step column for run '{run_name}' not found; using index.", file=sys.stderr)
                    x = pd.RangeIndex(len(df))
            else:
                x = pd.RangeIndex(len(df))

            y = pd.to_numeric(df[metric_col], errors="coerce").values
            x = np.asarray(x, dtype=float) / 3600

            if args.smooth and args.smooth > 1:
                y = roll_mean(y, args.smooth)

            # label = f"{os.path.splitext(os.path.basename(path))[0]} :: {run_name}"
            label = run_name
            plt.plot(x, y, label=label)
            any_plotted = True

    if not any_plotted:
        print("Nothing to plot. Check your inputs or --metric.", file=sys.stderr)
        sys.exit(1)

    plt.xlabel({"step": "Step", "time": "Wall clock time (h)", "index": "Index"}[args.x])
    plt.title(chosen_metric)

    if args.x_log:
        plt.xscale('log')
    if args.y_log:
        plt.yscale('log')

    ax = plt.gca()
    if args.x in ("step", "index"):
        ax.xaxis.set_major_locator(MaxNLocator(nbins=args.max_xticks, integer=True, prune='both'))
    else:
        ax.xaxis.set_major_locator(MaxNLocator(nbins=args.max_xticks, prune='both'))
    # show only x and y axes (no surrounding box)
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.yaxis.set_ticks_position("left")
    ax.xaxis.set_ticks_position("bottom")
    # if args.title:
        # plt.title(args.title)
    plt.legend(loc="best", fontsize=18)
    plt.tight_layout()

    if args.save:
        plt.savefig(args.save, dpi=200)
    else:
        plt.show()


if __name__ == "__main__":
    main()