FunctorGridFunction.hpp

#pragma once

#include <tuple>
#include <type_traits>

#include <dune/common/std/apply.hh>

#include <amdis/Operations.hpp>
#include <amdis/common/ForEach.hpp>
#include <amdis/common/Logical.hpp>
#include <amdis/gridfunctions/Derivative.hpp>
#include <amdis/gridfunctions/GridFunction.hpp>
#include <amdis/gridfunctions/Order.hpp>

namespace AMDiS
{
  namespace Impl
  {
    template <class T0, class... Ts>
    struct DomainType
    {
      using type = typename T0::Domain;
    };

    template <class T0, class... Ts>
    struct EntitySetType
    {
      using type = typename T0::EntitySet;
    };

  } // end namespace Impl


#ifndef DOXYGEN
  template <class Signatur, class Functor, class... LocalFunctions>
  class FunctorLocalFunction;
#endif

  // implementation
  template <class R, class D, class Functor, class... LocalFunctions>
  class FunctorLocalFunction<R(D), Functor, LocalFunctions...>
  {
  public:
    using Range = R;
    using Domain = D;

    enum { hasDerivative = true };

    template <class LocalFct>
    struct LocalFunctionWrapper
    {
      template <class GridFct>
      LocalFunctionWrapper(GridFct const& gridFct)
        : localFct_{localFunction(gridFct)}
      {}

      LocalFct const& operator*() const { return localFct_; }
      LocalFct      & operator*()       { return localFct_; }

      LocalFct localFct_;
    };

  public:
    /// Constructor. Stores copies of the functor and localFunction(gridfunction)s.
    template <class... GridFcts>
    FunctorLocalFunction(Functor const& fct, GridFcts&&... gridFcts)
      : fct_{fct}
      , localFcts_{FWD(gridFcts)...}
    {}

    /// Calls \ref bind for all localFunctions
    template <class Element>
    void bind(Element const& element)
    {
      Tools::for_each(localFcts_, [&](auto& localFct) {
        (*localFct).bind(element);
      });
    }

    /// Calls \ref unbind for all localFunctions
    void unbind()
    {
      Tools::for_each(localFcts_, [&](auto& localFct) {
        (*localFct).unbind();
      });
    }

    /// Applies the functor \ref fct_ to the evaluated localFunctions
    Range operator()(Domain const& x) const
    {
      return Tools::apply([&](auto&&... localFct) { return fct_((*localFct)(x)...); }, localFcts_);
    }

  public:
    /// Return the stored functor
    Functor const& fct() const
    {
      return fct_;
    }

    auto const& localFcts() const
    {
      return localFcts_;
    }

  private:
    Functor fct_;
    std::tuple<LocalFunctionWrapper<LocalFunctions>...> localFcts_;
  };


  /// \fn derivative
  /// \brief Derivative of the LocalFunction of a FunctorGridFunction, utilizing
  /// the chain-rule. Only available if the functor provides partial derivatives.
  /**
   * \f$ d_x(f(lf(x)...)) = \sum_i d_i(f)[lf(x)...] * derivative(lf[i]) \f$
   *
   * **Requirements:**
   * - The Functor `F` must model `Concepts::HasPartial`
   **/
  template <class Sig, class F, class... LFs, class Type,
    REQUIRES(Concepts::HasPartial<F>)>
  auto derivative(FunctorLocalFunction<Sig,F,LFs...> const& lf, Type const& type)
  {
    auto index_seq = std::index_sequence_for<LFs...>{};

    // d_i(f)[lgfs...] * lgfs_i
    auto term_i = [&](auto const _i)
    {
      auto di_f = Tools::apply([&](auto const&... lgfs) {
        return makeFunctorGridFunction(partial(lf.fct(), _i), (*lgfs)...);
      }, lf.localFcts());

      using std::get;
      auto const& lgfs_i = *get<VALUE(_i)>(lf.localFcts());
      return makeFunctorGridFunction(Operation::Multiplies{}, di_f, derivative(lgfs_i, type));
    };

    // sum_i [ d_i(f)[lgfs...] * derivative(lgfs_i)
    auto gridFct = Tools::apply([&](auto const... _i)
    {
      return makeFunctorGridFunction(Operation::Plus{}, term_i(_i)...);
    }, index_seq);

    return localFunction(gridFct);
  }


  /// \fn order
  /// \brief Calculate the polynomial order of a functor `F` that provides a free
  /// function order(f, degs...), where degs are the orders of the LocalFunctions.
  /**
   * **Requirements:**
   * - The functor `F` must model `Concepts::HasFunctorOrder`
   * - All localFunctions `LFs...` must model `Concepts::HasOrder`
   **/
  template <class Sig, class F, class... LFs,
    REQUIRES(Concepts::HasFunctorOrder<F,sizeof...(LFs)>
             && all_of_v<Concepts::HasOrder<LFs>...>)>
  int order(FunctorLocalFunction<Sig,F,LFs...> const& lf)
  {
    return Tools::apply([&lf](auto const&... lgfs) { return order(lf.fct(), order(*lgfs)...); },
      lf.localFcts());
  }


  /// \class FunctorGridFunction
  /// \brief A Gridfunction that applies a functor to the evaluated Gridfunctions
  /**
   * \ingroup GridFunctions
   * Composition of GridFunctions `g_i` by applying a functor `f` locally, i.e.
   * locally it is evaluated
   * \f$ f(g_0(x), g_1(x), ...) \f$
   *
   * \tparam Functor           The type of the outer functor `f`
   * \tparam GridFunctions...  The GridFunction types of `g_i`
   *
   * Requirements:
   * - `arity(f) == sizeof...(GridFunctions)`
   * - `arity(g_i) == arity(g_j) for i != j`
   * - `g_i` models concept \ref GridFunction
   **/
  template <class Functor, class... GridFunctions>
  class FunctorGridFunction
  {
  public:
    /// The result type of the functor when applied to the grid-functions
    using Range = decltype(std::declval<Functor>()(std::declval<typename GridFunctions::Range>()...));

    /// The argument type that can be applied to the grid-functions
    using Domain = typename Impl::DomainType<GridFunctions...>::type;

    /// The set of entities this grid-function binds to
    using EntitySet = typename Impl::EntitySetType<GridFunctions...>::type;

    enum { hasDerivative = false };

  private:
    template <class GridFct>
    using LocalFct = TYPEOF( localFunction(std::declval<GridFct>()) );

    using RawRange = remove_cvref_t<Range>;
    using LocalDomain = typename EntitySet::LocalCoordinate;

  public:
    using LocalFunction = FunctorLocalFunction<RawRange(LocalDomain), Functor, LocalFct<GridFunctions>...>;

    /// \brief Constructor. Stores copies of the functor and gridfunctions.
    template <class... GridFcts>
    explicit FunctorGridFunction(Functor const& fct, GridFcts&&... gridFcts)
      : fct_{fct}
      , gridFcts_{FWD(gridFcts)...}
    {}

    /// Applies the functor to the evaluated gridfunctions
    Range operator()(Domain const& x) const
    {
      return Tools::apply([&](auto&&... gridFct) { return fct_(gridFct(x)...); }, gridFcts_);
    }

    /// Return the stored \ref EntitySet of the first GridFunction
    EntitySet const& entitySet() const
    {
      return std::get<0>(gridFcts_).entitySet();
    }

    LocalFunction makeLocalFunction() const
    {
      return Tools::apply([&](auto const&... gridFcts) { return LocalFunction{fct_, gridFcts...}; }, gridFcts_);
    }

  private:
    Functor fct_;
    std::tuple<GridFunctions...> gridFcts_;
  };


  // Generator function for FunctorGridFunction expressions
  template <class Functor, class... GridFcts>
  auto makeFunctorGridFunction(Functor const& f, GridFcts const&... gridFcts)
  {
    static_assert(all_of_v<Concepts::GridFunction<GridFcts>...>,
      "All passed parameters must be GridFunctions.");
    static_assert(Concepts::Callable<Functor, typename GridFcts::Range...>,
      "Range types of grid functions are not compatible with the functor.");
    return FunctorGridFunction<Functor, GridFcts...>{f, gridFcts...};
  }


#ifndef DOXYGEN
  // PreGridFunction related to FunctorGridFunction.
  template <class Functor, class... PreGridFunctions>
  struct FunctorPreGridFunction
  {
    using Self = FunctorPreGridFunction;

    struct Creator
    {
      template <class GridView>
      static auto create(Self const& self, GridView const& gridView)
      {
        return Tools::apply([&](auto const&... pgf) {
          return makeFunctorGridFunction(self.fct_,
            makeGridFunction(pgf, gridView)...);
        }, self.preGridFcts_);
      }
    };

    template <class... PreGridFcts>
    explicit FunctorPreGridFunction(Functor const& fct, PreGridFcts&&... pgfs)
      : fct_{fct}
      , preGridFcts_{FWD(pgfs)...}
    {}

  private:
    Functor fct_;
    std::tuple<PreGridFunctions...> preGridFcts_;
  };

  namespace Traits
  {
    template <class Functor, class... PreGridFcts>
    struct IsPreGridFunction<FunctorPreGridFunction<Functor, PreGridFcts...>>
      : std::true_type {};
  }
#endif


  /// \brief Generator function for FunctorGridFunction.
  /// \relates FunctorGridFunction
  /**
   * \ingroup GridFunctions
   * Applies the functor `f` to the grid-functions `gridFcts...`.
   * See \ref FunctorGridFunction.
   *
   * **Examples:**
   * - `invokeAtQP([](Dune::FieldVector<double, 2> const& x) { return two_norm(x); }, X());`
   * - `invokeAtQP([](double u, auto const& x) { return u + x[0]; }, 1.0, X());`
   * - `invokeAtQP(Operation::Plus{}, X(0), X(1));`
   **/
  template <class Functor, class... PreGridFcts>
  auto invokeAtQP(Functor const& f, PreGridFcts&&... gridFcts)
  {
    return FunctorPreGridFunction<Functor, TYPEOF(gridFcts)...>{f, FWD(gridFcts)...};
  }

} // end namespace AMDiS