plots.py

from typing import Dict
import numpy as np
import matplotlib
matplotlib.use("TkAgg") 
import matplotlib.pyplot as plt
from matplotlib.ticker import PercentFormatter

plt.rcParams.update(
    {
        "text.usetex": True,
        "font.family": "serif",
        "axes.titlesize": 18,
        "legend.fontsize": 14,
        "lines.linewidth": 2,
    }
)
colors = {
    "ADMMslack": "b",
    "StandardSlack": "r",
}
mode_to_name = {
    "ADMMslack": "ADMMSlack",
    "StandardSlack": "ADMM",
}

def plot_cdf_sparse_timing(rows, modes_order=None, logx=False, title="Sparse timing CDFs (t_ms_cs)"):
    """
    Plot ECDFs of t_ms_cs for each mode on a single figure.

    rows: iterable of dicts with keys: "mode", "t_ms_cs"
    modes_order: optional list to control legend/order of modes
    logx: set True to log-scale the time axis
    """
    # group by mode
    by_mode = {}
    for r in rows:
        mode = r.get("mode")
        t = r.get("t_ms_cs", None)
        if mode is None or t is None:
            continue
        if not np.isfinite(t):
            continue
        by_mode.setdefault(mode, []).append(float(t))

    if not by_mode:
        print("No finite t_ms_cs data to plot.")
        return None, None

    # decide plotting order
    modes = modes_order if modes_order is not None else list(by_mode.keys())

    fig, ax = plt.subplots(figsize=(8, 5))
    for mode in modes:
        vals = np.array(by_mode.get(mode, []), dtype=float)
        vals = vals[np.isfinite(vals)]
        if vals.size == 0:
            continue
        x = np.sort(vals)
        y = np.arange(1, x.size + 1) / x.size * 100.0  # percent
        ax.step(x, y, where="post", label=str(mode))

    ax.set_xlabel("Time (ms) [sparse solver t_ms_cs]")
    ax.set_ylabel("Cumulative solves (%)")
    ax.yaxis.set_major_formatter(PercentFormatter(100))
    ax.set_ylim(0, 100)
    if logx:
        ax.set_xscale("log")

    ax.grid(True, alpha=0.3)
    ax.set_axisbelow(True)
    ax.legend(title="Mode")
    ax.set_title(title)
    fig.tight_layout()
    return fig, ax

def plot_compare_solvers(
    results: Dict,
    filename: str = "results/plot_compare_solvers"
):
    fig, axes = plt.subplots(1, 2, figsize=(14, 6), sharey=True)
    for i, ax in enumerate(axes.flatten()):
        res = "primal_residuals"
        if i > 0:
            res = "dual_residuals"
        for mode in reversed(results.keys()):
            d = results[mode]
            ax.plot(np.median(d[res], axis=0), color=colors[mode], linewidth=2, label=mode_to_name[mode])
            p_low = np.percentile(d[res], 2.5, axis=0)
            p_high = np.percentile(d[res], 97.5, axis=0)
            ax.fill_between(np.arange(d[res].shape[1]), p_low, p_high, color=colors[mode], alpha=0.2)
        ax.grid(True)
        ax.tick_params(axis="both", which="major", labelsize=26)
        ax.legend(fontsize=28)
        ax.set_yscale("log")

    axes[0].set_ylabel(r"Residual", fontsize=30)
    axes[0].set_title(r"Primal Residual", fontsize=30)
    axes[1].set_title(r"Dual Residual", fontsize=30)
    axes[0].set_xlabel(r"ADMM Iteration", fontsize=28)
    axes[1].set_xlabel(r"ADMM Iteration", fontsize=28)
    plt.tight_layout()
    plt.savefig(filename + ".jpg")

    fig = plt.figure(figsize=(7, 6))
    true_xs = results[list(results.keys())[0]]["xs"][:, -1, :]
    for mode in reversed(results.keys()):
        d = results[mode]
        delta_xs = np.zeros((d["xs"].shape[0], d["xs"].shape[1]))
        for i in range(delta_xs.shape[0]):
            delta_xs[i] = np.linalg.norm(d["xs"][i] - true_xs[i], axis=-1)
        plt.plot(np.mean(delta_xs, axis=0), color=colors[mode], linewidth=2, label=mode_to_name[mode])
        p_low = np.percentile(delta_xs, 2.5, axis=0)
        p_high = np.percentile(delta_xs, 97.5, axis=0)
        plt.fill_between(np.arange(d["xs"].shape[1]), p_low, p_high, color=colors[mode], alpha=0.2)
    plt.grid(True)
    plt.tick_params(axis="both", which="major", labelsize=20)
    plt.legend(fontsize=24)
    plt.yscale("log")
    plt.ylabel(r"$x_i - x^\star$", fontsize=26)
    plt.title(r"Distance to Solution", fontsize=26)
    plt.xlabel(r"ADMM Iteration", fontsize=24)
    plt.tight_layout()
    plt.savefig(filename + "_xs_distance.jpg")

    plt.show()

def plot_time_solvers(
    results: Dict,
    filename: str = "results/plot_time_solvers"
):
    tol_values = list(results["ADMMslack"].keys())

    fig, axes = plt.subplots(1, len(tol_values), figsize=(18, 6), sharey=True)
    for i, ax in enumerate(axes.flatten()):
        tol = tol_values[i]
        ax.set_title(fr"Solve Times ($\epsilon={tol}$)", fontsize=30)
        ax.set_xlabel(r"State Dimension $n_x$", fontsize=28)
        if i==0:
            ax.set_ylabel(r"Solve Time [ms]", fontsize=30)
        for mode in reversed(results.keys()):
            d = results[mode]
            nx_values = list(d[tol_values[0]].keys())
            solve_times = np.array([[d[tol][nx]["solve_time"] for tol in tol_values] for nx in nx_values]).T
            solve_times = solve_times[:, i, :]
            print("solve_times =", solve_times.shape)
            print("nx_values =", nx_values)
            ax.plot(nx_values, np.median(solve_times, axis=0), color=colors[mode], linewidth=2, label=mode_to_name[mode])
            p_low = np.percentile(solve_times, 2.5, axis=0)
            p_high = np.percentile(solve_times, 97.5, axis=0)
            ax.fill_between(nx_values, p_low, p_high, color=colors[mode], alpha=0.2)
        ax.grid(True)
        ax.tick_params(axis="both", which="both", labelsize=26)
        ax.legend(fontsize=28)
        ax.set_yscale("log")

    plt.tight_layout()
    plt.savefig(filename + ".jpg")

    fig, axes = plt.subplots(1, len(tol_values), figsize=(18, 6), sharey=True)
    for i, ax in enumerate(axes.flatten()):
        tol = tol_values[i]
        ax.set_title(r"Speedup " + f" ($\epsilon={tol}$)", fontsize=30)
        ax.set_xlabel(r"State Dimension $n_x$", fontsize=28)

        d = results["ADMMslack"]
        nx_values = list(d[tol_values[0]].keys())
        solve_times = np.array([[d[tol][nx]["solve_time"] for tol in tol_values] for nx in nx_values]).T
        solve_times = solve_times[:, i, :]
        print("solve_times =", solve_times.shape)

        d = results["StandardSlack"]
        solve_times_standard = np.array([[d[tol][nx]["solve_time"] for tol in tol_values] for nx in nx_values]).T
        print("solve_times_standard =", solve_times_standard.shape)
        solve_times_standard = solve_times_standard[:, i, :]
        print("solve_times_standard =", solve_times_standard.shape)

        speedups = solve_times_standard / solve_times
        
        ax.plot(nx_values, np.median(speedups, axis=0), color=colors[mode], linewidth=2)
        p_low = np.percentile(speedups, 2.5, axis=0)
        p_high = np.percentile(speedups, 97.5, axis=0)
        ax.fill_between(nx_values, p_low, p_high, color=colors[mode], alpha=0.2)
        ax.grid(True)
        ax.tick_params(axis="both", which="both", labelsize=26)
    plt.tight_layout()
    plt.savefig(filename + "_speedup.jpg")
    plt.show()