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

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:
            hits.append(("null", 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 format_x_axis(x, pos):
    """Format x-axis values with k notation for thousands"""
    if x >= 1000:
        return f"{x/1000:.0f}k"
    else:
        return f"{x:.0f}"


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", nargs="+", help="Metric name(s) to plot (e.g., cross_entropy_loss acc)")
    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=4, help="Maximum number of x-axis ticks")
    args = ap.parse_args()

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

    # Read all CSVs first so we can determine metrics before creating subplots
    data = []
    for path in args.csvs:
        try:
            df = pd.read_csv(path)
            data.append((path, df))
        except Exception as e:
            print(f"Failed to read {path}: {e}", file=sys.stderr)

    if not data:
        print("No CSVs could be read.", file=sys.stderr)
        sys.exit(1)

    # Determine metrics list
    if args.metric is None:
        auto_metric = pick_first_numeric_metric(data[0][1].columns)
        if auto_metric is None:
            print(f"{data[0][0]}: Could not auto-detect a metric. Use --metric.", file=sys.stderr)
            sys.exit(1)
        metrics = [auto_metric]
    else:
        metrics = args.metric

    nplots = len(metrics)
    fig, axes = plt.subplots(1, nplots, figsize=(6.0 * nplots, 6.0), squeeze=False)
    axes = axes[0]  # flatten single row

    any_plotted = False
    # Collect legend entries across all subplots
    label_to_handle = {}

    for i, metric in enumerate(metrics):
        ax = axes[i]
        for path, df in data:
            metric_hits = find_metric_columns(df.columns, metric)
            if not metric_hits:
                # Silent skip per-file if metric not found; warn once
                print(f"{path}: Metric '{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)

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

                label = run_name
                if run_name == "maskgit":
                    max_y = np.nanmax(y)
                    if np.isfinite(max_y):
                        line = ax.axhline(
                            y=max_y,
                            linestyle="--",
                            linewidth=3.5,
                            label=label,
                            color="#ff7f0e"  # standard matplotlib orange
                        )
                else: 
                    line, = ax.plot(x, y, label=label, linewidth=3.5, marker="o", markersize=7)
                # Record one handle per unique label for the common legend
                if label not in label_to_handle:
                    label_to_handle[label] = line
                any_plotted = True

        # Axis formatting per subplot
        ax.set_title(metric)
        ax.set_xlabel({"step": "Step", "time": "Relative Time (Process)", "index": "Index"}[args.x])
        if args.x_log:
            ax.set_xscale('log')
        if args.y_log:
            ax.set_yscale('log')

        if args.x in ("step", "index"):
            ax.xaxis.set_major_locator(MaxNLocator(nbins=args.max_xticks, integer=True, prune='both'))
            # Apply k-notation formatting for step/index axes
            ax.xaxis.set_major_formatter(FuncFormatter(format_x_axis))
        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 not any_plotted:
        print("Nothing to plot. Check your inputs or --metric.", file=sys.stderr)
        sys.exit(1)

    # Optional overall figure title
    if args.title:
        fig.suptitle(args.title, y=0.995)

    # Common legend above the plots and boxed
    handles = list(label_to_handle.values())
    labels = list(label_to_handle.keys())
    if handles:
        fig.legend(handles, labels, loc="upper center", fontsize=20, ncol=min(4, len(labels)), 
                   frameon=True, fancybox=False, shadow=False)
        plt.tight_layout(rect=[0, 0, 1, 0.90])  # Leave space at top for legend
    else:
        plt.tight_layout()

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


if __name__ == "__main__":
    main()