parallelhelpers.hh

Go to the documentation of this file.

// -*- mode: C++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*-
// vi: set et ts=4 sw=4 sts=4:
//
// SPDX-FileCopyrightText: Copyright © DuMux Project contributors, see AUTHORS.md in root folder
// SPDX-License-Identifier: GPL-3.0-or-later
//
#ifndef DUMUX_LINEAR_PARALLELHELPERS_HH
#define DUMUX_LINEAR_PARALLELHELPERS_HH
 
#include <algorithm>
#include <dune/common/exceptions.hh>
#include <dune/geometry/dimension.hh>
#include <dune/grid/common/datahandleif.hh>
#include <dune/grid/common/partitionset.hh>
#include <dune/istl/owneroverlapcopy.hh>
#include <dune/istl/paamg/pinfo.hh>
#include <dune/istl/bvector.hh>
#include <dune/istl/multitypeblockvector.hh>
#include <dumux/parallel/vectorcommdatahandle.hh>
#include <dumux/common/gridcapabilities.hh>
#include <dumux/common/multimapperview.hh>
#include <dumux/common/parameters.hh>
 
#include <string>
#include <utility>
 
namespace Dumux::Detail {
 
template<class LinearSolverTraits, bool canCommunicate = false>
class ParallelISTLHelperImpl {};
 
template<class LinearSolverTraits>
class ParallelISTLHelperImpl<LinearSolverTraits, true>
{
    using GridView = typename LinearSolverTraits::GridView;
    using DofMapper = typename LinearSolverTraits::DofMapper;
 
    // TODO: this is some large number (replace by limits?)
    static constexpr std::size_t ghostMarker_ = 1<<24;
 
    class BaseGatherScatter
    {
        static constexpr int numCodims_ = GridView::dimension + 1;
 
        static constexpr auto getActiveCodims_()
        {
            if constexpr (requires { LinearSolverTraits::dofCodims; })
                return LinearSolverTraits::dofCodims;
            else
                return std::bitset<numCodims_>{ 1UL << LinearSolverTraits::dofCodim };
        }
 
    public:
        BaseGatherScatter(const DofMapper& mapper)
        : mapper_(mapper) {}
 
        template<class EntityType>
        int index(const EntityType& e) const
        { return mapper_.index(e); }
 
        template<class EntityType>
        auto indices(const EntityType& e) const
        { return asMultiMapper(mapper_).indices(e); }
 
        bool contains(int dim, int codim) const
        { return codim >= 0 && codim < numCodims_ && activeCodims_.test(codim); }

        bool fixedSize(int dim, int codim) const
        {
            if (codim < 0 || codim >= numCodims_ || !activeCodims_.test(codim)) return true;
            const auto& gtypes = mapper_.types(codim);
            if (gtypes.empty()) return true;
            const auto nDofs = mapper_.size(gtypes[0]);
            for (const auto& gt : gtypes)
                if (mapper_.size(gt) != nDofs) return false;
            return true;
        }

        template<class EntityType>
        std::size_t size(EntityType& e) const
        { return asMultiMapper(mapper_).indices(e).size(); }
 
        template<class EntityType>
        bool isNeitherInteriorNorBorderEntity(EntityType& e) const
        { return e.partitionType() != Dune::InteriorEntity && e.partitionType() != Dune::BorderEntity; }
 
    protected:
        const DofMapper& mapper_;
    private:
        static constexpr auto activeCodims_ = getActiveCodims_();
    };

    class GhostGatherScatter
    : public BaseGatherScatter
    , public Dune::CommDataHandleIF<GhostGatherScatter, std::size_t>
    {
    public:
        using DataType = std::size_t;
        using BaseGatherScatter::contains;
        using BaseGatherScatter::fixedSize;
        using BaseGatherScatter::size;
 
        GhostGatherScatter(std::vector<std::size_t>& ranks, const DofMapper& mapper)
        : BaseGatherScatter(mapper)
        , ranks_(ranks)
        {}
 
        template<class MessageBuffer, class EntityType>
        void gather(MessageBuffer& buff, const EntityType& e) const
        {
            const bool isGhost = this->isNeitherInteriorNorBorderEntity(e);
            for (auto i : this->indices(e)) {
                auto& data = ranks_[i];
                if (isGhost) data = ghostMarker_;
                buff.write(data);
            }
        }
 
        template<class MessageBuffer, class EntityType>
        void scatter(MessageBuffer& buff, const EntityType& e, std::size_t n)
        {
            const bool isGhost = this->isNeitherInteriorNorBorderEntity(e);
            for (auto i : this->indices(e)) {
                std::size_t x;
                buff.read(x);
                if (isGhost) ranks_[i] = ghostMarker_;
            }
        }
    private:
        std::vector<std::size_t>& ranks_;
    };

    class InteriorBorderGatherScatter
    : public BaseGatherScatter
    , public Dune::CommDataHandleIF<InteriorBorderGatherScatter, std::size_t>
    {
    public:
        using DataType = std::size_t;
        using BaseGatherScatter::contains;
        using BaseGatherScatter::fixedSize;
        using BaseGatherScatter::size;
 
        InteriorBorderGatherScatter(std::vector<std::size_t>& ranks, const DofMapper& mapper)
        : BaseGatherScatter(mapper)
        , ranks_(ranks)
        {}

        template<class MessageBuffer, class EntityType>
        void gather(MessageBuffer& buff, const EntityType& e) const
        {
            const bool isGhost = this->isNeitherInteriorNorBorderEntity(e);
            for (auto i : this->indices(e)) {
                auto& data = ranks_[i];
                if (isGhost) data = ghostMarker_;
                buff.write(data);
            }
        }

        template<class MessageBuffer, class EntityType>
        void scatter(MessageBuffer& buff, const EntityType& e, std::size_t n)
        {
            const bool isGhost = this->isNeitherInteriorNorBorderEntity(e);
            for (auto i : this->indices(e)) {
                std::size_t x;
                buff.read(x);
                auto& data = ranks_[i];
                // we leave ghost unchanged
                // for other dofs, the process with the lowest rank
                // is assigned to be the (unique) owner
                using std::min;
                data = isGhost ? x : min(data, x);
            }
        }
    private:
        std::vector<std::size_t>& ranks_;
    };

    struct NeighbourGatherScatter
    : public BaseGatherScatter
    , public Dune::CommDataHandleIF<NeighbourGatherScatter, int>
    {
        using DataType = int;
        using BaseGatherScatter::contains;
        using BaseGatherScatter::fixedSize;
        using BaseGatherScatter::size;
 
        NeighbourGatherScatter(const DofMapper& mapper, int rank, std::set<int>& neighbours)
        : BaseGatherScatter(mapper)
        , rank_(rank)
        , neighbours_(neighbours)
        {}
 
        // NeighbourGatherScatter sends exactly 1 rank value per entity (not per DOF).
        template<class EntityType>
        std::size_t size(EntityType& e) const { return 1; }
        bool fixedSize(int dim, int codim) const { return true; }
 
        template<class MessageBuffer, class EntityType>
        void gather(MessageBuffer& buff, const EntityType& e) const
        {
            buff.write(rank_);
        }
 
        template<class MessageBuffer, class EntityType>
        void scatter(MessageBuffer& buff, const EntityType& e, std::size_t n)
        {
            int x;
            buff.read(x);
            neighbours_.insert(x);
        }
    private:
        int rank_;
        std::set<int>& neighbours_;
    };
 

    struct SharedGatherScatter
    : public BaseGatherScatter
    , public Dune::CommDataHandleIF<SharedGatherScatter, int>
    {
        using DataType = int;
        using BaseGatherScatter::contains;
        using BaseGatherScatter::fixedSize;
        using BaseGatherScatter::size;
 
        SharedGatherScatter(std::vector<int>& shared, const DofMapper& mapper)
        : BaseGatherScatter(mapper)
        , shared_(shared)
        {}
 
        template<class MessageBuffer, class EntityType>
        void gather(MessageBuffer& buff, EntityType& e) const
        {
            for (std::size_t k = 0; k < this->indices(e).size(); ++k)
                buff.write(int(true));
        }
 
        template<class MessageBuffer, class EntityType>
        void scatter(MessageBuffer& buff, const EntityType &e, std::size_t n)
        {
            for (auto i : this->indices(e)) {
                int x;
                buff.read(x);
                shared_[i] = shared_[i] || x;
            }
        }
    private:
        std::vector<int>& shared_;
    };

    template<typename GlobalIndex>
    struct GlobalIndexGatherScatter
    : public BaseGatherScatter
    , public Dune::CommDataHandleIF<GlobalIndexGatherScatter<GlobalIndex>, GlobalIndex>
    {
        using DataType = GlobalIndex;
        using BaseGatherScatter::contains;
        using BaseGatherScatter::fixedSize;
        using BaseGatherScatter::size;
 
        GlobalIndexGatherScatter(std::vector<GlobalIndex>& globalIndices, const DofMapper& mapper)
        : BaseGatherScatter(mapper)
        , globalIndices_(globalIndices)
        {}
 
        template<class MessageBuffer, class EntityType>
        void gather(MessageBuffer& buff, const EntityType& e) const
        {
            for (auto i : this->indices(e))
                buff.write(globalIndices_[i]);
        }
 
        template<class MessageBuffer, class EntityType>
        void scatter(MessageBuffer& buff, const EntityType& e, std::size_t n)
        {
            for (auto i : this->indices(e)) {
                DataType x;
                buff.read(x);
                using std::min;
                globalIndices_[i] = min(globalIndices_[i], x);
            }
        }
    private:
        std::vector<GlobalIndex>& globalIndices_;
    };
 
public:
 
    ParallelISTLHelperImpl(const GridView& gridView, const DofMapper& mapper)
    : gridView_(gridView), mapper_(mapper)
    {
        if constexpr (LinearSolverTraits::canCommunicate)
            initGhostsAndOwners_();
        else
            DUNE_THROW(Dune::InvalidStateException,
                "Cannot initialize parallel helper for a grid that cannot communicate."
            );
    }
 
    bool isGhost(std::size_t i) const
    { return isGhost_[i] == ghostMarker_; }
 
    bool isOwned(std::size_t i) const
    { return isOwned_[i] == 1; }

    template<class Comm>
    void createParallelIndexSet(Comm& comm) const
    {
        if constexpr (LinearSolverTraits::canCommunicate)
        {
            if (gridView_.comm().size() <= 1)
            {
                comm.remoteIndices().template rebuild<false>();
                return;
            }
 
            // First find out which dofs we share with other processors
            std::vector<int> isShared(mapper_.size(), false);
            SharedGatherScatter sgs(isShared, mapper_);
            gridView_.communicate(sgs, Dune::All_All_Interface, Dune::ForwardCommunication);
 
            // Count shared dofs that we own
            using GlobalIndex = typename Comm::ParallelIndexSet::GlobalIndex;
            GlobalIndex count = 0;
            for (std::size_t i = 0; i < isShared.size(); ++i)
                if (isShared[i] && isOwned_[i] == 1)
                    ++count;
 
            std::vector<GlobalIndex> counts(gridView_.comm().size());
            gridView_.comm().allgather(&count, 1, counts.data());
 
            // Compute start index start_p = \sum_{i=0}^{i<p} counts_i
            const int rank = gridView_.comm().rank();
            auto start = std::accumulate(counts.begin(), counts.begin() + rank, GlobalIndex(0));
 
            // starting from start (offset), number global indices owned by this process consecutively
            std::vector<GlobalIndex> globalIndices(mapper_.size(), std::numeric_limits<GlobalIndex>::max());
            for (std::size_t i = 0; i < globalIndices.size(); ++i)
            {
                if (isOwned_[i] == 1 && isShared[i])
                {
                    globalIndices[i] = start;
                    ++start;
                }
            }
 
            // publish global indices for the shared DOFS to other processors.
            GlobalIndexGatherScatter<GlobalIndex> gigs(globalIndices, mapper_);
            gridView_.communicate(gigs, Dune::All_All_Interface, Dune::ForwardCommunication);
 
            // put the information into the parallel index set
            resizeAndFillIndexSet_(comm, globalIndices);
 
            // Compute neighbours using communication
            std::set<int> neighbours;
            NeighbourGatherScatter ngs(mapper_, gridView_.comm().rank(), neighbours);
            gridView_.communicate(ngs, Dune::All_All_Interface, Dune::ForwardCommunication);
            comm.remoteIndices().setNeighbours(neighbours);
 
            comm.remoteIndices().template rebuild<false>();
        }
        else
            DUNE_THROW(Dune::InvalidStateException,
                "Cannot build parallel index set for a grid that cannot communicate."
            );
    }

    const DofMapper& dofMapper() const
    { return mapper_; }

    const GridView& gridView() const
    { return gridView_; }
 
private:
    void initGhostsAndOwners_()
    {
        const auto rank = gridView_.comm().rank();
        isOwned_.assign(mapper_.size(), rank);
 
        // find out about ghosts
        GhostGatherScatter ggs(isOwned_, mapper_);
 
        if (gridView_.comm().size() > 1)
            gridView_.communicate(ggs, Dune::InteriorBorder_All_Interface, Dune::ForwardCommunication);
 
        // isGhost_ contains rank when the dof is not on a ghost entity (interior or border)
        // and ghostMarker_ when the dof is on a ghost
        isGhost_ = isOwned_;
 
        // partition interior/border uniquely
        // after the communication each dof is assigned to one unique owner process rank
        InteriorBorderGatherScatter dh(isOwned_, mapper_);
 
        if (gridView_.comm().size() > 1)
            gridView_.communicate(dh, Dune::InteriorBorder_InteriorBorder_Interface, Dune::ForwardCommunication);
 
        // convert vector into mask vector
        for (auto& v : isOwned_)
            v = (v == rank) ? 1 : 0;
    }
 
    template<class Comm, class GlobalIndices>
    void resizeAndFillIndexSet_(Comm& comm, const GlobalIndices& globalIndices) const
    {
        comm.indexSet().beginResize();
 
        // add triplets characterizing each dof in the parallel index set
        // (globalIndex, (localIndex, attribute)) where attribute is owner, overlap or copy
        for (std::size_t localIndex = 0; localIndex < globalIndices.size(); ++localIndex)
        {
            const auto globalIndex = globalIndices[localIndex];
            if (globalIndex != std::numeric_limits<typename GlobalIndices::value_type>::max())
            {
                const bool isOwned = isOwned_[localIndex] > 0;
                const auto attr = getAttribute_(comm, isOwned, isGhost(localIndex));
                using LocalIndex = typename Comm::ParallelIndexSet::LocalIndex;
                comm.indexSet().add(globalIndex, LocalIndex{localIndex, attr});
            }
        }
 
        comm.indexSet().endResize();
    }
 
    template<class Comm>
    Dune::OwnerOverlapCopyAttributeSet::AttributeSet
    getAttribute_(const Comm& comm, const bool isOwned, const bool isGhost) const
    {
        if (isOwned) // this process owns the dof (uniquely)
            return Dune::OwnerOverlapCopyAttributeSet::owner;
        else if (isGhost && (comm.category() == Dune::SolverCategory::nonoverlapping) )
            return Dune::OwnerOverlapCopyAttributeSet::overlap;
        else
            return Dune::OwnerOverlapCopyAttributeSet::copy;
    }
 
    const GridView gridView_; 
    const DofMapper& mapper_; 
    std::vector<std::size_t> isOwned_;
    std::vector<std::size_t> isGhost_;
 
};
 
} // end namespace Dumux::Detail
 
namespace Dumux {

template<class LinearSolverTraits>
using ParallelISTLHelper =
    Detail::ParallelISTLHelperImpl<
        LinearSolverTraits, LinearSolverTraits::canCommunicate
    >;
 
 
template<class GridView, class DofMapper, int dofCodim>
class ParallelVectorHelper
{
public:
    ParallelVectorHelper(const GridView& gridView, const DofMapper& mapper)
    : gridView_(gridView), mapper_(mapper)
    {}

    template<class Block, class Alloc>
    void makeNonOverlappingConsistent(Dune::BlockVector<Block, Alloc>& v) const
    {
        if constexpr (Dune::Capabilities::canCommunicate<typename GridView::Traits::Grid, dofCodim>::v)
        {
            VectorCommDataHandleSum<DofMapper, Dune::BlockVector<Block, Alloc>, dofCodim, Block> gs(mapper_, v);
            if (gridView_.comm().size() > 1)
                gridView_.communicate(gs, Dune::InteriorBorder_InteriorBorder_Interface,
                                    Dune::ForwardCommunication);
        }
        else
            DUNE_THROW(Dune::InvalidStateException, "Cannot call makeNonOverlappingConsistent for a grid that cannot communicate codim-" << dofCodim << "-entities.");
    }

    template<class Block, class Alloc>
    void makeOverlappingConsistent(Dune::BlockVector<Block, Alloc>& v) const
    {
        if constexpr (Dune::Capabilities::canCommunicate<typename GridView::Traits::Grid, dofCodim>::v)
        {
            // Iterate over the gridView_ to make overlap entries 0. This is necessary to make
            // `VectorCommDataHandleSum` work properly.
            for (const auto& element : elements(gridView_))
                if (element.partitionType() == Dune::OverlapEntity)
                    v[gridView_.indexSet().index(element)] = 0;
 
            if (gridView_.comm().size() > 1)
            {
                VectorCommDataHandleSum<DofMapper, Dune::BlockVector<Block, Alloc>, dofCodim, Block> gs(mapper_, v);
                gridView_.communicate(gs, Dune::Overlap_All_Interface,
                                    Dune::ForwardCommunication);
            }
        }
        else
            DUNE_THROW(Dune::InvalidStateException, "Cannot call makeNonOverlappingConsistent for a grid that cannot communicate codim-" << dofCodim << "-entities.");
    }

    template<class... Blocks>
    void makeNonOverlappingConsistent(Dune::MultiTypeBlockVector<Blocks...>& v) const
    {
        DUNE_THROW(Dune::NotImplemented, "makeNonOverlappingConsistent for Dune::MultiTypeBlockVector");
    }
 
private:
    const GridView gridView_; 
    const DofMapper& mapper_; 
};
 
template<class GridView, class DofMapper>
class MultiCodimParallelVectorHelper
{
public:
    MultiCodimParallelVectorHelper(const GridView& gridView, const DofMapper& mapper)
    : gridView_(gridView), mapper_(mapper)
    {}

    template<class Block, class Alloc, std::size_t numCodims>
    void makeNonOverlappingConsistent(
        Dune::BlockVector<Block, Alloc>& v,
        const std::bitset<numCodims>& activeCodims
    ) const
    {
        static_assert(numCodims == GridView::dimension + 1,
            "Number of codims must match GridView::dimension+1");
 
        if (gridView_.comm().size() > 1)
        {
            MultiCodimVectorCommDataHandleSum<DofMapper, Dune::BlockVector<Block, Alloc>, GridView::dimension, Block> gs(
                mapper_, v, activeCodims
            );
            gridView_.communicate(gs, Dune::InteriorBorder_InteriorBorder_Interface, Dune::ForwardCommunication);
        }
    }

    template<class Block, class Alloc, std::size_t numCodims>
    void makeOverlappingConsistent(
        Dune::BlockVector<Block, Alloc>& v,
        const std::bitset<numCodims>& activeCodims
    ) const
    {
        static_assert(numCodims == GridView::dimension + 1,
            "Number of codims must match GridView::dimension+1");
 
        zeroOverlapEntries_<0>(v, activeCodims);
 
        if (gridView_.comm().size() > 1)
        {
            MultiCodimVectorCommDataHandleSum<DofMapper, Dune::BlockVector<Block, Alloc>, GridView::dimension, Block> gs(
                mapper_, v, activeCodims
            );
            gridView_.communicate(gs, Dune::Overlap_All_Interface, Dune::ForwardCommunication);
        }
    }
 
private:
    template<int codim, class Block, class Alloc, std::size_t numCodims>
    void zeroOverlapEntries_(Dune::BlockVector<Block, Alloc>& v,
                             const std::bitset<numCodims>& activeCodims) const
    {
        if constexpr (codim <= GridView::dimension)
        {
            if (activeCodims.test(codim))
            {
                for (const auto& element : elements(gridView_))
                {
                    for (int i = 0; i < element.subEntities(codim); ++i)
                    {
                        const auto entity = element.template subEntity<codim>(i);
                        if (entity.partitionType() == Dune::OverlapEntity)
                            for (auto i : asMultiMapper(mapper_).indices(entity))
                                v[i] = 0;
                    }
                }
            }
 
            zeroOverlapEntries_<codim + 1>(v, activeCodims);
        }
    }
 
    const GridView gridView_;
    const DofMapper& mapper_;
};

template<class Matrix, class GridView,
         class RowDofMapper, int rowDofCodim>
class ParallelMatrixHelper
{
    using IdType = typename GridView::Traits::Grid::Traits::GlobalIdSet::IdType;
 
    struct DofIdKey
    {
        int codim;
        IdType id;
 
        friend bool operator<(const DofIdKey& a, const DofIdKey& b)
        {
            if (a.codim != b.codim)
                return a.codim < b.codim;
            return a.id < b.id;
        }
    };

    template<class ColIsGhostFunc>
    struct MatrixPatternExchange
    : public Dune::CommDataHandleIF<MatrixPatternExchange<ColIsGhostFunc>, DofIdKey>
    {
        using DataType = DofIdKey;

        MatrixPatternExchange(const RowDofMapper& rowEntityMapper,
                              const std::map<DofIdKey,int>& globalToLocal,
                              const std::map<int,DofIdKey>& localToGlobal,
                              Matrix& A,
                              std::vector<std::set<int>>& sparsityPattern,
                              const ColIsGhostFunc& isGhostColumDof)
        : rowEntityMapper_(rowEntityMapper)
        , idToIndex_(globalToLocal), indexToID_(localToGlobal)
        , sparsityPattern_(sparsityPattern), A_(A), isGhostColumDof_(isGhostColumDof)
        {
            sparsityPattern_.resize(A.N());
        }

        bool contains (int dim, int codim) const
        { return (codim == rowDofCodim); }

        bool fixedSize(int dim, int codim) const
        { return false; }

        template<class EntityType>
        std::size_t size(EntityType& e) const
        {
            const auto rowIdx = rowEntityMapper_.index(e);
 
            // all column entity indices of this row that are in the index set
            std::size_t n = 0;
            for (auto colIt = A_[rowIdx].begin(); colIt != A_[rowIdx].end(); ++colIt)
            {
                if (indexToID_.count(colIt.index()))
                    ++n;
            }
 
            return n;
        }

        template<class MessageBuffer, class EntityType>
        void gather(MessageBuffer& buff, const EntityType& e) const
        {
            const auto rowIdx = rowEntityMapper_.index(e);
 
            // send all column entity ids of this row that are in the index set
            for (auto colIt = A_[rowIdx].begin(); colIt != A_[rowIdx].end(); ++colIt)
            {
                if (auto it = indexToID_.find(colIt.index()); it != indexToID_.end())
                    buff.write(it->second);
            }
        }

        template<class MessageBuffer, class EntityType>
        void scatter(MessageBuffer& buff, const EntityType& e, std::size_t n)
        {
            const auto rowIdx = rowEntityMapper_.index(e);
            for (std::size_t k = 0; k < n; k++)
            {
                // receive n column entity IDs
                DofIdKey key;
                buff.read(key);
 
                // only add entries that are contained in the index set
                if (const auto it = idToIndex_.find(key); it != idToIndex_.end())
                {
                    const auto colIdx = it->second;
                    if (!isGhostColumDof_(colIdx))
                        sparsityPattern_[rowIdx].insert(colIdx);
                }
            }
        }
 
    private:
        const RowDofMapper& rowEntityMapper_;
        const std::map<DofIdKey, int>& idToIndex_;
        const std::map<int, DofIdKey>& indexToID_;
        std::vector<std::set<int>>& sparsityPattern_;
        Matrix& A_;
        const ColIsGhostFunc& isGhostColumDof_;
 
    }; // class MatrixPatternExchange

    struct MatrixEntry
    {
        DofIdKey first;
        typename Matrix::block_type second;
    };

    struct MatrixEntryExchange
    : public Dune::CommDataHandleIF<MatrixEntryExchange, MatrixEntry>
    {
        using DataType = MatrixEntry;
 
        MatrixEntryExchange(const RowDofMapper& rowEntityMapper,
                            const std::map<DofIdKey, int>& globalToLocal,
                            const std::map<int, DofIdKey>& localToGlobal,
                            Matrix& A)
        : rowEntityMapper_(rowEntityMapper)
        , idToIndex_(globalToLocal), indexToID_(localToGlobal)
        , A_(A)
        {}

        bool contains(int dim, int codim) const
        { return (codim == rowDofCodim); }

        bool fixedSize(int dim, int codim) const
        { return false; }

        template<class EntityType>
        std::size_t size(EntityType& e) const
        {
            const auto rowIdx = rowEntityMapper_.index(e);
            std::size_t n = 0;
            for (auto colIt = A_[rowIdx].begin(); colIt != A_[rowIdx].end(); ++colIt)
            {
                if (indexToID_.count(colIt.index()))
                    ++n;
            }
 
            return n;
        }

        template<class MessageBuffer, class EntityType>
        void gather(MessageBuffer& buff, const EntityType& e) const
        {
            const auto rowIdx = rowEntityMapper_.index(e);
            // send all matrix entries for which the column index is in the index set
            for (auto colIt = A_[rowIdx].begin(); colIt != A_[rowIdx].end(); ++colIt)
            {
                if (auto it = indexToID_.find(colIt.index()); it != indexToID_.end())
                    buff.write(MatrixEntry{ it->second, *colIt });
            }
        }

        template<class MessageBuffer, class EntityType>
        void scatter(MessageBuffer& buff, const EntityType& e, std::size_t n)
        {
            const auto rowIdx = rowEntityMapper_.index(e);
            for (std::size_t k = 0; k < n; k++)
            {
                MatrixEntry m;
                buff.read(m);
                const auto& [colDofID, matrixBlock] = m; // unpack
 
                // only add entries in the index set and for which A has allocated memory
                if (auto it = idToIndex_.find(colDofID); it != idToIndex_.end())
                {
                    if (A_[rowIdx].find(it->second) != A_[rowIdx].end())
                        A_[rowIdx][it->second] += matrixBlock;
                }
            }
        }
 
    private:
        const RowDofMapper& rowEntityMapper_;
        const std::map<DofIdKey, int>& idToIndex_;
        const std::map<int, DofIdKey>& indexToID_;
        Matrix& A_;
 
    }; // class MatrixEntryExchange
 
public:
 
    ParallelMatrixHelper(const GridView& gridView, const RowDofMapper& mapper)
    : gridView_(gridView), mapper_(mapper)
    {
        idToIndex_.clear();
        indexToID_.clear();
 
        initMapsForCodim_<0>();
    }
 
    template<int codim>
    void initMapsForCodim_()
    {
        if constexpr (codim <= GridView::dimension)
        {
            for (const auto& element : elements(gridView_))
            {
                for (int i = 0; i < element.subEntities(codim); ++i)
                {
                    const auto entity = element.template subEntity<codim>(i);
                    if (entity.partitionType() == Dune::BorderEntity)
                    {
                        const auto localIdx = mapper_.index(entity);
                        const DofIdKey key{codim, gridView_.grid().globalIdSet().id(entity)};
                        idToIndex_.emplace(key, localIdx);
                        indexToID_.emplace(localIdx, key);
                    }
                }
            }
 
            initMapsForCodim_<codim + 1>();
        }
    }

    template<class IsGhostFunc>
    void extendMatrix(Matrix& A, const IsGhostFunc& isGhost)
    {
        if constexpr (Dune::Capabilities::canCommunicate<typename GridView::Grid, rowDofCodim>::v)
        {
            if (gridView_.comm().size() <= 1)
                return;
 
            // make a copy of the matrix as originally assembled
            Matrix matrixAsAssembled(A);
 
            // first get entries in sparsity pattern from other processes
            std::size_t numNonZeroEntries = 0;
            std::vector<std::set<int>> sparsityPattern; // column indices for every row
            MatrixPatternExchange<IsGhostFunc> dataHandle(mapper_, idToIndex_, indexToID_, A, sparsityPattern, isGhost);
            gridView_.communicate(
                dataHandle, Dune::InteriorBorder_InteriorBorder_Interface, Dune::ForwardCommunication
            );
 
            // add own entries to the sparsity pattern and count number of non-zeros
            for (auto rowIt = A.begin(); rowIt != A.end(); ++rowIt)
            {
                const auto colEndIt = A[rowIt.index()].end();
                for (auto colIt = rowIt->begin(); colIt != colEndIt; ++colIt)
                    if (!sparsityPattern[rowIt.index()].count(colIt.index()))
                        sparsityPattern[rowIt.index()].insert(colIt.index());
 
                numNonZeroEntries += sparsityPattern[rowIt.index()].size();
            }
 
            // insert any additional entries into the matrix
            A.setSize(matrixAsAssembled.N(), matrixAsAssembled.M(), numNonZeroEntries);
            A.setBuildMode(Matrix::row_wise);
            auto citer = A.createbegin();
            for (auto i = sparsityPattern.begin(), end = sparsityPattern.end(); i!=end; ++i, ++citer)
                for (auto si = i->begin(), send = i->end(); si!=send; ++si)
                    citer.insert(*si);
 
            // reset matrix to contain original values
            A = 0;
            const auto rowEndIt = matrixAsAssembled.end();
            for (auto rowIt = matrixAsAssembled.begin(); rowIt != rowEndIt; ++rowIt)
                for (auto colIt = matrixAsAssembled[rowIt.index()].begin(); colIt != matrixAsAssembled[rowIt.index()].end(); ++colIt)
                    A[rowIt.index()][colIt.index()] = *colIt;
 
            // The matrix A has now a possible extended sparsity pattern but any additional entries are
            // initialized to zero and have to be filled by the sumEntries function
        }
        else
            DUNE_THROW(Dune::InvalidStateException, "Cannot call extendMatrix for a grid that cannot communicate codim-" << rowDofCodim << "-entities.");
    }

    void sumEntries(Matrix& A)
    {
        if constexpr (Dune::Capabilities::canCommunicate<typename GridView::Grid, rowDofCodim>::v)
        {
            if (gridView_.comm().size() <= 1)
                return;
 
            MatrixEntryExchange dataHandle(mapper_, idToIndex_, indexToID_, A);
            gridView_.communicate(
                dataHandle, Dune::InteriorBorder_InteriorBorder_Interface, Dune::ForwardCommunication
            );
        }
        else
            DUNE_THROW(Dune::InvalidStateException,
                "Cannot call sumEntries for a grid that cannot communicate codim-" << rowDofCodim << "-entities."
            );
    }
 
private:
    const GridView gridView_;
    const RowDofMapper& mapper_;
    std::map<DofIdKey, int> idToIndex_;
    std::map<int, DofIdKey> indexToID_;
 
};

template<class Matrix, class GridView, class DofMapper, std::size_t numCodims>
class MultiCodimParallelMatrixHelper
{
    using IdType = typename GridView::Traits::Grid::Traits::GlobalIdSet::IdType;
 
    struct DofIdKey
    {
        int codim;
        IdType id;
        int subIndex; 
 
        friend bool operator<(const DofIdKey& a, const DofIdKey& b)
        {
            if (a.codim != b.codim)
                return a.codim < b.codim;
            if (a.id != b.id)
                return a.id < b.id;
            return a.subIndex < b.subIndex;
        }
    };
 
    // Each sent item encodes which row DOF within the entity and which column DOF.
    // This handles multi-DOF entities (e.g. PQ3 edges with 2 DOFs) correctly.
    struct PatternEntry
    {
        int rowSubIndex; 
        DofIdKey colKey; 
    };
 
    template<class ColIsGhostFunc>
    struct MatrixPatternExchange
    : public Dune::CommDataHandleIF<MatrixPatternExchange<ColIsGhostFunc>, PatternEntry>
    {
        using DataType = PatternEntry;
 
        MatrixPatternExchange(const DofMapper& dofMapper,
                              const std::map<DofIdKey, int>& globalToLocal,
                              const std::map<int, DofIdKey>& localToGlobal,
                              Matrix& A,
                              std::vector<std::set<int>>& sparsityPattern,
                              const ColIsGhostFunc& isGhostColumDof,
                              const std::bitset<numCodims>& activeCodims)
        : dofMapper_(dofMapper)
        , idToIndex_(globalToLocal)
        , indexToID_(localToGlobal)
        , sparsityPattern_(sparsityPattern)
        , A_(A)
        , isGhostColumDof_(isGhostColumDof)
        , activeCodims_(activeCodims)
        {
            sparsityPattern_.resize(A.N());
        }
 
        bool contains(int dim, int codim) const
        {
            if (codim < 0 || codim >= static_cast<int>(numCodims))
                return false;
            return activeCodims_.test(codim);
        }
 
        bool fixedSize(int dim, int codim) const
        { return false; }

        template<class EntityType>
        std::size_t size(EntityType& e) const
        {
            std::size_t n = 0;
            for (auto rowIdx : asMultiMapper(dofMapper_).indices(e))
                for (auto colIt = A_[rowIdx].begin(); colIt != A_[rowIdx].end(); ++colIt)
                    if (indexToID_.count(colIt.index()))
                        ++n;
            return n;
        }

        template<class MessageBuffer, class EntityType>
        void gather(MessageBuffer& buff, const EntityType& e) const
        {
            int rowSub = 0;
            for (auto rowIdx : asMultiMapper(dofMapper_).indices(e))
            {
                for (auto colIt = A_[rowIdx].begin(); colIt != A_[rowIdx].end(); ++colIt)
                    if (auto it = indexToID_.find(colIt.index()); it != indexToID_.end())
                        buff.write(PatternEntry{rowSub, it->second});
                ++rowSub;
            }
        }

        template<class MessageBuffer, class EntityType>
        void scatter(MessageBuffer& buff, const EntityType& e, std::size_t n)
        {
            const auto baseIdx = dofMapper_.index(e);
            for (std::size_t k = 0; k < n; k++)
            {
                PatternEntry entry;
                buff.read(entry);
                const auto rowIdx = baseIdx + entry.rowSubIndex;
                if (const auto it = idToIndex_.find(entry.colKey); it != idToIndex_.end())
                {
                    const auto colIdx = it->second;
                    if (!isGhostColumDof_(colIdx))
                        sparsityPattern_[rowIdx].insert(colIdx);
                }
            }
        }
 
    private:
        const DofMapper& dofMapper_;
        const std::map<DofIdKey, int>& idToIndex_;
        const std::map<int, DofIdKey>& indexToID_;
        std::vector<std::set<int>>& sparsityPattern_;
        Matrix& A_;
        const ColIsGhostFunc& isGhostColumDof_;
        const std::bitset<numCodims>& activeCodims_;
    };
 
    struct MatrixEntry
    {
        int rowSubIndex;                 
        DofIdKey colKey;                 
        typename Matrix::block_type block;
    };
 
    struct MatrixEntryExchange
    : public Dune::CommDataHandleIF<MatrixEntryExchange, MatrixEntry>
    {
        using DataType = MatrixEntry;
 
        MatrixEntryExchange(const DofMapper& dofMapper,
                            const std::map<DofIdKey, int>& globalToLocal,
                            const std::map<int, DofIdKey>& localToGlobal,
                            Matrix& A,
                            const std::bitset<numCodims>& activeCodims)
        : dofMapper_(dofMapper)
        , idToIndex_(globalToLocal)
        , indexToID_(localToGlobal)
        , A_(A)
        , activeCodims_(activeCodims)
        {}
 
        bool contains(int dim, int codim) const
        {
            if (codim < 0 || codim >= static_cast<int>(numCodims))
                return false;
            return activeCodims_.test(codim);
        }
 
        bool fixedSize(int dim, int codim) const
        { return false; }

        template<class EntityType>
        std::size_t size(EntityType& e) const
        {
            std::size_t n = 0;
            for (auto rowIdx : asMultiMapper(dofMapper_).indices(e))
                for (auto colIt = A_[rowIdx].begin(); colIt != A_[rowIdx].end(); ++colIt)
                    if (indexToID_.count(colIt.index()))
                        ++n;
            return n;
        }

        template<class MessageBuffer, class EntityType>
        void gather(MessageBuffer& buff, const EntityType& e) const
        {
            int rowSub = 0;
            for (auto rowIdx : asMultiMapper(dofMapper_).indices(e))
            {
                for (auto colIt = A_[rowIdx].begin(); colIt != A_[rowIdx].end(); ++colIt)
                    if (auto it = indexToID_.find(colIt.index()); it != indexToID_.end())
                        buff.write(MatrixEntry{rowSub, it->second, *colIt});
                ++rowSub;
            }
        }

        template<class MessageBuffer, class EntityType>
        void scatter(MessageBuffer& buff, const EntityType& e, std::size_t n)
        {
            const auto baseIdx = dofMapper_.index(e);
            for (std::size_t k = 0; k < n; k++)
            {
                MatrixEntry m;
                buff.read(m);
                const auto rowIdx = baseIdx + m.rowSubIndex;
                if (auto it = idToIndex_.find(m.colKey); it != idToIndex_.end())
                {
                    if (A_[rowIdx].find(it->second) != A_[rowIdx].end())
                        A_[rowIdx][it->second] += m.block;
                }
            }
        }
 
    private:
        const DofMapper& dofMapper_;
        const std::map<DofIdKey, int>& idToIndex_;
        const std::map<int, DofIdKey>& indexToID_;
        Matrix& A_;
        const std::bitset<numCodims>& activeCodims_;
    };
 
public:
    // Standard constructor: only communicates border-entity DOF couplings.
    // Suitable for iterative solvers (ISTL AMG) where ghost DOF equations
    // are handled separately by the parallel operator mechanism.
    MultiCodimParallelMatrixHelper(const GridView& gridView,
                                   const DofMapper& mapper,
                                   const std::bitset<numCodims>& activeCodims)
    : gridView_(gridView)
    , mapper_(mapper)
    , activeCodims_(activeCodims)
    , includeGhostAndAdjacent_(false)
    , commInterface_(Dune::InteriorBorder_InteriorBorder_Interface)
    {
        initMapsForCodim_<0>();
    }
 
    // Extended constructor for direct solvers (e.g. Trilinos/MUMPS):
    // also includes ghost entities and boundary-adjacent interior entities
    // so that element-interior DOF column entries in border DOF matrix rows
    // are correctly communicated across ranks.
    // Uses All_All_Interface so that ghost entities (GhostEntity on non-owner
    // rank) can participate in the scatter step and receive data from owners.
    MultiCodimParallelMatrixHelper(const GridView& gridView,
                                   const DofMapper& mapper,
                                   const std::bitset<numCodims>& activeCodims,
                                   bool includeGhostAndAdjacent)
    : gridView_(gridView)
    , mapper_(mapper)
    , activeCodims_(activeCodims)
    , includeGhostAndAdjacent_(includeGhostAndAdjacent)
    // InteriorBorder_All_Interface: owner (Interior/Border) gathers and
    // sends; ghost (All) scatters and receives.  This propagates the
    // owner's physical stiffness to ghost DOFs without sending ghost-
    // treatment (identity) entries from the non-owner back to the owner.
    , commInterface_(includeGhostAndAdjacent ? Dune::InteriorBorder_All_Interface
                                              : Dune::InteriorBorder_InteriorBorder_Interface)
    {
        initMapsForCodim_<0>();
    }
 
    template<class IsGhostFunc>
    void extendMatrix(Matrix& A, const IsGhostFunc& isGhost)
    {
        if (gridView_.comm().size() <= 1)
            return;
 
        Matrix matrixAsAssembled(A);
        std::size_t numNonZeroEntries = 0;
        std::vector<std::set<int>> sparsityPattern;
        MatrixPatternExchange<IsGhostFunc> dataHandle(mapper_, idToIndex_, indexToID_, A, sparsityPattern, isGhost, activeCodims_);
        gridView_.communicate(dataHandle, commInterface_, Dune::ForwardCommunication);
 
        for (auto rowIt = A.begin(); rowIt != A.end(); ++rowIt)
        {
            const auto colEndIt = A[rowIt.index()].end();
            for (auto colIt = rowIt->begin(); colIt != colEndIt; ++colIt)
            {
                if (!sparsityPattern[rowIt.index()].count(colIt.index()))
                    sparsityPattern[rowIt.index()].insert(colIt.index());
            }
 
            numNonZeroEntries += sparsityPattern[rowIt.index()].size();
        }
 
        A.setSize(matrixAsAssembled.N(), matrixAsAssembled.M(), numNonZeroEntries);
        A.setBuildMode(Matrix::row_wise);
        auto citer = A.createbegin();
        for (auto i = sparsityPattern.begin(), end = sparsityPattern.end(); i != end; ++i, ++citer)
            for (auto si = i->begin(), send = i->end(); si != send; ++si)
                citer.insert(*si);
 
        A = 0;
        const auto rowEndIt = matrixAsAssembled.end();
        for (auto rowIt = matrixAsAssembled.begin(); rowIt != rowEndIt; ++rowIt)
            for (auto colIt = matrixAsAssembled[rowIt.index()].begin(); colIt != matrixAsAssembled[rowIt.index()].end(); ++colIt)
                A[rowIt.index()][colIt.index()] = *colIt;
    }
 
    void sumEntries(Matrix& A)
    {
        if (gridView_.comm().size() <= 1)
            return;
 
        MatrixEntryExchange dataHandle(mapper_, idToIndex_, indexToID_, A, activeCodims_);
        gridView_.communicate(dataHandle, commInterface_, Dune::ForwardCommunication);
    }
 
private:
    template<std::size_t codim>
    void initMapsForCodim_()
    {
        if constexpr (codim <= GridView::dimension)
        {
            if (activeCodims_.test(codim))
            {
                for (const auto& element : elements(gridView_))
                {
                    for (int i = 0; i < element.subEntities(codim); ++i)
                    {
                        const auto entity = element.template subEntity<codim>(i);
                        const auto pt = entity.partitionType();
 
                        // Standard: always include BorderEntity DOFs.
                        bool shouldProcess = (pt == Dune::BorderEntity);
 
                        // Extended mode (for direct solvers): also include ghost entities
                        // and interior entities adjacent to the partition boundary so that
                        // element-interior DOF column entries in border rows are communicated.
                        if (!shouldProcess && includeGhostAndAdjacent_) {
                            if constexpr (codim == 0) {
                                // Elements are never BorderEntity; include GhostEntity elements
                                // AND InteriorEntity elements adjacent to a border sub-entity.
                                shouldProcess = (pt == Dune::GhostEntity)
                                             || (pt == Dune::InteriorEntity
                                                 && elementHasBorderSubEntity_(entity));
                            } else {
                                shouldProcess = (pt == Dune::GhostEntity);
                            }
                        }
 
                        if (shouldProcess)
                        {
                            // Use asMultiMapper(mapper_).indices(entity) to handle multi-DOF entities
                            // (e.g. PQ3 with 2 DOFs per edge). Each DOF gets a unique
                            // DofIdKey that encodes both the entity ID and the DOF sub-index.
                            const auto globalId = gridView_.grid().globalIdSet().id(entity);
                            int subIdx = 0;
                            for (auto dofIdx : asMultiMapper(mapper_).indices(entity))
                            {
                                const DofIdKey key{static_cast<int>(codim), globalId, subIdx};
                                idToIndex_.emplace(key, dofIdx);
                                indexToID_.emplace(dofIdx, key);
                                ++subIdx;
                            }
                        }
                    }
                }
            }
 
            initMapsForCodim_<codim + 1>();
        }
    }
 
    // Helper: returns true if element has at least one BorderEntity sub-entity
    template<class Element>
    static bool elementHasBorderSubEntity_(const Element& element)
    {
        for (int i = 0; i < element.subEntities(1); ++i)
            if (element.template subEntity<1>(i).partitionType() == Dune::BorderEntity)
                return true;
        if constexpr (Element::mydimension >= 2)
            for (int i = 0; i < element.subEntities(2); ++i)
                if (element.template subEntity<2>(i).partitionType() == Dune::BorderEntity)
                    return true;
        if constexpr (Element::mydimension >= 3)
            for (int i = 0; i < element.subEntities(3); ++i)
                if (element.template subEntity<3>(i).partitionType() == Dune::BorderEntity)
                    return true;
        return false;
    }
 
    const GridView gridView_;
    const DofMapper& mapper_;
    std::bitset<numCodims> activeCodims_;
    bool includeGhostAndAdjacent_;
    Dune::InterfaceType commInterface_;
    std::map<DofIdKey, int> idToIndex_;
    std::map<int, DofIdKey> indexToID_;
};

template<class LinearSolverTraits, class ParallelTraits,
         class Matrix, class ParallelHelper>
void prepareMatrixParallel(Matrix& A, ParallelHelper& pHelper)
{
    if constexpr (ParallelTraits::isNonOverlapping)
    {
        using GridView = typename LinearSolverTraits::GridView;
        using DofMapper = typename LinearSolverTraits::DofMapper;
 
        if constexpr (requires { LinearSolverTraits::dofCodims; })
        {
            constexpr auto dofCodims = LinearSolverTraits::dofCodims;
            MultiCodimParallelMatrixHelper<Matrix, GridView, DofMapper, GridView::dimension + 1> matrixHelper(
                pHelper.gridView(), pHelper.dofMapper(), dofCodims
            );
            matrixHelper.extendMatrix(A, [&pHelper](auto idx){ return pHelper.isGhost(idx); });
            matrixHelper.sumEntries(A);
        }
        else
        {
            static constexpr int dofCodim = LinearSolverTraits::dofCodim;
            ParallelMatrixHelper<Matrix, GridView, DofMapper, dofCodim> matrixHelper(pHelper.gridView(), pHelper.dofMapper());
            matrixHelper.extendMatrix(A, [&pHelper](auto idx){ return pHelper.isGhost(idx); });
            matrixHelper.sumEntries(A);
        }
    }
}

template<class LinearSolverTraits, class ParallelTraits,
         class Vector, class ParallelHelper>
void prepareVectorParallel(Vector& b, ParallelHelper& pHelper)
{
    if constexpr (ParallelTraits::isNonOverlapping)
    {
        using GridView = typename LinearSolverTraits::GridView;
        using DofMapper = typename LinearSolverTraits::DofMapper;
 
        if constexpr (requires { LinearSolverTraits::dofCodims; })
        {
            MultiCodimParallelVectorHelper<GridView, DofMapper> vectorHelper(pHelper.gridView(), pHelper.dofMapper());
            vectorHelper.makeNonOverlappingConsistent(b, LinearSolverTraits::dofCodims);
        }
        else
        {
            static constexpr int dofCodim = LinearSolverTraits::dofCodim;
            ParallelVectorHelper<GridView, DofMapper, dofCodim> vectorHelper(pHelper.gridView(), pHelper.dofMapper());
            vectorHelper.makeNonOverlappingConsistent(b);
        }
    }
}

template<class LinearSolverTraits, class ParallelTraits,
         class Matrix, class Vector, class ParallelHelper>
void prepareLinearAlgebraParallel(Matrix& A, Vector& b, ParallelHelper& pHelper)
{
    prepareMatrixParallel<LinearSolverTraits, ParallelTraits>(A, pHelper);
    prepareVectorParallel<LinearSolverTraits, ParallelTraits>(b, pHelper);
}

template<class LinearSolverTraits, class ParallelTraits,
         class Matrix, class Vector, class ParallelHelper>
void prepareLinearAlgebraParallel(Matrix& A, Vector& b,
                                  std::shared_ptr<typename ParallelTraits::Comm>& comm,
                                  std::shared_ptr<typename ParallelTraits::LinearOperator>& fop,
                                  std::shared_ptr<typename ParallelTraits::ScalarProduct>& sp,
                                  ParallelHelper& pHelper)
{
    prepareLinearAlgebraParallel<LinearSolverTraits, ParallelTraits>(A, b, pHelper);
    const auto category = ParallelTraits::isNonOverlapping ?
        Dune::SolverCategory::nonoverlapping : Dune::SolverCategory::overlapping;
 
    comm = std::make_shared<typename ParallelTraits::Comm>(pHelper.gridView().comm(), category);
    pHelper.createParallelIndexSet(*comm);
    fop = std::make_shared<typename ParallelTraits::LinearOperator>(A, *comm);
    sp = std::make_shared<typename ParallelTraits::ScalarProduct>(*comm);
}
 
} // end namespace Dumux
 
#endif