sparse_kkt_solver.hpp

// Copyright (c) Sleipnir contributors
 
#pragma once
 
#include <algorithm>
#include <limits>
 
#include <Eigen/Core>
#include <Eigen/SparseCholesky>
#include <Eigen/SparseCore>
 
#include "sleipnir/optimization/solver/util/inertia.hpp"
 
namespace slp {
 
template <typename Scalar>
class SparseKKTSolver {
 public:
  using DenseVector = Eigen::VectorX<Scalar>;
  using SparseMatrix = Eigen::SparseMatrix<Scalar>;
 
  SparseKKTSolver(int num_decision_variables, int num_equality_constraints)
      : m_num_decision_variables{num_decision_variables},
        m_num_equality_constraints{num_equality_constraints} {}
 
  SparseKKTSolver(int num_decision_variables, int num_equality_constraints,
                  Scalar γ_min)
      : m_num_decision_variables{num_decision_variables},
        m_num_equality_constraints{num_equality_constraints},
        m_γ_min{γ_min} {}
 
  Eigen::ComputationInfo info() const { return m_info; }
 
  SparseKKTSolver& compute(const SparseMatrix& lhs) {
    // Regularization with zeros ensures the pattern analysis in the sparse
    // solver is reused by all factorizations
    m_lhs = lhs + regularization(Scalar(0), Scalar(0));
 
    if (!m_analyzed_pattern) {
      m_solver.analyzePattern(m_lhs);
      m_analyzed_pattern = true;
    }
 
    // If the previous δ was small and nonzero, skip the δ = 0 attempt because
    // it's almost always rejected for small pivots. If the previous δ was
    // large, attempt δ = 0 so δ can drop to zero quickly.
    if (m_prev_δ == Scalar(0) || m_prev_δ > Scalar(1e-2)) {
      m_solver.factorize(m_lhs);
      m_info = m_solver.info();
 
      if (m_info == Eigen::Success) {
        // If the inertia is ideal and D from LDLT is sufficiently far from
        // zero, don't regularize the system
        if (inertia() == ideal_inertia &&
            (m_solver.vectorD().cwiseAbs().array() >= Scalar(1e-4)).all()) {
          m_prev_δ = Scalar(0);
          m_prev_γ = Scalar(0);
          return *this;
        }
      }
    }
 
    // We'll give lhs the correct inertia by adding [δI, 0; 0, −γI] where δ and
    // γ regularize the Hessian and equality constraint Jacobian respectively.
 
    // The minimum nonzero Hessian regularization
    constexpr Scalar δ_min(1e-10);
 
    // The maximum constraint regularization. Larger values distort the equality
    // constraint Jacobian more than they help a rank-deficient one.
    constexpr Scalar γ_max(1e-8);
 
    // The initial nonzero constraint regularization
    constexpr Scalar γ_initial(1e-10);
 
    // If the previous δ was zero, attempt a small value. Otherwise, attempt a
    // smaller value than the previous δ so δ trends downward. δ has a floor
    // because Hessian rows with no curvature (e.g., states in linear dynamics)
    // have only δ on their diagonal, and pivots that small cause element growth
    // that corrupts the inertia and solution.
    Scalar δ = m_prev_δ == Scalar(0) ? Scalar(1e-4)
                                     : std::max(m_prev_δ / Scalar(2), δ_min);
 
    // Start γ at the minimum to minimize equality constraint Jacobian
    // distortion.
    Scalar γ = m_γ_min;
 
    while (true) {
      m_lhs = lhs + regularization(δ, γ);
      m_solver.factorize(m_lhs);
      m_info = m_solver.info();
 
      if (m_info == Eigen::Success) {
        Inertia inertia = this->inertia();
 
        if (inertia == ideal_inertia) {
          // If the inertia is ideal, report success
          m_prev_δ = δ;
          m_prev_γ = γ;
          return *this;
        } else if (inertia.zero > 0) {
          if (γ == Scalar(0)) {
            // If there's zero eigenvalues and γ = 0, increase γ to potentially
            // compensate for a rank-deficient equality constraint Jacobian
            γ = γ_initial;
          } else {
            // If there's zero eigenvalues and γ > 0, increase δ and γ to drive
            // all eigenvalues away from zero
            δ *= Scalar(10);
            γ = std::min(γ * Scalar(10), γ_max);
          }
        } else if (inertia.negative > ideal_inertia.negative) {
          // If there's too many negative eigenvalues, increase δ to add more
          // positive eigenvalues
          δ *= Scalar(10);
        } else if (inertia.positive > ideal_inertia.positive) {
          // If there's too many positive eigenvalues and the smallest pivot is
          // from the equality constraint Jacobian block, increase γ to add more
          // negative eigenvalues.
          //
          // Otherwise, the smallest pivot is from the Hessian block, or γ is
          // already large enough that increasing it further would distort the
          // equality constraint Jacobian more than it helps. Increase δ to
          // reduce element growth from Hessian pivots instead.
          if (γ < γ_max && !smallest_pivot_is_hessian()) {
            γ = γ == Scalar(0) ? γ_initial : std::min(γ * Scalar(10), γ_max);
          } else {
            δ *= Scalar(10);
          }
        }
      } else {
        // If the decomposition failed, increase δ and γ to drive all
        // eigenvalues away from zero
        δ *= Scalar(10);
        γ = γ == Scalar(0) ? γ_initial : std::min(γ * Scalar(10), γ_max);
      }
 
      // If the lhs perturbation is too high, report failure. This can be caused
      // by ill-conditioning.
      if (δ > Scalar(1e20) || γ > Scalar(1e20)) {
        m_info = Eigen::NumericalIssue;
        m_prev_δ = δ;
        m_prev_γ = γ;
        return *this;
      }
    }
  }
 
  template <typename Rhs>
  DenseVector solve(const Eigen::MatrixBase<Rhs>& rhs) const {
    DenseVector x = m_solver.solve(rhs);
 
    // Element growth in the factorization can make the solution inaccurate
    // (e.g., from small Hessian pivots in rows with no curvature), so apply
    // iterative refinement with the same factorization
    constexpr int max_refinement_steps = 3;
    for (int i = 0; i < max_refinement_steps; ++i) {
      DenseVector r = rhs - m_lhs.template selfadjointView<Eigen::Lower>() * x;
      if (r.norm() <= Scalar(1e-12) * rhs.norm()) {
        break;
      }
      x += m_solver.solve(r);
    }
 
    return x;
  }
 
  template <typename Rhs>
  DenseVector solve(const Eigen::SparseMatrixBase<Rhs>& rhs) const {
    return solve(rhs.toDense());
  }
 
  Scalar hessian_regularization() const { return m_prev_δ; }
 
  Scalar constraint_jacobian_regularization() const { return m_prev_γ; }
 
 private:
  using Solver = Eigen::SimplicialLDLT<SparseMatrix>;
 
  Solver m_solver;
  bool m_analyzed_pattern = false;
 
  SparseMatrix m_lhs;
 
  Eigen::ComputationInfo m_info = Eigen::Success;
 
  int m_num_decision_variables = 0;
 
  int m_num_equality_constraints = 0;
 
  Scalar m_γ_min{1e-10};
 
  Inertia ideal_inertia{m_num_decision_variables, m_num_equality_constraints,
                        0};
 
  Scalar m_prev_δ{0};
 
  Scalar m_prev_γ{0};
 
  Inertia inertia() const {
    constexpr Scalar ε = std::numeric_limits<Scalar>::epsilon();
 
    Inertia inertia;
    for (const auto& d : m_solver.vectorD()) {
      if (d > ε) {
        ++inertia.positive;
      } else if (d < -ε) {
        ++inertia.negative;
      } else {
        ++inertia.zero;
      }
    }
 
    return inertia;
  }
 
  bool smallest_pivot_is_hessian() const {
    Eigen::Index k;
    m_solver.vectorD().cwiseAbs().minCoeff(&k);
 
    // The factorization is of PAP⁻¹, so the original row of pivot k is P⁻¹k
    return m_solver.permutationPinv().indices()(k) < m_num_decision_variables;
  }
 
  SparseMatrix regularization(Scalar δ, Scalar γ) const {
    DenseVector vec{m_num_decision_variables + m_num_equality_constraints};
    vec.segment(0, m_num_decision_variables).setConstant(δ);
    vec.segment(m_num_decision_variables, m_num_equality_constraints)
        .setConstant(-γ);
 
    return SparseMatrix{vec.asDiagonal()};
  }
};
 
}  // namespace slp