Particle Data Structure

Creating a general AoSoA data structure with general types.

packages/Core/src/Cabana_AoSoA.hpp

@@ -4,6 +4,7 @@
 #include <Cabana_SoA.hpp>
 
 #include <type_traits>
+#include <cmath>
 
 namespace Cabana
 {
@@ -100,16 +101,16 @@ class AoSoA
     // elements). This may or not be equal to the total number of allocated
     // elements in the AoSoA if the size of the struct arrays is not evenly
     // divisible by the requested size.
-    size_t size() const;
+    size_t size() const { return _size; }
 
     // Get the allocated size of the AoSoA (global number of allocated
     // elements). This may or not be equal to the total number of requested
     // elements in the AoSoA if the size of the struct arrays is not evenly
     // divisible by the requested size.
-    size_t allocatedSize() const;
+    size_t allocatedSize() const { return _num_struct * array_size; }
 
     // Get the number of structs in the array.
-    size_t numStruct() const;
+    size_t numStruct() const { return _num_struct; }
 
     // Get the beginning of the array.
     Index begin() const;
@@ -167,7 +168,7 @@ class AoSoA
     // Number of structs in the array.
     size_t _num_struct;
 
-    // Structs.
+    // Structs-of-Arrays.
     struct_type* _structs;
 };
 
@@ -175,4 +176,116 @@ class AoSoA
 
 } // end namespace Cabana
 
+//---------------------------------------------------------------------------//
+// Implementation.
+//---------------------------------------------------------------------------//
+
+namespace Cabana
+{
+//---------------------------------------------------------------------------//
+// Constructor.
+template<typename Device, size_t ArraySize, typename... Types>
+AoSoA<Device,ArraySize,Types...>::AoSoA( const size_t size )
+    : _size( size )
+    , _num_struct( std::ceil(size/ArraySize) )
+{
+    Device::allocate( _structs, _num_struct );
+}
+
+//---------------------------------------------------------------------------//
+// Get the beginning of the array.
+template<typename Device, size_t ArraySize, typename... Types>
+Index AoSoA<Device,ArraySize,Types...>::begin() const
+{
+    Index idx;
+    idx._s = 0;
+    idx._i = 0;
+    return idx;
+}
+
+//---------------------------------------------------------------------------//
+// Get the end of the array.
+template<typename Device, size_t ArraySize, typename... Types>
+Index AoSoA<Device,ArraySize,Types...>::end() const
+{
+    Index idx;
+    idx._s = _num_struct - 1;
+    idx._i = _size % ArraySize;
+    return idx;
+}
+
+//---------------------------------------------------------------------------//
+// Access the data array at a given struct member index.
+template<size_t I>
+template<typename Device, size_t ArraySize, typename... Types>
+inline struct_member_array_type<I>
+AoSoA<Device,ArraySize,Types...>::array( const ParticleIndex& idx )
+{
+    return getStructMember<I>( _structs[idx._s] );
+}
+
+//---------------------------------------------------------------------------//
+// Access the data value at a given struct member index and offset index
+// for Rank 0 data.
+template<size_t I>
+template<typename Device, size_t ArraySize, typename... Types>
+inline
+std::enable_if<(1==std::rank<struct_member_array_type<I>>),
+               struct_member_reference_type<I> >::type
+AoSoA<Device,ArraySize,Types...>::operator()( const ParticleIndex& idx )
+{
+    return array(idx)[idx._i];
+}
+
+//---------------------------------------------------------------------------//
+// Access the data value at a given struct member index and offset index
+// for Rank 1 data.
+template<size_t I>
+template<typename Device, size_t ArraySize, typename... Types>
+inline
+std::enable_if<(2==std::rank<struct_member_array_type<I>>),
+               struct_member_reference_type<I> >::type
+AoSoA<Device,ArraySize,Types...>::operator()( const ParticleIndex& idx,
+                                              const int d0 )
+{
+    return array(idx)[idx._i][d0];
+}
+
+//---------------------------------------------------------------------------//
+// Access the data value at a given struct member index and offset index
+// for Rank 2 data.
+template<size_t I>
+template<typename Device, size_t ArraySize, typename... Types>
+inline
+std::enable_if<(3==std::rank<struct_member_array_type<I>>),
+               struct_member_reference_type<I> >::type
+AoSoA<Device,ArraySize,Types...>::operator()( const ParticleIndex& idx,
+                                              const int d0,
+                                              const int d1 )
+{
+    return array(idx)[idx._i][d0][d1];
+}
+
+//---------------------------------------------------------------------------//
+// Access the data value at a given struct member index and offset index
+// for Rank 3 data.
+template<size_t I>
+template<typename Device, size_t ArraySize, typename... Types>
+inline
+std::enable_if<(4==std::rank<struct_member_array_type<I>>),
+               struct_member_reference_type<I> >::type
+AoSoA<Device,ArraySize,Types...>::operator()( const ParticleIndex& idx,
+                                              const int d0,
+                                              const int d1,
+                                              const int d2 )
+{
+    return array(idx)[idx._i][d0][d1][d2];
+}
+
+//---------------------------------------------------------------------------//
+
+} // end namespace Cabana
+
+//---------------------------------------------------------------------------//
+
 #endif // CABANA_AOSOA_HPP