AverageInterpolator.hpp

#pragma once

#include <map>
#include <vector>

#include <amdis/common/FakeContainer.hpp>
#include <amdis/common/FieldMatVec.hpp>
#include <amdis/functions/EntitySet.hpp>
#include <amdis/functions/HierarchicNodeToRangeMap.hpp>
#include <amdis/functions/NodeIndices.hpp>
#include <amdis/linearalgebra/Traits.hpp>
#include <amdis/operations/Assigner.hpp>
#include <amdis/typetree/Traversal.hpp>

namespace AMDiS
{
  namespace tag
  {
    struct average {};
  }

  template <class Basis,
            class TreePath = Dune::TypeTree::HybridTreePath<>>
  struct AverageInterpolator
  {
    template <class Coeff, class GridFct, class BitVector>
    void operator()(Coeff& coeff, GridFct const& gf, BitVector const& bitVec) const
    {
      // storage for values computed during the interpolation
      VectorType_t<typename Coeff::value_type, Coeff> coeff0(basis_);
      VectorType_t<std::uint8_t, Coeff> counter(basis_);
      using Counter = TYPEOF(counter);

      // Obtain a local view of the gridFunction
      auto lf = localFunction(gf);
      auto localView = basis_.localView();

      std::vector<typename Coeff::value_type> localCoeff;
      std::vector<typename Counter::value_type> localOnes(localView.maxSize(), 1);

      for (const auto& e : entitySet(basis_))
      {
        localView.bind(e);
        lf.bind(e);

        auto&& subTree = Dune::TypeTree::child(localView.tree(),treePath_);
        Traversal::forEachLeafNode(subTree, [&](auto const& node, auto const& tp)
        {
          auto bitVecRange = mappedRangeView(Dune::range(node.size()), [&](std::size_t i) -> bool {
            return bitVec[localView.index(node.localIndex(i))];
          });

          // check whether there is a local contribution
          std::vector<bool> mask(bitVecRange.begin(), bitVecRange.end());
          if (std::all_of(mask.begin(), mask.end(), [](bool m) { return !m; }))
            return;

          // compute local interpolation coefficients
          node.finiteElement().localInterpolation().interpolate(HierarchicNodeWrapper{tp,lf}, localCoeff);

          // assign local coefficients to global vector
          coeff0.scatter(localView, node, localCoeff, mask, Assigner::plus_assign{});

          localOnes.resize(node.size(), 1);
          counter.scatter(localView, node, localOnes, mask, Assigner::plus_assign{});
        });

        lf.unbind();
      }
      coeff0.finish();
      counter.finish();

      // storage for values computed during the interpolation
      std::map<typename Basis::MultiIndex, typename Coeff::value_type> coeffMap;
      std::map<typename Basis::MultiIndex, std::uint8_t> counterMap;

      std::vector<typename Coeff::value_type> localCoeffs;
      std::vector<typename Counter::value_type> localCounter;

      for (const auto& e : entitySet(basis_))
      {
        localView.bind(e);

        auto&& node = Dune::TypeTree::child(localView.tree(),treePath_);
        counter.gather(localView, node, localCounter);
        coeff0.gather(localView, node, localCoeffs);

        for (std::size_t i = 0; i < node.size(); ++i) {
          if (localCounter[i] > 0) {
            auto idx = localView.index(node.localIndex(i));
            coeffMap[idx] = localCoeffs[i];
            counterMap[idx] = localCounter[i];
          }
        }
      }

      assert(coeffMap.size() == counterMap.size());

      // assign computed values to target vector
      auto value_it = coeffMap.begin();
      auto count_it = counterMap.begin();
      for (; value_it != coeffMap.end(); ++value_it, ++count_it) {
        assert(count_it->second > 0);
        coeff.set(value_it->first, value_it->second/double(count_it->second));
      }
      coeff.finish();
    }

    template <class Coeff, class GridFct>
    void operator()(Coeff& coeff, GridFct const& gf) const
    {
      (*this)(coeff, gf, FakeContainer<bool,true>{});
    }

    AverageInterpolator(Basis const& basis, TreePath treePath = {})
      : basis_{basis}
      , treePath_{treePath}
    {}

    Basis const& basis_;
    TreePath treePath_ = {};
  };


  template <>
  struct InterpolatorFactory<tag::average>
  {
    template <class Basis,
              class TreePath = Dune::TypeTree::HybridTreePath<>>
    static auto create(Basis const& basis, TreePath treePath = {})
    {
      return AverageInterpolator<Basis,TreePath>{basis, treePath};
    }
  };

} // end namespace AMDiS