Cabana_NeighborList.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_NEIGHBORLIST_HPP
#define CABANA_NEIGHBORLIST_HPP
 
#include <Kokkos_Core.hpp>
 
#include <Cabana_Sort.hpp>
 
namespace Cabana
{
//---------------------------------------------------------------------------//
// Neighbor List Interface
//---------------------------------------------------------------------------//
class FullNeighborTag
{
};
 
//---------------------------------------------------------------------------//
class HalfNeighborTag
{
};
 
//---------------------------------------------------------------------------//
class SelfNeighborTag
{
};
 
//---------------------------------------------------------------------------//
template <class Tag>
class NeighborDiscriminator;

template <>
class NeighborDiscriminator<SelfNeighborTag>
{
  public:
    KOKKOS_INLINE_FUNCTION
    static bool isValid( const std::size_t p, const std::size_t n )
    {
        return ( p != n );
    }
};

template <>
class NeighborDiscriminator<FullNeighborTag>
{
  public:
    template <std::size_t NumSpaceDim>
    KOKKOS_INLINE_FUNCTION static bool
    isValid( const std::size_t p, const Kokkos::Array<double, NumSpaceDim>,
             const std::size_t n, const Kokkos::Array<double, NumSpaceDim> )
    {
        return ( p != n );
    }

    KOKKOS_INLINE_FUNCTION
    static bool isValid( const std::size_t p, const double, const double,
                         const double, const std::size_t n, const double,
                         const double, const double )
    {
        return ( p != n );
    }
};

template <>
class NeighborDiscriminator<HalfNeighborTag>
{
  public:
    template <std::size_t NumSpaceDim>
    KOKKOS_INLINE_FUNCTION static bool
    isValid( const std::size_t p, const Kokkos::Array<double, NumSpaceDim> xp,
             const std::size_t n, const Kokkos::Array<double, NumSpaceDim> xn )
    {
        return ( ( p != n ) &&
                 ( ( xn[0] > xp[0] ) ||
                   ( ( xn[0] == xp[0] ) &&
                     ( ( xn[1] > xp[1] ) ||
                       ( ( xn[1] == xp[1] ) && ( xn[2] > xp[2] ) ) ) ) ) );
    }

    KOKKOS_INLINE_FUNCTION static bool
    isValid( const std::size_t p, const Kokkos::Array<double, 3> xp,
             const std::size_t n, const Kokkos::Array<double, 3> xn )
    {
        return ( ( p != n ) &&
                 ( ( xn[0] > xp[0] ) ||
                   ( ( xn[0] == xp[0] ) &&
                     ( ( xn[1] > xp[1] ) ||
                       ( ( xn[1] == xp[1] ) && ( xn[2] > xp[2] ) ) ) ) ) );
    }

    KOKKOS_INLINE_FUNCTION static bool
    isValid( const std::size_t p, const Kokkos::Array<double, 2> xp,
             const std::size_t n, const Kokkos::Array<double, 2> xn )
    {
        return ( ( p != n ) &&
                 ( ( xn[0] > xp[0] ) ||
                   ( ( xn[0] == xp[0] ) && ( ( xn[1] > xp[1] ) ) ) ) );
    }

    KOKKOS_INLINE_FUNCTION static bool
    isValid( const std::size_t p, const double xp, const double yp,
             const double zp, const std::size_t n, const double xn,
             const double yn, const double zn )
    {
        return ( ( p != n ) &&
                 ( ( xn > xp ) ||
                   ( ( xn == xp ) &&
                     ( ( yn > yp ) || ( ( yn == yp ) && ( zn > zp ) ) ) ) ) );
    }
};
 
//---------------------------------------------------------------------------//
template <class NeighborListType>
class NeighborList
{
  public:
    using memory_space = typename NeighborListType::memory_space;

    KOKKOS_INLINE_FUNCTION
    static std::size_t totalNeighbor( const NeighborListType& list );

    KOKKOS_INLINE_FUNCTION
    static std::size_t maxNeighbor( const NeighborListType& list );

    KOKKOS_INLINE_FUNCTION
    static std::size_t numNeighbor( const NeighborListType& list,
                                    const std::size_t particle_index );

    KOKKOS_INLINE_FUNCTION
    static std::size_t getNeighbor( const NeighborListType& list,
                                    const std::size_t particle_index,
                                    const std::size_t neighbor_index );

    KOKKOS_INLINE_FUNCTION
    std::size_t setNeighbor( NeighborListType& list,
                             const std::size_t particle_index,
                             const std::size_t neighbor_index );
};
 
//---------------------------------------------------------------------------//
 
namespace Impl
{
template <class ListType>
KOKKOS_INLINE_FUNCTION std::size_t
totalNeighbor( const ListType& list, const std::size_t num_particles )
{
    std::size_t total_n = 0;
    // Sum neighbors across all particles.
    for ( std::size_t p = 0; p < num_particles; p++ )
        total_n += NeighborList<ListType>::numNeighbor( list, p );
    return total_n;
}

template <class ListType>
KOKKOS_INLINE_FUNCTION std::size_t
maxNeighbor( const ListType& list, const std::size_t num_particles )
{
    std::size_t max_n = 0;
    for ( std::size_t p = 0; p < num_particles; p++ )
        if ( NeighborList<ListType>::numNeighbor( list, p ) > max_n )
            max_n = NeighborList<ListType>::numNeighbor( list, p );
    return max_n;
}
} // namespace Impl

template <class ExecutionSpace, class ListType>
Kokkos::View<int* [2], typename ListType::memory_space>
neighborHistogram( ExecutionSpace exec_space, const std::size_t num_particles,
                   const ListType& list, const int num_bin )
{
    // Allocate View of neighbors per particle
    auto num_neigh = Kokkos::View<int*, typename ListType::memory_space>(
        "particle_count", num_particles );
 
    // Extract from neighbor list interface.
    auto extract_functor = KOKKOS_LAMBDA( const int p )
    {
        num_neigh( p ) = NeighborList<ListType>::numNeighbor( list, p );
    };
    Kokkos::RangePolicy<ExecutionSpace> particle_policy( exec_space, 0,
                                                         num_particles );
    Kokkos::parallel_for( "Cabana::NeighborList::neighborHistogram::extract",
                          particle_policy, extract_functor );
    Kokkos::fence();
 
    auto bin_data = Cabana::binByKey( num_neigh, num_bin );
 
    auto histogram = Kokkos::View<int* [2], typename ListType::memory_space>(
        "particle_count", num_bin, 2 );
    auto histogram_functor = KOKKOS_LAMBDA( const int b )
    {
        int max_neigh = NeighborList<ListType>::maxNeighbor( list );
        double bin_width =
            static_cast<double>( max_neigh ) / static_cast<double>( num_bin );
        if ( num_bin > max_neigh )
            bin_width = 1;
        // Wait to cast back to int to get the actual bin edge.
        histogram( b, 0 ) = static_cast<int>( ( b + 1 ) * bin_width );
        histogram( b, 1 ) = bin_data.binSize( b );
    };
    Kokkos::RangePolicy<ExecutionSpace> bin_policy( exec_space, 0, num_bin );
    Kokkos::parallel_for( "Cabana::NeighborList::neighborHistogram", bin_policy,
                          histogram_functor );
    Kokkos::fence();
 
    return histogram;
}
 
} // end namespace Cabana
 
#endif // end CABANA_NEIGHBORLIST_HPP