make_plots_modular.py

import numpy as np
import scipy as sp
import pickle
import matplotlib as mpl
import matplotlib.pyplot as plt
import seaborn as sns
import sys
from pathlib import Path
from copy import deepcopy

mpl.rcParams["text.usetex"] = True
mpl.rcParams["font.family"] = "serif"

PROJECT_ROOT = Path(__file__).resolve().parent.parent
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

fafw = True
T_conv = 10  # the "interesting T" where objective_diff_conv was stored

METHOD_STYLES = {
    "SDP": {"label": "MOSEK"},
    "FW": {"label": "FW"},
    "FAFW": {"label": "FA-FW"},
    "FW_AUTOGRAD": {"label": "FW auto"},
    "FAFW_AUTOGRAD": {"label": "FA-FW auto"},
}

RESULT_FILES = [
    "results_to_54.pkl",
    "results_to_86.pkl",
    "results_100.pkl",
]

# sns.set_theme(style="whitegrid", context="paper")
# palette = sns.color_palette("deep")
#COLOR_LQG = palette[0]
#COLOR_DRLQ = palette[1]
#COLOR_DRLQC = palette[2]
#COLOR_SDP = palette[3]


# ----------------------------
# Helpers
# ----------------------------
def pad_with_last(curves):
    """curves: list of 1D lists/arrays, possibly different lengths.
       Returns 2D numpy array padded by repeating last entry."""
    max_len = max(len(c) for c in curves)
    padded = np.array([list(c) + [c[-1]] * (max_len - len(c)) for c in curves])
    return padded


def mean_std_curves(curves, scale=1.0):
    """curves is list of 1D sequences."""
    padded = pad_with_last(curves)
    return scale * padded.mean(axis=0), scale * padded.std(axis=0)


def load_results(path):
    with open(path, "rb") as f:
        return pickle.load(f)


def merge_result_dicts(result_dicts, source_files):
    merged = {"meta": {"source_files": list(source_files)}}

    all_T_values = []
    upper_bounds = []
    interesting_T_values = []
    last_iters = []

    for results_i in result_dicts:
        meta_i = results_i.get("meta", {})
        all_T_values.extend(meta_i.get("T_values", []))
        if meta_i.get("upper_bound_SDP") is not None:
            upper_bounds.append(meta_i["upper_bound_SDP"])
        if meta_i.get("interesting_T") is not None:
            interesting_T_values.append(meta_i["interesting_T"])
        if meta_i.get("last_iter") is not None:
            last_iters.append(meta_i["last_iter"])

        for method, method_data in results_i.items():
            if method == "meta":
                continue

            if method not in merged:
                merged[method] = {"T_values": []}

            for T in method_data.get("T_values", []):
                merged[method][T] = deepcopy(method_data[T])
                if T not in merged[method]["T_values"]:
                    merged[method]["T_values"].append(T)

    merged["meta"]["T_values"] = sorted(set(all_T_values))
    merged["meta"]["upper_bound_SDP"] = max(upper_bounds) if upper_bounds else None
    merged["meta"]["interesting_T"] = interesting_T_values[0] if interesting_T_values else None
    merged["meta"]["last_iter"] = max(last_iters) if last_iters else None
    merged["meta"]["sample_size"] = 0

    for method, method_data in merged.items():
        if method == "meta":
            continue
        method_data["T_values"] = sorted(method_data["T_values"])
        if method_data["T_values"]:
            sample_sizes = [len(method_data[T]["time"]) for T in method_data["T_values"] if "time" in method_data[T]]
            if sample_sizes:
                merged["meta"]["sample_size"] = max(merged["meta"]["sample_size"], max(sample_sizes))

    return merged


# ----------------------------
# Load results (modular format)
# ----------------------------
result_dicts = [load_results(path) for path in RESULT_FILES]
results = merge_result_dicts(result_dicts, RESULT_FILES)

sample_size = results["meta"]["sample_size"]
T_values_all = np.array(results["meta"]["T_values"], dtype=int)

# method-specific T lists (in the modular runner)
methods_present = [m for m in METHOD_STYLES if m in results]
T_values_by_method = {
    m: np.array(results[m]["T_values"], dtype=int) for m in methods_present
}

# Compute N_xi for each T value (x-axis): N_xi = n_x + T_sys*(n_x + n_y)
def get_N_xi(method, T_vals):
    #return np.array([results[method][T]["lqg_system"][0].N_xi for T in T_vals])
    return T_vals

N_xi_by_method = {
    m: get_N_xi(m, T_values_by_method[m]) for m in methods_present
}

print("sample_size:", sample_size)
print("source_files:", results["meta"]["source_files"])
for method in methods_present:
    print(f"T_values_{method}:", T_values_by_method[method])


# ----------------------------
# (1) Run time comparison
# ----------------------------
fig_runtime, ax2 = plt.subplots(1, 1, figsize=(8, 2.5), constrained_layout=True)

runtime_log_scale = False  # True: geometric mean ± std in log space; False: arithmetic mean ± std

def runtime_mean_band(times_2d):
    if runtime_log_scale:
        log_t = [np.log(np.asarray(times_T, dtype=float)) for times_T in times_2d]
        mean_log = np.array([vals.mean() for vals in log_t])
        std_log = np.array([vals.std() for vals in log_t])
        return np.exp(mean_log), np.exp(mean_log - std_log), np.exp(mean_log + std_log)
    else:
        t = [np.asarray(times_T, dtype=float) for times_T in times_2d]
        mean = np.array([vals.mean() for vals in t])
        std = np.array([vals.std() for vals in t])
        return mean, np.maximum(mean - std, 1e-9), mean + std

runtime_methods = ["SDP", "FW", "FAFW", "FW_AUTOGRAD", "FAFW_AUTOGRAD"]
for method in runtime_methods:
    if method not in results:
        continue
    if method in {"FAFW", "FAFW_AUTOGRAD"} and not fafw:
        continue

    T_vals = T_values_by_method[method]
    run_times = [results[method][T]["time"] for T in T_vals]
    mean_t, lo_t, hi_t = runtime_mean_band(run_times)

    print(f"{method} run times per T:", [len(times_T) for times_T in run_times])
    line, = ax2.plot(
        N_xi_by_method[method],
        mean_t,
        label=METHOD_STYLES[method]["label"],
        linewidth=2.0,
    )
    ax2.fill_between(N_xi_by_method[method], lo_t, hi_t, alpha=0.2, color=line.get_color())

#ax2.set_xlabel(r"$N_\xi = n_x + T(n_x + n_y)$", fontsize=8)
# latex math font for d
ax2.set_xlabel(r"$d$", fontsize=15)
ax2.set_xscale('log')
ax2.set_ylabel("runtime [s]", fontsize=15)
ax2.set_yscale('log')
handles, labels = ax2.get_legend_handles_labels()
label_to_handle = dict(zip(labels, handles))
sorted_labels = [
    METHOD_STYLES["FAFW"]["label"],
    METHOD_STYLES["FAFW_AUTOGRAD"]["label"],
    METHOD_STYLES["FW"]["label"],
    METHOD_STYLES["FW_AUTOGRAD"]["label"],
    METHOD_STYLES["SDP"]["label"],
]
sorted_labels = [label for label in sorted_labels if label in label_to_handle]
sorted_handles = [label_to_handle[label] for label in sorted_labels]
ax2.legend(sorted_handles, sorted_labels, fontsize=12, loc='upper left', ncols=5, columnspacing=0.8, handletextpad=0.4)
ax2.set_ylim(bottom=1e-3, top=1e4)  # set lower and upper limits for better log scale visualizatio
ax2.set_xlim(left=1e0, right=1e2)  # set limits for x-axis as well
# ax2.set_title("Run time comparison", fontsize=10)
ax2.tick_params(axis='both', labelsize=13)
ax2.grid(which='both', linestyle='--', linewidth=0.5)
fig_runtime.tight_layout()
fig_runtime.savefig("results_runtime.pdf", format='pdf', bbox_inches='tight')


# ----------------------------
# (2) Gelbrich distance terms (FW)
# ----------------------------
# fig_gelbrich, ax3 = plt.subplots(figsize=(5.5, 4.0))

# if "FAFW" in results:
#     mu_term_fw = []
#     Sigma_term_fw = []
#     for T in T_values_by_method["FAFW"]:
#         mu_terms_T = []
#         Sig_terms_T = []
#         sample_size_T = len(results["FAFW"][T]["time"])
#         for n in range(sample_size_T):
#             lqg = results["FAFW"][T]["lqg_system"][n]
#             mu_opt = results["FAFW"][T]["mu_opt"][n]
#             Sig_opt = results["FAFW"][T]["Sigma_opt"][n]

#             mu_terms_T.append(np.linalg.norm(mu_opt - lqg.mu_hat)**2/lqg.rho**2) # normalize by rho^2 to keep scale O(1), since mu_hat norm grows with dimension

#             # Gelbrich/Bures-like trace term
#             Sh = lqg.Sigma_hat
#             # use sqrtm twice exactly as your original
#             Sig_terms_T.append(
#                 np.trace(Sig_opt + Sh - 2 * sp.linalg.sqrtm(sp.linalg.sqrtm(Sh) @ Sig_opt @ sp.linalg.sqrtm(Sh)))/lqg.rho**2
#             )

#         mu_term_fw.append(mu_terms_T)
#         Sigma_term_fw.append(Sig_terms_T)

#     mu_mean = np.array([np.mean(vals) for vals in mu_term_fw])
#     mu_std = np.array([np.std(vals) for vals in mu_term_fw])
#     sig_mean = np.array([np.mean(vals) for vals in Sigma_term_fw])
#     sig_std = np.array([np.std(vals) for vals in Sigma_term_fw])
#     g_mean = mu_mean + sig_mean
#     g_std = np.array([np.std(np.asarray(mu_vals) + np.asarray(sig_vals)) for mu_vals, sig_vals in zip(mu_term_fw, Sigma_term_fw)])

#     ax3.plot(N_xi_by_method["FAFW"], mu_mean, label=r"$\Vert \mu - \hat{\mu} \Vert_2^2$", marker='o', markersize=4)
#     ax3.fill_between(N_xi_by_method["FAFW"], mu_mean - mu_std, mu_mean + mu_std, alpha=0.2)

#     ax3.plot(N_xi_by_method["FAFW"], sig_mean, label=r"$\mathrm{Tr}(\Sigma + \hat{\Sigma} - 2(\hat{\Sigma}^{1/2}\Sigma \hat{\Sigma}^{1/2})^{1/2})$",
#             marker='o', markersize=4)
#     ax3.fill_between(N_xi_by_method["FAFW"], sig_mean - sig_std, sig_mean + sig_std, alpha=0.2)

#     ax3.plot(N_xi_by_method["FAFW"], g_mean, label=r"$G((\mu,\Sigma),(\hat{\mu},\hat{\Sigma}))$", marker='o', markersize=4)
#     ax3.fill_between(N_xi_by_method["FAFW"], g_mean - g_std, g_mean + g_std, alpha=0.2)

#     # ax3.set_xlabel(r"$N_\xi = n_x + T(n_x + n_y)$", fontsize=8)
#     ax3.set_xlabel(r"$d$", fontsize=8)
#     ax3.set_xscale('log')
#     ax3.set_ylim(-0.25, 1.7)
#     ax3.legend(fontsize=8, loc='upper left')
#     ax3.set_title("Frank-Wolfe Gelbrich distance terms", fontsize=10)
#     ax3.grid(which='both', linestyle='--', linewidth=0.5)
#     fig_gelbrich.tight_layout()
#     fig_gelbrich.savefig("results_gelbrich.pdf", format='pdf', bbox_inches='tight')


# ----------------------------
# (3) Convergence curves at T = 10
# # ----------------------------
# ax4 = fig.add_subplot(gs[0, 2])
# ax4.set_title(f"Convergence rate for T = {T_conv}", fontsize=10)

# # FW curves
# if "FW" in results:
#     fw_curves = results["FW"][T_conv].get("objective_diff_conv", None)
#     if fw_curves is not None:
#         mean_fw_conv, std_fw_conv = mean_std_curves(fw_curves)
#         x_fw = np.arange(1, len(mean_fw_conv) + 1)  # start at 1 for log scale
#         ax4.plot(x_fw, mean_fw_conv, markersize=4, label="FW", color='green')
#         ax4.fill_between(x_fw, mean_fw_conv - std_fw_conv, mean_fw_conv + std_fw_conv, alpha=0.2, color='green')
#     else:
#         ax4.text(0.1, 0.5, "No FW objective_diff_conv stored", transform=ax4.transAxes)

# # FAFW curves
# if fafw and ("FAFW" in results):
#     fafw_curves = results["FAFW"][T_conv].get("objective_diff_conv", None)
#     if fafw_curves is not None:
#         mean_fa_conv, std_fa_conv = mean_std_curves(fafw_curves)
#         x_fa = np.arange(1, len(mean_fa_conv) + 1)
#         ax4.plot(x_fa, mean_fa_conv, markersize=4, label="Fully adaptive FW", color='red')
#         ax4.fill_between(x_fa, mean_fa_conv - std_fa_conv, mean_fa_conv + std_fa_conv, alpha=0.2, color='red')
#     else:
#         ax4.text(0.1, 0.4, "No FAFW objective_diff_conv stored", transform=ax4.transAxes)

# ax4.set_xlabel("iteration", fontsize=8)
# ax4.set_ylabel("relative optimality gap", fontsize=8)
# ax4.set_yscale('log')
# ax4.set_xscale('log')
# ax4.grid(which='both', linestyle='--', linewidth=0.5)
# ax4.legend(fontsize=8, loc='upper right')


# # ----------------------------
# # (4) Heatmaps at T = 10 (hat moments, optimal moments, controller)
# # ----------------------------
# # Use FW’s stored lqg + solution at T_conv, N=0
# if "FW" in results:
#     lqg0 = results["FW"][T_conv]["lqg_system"][0]
#     Sigma_hat = lqg0.Sigma_hat
#     mu_hat = lqg0.mu_hat

#     # ax5 = fig.add_subplot(gs[1, 1])
#     # gap = np.full((mu_hat.shape[0], 1), np.nan)
#     # combined = np.hstack([mu_hat.reshape(-1, 1), gap, Sigma_hat])
#     # sns.heatmap(combined, cmap='coolwarm', cbar=True, annot=True, fmt=".1f",
#     #             linewidths=0.5, annot_kws={"fontsize": 4}, center=0)
#     # ax5.set_title(r"$\hat{\mu}$ and $\hat{\Sigma}$ heatmap", fontsize=10)
#     # ax5.set_xticks([0.5, 2 + mu_hat.shape[0]/2])
#     # ax5.set_xticklabels([r'$\hat{\mu}$', r'$\hat{\Sigma}$'], rotation=0)
#     # ax5.set_yticks([])
#     # ax5.tick_params(axis='x', length=0)

#     # ax6 = fig.add_subplot(gs[1, 2])
#     # mu_opt = results["FW"][T_conv]["mu_opt"][0]
#     # Sigma_opt = results["FW"][T_conv]["Sigma_opt"][0]
#     # gap = np.full((mu_opt.shape[0], 1), np.nan)
#     # combined = np.hstack([mu_opt.reshape(-1, 1), gap, Sigma_opt])
#     # sns.heatmap(combined, cmap='coolwarm', cbar=True, annot=True, fmt=".1f",
#     #             linewidths=0.5, annot_kws={"fontsize": 4}, center=0)
#     # ax6.set_title(r"Optimal $\mu$ and $\Sigma$ heatmap", fontsize=10)
#     # ax6.set_xticks([0.5, 2 + mu_opt.shape[0]/2])
#     # ax6.set_xticklabels([r'$\mu$', r'$\Sigma$'], rotation=0)
#     # ax6.set_yticks([])
#     # ax6.tick_params(axis='x', length=0)

#     ax1 = fig.add_subplot(gs[1, 0])
#     U_opt = results["FW"][T_conv]["U_opt"][0]
#     q_opt = results["FW"][T_conv]["q_opt"][0]
#     gap = np.full((q_opt.shape[0], 1), np.nan)
#     combined = np.hstack([q_opt.reshape(-1, 1), gap, U_opt])
#     sns.heatmap(combined, cmap='coolwarm', cbar=True, annot=True, fmt=".1f",
#                 linewidths=0.5, annot_kws={"fontsize": 7}, center=0)
#     ax1.set_title(r"Optimal $q$ and $U$ heatmap", fontsize=10)
#     ax1.set_xticks([0.5, 2 + q_opt.shape[0]/2])
#     ax1.set_xticklabels([r'$q$', r'$U$'], rotation=0)
#     ax1.set_yticks([])
#     ax1.tick_params(axis='x', length=0)

plt.show()