main.py

import pickle
import argparse
import time
import math
import copy
import types
from dataclasses import dataclass
from typing import Dict, Tuple, List
import numpy as np
import scipy.sparse as sp

import admm_core as ac
from admm import ADMM, ADMMProjection, ADMMRhoUpdate
from linear_control_problem import generate_problem
from plots import plot_time_solvers, plot_compare_solvers, plot_cdf_sparse_timing
    
# -------------------------- Run & compare solvers --------------------------
@dataclass
class RunCfg:
    nx: int = 8
    N: int = 20
    infeasible_start: bool = False
    seeds: List[int] = (0, 1, 2)
    alpha: float = 10.0
    rho: float = 1e-1
    max_iter: int = 2000
    primal_tol: float = 1e-6
    dual_tol: float = 1e-6
    rho_update_steps: int = 25
    rho_update_ratio: float = 10.0
    min_rho: float = 1e-6
    max_rho: float = 1e6
    # which projections to test
    modes: Tuple[str, ...] = ("Standard", "ADMMslack", "StandardSlack")
    # which Python rho update
    rho_update_py: ADMMRhoUpdate = ADMMRhoUpdate.SQRT
    # rho_update_py: ADMMRhoUpdate = ADMMRhoUpdate.NONE
    # rho_update_py: ADMMRhoUpdate = ADMMRhoUpdate.BASIC
    # rho_update_py: ADMMRhoUpdate = ADMMRhoUpdate.OSQP

def run_once(cfg: RunCfg, seed: int, time_python = False, time_dense = False) -> List[Dict]:
    Q_dn, q, A_dn, l, u = generate_problem(cfg.nx, cfg.N, cfg.infeasible_start, seed)
    A_sp = sp.csr_matrix(A_dn)  # identical entries to dense A
    Q_sp = sp.csr_matrix(Q_dn)  # identical entries to dense Q
    x0 = np.zeros(Q_dn.shape[0])

    # Common kwargs for C++
    cpp_kwargs = dict(
        alpha=cfg.alpha, rho=cfg.rho, max_iter=cfg.max_iter,
        primal_tol=cfg.primal_tol, dual_tol=cfg.dual_tol,
        rho_update_steps=cfg.rho_update_steps, rho_update_ratio=cfg.rho_update_ratio,
        min_rho=cfg.min_rho, max_rho=cfg.max_rho,
    )

    # Python solver options
    py_opts = {
        "alpha": cfg.alpha, "rho": cfg.rho, "max_iter": cfg.max_iter,
        "primal_tol": cfg.primal_tol, "dual_tol": cfg.dual_tol,
        "rho_update_steps": cfg.rho_update_steps, "rho_update_ratio": cfg.rho_update_ratio,
        "min_rho": cfg.min_rho, "max_rho": cfg.max_rho,
    }

    results: List[Dict] = []
    for mode in cfg.modes:
        # ---- Python ADMM ----
        py_proj = {
            "Standard": ADMMProjection.STANDARD,
            "ADMMslack": ADMMProjection.ADMMSLACK,
            "StandardSlack": ADMMProjection.STANDARDSLACK,
        }[mode]

        if time_python:
            py_solver = ADMM(Q=Q_dn, q=q, A=A_dn, l=l, u=u, options=py_opts)
            t0 = time.perf_counter()
            x_py, it_py, prim_py, dual_py = py_solver.solve(
                projection=py_proj,
                rho_update=cfg.rho_update_py,
                return_traces=False
            )
            t1 = time.perf_counter()
            t_ms_py = (t1 - t0) * 1e3
        else:
            t_ms_py = math.nan
            it_py, prim_py, dual_py = -1, math.nan, math.nan

        # ---- C++ dense ----
        if time_dense:
            t2 = time.perf_counter()
            res_dense = ac.solve_dense(Q_dn, A_dn, q, l, u, Q_dn.shape[0], x0, mode=mode, **cpp_kwargs)
            t3 = time.perf_counter()
            t_ms_cd = (t3 - t2) * 1e3
        else:
            t_ms_cd = math.nan
            res_dense = types.SimpleNamespace(primal_inf_norm=math.nan, dual_inf_norm=math.nan, iters=-1, converged=False)

        # ---- C++ sparse (SciPy CSR directly) ----
        Q_sp = sp.csr_matrix(Q_dn)  # identical entries to dense Q
        t4 = time.perf_counter()
        res_sparse = ac.solve_sparse_scipy(Q_sp, A_sp, q, l, u, Q_dn.shape[0], x0, mode=mode, **cpp_kwargs)
        t5 = time.perf_counter()
        t_ms_cs = (t5 - t4) * 1e3
        if not bool(res_sparse.converged):
            print("WARNING: C++ sparse did not converge for seed", seed, "mode", mode)
            t_ms_cs = math.inf

        # deltas (∞-norm) w.r.t. C++ dense
        x_s = np.asarray(res_sparse.x)
        if time_dense and time_python:
            x_d = np.asarray(res_dense.x)
            d_py = np.linalg.norm(x_py - x_d, ord=np.inf)
            d_sp = np.linalg.norm(x_s - x_d, ord=np.inf)
        else:
            d_py = math.nan
            d_sp = math.nan

        results.append(dict(
            seed=seed, mode=mode,
            t_ms_py=t_ms_py,
            t_ms_cd=t_ms_cd,
            t_ms_cs=t_ms_cs,
            it_py=it_py, it_cd=res_dense.iters, it_cs=res_sparse.iters,
            prim_py=prim_py, prim_cd=res_dense.primal_inf_norm, prim_cs=res_sparse.primal_inf_norm,
            dual_py=dual_py, dual_cd=res_dense.dual_inf_norm, dual_cs=res_sparse.dual_inf_norm,
            dx_inf_py_vs_cd=d_py, dx_inf_cs_vs_cd=d_sp,
            conv_py=bool(prim_py < cfg.primal_tol and dual_py < cfg.dual_tol),
            conv_cd=bool(res_dense.converged),
            conv_cs=bool(res_sparse.converged),
        ))
    return results

def compare_solvers(cfg: RunCfg) -> Dict:
    Q, _, _, _, _ = generate_problem(
        cfg.nx,
        cfg.N,
        cfg.infeasible_start,
        seed=0
    )

    options = {
        "alpha": cfg.alpha, 
        "rho": cfg.rho,
        "max_iter": cfg.max_iter,
        "primal_tol": cfg.primal_tol,
        "dual_tol": cfg.dual_tol,
        "rho_update_steps": cfg.rho_update_steps,
        "rho_update_ratio": cfg.rho_update_ratio,
        "min_rho": cfg.min_rho,
        "max_rho": cfg.max_rho,
        "check_for_exit": False
    }
    results = {
        mode: {
            "xs": np.zeros((len(cfg.seeds), 1+cfg.max_iter, Q.shape[0])),
            "primal_residuals": np.zeros((len(cfg.seeds), 1+cfg.max_iter)),
            "dual_residuals": np.zeros((len(cfg.seeds), 1+cfg.max_iter)),
            "num_iters": (1+cfg.max_iter) * np.ones(len(cfg.seeds)),
        } for mode in cfg.modes
    }
    for i, seed in enumerate(cfg.seeds):
        Q, q, A, l, u = generate_problem(
            cfg.nx,
            cfg.N,
            cfg.infeasible_start,
            seed=seed
        )
        '''
        Uncomment to see spartisty pattern
        '''
        # plt.spy(A)
        # plt.show()
        for mode in cfg.modes:
            print(f"Python ADMM: mode = {mode}, seed={seed}")
            projection = {
                "Standard": ADMMProjection.STANDARD,
                "ADMMslack": ADMMProjection.ADMMSLACK,
                "StandardSlack": ADMMProjection.STANDARDSLACK,
            }[mode]
            solver = ADMM(Q=Q, q=q, A=A, l=l, u=u, options=options)

            x, num_iter, prim, dual, traces = solver.solve(
                projection=projection,
                rho_update=cfg.rho_update_py,
                return_traces=True
            )

            results[mode]["xs"][i][:(1+num_iter)] = np.array(traces["X"])
            results[mode]["primal_residuals"][i][:(1+num_iter)] = np.array(traces["primal"])
            results[mode]["dual_residuals"][i][:(1+num_iter)] = np.array(traces["dual"])
            results[mode]["num_iters"][i] = 1+num_iter
    return results

def summarize(rows: List[Dict]) -> None:
    by_mode: Dict[str, List[Dict]] = {}
    for r in rows:
        by_mode.setdefault(r["mode"], []).append(r)

    def avg(xs): return sum(xs) / len(xs) if xs else float("nan")

    print("\n=== Linear Control ADMM Ablations (avg over seeds) ===")
    for mode, group in by_mode.items():
        print(f"\nMode: {mode}")
        print(f"  time_ms:  py={avg([g['t_ms_py'] for g in group]):7.1f}   "
              f"cpp_dense={avg([g['t_ms_cd'] for g in group]):7.1f}   "
              f"cpp_sparse={avg([g['t_ms_cs'] for g in group]):7.1f}")
        print(f"  iters:    py={avg([g['it_py'] for g in group]):7.1f}   "
              f"dense={avg([g['it_cd'] for g in group]):7.1f}   "
              f"sparse={avg([g['it_cs'] for g in group]):7.1f}")
        print(f"  prim_inf: py={avg([g['prim_py'] for g in group]):.2e}  "
              f"dense={avg([g['prim_cd'] for g in group]):.2e}  "
              f"sparse={avg([g['prim_cs'] for g in group]):.2e}")
        print(f"  dual_inf: py={avg([g['dual_py'] for g in group]):.2e}  "
              f"dense={avg([g['dual_cd'] for g in group]):.2e}  "
              f"sparse={avg([g['dual_cs'] for g in group]):.2e}")
        print(f"  ||x_py-x_cd||_inf = {avg([g['dx_inf_py_vs_cd'] for g in group]):.2e}   "
              f"||x_cs-x_cd||_inf = {avg([g['dx_inf_cs_vs_cd'] for g in group]):.2e}")
        conv_rate_py = sum(g["conv_py"] for g in group) / len(group)
        conv_rate_cd = sum(g["conv_cd"] for g in group) / len(group)
        conv_rate_cs = sum(g["conv_cs"] for g in group) / len(group)
        print(f"  converge%: py={conv_rate_py*100:5.1f}  dense={conv_rate_cd*100:5.1f}  sparse={conv_rate_cs*100:5.1f}")

def main():
    parser = argparse.ArgumentParser(description="Comparison")
    parser.add_argument("--nx", type=int, default=12, help="State dimension")
    parser.add_argument("--N", type=int, default=20, help="MPC horizon")
    parser.add_argument(
        "--num_repeats",
        type=int,
        default=20, # 100,
        help="Number of repeats of the experimrents",
    )
    parser.add_argument(
        "--infeasible_start",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="If True, solves with infeasible initial condition.",
    )
    parser.add_argument(
        "--compute_results_file",
        type=str,
        default="compare_solvers.pkl",
        help="File name for storing comparison results (.pkl).",
    )
    parser.add_argument(
        "--compute_results",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="If True, plot results.",
    )
    parser.add_argument(
        "--plot",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="If True, plot results.",
    )

    parser.add_argument(
        "--time_results_file",
        type=str,
        default="time_solvers.pkl",
        help="File name for storing timing results (.pkl).",
    )
    parser.add_argument(
        "--time_solvers",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="If True, time solvers.",
    )
    parser.add_argument(
        "--plot_time_solvers",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="If True, plot timing results.",
    )

    args = parser.parse_args()


    compute_results_file = (
        "results/" +
        args.compute_results_file[:-4] +
        "infeasible_start=" +
        str(args.infeasible_start) +
        ".pkl"
    )
    time_results_file = (
        "results/" +
        args.time_results_file[:-4] +
        "infeasible_start=" +
        str(args.infeasible_start) +
        ".pkl"
    )
    if args.compute_results:
        cfg = RunCfg(
            nx=args.nx,
            N=args.N,
            infeasible_start=args.infeasible_start,
            max_iter=500,
            seeds=list(range(args.num_repeats)),
            modes=("ADMMslack", "StandardSlack")
        )
        results = compare_solvers(cfg)
        with open(compute_results_file, "wb") as file:
            pickle.dump(results, file)

    if args.plot:
        with open(compute_results_file, "rb") as file:
            results = pickle.load(file)
        plot_compare_solvers(results, compute_results_file[:-4])

    if args.time_solvers:
        print(">>> Timing solvers")
        nx_values = [4, 8, 16, 32]
        tol_values = [1e-3, 1e-4, 1e-5]
        cfg = RunCfg(
            nx=args.nx,
            N=args.N,
            seeds=list(range(args.num_repeats)),
            modes=("ADMMslack", "StandardSlack")
        )
        results = {
            mode: {
                tol: {
                    nx: {
                        "num_iters": (1+cfg.max_iter) * np.ones(len(cfg.seeds)),
                        "solve_time": np.zeros(len(cfg.seeds))
                    }
                    for nx in nx_values
                } for tol in tol_values
            } for mode in cfg.modes
        }
        for tol in tol_values:
            for nx in nx_values:
                print(f"Solving for (tol, nx) = ({tol}, {nx})")
                cfg = RunCfg(
                    nx=nx,
                    N=args.N,
                    primal_tol=tol,
                    dual_tol=tol,
                    seeds=list(range(args.num_repeats)),
                    modes=("ADMMslack", "StandardSlack")
                )
                for i, s in enumerate(cfg.seeds):
                    rows = run_once(cfg, s)
                    rows = run_once(cfg, s)  # repeat to avoid cold-start effects
                    for r in rows:
                        results[r['mode']][tol][nx]["num_iters"][i] = r['it_cs']
                        results[r['mode']][tol][nx]["solve_time"][i] = r["t_ms_cs"]

        with open(time_results_file, "wb") as file:
            pickle.dump(results, file)

    if args.plot_time_solvers:
        with open(time_results_file, "rb") as file:
            results = pickle.load(file)
        plot_time_solvers(results, time_results_file[:-4])

if __name__ == "__main__":
    main()