run_power_network_compare.py

import argparse
import pickle
from typing import Tuple

import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.colors import TwoSlopeNorm
import numpy as np
import seaborn as sns
from scipy.linalg import sqrtm

from LQGSystem import LQGSystem
from LQGSystem_independent import LQGSystemIndependent
from dual_frankwolfe_dantzig import FrankWolfeOptimizer as CorrelatedFW
from dual_frankwolfe_dantzig_independent import FrankWolfeOptimizerIndependent
from power_network_lqg import build_nominal_independent_lti_data
from power_network_lqg import build_power_network_lti_data_with_covariances

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

# Do we need this function?
def build_interleaved_sigma_from_independent(
    X0: np.ndarray,
    W: np.ndarray,
    V: np.ndarray,
) -> np.ndarray:
    n_x = X0.shape[0]
    T = W.shape[2]
    n_y = V.shape[0]
    N_xi = n_x + T * (n_x + n_y)
    Sigma = np.zeros((N_xi, N_xi), dtype=float)
    Sigma[:n_x, :n_x] = 0.5 * (X0 + X0.T)
    for t in range(T):
        base = n_x + t * (n_x + n_y)
        w_slice = slice(base, base + n_x)
        v_slice = slice(base + n_x, base + n_x + n_y)
        Sigma[w_slice, w_slice] = 0.5 * (W[:, :, t] + W[:, :, t].T)
        Sigma[v_slice, v_slice] = 0.5 * (V[:, :, t] + V[:, :, t].T)
    return Sigma


def cost_matrix_for_controller(lqg_corr: LQGSystem, U_eta: np.ndarray) -> np.ndarray:
    Umap = U_eta @ lqg_corr.F
    Xmap = lqg_corr.D + lqg_corr.H @ Umap
    J = Xmap.T @ lqg_corr.Q @ Xmap + Umap.T @ lqg_corr.R @ Umap
    return 0.5 * (J + J.T)


def expected_cost_from_sigma(lqg_corr: LQGSystem, U_eta: np.ndarray, Sigma_xi: np.ndarray) -> float:
    J = cost_matrix_for_controller(lqg_corr, U_eta)
    return float(np.trace(J @ Sigma_xi))


def linearization_oracle_cov_bures(
    D: np.ndarray,
    Sigma_cur: np.ndarray,
    Sigma_nom: np.ndarray,
    rho: float,
    delta: float,
    max_bisect_iter: int = 200,
    tol: float = 1e-9,
) -> np.ndarray:
    """
    Covariance-only OT/Gelbrich oracle (mean fixed to zero).
    """
    D = 0.5 * (D + D.T)
    Sigma_nom = 0.5 * (Sigma_nom + Sigma_nom.T)
    Sigma_cur = 0.5 * (Sigma_cur + Sigma_cur.T)
    d = D.shape[0]

    eig_vals, eig_vecs = np.linalg.eigh(D)
    lambda1 = eig_vals[-1]
    p1 = eig_vecs[:, -1]

    LB = float(lambda1 * (1.0 + np.sqrt(max(p1.T @ Sigma_nom @ p1, 0.0)) / rho))
    UB = float(lambda1 * (1.0 + np.sqrt(max(np.trace(Sigma_nom), 0.0)) / rho))
    if UB <= LB:
        UB = LB + 1e-6

    offset_D = Sigma_cur.T @ D
    Sigma_nom_sqrt = 0.5 * (np.real_if_close(sqrtm(Sigma_nom)) + np.real_if_close(sqrtm(Sigma_nom)).T)

    def phi(gamma: float, D_shifted_inv: np.ndarray) -> float:
        return float(
            gamma * (rho ** 2 - np.trace(Sigma_nom))
            + gamma * np.sum(D_shifted_inv * Sigma_nom)
            - np.trace(offset_D)
        )

    def Delta(Sigma_tilde: np.ndarray) -> float:
        return float(np.trace(Sigma_tilde.T @ D - offset_D))

    Sigma_tilde = Sigma_nom.copy()
    for _ in range(max_bisect_iter):
        gamma = 0.5 * (LB + UB)
        D_shifted_inv = gamma * np.linalg.inv(gamma * np.eye(d) - D)
        Sigma_tilde = 0.5 * (D_shifted_inv.T @ Sigma_nom @ D_shifted_inv + (D_shifted_inv.T @ Sigma_nom @ D_shifted_inv).T)

        mid = 0.5 * (Sigma_nom_sqrt @ Sigma_tilde @ Sigma_nom_sqrt + (Sigma_nom_sqrt @ Sigma_tilde @ Sigma_nom_sqrt).T)
        mid_sqrt = 0.5 * (np.real_if_close(sqrtm(mid)) + np.real_if_close(sqrtm(mid)).T)
        dphi = float(rho ** 2 - np.trace(Sigma_tilde + Sigma_nom - 2.0 * mid_sqrt))

        if dphi < 0:
            LB = gamma
        else:
            UB = gamma

        if ((dphi >= 0) and (Delta(Sigma_tilde) > delta * phi(gamma, D_shifted_inv))) or (abs(UB - LB) <= tol):
            break

    return Sigma_tilde


def worst_case_cost_for_fixed_controller(
    lqg_corr: LQGSystem,
    U_eta: np.ndarray,
    Sigma_nom: np.ndarray,
    rho: float,
    lambda_min: float,  # kept for API compatibility; not used in covariance-only Bures FW
    delta: float = 1.0 - 1e-3,
    max_iter: int = 2000,
    tol: float = 1e-4,
) -> Tuple[float, np.ndarray]:
    """
    Covariance-only worst-case for a fixed controller over Bures ball:
      max_{Sigma} tr(J Sigma)
      s.t. Gelbrich/OT covariance distance(Sigma, Sigma_nom) <= rho
    solved with Frank-Wolfe (mean fixed to zero).
    """
    J = cost_matrix_for_controller(lqg_corr, U_eta)
    Sigma_k = 0.5 * (Sigma_nom + Sigma_nom.T)

    for k in range(1, max_iter + 1):
        Sigma_tilde = linearization_oracle_cov_bures(
            D=J,
            Sigma_cur=Sigma_k,
            Sigma_nom=Sigma_nom,
            rho=rho,
            delta=delta,
        )
        gap = float(np.trace((Sigma_tilde - Sigma_k).T @ J))
        if abs(gap) / max(abs(np.trace(J @ Sigma_k)), 1e-12) < tol:
            break
        eta = 2.0 / (k + 2.0)
        Sigma_k = 0.5 * ((1.0 - eta) * Sigma_k + eta * Sigma_tilde + ((1.0 - eta) * Sigma_k + eta * Sigma_tilde).T)

    return float(np.trace(J @ Sigma_k)), Sigma_k


def plot_sigmas(
    S_nom: np.ndarray,
    S_ind: np.ndarray,
    S_corr: np.ndarray,
    save_path: str,
    n_x: int,
    n_y: int,
    T: int,
    dpi: int = 300,
):
    fig, axes = plt.subplots(1, 2, figsize=(13, 5), constrained_layout=True)
    mats = [S_nom, S_corr]
    # titles = ["Nominal Sigma", "Independent Worst-Case Sigma", "Correlated Worst-Case Sigma"]
    titles = ["", ""]  # no titles for cleaner look in paper figure
    # Shared, zero-centered color normalization across all panels.
    all_vals = np.concatenate([M.ravel() for M in mats])
    vabs = float(np.max(np.abs(all_vals)))
    if vabs <= 0:
        vabs = 1.0
    norm = TwoSlopeNorm(vmin=-vabs, vcenter=0.0, vmax=vabs)
    last_mesh = None

    # xi = [x0, w0, v0, w1, v1, ..., w_{T-1}, v_{T-1}]
    # Major bounds separate timesteps, with the 0th timestep containing x0 as well.
    # Minor bounds mark the internal splits within each timestep block.
    major_bounds = [n_x + (t + 1) * (n_x + n_y) for t in range(T)]
    minor_bounds = [n_x] + [n_x + t * (n_x + n_y) + n_x for t in range(T)]
    block_centers = [0.5 * n_x]
    block_labels = [r"$x_0$"]
    for t in range(T):
        w_start = n_x + t * (n_x + n_y)
        v_start = w_start + n_x
        block_centers.extend([w_start + 0.5 * n_x, v_start + 0.5 * n_y])
        block_labels.extend([rf"$w_{t}$", rf"$v_{t}$"])

    label_centers = block_centers[:5]
    label_text = block_labels[:5]

    # Force a pure white midpoint at zero.
    # cmap = sns.blend_palette(["#2166ac", "#ffffff", "#b2182b"], as_cmap=True)

    for ax, M, title in zip(axes, mats, titles):
        sns.heatmap(
            M,
            ax=ax,
            cmap="seismic",
            vmin=-vabs,
            vmax=vabs,
            center=0.0,
            cbar=False,
            xticklabels=False,
            yticklabels=False,
            square=False,
        )
        last_mesh = ax.collections[0]
        for b in major_bounds:
            if 0 < b < M.shape[0]:
                ax.axhline(b, color="k", linewidth=3, alpha=0.6, linestyle="--")
                ax.axvline(b, color="k", linewidth=3, alpha=0.6, linestyle="--")
        for b in minor_bounds:
            if 0 < b < M.shape[0]:
                ax.axhline(b, color="k", linewidth=3, alpha=0.35, linestyle=":")
                ax.axvline(b, color="k", linewidth=3, alpha=0.35, linestyle=":")
        ax.set_title(title)
        ax.set_xticks(label_centers)
        ax.set_xticklabels(label_text, rotation=0, fontsize=23)
        ax.set_yticks(label_centers)
        ax.set_yticklabels(label_text, rotation=0, fontsize=23)
        ax.tick_params(top=True, bottom=False, labeltop=True, labelbottom=False, length=0)
        for tick in ax.get_yticklabels():
            tick.set_verticalalignment("center")
        # Solid outline for each panel.
        for spine in ax.spines.values():
            spine.set_visible(True)
            spine.set_linewidth(3)
            spine.set_edgecolor("black")
            spine.set_linestyle("-")
    # One shared colorbar on the right for all three panels.
    # cbar = fig.colorbar(last_mesh, ax=axes.ravel().tolist(), location="right", fraction=0.035, pad=0.03)
    cbar =fig.colorbar(last_mesh, ax=axes.ravel().tolist(), location="left", fraction=0.035, pad=0.03)

    cbar.ax.tick_params(labelsize=23)
    plt.savefig(save_path, format="pdf", bbox_inches="tight", dpi=dpi)
    plt.close(fig)


def plot_costs_vs_alpha(
    alphas: np.ndarray,
    costs_nom_mean: np.ndarray,
    costs_ind_mean: np.ndarray,
    costs_corr_mean: np.ndarray,
    costs_nom_std: np.ndarray,
    costs_ind_std: np.ndarray,
    costs_corr_std: np.ndarray,
    save_path: str,
    ):
    fig, ax = plt.subplots(1, 1, figsize=(8, 2.5), constrained_layout=True)
    ax.plot(alphas, costs_nom_mean, marker="", label="LQG", linewidth=2.0)
    ax.plot(alphas, costs_ind_mean, marker="", label="DRLQ", linewidth=2.0)
    ax.plot(alphas, costs_corr_mean, marker="", label="DRLQ-C (base)", linewidth=2.0)
    ax.fill_between(alphas, costs_nom_mean - costs_nom_std, costs_nom_mean + costs_nom_std, alpha=0.2)
    ax.fill_between(alphas, costs_ind_mean - costs_ind_std, costs_ind_mean + costs_ind_std, alpha=0.2)
    ax.fill_between(alphas, costs_corr_mean - costs_corr_std, costs_corr_mean + costs_corr_std, alpha=0.2)
    ax.set_xscale("log")
    ax.set_xlabel(r"$\kappa$", fontsize=15)
    ax.set_ylabel("cost ratio", fontsize=15)
    ax.set_ylim(top=1.2, bottom=0.95)
    ax.set_xlim(left=alphas[0], right=alphas[-1])
    ax.tick_params(axis='both', labelsize=13)
    ax.grid(True, which="both", alpha=0.3)
    ax.legend(fontsize=12)
    fig.tight_layout()
    plt.savefig(save_path, format="pdf", bbox_inches="tight")
    plt.close(fig)


def render_plots_from_saved_data(data: dict, out_prefix: str, n_buses_fallback: int, T_fallback: int):
    sigma_nom = data.get("Sigma_nom")
    sigma_ind = data.get("Sigma_ind_wc")
    sigma_corr = data.get("Sigma_corr_wc")
    if sigma_nom is not None and sigma_ind is not None and sigma_corr is not None:
        n_x = int(data.get("n_x", 2 * n_buses_fallback))
        n_y = int(data.get("n_y", np.arange(0, n_buses_fallback, 5).size))
        T = int(data.get("T", T_fallback))
        plot_sigmas(
            sigma_nom,
            sigma_ind,
            sigma_corr,
            save_path=f"{out_prefix}_sigmas.pdf",
            n_x=n_x,
            n_y=n_y,
            T=T,
        )
    else:
        print("Skipping sigma plot: missing one of Sigma_nom/Sigma_ind_wc/Sigma_corr_wc in data file.")

    alpha_grid = data.get("alpha_grid")
    nom_mean = data.get("cost_alpha_nom_rel_mean")
    ind_mean = data.get("cost_alpha_ind_rel_mean")
    corr_mean = data.get("cost_alpha_corr_rel_mean")
    nom_std = data.get("cost_alpha_nom_rel_std")
    ind_std = data.get("cost_alpha_ind_rel_std")
    corr_std = data.get("cost_alpha_corr_rel_std")

    # Backward compatibility: build mean/std from sample arrays if needed.
    if (nom_mean is None or ind_mean is None or corr_mean is None) and alpha_grid is not None:
        nom_rel = data.get("cost_alpha_nom_rel")
        ind_rel = data.get("cost_alpha_ind_rel")
        corr_rel = data.get("cost_alpha_corr_rel")
        if nom_rel is not None and ind_rel is not None and corr_rel is not None:
            nom_rel = np.asarray(nom_rel)
            ind_rel = np.asarray(ind_rel)
            corr_rel = np.asarray(corr_rel)
            if nom_rel.ndim == 1:
                nom_mean, ind_mean, corr_mean = nom_rel, ind_rel, corr_rel
                nom_std = np.zeros_like(nom_mean)
                ind_std = np.zeros_like(ind_mean)
                corr_std = np.zeros_like(corr_mean)
            else:
                nom_mean, ind_mean, corr_mean = np.mean(nom_rel, axis=0), np.mean(ind_rel, axis=0), np.mean(corr_rel, axis=0)
                nom_std, ind_std, corr_std = np.std(nom_rel, axis=0), np.std(ind_rel, axis=0), np.std(corr_rel, axis=0)

    if alpha_grid is not None and nom_mean is not None and ind_mean is not None and corr_mean is not None:
        if nom_std is None:
            nom_std = np.zeros_like(np.asarray(nom_mean))
        if ind_std is None:
            ind_std = np.zeros_like(np.asarray(ind_mean))
        if corr_std is None:
            corr_std = np.zeros_like(np.asarray(corr_mean))
        plot_costs_vs_alpha(
            np.asarray(alpha_grid),
            np.asarray(nom_mean),
            np.asarray(ind_mean),
            np.asarray(corr_mean),
            np.asarray(nom_std),
            np.asarray(ind_std),
            np.asarray(corr_std),
            save_path=f"{out_prefix}_cost_vs_alpha.pdf",
        )
    else:
        print("Skipping cost-vs-alpha plot: missing alpha grid and/or relative cost arrays in data file.")


def main():
    parser = argparse.ArgumentParser()
    n_buses = 20
    parser.add_argument("--n_buses", type=int, default=n_buses)
    parser.add_argument("--dt", type=float, default=0.5)
    T = 2
    parser.add_argument("--T", type=int, default=T)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--max_iter_corr", type=int, default=30000)
    parser.add_argument("--max_iter_ind", type=int, default=30000)
    parser.add_argument(
        "--fw_mode",
        type=str,
        default="fully_adaptive",
        choices=["standard", "fully_adaptive"],
        help="Frank-Wolfe mode used for both correlated and independent methods.",
    )
    parser.add_argument("--delta", type=float, default=1 - 1e-3)
    parser.add_argument("--tol", type=float, default=1e-6)
    parser.add_argument("--corr_alpha", type=float, default=10)
    parser.add_argument("--n_sigma_samples", type=int, default=10, help="Number of random covariance directions Sigma_pret.")
    rho_common = 1.0
    rho_x0 = rho_common
    rho_w = rho_common
    rho_v = rho_common/4 # define 
    parser.add_argument("--rho_x0", type=float, default=rho_x0, help="Independent method radius for X0 block.")
    parser.add_argument("--rho_w", type=float, default=rho_w, help="Independent method radius for W blocks.")
    parser.add_argument("--rho_v", type=float, default=rho_v, help="Independent method radius for V blocks.")
    parser.add_argument("--rho_corr", type=float, default=np.sqrt(rho_x0**2 + T*rho_w**2 + T*rho_v**2), help="Radius for correlated method and fixed-controller worst-case evaluation.")
    parser.add_argument("--out_prefix", type=str, default="power_compare")
    parser.add_argument("--plot_only", type=bool, default=False, help="Only load saved data and regenerate plots.")
    parser.add_argument("--data_path", type=str, default=None, help="Path to saved pickle data (defaults to {out_prefix}_data.pkl).")
    args = parser.parse_args()

    if args.plot_only:
        data_path = args.data_path if args.data_path is not None else f"{args.out_prefix}_data.pkl"
        with open(data_path, "rb") as f:
            data = pickle.load(f)
        render_plots_from_saved_data(data, args.out_prefix, args.n_buses, args.T)
        print(
            f"Regenerated plots from {data_path}: "
            f"{args.out_prefix}_sigmas.pdf and {args.out_prefix}_cost_vs_alpha.pdf"
        )
        return

    n = args.n_buses
    m = int(np.arange(0, n, 2).size)
    nx = 2 * n
    nu = n

    # model_data = build_nominal_independent_lti_data(n=n, m=m, dt=args.dt, seed=args.seed)
    rng = np.random.default_rng(args.seed)

    def rand_spd_less_diag_dom(d, rank_frac=0.2, jitter=1e-2, scale=1) -> np.ndarray:
        """
        Dense SPD matrix with weaker diagonal dominance.
        """
        r = max(1, int(rank_frac * d))
        U = rng.standard_normal((d, r))
        S = U @ U.T + jitter * np.eye(d)  # SPD
        # diag is used for scaling to reduce diagonal dominance; we want the final matrix to have more balanced off-diagonal structure, which is more challenging for the independent method. The scaling transforms S into a correlation-like matrix where the diagonal entries are 1, and then we scale the whole matrix by 'scale' to set the overall magnitude.
        diag = np.sqrt(np.diag(S))
        assert np.all(diag > 1e-12), "Diagonal entries must be positive for scaling."
        # explanation of the scaling: we want to scale the matrix so that the Frobenius norm is approximately 'scale', but we also want to reduce diagonal dominance. The original matrix S has a certain Frobenius norm, and we want to apply a scaling factor to it. The term (S / diag[:, None]) / diag[None, :] scales the matrix S by dividing each element S[i,j] by sqrt(S[i,i]) * sqrt(S[j,j]), which reduces diagonal dominance. Then we multiply by 'scale' to set the overall magnitude.
        C = (S / diag[:, None]) / diag[None, :]  # correlation-like SPD, diag=1
        return scale * C

    X0_nom=rand_spd_less_diag_dom(nx)
    W_nom=rand_spd_less_diag_dom(nx)
    V_nom=rand_spd_less_diag_dom(m)

    model_data = build_power_network_lti_data_with_covariances(
        n=n, m=m, dt=args.dt, seed=args.seed,
        X0=X0_nom,
        W=W_nom,
        V=V_nom,
    )

    # Build both LQG systems from the same standardized model.
    lqg_corr = LQGSystem(n_x=nx, n_y=m, n_u=nu, T=args.T, model_data=model_data)
    if args.rho_corr is not None:
        lqg_corr.rho = float(args.rho_corr)

    lqg_ind = LQGSystemIndependent(
        n_x=nx,
        n_y=m,
        n_u=nu,
        T=args.T,
        seed=args.seed,
        model_data=model_data,
        rho=(float(args.rho_corr) if args.rho_corr is not None else None),
        rho_x0=args.rho_x0,
        rho_w=args.rho_w,
        rho_v=args.rho_v,
    )

    # Correlated robust optimization -> U_eta from correlated method.
    fw_corr = CorrelatedFW(
        lqg_corr,
        max_iter=args.max_iter_corr,
        delta=args.delta,
        tol=args.tol,
        conv_tol=args.tol,
        verbose=True,
        convergence_plot=False,
    )
    if args.fw_mode == "fully_adaptive":
        val_corr, mu_corr, Sigma_corr_wc, n_iter_corr = fw_corr.optimize_fully_adaptive()
        U_eta_corr, q_corr = fw_corr.compute_controller()
        val_corr, _, _ = fw_corr.compute_gradients(U_eta_corr, q_corr)
    else:
        corr_out = fw_corr.optimize()
        if corr_out is None:
            U_eta_corr, q_corr = fw_corr.compute_controller()
            val_corr, _, _ = fw_corr.compute_gradients(U_eta_corr, q_corr)
            mu_corr = fw_corr.mu
            Sigma_corr_wc = fw_corr.Sigma
            n_iter_corr = fw_corr.max_iter
        else:
            val_corr, mu_corr, Sigma_corr_wc, U_eta_corr, q_corr, n_iter_corr = corr_out

    

    # Independent robust optimization -> U_eta from independent method.
    fw_ind = FrankWolfeOptimizerIndependent(
        lqg_ind,
        max_iter=args.max_iter_ind,
        delta=args.delta,
        tol=args.tol,
        conv_tol=args.tol,
        verbose=True,
        convergence_plot=False,
    )
    if args.fw_mode == "fully_adaptive":
        val_ind, X0_ind_wc, W_ind_wc, V_ind_wc, n_iter_ind = fw_ind.optimize_fully_adaptive()
    else:
        val_ind, X0_ind_wc, W_ind_wc, V_ind_wc, n_iter_ind = fw_ind.optimize()
    ctrl_ind = fw_ind.compute_worst_case_controller()
    U_eta_ind = ctrl_ind["U_eta"]

    Sigma_nom = lqg_corr.Sigma_hat
    Sigma_ind_wc = build_interleaved_sigma_from_independent(X0_ind_wc, W_ind_wc, V_ind_wc)
    K_nom = lqg_ind.calculate_state_feedback_gains(lqg_ind.P)
    L_nom, _, _ = lqg_ind.calculate_filter_gains(lqg_ind.X0_hat, lqg_ind.W_hat, lqg_ind.V_hat)
    U_y_nom = lqg_ind.observer_gains_to_output_feedback(K_nom, L_nom)
    H_ind, D_ind, C_bar_ind, E_ind, CH_ind, F_ind = lqg_ind.build_stacked_trajectory_matrices()
    U_eta_nom = lqg_ind.output_feedback_to_purified(U_y_nom, CH_ind)
    cost_nominal_opt = expected_cost_from_sigma(lqg_corr, U_eta_nom, Sigma_nom)

    print(f"np.linalg.norm(X0_nom): {np.linalg.norm(X0_nom)}")
    print(f"np.linalg.norm(W_nom): {np.linalg.norm(W_nom)}")
    print(f"np.linalg.norm(V_nom): {np.linalg.norm(V_nom)}")

    print(f"np.linalg.norm(mu_corr): {np.linalg.norm(mu_corr)}")
    print(f"np.linalg.norm(Sigma_nom): {np.linalg.norm(Sigma_nom)}")
    print(f"np.linalg.norm(Sigma_corr_wc): {np.linalg.norm(Sigma_corr_wc)}")
    print(f"np.linalg.norm(Sigma_corr_wc - lqg_corr.Sigma_hat): {np.linalg.norm(Sigma_corr_wc - lqg_corr.Sigma_hat)}")
    print(f"np.linalg.norm(Sigma_corr_wc - Sigma_ind_wc): {np.linalg.norm(Sigma_corr_wc - Sigma_ind_wc)}")
    print(f"np.linalg.norm(U_eta_corr): {np.linalg.norm(U_eta_corr)}")
    print(f"np.linalg.norm(U_eta_ind): {np.linalg.norm(U_eta_ind)}")
    print(f"np.linalg.norm(U_eta_corr - U_eta_ind): {np.linalg.norm(U_eta_corr - U_eta_ind)}")

    # compare system differences in system matrices: H, D, F, Q, R
    Q_ind, R_ind = lqg_ind.build_stacked_cost_matrices()
    print(f"np.linalg.norm(lqg_corr.H - lqg_ind.H): {np.linalg.norm(lqg_corr.H - H_ind)}")
    print(f"np.linalg.norm(lqg_corr.D - lqg_ind.D): {np.linalg.norm(lqg_corr.D - D_ind)}")
    print(f"np.linalg.norm(lqg_corr.F - lqg_ind.F): {np.linalg.norm(lqg_corr.F - F_ind)}")
    print(f"np.linalg.norm(lqg_corr.Q - lqg_ind.Q): {np.linalg.norm(lqg_corr.Q - Q_ind)}")
    print(f"np.linalg.norm(lqg_corr.R - lqg_ind.R): {np.linalg.norm(lqg_corr.R - R_ind)}")
    print(f"np.linalg.norm(lqg_corr.C_bar - C_bar_ind): {np.linalg.norm(lqg_corr.C - C_bar_ind)}")
    print(f"np.linalg.norm(lqg_corr.E - E_ind): {np.linalg.norm(lqg_corr.E - E_ind)}")

    r_x0 = max(1, int(0.2 * nx))
    r_w = max(1, int(0.2 * nx))
    r_v = max(1, int(0.2 * m))
    r_sigma_pret = max(1, int(0.1 * Sigma_nom.shape[0]))
    print("=" * 80)
    print("Problem dimensions and radii")
    print(f"n_x={nx}, n_u={nu}, m={m}")
    print(f"r_X0={r_x0}, r_W={r_w}, r_V={r_v}")
    print(
        f"rho_corr={lqg_corr.rho:.6g}, "
        f"rho_x0={lqg_ind.rho_x0:.6g}, "
        f"rho_w={lqg_ind.rho_w:.6g}, "
        f"rho_v={lqg_ind.rho_v:.6g}"
    )
    print(f"Sigma_pret full dimension={Sigma_nom.shape[0]}, r_Sigma_pret={r_sigma_pret}")

    print("Plotting covariances...")
    # Plot nominal, independent-worst, correlated-worst.
    plot_sigmas(
        Sigma_nom,
        Sigma_ind_wc,
        Sigma_corr_wc,
        save_path=f"{args.out_prefix}_sigmas.pdf",
        n_x=nx,
        n_y=m,
        T=args.T,
    )

    # Out-of-sample correlated true covariance sweep with multiple random directions.
    alpha_grid = np.logspace(-2, 1, num=100, dtype=float)
    n_alpha = alpha_grid.size
    n_sigma_samples = max(1, int(args.n_sigma_samples))
    cost_alpha_nom = np.zeros((n_sigma_samples, n_alpha), dtype=float)
    cost_alpha_ind = np.zeros((n_sigma_samples, n_alpha), dtype=float)
    cost_alpha_corr = np.zeros((n_sigma_samples, n_alpha), dtype=float)
    cost_alpha_nom_rel = np.zeros((n_sigma_samples, n_alpha), dtype=float)
    cost_alpha_ind_rel = np.zeros((n_sigma_samples, n_alpha), dtype=float)
    cost_alpha_corr_rel = np.zeros((n_sigma_samples, n_alpha), dtype=float)

    print("Computing expected costs under Sigma_true(alpha) = Sigma_nom + alpha * (Sigma_pret - Sigma_nom)...")
    for s in range(n_sigma_samples):
        Sigma_pret = rand_spd_less_diag_dom(Sigma_nom.shape[0], rank_frac=0.1, jitter=1e-2, scale=1.0)
        # Sigma_pret = Sigma_corr_wc
        direction = Sigma_pret
        for i, alpha in enumerate(alpha_grid):
            Sigma_true_alpha = Sigma_nom + alpha * direction
            assert np.linalg.eigvalsh(Sigma_true_alpha).min() >= 0, "Direction is not positive semidefinite."
            cost_alpha_nom[s, i] = expected_cost_from_sigma(lqg_corr, U_eta_nom, Sigma_true_alpha)
            cost_alpha_ind[s, i] = expected_cost_from_sigma(lqg_corr, U_eta_ind, Sigma_true_alpha)
            cost_alpha_corr[s, i] = expected_cost_from_sigma(lqg_corr, U_eta_corr, Sigma_true_alpha)

        denom = cost_alpha_corr[s]
        # cost_alpha_nom_rel[s] = (cost_alpha_nom[s] - cost_alpha_corr[s]) / denom
        # cost_alpha_ind_rel[s] = (cost_alpha_ind[s] - cost_alpha_corr[s]) / denom
        # cost_alpha_corr_rel[s] = (cost_alpha_corr[s] - cost_alpha_corr[s]) / denom

        cost_alpha_nom_rel[s] = cost_alpha_nom[s] / denom
        cost_alpha_ind_rel[s] = cost_alpha_ind[s] / denom
        cost_alpha_corr_rel[s] = cost_alpha_corr[s] / denom

    # Mean/std across random perturbation directions.
    cost_alpha_nom_rel_mean = np.mean(cost_alpha_nom_rel, axis=0)
    cost_alpha_ind_rel_mean = np.mean(cost_alpha_ind_rel, axis=0)
    cost_alpha_corr_rel_mean = np.mean(cost_alpha_corr_rel, axis=0)
    cost_alpha_nom_rel_std = np.std(cost_alpha_nom_rel, axis=0)
    cost_alpha_ind_rel_std = np.std(cost_alpha_ind_rel, axis=0)
    cost_alpha_corr_rel_std = np.std(cost_alpha_corr_rel, axis=0)

    # print("Computing worst-case costs for fixed controllers...")
    # wc_cost_ind, Sigma_wc_ind_for_ctrl = worst_case_cost_for_fixed_controller(
    #     lqg_corr=lqg_corr,
    #     U_eta=U_eta_ind,
    #     Sigma_nom=Sigma_nom,
    #     rho=lqg_corr.rho,
    #     lambda_min=lqg_corr.lambda_min,
    # )

    # wc_cost_corr, Sigma_wc_corr_for_ctrl = worst_case_cost_for_fixed_controller(
    #     lqg_corr=lqg_corr,
    #     U_eta=U_eta_corr,
    #     Sigma_nom=Sigma_nom,
    #     rho=lqg_corr.rho,
    #     lambda_min=lqg_corr.lambda_min,
    # )

    print("=" * 80)
    print("Optimization summaries")
    print(f"FW mode: {args.fw_mode}")
    print(f"Nominal objective: {cost_nominal_opt:.6e}")
    print(f"Independent FW objective: {val_ind:.6e}, iterations: {n_iter_ind}")
    print(f"Correlated FW objective: {val_corr:.6e}, iterations: {n_iter_corr}")
    print(f"Radii used: rho_corr={lqg_corr.rho:.6g}, rho_x0={lqg_ind.rho_x0:.6g}, rho_w={lqg_ind.rho_w:.6g}, rho_v={lqg_ind.rho_v:.6g}")

    print("-" * 80)
    print("Relative out-of-sample cost difference under random covariance directions")
    print("(mean over Sigma_pret samples; each sample uses (cost - correlated_cost) / correlated_cost)")
    for i, alpha in enumerate(alpha_grid):
        print(
            f"alpha={alpha:.0e} | nominal={cost_alpha_nom_rel_mean[i]:.6e}, "
            f"independent={cost_alpha_ind_rel_mean[i]:.6e}, correlated={cost_alpha_corr_rel_mean[i]:.6e}"
        )
    print("-" * 80)
    # print("Worst-case fixed-controller cost over covariance ambiguity set")
    # print(f"Worst-case cost (ind controller): {wc_cost_ind:.6e}")
    # print(f"Worst-case cost (corr controller): {wc_cost_corr:.6e}")
    # print("=" * 80)

    plot_costs_vs_alpha(
        alpha_grid,
        cost_alpha_nom_rel_mean,
        cost_alpha_ind_rel_mean,
        cost_alpha_corr_rel_mean,
        cost_alpha_nom_rel_std,
        cost_alpha_ind_rel_std,
        cost_alpha_corr_rel_std,
        save_path=f"{args.out_prefix}_cost_vs_alpha.pdf",
    )

    with open(f"{args.out_prefix}_data.pkl", "wb") as f:
        pickle.dump(
            {
                "U_eta_ind": U_eta_ind,
                "U_eta_corr": U_eta_corr,
                "U_eta_nom": U_eta_nom,
                "q_corr": q_corr,
                "Sigma_nom": Sigma_nom,
                "Sigma_ind_wc": Sigma_ind_wc,
                "Sigma_corr_wc": Sigma_corr_wc,
                "n_x": nx,
                "n_y": m,
                "T": args.T,
                "cost_nominal_opt": cost_nominal_opt,
                # "Sigma_wc_ind_for_ctrl": Sigma_wc_ind_for_ctrl,
                # "Sigma_wc_corr_for_ctrl": Sigma_wc_corr_for_ctrl,
                "alpha_grid": alpha_grid,
                "n_sigma_samples": n_sigma_samples,
                "cost_alpha_nom": cost_alpha_nom,
                "cost_alpha_ind": cost_alpha_ind,
                "cost_alpha_corr": cost_alpha_corr,
                "cost_alpha_nom_rel": cost_alpha_nom_rel,
                "cost_alpha_ind_rel": cost_alpha_ind_rel,
                "cost_alpha_corr_rel": cost_alpha_corr_rel,
                "cost_alpha_nom_rel_mean": cost_alpha_nom_rel_mean,
                "cost_alpha_ind_rel_mean": cost_alpha_ind_rel_mean,
                "cost_alpha_corr_rel_mean": cost_alpha_corr_rel_mean,
                "cost_alpha_nom_rel_std": cost_alpha_nom_rel_std,
                "cost_alpha_ind_rel_std": cost_alpha_ind_rel_std,
                "cost_alpha_corr_rel_std": cost_alpha_corr_rel_std,
                # "wc_cost_ind": wc_cost_ind,
                # "wc_cost_corr": wc_cost_corr,
            },
            f,
        )
    print(
        f"Saved outputs to {args.out_prefix}_sigmas.pdf, "
        f"{args.out_prefix}_cost_vs_alpha.pdf and {args.out_prefix}_data.pkl"
    )


if __name__ == "__main__":
    main()