autodifffe.hh

Go to the documentation of this file.

// SPDX-FileCopyrightText: 2021-2024 The Ikarus Developers mueller@ibb.uni-stuttgart.de
// SPDX-License-Identifier: LGPL-3.0-or-later
 
#pragma once
 
#include <autodiff/forward/dual/dual.hpp>
#include <autodiff/forward/dual/eigen.hpp>
 
#include <ikarus/finiteelements/ferequirements.hh>
#include <ikarus/finiteelements/physicshelper.hh>
#include <ikarus/utils/traits.hh>
 
namespace Ikarus {
 
  template <typename RealFE_, typename FERequirementType_ = FERequirements<>, bool useEigenRef = false,
            bool forceAutoDiff = false>
  class AutoDiffFE : public RealFE_ {
  public:
    using RealFE            = RealFE_;                 
    using Basis             = typename RealFE::Basis;  
    using Traits            = TraitsFromFE<Basis, FERequirementType_, useEigenRef>;  
    using LocalView         = typename Traits::LocalView;                            
    using Element           = typename Traits::Element;                              
    using FERequirementType = typename Traits::FERequirementType;  
 
    void calculateMatrix(const FERequirementType& req, typename Traits::template MatrixType<> h) const {
      // real element implements calculateMatrix by itself, then we simply forward the call
      if constexpr (requires { RealFE::calculateMatrix(req, h); } and not forceAutoDiff) {
        RealFE::calculateMatrix(req, h);
      } else if constexpr (requires {
                             this->template calculateVectorImpl<autodiff::dual>(
                                 req, std::declval<typename Traits::template VectorType<autodiff::dual>>(),
                                 std::declval<const Eigen::VectorXdual&>());
                           }) {
        // real element implements calculateVector by itself, therefore we only need first order derivatives
        Eigen::VectorXdual dx(this->localView().size());
        Eigen::VectorXdual g(this->localView().size());
        dx.setZero();
        auto f = [&](auto& x) -> auto& {
          g.setZero();
          this->template calculateVectorImpl<autodiff::dual>(req, g, x);
          return g;
        };
        jacobian(f, autodiff::wrt(dx), at(dx), g, h);
      } else if constexpr (requires {
                             this->template calculateScalarImpl<autodiff::dual2nd>(
                                 req, std::declval<typename Traits::template VectorType<autodiff::dual2nd>>());
                           }) {
        // real element implements calculateScalar by itself, therefore we need second order derivatives
        Eigen::VectorXdual2nd dx(this->localView().size());
        Eigen::VectorXd g;
        autodiff::dual2nd e;
        dx.setZero();
        auto f = [&](auto& x) { return this->template calculateScalarImpl<autodiff::dual2nd>(req, x); };
        hessian(f, autodiff::wrt(dx), at(dx), e, g, h);
      } else
        static_assert(Dune::AlwaysFalse<AutoDiffFE>::value,
                      "Appropriate calculateScalarImpl or calculateVectorImpl functions are not implemented for the "
                      "chosen element.");
    }
 
    void calculateVector(const FERequirementType& req, typename Traits::template VectorType<> g) const {
      // real element implements calculateVector by itself, then we simply forward the call
      if constexpr (requires {
                      this->template calculateVectorImpl<double>(
                          req, std::declval<typename Traits::template VectorType<double>>(),
                          std::declval<const Eigen::VectorXd&>());
                    }
                    and not forceAutoDiff) {
        return this->template calculateVectorImpl<double>(req, g);
      } else if constexpr (requires {
                             this->template calculateScalarImpl<autodiff::dual>(
                                 req, std::declval<const Eigen::VectorXdual&>());
                           }) {
        // real element implements calculateScalar by itself, therefore we need first order derivatives
        Eigen::VectorXdual dx(this->localView().size());
        dx.setZero();
        autodiff::dual e;
        auto f = [&](auto& x) { return this->template calculateScalarImpl<autodiff::dual>(req, x); };
        gradient(f, autodiff::wrt(dx), at(dx), e, g);
      } else
        static_assert(Dune::AlwaysFalse<AutoDiffFE>::value,
                      "Appropriate calculateScalarImpl function is not implemented for the "
                      "chosen element.");
    }
 
    void calculateLocalSystem(const FERequirementType& req, typename Traits::template MatrixType<> h,
                              typename Traits::template VectorType<> g) const {
      Eigen::VectorXdual2nd dx(this->localView().size());
      dx.setZero();
      auto f = [&](auto& x) { return this->calculateScalarImpl(req, x); };
      hessian(f, autodiff::wrt(dx), at(dx), g, h);
    }
 
    [[nodiscard]] double calculateScalar(const FERequirementType& par) const {
      // real element implements calculateScalar by itself, then we simply forward the call
      if constexpr (requires { RealFE::calculateScalar(par); }) {
        return RealFE::calculateScalar(par);
      } else if constexpr (requires { this->calculateScalarImpl(par); }) {
        // real element only implements the protected calculateScalarImpl by itself, thus we call that one.
        return this->calculateScalarImpl(par);
      } else {
        static_assert(Dune::AlwaysFalse<AutoDiffFE>::value,
                      "Appropriate calculateScalar and calculateScalarImpl functions are not implemented for the "
                      "chosen element.");
      }
    }
 
    const RealFE& realFE() const { return *this; }
 
    template <typename... Args>
    explicit AutoDiffFE(Args&&... args) : RealFE{std::forward<Args>(args)...} {}
  };
}  // namespace Ikarus