ParticleSpatialOverlapLayout.h

//
// Class ParticleSpatialOverlapLayout
//   Particle layout based on spatial decomposition.
//
//   This is a specialized version of ParticleSpatialLayout, which allows
//   particles to be on multiple processes if they are close to the respective
//   region.
//
//   After each 'time step' in a calculation, which is defined as a period
//   in which the particle positions may change enough to affect the global
//   layout, the user must call the 'update' routine, which will move
//   particles between processors, etc.  After the Nth call to update, a
//   load balancing routine will be called instead.  The user may set the
//   frequency of load balancing (N), or may supply a function to
//   determine if load balancing should be done or not.
//
#ifndef IPPL_PARTICLE_SPATIAL_OVERLAP_LAYOUT_H
#define IPPL_PARTICLE_SPATIAL_OVERLAP_LAYOUT_H
 
#include "Types/IpplTypes.h"
 
#include "FieldLayout/FieldLayout.h"
#include "Particle/ParticleBase.h"
#include "Particle/ParticleLayout.h"
#include "Region/RegionLayout.h"
 
namespace ippl {
    // This additional namespace is needed to allow for Mesh to have default argument
    // UniformCartesian<T, Dim> via typedef at the end of the file. Compiling for cuda led to error
    // when trying to do so directly...
    template <typename T, unsigned Dim, class Mesh = UniformCartesian<T, Dim>,
              typename... PositionProperties>
    class ParticleSpatialOverlapLayout
        : public ParticleSpatialLayout<T, Dim, Mesh, PositionProperties...> {
    public:
        using Base = ParticleSpatialLayout<T, Dim, Mesh, PositionProperties...>;
        using typename Base::position_execution_space;
        using typename Base::position_memory_space;
 
        using typename Base::bool_type;
        using typename Base::hash_type;
        using typename Base::locate_type;
 
        using typename Base::FieldLayout_t;
        using typename Base::RegionLayout_t;
        using typename Base::vector_type;
 
        using size_type = detail::size_type;
 
        using index_t                     = typename hash_type::value_type;
        using particle_neighbor_list_type = hash_type;
        using typename Base::particle_position_type;
 
        using typename Base::region_view_type;
        using typename Base::region_type;
        using neighbor_list = typename FieldLayout_t::neighbor_list;
 
    public:
        // the maximum number of overlapping ranks
        using locate_type_nd = Kokkos::View<index_t * [1 << Dim], position_memory_space>;
 
        class ParticleNeighborData {
        private:
            friend class ParticleSpatialOverlapLayout;
 
            ParticleNeighborData(size_type numLocalParticles, Vector<size_type, Dim> cellStrides,
                                 Vector<size_type, Dim> numCells, Vector<T, Dim> cellWidth,
                                 region_type region, hash_type cellStartingIdx, hash_type cellIndex,
                                 hash_type cellParticleCount, hash_type cellPermutationForward,
                                 hash_type cellPermutationBackward)
                : numLocalParticles(numLocalParticles)
                , cellStrides(cellStrides)
                , numCells(numCells)
                , cellWidth(cellWidth)
                , region(region)
                , cellStartingIdx(cellStartingIdx)
                , cellIndex(cellIndex)
                , cellParticleCount(cellParticleCount)
                , cellPermutationForward(cellPermutationForward)
                , cellPermutationBackward(cellPermutationBackward) {}
 
            size_type numLocalParticles;
            Vector<size_type, Dim> cellStrides;
            Vector<size_type, Dim> numCells;
            Vector<T, Dim> cellWidth;
            region_type region;
            hash_type cellStartingIdx;
            hash_type cellIndex;
            hash_type cellParticleCount;
            hash_type cellPermutationForward;
            hash_type cellPermutationBackward;
        };
 
    public:
        ParticleSpatialOverlapLayout(FieldLayout<Dim>& fl, Mesh& mesh, const T& rcutoff);
 
        ParticleSpatialOverlapLayout()
            : ParticleSpatialLayout<T, Dim, PositionProperties...>() {}
 
        ~ParticleSpatialOverlapLayout() = default;
 
        void updateLayout(FieldLayout<Dim>&, Mesh&);
 
        template <class ParticleContainer>
        void update(ParticleContainer& pc);
 
        template <typename ExecutionSpace, typename Functor>
        void forEachPair(Functor&& f) const;
 
        ParticleNeighborData getParticleNeighborData() const;
 
        KOKKOS_FUNCTION static particle_neighbor_list_type getParticleNeighbors(
            index_t particleIndex, const ParticleNeighborData& particleNeighborData);
 
        KOKKOS_FUNCTION static particle_neighbor_list_type getParticleNeighbors(
            const vector_type& pos, const ParticleNeighborData& particleNeighborData);
 
        size_type numberOfSends(int rank, const locate_type& ranks);
 
        void fillHash(int rank, const locate_type& ranks, const locate_type& offsets,
                      hash_type& hash);
 
        template <typename ParticleContainer>
        std::pair<detail::size_type, detail::size_type> locateParticles(
            const ParticleContainer& pc, locate_type& ranks, locate_type& rankOffsets,
            bool_type& invalid, locate_type& nSends_dview, locate_type& sends_dview) const;
 
        locate_type getFlatNeighbors(const neighbor_list& neighbors) const;
 
        locate_type getNonNeighborRanks(const locate_type& neighbors_view) const;
 
    protected:
        const T rcutoff_m;
        Vector<size_type, Dim> numCells_m;
        Vector<size_type, Dim> cellStrides_m;
        Vector<T, Dim> cellWidth_m;
        size_type totalCells_m;
        size_type numGhostCells_m;
        size_type numLocalCells_m;
        size_type numLocalParticles_m;
        static constexpr size_type numGhostCellsPerDim_m = 1;
        hash_type cellPermutationForward_m;
        hash_type cellPermutationBackward_m;
        hash_type cellStartingIdx_m;
        hash_type cellIndex_m;
        hash_type cellParticleCount_m;
 
        using CellIndex_t     = Vector<size_type, Dim>;
        using FlatCellIndex_t = typename CellIndex_t::value_type;
 
    public:
        void initializeCells();
 
        template <class ParticleContainer>
        void particleExchange(ParticleContainer& pc);
 
        template <class ParticleContainer>
        void buildCells(ParticleContainer& pc);
 
        template <class ParticleContainer>
        void createPeriodicGhostParticles(ParticleContainer& pc);
 
    protected:
        template <std::size_t... Idx>
        KOKKOS_INLINE_FUNCTION constexpr static bool isCloseToBoundary(
            const std::index_sequence<Idx...>&, const vector_type& pos, const region_type& region,
            Kokkos::Array<bool, Dim> periodic, T overlap);
 
        KOKKOS_INLINE_FUNCTION constexpr static FlatCellIndex_t toFlatCellIndex(
            const CellIndex_t& cellIndex, const Vector<size_type, Dim>& cellStrides,
            hash_type cellPermutationForward);
 
        KOKKOS_INLINE_FUNCTION constexpr static CellIndex_t toCellIndex(
            FlatCellIndex_t nonPermutedIndex, const Vector<size_type, Dim>& numCells);
 
        template <std::size_t... Idx>
        KOKKOS_INLINE_FUNCTION constexpr static bool isLocalCellIndex(
            const std::index_sequence<Idx...>&, const CellIndex_t& index,
            const Vector<size_type, Dim>& numCells);
 
        template <std::size_t... Idx>
        KOKKOS_INLINE_FUNCTION constexpr static bool positionInRegion(
            const std::index_sequence<Idx...>&, const vector_type& pos, const region_type& region,
            T overlap);
 
        KOKKOS_INLINE_FUNCTION constexpr static CellIndex_t getCellIndex(
            const vector_type& pos, const region_type& region, const Vector<T, Dim>& cellWidth);
 
        using cell_particle_neighbor_list_type =
            Kokkos::Array<size_type, detail::countHypercubes(Dim)>;
 
        KOKKOS_INLINE_FUNCTION constexpr static cell_particle_neighbor_list_type getCellNeighbors(
            const CellIndex_t& cellIndex, const Vector<size_type, Dim>& cellStrides,
            const hash_type& cellPermutationForward);
    };
}  // namespace ippl
 
#include "Particle/ParticleSpatialOverlapLayout.hpp"
 
#endif  // IPPL_PARTICLE_SPATIAL_OVERLAP_LAYOUT_H