Jacobian.h

Go to the documentation of this file.

 
#ifndef PBAT_FEM_JACOBIAN_H
#define PBAT_FEM_JACOBIAN_H
 
#include "Concepts.h"
#include "pbat/common/Concepts.h"
 
#include <Eigen/Cholesky>
#include <Eigen/LU>
#include <Eigen/SVD>
#include <fmt/core.h>
#include <pbat/Aliases.h>
#include <pbat/common/Eigen.h>
#include <pbat/profiling/Profiling.h>
#include <string>
#include <tbb/parallel_for.h>
 
namespace pbat::fem {

template <
    CElement TElement,
    class TDerivedX,
    class TDerivedx,
    common::CFloatingPoint TScalar = typename TDerivedX::Scalar>
[[maybe_unused]] auto
Jacobian(Eigen::MatrixBase<TDerivedX> const& X, Eigen::MatrixBase<TDerivedx> const& x)
    -> Eigen::Matrix<TScalar, TDerivedx::RowsAtCompileTime, TElement::kDims>
{
    static_assert(
        TDerivedx::ColsAtCompileTime == TElement::kNodes,
        "x must have the same number of columns as the number of nodes in the element.");
    auto constexpr kDimsOut                                   = TDerivedx::RowsAtCompileTime;
    Eigen::Matrix<TScalar, kDimsOut, TElement::kDims> const J = x * TElement::GradN(X);
    return J;
}

template <class TDerived>
[[maybe_unused]] auto DeterminantOfJacobian(Eigen::MatrixBase<TDerived> const& J) ->
    typename TDerived::Scalar
{
    using ScalarType         = typename TDerived::Scalar;
    bool const bIsSquare     = J.rows() == J.cols();
    auto const detJ          = bIsSquare ? J.determinant() : J.jacobiSvd().singularValues().prod();
    ScalarType constexpr eps = ScalarType(1e-10);
    if (detJ <= eps)
    {
        throw std::invalid_argument("Inverted or singular jacobian");
    }
    return detJ;
}

template <CElement TElement, int QuadratureOrder, class TDerivedE, class TDerivedX>
[[maybe_unused]] auto
DeterminantOfJacobian(Eigen::DenseBase<TDerivedE> const& E, Eigen::MatrixBase<TDerivedX> const& X)
    -> Eigen::Matrix<
        typename TDerivedX::Scalar,
        TElement::template QuadratureType<QuadratureOrder, typename TDerivedX::Scalar>::kPoints,
        Eigen::Dynamic>
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.fem.DeterminantOfJacobian");
    using ScalarType        = typename TDerivedX::Scalar;
    using ElementType       = TElement;
    using AffineElementType = typename ElementType::AffineBaseType;
    using QuadratureRuleType =
        typename ElementType::template QuadratureType<QuadratureOrder, ScalarType>;
    using MatrixType = Eigen::Matrix<ScalarType, QuadratureRuleType::kPoints, Eigen::Dynamic>;
    if (X.rows() < ElementType::kDims)
    {
        throw std::invalid_argument(
            fmt::format(
                "Invalid number of dimensions, expected at least {}, but got {}",
                ElementType::kDims,
                X.rows()));
    }
    auto const Xg = common::ToEigen(QuadratureRuleType::points)
                        .reshaped(QuadratureRuleType::kDims + 1, QuadratureRuleType::kPoints)
                        .template bottomRows<ElementType::kDims>();
    auto const numberOfElements = E.cols();
    MatrixType detJe(QuadratureRuleType::kPoints, numberOfElements);
    tbb::parallel_for(Index{0}, Index{numberOfElements}, [&](Index e) {
        auto const nodes                                   = E.col(e);
        auto const vertices                                = nodes(ElementType::Vertices);
        auto constexpr kRowsJ                              = TDerivedX::RowsAtCompileTime;
        auto constexpr kColsJ                              = AffineElementType::kNodes;
        Eigen::Matrix<ScalarType, kRowsJ, kColsJ> const Ve = X(Eigen::placeholders::all, vertices);
        if constexpr (AffineElementType::bHasConstantJacobian)
        {
            ScalarType const detJ = DeterminantOfJacobian(
                Jacobian<AffineElementType>(Xg.col(0) /*arbitrary eval.pt.*/, Ve));
            detJe.col(e).setConstant(detJ);
        }
        else
        {
            for (auto g = 0; g < QuadratureRuleType::kPoints; ++g)
            {
                auto const detJ = DeterminantOfJacobian(Jacobian<AffineElementType>(Xg.col(g), Ve));
                detJe(g, e)     = detJ;
            }
        }
    });
    return detJe;
}

template <int QuadratureOrder, CMesh TMesh>
[[maybe_unused]] auto DeterminantOfJacobian(TMesh const& mesh) -> Eigen::Matrix<
    typename TMesh::ScalarType,
    TMesh::ElementType::template QuadratureType<QuadratureOrder, typename TMesh::ScalarType>::
        kPoints,
    Eigen::Dynamic>
{
    return DeterminantOfJacobian<typename TMesh::ElementType, QuadratureOrder>(mesh.E, mesh.X);
}

template <CElement TElement, class TDerivedEg, class TDerivedX, class TDerivedXi>
[[maybe_unused]] auto DeterminantOfJacobianAt(
    Eigen::DenseBase<TDerivedEg> const& Eg,
    Eigen::MatrixBase<TDerivedX> const& X,
    Eigen::MatrixBase<TDerivedXi> const& Xi)
    -> Eigen::Vector<typename TDerivedX::Scalar, Eigen::Dynamic>
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.fem.DeterminantOfJacobian");
    using ScalarType        = typename TDerivedX::Scalar;
    using ElementType       = TElement;
    using AffineElementType = typename ElementType::AffineBaseType;
    Eigen::Vector<ScalarType, Eigen::Dynamic> detJe(Eg.cols());
    tbb::parallel_for(Index{0}, Index{Eg.cols()}, [&](Index g) {
        auto const nodes                                   = Eg.col(g);
        auto const vertices                                = nodes(ElementType::Vertices);
        auto constexpr kRowsJ                              = TDerivedX::RowsAtCompileTime;
        auto constexpr kColsJ                              = AffineElementType::kNodes;
        Eigen::Matrix<ScalarType, kRowsJ, kColsJ> const Ve = X(Eigen::placeholders::all, vertices);
        ScalarType const detJ = DeterminantOfJacobian(Jacobian<AffineElementType>(Xi.col(g), Ve));
        detJe(g)              = detJ;
    });
    return detJe;
}

template <CElement TElement, class TDerivedE, class TDerivedX, class TDerivedeg, class TDerivedXi>
[[maybe_unused]] auto DeterminantOfJacobianAt(
    Eigen::DenseBase<TDerivedE> const& E,
    Eigen::MatrixBase<TDerivedX> const& X,
    Eigen::MatrixBase<TDerivedeg> const& eg,
    Eigen::MatrixBase<TDerivedXi> const& Xi)
    -> Eigen::Vector<typename TDerivedX::Scalar, Eigen::Dynamic>
{
    return DeterminantOfJacobianAt<TElement>(
        E(Eigen::placeholders::all, eg.derived()),
        X.derived(),
        Xi.derived());
}

template <
    CElement TElement,
    class TDerivedX,
    class TDerivedx,
    common::CFloatingPoint TScalar = typename TDerivedX::Scalar>
auto ReferencePosition(
    Eigen::MatrixBase<TDerivedX> const& X,
    Eigen::MatrixBase<TDerivedx> const& x,
    int const maxIterations = 5,
    TScalar const eps       = 1e-10) -> Eigen::Vector<TScalar, TElement::kDims>
{
    using ScalarType = TScalar;
    static_assert(
        std::is_same_v<ScalarType, typename TDerivedX::Scalar>,
        "X and x must have given scalar type");
    static_assert(
        std::is_same_v<ScalarType, typename TDerivedx::Scalar>,
        "X and x must have given scalar type");
    using ElementType = TElement;
    assert(x.cols() == ElementType::kNodes);
    auto constexpr kDims    = ElementType::kDims;
    auto constexpr kNodes   = ElementType::kNodes;
    auto constexpr kOutDims = Eigen::MatrixBase<TDerivedx>::RowsAtCompileTime;
    // Initial guess is element's barycenter.
    auto const coords = common::ToEigen(ElementType::Coordinates).reshaped(kDims, kNodes);
    auto const vertexLagrangePositions =
        (coords(Eigen::placeholders::all, ElementType::Vertices).template cast<ScalarType>() /
         static_cast<ScalarType>(ElementType::kOrder));
    Eigen::Vector<ScalarType, kDims> Xik = vertexLagrangePositions.rowwise().sum() /
                                           static_cast<ScalarType>(ElementType::Vertices.size());
 
    Eigen::Vector<ScalarType, kOutDims> rk = x * ElementType::N(Xik) - X;
    Eigen::Matrix<ScalarType, kOutDims, kDims> J;
    Eigen::LDLT<Eigen::Matrix<ScalarType, kDims, kDims>> LDLT;
 
    // If Jacobian is constant, let's precompute it
    if constexpr (ElementType::bHasConstantJacobian)
    {
        J = Jacobian<ElementType>(Xik, x);
        LDLT.compute(J.transpose() * J);
    }
    // Do up to maxIterations Gauss Newton iterations
    for (auto k = 0; k < maxIterations; ++k)
    {
        if (rk.template lpNorm<1>() <= eps)
            break;
 
        // Non-constant jacobians need to be updated
        if constexpr (not ElementType::bHasConstantJacobian)
        {
            J = Jacobian<ElementType>(Xik, x);
            LDLT.compute(J.transpose() * J);
        }
 
        Xik -= LDLT.solve(J.transpose() * rk);
        rk = x * ElementType::N(Xik) - X;
    }
    return Xik;
}

template <
    CElement TElement,
    class TDerivedEg,
    class TDerivedX,
    class TDerivedXg,
    common::CFloatingPoint TScalar = typename TDerivedXg::Scalar>
auto ReferencePositions(
    Eigen::DenseBase<TDerivedEg> const& Eg,
    Eigen::MatrixBase<TDerivedX> const& X,
    Eigen::MatrixBase<TDerivedXg> const& Xg,
    int maxIterations = 5,
    TScalar eps       = 1e-10) -> Eigen::Matrix<TScalar, TElement::kDims, Eigen::Dynamic>
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.fem.ReferencePositions");
    static_assert(
        std::is_same_v<TScalar, typename TDerivedX::Scalar>,
        "X must have given scalar type");
    static_assert(
        std::is_same_v<TScalar, typename TDerivedXg::Scalar>,
        "Xg must have given scalar type");
    using ElementType = TElement;
    using ScalarType  = TScalar;
    using MatrixType  = Eigen::Matrix<ScalarType, ElementType::kDims, Eigen::Dynamic>;
    MatrixType Xi(ElementType::kDims, Eg.cols());
    tbb::parallel_for(Index{0}, Index{Eg.cols()}, [&](Index g) {
        auto const nodes                                     = Eg.col(g);
        auto constexpr kOutDims                              = TDerivedX::RowsAtCompileTime;
        auto constexpr kNodes                                = ElementType::kNodes;
        Eigen::Matrix<ScalarType, kOutDims, kNodes> const Xe = X(Eigen::placeholders::all, nodes);
        Eigen::Vector<ScalarType, kOutDims> const Xk         = Xg.col(g);
        Xi.col(g) = ReferencePosition<ElementType>(Xk, Xe, maxIterations, eps);
    });
    return Xi;
}

template <CElement TElement, class TDerivedE, class TDerivedX, class TDerivedEg, class TDerivedXg>
auto ReferencePositions(
    Eigen::DenseBase<TDerivedE> const& E,
    Eigen::MatrixBase<TDerivedX> const& X,
    Eigen::DenseBase<TDerivedEg> const& eg,
    Eigen::MatrixBase<TDerivedXg> const& Xg,
    int maxIterations               = 5,
    typename TDerivedXg::Scalar eps = 1e-10)
    -> Eigen::Matrix<typename TDerivedXg::Scalar, TElement::kDims, Eigen::Dynamic>
{
    return ReferencePositions<TElement>(
        E(Eigen::placeholders::all, eg.derived()),
        X.derived(),
        Xg.derived(),
        maxIterations,
        eps);
}

template <CMesh TMesh, class TDerivedEg, class TDerivedXg>
auto ReferencePositions(
    TMesh const& mesh,
    Eigen::DenseBase<TDerivedEg> const& eg,
    Eigen::MatrixBase<TDerivedXg> const& Xg,
    int maxIterations              = 5,
    typename TMesh::ScalarType eps = 1e-10)
    -> Eigen::Matrix<typename TMesh::ScalarType, TMesh::ElementType::kDims, Eigen::Dynamic>
{
    return ReferencePositions<typename TMesh::ElementType>(
        mesh.E(Eigen::placeholders::all, eg.derived()),
        mesh.X,
        Xg.derived(),
        maxIterations,
        eps);
}
 
} // namespace pbat::fem
 
#endif // PBAT_FEM_JACOBIAN_H