partitionhelper.hh

// -*- tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 2 -*-
// vi: set et ts=4 sw=2 sts=2:
#ifndef DUNE_MMESH_MISC_PARTITIONHELPER_HH
#define DUNE_MMESH_MISC_PARTITIONHELPER_HH
 
#include <dune/grid/common/partitionset.hh>
#include <dune/mmesh/grid/rangegenerators.hh>
 
namespace Dune {
 
// PartitionHelper
template <class Grid>
struct PartitionHelper {
  static constexpr int dim = Grid::dimension;
  using IdType = typename Grid::IdType;
  using ConnectivityType = std::unordered_set<int>;
  using LinksType = std::vector<int>;
  using LeafIterator =
      typename Grid::LeafIterator::Implementation::HostGridLeafIterator;
 
  PartitionHelper(const Grid& grid) : grid_(grid) {}
 
  void distribute() {
    // Initialize leafBegin_ and leafEnd_
    if constexpr (dim == 2) {
      leafBegin_ = grid().getHostGrid().finite_faces_begin();
      leafEnd_ = grid().getHostGrid().finite_faces_end();
    } else if constexpr (dim == 3) {
      leafBegin_ = grid().getHostGrid().finite_cells_begin();
      leafEnd_ = grid().getHostGrid().finite_cells_end();
    }
 
    if (grid().comm().size() == 1) return;
 
    setRanks();
    computePartitions();
  }
 
  template <class Entity>
  bool contains(PartitionIteratorType pitype, const Entity& e) const {
    return contains(pitype, partition(e));
  }
 
  bool contains(PartitionIteratorType pitype, int partitionType) const {
    if (partitionType == -1) return false;
 
    switch (pitype) {
      case All_Partition:
        return true;
 
      case Interior_Partition:
        return partitionType == 0;
 
      case InteriorBorder_Partition:
      case Overlap_Partition:
      case OverlapFront_Partition:
        return partitionType != 2;
 
      case Ghost_Partition:
        return partitionType == 2;
    }
    return false;
  }
 
  template <class Entity>
  PartitionType partitionType(const Entity& e) const {
    return partitionType(partition(e));
  }
 
  PartitionType partitionType(int partition) const {
    if (partition == 0) return InteriorEntity;
 
    if (partition == 1)
      return BorderEntity;
 
    else
      return GhostEntity;
  }
 
  void updatePartitions() { computePartitions(); }
 
  const LinksType& links() const { return links_; }
 
  auto& comm() const { return grid().comm(); }
 
  template <class Entity>
  ConnectivityType connectivity(const Entity& e) const {
    if constexpr (Entity::dimension == dim)
      return e.impl().hostEntity()->info().connectivity;
    else
      try {
        return interfaceConnectivity_.at(e.impl().id());
      } catch (std::out_of_range&) {
        return ConnectivityType();
      }
  }
 
  template <class Entity>
  int rank(const Entity& e) const {
    if constexpr (Entity::dimension == dim)
      return e.impl().hostEntity()->info().rank;
    else
      return grid().asIntersection(e).inside().impl().hostEntity()->info().rank;
  }
 
  const LeafIterator& leafInteriorBegin() const { return leafBegin_; }
 
  const LeafIterator& leafInteriorEnd() const { return leafEnd_; }
 
 private:
  template <class Entity>
  int partition(const Entity& e) const {
    static constexpr int cd = Entity::codimension;
    if (grid().comm().size() == 1) return 0;  // interior
 
    if constexpr (Entity::dimension == dim) {
      if constexpr (Entity::codimension == 0 || Entity::codimension == dim)
        return e.impl().hostEntity()->info().partition;
      else {
        auto entry = partition_[cd - 1].find(e.impl().id());
        if (entry != partition_[cd - 1].end()) return entry->second;
      }
    } else {
      auto entry = interfacePartition_[cd].find(e.impl().id());
      if (entry != interfacePartition_[cd].end()) return entry->second;
 
      // handle empty interface
      return -1;
    }
 
    DUNE_THROW(InvalidStateException, "Partition not set yet!");
    return 0;
  }
 
  template <class Entity>
  void setPartition(const Entity& e, int partition) {
    if constexpr (Entity::dimension == dim) {
      if constexpr (Entity::codimension == 0 || Entity::codimension == dim)
        e.impl().hostEntity()->info().partition = partition;
      else
        partition_[Entity::codimension - 1][e.impl().id()] = partition;
    } else
      interfacePartition_[Entity::codimension][e.impl().id()] = partition;
  }
 
  template <class Entity>
  void addConnectivity(const Entity& e, int rank) {
    if constexpr (Entity::dimension == dim)
      e.impl().hostEntity()->info().connectivity.insert(rank);
    else
      interfaceConnectivity_[e.impl().id()].insert(rank);
 
    if (partition(e) == 0)
      if (std::find(links_.begin(), links_.end(), rank) == links_.end())
        links_.push_back(rank);
  }
 
  template <class Entity>
  void clearConnectivity(const Entity& e) {
    if constexpr (Entity::dimension == dim)
      e.impl().hostEntity()->info().connectivity.clear();
    else
      interfaceConnectivity_[e.impl().id()].clear();
  }
 
  void setRanks() {
    int rank = grid().comm().rank();
    int size = grid().comm().size();
 
    std::size_t N;
    if constexpr (dim == 2)
      N = grid().getHostGrid().number_of_faces();
    else
      N = grid().getHostGrid().number_of_cells();
 
    std::size_t i = 0;
    forEntityDim<dim>([this, &i, &rank, &size, &N](const auto& fc) {
      auto getRank = [&size, &N](std::size_t i) { return i * size / N; };
      int r = getRank(i);
 
      // store begin iterator
      if (r == rank && getRank(i - 1) == rank - 1) this->leafBegin_ = fc;
 
      // store end iterator
      if (r == rank + 1 && getRank(i - 1) == rank) this->leafEnd_ = fc;
 
      auto e = grid().entity(fc);
      e.impl().hostEntity()->info().rank = r;
      i++;
    });
  }
 
  void computePartitions() {
    links_.clear();
 
    for (int i = 0; i <= dim; ++i) partition_[i].clear();
 
    // Set interior elements
    forEntityDim<dim>([this](const auto& fc) {
      const auto e = grid().entity(fc);
      if (rank(e) == grid().comm().rank())
        setPartition(e, 0);  // interior
      else
        setPartition(e, -1);  // none
 
      clearConnectivity(e);
    });
 
    // Set ghosts
    forEntityDim<dim - 1>([this](const auto& fc) {
      const auto e = grid().entity(fc);
      const auto is = grid().asIntersection(e);
      if (is.neighbor()) {
        const auto& inside = is.inside();
        const auto& outside = is.outside();
        int pIn = partition(inside);
        int pOut = partition(outside);
 
        // border
        if ((pIn == 0 && pOut != 0) or (pIn != 0 && pOut == 0)) {
          setPartition(pIn == 0 ? outside : inside, 2);  // ghost
          addConnectivity(inside, rank(outside));
          addConnectivity(outside, rank(inside));
        }
      }
    });
 
    // Compute interface partitions based on Codim 0 entities, this might add
    // further ghost elements to the bulk
    computeInterfacePartitions();
 
    // Set facets
    forEntityDim<dim - 1>([this](const auto& fc) {
      const auto e = grid().entity(fc);
      const auto is = grid().asIntersection(e);
 
      int pIn = partition(is.inside());
 
      if (is.neighbor()) {
        int pOut = partition(is.outside());
 
        // interior
        if (pIn == 0 && pOut == 0) setPartition(e, 0);
 
        // border
        else if ((pIn == 0 && pOut == 2) or (pIn == 2 && pOut == 0))
          setPartition(e, 1);
 
        // ghost
        else if (pIn == 2 || pOut == 2)
          setPartition(e, 2);
 
        // none
        else
          setPartition(e, -1);
      } else {
        // interior
        if (pIn == 0) setPartition(e, 0);
 
        // ghost
        else if (pIn == 2)
          setPartition(e, 2);
 
        // none
        else
          setPartition(e, -1);
      }
    });
 
    // Codim 2 in 3D
    if constexpr (dim == 3) {
      forEntityDim<1>([this](const auto& fc) {
        const auto edge = grid().entity(fc);
        std::size_t count = 0, interior = 0, ghost = 0;
        for (const auto e : incidentElements(edge)) {
          count++;
          if (partition(e) == 0) interior++;
          if (partition(e) == 2) ghost++;
        }
 
        // interior
        if (interior == count) setPartition(edge, 0);
 
        // border
        else if (interior > 0 and ghost > 0)
          setPartition(edge, 1);
 
        // ghost
        else if (interior == 0 and ghost > 0)
          setPartition(edge, 2);
 
        // none
        else
          setPartition(edge, -1);
      });
    }
 
    // Set vertices
    forEntityDim<0>([this](const auto& fc) {
      const auto v = grid().entity(fc);
      std::size_t count = 0, interior = 0, ghost = 0;
      for (const auto e : incidentElements(v)) {
        count++;
        if (partition(e) == 0) interior++;
        if (partition(e) == 2) ghost++;
      }
 
      // interior
      if (interior == count) setPartition(v, 0);
 
      // border
      else if (interior > 0 and ghost > 0)
        setPartition(v, 1);
 
      // ghost
      else if (interior == 0 and ghost > 0)
        setPartition(v, 2);
 
      // none
      else
        setPartition(v, -1);
    });
  }
 
 public:
  void computeInterfacePartitions() {
    static constexpr int idim = dim - 1;
 
    for (int i = 0; i <= idim; ++i) interfacePartition_[i].clear();
 
    // Set interior elements
    forEntityDim<idim>([this](const auto& fc) {
      if (!grid().isInterface(grid().entity(fc))) return;
 
      const auto e = grid().interfaceGrid().entity(fc);
      const auto is = grid().asIntersection(e);
 
      clearConnectivity(e);
 
      if (rank(e) == grid().comm().rank()) {
        setPartition(e, 0);  // interior
 
        // add connectivity to this entity being ghost on outside rank
        if (is.neighbor())
          if (rank(is.outside()) != grid().comm().rank())
            addConnectivity(e, rank(is.outside()));
      } else
        setPartition(e, -1);  // none
 
      // if outside bulk entity is interior, the interface element is at least
      // ghost
      if (is.neighbor())
        if (rank(is.outside()) == grid().comm().rank())
          if (partition(e) == -1) {
            setPartition(e, 2);  // ghost
            addConnectivity(e, rank(is.inside()));
          }
    });
 
    // Set ghosts
    forEntityDim<idim - 1>([this](const auto& fc) {
      if (!grid().isInterface(grid().entity(fc))) return;
 
      // convert to interface vertex
      const auto e = grid().interfaceGrid().entity(fc);
 
      std::size_t count = 0, interior = 0, other = 0;
      for (const auto& incident : incidentInterfaceElements(e)) {
        count++;
        if (partition(incident) == 0) interior++;
        if (partition(incident) != 0) other++;
      }
 
      // if we find both interior and other entities we set none to ghost
      if (interior > 0 and other > 0)
        for (const auto& incident : incidentInterfaceElements(e))
          if (partition(incident) == -1) {
            setPartition(incident, 2);  // ghost
 
            // make sure that both adjacent bulk entities are at least ghost
            auto intersection = grid().asIntersection(incident);
            if (partition(intersection.inside()) == -1)
              setPartition(intersection.inside(), 2);  // ghost
            if (intersection.neighbor())
              if (partition(intersection.outside()) == -1)
                setPartition(intersection.outside(), 2);  // ghost
          }
    });
 
    // Set facets
    forEntityDim<idim - 1>([this](const auto& fc) {
      if (!grid().isInterface(grid().entity(fc))) return;
 
      // convert to interface vertex
      const auto e = grid().interfaceGrid().entity(fc);
 
      std::size_t count = 0, interior = 0, ghost = 0;
      std::unordered_set<int> connectivity;
      for (const auto& incident : incidentInterfaceElements(e)) {
        count++;
        if (partition(incident) == 0) interior++;
        if (partition(incident) == 2) ghost++;
        connectivity.insert(rank(incident));
      }
 
      // interior
      if (interior == count) setPartition(e, 0);
 
      // border
      else if (interior > 0 and ghost > 0) {
        setPartition(e, 1);
 
        for (const auto& incident : incidentInterfaceElements(e))
          for (auto r : connectivity)
            if (r != rank(incident)) addConnectivity(incident, r);
      }
 
      // ghost
      else if (ghost > 0)
        setPartition(e, 2);
 
      // none
      else
        setPartition(e, -1);
    });
 
    // Set vertices in 3d
    if constexpr (dim == 3) {
      forEntityDim<idim - 2>([this](const auto& fc) {
        if (!grid().isInterface(grid().entity(fc))) return;
 
        const auto v = grid().interfaceGrid().entity(fc);
 
        std::size_t count = 0, interior = 0, ghost = 0;
        for (const auto& incident : incidentInterfaceElements(v)) {
          count++;
          if (partition(incident) == 0) interior++;
          if (partition(incident) == 2) ghost++;
        }
 
        // interior
        if (interior == count) setPartition(v, 0);
 
        // border
        else if (interior > 0 and ghost > 0)
          setPartition(v, 1);
 
        // ghost
        else if (ghost > 0)
          setPartition(v, 2);
 
        // none
        else
          setPartition(v, -1);
      });
    }
  }
 
 private:
  template <int edim, class F>
  void forEntityDim(const F& f) {
    if constexpr (edim == dim) {
      if constexpr (dim == 2)
        for (auto fc = grid().getHostGrid().finite_faces_begin();
             fc != grid().getHostGrid().finite_faces_end(); ++fc)
          f(fc);
 
      if constexpr (dim == 3)
        for (auto fc = grid().getHostGrid().finite_cells_begin();
             fc != grid().getHostGrid().finite_cells_end(); ++fc)
          f(fc);
    }
 
    if constexpr (edim == 0)
      for (auto fc = grid().getHostGrid().finite_vertices_begin();
           fc != grid().getHostGrid().finite_vertices_end(); ++fc)
        f(fc);
 
    if constexpr (edim == 1)
      for (auto fc = grid().getHostGrid().finite_edges_begin();
           fc != grid().getHostGrid().finite_edges_end(); ++fc)
        f(*fc);
 
    if constexpr (dim == 3 && edim == 2)
      for (auto fc = grid().getHostGrid().finite_facets_begin();
           fc != grid().getHostGrid().finite_facets_end(); ++fc)
        f(*fc);
  }
 
  const Grid& grid() const { return grid_; }
 
  std::array<double, 2> xbounds_;
  std::array<std::unordered_map<IdType, int>, dim - 1> partition_;
  std::array<std::unordered_map<IdType, int>, dim> interfacePartition_;
  std::unordered_map<IdType, ConnectivityType> interfaceConnectivity_;
  LeafIterator leafBegin_, leafEnd_;
  LinksType links_;
  const Grid& grid_;
};
 
}  // end namespace Dune
 
#endif