IndexDistribution.hpp

#pragma once
 
#include <array>
#include <algorithm>
#include <iterator>
#include <numeric>
#include <vector>
 
#include <petsc.h>
 
#include <dune/common/typeutilities.hh>
#include <dune/common/parallel/communication.hh>
 
#if HAVE_MPI
#include <dune/common/parallel/indexset.hh>
#include <dune/common/parallel/localindex.hh>
#include <dune/common/parallel/plocalindex.hh>
#include <dune/common/parallel/remoteindices.hh>
#include <dune/common/parallel/remoteindices.hh>
 
#include <amdis/common/parallel/Communicator.hpp>
#include <amdis/linearalgebra/ParallelIndexSet.hpp>
#endif
 
#include <amdis/functions/GlobalIdSet.hpp>
#include <amdis/linearalgebra/AttributeSet.hpp>
 
namespace AMDiS
{
  template <class LI>
  class PetscSequentialIndexDistribution
  {
    using Self = PetscSequentialIndexDistribution;
 
  public:
    using size_type = LI;
    using DofIndex = size_type;
    using LocalIndex = size_type;
    using GlobalIndex = PetscInt;
 
  public:
    template <class Basis,
      Dune::disableCopyMove<Self, Basis> = 0>
    PetscSequentialIndexDistribution(Basis const& basis)
    {
      update(basis);
    }
 
    MPI_Comm comm() const
    {
      return PETSC_COMM_SELF;
    }
 
    size_type localSize() const
    {
      return localSize_;
    }
 
    std::array<size_type,1> localSizes() const
    {
      return {localSize_};
    }
 
    size_type globalSize() const
    {
      return globalSize_;
    }
 
    std::array<GlobalIndex,1> globalStarts() const
    {
      return {0u};
    }
 
    std::vector<GlobalIndex> const& globalIndices() const
    {
      return indices_;
    }
 
    GlobalIndex ghostSize() const
    {
      return 0;
    }
 
    std::array<GlobalIndex,0> ghostIndices() const
    {
      return {};
    }
 
    LocalIndex globalToGhost(GlobalIndex const& n) const
    {
      assert(false && "There are no ghost indices in sequential dofmappings");
      return 0;
    }
 
    LocalIndex dofToGhost(DofIndex const& n) const
    {
      assert(false && "There are no ghost indices in sequential dofmappings");
      return 0;
    }
 
    GlobalIndex global(LocalIndex const& n) const
    {
      return n;
    }
 
    LocalIndex globalToLocal(GlobalIndex const& n) const
    {
      return n;
    }
 
    LocalIndex dofToLocal(DofIndex const& n) const
    {
      return n;
    }
 
 
    bool owner(DofIndex const& n) const
    {
      assert(n < localSize());
      return true;
    }
 
    bool globalOwner(GlobalIndex const& n) const
    {
      return globalOwner(0, n);
    }
 
    bool globalOwner(int p, GlobalIndex const& n) const
    {
      assert(p == 0);
      assert(n < globalSize());
      return true;
    }
 
    template <class Basis>
    void update(Basis const& basis)
    {
      localSize_ = basis.dimension();
      globalSize_ = basis.dimension();
      indices_.resize(globalSize_);
      std::iota(indices_.begin(), indices_.end(), size_type(0));
    }
 
    void debug() const {}
 
  private:
    size_type localSize_ = 0;
    size_type globalSize_ = 0;
    std::vector<GlobalIndex> indices_;
  };
 
  template <class GlobalId, class LI, class Comm>
  struct PetscIndexDistributionType
  {
    using type = PetscSequentialIndexDistribution<LI>;
  };
 
#if HAVE_MPI
  template <class GlobalId, class LI>
  class PetscParallelIndexDistribution
  {
    friend class MatrixNnzStructure;
 
    using Self = PetscParallelIndexDistribution;
    using PLocalIndex = Dune::ParallelLocalIndex<DefaultAttributeSet::Type>;
    using IndexSet = Dune::ParallelIndexSet<GlobalId, PLocalIndex, 512>;
    using Attribute = typename AttributeSet<IndexSet>::type;
    using RemoteIndices = Dune::RemoteIndices<IndexSet>;
 
  public:
    using size_type = LI;
    using DofIndex = size_type;
    using LocalIndex = size_type;
    using GlobalIndex = PetscInt;
 
  public:
    template <class Basis,
      Dune::disableCopyMove<Self, Basis> = 0>
    PetscParallelIndexDistribution(Basis const& basis)
      : world_(basis.gridView().comm())
    {
      update(basis);
    }
 
    MPI_Comm comm() const
    {
      return world_;
    }
 
    size_type localSize() const
    {
      return localSize_;
    }
 
    std::vector<size_type> const& localSizes() const
    {
      return sizes_;
    }
 
    size_type globalSize() const
    {
      return globalSize_;
    }
 
    std::vector<GlobalIndex> const& globalStarts() const
    {
      return starts_;
    }
 
    std::vector<GlobalIndex> const& globalIndices() const
    {
      return globalIndices_;
    }
 
    size_type ghostSize() const
    {
      return ghostSize_;
    }
 
    std::vector<GlobalIndex> const& ghostIndices() const
    {
      return ghostIndices_;
    }
 
    LocalIndex globalToGhost(GlobalIndex const& n) const
    {
      auto it = std::find(ghostIndices_.begin(), ghostIndices_.end(), n);
      assert(it != ghostIndices_.end());
      return std::distance(ghostIndices_.begin(), it);
    }
 
    LocalIndex dofToGhost(DofIndex const& n) const
    {
      assert(!owner(n));
      assert(n < ghostLocalIndices_.size());
      assert(ghostLocalIndices_[n] < ghostSize_);
 
      return ghostLocalIndices_[n];
    }
 
    GlobalIndex global(DofIndex const& n) const
    {
      assert(n < globalIndices_.size());
      return globalIndices_[n];
    }
 
 
    LocalIndex globalToLocal(GlobalIndex const& n) const
    {
      return n - starts_[world_.rank()];
    }
 
    LocalIndex dofToLocal(DofIndex const& n) const
    {
      assert(n < globalIndices_.size());
      return globalToLocal(globalIndices_[n]);
    }
 
    bool owner(DofIndex const& n) const
    {
      assert(n < owner_.size());
      return owner_[n];
    }
 
    bool globalOwner(GlobalIndex const& n) const
    {
      return globalOwner(world_.rank(), n);
    }
 
    bool globalOwner(int p, GlobalIndex const& n) const
    {
      assert(p < world_.size());
      return n >= starts_[p] && n < starts_[p+1];
    }
 
    RemoteIndices const& remoteIndices() const
    {
      return *remoteIndices_;
    }
 
    template <class Basis>
    void update(Basis const& basis);
 
    void debug() const;
 
  private:
    void reset()
    {
      sizes_.clear();
      starts_.clear();
      localSize_ = 0;
      globalSize_ = 0;
      ghostSize_ = 0;
 
      globalIndices_.clear();
      ghostIndices_.clear();
      ghostLocalIndices_.clear();
      owner_.clear();
      indexSet_ = std::make_unique<IndexSet>();
      remoteIndices_ = std::make_unique<RemoteIndices>();
    }
 
  private:
    Mpi::Communicator world_;
 
    std::vector<size_type> sizes_;
    std::vector<GlobalIndex> starts_;
    size_type localSize_;
    size_type globalSize_;
    size_type ghostSize_;
 
    std::vector<GlobalIndex> globalIndices_; // indexed by LocalIndex
    std::vector<GlobalIndex> ghostIndices_;
    std::vector<LocalIndex> ghostLocalIndices_;
    std::vector<bool> owner_;
 
    std::unique_ptr<IndexSet> indexSet_ = nullptr;
    std::unique_ptr<RemoteIndices> remoteIndices_ = nullptr;
 
    const Mpi::Tag tag_{7513};
  };
 
  template <class GlobalId, class LI>
  struct PetscIndexDistributionType<GlobalId, LI, Dune::Communication<MPI_Comm>>
  {
    using type = PetscParallelIndexDistribution<GlobalId, LI>;
  };
#endif
 
  template <class B>
  using PetscIndexDistribution_t
    = typename PetscIndexDistributionType<GlobalIdType_t<B>, typename B::size_type, typename B::GridView::CollectiveCommunication>::type;
 
} // end namespace AMDiS
 
#include <amdis/linearalgebra/petsc/IndexDistribution.inc.hpp>