HessianProduct.h

Go to the documentation of this file.

 
#ifndef PBAT_SIM_ALGORITHM_NEWTON_HESSIANPRODUCT_H
#define PBAT_SIM_ALGORITHM_NEWTON_HESSIANPRODUCT_H
 
#include "pbat/Aliases.h"
#include "pbat/common/ConstexprFor.h"
#include "pbat/math/linalg/SelectionMatrix.h"
#include "pbat/math/linalg/SparsityPattern.h"
 
#include <vector>
 
namespace pbat::sim::algorithm::newton {
struct HessianOperator;
} // namespace pbat::sim::algorithm::newton
 
namespace Eigen {
namespace internal {
template <>
struct traits<pbat::sim::algorithm::newton::HessianOperator>
    : public Eigen::internal::traits<pbat::CSCMatrix>
{
};
} // namespace internal
} // namespace Eigen
 
namespace pbat::sim::algorithm::newton {
template <class T>
concept CElasticPotential = requires(T U)
{
    { U.ComputeElementElasticity(std::declval<VectorX>(), false, true, true) };
    { U.ToMatrix(std::declval<CSCMatrix&>()) };
    { U.GH } -> std::convertible_to<math::linalg::SparsityPattern<>>;
};

struct Hessian
{
    static auto constexpr kDims = 3; 
    static auto constexpr kContactInds = 4; 
    using TripletType = Eigen::Triplet<Scalar, CSCMatrix::StorageIndex>; 
    void AllocateContacts(std::size_t nContacts);
    void SetSparsityPattern(CSCMatrix const& HS);
    template <class TDerivedF>
    void ImposeDirichletConstraints(Eigen::DenseBase<TDerivedF> const& F);
    template <class TDerivedHCB, class TDerivedHCIB>
    void AddContactHessianBlock(
        Eigen::DenseBase<TDerivedHCB> const& HCB,
        Eigen::DenseBase<TDerivedHCIB> const& HCIB);
    template <CElasticPotential TElasticPotential, class TDerivedM>
    void ConstructContactLessHessian(
        Eigen::DenseBase<TDerivedM> const& diagM,
        Scalar bt2,
        TElasticPotential const& U);
    void ConstructContactHessian();
 
    IndexVectorX diag; 
    CSCMatrix HNC;     
    std::vector<TripletType>
    HCij;
    CSCMatrix HC; 
    CSCMatrix S;
};

struct HessianOperator : public Eigen::EigenBase<HessianOperator>
{
    using Scalar       = pbat::Scalar;                     
    using RealScalar   = pbat::Scalar;                     
    using StorageIndex = typename CSCMatrix::StorageIndex; 
    enum
    {
        ColsAtCompileTime = Eigen::Dynamic,
        MaxColsAtCompileTime = Eigen::Dynamic,
        IsRowMajor = false
    };
 
    using SelfType = HessianOperator; 
    HessianOperator() = default;
    HessianOperator(Hessian* data);
    [[maybe_unused]] Eigen::Index rows() const { return mData->HNC.rows(); }
    [[maybe_unused]] Eigen::Index cols() const { return mData->HNC.cols(); }
    template <class Rhs>
    Eigen::Product<SelfType, Rhs, Eigen::AliasFreeProduct>
    operator*(Eigen::MatrixBase<Rhs> const& x) const
    {
        return Eigen::Product<SelfType, Rhs, Eigen::AliasFreeProduct>(*this, x.derived());
    }
 
    Hessian* mData; 
};
 
template <class TDerivedF>
inline void Hessian::ImposeDirichletConstraints(Eigen::DenseBase<TDerivedF> const& F)
{
    S   = math::linalg::SelectionMatrix(F, diag.size());
    HNC = S.transpose() * (HNC * S);
    HC  = S.transpose() * (HC * S);
}
 
template <class TDerivedHCB, class TDerivedHCIB>
inline void Hessian::AddContactHessianBlock(
    Eigen::DenseBase<TDerivedHCB> const& HCB,
    Eigen::DenseBase<TDerivedHCIB> const& HCIB)
{
    pbat::common::ForRange<0, kContactInds>(
        [&]<auto kj>() {
            pbat::common::ForRange<0, kContactInds>(
                [&]<auto ki>() {
                    auto const bi = HCIB(ki);
                    auto const bj = HCIB(kj);
                    auto const istart = bi * kDims;
                    auto const jstart = bj * kDims;
                    auto const HCbibj = HCB.template block<kDims, kDims>(ki, kj);
                    using IndexType = std::remove_cvref_t<decltype(std::declval<TripletType>().row()
                    )>;
                    pbat::common::ForRange<0, kDims>(
                        [&]<auto j>() {
                            pbat::common::ForRange<0, kDims>(
                                [&]<auto i>() {
                                    HCij.emplace_back(
                                        static_cast<IndexType>(istart + i),
                                        static_cast<IndexType>(jstart + j),
                                        HCbibj(i, j));
                                });
                        });
                });
        });
}
 
template <CElasticPotential TElasticPotential, class TDerivedM>
inline void Hessian::ConstructContactLessHessian(
    Eigen::DenseBase<TDerivedM> const& diagM,
    Scalar bt2,
    TElasticPotential const& U)
{
    U.ToMatrix(HNC);
    HNC *= bt2;
    HNC.coeffs()(diag) += diagM;
}
} // namespace pbat::sim::algorithm::newton
 
namespace Eigen {
namespace internal {
template <typename Rhs, int ProductType>
struct generic_product_impl<
        pbat::sim::algorithm::newton::HessianOperator,
        Rhs,
        SparseShape,
        DenseShape,
        ProductType>
    : generic_product_impl_base<
        pbat::sim::algorithm::newton::HessianOperator,
        Rhs,
        generic_product_impl<pbat::sim::algorithm::newton::HessianOperator, Rhs>>
{
    using HessianOperator = pbat::sim::algorithm::newton::HessianOperator;
    using Scalar          = typename Product<HessianOperator, Rhs>::Scalar;
    using Hessian         = pbat::sim::algorithm::newton::Hessian;
    template <typename Dst>
    static void scaleAndAddTo(
        Eigen::MatrixBase<Dst>& dst,
        HessianOperator const& lhs,
        Eigen::MatrixBase<Rhs> const& rhs,
        Scalar const& alpha)
    {
        dst += alpha * (lhs.mData->HNC * rhs + lhs.mData->HC * rhs);
    }
};
} // namespace internal
} // namespace Eigen
 
#endif // PBAT_SIM_ALGORITHM_NEWTON_HESSIANPRODUCT_H