Interpolate.hpp

#pragma once

#include <memory>
#include <vector>

#include <dune/common/tuplevector.hh>
#include <dune/common/version.hh>
#include <dune/common/std/optional.hh>
#include <dune/functions/functionspacebases/interpolate.hh>

#include <amdis/common/Concepts.hpp>
#include <amdis/common/FakeContainer.hpp>
#include <amdis/common/FieldMatVec.hpp>
#include <amdis/common/Logical.hpp>
#include <amdis/common/Tags.hpp>
#include <amdis/functions/FunctionFromCallable.hpp>
#include <amdis/functions/HierarchicNodeToRangeMap.hpp>
#include <amdis/functions/NodeIndices.hpp>
#include <amdis/gridfunctions/GridFunction.hpp>
#include <amdis/linearalgebra/Traits.hpp>
#include <amdis/operations/Assigner.hpp>
#include <amdis/typetree/Traversal.hpp>

namespace AMDiS
{
  namespace Impl
  {
    template <class B, class Vec, class GF, class TP, class C, class BV, class NTRE, class Assign>
    void interpolateTreeSubset(B const& basis, Vec& vector, GF const& gf, TP const& treePath,
                               C& counter, BV const& bitVec, NTRE const& nodeToRangeEntry, Assign assign)
    {
      auto localView = basis.localView();

      // set vector to zero at subtree
      if (! std::is_same<Assign, Assigner::assign>::value) {
        for (const auto& e : elements(basis.gridView(), typename BackendTraits<B>::PartitionSet{}))
        {
          localView.bind(e);
          auto&& subTree = Dune::TypeTree::child(localView.tree(), treePath);
          vector.forEach(nodeIndices(localView, subTree), [&](auto dof, auto& coeff) {
            if (bitVec[dof]) { coeff = 0; }
          });
        }
      }

      // Obtain a local view of f
      auto lf = localFunction(gf);

      vector.init(sizeInfo(basis), false);
      counter.init(sizeInfo(basis), true); // set to zero
      for (const auto& e : elements(basis.gridView(), typename BackendTraits<B>::PartitionSet{}))
      {
        localView.bind(e);
        lf.bind(e);

        auto&& subTree = Dune::TypeTree::child(localView.tree(),treePath);
        for_each_leaf_node(subTree, [&](auto const& node, auto const& tp)
        {
          using Traits = typename TYPEOF(node)::FiniteElement::Traits::LocalBasisType::Traits;
          using RangeField = typename Traits::RangeFieldType;

          auto&& fe = node.finiteElement();
          std::size_t feSize = fe.localBasis().size();

          auto bitVecRange = mappedRangeView(Dune::range(feSize), [&](std::size_t i) -> bool {
            return bitVec[localView.index(node.localIndex(i))];
          });

          std::vector<RangeField> visit(bitVecRange.begin(), bitVecRange.end());
          if (std::all_of(visit.begin(), visit.end(), [](auto i) { return i == 0; }))
            return;

          // extract component of local function result corresponding to node in tree
          auto localFj = functionFromCallable<Traits>([&](auto const& local)
          {
            const auto& tmp = lf(local);
            return nodeToRangeEntry(node, tp, Dune::MatVec::as_vector(tmp));
          });

          thread_local std::vector<RangeField> interpolationCoeff;
          fe.localInterpolation().interpolate(localFj, interpolationCoeff);

          counter.scatter(localView, node, visit, assign);
          vector.scatter(localView, node, interpolationCoeff, visit, assign);
        });
      }
      vector.finish();
      counter.finish();
    }

  } // namespace Impl


  template <class T, class Default>
  decltype(auto) value_or(T&& value, Default&&) { return FWD(value); }

  template <class Default>
  decltype(auto) value_or(tag::defaulted, Default&& def) { return FWD(def); }


  /// \brief Interpolate given function in discrete function space
  /**
   * Interpolation is done wrt the leaf node of the ansatz tree
   * corresponding to the given tree path.
   *
   * Notice that this will only work if the range type of f and
   * the block type of coeff are compatible and supported by
   * flatVectorView.
   *
   * \param basis  Global function space basis of discrete function space
   * \param vec    Coefficient vector to represent the interpolation
   * \param gf     GridFunction to interpolate
   * \param tp_    Tree path specifying the part of the ansatz tree to use [RootTreePath]
   * \param c_     Vector that counts for the number of value assignments [FakeContainer]
   * \param bv_    A vector with flags marking all DOFs that should be interpolated [FakeContainer]
   * \param a_     Assignment mode [Assigner::assign]
   */
  template <class Basis, class Vec, class GF, class TP, class C, class BV, class Assign>
  void interpolate(Basis const& basis, Vec& vec, GF const& gf, TP const& tp_, C&& c_, BV&& bv_, Assign&& a_)
  {
    auto&& tp = value_or(FWD(tp_), Dune::TypeTree::hybridTreePath());
    auto&& c  = value_or(FWD(c_),  FakeContainer<int,1>());
    auto&& bv = value_or(FWD(bv_), FakeContainer<bool,true>());
    auto&& a  = value_or(FWD(a_),  Assigner::assign());

    auto ntrm = AMDiS::HierarchicNodeToRangeMap();
    AMDiS::Impl::interpolateTreeSubset(basis, vec, gf, tp, c, bv, ntrm, a);
  }

  template <class B, class Vec, class GF, class TP, class C, class BV>
  void interpolate(B const& basis, Vec& vec, GF const& gf, TP&& tp, C&& c, BV const& bitVec)
  {
    static_assert(not std::is_same<BV, tag::defaulted>::value, "");
    AMDiS::interpolate(basis, vec, gf, FWD(tp), FWD(c), bitVec, tag::defaulted{});
  }

  template <class B, class Vec, class GF, class TP, class C>
  void interpolate(B const& basis, Vec& vec, GF const& gf, TP&& tp, C& counter)
  {
    static_assert(not std::is_same<C, tag::defaulted>::value, "");
    AMDiS::interpolate(basis, vec, gf, FWD(tp), counter, tag::defaulted{}, Assigner::plus_assign{});
  }

  template <class B, class Vec, class GF, class TreePath>
  void interpolate(B const& basis, Vec& vec, GF const& gf, TreePath const& treePath)
  {
    static_assert(not std::is_same<TreePath, tag::defaulted>::value, "");
    AMDiS::interpolate(basis, vec, gf, treePath, tag::defaulted{}, tag::defaulted{}, Assigner::assign{});
  }

  template <class B, class Vec, class GF>
  void interpolate(B const& basis, Vec& vec, GF const& gf)
  {
    AMDiS::interpolate(basis, vec, gf, tag::defaulted{}, tag::defaulted{}, tag::defaulted{}, Assigner::assign{});
  }

} // end namespace AMDiS