Cabana_LinkedCellList.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_LINKEDCELLLIST_HPP
#define CABANA_LINKEDCELLLIST_HPP
 
#include <Cabana_NeighborList.hpp>
#include <Cabana_Slice.hpp>
#include <Cabana_Sort.hpp>
#include <Cabana_Utils.hpp>
#include <impl/Cabana_CartesianGrid.hpp>
 
#include <Kokkos_Core.hpp>
#include <Kokkos_Profiling_ScopedRegion.hpp>
#include <Kokkos_ScatterView.hpp>
 
#include <cassert>
 
namespace Cabana
{
//---------------------------------------------------------------------------//
 
template <class Scalar>
struct LinkedCellStencil
{
    Impl::CartesianGrid<Scalar> grid;
    int max_cells_dir;
    int max_cells;
    int cell_range;

    LinkedCellStencil() = default;

    LinkedCellStencil( const Scalar neighborhood_radius,
                       const Scalar cell_size_ratio, const Scalar grid_min[3],
                       const Scalar grid_max[3] )
    {
        Scalar dx = neighborhood_radius * cell_size_ratio;
        grid = Impl::CartesianGrid<Scalar>(
            grid_min[0], grid_min[1], grid_min[2], grid_max[0], grid_max[1],
            grid_max[2], dx, dx, dx );
        cell_range = std::ceil( 1 / cell_size_ratio );
        max_cells_dir = 2 * cell_range + 1;
        max_cells = max_cells_dir * max_cells_dir * max_cells_dir;
    }

    KOKKOS_INLINE_FUNCTION
    void getCells( const int cell, int& imin, int& imax, int& jmin, int& jmax,
                   int& kmin, int& kmax ) const
    {
        int i, j, k;
        grid.ijkBinIndex( cell, i, j, k );
 
        kmin = ( k - cell_range > 0 ) ? k - cell_range : 0;
        kmax =
            ( k + cell_range + 1 < grid._nz ) ? k + cell_range + 1 : grid._nz;
 
        jmin = ( j - cell_range > 0 ) ? j - cell_range : 0;
        jmax =
            ( j + cell_range + 1 < grid._ny ) ? j + cell_range + 1 : grid._ny;
 
        imin = ( i - cell_range > 0 ) ? i - cell_range : 0;
        imax =
            ( i + cell_range + 1 < grid._nx ) ? i + cell_range + 1 : grid._nx;
    }
};
 
//---------------------------------------------------------------------------//
template <class MemorySpace, class Scalar = double>
class LinkedCellList
{
  public:
    using memory_space = MemorySpace;
    static_assert( Kokkos::is_memory_space<MemorySpace>() );

    using execution_space = typename memory_space::execution_space;
    using size_type = typename memory_space::size_type;

    using CountView = Kokkos::View<int*, memory_space>;
    using OffsetView = Kokkos::View<size_type*, memory_space>;
    using stencil_type = Cabana::LinkedCellStencil<Scalar>;

    LinkedCellList() = default;

    template <class PositionType>
    LinkedCellList(
        PositionType positions, const Scalar grid_delta[3],
        const Scalar grid_min[3], const Scalar grid_max[3],
        typename std::enable_if<( is_slice<PositionType>::value ||
                                  Kokkos::is_view<PositionType>::value ),
                                int>::type* = 0 )
        : _begin( 0 )
        , _end( size( positions ) )
        , _grid( grid_min[0], grid_min[1], grid_min[2], grid_max[0],
                 grid_max[1], grid_max[2], grid_delta[0], grid_delta[1],
                 grid_delta[2] )
        , _cell_stencil( grid_delta[0], 1.0, grid_min, grid_max )
        , _sorted( false )
    {
        std::size_t np = size( positions );
        allocate( totalBins(), np );
        build( positions, 0, np );
    }

    template <class PositionType>
    LinkedCellList(
        PositionType positions, const std::size_t begin, const std::size_t end,
        const Scalar grid_delta[3], const Scalar grid_min[3],
        const Scalar grid_max[3],
        typename std::enable_if<( is_slice<PositionType>::value ||
                                  Kokkos::is_view<PositionType>::value ),
                                int>::type* = 0 )
        : _begin( begin )
        , _end( end )
        , _grid( grid_min[0], grid_min[1], grid_min[2], grid_max[0],
                 grid_max[1], grid_max[2], grid_delta[0], grid_delta[1],
                 grid_delta[2] )
        , _cell_stencil( grid_delta[0], 1.0, grid_min, grid_max )
        , _sorted( false )
    {
        allocate( totalBins(), end - begin );
        build( positions, begin, end );
    }

    template <class PositionType>
    LinkedCellList(
        PositionType positions, const Scalar grid_delta[3],
        const Scalar grid_min[3], const Scalar grid_max[3],
        const Scalar neighborhood_radius, const Scalar cell_size_ratio = 1,
        typename std::enable_if<( is_slice<PositionType>::value ||
                                  Kokkos::is_view<PositionType>::value ),
                                int>::type* = 0 )
        : _begin( 0 )
        , _end( size( positions ) )
        , _grid( grid_min[0], grid_min[1], grid_min[2], grid_max[0],
                 grid_max[1], grid_max[2], grid_delta[0], grid_delta[1],
                 grid_delta[2] )
        , _cell_stencil( neighborhood_radius, cell_size_ratio, grid_min,
                         grid_max )
        , _sorted( false )
    {
        std::size_t np = size( positions );
        allocate( totalBins(), np );
        build( positions, 0, np );
    }

    template <class PositionType>
    LinkedCellList(
        PositionType positions, const std::size_t begin, const std::size_t end,
        const Scalar grid_delta[3], const Scalar grid_min[3],
        const Scalar grid_max[3], const Scalar neighborhood_radius,
        const Scalar cell_size_ratio = 1,
        typename std::enable_if<( is_slice<PositionType>::value ||
                                  Kokkos::is_view<PositionType>::value ),
                                int>::type* = 0 )
        : _begin( begin )
        , _end( end )
        , _grid( grid_min[0], grid_min[1], grid_min[2], grid_max[0],
                 grid_max[1], grid_max[2], grid_delta[0], grid_delta[1],
                 grid_delta[2] )
        , _cell_stencil( neighborhood_radius, cell_size_ratio, grid_min,
                         grid_max )
        , _sorted( false )
    {
        allocate( totalBins(), end - begin );
        build( positions, begin, end );
    }

    KOKKOS_INLINE_FUNCTION
    int numParticles() const { return _permutes.extent( 0 ); }

    KOKKOS_INLINE_FUNCTION
    std::size_t getParticleBegin() const { return _begin; }

    KOKKOS_INLINE_FUNCTION
    std::size_t getParticleEnd() const { return _end; }

    KOKKOS_INLINE_FUNCTION
    int totalBins() const { return _grid.totalNumCells(); }

    KOKKOS_INLINE_FUNCTION
    int numBin( const int dim ) const { return _grid.numBin( dim ); }

    KOKKOS_INLINE_FUNCTION
    size_type cardinalBinIndex( const int i, const int j, const int k ) const
    {
        return _grid.cardinalCellIndex( i, j, k );
    }

    KOKKOS_INLINE_FUNCTION
    void ijkBinIndex( const int cardinal, int& i, int& j, int& k ) const
    {
        _grid.ijkBinIndex( cardinal, i, j, k );
    }

    KOKKOS_INLINE_FUNCTION
    int binSize( const int i, const int j, const int k ) const
    {
        return _bin_data.binSize( cardinalBinIndex( i, j, k ) );
    }

    KOKKOS_INLINE_FUNCTION
    size_type binOffset( const int i, const int j, const int k ) const
    {
        return _bin_data.binOffset( cardinalBinIndex( i, j, k ) );
    }

    KOKKOS_INLINE_FUNCTION
    size_type permutation( const int particle_id ) const
    {
        return _bin_data.permutation( particle_id );
    }

    KOKKOS_INLINE_FUNCTION
    std::size_t rangeBegin() const { return _bin_data.rangeBegin(); }

    KOKKOS_INLINE_FUNCTION
    std::size_t rangeEnd() const { return _bin_data.rangeEnd(); }

    KOKKOS_INLINE_FUNCTION
    stencil_type cellStencil() const { return _cell_stencil; }

    BinningData<MemorySpace> binningData() const { return _bin_data; }

    template <class ExecutionSpace, class PositionType>
    void build( ExecutionSpace, PositionType positions, const std::size_t begin,
                const std::size_t end )
    {
        Kokkos::Profiling::ScopedRegion region(
            "Cabana::LinkedCellList::build" );
 
        static_assert( is_accessible_from<memory_space, ExecutionSpace>{}, "" );
        assert( end >= begin );
        assert( end <= size( positions ) );
 
        // Resize the binning data. Note that the permutation vector spans
        // only the length of begin-end;
        std::size_t ncell = totalBins();
        if ( _counts.extent( 0 ) != ncell )
        {
            Kokkos::resize( _counts, ncell );
            Kokkos::resize( _offsets, ncell );
        }
        std::size_t nparticles = end - begin;
        if ( _permutes.extent( 0 ) != nparticles )
        {
            Kokkos::resize( _permutes, nparticles );
        }
 
        // Get local copies of class data for lambda function capture.
        auto grid = _grid;
        auto counts = _counts;
        auto offsets = _offsets;
        auto permutes = _permutes;
 
        // Count.
        Kokkos::RangePolicy<ExecutionSpace> particle_range( begin, end );
        Kokkos::deep_copy( _counts, 0 );
        auto counts_sv = Kokkos::Experimental::create_scatter_view( _counts );
        auto cell_count = KOKKOS_LAMBDA( const std::size_t p )
        {
            int i, j, k;
            grid.locatePoint( positions( p, 0 ), positions( p, 1 ),
                              positions( p, 2 ), i, j, k );
            auto counts_data = counts_sv.access();
            counts_data( grid.cardinalCellIndex( i, j, k ) ) += 1;
        };
        Kokkos::parallel_for( "Cabana::LinkedCellList::build::cell_count",
                              particle_range, cell_count );
        Kokkos::fence();
        Kokkos::Experimental::contribute( _counts, counts_sv );
 
        // Compute offsets.
        Kokkos::RangePolicy<ExecutionSpace> cell_range( 0, ncell );
        auto offset_scan = KOKKOS_LAMBDA( const std::size_t c, int& update,
                                          const bool final_pass )
        {
            if ( final_pass )
                offsets( c ) = update;
            update += counts( c );
        };
        Kokkos::parallel_scan( "Cabana::LinkedCellList::build::offset_scan",
                               cell_range, offset_scan );
        Kokkos::fence();
 
        // Reset counts.
        Kokkos::deep_copy( _counts, 0 );
 
        // Compute the permutation vector.
        auto create_permute = KOKKOS_LAMBDA( const std::size_t p )
        {
            int i, j, k;
            grid.locatePoint( positions( p, 0 ), positions( p, 1 ),
                              positions( p, 2 ), i, j, k );
            auto cell_id = grid.cardinalCellIndex( i, j, k );
            int c = Kokkos::atomic_fetch_add( &counts( cell_id ), 1 );
            permutes( offsets( cell_id ) + c ) = p;
        };
        Kokkos::parallel_for( "Cabana::LinkedCellList::build::create_permute",
                              particle_range, create_permute );
        Kokkos::fence();
 
        // Create the binning data.
        _bin_data = BinningData<MemorySpace>( begin, end, _counts, _offsets,
                                              _permutes );
 
        // Store the bin per particle (for neighbor iteration).
        storeParticleBins();
    }

    template <class PositionType>
    void build( PositionType positions, const std::size_t begin,
                const std::size_t end )
    {
        // Use the default execution space.
        build( execution_space{}, positions, begin, end );
    }

    template <class PositionType>
    void build( PositionType positions )
    {
        build( positions, 0, size( positions ) );
    }

    void storeParticleBins()
    {
        Kokkos::parallel_for(
            "Cabana::LinkedCellList::storeBinIndices",
            Kokkos::RangePolicy<execution_space>( 0, totalBins() ), *this );
    }

    auto getParticleBins() const { return _particle_bins; }

    KOKKOS_INLINE_FUNCTION
    auto getParticleBin( const int particle_index ) const
    {
        assert( particle_index >= static_cast<int>( _begin ) );
        assert( particle_index < static_cast<int>( _end ) );
        return _particle_bins( particle_index - _begin );
    }

    KOKKOS_FUNCTION void operator()( const int i ) const
    {
        int bin_ijk[3];
        ijkBinIndex( i, bin_ijk[0], bin_ijk[1], bin_ijk[2] );
        auto offset = binOffset( bin_ijk[0], bin_ijk[1], bin_ijk[2] );
        auto size = binSize( bin_ijk[0], bin_ijk[1], bin_ijk[2] );
        for ( size_t p = offset; p < offset + size; ++p )
        {
            if ( _sorted )
            {
                _particle_bins( p ) = i;
            }
            else
            {
                _particle_bins( permutation( p ) - _begin ) = i;
            }
        }
    }

    void update( const bool sorted ) { _sorted = sorted; }

    KOKKOS_INLINE_FUNCTION
    auto sorted() const { return _sorted; }

    KOKKOS_INLINE_FUNCTION
    void getStencilCells( const int cell, int& imin, int& imax, int& jmin,
                          int& jmax, int& kmin, int& kmax ) const
    {
        _cell_stencil.getCells( cell, imin, imax, jmin, jmax, kmin, kmax );
    }

    KOKKOS_INLINE_FUNCTION
    auto getParticle( const int offset ) const
    {
        std::size_t j;
        if ( !sorted() )
            j = permutation( offset );
        else
            j = offset + getParticleBegin();
        return j;
    }
 
  private:
    std::size_t _begin;
    std::size_t _end;
 
    // Building the linked cell.
    BinningData<MemorySpace> _bin_data;
    Impl::CartesianGrid<Scalar> _grid;
 
    CountView _counts;
    OffsetView _offsets;
    OffsetView _permutes;
 
    // Iterating over the linked cell.
    stencil_type _cell_stencil;
 
    bool _sorted;
    CountView _particle_bins;
 
    void allocate( const int ncell, const int nparticles )
    {
        _counts = CountView(
            Kokkos::view_alloc( Kokkos::WithoutInitializing, "counts" ),
            ncell );
        _offsets = OffsetView(
            Kokkos::view_alloc( Kokkos::WithoutInitializing, "offsets" ),
            ncell );
        _permutes = OffsetView(
            Kokkos::view_alloc( Kokkos::WithoutInitializing, "permutes" ),
            nparticles );
        // This is only used for iterating over the particles (not building the
        // permutaion vector).
        _particle_bins = CountView(
            Kokkos::view_alloc( Kokkos::WithoutInitializing, "counts" ),
            nparticles );
    }
};

template <class MemorySpace, class PositionType, class Scalar>
auto createLinkedCellList( PositionType positions, const Scalar grid_delta[3],
                           const Scalar grid_min[3], const Scalar grid_max[3] )
{
    return LinkedCellList<MemorySpace, Scalar>( positions, grid_delta, grid_min,
                                                grid_max );
}

template <class MemorySpace, class PositionType, class Scalar>
auto createLinkedCellList( PositionType positions, const std::size_t begin,
                           const std::size_t end, const Scalar grid_delta[3],
                           const Scalar grid_min[3], const Scalar grid_max[3] )
{
    return LinkedCellList<MemorySpace, Scalar>(
        positions, begin, end, grid_delta, grid_min, grid_max );
}

template <class MemorySpace, class PositionType, class Scalar>
auto createLinkedCellList( PositionType positions, const Scalar grid_delta[3],
                           const Scalar grid_min[3], const Scalar grid_max[3],
                           const Scalar neighborhood_radius,
                           const Scalar cell_size_ratio = 1.0 )
{
    return LinkedCellList<MemorySpace, Scalar>( positions, grid_delta, grid_min,
                                                grid_max, neighborhood_radius,
                                                cell_size_ratio );
}

template <class MemorySpace, class PositionType, class Scalar>
auto createLinkedCellList( PositionType positions, const std::size_t begin,
                           const std::size_t end, const Scalar grid_delta[3],
                           const Scalar grid_min[3], const Scalar grid_max[3],
                           const Scalar neighborhood_radius,
                           const Scalar cell_size_ratio = 1.0 )
{
    return LinkedCellList<MemorySpace, Scalar>(
        positions, begin, end, grid_delta, grid_min, grid_max,
        neighborhood_radius, cell_size_ratio );
}
 
//---------------------------------------------------------------------------//
template <typename>
struct is_linked_cell_list_impl : public std::false_type
{
};
 
template <typename MemorySpace, typename Scalar>
struct is_linked_cell_list_impl<LinkedCellList<MemorySpace, Scalar>>
    : public std::true_type
{
};

template <class T>
struct is_linked_cell_list
    : public is_linked_cell_list_impl<typename std::remove_cv<T>::type>::type
{
};
 
//---------------------------------------------------------------------------//
template <class LinkedCellListType, class PositionType>
void permute(
    LinkedCellListType& linked_cell_list, PositionType& positions,
    typename std::enable_if<( is_linked_cell_list<LinkedCellListType>::value &&
                              ( is_aosoa<PositionType>::value ||
                                is_slice<PositionType>::value ||
                                Kokkos::is_view<PositionType>::value ) ),
                            int>::type* = 0 )
{
    permute( linked_cell_list.binningData(), positions );
 
    // Update internal state.
    linked_cell_list.update( true );
 
    linked_cell_list.storeParticleBins();
}
 
//---------------------------------------------------------------------------//
template <class MemorySpace, typename Scalar>
class NeighborList<LinkedCellList<MemorySpace, Scalar>>
{
  public:
    using memory_space = MemorySpace;
    using list_type = LinkedCellList<MemorySpace, Scalar>;

    KOKKOS_INLINE_FUNCTION static std::size_t
    totalNeighbor( const list_type& list )
    {
        std::size_t total_n = 0;
        // Sum neighbors across all particles in range.
        for ( std::size_t p = list.getParticleBegin();
              p < list.getParticleEnd(); p++ )
            total_n += numNeighbor( list, p );
        return total_n;
    }

    KOKKOS_INLINE_FUNCTION
    static std::size_t maxNeighbor( const list_type& list )
    {
        std::size_t max_n = 0;
        // Max neighbors across all particles in range.
        for ( std::size_t p = list.getParticleBegin();
              p < list.getParticleEnd(); p++ )
            if ( numNeighbor( list, p ) > max_n )
                max_n = numNeighbor( list, p );
        return max_n;
    }

    KOKKOS_INLINE_FUNCTION static std::size_t
    numNeighbor( const list_type& list, const std::size_t particle_index )
    {
        int total_count = 0;
        int imin, imax, jmin, jmax, kmin, kmax;
        list.getStencilCells( list.getParticleBin( particle_index ), imin, imax,
                              jmin, jmax, kmin, kmax );
 
        // Loop over the cell stencil.
        for ( int i = imin; i < imax; ++i )
            for ( int j = jmin; j < jmax; ++j )
                for ( int k = kmin; k < kmax; ++k )
                {
                    total_count += list.binSize( i, j, k );
                }
        return total_count;
    }

    KOKKOS_INLINE_FUNCTION static std::size_t
    getNeighbor( const list_type& list, const std::size_t particle_index,
                 const std::size_t neighbor_index )
    {
        std::size_t total_count = 0;
        std::size_t previous_count = 0;
        int imin, imax, jmin, jmax, kmin, kmax;
        list.getStencilCells( list.getParticleBin( particle_index ), imin, imax,
                              jmin, jmax, kmin, kmax );
 
        // Loop over the cell stencil.
        for ( int i = imin; i < imax; ++i )
            for ( int j = jmin; j < jmax; ++j )
                for ( int k = kmin; k < kmax; ++k )
                {
                    total_count += list.binSize( i, j, k );
                    // This neighbor is in this bin.
                    if ( total_count > neighbor_index )
                    {
                        int particle_id = list.binOffset( i, j, k ) +
                                          ( neighbor_index - previous_count );
                        return list.getParticle( particle_id );
                    }
                    previous_count = total_count;
                }
 
        assert( total_count <= totalNeighbor( list ) );
 
        // Should never make it to this point.
        return 0;
    }
};
 
} // end namespace Cabana
 
#endif // end CABANA_SORT_HPP