GlobalIdSet.hpp

#pragma once

#include <tuple>
#include <utility>
#include <vector>

#include <dune/common/exceptions.hh>
#include <dune/common/hybridutilities.hh>
#include <dune/common/version.hh>
#include <dune/functions/functionspacebases/nodes.hh>
#include <dune/functions/functionspacebases/lagrangebasis.hh>
#include <dune/grid/common/gridenums.hh>
#include <dune/localfunctions/common/localkey.hh>

#include <amdis/Output.hpp>
#include <amdis/common/Apply.hpp>
#include <amdis/common/ConceptsBase.hpp>
#include <amdis/common/ForEach.hpp>
#include <amdis/common/TupleUtility.hpp>
#include <amdis/functions/Nodes.hpp>
#include <amdis/utility/Twist.hpp>

namespace Dune
{
  namespace Functions
  {
    // forward declarations...
    template <class GV, int k, class MI>
    class LagrangeDGPreBasis;

    template<typename GV, class MI, bool HI>
    class TaylorHoodPreBasis;
  }
}


namespace AMDiS
{
  // forward declaration
  template <class PreBasis, class TP, class NodeTag = typename Node_t<PreBasis,TP>::NodeTag>
  class NodeIdSet;


  template <class EntityIdType>
  class GlobalIdType : public std::pair<EntityIdType, std::size_t>
  {
    using Super = std::pair<EntityIdType, std::size_t>;

  public:
    GlobalIdType(std::size_t i = 0)
      : Super(EntityIdType{}, i)
    {}

    using Super::Super;

    GlobalIdType& operator++()
    {
      ++this->second;
      return *this;
    }

    // Add to the second index only
    GlobalIdType operator+(std::size_t shift) const
    {
      return GlobalIdType{this->first, this->second + shift};
    }

    friend std::ostream& operator<<(std::ostream& os, GlobalIdType const& id)
    {
      os << "(" << id.first << "," << id.second << ")";
      return os;
    }
  };

  template <class Basis>
  using GlobalIdType_t
    = GlobalIdType<typename Basis::GridView::Grid::GlobalIdSet::IdType>;



  template <class GB, class = std::void_t<> >
  class GlobalBasisIdSet
  {
  public:
    using GlobalBasis = GB;
    using Tree = typename GB::LocalView::Tree;
    using size_type = std::size_t;

    using GridView = typename GlobalBasis::GridView;
    using Grid = typename GridView::Grid;
    using Element = typename GridView::template Codim<0>::Entity;
    using EntityIdType = typename Grid::GlobalIdSet::IdType;
    using PartitionType = Dune::PartitionType;
    using IdType = GlobalIdType_t<GB>;

    using PreBasis = typename GlobalBasis::PreBasis;
    using TreePath = typename GlobalBasis::PrefixPath;
    using NodeIdSet = AMDiS::NodeIdSet<PreBasis, TreePath>;
    using Twist = AMDiS::Twist<typename Grid::GlobalIdSet, GridView::dimension>;

  public:
    GlobalBasisIdSet(GlobalBasis const& globalBasis)
      : tree_(globalBasis.localView().tree())
      , nodeIdSet_(globalBasis.gridView())
      , twist_(globalBasis.gridView().grid().globalIdSet())
    {
      initializeTree(tree_);
    }


    void bind(Element const& element)
    {
      bindTree(tree_, element);
      nodeIdSet_.bind(tree_);
      twist_.bind(element);
      data_.resize(size());
      nodeIdSet_.fillIn(twist_, data_.begin());
    }

    void unbind()
    {
      nodeIdSet_.unbind();
    }

    size_type size() const
    {
      return tree_.size();
    }

    IdType id(size_type i) const
    {
      assert( i < data_.size() );
      return data_[i].first;
    }

    EntityIdType entityId(size_type i) const
    {
      return id(i).first;
    }

    PartitionType partitionType(size_type i) const
    {
      assert( i < data_.size() );
      return data_[i].second;
    }

  protected:
    Tree tree_;
    NodeIdSet nodeIdSet_;
    Twist twist_;
    using Data = std::pair<IdType, PartitionType>;
    std::vector<Data> data_;
  };


  // Specialization for SubspaceBasis
  template <class Basis>
  class GlobalBasisIdSet<Basis, std::void_t<typename Basis::RootBasis>>
      : public GlobalBasisIdSet<typename Basis::RootBasis>
  {
  public:
    GlobalBasisIdSet(Basis const& basis)
      : GlobalBasisIdSet<typename Basis::RootBasis>(basis.rootBasis())
    {}
  };


  template <class PB, class TP, class NodeTag>
  class NodeIdSet
  {
  public:
    using PreBasis = PB;
    using Node = Node_t<PreBasis,TP>;
    using GridView = typename PreBasis::GridView;
    using size_type = std::size_t;

    static_assert(Node::isLeaf, "Generic NodeIdSet implemented for LeafNodes only. Provide a specialization for your node!");

  private:
    static constexpr int dim = GridView::template Codim<0>::Entity::mydimension;

  public:
    NodeIdSet(GridView const& gridView)
      : gridView_(gridView)
      , sizes_{}
    {}

    void bind(const Node& node)
    {
      node_ = &node;
      size_ = node_->finiteElement().size();

      std::fill(std::begin(sizes_), std::end(sizes_), 0u);
      for (size_type i = 0; i < size_ ; ++i) {
        Dune::LocalKey const& localKey = node_->finiteElement().localCoefficients().localKey(i);
        sizes_[localKey.codim()]++;
      }
      auto refElem = Dune::referenceElement<double,GridView::dimension>(node_->element().type());
      for (size_type c = 0; c <= GridView::dimension ; ++c)
        sizes_[c] /= refElem.size(c);
    }

    void unbind()
    {
      node_ = nullptr;
    }

    size_type size() const
    {
      return size_;
    }

    template <class Twist, class It>
    It fillIn(Twist const& twist, It it, size_type shift = 0) const
    {
      assert(node_ != nullptr);
      const auto& gridIdSet = gridView_.grid().globalIdSet();

      for (size_type i = 0; i < size_ ; ++i, ++it) {
        Dune::LocalKey localKey = node_->finiteElement().localCoefficients().localKey(i);
        unsigned int s = localKey.subEntity();
        unsigned int c = localKey.codim();
        it->first = {gridIdSet.subId(node_->element(), s, c), shift + twist.get(localKey,sizes_[c])};

        it->second = Dune::Hybrid::switchCases(std::make_index_sequence<dim+1>{}, c,
          [&](auto codim) { return node_->element().template subEntity<codim>(s).partitionType(); },
          [&]() {
            error_exit("Invalid codimension c = {}", c);
            return Dune::PartitionType{};
          });
      }

      return it;
    }

  protected:
    GridView gridView_;
    const Node* node_ = nullptr;
    size_type size_ = 0;
    std::array<unsigned int,GridView::dimension+1> sizes_;
  };


  // Specialization for PowerBasis
  template <class PreBasis, class TP>
  class NodeIdSet<PreBasis, TP, Dune::TypeTree::PowerNodeTag>
  {
  public:
    using Node = Node_t<PreBasis,TP>;
    using GridView = typename PreBasis::GridView;
    using size_type = std::size_t;

  protected:
    using SubPreBasis = typename PreBasis::SubPreBasis;
    using SubTreePath = decltype(Dune::TypeTree::push_back(std::declval<TP>(), std::size_t(0)));
    using SubNodeIdSet = NodeIdSet<SubPreBasis, SubTreePath>;
    static const std::size_t children = Node::CHILDREN;

  public:
    NodeIdSet(GridView const& gridView)
      : subIds_(gridView)
    {}

    void bind(const Node& node)
    {
      node_ = &node;
      subIds_.bind(node.child(0));
    }

    void unbind()
    {
      node_ = nullptr;
      subIds_.unbind();
    }

    size_type size() const
    {
      return node_->size();
    }

    template <class Twist, class It>
    It fillIn(Twist const& twist, It it, size_type shift = 0) const
    {
      assert(node_ != nullptr);
      for (std::size_t child = 0; child < children; ++child)
      {
        size_type subTreeSize = subIds_.size();
        it = subIds_.fillIn(twist, it, shift);
        shift += subTreeSize;
      }
      return it;
    }

  protected:
    SubNodeIdSet subIds_;
    const Node* node_ = nullptr;
  };


  // Specialization for CompositePreBasis
  template <class PreBasis, class TP>
  class NodeIdSet<PreBasis, TP, Dune::TypeTree::CompositeNodeTag>
  {
  public:
    using Node = Node_t<PreBasis,TP>;
    using GridView = typename PreBasis::GridView;
    using size_type = std::size_t;

  protected:
    static const std::size_t children = Node::CHILDREN;
    using ChildIndices = std::make_index_sequence<children>;

    // The I'th SubPreBasis
    template <std::size_t I>
    using SubPreBasis = typename PreBasis::template SubPreBasis<I>;

    template <std::size_t I>
    using SubTreePath = decltype(Dune::TypeTree::push_back(std::declval<TP>(), Dune::index_constant<I>{}));

    template <std::size_t I>
    using SubNodeIdSet = NodeIdSet<SubPreBasis<I>, SubTreePath<I>>;

    // A tuple of NodeIdSets for the SubPreBases
    using IdsTuple = IndexMapTuple_t<std::make_index_sequence<children>, SubNodeIdSet>;

  public:
    NodeIdSet(GridView const& gridView)
      : idsTuple_(Ranges::applyIndices<children>([&](auto... i) {
          return std::make_tuple(SubNodeIdSet<VALUE(i)>(gridView)...);
          }))
    {}

    void bind(const Node& node)
    {
      node_ = &node;
      Ranges::forIndices<0,children>([&](auto i) {
        std::get<VALUE(i)>(idsTuple_).bind(node.child(i));
      });
    }

    void unbind()
    {
      node_ = nullptr;
      Ranges::forEach(idsTuple_, [](auto& ids) {
        ids.unbind();
      });
    }

    size_type size() const
    {
      return node_->size();
    }

    template <class Twist, class It>
    It fillIn(Twist const& twist, It it, size_type shift = 0) const
    {
      assert(node_ != nullptr);
      Ranges::forEach(idsTuple_, [&](auto const& ids)
      {
        size_type subTreeSize = ids.size();
        it = ids.fillIn(twist, it, shift);
        shift += subTreeSize;
      });
      return it;
    }

  private:
    IdsTuple idsTuple_;
    const Node* node_ = nullptr;
  };


  template <class GV, class MI, bool HI, class TP>
  class NodeIdSet<Dune::Functions::TaylorHoodPreBasis<GV,MI,HI>, TP, Dune::TypeTree::CompositeNodeTag>
  {
  public:
    using PreBasis = Dune::Functions::TaylorHoodPreBasis<GV,MI,HI>;
    using Node = Node_t<PreBasis,TP>;
    using GridView = typename PreBasis::GridView;
    using size_type = std::size_t;

  private:
    using PTP = decltype(Dune::TypeTree::push_back(std::declval<TP>(), Dune::index_constant<1>{}));
    using VTP = decltype(Dune::TypeTree::push_back(std::declval<TP>(), Dune::index_constant<0>{}));
    using V0TP = decltype(Dune::TypeTree::push_back(std::declval<VTP>(), std::size_t(0)));

    // Note: PreBasis::PQ1PreBasis is private
    using PQ1PreBasis = typename NodeIndexSet_t<PreBasis,TP>::PQ1NodeIndexSet::PreBasis;
    using PQ2PreBasis = typename NodeIndexSet_t<PreBasis,TP>::PQ2NodeIndexSet::PreBasis;

    using PQ1NodeIdSet = NodeIdSet<PQ1PreBasis, PTP>; // pressure
    using PQ2NodeIdSet = NodeIdSet<PQ2PreBasis, V0TP>; // velocity

  private:
    static constexpr int dow = GridView::dimensionworld;

  public:
    NodeIdSet(GridView const& gridView)
      : pq1NodeIdSet_(gridView)
      , pq2NodeIdSet_(gridView)
    {}

    void bind(const Node& node)
    {
      node_ = &node;
      using namespace Dune::Indices;
      pq1NodeIdSet_.bind(node.child(_1));
      pq2NodeIdSet_.bind(node.child(_0, 0));
    }

    void unbind()
    {
      pq1NodeIdSet_.unbind();
      pq2NodeIdSet_.unbind();
      node_ = nullptr;
    }

    size_type size() const
    {
      return node_->size();
    }

    template <class Twist, class It>
    It fillIn(Twist const& twist, It it, size_type shift = 0) const
    {
      assert(node_ != nullptr);
      for (int child = 0; child < dow; ++child) {
        size_type subTreeSize = pq2NodeIdSet_.size();
        it = pq2NodeIdSet_.fillIn(twist, it, shift);
        shift += subTreeSize;
      }
      it = pq1NodeIdSet_.fillIn(twist, it, shift);
      return it;
    }

  protected:
    PQ1NodeIdSet pq1NodeIdSet_;
    PQ2NodeIdSet pq2NodeIdSet_;
    const Node* node_ = nullptr;
  };


  template <class GV, int k, class MI, class TP>
  class NodeIdSet<Dune::Functions::LagrangeDGPreBasis<GV, k, MI>, TP>
  {
  public:
    using PreBasis = Dune::Functions::LagrangeDGPreBasis<GV, k, MI>;
    using Node = Node_t<PreBasis, TP>;
    using GridView = typename PreBasis::GridView;
    using size_type = std::size_t;

    NodeIdSet(GridView const& gridView)
      : gridView_(gridView)
    {}

    void bind(const Node& node)
    {
      node_ = &node;
      size_ = node_->finiteElement().size();
    }

    void unbind()
    {
      node_ = nullptr;
    }

    size_type size() const
    {
      return size_;
    }

    template <class Twist, class It>
    It fillIn(Twist const& /*twist*/, It it, size_type shift = 0) const
    {
      assert(node_ != nullptr);
      const auto& gridIdSet = gridView_.grid().globalIdSet();
      auto elementId = gridIdSet.id(node_->element());

      for (size_type i = 0; i < size_; ++i, ++it)
      {
        it->first = {elementId, shift + i};
        it->second = node_->element().partitionType();
      }

      return it;
    }

  protected:
    GridView gridView_;
    const Node* node_ = nullptr;
    size_type size_ = 0;
  };


} // end namespace AMDiS