Cabana_Grid_ParticleDistributor.hpp

Go to the documentation of this file.

/****************************************************************************
 * Copyright (c) 2018-2023 by the Cabana authors                            *
 * All rights reserved.                                                     *
 *                                                                          *
 * This file is part of the Cabana library. Cabana is distributed under a   *
 * BSD 3-clause license. For the licensing terms see the LICENSE file in    *
 * the top-level directory.                                                 *
 *                                                                          *
 * SPDX-License-Identifier: BSD-3-Clause                                    *
 ****************************************************************************/

#ifndef CABANA_PARTICLEGRIDDISTRIBUTOR_HPP
#define CABANA_PARTICLEGRIDDISTRIBUTOR_HPP
 
#include <Cabana_DeepCopy.hpp>
#include <Cabana_Distributor.hpp>
 
#include <Cabana_Grid_GlobalGrid.hpp>
#include <Cabana_Grid_GlobalMesh.hpp>
#include <Cabana_Grid_LocalGrid.hpp>
#include <Cabana_Grid_LocalMesh.hpp>
 
#include <Kokkos_Core.hpp>
 
#include <vector>
 
namespace Cabana
{
namespace Grid
{
 
//---------------------------------------------------------------------------//
// Particle Grid Distributor
//---------------------------------------------------------------------------//

template <class LocalGridType, std::size_t NSD = LocalGridType::num_space_dim>
std::enable_if_t<3 == NSD, std::vector<int>>
getTopology( const LocalGridType& local_grid )
{
    std::vector<int> topology( 27, -1 );
    int nr = 0;
    for ( int k = -1; k < 2; ++k )
        for ( int j = -1; j < 2; ++j )
            for ( int i = -1; i < 2; ++i, ++nr )
                topology[nr] = local_grid.neighborRank( i, j, k );
    return topology;
}

template <class LocalGridType, std::size_t NSD = LocalGridType::num_space_dim>
std::enable_if_t<2 == NSD, std::vector<int>>
getTopology( const LocalGridType& local_grid )
{
    std::vector<int> topology( 9, -1 );
    int nr = 0;
    for ( int j = -1; j < 2; ++j )
        for ( int i = -1; i < 2; ++i, ++nr )
            topology[nr] = local_grid.neighborRank( i, j );
    return topology;
}
 
namespace Impl
{
 
// Locate the particles in the local grid and get their destination rank.
// Particles are assumed to only migrate to a location in the nearest
// neighbor halo or stay on this rank. If the particle crosses a global
// periodic boundary, wrap it's coordinates back into the domain.
template <class LocalGridType, class PositionSliceType, class NeighborRankView,
          class DestinationRankView>
void getMigrateDestinations( const LocalGridType& local_grid,
                             const NeighborRankView& neighbor_ranks,
                             DestinationRankView& destinations,
                             PositionSliceType& positions )
{
    static constexpr std::size_t num_space_dim = LocalGridType::num_space_dim;
    using execution_space = typename PositionSliceType::execution_space;
 
    // Check within the local domain.
    const auto& local_mesh = createLocalMesh<Kokkos::HostSpace>( local_grid );
 
    // Use global domain for periodicity.
    const auto& global_grid = local_grid.globalGrid();
    const auto& global_mesh = global_grid.globalMesh();
 
    Kokkos::Array<double, num_space_dim> local_low{};
    Kokkos::Array<double, num_space_dim> local_high{};
    Kokkos::Array<bool, num_space_dim> periodic{};
    Kokkos::Array<double, num_space_dim> global_low{};
    Kokkos::Array<double, num_space_dim> global_high{};
    Kokkos::Array<double, num_space_dim> global_extent{};
 
    for ( std::size_t d = 0; d < num_space_dim; ++d )
    {
        local_low[d] = local_mesh.lowCorner( Own(), d );
        local_high[d] = local_mesh.highCorner( Own(), d );
        periodic[d] = global_grid.isPeriodic( d );
        global_low[d] = global_mesh.lowCorner( d );
        global_high[d] = global_mesh.highCorner( d );
        global_extent[d] = global_mesh.extent( d );
    }
 
    Kokkos::parallel_for(
        "Cabana::Grid::ParticleGridMigrate::get_destinations",
        Kokkos::RangePolicy<execution_space>( 0, positions.size() ),
        KOKKOS_LAMBDA( const int p ) {
            // Compute the logical index of the neighbor we are sending to.
            int nid[num_space_dim];
            for ( std::size_t d = 0; d < num_space_dim; ++d )
            {
                nid[d] = 1;
                if ( positions( p, d ) < local_low[d] )
                    nid[d] = 0;
                else if ( positions( p, d ) > local_high[d] )
                    nid[d] = 2;
            }
 
            // Compute the destination MPI rank [ni + 3*(nj + 3*nk) in 3d].
            int neighbor_index = nid[0];
            for ( std::size_t d = 1; d < num_space_dim; ++d )
            {
                int npower = 1;
                for ( std::size_t dp = 1; dp <= d; ++dp )
                    npower *= 3;
                neighbor_index += npower * nid[d];
            }
            destinations( p ) = neighbor_ranks( neighbor_index );
 
            // Shift particles through periodic boundaries.
            for ( std::size_t d = 0; d < num_space_dim; ++d )
            {
                if ( periodic[d] )
                {
                    if ( positions( p, d ) > global_high[d] )
                        positions( p, d ) -= global_extent[d];
                    else if ( positions( p, d ) < global_low[d] )
                        positions( p, d ) += global_extent[d];
                }
            }
        } );
}
} // namespace Impl
 
//-----------------------------------------------------------------------//
template <class LocalGridType, class PositionSliceType>
int migrateCount( const LocalGridType& local_grid,
                  const PositionSliceType& positions,
                  const int minimum_halo_width )
{
    using grid_type = LocalGridType;
    static constexpr std::size_t num_space_dim = grid_type::num_space_dim;
    using mesh_type = typename grid_type::mesh_type;
    using scalar_type = typename mesh_type::scalar_type;
    using uniform_type = UniformMesh<scalar_type, num_space_dim>;
    static_assert( std::is_same<mesh_type, uniform_type>::value,
                   "Migrate count requires a uniform mesh." );
 
    using execution_space = typename PositionSliceType::execution_space;
 
    // Check within the halo width, within the ghosted domain.
    const auto& local_mesh = createLocalMesh<Kokkos::HostSpace>( local_grid );
 
    Kokkos::Array<double, num_space_dim> local_low{};
    Kokkos::Array<double, num_space_dim> local_high{};
    for ( std::size_t d = 0; d < num_space_dim; ++d )
    {
        auto dx = local_grid.globalGrid().globalMesh().cellSize( d );
        local_low[d] =
            local_mesh.lowCorner( Ghost(), d ) + minimum_halo_width * dx;
        local_high[d] =
            local_mesh.highCorner( Ghost(), d ) - minimum_halo_width * dx;
    }
 
    int comm_count = 0;
    Kokkos::parallel_reduce(
        "Cabana::Grid::ParticleGridMigrate::count",
        Kokkos::RangePolicy<execution_space>( 0, positions.size() ),
        KOKKOS_LAMBDA( const int p, int& result ) {
            for ( std::size_t d = 0; d < num_space_dim; ++d )
                if ( positions( p, d ) < local_low[d] ||
                     positions( p, d ) > local_high[d] )
                {
                    result += 1;
                    break;
                }
        },
        comm_count );
 
    MPI_Allreduce( MPI_IN_PLACE, &comm_count, 1, MPI_INT, MPI_SUM,
                   local_grid.globalGrid().comm() );
 
    return comm_count;
}
 
//---------------------------------------------------------------------------//
template <class CommSpaceType = Mpi, class LocalGridType,
          class PositionSliceType>
Cabana::Distributor<typename PositionSliceType::memory_space>
createParticleDistributor( const LocalGridType& local_grid,
                           PositionSliceType& positions )
{
    using memory_space = typename PositionSliceType::memory_space;
 
    // Get all 26 neighbor ranks.
    auto topology = getTopology( local_grid );
 
    Kokkos::View<int*, Kokkos::HostSpace, Kokkos::MemoryUnmanaged>
        topology_host( topology.data(), topology.size() );
    auto topology_mirror =
        Kokkos::create_mirror_view_and_copy( memory_space(), topology_host );
    Kokkos::View<int*, memory_space> destinations(
        Kokkos::ViewAllocateWithoutInitializing( "destinations" ),
        positions.size() );
 
    // Determine destination ranks for all particles and wrap positions across
    // periodic boundaries.
    Impl::getMigrateDestinations( local_grid, topology_mirror, destinations,
                                  positions );
 
    // Create the Cabana distributor.
    Cabana::Distributor<memory_space, CommSpaceType> distributor(
        local_grid.globalGrid().comm(), destinations, topology );
    return distributor;
}
 
//---------------------------------------------------------------------------//
template <class CommSpaceType = Mpi, class LocalGridType,
          class ParticlePositions, class ParticleContainer>
bool particleMigrate( const LocalGridType& local_grid,
                      const ParticlePositions& positions,
                      ParticleContainer& particles, const int min_halo_width,
                      const bool force_migrate = false )
{
    // When false, this option checks that any particles are nearly outside the
    // ghosted halo region (outside the min_halo_width) before initiating
    // migration. Otherwise, anything outside the local domain is migrated
    // regardless of position in the halo.
    if ( !force_migrate )
    {
        // Check to see if we need to communicate.
        auto comm_count = migrateCount( local_grid, positions, min_halo_width );
 
        // If we have no particles near the ghosted boundary, then exit.
        if ( 0 == comm_count )
            return false;
    }
 
    auto distributor =
        createParticleDistributor<CommSpaceType>( local_grid, positions );
 
    // Redistribute the particles.
    migrate( distributor, particles );
    return true;
}
 
//---------------------------------------------------------------------------//
template <class CommSpaceType = Mpi, class LocalGridType,
          class ParticlePositions, class ParticleContainer>
bool particleMigrate( const LocalGridType& local_grid,
                      const ParticlePositions& positions,
                      const ParticleContainer& src_particles,
                      ParticleContainer& dst_particles,
                      const int min_halo_width,
                      const bool force_migrate = false )
{
    // When false, this option checks that any particles are nearly outside the
    // ghosted halo region (outside the  min_halo_width) before initiating
    // migration. Otherwise, anything outside the local domain is migrated
    // regardless of position in the halo.
    if ( !force_migrate )
    {
        // Check to see if we need to communicate.
        auto comm_count = migrateCount( local_grid, positions, min_halo_width );
 
        // If we have no particles near the ghosted boundary, copy, then exit.
        if ( 0 == comm_count )
        {
            Cabana::deep_copy( dst_particles, src_particles );
            return false;
        }
    }
 
    auto distributor =
        createParticleDistributor<CommSpaceType>( local_grid, positions );
 
    // Resize as needed.
    dst_particles.resize( distributor.totalNumImport() );
 
    // Redistribute the particles.
    migrate( distributor, src_particles, dst_particles );
    return true;
}

template <class... Args>
[[deprecated( "Cabana::Grid::particleGridMigrate is now "
              "Cabana::Grid::particleMigrate. This function wrapper will be "
              "removed in a future release." )]] auto
particleGridMigrate( Args&&... args )
{
    return Cabana::Grid::particleMigrate( std::forward<Args>( args )... );
}
 
template <class... Args>
[[deprecated(
    "Cabana::Grid::createParticleGridDistributor is now "
    "Cabana::Grid::createParticleDistributor. This function wrapper will be "
    "removed in a future release." )]] auto
createParticleGridDistributor( Args&&... args )
{
    return Cabana::Grid::createParticleDistributor(
        std::forward<Args>( args )... );
}
 
} // namespace Grid
} // namespace Cabana
 
#endif // end CABANA_PARTICLEGRIDDISTRIBUTOR_HPP