ComposerGridFunction.hpp

#pragma once

#include <tuple>
#include <type_traits>

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

namespace AMDiS
{
#ifndef DOXYGEN
  template <class Signatur, class Element, class Functor, class... LocalFunctions>
  class ComposerLocalFunction;
#endif

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

    enum { hasDerivative = true };

  public:
    template <class... LocalFcts>
    ComposerLocalFunction(Functor const& fct, LocalFcts&&... localFcts)
      : fct_{fct}
      , localFcts_{FWD(localFcts)...}
    {}

    void bind(Element const& element)
    {
      std::apply([&](auto &...lf) { (lf.bind(element), ...); }, localFcts_);
    }

    void unbind()
    {
      std::apply([](auto&... lf) { (lf.unbind(),...); },
        localFcts_);
    }

    bool bound() const
    {
      return std::get<0>(localFcts_).bound();
    }

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

    Element const& localContext() const
    {
      return std::get<0>(localFcts_).localContext();
    }


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

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

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

  template <class Element, class Functor, class... LocalFunctions>
  auto makeComposerLocalFunction(Functor const& f, LocalFunctions const&... lf)
  {
    using D = typename Element::Geometry::LocalCoordinate;
    using R = TYPEOF(f(lf(std::declval<D>())...));
    return ComposerLocalFunction<R(D), Element, Functor, LocalFunctions...>{f, lf...};
  }




  template <class Sig, class E, class F, class... LFs, class Type,
    REQUIRES(Concepts::HasPartial<F>)>
  auto derivativeOf(ComposerLocalFunction<Sig,E,F,LFs...> const& composed, Type const& type)
  {
    // d_i(f)[lgfs...] * lgfs_i
    auto term_i = [&](auto ii)
    {
      auto di_f = std::apply([&](auto const&... lf) {
        return makeComposerLocalFunction<E>(partial(composed.fct(), ii), lf...);
      }, composed.localFunctions());

      auto const& lf_i = std::get<ii>(composed.localFunctions());
      auto df
          = makeComposerLocalFunction<E>(Operation::Multiplies{}, di_f, derivativeOf(lf_i, type));
      if (composed.bound())
        df.bind(composed.localContext());
      return df;
    };

    // sum_i [ d_i(f)[lgfs...] * derivativeOf(lgfs_i)
    auto localFct = Ranges::applyIndices<sizeof...(LFs)>([&](auto... ii)
    {
      return makeComposerLocalFunction<E>(Operation::Plus{}, term_i(ii)...);
    });

    return localFct;
  }



  template <class Sig, class E, class F, class... LFs,
    REQUIRES(Concepts::HasFunctorOrder<F,sizeof...(LFs)>
             && (Concepts::Polynomial<LFs> &&...))>
  int order(ComposerLocalFunction<Sig,E,F,LFs...> const& composed)
  {
    return Ranges::apply([&](auto const&... lf) {
      return order(composed.fct(), order(lf)...);
    }, composed.localFunctions());
  }



  template <class Sig, class EntitySet, class Functor, class... GridFunctions>
  class ComposerGridFunction;

  template <class R, class D, class ES, class Functor, class... GridFunctions>
  class ComposerGridFunction<R(D), ES, Functor, GridFunctions...>
  {
  public:
    using Range = R;
    using Domain = D;
    using EntitySet = ES;

    enum { hasDerivative = false };

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

    using LocalDomain = typename EntitySet::LocalCoordinate;
    using Element = typename EntitySet::Element;

  public:
    using LocalFunction
      = ComposerLocalFunction<Range(LocalDomain), Element, Functor, LocalFct<GridFunctions>...>;

    template <class... GridFcts>
    ComposerGridFunction(EntitySet const& entitySet, Functor const& fct, GridFcts&&... gridFcts)
      : entitySet_{entitySet}
      , fct_{fct}
      , gridFcts_{FWD(gridFcts)...}
    {}

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

    EntitySet const& entitySet() const
    {
      return entitySet_;
    }

    LocalFunction makeLocalFunction() const
    {
      return std::apply([&](auto const&... gf) { return LocalFunction{fct_, localFunction(underlying(gf))...}; },
        gridFcts_);
    }

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


  // Generator function for ComposerGridFunction expressions
  template <class Functor, class GridView, class... GridFcts>
  auto makeComposerGridFunction(Functor const& f, GridView const& gridView,
                                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.");

    using EntitySet = Dune::Functions::GridViewEntitySet<GridView, 0>;
    using Domain = typename EntitySet::GlobalCoordinate;
    using Range = TYPEOF(f(underlying(gridFcts)(std::declval<Domain>())...));

    using FGF = ComposerGridFunction<Range(Domain), EntitySet, Functor, GridFcts...>;
    return FGF{EntitySet{gridView},f, gridFcts...};
  }


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

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

    template <class... PreGridFcts>
    explicit ComposerPreGridFunction(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<ComposerPreGridFunction<Functor, PreGridFcts...>>
      : std::true_type {};
  }
#endif



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

} // end namespace AMDiS