kkt_solver.hpp

// Copyright (c) Sleipnir contributors
 
#pragma once
 
#include <Eigen/Core>
#include <Eigen/SparseCore>
 
#include "sleipnir/optimization/solver/util/dense_kkt_solver.hpp"
#include "sleipnir/optimization/solver/util/sparse_kkt_solver.hpp"
 
namespace slp {
 
template <typename Scalar>
class KKTSolver {
 public:
  using DenseVector = Eigen::Vector<Scalar, Eigen::Dynamic>;
  using SparseMatrix = Eigen::SparseMatrix<Scalar>;
 
  KKTSolver(bool use_sparse_solver, int num_decision_variables,
            int num_equality_constraints)
      : m_use_sparse_solver{use_sparse_solver},
        m_sparse_solver{num_decision_variables, num_equality_constraints},
        m_dense_solver{num_decision_variables, num_equality_constraints} {}
 
  KKTSolver(bool use_sparse_solver, int num_decision_variables,
            int num_equality_constraints, Scalar γ_min)
      : m_use_sparse_solver{use_sparse_solver},
        m_sparse_solver{num_decision_variables, num_equality_constraints,
                        γ_min},
        m_dense_solver{num_decision_variables, num_equality_constraints,
                       γ_min} {}
 
  Eigen::ComputationInfo info() const {
    if (m_use_sparse_solver) {
      return m_sparse_solver.info();
    } else {
      return m_dense_solver.info();
    }
  }
 
  KKTSolver& compute(const SparseMatrix& lhs) {
    if (m_use_sparse_solver) {
      m_sparse_solver.compute(lhs);
    } else {
      m_dense_solver.compute(lhs);
    }
 
    return *this;
  }
 
  template <typename Rhs>
  DenseVector solve(const Eigen::MatrixBase<Rhs>& rhs) const {
    if (m_use_sparse_solver) {
      return m_sparse_solver.solve(rhs);
    } else {
      return m_dense_solver.solve(rhs);
    }
  }
 
  template <typename Rhs>
  DenseVector solve(const Eigen::SparseMatrixBase<Rhs>& rhs) const {
    if (m_use_sparse_solver) {
      return m_sparse_solver.solve(rhs);
    } else {
      return m_dense_solver.solve(rhs);
    }
  }
 
  Scalar hessian_regularization() const {
    if (m_use_sparse_solver) {
      return m_sparse_solver.hessian_regularization();
    } else {
      return m_dense_solver.hessian_regularization();
    }
  }
 
  Scalar constraint_jacobian_regularization() const {
    if (m_use_sparse_solver) {
      return m_sparse_solver.constraint_jacobian_regularization();
    } else {
      return m_dense_solver.constraint_jacobian_regularization();
    }
  }
 
 private:
  bool m_use_sparse_solver;
  SparseKKTSolver<Scalar> m_sparse_solver;
  DenseKKTSolver<Scalar> m_dense_solver;
};
 
}  // namespace slp