File size: 15,806 Bytes
import os, re, glob
import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import warnings
from matplotlib.patches import Patch
from matplotlib.lines import Line2D
import matplotlib.ticker as mticker


warnings.filterwarnings("ignore")

sns.set_theme(style="whitegrid", context="talk")
palette = sns.color_palette()

batch_size = 8

BASE_DIR = f"../sudoku_results_{batch_size}"
# METHODS = [
#     "cvxpylayer",
#     "lpgd",
#     "bpqp",
#     "ffocp_eq",
#     "dqp",
#     "qpth",
#     "ffoqp_eq",
# ]
# METHODS_LEGEND = {
#     "cvxpylayer": "CvxpyLayer",
#     "lpgd": "LPGD",
#     "bpqp": "BPQP",
#     "ffocp_eq": "FFOCP",
#     "dqp": "dQP",
#     "qpth": "qpth",
#     "ffoqp_eq": "FFOQP",
# }

METHODS = [
    "cvxpylayer",
    "qpth",
    "lpgd",
    "bpqp",
    "dqp",
    "ffocp_eq",
    "ffoqp_eq_schur",
    "ffocp_eq_warm",
    "ffocp_eq_no"
]
METHODS_LEGEND = {
    "cvxpylayer": "CvxpyLayer",
    "lpgd": "LPGD",
    "bpqp": "BPQP",
    "ffocp_eq": "FFOCP",
    "dqp": "dQP",
    "qpth": "qpth",
    "ffoqp_eq_schur": "FFOQP",
    "ffocp_eq_warm": "FFOCP (warm start)",
    "ffocp_eq_no": "FFOCP (no warm start)",
}

METHODS_STEPS = [method+"_steps" for method in METHODS]

method_order = [METHODS_LEGEND[m] for m in METHODS]

markers = ["o", "s", "D", "^", "v", "x", "P", "s", "D"]
# markers_dict = {method: markers[i] for i, method in enumerate(method_order)}

LINEWIDTH = 1.5


def is_subsequence(a, b):
    it = iter(b)
    return all(x in it for x in a)

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)
            m = m.removesuffix("_steps")
            df["method"] = METHODS_LEGEND[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_time_vs_method(df, time_names=['forward_time', 'backward_time'], plot_path=BASE_DIR, plot_name_tag=""):
    df_avg_method = df.groupby('method')[time_names].mean().reset_index()

    # Convert wide → long format so Seaborn can handle grouped bars
    df_long = df_avg_method.melt(id_vars='method', 
                                value_vars=time_names,
                                var_name='Metrics', 
                                value_name='Time')

    plt.figure(figsize=(8,5))
    sns.barplot(data=df_long, x='method', y='Time', hue='Metrics')
    plt.ylabel("Time")
    plt.title("Forward and Backward Time")
    plt.savefig(f"{plot_path}/{plot_name_tag}_time_vs_method.pdf", dpi=300, bbox_inches='tight')
    plt.close()

def plot_time_vs_epoch(df, time_names=['forward_time', 'backward_time'], iteration_name='epoch', plot_path=BASE_DIR, plot_name_tag=""):
    df_avg_epoch = df.groupby(['method', iteration_name])[time_names].mean().reset_index()

    # --- Forward Time Figure ---
    plt.figure(figsize=(8,5))
    sns.lineplot(data=df_avg_epoch, x=iteration_name, y=time_names[0], hue='method', marker=None, dashes=False, linewidth=LINEWIDTH)
    plt.ylabel("Forward Time")
    plt.title(f"Forward Time vs {iteration_name}")
    
    plt.savefig(f"{plot_path}/{plot_name_tag}_forward_time_vs_{iteration_name}.pdf", dpi=300, bbox_inches='tight')
    plt.close()

    # --- Backward Time Figure ---
    plt.figure(figsize=(8,5))
    sns.lineplot(data=df_avg_epoch, x=iteration_name, y=time_names[1], hue='method', marker=None, dashes=False, linewidth=LINEWIDTH)
    plt.ylabel("Backward Time")
    plt.title(f"Backward Time vs {iteration_name}")
    
    plt.savefig(f"{plot_path}/{plot_name_tag}_backward_time_vs_{iteration_name}.pdf", dpi=300, bbox_inches='tight')
    plt.close()

# def plot_total_time_vs_method(df, time_names=['forward_time', 'backward_time'], plot_path=BASE_DIR, plot_name_tag=""):
#     # Group by method, average over epochs and seeds
#     df_avg_method = df.groupby('method')[time_names].mean().reset_index()

#     # --- Stacked Bar Chart ---
#     methods = df_avg_method['method']
#     forward = df_avg_method[time_names[0]]
#     backward = df_avg_method[time_names[1]]

#     plt.figure(figsize=(8,5))
#     plt.bar(methods, forward, label=time_names[0], color=palette[0])
#     plt.bar(methods, backward, bottom=forward, label=time_names[1], color=palette[1])
#     plt.ylabel("Time")
#     plt.title("Total Time vs Method")
#     plt.legend()
#     plt.savefig(f"{plot_path}/{plot_name_tag}_total_time_vs_method.pdf", dpi=300, bbox_inches='tight')
#     plt.close()

# def plot_total_time_vs_method(
#     df,
#     time_names=["forward_time", "backward_time"],
#     plot_path=BASE_DIR,
#     plot_name_tag="",
# ):
#     df_avg_method = df.groupby("method")[time_names].mean().reindex(method_order)

#     methods = df_avg_method.index.tolist()
#     forward = df_avg_method[time_names[0]].to_numpy()
#     backward = df_avg_method[time_names[1]].to_numpy()

#     total_times = forward + backward
#     max_time = np.nanmax(total_times)
#     min_time = 0
#     y_max = max_time * 1.05

#     dashed_methods = {"CvxpyLayer", "LPGD", "BPQP", "FFOCP"}  # CP
#     x = np.arange(len(methods))
#     width = 0.75

#     fig, ax = plt.subplots(figsize=(10, 5))
#     sf = mticker.ScalarFormatter(useOffset=False)
#     sf.set_scientific(False)
#     ax.yaxis.set_major_formatter(sf)
#     ax.yaxis.get_offset_text().set_visible(False)

#     for i, m in enumerate(methods):
#         is_dashed = (m in dashed_methods)
#         ls = (0, (4, 2)) if is_dashed else "solid"
#         lw = 2.0 if is_dashed else 1.5

#         ax.bar(x[i], forward[i], width=width,
#                color=palette[0], edgecolor="black", linewidth=lw, linestyle=ls)
#         ax.bar(x[i], backward[i], bottom=forward[i], width=width,
#                color=palette[1], edgecolor="black", linewidth=lw, linestyle=ls)

#     ax.set_xticks(x)
#     ax.set_xticklabels(methods)
#     ax.set_ylabel("Time")
#     ax.set_title("Total Time vs Methods")
#     ax.set_ylim(min_time, y_max)

#     handles = [
#         Line2D([], [], linestyle="none", label="Comp. Time"),
#         Patch(facecolor=palette[0], edgecolor="black", label="Forward"),
#         Patch(facecolor=palette[1], edgecolor="black", label="Backward"),
#         Line2D([], [], linestyle="none", label=""),
#         Line2D([], [], linestyle="none", label="Solvers"),
#         Line2D([0], [0], color="black", linewidth=2, linestyle=(0, (4, 2)), label="CP methods"),
#         Line2D([0], [0], color="black", linewidth=2, linestyle="solid", label="QP methods"),
#     ]

#     leg = ax.legend(
#         handles=handles,
#         loc="upper right",
#         frameon=True,
#         handlelength=2.2,
#         handletextpad=0.8,
#         borderpad=0.7,
#         labelspacing=0.1,
#     )
#     leg._legend_box.align = "left"

#     for h, t in zip(leg.legend_handles, leg.get_texts()):
#         txt = t.get_text()
#         if txt in ("Comp. Time", "Solvers"):
#             t.set_weight("bold")
#             if hasattr(h, "set_visible"):
#                 h.set_visible(False)
#         if txt == "":
#             if hasattr(h, "set_visible"):
#                 h.set_visible(False)
#             t.set_color((0, 0, 0, 0))

#     fig.savefig(f"{plot_path}/{plot_name_tag}_total_time_vs_method.pdf", dpi=300, bbox_inches="tight")
#     plt.close(fig)

def plot_total_time_vs_method(
    df,
    time_names=["forward_time", "backward_time"],
    plot_path=BASE_DIR,
    plot_name_tag="",
    task_title_tag="Sudoku"
):
    available_methods = set(df["method"].unique())
    # print("hhhhhhh\n", df[df['method']=='FFOQP'])
    # assert("FFOQP" in available_methods)
    filtered_method_order = [m for m in method_order if m in available_methods]
    print("filtered_method_order:", filtered_method_order)
    print("available method: ", available_methods)

    df_avg_method = df.groupby("method")[time_names].mean().reindex(filtered_method_order)

    methods = df_avg_method.index.tolist()
    forward = df_avg_method[time_names[0]].to_numpy()
    backward = df_avg_method[time_names[1]].to_numpy()

    total_times = forward + backward
    max_time = np.nanmax(total_times)
    min_time = 0
    y_max = max_time * 1.05
    
    #### get groups of method to plot them close together
    groups = [
        ["CvxpyLayer", "qpth"], # KKT based
        ["LPGD","BPQP", "dQP"], # optimization based
        ["FFOCP", "FFOQP", "FFOCP (warm start)", "FFOCP (no warm start)"], # our methods
    ]
    
    # assert groups respect the filtered method order
    flat_groups = [m for g in groups for m in g]
   
    groups = [[m for m in g if m in filtered_method_order] for g in groups]
    groups = [g for g in groups if len(g) > 0]
    
    flat_groups = [m for g in groups for m in g]
    print("groups:", flat_groups)
    print("methods:", methods)
    assert is_subsequence(flat_groups, methods), "Method groups do not respect the method order."

    # dashed_methods = {"CvxpyLayer", "LPGD", "BPQP", "FFOCP"}  # CP
    
    ######### set bar positions with custom gaps
    width = 0.08 #0.75 # bar width
    inner_gap = 0.08 #1.5     # spacing between bars inside a group
    group_gap = len(methods)*0.125/7   ##so that when num methods is 7, group gap is 3.0   # extra spacing between groups

    x = []
    # labels = []
    pos = 0.0

    for g in groups:
        for m in g:
            x.append(pos)
            # labels.append(m)
            pos += inner_gap
        pos += group_gap  # extra space after each group
        
    ########## set bar label positions with custom gaps
    # inner_gap = 1.5     # spacing between bars inside a group
    # group_gap = len(methods)*2/7   ##so that when num methods is 7, group gap is 3.0   # extra spacing between groups
    # label_x = []
    # pos = 0.0
    # for g in groups:
    #     for m in g:
    #         label_x.append(pos)
    #         pos += inner_gap
    #     pos += group_gap  # extra space after each group
    label_x = x

    
    # x = np.arange(len(methods))
    

    fig, ax = plt.subplots(figsize=(14, 7))
    sf = mticker.ScalarFormatter(useOffset=False)
    sf.set_scientific(False)
    ax.yaxis.set_major_formatter(sf)
    ax.yaxis.get_offset_text().set_visible(False)

    for i, m in enumerate(methods):
        # is_dashed = (m in dashed_methods)
        # ls = (0, (4, 2)) if is_dashed else "solid"
        # lw = 2.0 if is_dashed else 1.5

        # ax.bar(x[i], forward[i], width=width,
        #        color=palette[0], edgecolor="black", linewidth=lw, linestyle=ls)
        # ax.bar(x[i], backward[i], bottom=forward[i], width=width,
        #        color=palette[1], edgecolor="black", linewidth=lw, linestyle=ls)
        
        ax.bar(x[i], forward[i], width=width,
               color=palette[0], edgecolor="black")
        ax.bar(x[i], backward[i], bottom=forward[i], width=width,
               color=palette[1], edgecolor="black")


    ax.set_xticks(label_x)
    methods = ["Cvxpy\nLayer" if m == "CvxpyLayer" else m for m in methods]
    ax.set_xticklabels(methods)
    ax.set_ylabel("Time")
    ax.set_title(f"{task_title_tag}: Total Time vs Methods")
    ax.set_ylim(min_time, y_max)
    min_x = min([min(x), min(label_x)])
    max_x = max([max(x), max(label_x)])
    ax.set_xlim(min_x - width, max_x + width)

    handles = [
        Line2D([], [], linestyle="none", label="Comp. Time"),
        Patch(facecolor=palette[0], edgecolor="black", label="Forward"),
        Patch(facecolor=palette[1], edgecolor="black", label="Backward"),
        # Line2D([], [], linestyle="none", label=""),
        # Line2D([], [], linestyle="none", label="Solvers"),
        # Line2D([0], [0], color="black", linewidth=2, linestyle=(0, (4, 2)), label="CP methods"),
        # Line2D([0], [0], color="black", linewidth=2, linestyle="solid", label="QP methods"),
    ]

    leg = ax.legend(
        handles=handles,
        loc="upper right",
        frameon=True,
        handlelength=2.2,
        handletextpad=0.8,
        borderpad=0.7,
        labelspacing=0.1,
    )
    leg._legend_box.align = "left"

    for h, t in zip(leg.legend_handles, leg.get_texts()):
        txt = t.get_text()
        if txt in ("Comp. Time", "Solvers"):
            t.set_weight("bold")
            if hasattr(h, "set_visible"):
                h.set_visible(False)
        if txt == "":
            if hasattr(h, "set_visible"):
                h.set_visible(False)
            t.set_color((0, 0, 0, 0))

    fig.savefig(f"{plot_path}/{plot_name_tag}_total_time_vs_method.pdf", dpi=300, bbox_inches="tight")
    plt.close(fig)

def plot_losse_vs_epoch(df, loss_metric_name, iteration_name='epoch', plot_path=BASE_DIR, plot_name_tag="", loss_range=None, stride=50,
    task_title_tag="Sudoku"):
    df_avg_epoch = df.groupby(['method', iteration_name], observed=True)[[loss_metric_name]].mean().reset_index()
    # df_avg_epoch = df_avg_epoch[df_avg_epoch[iteration_name] % stride == 0]

    observed_methods = df_avg_epoch['method'].dropna().unique().tolist()
    
    # --- Forward Time Figure ---
    plt.figure(figsize=(8,5))
    ax = sns.lineplot(data=df_avg_epoch, x=iteration_name, y=loss_metric_name, hue='method', dashes=False, linewidth=LINEWIDTH, hue_order=observed_methods)
    plt.ylabel("loss")
    plt.title(f"{task_title_tag}: Loss vs iteration")
    
    # plt.legend(
    #     title=None,
    #     loc="lower center",
    #     bbox_to_anchor=(0.5, -0.25),
    #     ncol=(df["method"].nunique()),           
    #     frameon=False
    # )
    
    if loss_range is not None:
        ax.set_ylim(loss_range)
    
    plt.savefig(f"{plot_path}/{plot_name_tag}_{loss_metric_name}_vs_{iteration_name}.pdf", dpi=300, bbox_inches='tight')
    plt.close()
        

if __name__=="__main__":
    df = load_results()
    df = df.rename(columns=lambda c: c.strip() if isinstance(c, str) else c)
    df["method"] = pd.Categorical(df["method"], categories=method_order, ordered=True)
    
    # plot_time_vs_method(df, time_names=['forward_time', 'backward_time'], plot_path=BASE_DIR, plot_name_tag="sudoku")
    # plot_time_vs_epoch(df, time_names=['forward_time', 'backward_time'], iteration_name='epoch', plot_path=BASE_DIR)
    plot_total_time_vs_method(df, time_names=['forward_time', 'backward_time'], plot_path=BASE_DIR, plot_name_tag="sudoku")
    
    
    df = load_results(methods=METHODS_STEPS)
    df = df.rename(columns=lambda c: c.strip() if isinstance(c, str) else c)
    df["method"] = pd.Categorical(df["method"], categories=method_order, ordered=True)

    # plot_time_vs_method(df, time_names=['iter_forward_time', 'iter_backward_time'], plot_path=BASE_DIR, plot_name_tag="sudoku_steps")
    # plot_time_vs_epoch(df, time_names=['iter_forward_time', 'iter_backward_time'], iteration_name='iter', plot_path=BASE_DIR, plot_name_tag="sudoku_steps")
    # plot_total_time_vs_method(df, time_names=['iter_forward_time', 'iter_backward_time'], plot_path=BASE_DIR, plot_name_tag="sudoku_steps")

    plot_losse_vs_epoch(df, "train_loss", iteration_name='iter', plot_path=BASE_DIR, plot_name_tag="sudoku_steps", loss_range=(0, 2.4))
    # plot_losse_vs_epoch(df, "train_error", iteration_name='iter', plot_path=BASE_DIR, plot_name_tag="sudoku_steps")