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, std::size_t NumSpaceDim = 3>
struct LinkedCellStencil
{
    static constexpr std::size_t num_space_dim = NumSpaceDim;

    Impl::CartesianGrid<Scalar, num_space_dim> grid;
    int max_cells_dir;
    int max_cells;
    int cell_range;

    LinkedCellStencil() = default;

    template <class ArrayType>
    LinkedCellStencil( const Scalar neighborhood_radius,
                       const Scalar cell_size_ratio, const ArrayType& grid_min,
                       const ArrayType& grid_max )
    {
        Scalar dx = neighborhood_radius * cell_size_ratio;
        grid = Impl::CartesianGrid<Scalar, num_space_dim>( grid_min, grid_max,
                                                           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;
    }

    LinkedCellStencil( const Scalar neighborhood_radius,
                       const Scalar cell_size_ratio,
                       const Scalar grid_min[num_space_dim],
                       const Scalar grid_max[num_space_dim] )
    {
        Scalar dx = neighborhood_radius * cell_size_ratio;
        grid = Impl::CartesianGrid<Scalar, num_space_dim>( grid_min, grid_max,
                                                           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;
    }

    template <std::size_t NSD = num_space_dim>
    KOKKOS_INLINE_FUNCTION std::enable_if_t<3 == NSD, 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._nx[2] ) ? k + cell_range + 1
                                                    : grid._nx[2];
 
        jmin = ( j - cell_range > 0 ) ? j - cell_range : 0;
        jmax = ( j + cell_range + 1 < grid._nx[1] ) ? j + cell_range + 1
                                                    : grid._nx[1];
 
        imin = ( i - cell_range > 0 ) ? i - cell_range : 0;
        imax = ( i + cell_range + 1 < grid._nx[0] ) ? i + cell_range + 1
                                                    : grid._nx[0];
    }

    KOKKOS_INLINE_FUNCTION void
    getCells( const int cell, Kokkos::Array<int, num_space_dim>& min,
              Kokkos::Array<int, num_space_dim>& max ) const
    {
        Kokkos::Array<int, num_space_dim> ijk;
        grid.ijkBinIndex( cell, ijk );
 
        for ( std::size_t d = 0; d < num_space_dim; ++d )
        {
            min[d] = ( ijk[d] - cell_range > 0 ) ? ijk[d] - cell_range : 0;
            max[d] = ( ijk[d] + cell_range + 1 < grid._nx[d] )
                         ? ijk[d] + cell_range + 1
                         : grid._nx[d];
        }
    }
};
 
//---------------------------------------------------------------------------//
template <class MemorySpace, class Scalar = double, std::size_t NumSpaceDim = 3>
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 = LinkedCellStencil<Scalar, NumSpaceDim>;

    static constexpr std::size_t num_space_dim = NumSpaceDim;

    LinkedCellList() = default;

    template <class PositionType,
              template <class, std::size_t, class...> class ArrayType,
              class... Args>
    LinkedCellList(
        PositionType positions,
#ifdef KOKKOS_ENABLE_DEPRECATED_CODE_4
        const ArrayType<Scalar, num_space_dim, Args...> grid_delta,
        const ArrayType<Scalar, num_space_dim, Args...> grid_min,
        const ArrayType<Scalar, num_space_dim, Args...> grid_max,
#else
        const ArrayType<Scalar, num_space_dim> grid_delta,
        const ArrayType<Scalar, num_space_dim> grid_min,
        const ArrayType<Scalar, num_space_dim> grid_max,
#endif
        typename std::enable_if<( is_slice<PositionType>::value ||
                                  Kokkos::is_view<PositionType>::value ),
                                int>::type* = 0 )
        : _begin( 0 )
        , _end( size( positions ) )
        , _grid( grid_min, grid_max, grid_delta )
        , _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,
              template <class, std::size_t, class...> class ArrayType,
              class... Args>
    LinkedCellList(
        PositionType positions, const std::size_t begin, const std::size_t end,
#ifdef KOKKOS_ENABLE_DEPRECATED_CODE_4
        const ArrayType<Scalar, num_space_dim, Args...> grid_delta,
        const ArrayType<Scalar, num_space_dim, Args...> grid_min,
        const ArrayType<Scalar, num_space_dim, Args...> grid_max,
#else
        const ArrayType<Scalar, num_space_dim> grid_delta,
        const ArrayType<Scalar, num_space_dim> grid_min,
        const ArrayType<Scalar, num_space_dim> grid_max,
#endif
        typename std::enable_if<( is_slice<PositionType>::value ||
                                  Kokkos::is_view<PositionType>::value ),
                                int>::type* = 0 )
        : _begin( begin )
        , _end( end )
        , _grid( grid_min, grid_max, grid_delta )
        , _cell_stencil( grid_delta[0], 1.0, grid_min, grid_max )
        , _sorted( false )
    {
        allocate( totalBins(), end - begin );
        build( positions, begin, end );
    }

    template <class PositionType,
              template <class, std::size_t, class...> class ArrayType,
              class... Args>
    LinkedCellList(
        PositionType positions,
#ifdef KOKKOS_ENABLE_DEPRECATED_CODE_4
        const ArrayType<Scalar, num_space_dim, Args...> grid_delta,
        const ArrayType<Scalar, num_space_dim, Args...> grid_min,
        const ArrayType<Scalar, num_space_dim, Args...> grid_max,
#else
        const ArrayType<Scalar, num_space_dim> grid_delta,
        const ArrayType<Scalar, num_space_dim> grid_min,
        const ArrayType<Scalar, num_space_dim> grid_max,
#endif
        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, grid_max, grid_delta )
        , _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,
              template <class, std::size_t, class...> class ArrayType,
              class... Args>
    LinkedCellList(
        PositionType positions, const std::size_t begin, const std::size_t end,
#ifdef KOKKOS_ENABLE_DEPRECATED_CODE_4
        const ArrayType<Scalar, num_space_dim, Args...> grid_delta,
        const ArrayType<Scalar, num_space_dim, Args...> grid_min,
        const ArrayType<Scalar, num_space_dim, Args...> grid_max,
#else
        const ArrayType<Scalar, num_space_dim> grid_delta,
        const ArrayType<Scalar, num_space_dim> grid_min,
        const ArrayType<Scalar, num_space_dim> grid_max,
#endif
        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, grid_max, grid_delta )
        , _cell_stencil( neighborhood_radius, cell_size_ratio, grid_min,
                         grid_max )
        , _sorted( false )
    {
        allocate( totalBins(), end - begin );
        build( positions, begin, end );
    }

    template <class PositionType>
    LinkedCellList(
        PositionType positions, const Scalar grid_delta[num_space_dim],
        const Scalar grid_min[num_space_dim],
        const Scalar grid_max[num_space_dim],
        typename std::enable_if<( is_slice<PositionType>::value ||
                                  Kokkos::is_view<PositionType>::value ),
                                int>::type* = 0 )
        : _begin( 0 )
        , _end( size( positions ) )
        , _grid( grid_min, grid_max, grid_delta )
        , _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[num_space_dim],
        const Scalar grid_min[num_space_dim],
        const Scalar grid_max[num_space_dim],
        typename std::enable_if<( is_slice<PositionType>::value ||
                                  Kokkos::is_view<PositionType>::value ),
                                int>::type* = 0 )
        : _begin( begin )
        , _end( end )
        , _grid( grid_min, grid_max, grid_delta )
        , _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[num_space_dim],
        const Scalar grid_min[num_space_dim],
        const Scalar grid_max[num_space_dim], 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, grid_max, grid_delta )
        , _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[num_space_dim],
        const Scalar grid_min[num_space_dim],
        const Scalar grid_max[num_space_dim], 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, grid_max, grid_delta )
        , _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 ); }

    template <std::size_t NSD = num_space_dim>
    KOKKOS_INLINE_FUNCTION std::enable_if_t<3 == NSD, size_type>
    cardinalBinIndex( const int i, const int j, const int k ) const
    {
        return _grid.cardinalCellIndex( i, j, k );
    }

    KOKKOS_INLINE_FUNCTION
    size_type
    cardinalBinIndex( const Kokkos::Array<int, num_space_dim>& ijk ) const
    {
        return _grid.cardinalCellIndex( ijk );
    }

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

    KOKKOS_INLINE_FUNCTION
    void ijkBinIndex( const int cardinal,
                      Kokkos::Array<int, num_space_dim>& ijk ) const
    {
        _grid.ijkBinIndex( cardinal, ijk );
    }

    KOKKOS_INLINE_FUNCTION
    int binSize( const Kokkos::Array<int, num_space_dim> ijk ) const
    {
        return _bin_data.binSize( cardinalBinIndex( ijk ) );
    }

    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 binOffset( const Kokkos::Array<int, num_space_dim> ijk ) const
    {
        return _bin_data.binOffset( cardinalBinIndex( ijk ) );
    }

    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 )
        {
            Kokkos::Array<int, num_space_dim> ijk;
            Kokkos::Array<Scalar, num_space_dim> pos;
            for ( std::size_t d = 0; d < num_space_dim; ++d )
                pos[d] = positions( p, d );
            grid.locatePoint( pos, ijk );
            auto counts_data = counts_sv.access();
            counts_data( grid.cardinalCellIndex( ijk ) ) += 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 )
        {
            Kokkos::Array<int, num_space_dim> ijk;
            Kokkos::Array<Scalar, num_space_dim> pos;
            for ( std::size_t d = 0; d < num_space_dim; ++d )
                pos[d] = positions( p, d );
            grid.locatePoint( pos, ijk );
            auto cell_id = grid.cardinalCellIndex( ijk );
            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
    {
        Kokkos::Array<int, num_space_dim> bin_ijk;
        ijkBinIndex( i, bin_ijk );
        auto offset = binOffset( bin_ijk );
        auto size = binSize( bin_ijk );
        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; }

    template <std::size_t NSD = num_space_dim>
    KOKKOS_INLINE_FUNCTION std::enable_if_t<3 == NSD, 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
    void getStencilCells( const int cell,
                          Kokkos::Array<int, num_space_dim>& min,
                          Kokkos::Array<int, num_space_dim>& max ) const
    {
        _cell_stencil.getCells( cell, min, max );
    }

    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, num_space_dim> _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 PositionType,
          template <class, std::size_t, class...> class ArrayType,
          class... Args, class Scalar, std::size_t NumSpaceDim>
auto createLinkedCellList(
    PositionType positions,
#ifdef KOKKOS_ENABLE_DEPRECATED_CODE_4
    const ArrayType<Scalar, NumSpaceDim, Args...> grid_delta,
    const ArrayType<Scalar, NumSpaceDim, Args...> grid_min,
    const ArrayType<Scalar, NumSpaceDim, Args...> grid_max
#else
    const ArrayType<Scalar, NumSpaceDim> grid_delta,
    const ArrayType<Scalar, NumSpaceDim> grid_min,
    const ArrayType<Scalar, NumSpaceDim> grid_max
#endif
)
{
    using memory_space = typename PositionType::memory_space;
    using scalar_type = typename PositionType::value_type;
    return LinkedCellList<memory_space, scalar_type, NumSpaceDim>(
        positions, grid_delta, grid_min, grid_max );
}

template <class PositionType,
          template <class, std::size_t, class...> class ArrayType,
          class... Args, class Scalar, std::size_t NumSpaceDim>
auto createLinkedCellList(
    PositionType positions, const std::size_t begin, const std::size_t end,
#ifdef KOKKOS_ENABLE_DEPRECATED_CODE_4
    const ArrayType<Scalar, NumSpaceDim, Args...> grid_delta,
    const ArrayType<Scalar, NumSpaceDim, Args...> grid_min,
    const ArrayType<Scalar, NumSpaceDim, Args...> grid_max
#else
    const ArrayType<Scalar, NumSpaceDim> grid_delta,
    const ArrayType<Scalar, NumSpaceDim> grid_min,
    const ArrayType<Scalar, NumSpaceDim> grid_max
#endif
)
{
    using memory_space = typename PositionType::memory_space;
    using scalar_type = typename PositionType::value_type;
    return LinkedCellList<memory_space, scalar_type, NumSpaceDim>(
        positions, begin, end, grid_delta, grid_min, grid_max );
}

template <class PositionType,
          template <class, std::size_t, class...> class ArrayType,
          class... Args, class Scalar, std::size_t NumSpaceDim>
auto createLinkedCellList(
    PositionType positions,
#ifdef KOKKOS_ENABLE_DEPRECATED_CODE_4
    const ArrayType<Scalar, NumSpaceDim, Args...> grid_delta,
    const ArrayType<Scalar, NumSpaceDim, Args...> grid_min,
    const ArrayType<Scalar, NumSpaceDim, Args...> grid_max,
#else
    const ArrayType<Scalar, NumSpaceDim> grid_delta,
    const ArrayType<Scalar, NumSpaceDim> grid_min,
    const ArrayType<Scalar, NumSpaceDim> grid_max,
#endif
    const typename PositionType::value_type neighborhood_radius,
    const typename PositionType::value_type cell_size_ratio = 1.0 )
{
    using memory_space = typename PositionType::memory_space;
    using scalar_type = typename PositionType::value_type;
    return LinkedCellList<memory_space, scalar_type, NumSpaceDim>(
        positions, grid_delta, grid_min, grid_max, neighborhood_radius,
        cell_size_ratio );
}

template <class PositionType,
          template <class, std::size_t, class...> class ArrayType,
          class... Args, class Scalar, std::size_t NumSpaceDim>
auto createLinkedCellList(
    PositionType positions, const std::size_t begin, const std::size_t end,
#ifdef KOKKOS_ENABLE_DEPRECATED_CODE_4
    const ArrayType<Scalar, NumSpaceDim, Args...> grid_delta,
    const ArrayType<Scalar, NumSpaceDim, Args...> grid_min,
    const ArrayType<Scalar, NumSpaceDim, Args...> grid_max,
#else
    const ArrayType<Scalar, NumSpaceDim> grid_delta,
    const ArrayType<Scalar, NumSpaceDim> grid_min,
    const ArrayType<Scalar, NumSpaceDim> grid_max,
#endif
    const typename PositionType::value_type neighborhood_radius,
    const typename PositionType::value_type cell_size_ratio = 1.0 )
{
    using memory_space = typename PositionType::memory_space;
    using scalar_type = typename PositionType::value_type;
    return LinkedCellList<memory_space, scalar_type, NumSpaceDim>(
        positions, begin, end, grid_delta, grid_min, grid_max,
        neighborhood_radius, cell_size_ratio );
}

template <class PositionType>
auto createLinkedCellList(
    PositionType positions,
    const typename PositionType::value_type grid_delta[3],
    const typename PositionType::value_type grid_min[3],
    const typename PositionType::value_type grid_max[3] )
{
    using memory_space = typename PositionType::memory_space;
    using scalar_type = typename PositionType::value_type;
    return LinkedCellList<memory_space, scalar_type, 3>( positions, grid_delta,
                                                         grid_min, grid_max );
}

template <class PositionType>
auto createLinkedCellList(
    PositionType positions, const std::size_t begin, const std::size_t end,
    const typename PositionType::value_type grid_delta[3],
    const typename PositionType::value_type grid_min[3],
    const typename PositionType::value_type grid_max[3] )
{
    using memory_space = typename PositionType::memory_space;
    using scalar_type = typename PositionType::value_type;
    return LinkedCellList<memory_space, scalar_type, 3>(
        positions, begin, end, grid_delta, grid_min, grid_max );
}

template <class PositionType>
auto createLinkedCellList(
    PositionType positions,
    const typename PositionType::value_type grid_delta[3],
    const typename PositionType::value_type grid_min[3],
    const typename PositionType::value_type grid_max[3],
    const typename PositionType::value_type neighborhood_radius,
    const typename PositionType::value_type cell_size_ratio = 1.0 )
{
    using memory_space = typename PositionType::memory_space;
    using scalar_type = typename PositionType::value_type;
    return LinkedCellList<memory_space, scalar_type, 3>(
        positions, grid_delta, grid_min, grid_max, neighborhood_radius,
        cell_size_ratio );
}

template <class PositionType>
auto createLinkedCellList(
    PositionType positions, const std::size_t begin, const std::size_t end,
    const typename PositionType::value_type grid_delta[3],
    const typename PositionType::value_type grid_min[3],
    const typename PositionType::value_type grid_max[3],
    const typename PositionType::value_type neighborhood_radius,
    const typename PositionType::value_type cell_size_ratio = 1.0 )
{
    using memory_space = typename PositionType::memory_space;
    using scalar_type = typename PositionType::value_type;
    return LinkedCellList<memory_space, scalar_type, 3>(
        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, std::size_t Dim>
struct is_linked_cell_list_impl<LinkedCellList<MemorySpace, Scalar, Dim>>
    : 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, std::size_t NumSpaceDim>
class NeighborList<LinkedCellList<MemorySpace, Scalar, NumSpaceDim>>
{
  public:
    using memory_space = MemorySpace;
    using list_type = LinkedCellList<MemorySpace, Scalar, NumSpaceDim>;
    static constexpr std::size_t num_space_dim = NumSpaceDim;

    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;
    }

    template <std::size_t NSD = num_space_dim>
    KOKKOS_INLINE_FUNCTION static std::enable_if_t<3 == NSD, std::size_t>
    numNeighbor( const list_type& list, const std::size_t particle_index )
    {
        int total_count = 0;
        Kokkos::Array<int, 3> min;
        Kokkos::Array<int, 3> max;
        list.getStencilCells( list.getParticleBin( particle_index ), min, max );
 
        Kokkos::Array<int, 3> ijk;
        for ( int i = min[0]; i < max[0]; ++i )
            for ( int j = min[1]; j < max[1]; ++j )
                for ( int k = min[2]; k < max[2]; ++k )
                {
                    ijk = { i, j, k };
                    total_count += list.binSize( ijk );
                }
 
        return total_count;
    }

    template <std::size_t NSD = num_space_dim>
    KOKKOS_INLINE_FUNCTION static std::enable_if_t<2 == NSD, std::size_t>
    numNeighbor( const list_type& list, const std::size_t particle_index )
    {
        int total_count = 0;
        Kokkos::Array<int, 2> min;
        Kokkos::Array<int, 2> max;
        list.getStencilCells( list.getParticleBin( particle_index ), min, max );
 
        Kokkos::Array<int, 2> ij;
        for ( int i = min[0]; i < max[0]; ++i )
            for ( int j = min[1]; j < max[1]; ++j )
            {
                ij = { i, j };
                total_count += list.binSize( ij );
            }
 
        return total_count;
    }

    template <std::size_t NSD = num_space_dim>
    KOKKOS_INLINE_FUNCTION static std::enable_if_t<3 == NSD, 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;
        Kokkos::Array<int, 3> min;
        Kokkos::Array<int, 3> max;
        list.getStencilCells( list.getParticleBin( particle_index ), min, max );
 
        Kokkos::Array<int, 3> ijk;
        for ( int i = min[0]; i < max[0]; ++i )
            for ( int j = min[1]; j < max[1]; ++j )
                for ( int k = min[2]; k < max[2]; ++k )
                {
                    ijk = { i, j, k };
 
                    total_count += list.binSize( ijk );
                    // This neighbor is in this bin.
                    if ( total_count > neighbor_index )
                    {
                        int particle_id = list.binOffset( ijk ) +
                                          ( neighbor_index - previous_count );
                        return list.getParticle( particle_id );
                    }
                    // Update previous to all bins so far.
                    previous_count = total_count;
                }
 
        assert( total_count <= totalNeighbor( list ) );
 
        // Should never make it to this point.
        return 0;
    }

    template <std::size_t NSD = num_space_dim>
    KOKKOS_INLINE_FUNCTION static std::enable_if_t<2 == NSD, 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;
        Kokkos::Array<int, 2> min;
        Kokkos::Array<int, 2> max;
        list.getStencilCells( list.getParticleBin( particle_index ), min, max );
 
        // Loop over the cell stencil.
        Kokkos::Array<int, 2> ij;
        for ( int i = min[0]; i < max[0]; ++i )
            for ( int j = min[1]; j < max[1]; ++j )
            {
                ij = { i, j };
 
                total_count += list.binSize( ij );
                // This neighbor is in this bin.
                if ( total_count > neighbor_index )
                {
                    int particle_id = list.binOffset( ij ) +
                                      ( neighbor_index - previous_count );
                    return list.getParticle( particle_id );
                }
                // Update previous to all bins so far.
                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