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/common/Order.hpp>
#include <amdis/gridfunctions/Derivative.hpp>
#include <amdis/gridfunctions/GridFunction.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:
    template <class... GridFcts>
    FunctorLocalFunction(Functor const& fct, GridFcts&&... gridFcts)
      : fct_{fct}
      , localFcts_{FWD(gridFcts)...}
    {}

    template <class Element>
    void bind(Element const& element)
    {
      Tools::for_each(localFcts_, [&](auto& localFct) {
        (*localFct).bind(element);
      });
    }

    void unbind()
    {
      Tools::for_each(localFcts_, [&](auto& localFct) {
        (*localFct).unbind();
      });
    }

    Range operator()(Domain const& x) const
    {
      return Tools::apply([&](auto&&... localFct) { return fct_((*localFct)(x)...); }, localFcts_);
    }

  public:
    Functor const& fct() const
    {
      return fct_;
    }

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

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



  template <class Sig, class F, class... LFs, class Type,
    REQUIRES(Concepts::HasPartial<F>)>
  auto derivativeOf(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, derivativeOf(lgfs_i, type));
    };

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

    return localFunction(gridFct);
  }



  template <class Sig, class F, class... LFs,
    REQUIRES(Concepts::HasFunctorOrder<F,sizeof...(LFs)>
             && (Concepts::Polynomial<LFs> &&...))>
  int order(FunctorLocalFunction<Sig,F,LFs...> const& lf)
  {
    return Tools::apply([&lf](auto const&... lgfs) {
      return order(lf.fct(), order(*lgfs)...);
    }, lf.localFcts());
  }



  template <class Functor, class... GridFunctions>
  class FunctorGridFunction
  {
  public:
    using Range = decltype(std::declval<Functor>()(std::declval<typename GridFunctions::Range>()...));

    using Domain = typename Impl::DomainType<GridFunctions...>::type;

    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>...>;

    template <class... GridFcts>
    explicit FunctorGridFunction(Functor const& fct, GridFcts&&... gridFcts)
      : fct_{fct}
      , gridFcts_{FWD(gridFcts)...}
    {}

    Range operator()(Domain const& x) const
    {
      return Tools::apply([&](auto&&... gridFct) { return fct_(gridFct(x)...); }, gridFcts_);
    }

    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((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



  template <class Functor, class... PreGridFcts>
  auto invokeAtQP(Functor const& f, PreGridFcts&&... gridFcts)
  {
    return FunctorPreGridFunction<Functor, TYPEOF(gridFcts)...>{f, FWD(gridFcts)...};
  }

} // end namespace AMDiS