qoco-gpu-benchmarks/plot_benchmarks.py at main · qoco-org/qoco-gpu-benchmarks · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
import os
import pandas as pd
import matplotlib.pyplot as plt

# Enable LaTeX
plt.rcParams["text.usetex"] = True
plt.rcParams["font.family"] = "serif"

COLOR = {
    "qoco": "royalblue",
    "qoco_cuda": "mediumseagreen",
    "cuclarabel": "darkviolet",
    "mosek": "firebrick",
    "gurobi": "coral",
}

PROBLEMS = [
    "portfolio",
    "huber",
    "group_lasso",
    "multiperiod_portfolio",
    "tv_denoising",
]

SOLVERS = {
    "qoco": r"\textsc{QOCO}",
    "qoco_cuda": r"\textsc{QOCO-GPU}",
    "cuclarabel": r"\textsc{CuClarabel}",
    "mosek": r"\textsc{Mosek}",
    "gurobi": r"\textsc{Gurobi}",
}


def plot_problem(ax, prob_name):
    """Plot one benchmark problem on an axis."""

    if not os.path.exists(prob_name):
        print(f"Warning: {prob_name} directory not found")
        return

    for solver_name, display_name in SOLVERS.items():

        csv_path = os.path.join(prob_name, f"{solver_name}_results.csv")

        if not os.path.exists(csv_path):
            print(f"Warning: {csv_path} not found")
            continue

        df = pd.read_csv(csv_path)

        df = df.dropna(subset=["size", "setup_time", "solve_time"])

        if len(df) == 0:
            continue

        df["runtime"] = df["setup_time"] + df["solve_time"]
        df = df[df["runtime"] <= 3600.0]

        df = df.sort_values("size")

        ax.plot(
            df["size"],
            df["runtime"],
            "o-",
            linewidth=2,
            markersize=5,
            label=display_name,
            color=COLOR[solver_name],
        )

    ax.set_xscale("log")
    ax.set_yscale("log")

    # ax.set_xlabel("Problem Size")
    # ax.set_ylabel("Runtime (seconds)")

    ax.set_title(prob_name.replace("_", " ").title(), usetex=True)

    ax.grid(True, alpha=0.3)


def main():

    plt.figure(figsize=(8.5, 11))

    axes = []

    # create subplots
    for i, prob in enumerate(PROBLEMS):
        ax = plt.subplot(3, 2, i + 1)
        axes.append(ax)
        plot_problem(ax, prob)

    # center the fifth plot
    ax5 = axes[4]
    pos = ax5.get_position()

    xright = axes[3].get_position().x0

    ax5.set_position(
        [0.5 * (pos.x0 + xright), pos.y0, pos.width, pos.height,]
    )

    # legend
    handles, labels = plt.gca().get_legend_handles_labels()
    plt.legend(handles, labels, loc="center right", bbox_to_anchor=(1.5, 0.5))

    plt.savefig(
        "figures/benchmark_runtime.pdf", dpi=300, bbox_inches="tight",
    )


if __name__ == "__main__":
    main()