source_current/sudoku/plot_results.py

import os, re, glob
import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt

TASK = "sudoku"
batch_size = 32
BASE_DIR = f"../{TASK}_results_{batch_size}"
PLOT_PATH=os.path.join(BASE_DIR, "figures")
os.makedirs(PLOT_PATH, exist_ok=True)

# METHODS = [
#     "cvxpylayer",
#     "ffoqp_eq",
#     "ffocp_eq",
#     "lpgd"
# ]

METHODS = [
    "lpgd"
]

def load_results(base_dir=BASE_DIR, methods=METHODS):
    dfs = []
    for m in methods:
        pattern = os.path.join(base_dir, m, "*.csv")
        for fp in sorted(glob.glob(pattern)):
            df = pd.read_csv(fp)
            df["method"] = m

            fname = os.path.basename(fp)
            def grab(pat, cast=float):
                mo = re.search(pat, fname)
                return cast(mo.group(1)) if mo else np.nan

            df["seed"] = grab(r"_seed(\d+)", int)
            df["n"] = grab(r"n(\d+)", int)
            df["lr"]   = grab(r"lr([0-9eE\.\-]+)", float)
            #df["eps"]  = grab(r"eps([0-9eE\.\-]+)", float)
            dfs.append(df)
        
        print("method: ", m)
        # print(df)
        
    if not dfs:
        raise FileNotFoundError(f"No CSVs found under {base_dir}.")
    return pd.concat(dfs, ignore_index=True, sort=False)




def plot_metric_curve(df, metric_type, fig_path=PLOT_PATH):
    
    sns.set_theme(style="whitegrid", context="talk")
    assert(metric_type in ["error", "loss"])
    fig_name=f"{metric_type}_curve.png"

    metric_cols = [c for c in [f"test_{metric_type}", f"train_{metric_type}"] if c in df.columns]
    id_vars = [c for c in ["epoch","method","seed","ydim"] if c in df.columns]

    ## collapse metric columns into a value column
    long_curves = df.melt(
        id_vars=id_vars,
        value_vars=metric_cols,
        var_name="metric", value_name="value"
    ).dropna(subset=["value","epoch","method"])

    ## for each metric (col) and each dimensionality (ydim), create a plot of value vs epoch for different methods
    g = sns.relplot(
        data=long_curves, x="epoch", y="value",
        hue="method", style="method",
        markers=True, dashes=True,
        kind="line", ci=None,
        linewidth=1.5, alpha=0.9,
        col="metric", col_wrap=2, height=4, aspect=1.3,
        facet_kws=dict(sharey=False)
    )
    
    # g = sns.relplot(
    #     data=long_curves, x="epoch", y="value",
    #     hue="method", style="method",
    #     markers=True, dashes=True,
    #     kind="line", ci=None,
    #     linewidth=1.5, alpha=0.9,
    #     col="metric", row='ydim',  # remove col_wrap
    #     height=4, aspect=1.3,
    #     facet_kws=dict(sharey=False)
    # )

    g.set_titles("{col_name}")
    g.set_xlabels("Epoch"); g.set_ylabels("Loss")
    plt.suptitle("Loss vs Epoch by Method", y=1.02)

    ## reset legend
    handles, labels = g.axes.flat[0].get_legend_handles_labels()
    if g._legend is not None:
        g._legend.remove()

    g.figure.legend(
        handles, labels,
        loc="lower center",
        bbox_to_anchor=(0.5, -0.10),
        ncol=max(1, len(labels)),
        frameon=False
    )
    g.figure.subplots_adjust(bottom=0.50)

    plt.tight_layout()
    plt.savefig(os.path.join(fig_path, fig_name), dpi=300, bbox_inches="tight")
    
    
    
    
def plot_final_metric(df, metric_type, fig_path=PLOT_PATH):
    keys = ["method","seed"] if "seed" in df.columns else ["method"]
    assert(metric_type in ["loss", "error"])
    fig_name=f"{metric_type}_final.png"

    if "epoch" in df.columns and df["epoch"].notna().any():
        last_rows = df.sort_values("epoch").groupby(keys, dropna=False).tail(1)
    else:
        last_rows = df.groupby(keys, dropna=False).tail(1)
        
    #print(last_rows)

    loss_final = last_rows.melt(
        id_vars=["method","seed","ydim"],
        value_vars=[f"train_{metric_type}",f"test_{metric_type}"],
        var_name="metric", value_name="value"
    ).dropna(subset=["value","method"])
    
    #print(loss_final)

    plt.figure(figsize=(10,5))
    
    ax = sns.barplot(data=loss_final, x="method", y="value", hue="metric", errorbar=("ci",95))
    ax.set_title("Final Loss by Method")
    ax.set_xlabel(""); ax.set_ylabel("Loss")
    ax.tick_params(axis="x", rotation=20)
    
    # g = sns.catplot(
    #     data=loss_final,
    #     x="method", y="value", hue="metric",
    #     col="ydim", kind="bar", ci=95,
    #     height=4, aspect=1.2
    # )
    # g.set_titles("ydim = {col_name}")
    # g.set_axis_labels("", "Loss")
    # g.set_xticklabels(rotation=20)
    
    plt.tight_layout()
    plt.savefig(os.path.join(fig_path, fig_name), dpi=300, bbox_inches="tight")
    
    
    

def plot_forward_backward_time(df, fig_path=PLOT_PATH, fig_name="forward_backward_time.png"):
    keys = ["method","seed"] if "seed" in df.columns else ["method"]
    
    time_median = (
        df.groupby(keys, dropna=False)[["forward_time","backward_time"]]
        .median().reset_index()
    )
    
    # print(time_median)

    time_long = time_median.melt(id_vars=["method","seed"], value_vars=["forward_time","backward_time"],
                             var_name="phase", value_name="seconds").dropna(subset=["seconds","method"])
    
    # print(time_long)

    # method_order_all = [
    #     "cvxpylayer","qpth","ffocp_eq","lpgd"
    # ]
    method_order_all = METHODS
    
    method_order = [m for m in method_order_all if m in time_long["method"].unique()]
    phase_order = ["forward_time", "backward_time"]
    time_long["method"] = pd.Categorical(time_long["method"],
                                        categories=method_order, ordered=True)
    
    # print(time_long)

    plt.figure(figsize=(10,5))
    ax = sns.barplot(
        data=time_long,
        x="method", y="seconds",
        hue="phase",
        order=method_order,
        hue_order=phase_order,
        errorbar=("ci",95)
    )
    
    ax.set_title("Per-Epoch Time by Method")
    ax.set_xlabel(""); ax.set_ylabel("Seconds")
    ax.tick_params(axis="x", rotation=20)
    plt.tight_layout()
    plt.savefig(os.path.join(fig_path, fig_name), dpi=300, bbox_inches="tight")



if __name__=="__main__":
    df = load_results()
    df = df.rename(columns=lambda c: c.strip() if isinstance(c, str) else c)
    df['ydim'] = df['n'] ** 6
    print(df)
    
    plot_metric_curve(df, metric_type="loss")
    plot_metric_curve(df, metric_type="error")
    plot_final_metric(df, metric_type="loss")
    plot_final_metric(df, metric_type="error")
    plot_forward_backward_time(df)