[tab] slurm/utils/mihir/plotting/causal_plot copy.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
from brokenaxes import brokenaxes
import pandas as pd

# =============================================================================
# BROKEN Y-AXIS CONFIGURATION
# =============================================================================
# Set to None to disable broken y-axis, or specify ranges as tuples
# Example: ((0, 0.1), (0.4, 0.7)) creates a break between 0.1 and 0.4
BROKEN_Y_LIMS = None  # Change this to enable broken y-axis
# BROKEN_Y_LIMS = ((0, 0.1), (0.4, 0.7))  # Example configuration

# Space between broken axis segments (adjust as needed)
BROKEN_Y_HSPACE = 0.05


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 = plt.figure(figsize=(6.0 * nplots, 6.0))
    
    any_plotted = False
    # Collect legend entries across all subplots
    label_to_handle = {}

    for i, metric in enumerate(metrics):
        # Create subplot position
        if nplots == 1:
            subplot_pos = 111
        else:
            subplot_pos = 1, nplots, i + 1
        
        # Use brokenaxes if configured, otherwise regular subplot
        if BROKEN_Y_LIMS is not None:
            ax = brokenaxes(
                ylims=BROKEN_Y_LIMS,
                hspace=BROKEN_Y_HSPACE,
                subplot_spec=subplot_pos
            )
        else:
            ax = fig.add_subplot(subplot_pos)
        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)
                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])
        
        # Handle log scales
        if args.x_log:
            ax.set_xscale('log')
        if args.y_log:
            ax.set_yscale('log')

        # Format x-axis ticks
        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'))

        # Apply styling to individual axes if using brokenaxes
        if BROKEN_Y_LIMS is not None:
            # For brokenaxes, we need to style each individual axis
            for individual_ax in ax.axs:
                # show only x and y axes (no surrounding box)
                individual_ax.spines["top"].set_visible(False)
                individual_ax.spines["right"].set_visible(False)
                individual_ax.yaxis.set_ticks_position("left")
                individual_ax.xaxis.set_ticks_position("bottom")
                
                # Apply log scales to individual axes
                if args.x_log:
                    individual_ax.set_xscale('log')
                if args.y_log:
                    individual_ax.set_yscale('log')
                
                # Format x-axis ticks for individual axes
                if args.x in ("step", "index"):
                    individual_ax.xaxis.set_major_locator(MaxNLocator(nbins=args.max_xticks, integer=True, prune='both'))
                    individual_ax.xaxis.set_major_formatter(FuncFormatter(format_x_axis))
                else:
                    individual_ax.xaxis.set_major_locator(MaxNLocator(nbins=args.max_xticks, prune='both'))
        else:
            # Regular axis styling for non-broken axes
            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()