Gradient.h

Go to the documentation of this file.

 
#ifndef PBAT_FEM_GRADIENT_H
#define PBAT_FEM_GRADIENT_H
 
#include "Concepts.h"
#include "ShapeFunctions.h"
#include "pbat/Aliases.h"
#include "pbat/common/Eigen.h"
#include "pbat/profiling/Profiling.h"
 
#include <exception>
#include <fmt/core.h>
#include <tbb/parallel_for.h>
 
namespace pbat::fem {

template <
    CElement TElement,
    int Dims,
    class TDerivedE,
    class TDerivedeg,
    class TDerivedGNeg,
    class TDerivedIn,
    class TDerivedOut>
void GemmGradient(
    Eigen::DenseBase<TDerivedE> const& E,
    Eigen::Index nNodes,
    Eigen::DenseBase<TDerivedeg> const& eg,
    Eigen::MatrixBase<TDerivedGNeg> const& GNeg,
    Eigen::MatrixBase<TDerivedIn> const& X,
    Eigen::DenseBase<TDerivedOut>& Y);

template <CMesh TMesh, class TDerivedeg, class TDerivedGNeg, class TDerivedIn, class TDerivedOut>
inline void GemmGradient(
    TMesh const& mesh,
    Eigen::DenseBase<TDerivedeg> const& eg,
    Eigen::MatrixBase<TDerivedGNeg> const& GNeg,
    Eigen::MatrixBase<TDerivedIn> const& X,
    Eigen::DenseBase<TDerivedOut>& Y)
{
    GemmGradient<typename TMesh::ElementType, TMesh::kDims>(
        mesh.E,
        static_cast<typename TMesh::IndexType>(mesh.X.cols()),
        eg.derived(),
        GNeg.derived(),
        X.derived(),
        Y.derived());
}

template <
    CElement TElement,
    int Dims,
    Eigen::StorageOptions Options,
    class TDerivedE,
    class TDerivedeg,
    class TDerivedGNeg>
auto GradientMatrix(
    Eigen::DenseBase<TDerivedE> const& E,
    Eigen::Index nNodes,
    Eigen::DenseBase<TDerivedeg> const& eg,
    Eigen::MatrixBase<TDerivedGNeg> const& GNeg)
    -> Eigen::SparseMatrix<typename TDerivedGNeg::Scalar, Options, typename TDerivedE::Scalar>;

template <Eigen::StorageOptions Options, CMesh TMesh, class TDerivedeg, class TDerivedGNeg>
auto GradientMatrix(
    TMesh const& mesh,
    Eigen::DenseBase<TDerivedeg> const& eg,
    Eigen::MatrixBase<TDerivedGNeg> const& GNeg)
    -> Eigen::SparseMatrix<typename TDerivedGNeg::Scalar, Options, typename TMesh::IndexType>
{
    return GradientMatrix<typename TMesh::ElementType, TMesh::kDims, Options>(
        mesh.E,
        static_cast<typename TMesh::IndexType>(mesh.X.cols()),
        eg.derived(),
        GNeg.derived());
}
 
template <
    CElement TElement,
    int Dims,
    class TDerivedE,
    class TDerivedeg,
    class TDerivedGNeg,
    class TDerivedIn,
    class TDerivedOut>
inline void GemmGradient(
    Eigen::DenseBase<TDerivedE> const& E,
    Eigen::Index nNodes,
    Eigen::DenseBase<TDerivedeg> const& eg,
    Eigen::MatrixBase<TDerivedGNeg> const& GNeg,
    Eigen::MatrixBase<TDerivedIn> const& X,
    Eigen::DenseBase<TDerivedOut>& Y)
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.fem.GemmGradient");
    // Check inputs
    auto const rows = Dims * eg.size();
    auto const cols = nNodes;
    bool const bDimensionsMatch =
        (X.cols() == Y.cols()) and (X.rows() == cols) and (Y.rows() == rows);
    if (not bDimensionsMatch)
    {
        std::string const what = fmt::format(
            "Expected input to have rows={} and output to have rows={}, and same number of "
            "columns, but got dimensions "
            "X,Y=({} x {}), ({} x {})",
            cols,
            rows,
            X.rows(),
            X.cols(),
            Y.rows(),
            Y.cols());
        throw std::invalid_argument(what);
    }
    // Compute gradient
    auto constexpr kNodesPerElement = TElement::kNodes;
    auto constexpr kDims            = Dims;
    auto const nQuadPts             = eg.size();
    for (auto c = 0; c < X.cols(); ++c)
    {
        for (auto g = 0; g < nQuadPts; ++g)
        {
            auto const e     = eg(g);
            auto const nodes = E.col(e);
            auto const Xe    = X.col(c)(nodes);
            auto const Geg   = GNeg.template block<kNodesPerElement, kDims>(0, g * kDims);
            for (auto d = 0; d < kDims; ++d)
            {
                Y(d * nQuadPts + g) += Geg(Eigen::placeholders::all, d).transpose() * Xe;
            }
        }
    }
}
 
template <
    CElement TElement,
    int Dims,
    Eigen::StorageOptions Options,
    class TDerivedE,
    class TDerivedeg,
    class TDerivedGNeg>
auto GradientMatrix(
    Eigen::DenseBase<TDerivedE> const& E,
    Eigen::Index nNodes,
    Eigen::DenseBase<TDerivedeg> const& eg,
    Eigen::MatrixBase<TDerivedGNeg> const& GNeg)
    -> Eigen::SparseMatrix<typename TDerivedGNeg::Scalar, Options, typename TDerivedE::Scalar>
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.fem.GradientMatrix");
    using ScalarType       = typename TDerivedGNeg::Scalar;
    using IndexType        = typename TDerivedE::Scalar;
    using SparseMatrixType = Eigen::SparseMatrix<ScalarType, Options, IndexType>;
    using Triplet          = Eigen::Triplet<ScalarType, IndexType>;
 
    // Compile-time constants
    static_assert(
        TElement::kNodes != Eigen::Dynamic,
        "Element nodes must be known at compile time");
    auto constexpr kNodesPerElement = TElement::kNodes;
    auto constexpr kDims            = Dims;
 
    auto const nQuadPts = eg.size();
    std::vector<Triplet> triplets{};
    triplets.reserve(static_cast<std::size_t>(kNodesPerElement * kDims * nQuadPts));
    SparseMatrixType G(kDims * nQuadPts, nNodes);
 
    for (auto g = 0; g < nQuadPts; ++g)
    {
        auto const e     = eg(g);
        auto const nodes = E.col(e);
        auto const Geg   = GNeg.template block<kNodesPerElement, kDims>(0, g * kDims);
        for (auto d = 0; d < kDims; ++d)
        {
            for (auto j = 0; j < kNodesPerElement; ++j)
            {
                auto const ni = static_cast<IndexType>(d * nQuadPts + g);
                auto const nj = static_cast<IndexType>(nodes(j));
                triplets.emplace_back(ni, nj, Geg(j, d));
            }
        }
    }
    G.setFromTriplets(triplets.begin(), triplets.end());
    return G;
}
 
} // namespace pbat::fem
 
#endif // PBAT_FEM_GRADIENT_H