SimpleDataTransfer.hpp

#pragma once

#include <map>
#include <optional>
#include <vector>

#include <amdis/DataTransfer.hpp>
#include <amdis/functions/EntitySet.hpp>
#include <amdis/gridfunctions/CoarsenedGridFunction.hpp>
#include <amdis/gridfunctions/DiscreteFunction.hpp>

#include <dune/grid/common/capabilities.hh>

namespace AMDiS {

namespace tag {

  struct simple_datatransfer
      : datatransfer
  {};

} // end namespace tag


template <class Basis, class Container>
class SimpleDataTransfer
{
  using CoefficientVector = std::vector<typename Container::value_type>;
  using EntityId = typename Basis::GridView::Grid::LocalIdSet::IdType;
  using PersistentStorage =  std::map<EntityId, CoefficientVector>;
  using Domain = typename Basis::GridView::template Codim<0>::Entity::Geometry::LocalCoordinate;

private:
  PersistentStorage persistentStorage_;

public:
  SimpleDataTransfer()
  {
    static_assert(Dune::Capabilities::hasSingleGeometryType<typename Basis::GridView::Grid>::v);
  }

  void preAdapt(Basis const& basis, Container const& coeff, bool mightCoarsen)
  {
    CoefficientVector interpolationCoeff;
    auto localView = basis.localView();
    auto interpol = [&](auto const& lf, auto& localCoeff) {
      localCoeff.resize(localView.size());
      Traversal::forEachLeafNode(localView.tree(), [&](auto const& node, auto const& tp)
      {
        auto lf_at_tp = HierarchicNodeWrapper{tp,lf};
        node.finiteElement().localInterpolation().interpolate(lf_at_tp, interpolationCoeff);
        assert(node.size() == interpolationCoeff.size());
        for (std::size_t i = 0; i < node.size(); ++i)
          localCoeff[node.localIndex(i)] = interpolationCoeff[i];
      });
    };

    DiscreteFunction coeffFct{coeff, basis};
    CoarsenedGridFunction c_gf{basis.gridView(), coeffFct};
    auto c_lf = localFunction(c_gf);
    auto const& idSet = basis.gridView().grid().localIdSet();
    for (auto const& e : entitySet(basis))
    {
      localView.bind(e);

      // store the coefficients on the current element
      coeff.gather(localView, persistentStorage_[idSet.id(e)]);

      // if the element might be coarsened, extract also interpolated coefficients on
      // coarser elements
      if (e.mightVanish()) {
        auto father = e;
        while (father.hasFather()) {
          father = father.father();
          auto [it,inserted] = persistentStorage_.try_emplace(idSet.id(father));
          if (inserted) {
            c_lf.bind(father);
            interpol(c_lf, it->second);
          }

          if (!inserted || !father.mightVanish())
            break;
        }
      }
    }
  }


  void adapt(Basis const& basis, Container& coeff)
  {
    coeff.init(basis, false);
    auto localView = basis.localView();

    CoefficientVector interpolationCoeff;
    auto interpol = [&](auto const& local, auto const& localCoeff) {
      Traversal::forEachLeafNode(localView.tree(), [&](auto const& node, auto const& tp)
      {
        auto const& localFE = node.finiteElement();
        auto const& localBasis = localFE.localBasis();
        auto const& localInterpolation = localFE.localInterpolation();

        using LocalBasis = TYPEOF(localBasis);
        using RangeType = typename LocalBasis::Traits::RangeType;
        std::vector<RangeType> shapeValues;
        localInterpolation.interpolate([&](auto const& x) {
          localBasis.evaluateFunction(local(x), shapeValues);
          auto range = localCoeff[node.localIndex(0)] * shapeValues[0];
          assert(node.size() == shapeValues.size());
          for (std::size_t i = 1; i < shapeValues.size(); ++i)
            range.axpy(localCoeff[node.localIndex(i)], shapeValues[i]);
          return range;
        }, interpolationCoeff);
        coeff.scatter(localView, node, interpolationCoeff, Assigner::assign{});
      });
    };

    auto const& idSet = basis.gridView().grid().localIdSet();
    for (const auto& e : entitySet(basis))
    {
      localView.bind(e);
      if (e.isNew()) {
        // interpolate from coarser elements to the fine element
        auto father = e;
        std::function<Domain(Domain const&)> local = [](Domain const& x) { return x; };
        [[maybe_unused]] bool found = false;
        while (father.hasFather()) {
          // store the coordinate transform
          local = [local, geo=father.geometryInFather()](Domain const& x) {
            return geo.global(local(x));
          };
          father = father.father();

          // if coarse element is found, interpolate to fine element
          auto it = persistentStorage_.find(idSet.id(father));
          if (it != persistentStorage_.end()) {
            interpol(local, it->second);
            found = true;
            break;
          }
        }
        assert(found);
      } else {
        // just extract the stored coefficients
        auto it = persistentStorage_.find(idSet.id(e));
        test_exit(it != persistentStorage_.end(),
          "Element should be in persistent storage but isn't.");
        coeff.scatter(localView, it->second, Assigner::assign{});
      }
    }

    coeff.finish();
  }

  void postAdapt(Container& coeff)
  {
    persistentStorage_.clear();
  };
};

template <>
struct DataTransferFactory<tag::simple_datatransfer>
{
  template <class Basis, class T, template <class> class Impl,
    REQUIRES(Concepts::GlobalBasis<Basis>)>
  static DataTransfer<Basis,VectorFacade<T,Impl>> create(Basis const&, VectorFacade<T,Impl> const&)
  {
    return SimpleDataTransfer<Basis, VectorFacade<T,Impl>>{};
  }
};


} // end namespace AMDiS