Skip to content

Commit

Permalink
SapModel caches the Hessian factorization.
Browse files Browse the repository at this point in the history
  • Loading branch information
amcastro-tri committed May 12, 2024
1 parent 2b46e26 commit 0c829e5
Show file tree
Hide file tree
Showing 6 changed files with 360 additions and 3 deletions.
14 changes: 14 additions & 0 deletions multibody/contact_solvers/sap/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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"],
Expand Down Expand Up @@ -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",
],
Expand Down
73 changes: 73 additions & 0 deletions multibody/contact_solvers/sap/dense_supernodal_solver.cc
Original file line number Diff line number Diff line change
@@ -0,0 +1,73 @@
#include "drake/multibody/contact_solvers/sap/dense_supernodal_solver.h"

#include <vector>

namespace drake {
namespace multibody {
namespace contact_solvers {
namespace internal {

DenseSuperNodalSolver::DenseSuperNodalSolver(
const std::vector<MatrixX<double>>* A, const BlockSparseMatrix<double>* J)
: A_(A), J_(J) {
DRAKE_THROW_UNLESS(A != nullptr);
DRAKE_THROW_UNLESS(J != nullptr);
const int nv = [A]() {
int size = 0;
for (const auto& Ai : *A) {
DRAKE_THROW_UNLESS(Ai.rows() == Ai.cols());
size += Ai.rows();
}
return size;
}();
H_.resize(nv, nv);
}

bool DenseSuperNodalSolver::DoSetWeightMatrix(
const std::vector<Eigen::MatrixXd>& G) {
const int nv = size();

// 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 Hessian 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;
}

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
58 changes: 58 additions & 0 deletions multibody/contact_solvers/sap/dense_supernodal_solver.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,58 @@
#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::runtime_error if A or J is nullptr.
// @throws std::runtime_error 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);

// @returns the size of the system.
int size() const { return H_.rows(); }

private:
// Implementations of SuperNodalSolver NVIs.
bool DoSetWeightMatrix(
const std::vector<Eigen::MatrixXd>& block_diagonal_G) final;
Eigen::MatrixXd DoMakeFullMatrix() const final { return H_; }
bool DoFactor() final;
void DoSolveInPlace(Eigen::VectorXd* b) const final;
int DoGetSize() const final { return size(); }

const std::vector<MatrixX<double>>* A_{nullptr};
const BlockSparseMatrix<double>* J_{nullptr};
// 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<MatrixX<double>> Hldlt_{H_};
};

} // namespace internal
} // namespace contact_solvers
} // namespace multibody
} // namespace drake
86 changes: 84 additions & 2 deletions multibody/contact_solvers/sap/sap_model.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -15,9 +18,56 @@ 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.");
return nullptr;
}

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);
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();
Expand Down Expand Up @@ -137,6 +187,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 dealy its construction
// to the very first time 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>
Expand Down Expand Up @@ -256,6 +317,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();
Expand Down
Loading

0 comments on commit 0c829e5

Please sign in to comment.