Cabana_ParticleList.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_PARTICLELIST_HPP
#define CABANA_PARTICLELIST_HPP
 
#include <Cabana_AoSoA.hpp>
#include <Cabana_Fields.hpp>
#include <Cabana_MemberTypes.hpp>
#include <Cabana_SoA.hpp>
#include <Cabana_Tuple.hpp>
 
#include <memory>
#include <string>
#include <type_traits>
 
namespace Cabana
{
//---------------------------------------------------------------------------//
template <class... FieldTags>
struct ParticleTraits
{
    using member_types = Cabana::MemberTypes<typename FieldTags::data_type...>;
};
 
//---------------------------------------------------------------------------//
template <class... FieldTags>
struct Particle
{
    using traits = ParticleTraits<FieldTags...>;
    using tuple_type = Cabana::Tuple<typename traits::member_types>;
    static constexpr int vector_length = 1;

    Particle() = default;

    KOKKOS_FORCEINLINE_FUNCTION
    Particle( const tuple_type& tuple )
        : _tuple( tuple )
    {
    }

    KOKKOS_FORCEINLINE_FUNCTION
    tuple_type& tuple() { return _tuple; }

    KOKKOS_FORCEINLINE_FUNCTION
    const tuple_type& tuple() const { return _tuple; }

    tuple_type _tuple;
};
 
//---------------------------------------------------------------------------//
template <int VectorLength, class... FieldTags>
struct ParticleView
{
    using traits = ParticleTraits<FieldTags...>;
    using soa_type = Cabana::SoA<typename traits::member_types, VectorLength>;
    static constexpr int vector_length = VectorLength;

    ParticleView() = default;

    KOKKOS_FORCEINLINE_FUNCTION
    ParticleView( soa_type& soa, const int vector_index )
        : _soa( soa )
        , _vector_index( vector_index )
    {
    }

    KOKKOS_FORCEINLINE_FUNCTION
    soa_type& soa() { return _soa; }

    KOKKOS_FORCEINLINE_FUNCTION
    const soa_type& soa() const { return _soa; }

    KOKKOS_FORCEINLINE_FUNCTION
    int vectorIndex() const { return _vector_index; }

    soa_type& _soa;

    int _vector_index;
};
 
//---------------------------------------------------------------------------//
// Particle accessors.
//---------------------------------------------------------------------------//
template <class FieldTag, class... FieldTags, class... IndexTypes>
KOKKOS_FORCEINLINE_FUNCTION typename std::enable_if<
    sizeof...( IndexTypes ) == FieldTag::rank,
    typename Particle<FieldTags...>::tuple_type::
        template member_const_reference_type<
            TypeIndexer<FieldTag, FieldTags...>::index>>::type
get( const Particle<FieldTags...>& particle, FieldTag, IndexTypes... indices )
{
    return Cabana::get<TypeIndexer<FieldTag, FieldTags...>::index>(
        particle.tuple(), indices... );
}

template <class FieldTag, class... FieldTags, class... IndexTypes>
KOKKOS_FORCEINLINE_FUNCTION typename std::enable_if<
    sizeof...( IndexTypes ) == FieldTag::rank,
    typename Particle<FieldTags...>::tuple_type::template member_reference_type<
        TypeIndexer<FieldTag, FieldTags...>::index>>::type
get( Particle<FieldTags...>& particle, FieldTag, IndexTypes... indices )
{
    return Cabana::get<TypeIndexer<FieldTag, FieldTags...>::index>(
        particle.tuple(), indices... );
}
 
//---------------------------------------------------------------------------//
template <class FieldTag, class... FieldTags, class... IndexTypes,
          int VectorLength>
KOKKOS_FORCEINLINE_FUNCTION typename std::enable_if<
    sizeof...( IndexTypes ) == FieldTag::rank,
    typename ParticleView<VectorLength, FieldTags...>::soa_type::
        template member_const_reference_type<
            TypeIndexer<FieldTag, FieldTags...>::index>>::type
get( const ParticleView<VectorLength, FieldTags...>& particle, FieldTag,
     IndexTypes... indices )
{
    return Cabana::get<TypeIndexer<FieldTag, FieldTags...>::index>(
        particle.soa(), particle.vectorIndex(), indices... );
}

template <class FieldTag, class... FieldTags, class... IndexTypes,
          int VectorLength>
KOKKOS_FORCEINLINE_FUNCTION typename std::enable_if<
    sizeof...( IndexTypes ) == FieldTag::rank,
    typename ParticleView<VectorLength, FieldTags...>::soa_type::
        template member_reference_type<
            TypeIndexer<FieldTag, FieldTags...>::index>>::type
get( ParticleView<VectorLength, FieldTags...>& particle, FieldTag,
     IndexTypes... indices )
{
    return Cabana::get<TypeIndexer<FieldTag, FieldTags...>::index>(
        particle.soa(), particle.vectorIndex(), indices... );
}
 
//---------------------------------------------------------------------------//
template <class MemorySpace, int VectorLength, class... FieldTags>
class ParticleList
{
  public:
    using memory_space = MemorySpace;
    using traits = ParticleTraits<FieldTags...>;
    using member_types = typename traits::member_types;
    using aosoa_type = Cabana::AoSoA<member_types, memory_space, VectorLength>;
    using tuple_type = typename aosoa_type::tuple_type;
    template <std::size_t M>
    using slice_type = typename aosoa_type::template member_slice_type<M>;
    using particle_type = Particle<FieldTags...>;
    using particle_view_type =
        ParticleView<aosoa_type::vector_length, FieldTags...>;

    ParticleList( const std::string& label )
        : _aosoa( label )
    {
    }

    ParticleList( const aosoa_type& aosoa )
        : _aosoa( aosoa )
    {
    }

    KOKKOS_INLINE_FUNCTION std::size_t size() const { return _aosoa.size(); }
    void resize( const std::size_t n ) { _aosoa.resize( n ); }

    aosoa_type& aosoa() { return _aosoa; }
    KOKKOS_INLINE_FUNCTION const aosoa_type& aosoa() const { return _aosoa; }

    template <class IndexType>
    KOKKOS_INLINE_FUNCTION auto getParticle( IndexType p ) const
    {
        return particle_type( _aosoa.getTuple( p ) );
    }

    template <class ParticleType, class IndexType>
    KOKKOS_INLINE_FUNCTION void setParticle( ParticleType particle,
                                             IndexType p ) const
    {
        _aosoa.setTuple( p, particle.tuple() );
    }

    template <class IndexType>
    KOKKOS_INLINE_FUNCTION auto getParticleView( IndexType p ) const
    {
        auto s = Cabana::Impl::Index<particle_view_type::vector_length>::s( p );
        auto a = Cabana::Impl::Index<particle_view_type::vector_length>::a( p );
        return particle_view_type( _aosoa.access( s ), a );
    }

    const std::string& label() const { return _aosoa.label(); }

    template <class FieldTag>
    slice_type<TypeIndexer<FieldTag, FieldTags...>::index>
    slice( FieldTag ) const
    {
        return Cabana::slice<TypeIndexer<FieldTag, FieldTags...>::index>(
            _aosoa, FieldTag::label() );
    }
 
  protected:
    aosoa_type _aosoa;
};
 
template <class>
struct is_particle_list_impl : public std::false_type
{
};
 
template <class MemorySpace, int VectorLength, class... FieldTags>
struct is_particle_list_impl<
    ParticleList<MemorySpace, VectorLength, FieldTags...>>
    : public std::true_type
{
};

template <class T>
struct is_particle_list
    : public is_particle_list_impl<typename std::remove_cv<T>::type>::type
{
};
 
//---------------------------------------------------------------------------//
template <class MemorySpace, int VectorLength, class... FieldTags>
auto createParticleList( const std::string& label,
                         ParticleTraits<FieldTags...> )
{
    return ParticleList<MemorySpace, VectorLength, FieldTags...>( label );
}

template <class MemorySpace, class... FieldTags>
auto createParticleList( const std::string& label,
                         ParticleTraits<FieldTags...> )
{
    return ParticleList<
        MemorySpace,
        Impl::PerformanceTraits<
            typename MemorySpace::execution_space>::vector_length,
        FieldTags...>( label );
}
 
//---------------------------------------------------------------------------//
 
} // end namespace Cabana
 
#endif // end CABANA_PARTICLELIST_HPP