admm.hpp

#pragma once
#include <Eigen/Dense>
#include <Eigen/Sparse>
#include <limits>
#include <string>
#include <cassert>
#include <cmath>

namespace admm {

//------------------------------- Options / Result --------------------------------
enum class Mode { Standard, ADMMslack, StandardSlack };

struct Options {
    double alpha = 1.0;          // slack penalty (only used in ADMMslack)
    double rho   = 1e-1;         // initial rho
    int    max_iter = 1000;
    double primal_tol = 1e-4;
    double dual_tol   = 1e-4;

    int    rho_update_steps = 25;     // 0 disables updates
    double rho_update_ratio = 10.0;   // imbalance threshold
    double min_rho = 1e-9;
    double max_rho = 1e9;

    Mode   mode = Mode::Standard;
};

struct Result {
    Eigen::VectorXd x;      // first n entries (for augmented systems)
    int iters = 0;
    double primal_inf_norm = std::numeric_limits<double>::infinity();
    double dual_inf_norm   = std::numeric_limits<double>::infinity();
    double rho_out = 0.0;
    bool converged = false;
    // Eigen::VectorXd primal_inf_norm_vec;
    // Eigen::VectorXd dual_inf_norm;
};

//----------------------------------- helpers ------------------------------------
inline double clamp(double v, double lo, double hi) {
    return std::max(lo, std::min(hi, v));
}

inline void project_box_hard(Eigen::VectorXd& z,
                             const Eigen::VectorXd& lo,
                             const Eigen::VectorXd& hi) {
    z = z.cwiseMax(lo).cwiseMin(hi);
}

inline void project_box_slack(Eigen::VectorXd& z,
                              const Eigen::VectorXd& z_tilde,
                              const Eigen::VectorXd& lo,
                              const Eigen::VectorXd& hi,
                              double rho, double alpha)
{
    // Precompute divison operation
    const double tau = alpha / (rho + alpha);
    // Array views (no allocation)
    auto z_a  = z.array();
    const auto zt = z_tilde.array();
    const auto lo_a = lo.array();
    const auto hi_a = hi.array();

    // Start from z_tilde everywhere
    z_a = zt;
    // Where below bound: z = z_tilde + tau*(lo - z_tilde)
    z_a = (zt < lo_a).select( zt + tau * (lo_a - zt), z_a );
    // Where above bound: z = z_tilde + tau*(hi - z_tilde)
    z_a = (zt > hi_a).select( zt + tau * (hi_a - zt), z_a );
}

//----------------------- Linear system adapters (C++17) --------------------------
//
// These wrap the "P x = rhs" solves so the core loop is shared between dense/sparse.
//
// DenseLinSys:  P = Q + rho * (A^T A), factor with LLT if possible, else LDLT.
// SparseLinSys: P = Q + rho * (A^T A), factor with SimplicialLDLT.
//
struct DenseLinSys {
    // references to precomputed pieces (not owned)
    const Eigen::MatrixXd& Q;
    const Eigen::MatrixXd& AtA;

    // owned
    Eigen::MatrixXd P;
    Eigen::LLT<Eigen::MatrixXd>  llt;
    Eigen::LDLT<Eigen::MatrixXd> ldlt;
    bool use_llt = true;

    DenseLinSys(const Eigen::MatrixXd& Q_,
                const Eigen::MatrixXd& AtA_,
                double rho)
        : Q(Q_), AtA(AtA_), P(Q_ + rho * AtA_)  // build P once
    {
        refactor(rho); // factor P
    }

    void refactor(double /*rho*/) {
        // P already set by caller (or by constructor); (re)factor it
        llt.compute(P);
        if (llt.info() == Eigen::Success) { use_llt = true; return; }

        ldlt.compute(P);
        if (ldlt.info() == Eigen::Success) { use_llt = false; return; }

        // Tiny diagonal shift fallback
        Eigen::MatrixXd P_shift = P;
        P_shift.diagonal().array() += 1e-12;
        ldlt.compute(P_shift);
        use_llt = false; // we fall back to LDLT
    }

    void rebuild_and_refactor(double rho) {
        P = Q + rho * AtA;
        refactor(rho);
    }

    void solve_into(const Eigen::VectorXd& rhs, Eigen::VectorXd& x) {
        if (use_llt) {
            x = llt.solve(rhs);
            if (llt.info() == Eigen::Success) return;
            use_llt = false;
            ldlt.compute(P); // fallback
        }
        x = ldlt.solve(rhs);
    }
};

struct SparseLinSys {
    const Eigen::SparseMatrix<double>& Q;
    const Eigen::SparseMatrix<double>& AtA;

    Eigen::SparseMatrix<double> P;
    Eigen::SimplicialLDLT<Eigen::SparseMatrix<double>> ldlt;

    SparseLinSys(const Eigen::SparseMatrix<double>& Q_,
                 const Eigen::SparseMatrix<double>& AtA_,
                 double rho)
        : Q(Q_), AtA(AtA_), P(Q_ + rho * AtA_)
    {
        refactor();
    }

    void refactor() {
        ldlt.compute(P);
        if (ldlt.info() != Eigen::Success) {
            Eigen::SparseMatrix<double> I(P.rows(), P.cols());
            I.setIdentity();
            ldlt.compute(P + 1e-12 * I);
        }
    }

    void rebuild_and_refactor(double rho) {
        P = Q + rho * AtA;
        refactor();
    }

    void solve_into(const Eigen::VectorXd& rhs, Eigen::VectorXd& x) {
        x = ldlt.solve(rhs);
    }
};

//------------------------------ core ADMM loop ----------------------------------
//
// Works for both DenseLinSys and SparseLinSys.
// A, At, Q multiply naturally in both dense and sparse Eigen.
//
template <class MatA, class MatAT, class QMat, class LinSys>
inline Result solve_core(
    const QMat& Q,          // (N x N)
    const MatA& A,          // (m x N)
    const MatAT& At,        // (N x m)
    const Eigen::VectorXd& q,   // (N)
    const Eigen::VectorXd& l,   // (m)
    const Eigen::VectorXd& u,   // (m)
    int n,                      // number of original x entries to return
    Eigen::VectorXd x,          // initial x (N) (by value on purpose; we modify it)
    const Options& opts,
    LinSys& linsys              // pre-built with Q, AtA, rho
) {
    const int N = static_cast<int>(x.size());
    const int m = static_cast<int>(A.rows());
    assert(Q.rows() == Q.cols() && Q.rows() == N);
    assert(q.size() == N);
    assert(A.cols() == N && l.size() == m && u.size() == m);
    assert(n <= N);

    double rho = opts.rho;
    const double eps = 1e-12;

    Eigen::VectorXd z  = Eigen::VectorXd::Zero(m);
    Eigen::VectorXd mu = Eigen::VectorXd::Zero(m);

    // Pre-allocate temporaries
    Eigen::VectorXd rhs(N), tmp(N), Ax(m), Axmz(m);

    double prim = std::numeric_limits<double>::infinity();
    double dual = std::numeric_limits<double>::infinity();

    int k = 0;
    for (k = 0; k < opts.max_iter; ++k) {
        // rhs = -q + At * (-mu + rho*z)
        rhs = -q;
        rhs.noalias() += At * (-mu + rho * z);

        linsys.solve_into(rhs, x);

        // z-update: z_tilde = A*x + mu/rho
        Ax.noalias() = A * x;
        Eigen::VectorXd z_tilde = Ax + mu / rho;

        if (opts.mode == Mode::ADMMslack) {
            project_box_slack(z, z_tilde, l, u, rho, opts.alpha);
        } else {
            z = z_tilde.cwiseMax(l).cwiseMin(u);
        }

        // mu update and residuals
        Axmz.noalias() = Ax - z;
        mu.noalias()  += rho * Axmz;

        prim = Axmz.lpNorm<Eigen::Infinity>();

        tmp.noalias() = Q * x;
        tmp.noalias() += q;
        tmp.noalias() += At * mu;
        dual = tmp.lpNorm<Eigen::Infinity>();

        if (prim < opts.primal_tol && dual < opts.dual_tol) {
            Result r;
            r.x = x.head(n);
            r.iters = k + 1;
            r.primal_inf_norm = prim;
            r.dual_inf_norm   = dual;
            r.rho_out = rho;
            r.converged = true;
            return r;
        }

        // rho update
        if (opts.rho_update_steps > 0 && (k % opts.rho_update_steps) == 0) {
            if ((prim > opts.rho_update_ratio * dual) || (dual > opts.rho_update_ratio * prim)) {
                double new_rho = clamp(rho * std::sqrt(prim / (dual + eps)), opts.min_rho, opts.max_rho);
                if (std::abs(new_rho - rho) > 0) {
                    rho = new_rho;
                    linsys.rebuild_and_refactor(rho);
                }
            }
        }
    }

    Result r;
    r.x = x.head(n);
    r.iters = k;
    r.primal_inf_norm = prim;
    r.dual_inf_norm   = dual;
    r.rho_out = rho;
    r.converged = (prim < opts.primal_tol && dual < opts.dual_tol);
    return r;
}

//------------------------------ public entrypoints -------------------------------

inline Result solve_dense(
    const Eigen::MatrixXd& Q,     // (N x N)
    const Eigen::MatrixXd& A,     // (m x N)
    const Eigen::VectorXd& q,     // (N)
    const Eigen::VectorXd& l,     // (m)
    const Eigen::VectorXd& u,     // (m)
    int n,                        // # original x entries to return
    Eigen::VectorXd x0,           // (N)
    const Options& opts = Options()
) {
    const int N = static_cast<int>(x0.size());
    const int m = static_cast<int>(A.rows());
    assert(Q.rows() == Q.cols() && Q.rows() == N);
    assert(A.cols() == N && q.size() == N && l.size() == m && u.size() == m);

    if (opts.mode == Mode::StandardSlack) {
        // Build augmented dense problem
        const int Nbar = N + m;

        Eigen::MatrixXd Qbar = Eigen::MatrixXd::Zero(Nbar, Nbar);
        Qbar.topLeftCorner(N, N) = Q;
        Qbar.bottomRightCorner(m, m).setIdentity();
        Qbar.bottomRightCorner(m, m) *= opts.alpha;     // alpha*I_m

        Eigen::VectorXd qbar(Nbar);
        qbar << q, Eigen::VectorXd::Zero(m);

        Eigen::MatrixXd Abar(m, Nbar);
        Abar.leftCols(N)  = A;
        Abar.rightCols(m).setIdentity();

        Eigen::VectorXd x0bar = Eigen::VectorXd::Zero(Nbar);
        x0bar.head(N) = x0; // slack part starts at 0

        const Eigen::MatrixXd At  = Abar.transpose();
        const Eigen::MatrixXd AtA = At * Abar;

        // Use Standard projection in core loop
        DenseLinSys linsys(Qbar, AtA, opts.rho);
        Options local = opts; local.mode = Mode::Standard;
        return solve_core(Qbar, Abar, At, qbar, l, u, n, std::move(x0bar), local, linsys);
    }

    // Regular (non-augmented) path
    const Eigen::MatrixXd At  = A.transpose();
    const Eigen::MatrixXd AtA = At * A;

    DenseLinSys linsys(Q, AtA, opts.rho);
    return solve_core(Q, A, At, q, l, u, n, std::move(x0), opts, linsys);
}

inline Result solve_sparse(
    const Eigen::SparseMatrix<double>& Q,   // (N x N)
    const Eigen::SparseMatrix<double>& A,   // (m x N)
    const Eigen::VectorXd& q,               // (N)
    const Eigen::VectorXd& l,               // (m)
    const Eigen::VectorXd& u,               // (m)
    int n,                                  // # original x entries to return
    Eigen::VectorXd x0,                     // (N)
    const Options& opts = Options()
) {
    const int N = static_cast<int>(x0.size());
    const int m = static_cast<int>(A.rows());
    assert(Q.rows() == Q.cols() && Q.rows() == N);
    assert(A.cols() == N && q.size() == N && l.size() == m && u.size() == m);

    if (opts.mode == Mode::StandardSlack) {
        // Build augmented sparse problem
        const int Nbar = N + m;

        // Qbar = blockdiag(Q, alpha*I_m)
        std::vector<Eigen::Triplet<double>> trips;
        trips.reserve(Q.nonZeros() + m);
        for (int k = 0; k < Q.outerSize(); ++k)
            for (Eigen::SparseMatrix<double>::InnerIterator it(Q, k); it; ++it)
                trips.emplace_back(it.row(), it.col(), it.value());
        for (int i = 0; i < m; ++i)
            trips.emplace_back(N + i, N + i, opts.alpha);

        Eigen::SparseMatrix<double> Qbar(Nbar, Nbar);
        Qbar.setFromTriplets(trips.begin(), trips.end());
        Qbar.makeCompressed();

        // Abar = [A  I_m]
        std::vector<Eigen::Triplet<double>> Atrips;
        Atrips.reserve(A.nonZeros() + m);
        for (int k = 0; k < A.outerSize(); ++k)
            for (Eigen::SparseMatrix<double>::InnerIterator it(A, k); it; ++it)
                Atrips.emplace_back(it.row(), it.col(), it.value());
        for (int i = 0; i < m; ++i)
            Atrips.emplace_back(i, N + i, 1.0);

        Eigen::SparseMatrix<double> Abar(m, Nbar);
        Abar.setFromTriplets(Atrips.begin(), Atrips.end());
        Abar.makeCompressed();

        Eigen::VectorXd qbar(Nbar);
        qbar << q, Eigen::VectorXd::Zero(m);

        Eigen::VectorXd x0bar = Eigen::VectorXd::Zero(Nbar);
        x0bar.head(N) = x0;

        Eigen::SparseMatrix<double> At = Abar.transpose();
        Eigen::SparseMatrix<double> AtA = (At * Abar).pruned();

        SparseLinSys linsys(Qbar, AtA, opts.rho);
        Options local = opts; local.mode = Mode::Standard;  // standard projection
        return solve_core(Qbar, Abar, At, qbar, l, u, n, std::move(x0bar), local, linsys);
    }

    // Regular (non-augmented) path
    Eigen::SparseMatrix<double> At  = A.transpose();
    Eigen::SparseMatrix<double> AtA = (At * A).pruned();

    SparseLinSys linsys(Q, AtA, opts.rho);
    return solve_core(Q, A, At, q, l, u, n, std::move(x0), opts, linsys);
}

} // namespace admm