SapModel caches the Hessian factorization

multibody/contact_solvers/sap/BUILD.bazel

@@ -15,6 +15,7 @@ drake_cc_package_library(
     visibility = ["//visibility:public"],
     deps = [
         ":contact_problem_graph",
+        ":dense_supernodal_solver",
         ":partial_permutation",
         ":sap_ball_constraint",
         ":sap_constraint",
@@ -47,6 +48,16 @@ drake_cc_library(
     ],
 )
 
+drake_cc_library(
+    name = "dense_supernodal_solver",
+    srcs = ["dense_supernodal_solver.cc"],
+    hdrs = ["dense_supernodal_solver.h"],
+    deps = [
+        "//common:essential",
+        "//multibody/contact_solvers:supernodal_solver",
+    ],
+)
+
 drake_cc_library(
     name = "partial_permutation",
     srcs = ["partial_permutation.cc"],
@@ -245,13 +256,16 @@ drake_cc_library(
     hdrs = ["sap_model.h"],
     deps = [
         ":contact_problem_graph",
+        ":dense_supernodal_solver",
         ":partial_permutation",
         ":sap_constraint_bundle",
         ":sap_contact_problem",
         "//common:default_scalars",
         "//common:essential",
         "//math:linear_solve",
         "//multibody/contact_solvers:block_sparse_matrix",
+        "//multibody/contact_solvers:block_sparse_supernodal_solver",
+        "//multibody/contact_solvers:conex_supernodal_solver",
         "//systems/framework:context",
         "//systems/framework:leaf_system",
     ],
@@ -284,6 +298,15 @@ drake_cc_library(
     deps = ["//common:default_scalars"],
 )
 
+drake_cc_googletest(
+    name = "dense_supernodal_solver_test",
+    deps = [
+        ":dense_supernodal_solver",
+        "//common/test_utilities:eigen_matrix_compare",
+        "//common/test_utilities:expect_throws_message",
+    ],
+)
+
 drake_cc_library(
     name = "validate_constraint_gradients",
     testonly = 1,

multibody/contact_solvers/sap/dense_supernodal_solver.cc

@@ -0,0 +1,89 @@
+#include "drake/multibody/contact_solvers/sap/dense_supernodal_solver.h"
+
+#include <vector>
+
+namespace drake {
+namespace multibody {
+namespace contact_solvers {
+namespace internal {
+
+int SafeGetCols(const BlockSparseMatrix<double>* J) {
+  DRAKE_THROW_UNLESS(J != nullptr);
+  return J->cols();
+}
+
+template <typename T>
+const T& SafeGetReference(std::string_view variable_name, const T* ptr) {
+  if (ptr == nullptr) {
+    throw std::runtime_error(
+        fmt::format("Condition '{} != nullptr' failed.", variable_name));
+  }
+  return *ptr;
+}
+
+DenseSuperNodalSolver::DenseSuperNodalSolver(
+    const std::vector<MatrixX<double>>* A, const BlockSparseMatrix<double>* J)
+    : A_(SafeGetReference("A", A)),
+      J_(SafeGetReference("J", J)),
+      H_(SafeGetCols(J), SafeGetCols(J)),
+      Hldlt_(H_) {
+  const int nv = [this]() {
+    int size = 0;
+    for (const auto& Ai : A_) {
+      DRAKE_THROW_UNLESS(Ai.rows() == Ai.cols());
+      size += Ai.rows();
+    }
+    return size;
+  }();
+  DRAKE_THROW_UNLESS(nv == J->cols());
+}
+
+bool DenseSuperNodalSolver::DoSetWeightMatrix(
+    const std::vector<Eigen::MatrixXd>& G) {
+  const int nv = H_.rows();
+
+  // Make dense dynamics matrix.
+  MatrixX<double> Adense = MatrixX<double>::Zero(nv, nv);
+  int offset = 0;
+  for (const auto& Ac : A_) {
+    const int nv_clique = Ac.rows();
+    Adense.block(offset, offset, nv_clique, nv_clique) = Ac;
+    offset += nv_clique;
+  }
+
+  // Make dense Jacobian matrix.
+  const MatrixX<double> Jdense = J_.MakeDenseMatrix();
+
+  // Make dense weight matrix G.
+  const int nk = Jdense.rows();
+  MatrixX<double> Gdense = MatrixX<double>::Zero(nk, nk);
+  offset = 0;
+  for (const auto& Gi : G) {
+    const int ni = Gi.rows();
+    if (offset + ni > nk) {
+      // Weight matrix G is incompatible with the Jacobian matrix J.
+      return false;
+    }
+    Gdense.block(offset, offset, ni, ni) = Gi;
+    offset += ni;
+  }
+  if (offset != nk) return false;  // G might miss some rows.
+
+  H_ = Adense + Jdense.transpose() * Gdense * Jdense;
+
+  return true;
+}
+
+bool DenseSuperNodalSolver::DoFactor() {
+  Hldlt_.compute(H_);
+  return Hldlt_.info() == Eigen::Success;
+}
+
+void DenseSuperNodalSolver::DoSolveInPlace(Eigen::VectorXd* b) const {
+  *b = Hldlt_.solve(*b);
+}
+
+}  // namespace internal
+}  // namespace contact_solvers
+}  // namespace multibody
+}  // namespace drake

multibody/contact_solvers/sap/dense_supernodal_solver.h

@@ -0,0 +1,60 @@
+#pragma once
+
+#include <vector>
+
+#include "drake/multibody/contact_solvers/supernodal_solver.h"
+
+namespace drake {
+namespace multibody {
+namespace contact_solvers {
+namespace internal {
+
+// A class that implements the SuperNodalSolver interface using dense algebra.
+// As with every other SuperNodalSolver, this class implements a Cholesky
+// factorization of a Hessian matrix H of the form H = A + Jᵀ⋅G⋅J, where A is
+// referred to as the dynamics matrix and J is the Jacobian matrix.
+class DenseSuperNodalSolver final : public SuperNodalSolver {
+ public:
+  DRAKE_NO_COPY_NO_MOVE_NO_ASSIGN(DenseSuperNodalSolver)
+
+  // Constructs a dense solver for dynamics matrix A and Jacobian matrix J.
+  // This class holds references to matrices A and J, and therefore they must
+  // outlive this object.
+  // @throws std::exception if A or J is nullptr.
+  // @throws std::exception if a block in A is not square.
+  // @pre each block in A is SPD.
+  // @pre The number of columns of J equals the size of A.
+  // @note The sizes in J are not checked here, but during SetWeightMatrix().
+  DenseSuperNodalSolver(const std::vector<MatrixX<double>>* A,
+                        const BlockSparseMatrix<double>* J);
+
+ private:
+  // Implementations of SuperNodalSolver NVIs. NVIs perform basic checks.
+  bool DoSetWeightMatrix(
+      const std::vector<Eigen::MatrixXd>& block_diagonal_G) final;
+  Eigen::MatrixXd DoMakeFullMatrix() const final {
+    // N.B. SuperNodalSolver's NVI already checked that SetWeightMatrix() was
+    // called and the matrix was not yet factored via Factor(). Therefore no
+    // additional checks are needed here.
+    return H_;
+  }
+  bool DoFactor() final;
+  void DoSolveInPlace(Eigen::VectorXd* b) const final;
+  int DoGetSize() const final { return H_.rows(); }
+
+  const std::vector<MatrixX<double>>& A_;
+  const BlockSparseMatrix<double>& J_;
+  // The support for dense algebra is mostly for testing purposes, even
+  // though the computation of the dense H (and in particular of the Jᵀ⋅G⋅J
+  // term) is very costly. Therefore below we decided to trade off speed for
+  // stability when choosing to use an LDLT decomposition instead of a slightly
+  // faster, though less stable, LLT decomposition.
+  MatrixX<double> H_;
+  // N.B. This instantiates an in-place LDLT solver that uses the memory in H_.
+  Eigen::LDLT<Eigen::Ref<MatrixX<double>>> Hldlt_;
+};
+
+}  // namespace internal
+}  // namespace contact_solvers
+}  // namespace multibody
+}  // namespace drake

multibody/contact_solvers/sap/sap_model.cc

@@ -6,7 +6,10 @@
 #include "drake/common/default_scalars.h"
 #include "drake/math/linear_solve.h"
 #include "drake/multibody/contact_solvers/block_sparse_matrix.h"
+#include "drake/multibody/contact_solvers/block_sparse_supernodal_solver.h"
+#include "drake/multibody/contact_solvers/conex_supernodal_solver.h"
 #include "drake/multibody/contact_solvers/sap/contact_problem_graph.h"
+#include "drake/multibody/contact_solvers/sap/dense_supernodal_solver.h"
 
 namespace drake {
 namespace multibody {
@@ -15,9 +18,60 @@ namespace internal {
 
 using systems::Context;
 
+HessianFactorizationCache::HessianFactorizationCache(
+    SapHessianFactorizationType type, const std::vector<MatrixX<double>>* A,
+    const BlockSparseMatrix<double>* J) {
+  DRAKE_DEMAND(A != nullptr);
+  DRAKE_DEMAND(J != nullptr);
+  switch (type) {
+    case SapHessianFactorizationType::kConex:
+      factorization_ = std::make_unique<ConexSuperNodalSolver>(
+          J->block_rows(), J->get_blocks(), *A);
+      break;
+    case SapHessianFactorizationType::kBlockSparseCholesky:
+      factorization_ = std::make_unique<BlockSparseSuperNodalSolver>(
+          J->block_rows(), J->get_blocks(), *A);
+      break;
+    case SapHessianFactorizationType::kDense:
+      factorization_ = std::make_unique<DenseSuperNodalSolver>(A, J);
+      break;
+  }
+}
+
+std::unique_ptr<HessianFactorizationCache> HessianFactorizationCache::Clone()
+    const {
+  throw std::runtime_error(
+      "Attempting to clone an expensive Hessian factorization.");
+}
+
+void HessianFactorizationCache::UpdateWeightMatrixAndFactor(
+    const std::vector<MatrixX<double>>& G) {
+  DRAKE_DEMAND(!is_empty());
+  mutable_factorization()->SetWeightMatrix(G);
+  mutable_factorization()->Factor();
+}
+
+void HessianFactorizationCache::SolveInPlace(
+    EigenPtr<MatrixX<double>> rhs) const {
+  DRAKE_DEMAND(!is_empty());
+  // TODO(amcastro-tri): SuperNodalSolver should provide this in-place signature
+  // for multiple RHSs.
+  const int num_rhs = rhs->cols();
+  for (int i = 0; i < num_rhs; ++i) {
+    auto rhs_i = rhs->col(i);
+    // N.B. Unfortunately SolveInPlace() only accepts VectorXd*, while here the
+    // return of col() is a block object. Therefore, we are forced to make a
+    // copy.
+    // TODO(amcastro-tri): Update SuperNodalSolver::SolveInPlace() to accept
+    // EigenPtr instead.
+    rhs_i = factorization()->Solve(rhs_i);
+  }
+}
+
 template <typename T>
-SapModel<T>::SapModel(const SapContactProblem<T>* problem_ptr)
-    : problem_(problem_ptr) {
+SapModel<T>::SapModel(const SapContactProblem<T>* problem_ptr,
+                      SapHessianFactorizationType hessian_type)
+    : problem_(problem_ptr), hessian_type_(hessian_type) {
   // Graph to the original contact problem, including all cliques
   // (participating and non-participating).
   const ContactProblemGraph& graph = problem().graph();
@@ -137,6 +191,17 @@ void SapModel<T>::DeclareCacheEntries() {
       {system_->cache_entry_ticket(
           system_->cache_indexes().constraint_velocities)});
   system_->mutable_cache_indexes().hessian = hessian_cache_entry.cache_index();
+
+  // N.B. The default constructible SapHessianFactorizationCache is empty. Since
+  // the computation of the factorization is costly, we delay its construction
+  // to the very first time it is computed.
+  const auto& hessian_factorization_cache_entry = system_->DeclareCacheEntry(
+      "Hessian factorization cache.",
+      systems::ValueProducer(this, &SapModel<T>::CalcHessianFactorizationCache),
+      {system_->cache_entry_ticket(
+          system_->cache_indexes().constraint_velocities)});
+  system_->mutable_cache_indexes().hessian_factorization =
+      hessian_factorization_cache_entry.cache_index();
 }
 
 template <typename T>
@@ -256,6 +321,27 @@ void SapModel<T>::CalcHessianCache(const systems::Context<T>& context,
       bundle_data, &cache->impulses.gamma, &cache->G);
 }
 
+template <typename T>
+void SapModel<T>::CalcHessianFactorizationCache(
+    const systems::Context<T>&, HessianFactorizationCache*) const {
+  throw std::runtime_error(
+      "Hessian computation is only supported for T = double");
+}
+
+template <>
+void SapModel<double>::CalcHessianFactorizationCache(
+    const systems::Context<double>& context,
+    HessianFactorizationCache* hessian) const {
+  // Make only for the very first time. This can be an expensive computation for
+  // sparse Hessians even when the factorization is not yet computed.
+  if (hessian->is_empty()) {
+    *hessian = HessianFactorizationCache(hessian_type_, &dynamics_matrix(),
+                                         &constraints_bundle().J());
+  }
+  const std::vector<MatrixX<double>>& G = EvalConstraintsHessian(context);
+  hessian->UpdateWeightMatrixAndFactor(G);
+}
+
 template <typename T>
 int SapModel<T>::num_cliques() const {
   return problem().graph().participating_cliques().permuted_domain_size();