MeshVertexTetrahedronDcd.h

#ifndef PBAT_SIM_CONTACT_MESHVERTEXTETRAHEDRONDCD_H
#define PBAT_SIM_CONTACT_MESHVERTEXTETRAHEDRONDCD_H
 
#include "MultibodyTetrahedralMeshSystem.h"
#include "PhysicsBasedAnimationToolkitExport.h"
#include "pbat/Aliases.h"
#include "pbat/geometry/AabbKdTreeHierarchy.h"
#include "pbat/geometry/AabbRadixTreeHierarchy.h"
#include "pbat/geometry/DistanceQueries.h"
#include "pbat/geometry/HierarchicalHashGrid.h"
#include "pbat/geometry/OverlapQueries.h"
#include "pbat/math/linalg/mini/Eigen.h"
 
#include <exception>
#include <unordered_map>
#include <vector>
 
namespace pbat::sim::contact {

class MeshVertexTetrahedronDcd
{
  public:
    using ScalarType = Scalar; 
    using IndexType  = Index;  
    static constexpr int kMaxVertexTriangleContacts =
        8;                          
    static constexpr int kDims = 3; 

    PBAT_API MeshVertexTetrahedronDcd(
        Eigen::Ref<Eigen::Matrix<ScalarType, kDims, Eigen::Dynamic>> X,
        Eigen::Ref<Eigen::Matrix<IndexType, 4, Eigen::Dynamic>> T);
    PBAT_API void UpdateActiveSet(
        Eigen::Ref<Eigen::Matrix<ScalarType, kDims, Eigen::Dynamic> const> const& X,
        Eigen::Ref<Eigen::Matrix<IndexType, 4, Eigen::Dynamic> const> const& T);
    template <class FOnVertexTriangleContactPair>
    void ForEachActiveVertexTriangleContact(
        IndexType v,
        FOnVertexTriangleContactPair fOnVertexTriangleContactPair) const;
    PBAT_API MultibodyTetrahedralMeshSystem<IndexType> const& MultibodySystem() const;
 
  protected:
    template <class TDerivedX>
    void ComputeTriangleAabbs(Eigen::DenseBase<TDerivedX> const& X);
    template <class TDerivedX, class TDerivedT>
    void ComputeTetrahedronMeshAabbs(
        IndexType o,
        Eigen::DenseBase<TDerivedX> const& X,
        Eigen::DenseBase<TDerivedT> const& T);
    PBAT_API void ComputeBodyAabbs();
    PBAT_API void UpdateMeshTriangleBvhs();
    PBAT_API void UpdateMeshTetrahedronHashGrid(IndexType bodyIndex);
    PBAT_API void RecomputeBodiesBvh();
    template <class FOnBodyPair>
    void ForEachBodyPair(FOnBodyPair&& fOnBodyPair) const;
    template <class TDerivedX, class TDerivedT>
    void MarkPenetratingVertices(
        IndexType ov,
        IndexType ot,
        Eigen::DenseBase<TDerivedX> const& X,
        Eigen::DenseBase<TDerivedT> const& T);
    template <class TDerivedX>
    void FindNearestTrianglesToPenetratingVertices(
        IndexType ov,
        IndexType of,
        Eigen::DenseBase<TDerivedX> const& X);
 
  private:
    sim::contact::MultibodyTetrahedralMeshSystem<IndexType>
        mMultibodySystem; 
 
    Eigen::Matrix<ScalarType, 2 * kDims, Eigen::Dynamic>
        mTriangleAabbs; 
    Eigen::Matrix<ScalarType, 2 * kDims, Eigen::Dynamic>
        mTetrahedronAabbs; 
 
    Eigen::Matrix<ScalarType, 2 * kDims, Eigen::Dynamic>
        mBodyAabbs; 
 
    std::vector<geometry::AabbKdTreeHierarchy<kDims>>
        mTriangleMeshBvhs; 
    std::vector<geometry::HierarchicalHashGrid<kDims, ScalarType, IndexType>>
        mTetrahedronMeshHashGrids; 
 
    geometry::AabbRadixTreeHierarchy<kDims> mBodiesBvh; 
 
    Eigen::Vector<bool, Eigen::Dynamic>
        mIsTetrahedronMeshHashGridDirty; 
    std::vector<std::unordered_map<IndexType, IndexType>>
        mBodyVertexToOtherBodiesValence; 
    Eigen::Matrix<IndexType, kMaxVertexTriangleContacts, Eigen::Dynamic>
        mVFC; 
};
 
template <class FOnVertexTriangleContactPair>
inline void MeshVertexTetrahedronDcd::ForEachActiveVertexTriangleContact(
    IndexType v,
    FOnVertexTriangleContactPair fOnVertexTriangleContactPair) const
{
    for (auto k = 0; k < kMaxVertexTriangleContacts; ++k)
    {
        IndexType const f = mVFC(k, v);
        if (f == -1)
            break;
        fOnVertexTriangleContactPair(f, k);
    }
}
 
template <class TDerivedX>
inline void MeshVertexTetrahedronDcd::ComputeTriangleAabbs(Eigen::DenseBase<TDerivedX> const& X)
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.sim.contact.MeshVertexTetrahedronDcd.ComputeTriangleAabbs");
    auto const nTriangles = mMultibodySystem.NumContactTriangles();
    for (auto f = 0; f < nTriangles; ++f)
    {
        Matrix<kDims, 3> XF;
        XF = X(Eigen::placeholders::all, mMultibodySystem.F.col(f)).template block<kDims, 3>(0, 0);
        auto L = mTriangleAabbs.col(f).head<kDims>();
        auto U = mTriangleAabbs.col(f).tail<kDims>();
        L      = XF.rowwise().minCoeff();
        U      = XF.rowwise().maxCoeff();
    }
}
 
template <class TDerivedX, class TDerivedT>
inline void MeshVertexTetrahedronDcd::ComputeTetrahedronMeshAabbs(
    IndexType o,
    Eigen::DenseBase<TDerivedX> const& X,
    Eigen::DenseBase<TDerivedT> const& T)
{
    PBAT_PROFILE_NAMED_SCOPE(
        "pbat.sim.contact.MeshVertexTetrahedronDcd.ComputeTetrahedronMeshAabbs");
    auto const [tbegin, tend] = mMultibodySystem.TetrahedraRangeFor(o);
    auto TB = mTetrahedronAabbs(Eigen::placeholders::all, Eigen::seq(tbegin, tend - 1));
    for (auto t = tbegin; t < tend; ++t)
    {
        Matrix<kDims, 4> XT = X(Eigen::placeholders::all, T.col(t)).template block<kDims, 4>(0, 0);
        auto L              = TB.col(t - tbegin).template head<kDims>();
        auto U              = TB.col(t - tbegin).template tail<kDims>();
        L                   = XT.rowwise().minCoeff();
        U                   = XT.rowwise().maxCoeff();
    }
}
 
template <class FOnBodyPair>
inline void MeshVertexTetrahedronDcd::ForEachBodyPair(FOnBodyPair&& fOnBodyPair) const
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.sim.contact.MeshVertexTetrahedronDcd.ForEachBodyPair");
    mBodiesBvh.SelfOverlaps(
        [this](IndexType oi, IndexType oj) {
            if (oi == oj)
            {
                throw std::runtime_error(
                    "MeshVertexTetrahedronDcd::UpdateActiveSet: oi == oj, which is not allowed.");
            }
            auto L1 = mBodyAabbs.col(oi).template head<kDims>();
            auto U1 = mBodyAabbs.col(oi).template tail<kDims>();
            auto L2 = mBodyAabbs.col(oj).template head<kDims>();
            auto U2 = mBodyAabbs.col(oj).template tail<kDims>();
            using math::linalg::mini::FromEigen;
            return geometry::OverlapQueries::AxisAlignedBoundingBoxes(
                FromEigen(L1),
                FromEigen(U1),
                FromEigen(L2),
                FromEigen(U2));
        },
        [fOnBodyPair = std::forward<FOnBodyPair>(
             fOnBodyPair)](IndexType oi, IndexType oj, [[maybe_unused]] IndexType k) {
            fOnBodyPair(oi, oj);
        });
}
 
template <class TDerivedX, class TDerivedT>
inline void MeshVertexTetrahedronDcd::MarkPenetratingVertices(
    IndexType ov,
    IndexType ot,
    Eigen::DenseBase<TDerivedX> const& X,
    Eigen::DenseBase<TDerivedT> const& T)
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.sim.contact.MeshVertexTetrahedronDcd.MarkPenetratingVertices");
    auto const Vov            = mMultibodySystem.ContactVerticesOf(ov);
    auto Xv                   = X(Eigen::placeholders::all, Vov);
    auto const [tbegin, tend] = mMultibodySystem.TetrahedraRangeFor(ot);
    auto Lt =
        mTetrahedronAabbs(Eigen::placeholders::all, Eigen::seq(tbegin, tend - 1)).topRows<kDims>();
    auto Ut = mTetrahedronAabbs(Eigen::placeholders::all, Eigen::seq(tbegin, tend - 1))
                  .bottomRows<kDims>();
    std::unordered_map<IndexType, IndexType>& mBodyVertexToOtherBodiesValenceMap =
        mBodyVertexToOtherBodiesValence[static_cast<std::size_t>(ov)];
    mTetrahedronMeshHashGrids[static_cast<std::size_t>(ot)]
        .BroadPhase(Lt, Ut, Xv, [&](IndexType vq, IndexType tp) {
            auto const i  = Vov(vq);
            auto const XT = X(Eigen::placeholders::all, T.col(tbegin + tp));
            using math::linalg::mini::FromEigen;
            if (geometry::OverlapQueries::PointTetrahedron3D(
                    FromEigen(X.col(i).template head<kDims>()),
                    FromEigen(XT.col(0).template head<kDims>()),
                    FromEigen(XT.col(1).template head<kDims>()),
                    FromEigen(XT.col(2).template head<kDims>()),
                    FromEigen(XT.col(3).template head<kDims>())))
            {
                auto it  = mBodyVertexToOtherBodiesValenceMap.find(vq);
                auto end = mBodyVertexToOtherBodiesValenceMap.end();
                if (it == end)
                {
                    mBodyVertexToOtherBodiesValenceMap.insert({vq, IndexType(1)});
                }
                else
                {
                    auto& valence = it->second;
                    ++valence;
                }
            }
        });
}
 
template <class TDerivedX>
inline void MeshVertexTetrahedronDcd::FindNearestTrianglesToPenetratingVertices(
    IndexType ov,
    IndexType of,
    Eigen::DenseBase<TDerivedX> const& X)
{
    PBAT_PROFILE_NAMED_SCOPE(
        "pbat.sim.contact.MeshVertexTetrahedronDcd.FindNearestTrianglesToPenetratingVertices");
    IndexType const vbegin = mMultibodySystem.ContactVerticesRangeFor(ov).first;
    auto const fbegin      = mMultibodySystem.ContactTrianglesRangeFor(of).first;
    std::unordered_map<IndexType, IndexType>& mBodyVertexToOtherBodiesValenceMap =
        mBodyVertexToOtherBodiesValence[static_cast<std::size_t>(ov)];
    for (auto& [vq, valence] : mBodyVertexToOtherBodiesValenceMap)
    {
        assert(valence > 0);
        --valence;
        auto const v = vbegin + vq;
        auto const i = mMultibodySystem.V(v);
        auto const P = X.col(i).template head<kDims>();
        using math::linalg::mini::FromEigen;
        mTriangleMeshBvhs[static_cast<std::size_t>(of)].NearestNeighbours(
            [&]<class TL, class TU>(TL const& L, TU const& U) {
                return geometry::DistanceQueries::PointAxisAlignedBoundingBox(
                    FromEigen(P),
                    FromEigen(L),
                    FromEigen(U));
            },
            [&](IndexType fq) { 
                auto const f = fbegin + fq;
                auto const iF = mMultibodySystem.F.col(f);
                auto const XT = X(Eigen::placeholders::all, iF);
                return geometry::DistanceQueries::PointTriangle(
                    FromEigen(P),
                    FromEigen(XT.col(0).template head<kDims>()),
                    FromEigen(XT.col(1).template head<kDims>()),
                    FromEigen(XT.col(2).template head<kDims>()));
            },
            [&](IndexType fq, [[maybe_unused]] ScalarType d, IndexType k) {
                if (k >= kMaxVertexTriangleContacts)
                    return;
                mVFC(k, v) = fbegin + fq;
            }/*,
            radius,
            eps*/);
    }
    std::erase_if(mBodyVertexToOtherBodiesValenceMap, [](auto const& keyValuePair) {
        auto const& valence = keyValuePair.second;
        assert(valence >= 0);
        return valence == 0;
    });
}
 
} // namespace pbat::sim::contact
 
#endif // PBAT_SIM_CONTACT_MESHVERTEXTETRAHEDRONDCD_H