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/utils/traits.hh>
 
namespace Ikarus {
 
template <typename FEImpl, bool forceAutoDiff = false>
class AutoDiffFE : public FEImpl
{
public:
  using RealFE       = FEImpl;                        
  using BasisHandler = typename RealFE::BasisHandler; 
  using Traits       = typename RealFE::Traits;       
  using LocalView    = typename Traits::LocalView;    
  using Element      = typename Traits::Element;      
  using Requirement  = typename RealFE::Requirement;  
private:
  using Mixin = FEImpl::Mixin;
 
public:
  friend void calculateMatrix(const AutoDiffFE& self, const Requirement& par, const MatrixAffordance& affordance,
                              typename Traits::template MatrixType<> h) {
    self.calculateMatrix(par, affordance, h);
  }
 
  friend void calculateVector(const AutoDiffFE& self, const Requirement& par, VectorAffordance affordance,
                              typename Traits::template VectorType<double> g) {
    self.calculateVector(par, affordance, g);
  }
 
  friend void calculateLocalSystem(const AutoDiffFE& self, const Requirement& par, const MatrixAffordance& affordanceM,
                                   VectorAffordance affordanceV, typename Traits::template MatrixType<> h,
                                   typename Traits::template VectorType<> g) {
    self.calculateLocalSystem(par, affordanceM, affordanceV, h, g);
  }
 
  friend auto calculateScalar(const AutoDiffFE& self, const Requirement& par, ScalarAffordance affordance) {
    return self.calculateScalar(par, affordance);
  }
 
  const RealFE& realFE() const { return *this; }
 
  template <typename... Args>
  explicit AutoDiffFE(Args&&... args)
      : RealFE{std::forward<Args>(args)...} {}
 
private:
  void calculateMatrix(const Requirement& req, const MatrixAffordance& affordance,
                       typename Traits::template MatrixType<> h) const {
    // real element implements calculateMatrix by itself, then we simply forward the call
 
    if constexpr (requires(Eigen::VectorXd v) {
                    static_cast<const Mixin&>(std::declval<AutoDiffFE>())
                        .template calculateMatrixImpl<double>(req, affordance, h, v);
                  } and not forceAutoDiff) {
      return Mixin::template calculateMatrixImpl<double>(req, affordance, h);
    } else if constexpr (requires(Eigen::VectorXdual v) {
                           static_cast<const Mixin&>(std::declval<AutoDiffFE>())
                               .template calculateVectorImpl<autodiff::dual>(
                                   req, vectorAffordance(affordance),
                                   std::declval<typename Traits::template VectorType<autodiff::dual>>(), v);
                         }) {
      // 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 = [this, &req, affordance, &g](auto& x) -> auto& {
        // Since req is const as a function argument, we can not make this lambda capture by mutable reference
        // But we have to do this since for efficiency reason we reuse the g vector
        // therefore, the only remaining option is to cast the const away from g
        Eigen::VectorXdual& gref = const_cast<Eigen::VectorXdual&>(g);
        gref.setZero();
        Mixin::template calculateVectorImpl<autodiff::dual>(req, vectorAffordance(affordance), gref, x);
        return g;
      };
      jacobian(f, autodiff::wrt(dx), at(dx), g, h);
    } else if constexpr (requires(typename Traits::template VectorType<autodiff::dual2nd> v) {
                           static_cast<const Mixin&>(std::declval<AutoDiffFE>())
                               .template calculateScalarImpl<autodiff::dual2nd>(req, scalarAffordance(affordance), v);
                         }) {
      // 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 = [this, &req, affordance](auto& x) {
        return Mixin::template calculateScalarImpl<autodiff::dual2nd>(req, scalarAffordance(affordance), 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 Requirement& req, VectorAffordance affordance,
                       typename Traits::template VectorType<> g) const {
    // real element implements calculateVector by itself, then we simply forward the call
    if constexpr (requires {
                    static_cast<const Mixin&>(std::declval<AutoDiffFE>())
                        .template calculateVectorImpl<double>(
                            req, affordance, std::declval<typename Traits::template VectorType<double>>(),
                            std::declval<const Eigen::VectorXd&>());
                  } and not forceAutoDiff) {
      return Mixin::template calculateVectorImpl<double>(req, affordance, g);
    } else if constexpr (requires {
                           static_cast<const Mixin&>(std::declval<AutoDiffFE>())
                               .template calculateScalarImpl<autodiff::dual>(req, scalarAffordance(affordance),
                                                                             std::declval<const Eigen::VectorXdual&>());
                         }) {
      // real element implements calculateScalar by itself but no calculateVectorImpl, therefore we need first order
      // derivatives
      Eigen::VectorXdual dx(this->localView().size());
      dx.setZero();
      autodiff::dual e;
      auto f = [this, &req, affordance](auto& x) {
        return Mixin::template calculateScalarImpl<autodiff::dual>(req, scalarAffordance(affordance), 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.");
  }
 
  [[nodiscard]] double calculateScalar(const Requirement& par, ScalarAffordance affordance) const {
    // real element implements calculateScalar by itself, then we simply forward the call
    if constexpr (requires {
                    static_cast<const Mixin&>(std::declval<AutoDiffFE>())
                        .template calculateScalarImpl<double>(par, affordance);
                  }) {
      // real element only implements the protected calculateScalarImpl by itself, thus we call that one.
      return Mixin::template calculateScalarImpl<double>(par, affordance);
    } else {
      static_assert(Dune::AlwaysFalse<AutoDiffFE>::value,
                    "Appropriate calculateScalar and calculateScalarImpl functions are not implemented for the "
                    "chosen element.");
    }
  }
 
  void calculateLocalSystem(const Requirement& req, const MatrixAffordance& affordanceM, VectorAffordance affordanceV,
                            typename Traits::template MatrixType<> h, typename Traits::template VectorType<> g) const {
    assert(scalarAffordance(affordanceM) == scalarAffordance(affordanceV));
    Eigen::VectorXdual2nd dx(this->localView().size());
    dx.setZero();
    auto f = [&](auto& x) { return Mixin::calculateScalarImpl(req, scalarAffordance(affordanceV), x); };
    hessian(f, autodiff::wrt(dx), at(dx), g, h);
  }
};
} // namespace Ikarus