Particle Data Structure

Creating a general AoSoA data structure with general types.

packages/Core/src/Cabana_AoSoA.hpp

@@ -10,6 +10,10 @@
 namespace Cabana
 {
 //---------------------------------------------------------------------------//
+/*!
+  \class AoSoA
+  \brief Array-of-Structs-of-Arrays
+*/
 template<typename Device, std::size_t ArraySize, typename... Types>
 class AoSoA
 {
@@ -19,15 +23,15 @@ class AoSoA
     using device_type = Device;
 
     // Inner array size (size of the arrays held by the structs).
-    static constexpr std::size_t array_size = ArraySize;
+    static constexpr std::size_t soa_array_size = ArraySize;
 
     // SoA type.
-    using struct_type = SoA<array_size,Types...>;
+    using soa_type = SoA<soa_array_size,Types...>;
 
     // Struct member array return type at a given index I.
     template<std::size_t I>
     using struct_member_array_type =
-        typename ArrayTypeAtIndex<I,array_size,Types...>::return_type;
+        typename ArrayTypeAtIndex<I,soa_array_size,Types...>::return_type;
 
     // Struct member array const return type at a given index I.
     template<std::size_t I>
@@ -86,7 +90,7 @@ class AoSoA
         // Prefix increment operator.
         Index& operator++()
         {
-            _i = ( array_size - 1 == _i ) ? 0 : _i + 1;
+            _i = ( soa_array_size - 1 == _i ) ? 0 : _i + 1;
             _s = ( 0 == _i ) ? _s + 1: _s;
             return *this;
         };
@@ -118,16 +122,15 @@ class AoSoA
     AoSoA( const std::size_t size )
         : _size( size )
     {
-        _num_struct = _size / ArraySize;
-        if ( 0 < _size % ArraySize ) ++_num_struct;
-
-        Device::allocate( _structs, _num_struct );
+        _num_struct = std::floor(_size / soa_array_size);
+        if ( 0 < _size % soa_array_size ) ++_num_struct;
+        Device::allocate( _data, _num_struct );
     }
 
     // Destructor.
     ~AoSoA()
     {
-        Device::deallocate( _structs );
+        Device::deallocate( _data );
     }
 
     // Get the requested size of the AoSoA (global number of requested
@@ -140,34 +143,34 @@ class AoSoA
     // 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.
-    std::size_t allocatedSize() const { return _num_struct * array_size; }
+    std::size_t allocatedSize() const { return _num_struct * soa_array_size; }
 
     // Get the number of structs in the array.
     std::size_t numStruct() const { return _num_struct; }
 
-    // Get the beginning of the array.
-    Index begin() const
+    // Access the data array at a given struct member index.
+    template<std::size_t I>
+    inline struct_member_array_type<I> array( const std::size_t s )
     {
-        return Index( 0, 0 );
+        return getStructMember<I>( _data[s] );
     }
 
-    // Get the end of the array.
-    Index end() const
+    template<std::size_t I>
+    inline struct_member_const_array_type<I> array( const std::size_t s ) const
     {
-        return Index( _num_struct - 1, _size % ArraySize );
+        return getStructMember<I>( _data[s] );
     }
 
-    // Access the data array at a given struct member index.
-    template<std::size_t I>
-    inline struct_member_array_type<I> array( const Index& idx )
+    // Get the index at the beginning of the array.
+    Index begin() const
     {
-        return getStructMember<I>( _structs[idx.s()] );
+        return Index( 0, 0 );
     }
 
-    template<std::size_t I>
-    inline struct_member_const_array_type<I> array( const Index& idx ) const
+    // Get the index at end of the array.
+    Index end() const
     {
-        return getStructMember<I>( _structs[idx.s()] );
+        return Index( _num_struct - 1, _size % soa_array_size );
     }
 
     // Access the data value at a given struct member index and offset index
@@ -178,7 +181,7 @@ class AoSoA
                             struct_member_reference_type<I> >::type
     get( const Index& idx )
     {
-        return array<I>(idx)[idx.i()];
+        return array<I>(idx.s())[idx.i()];
     }
 
     template<std::size_t I>
@@ -187,7 +190,7 @@ class AoSoA
                             struct_member_const_reference_type<I> >::type
     get( const Index& idx ) const
     {
-        return array<I>(idx)[idx.i()];
+        return array<I>(idx.s())[idx.i()];
     }
 
     // Access the data value at a given struct member index and offset index
@@ -199,7 +202,7 @@ class AoSoA
     get( const Index& idx,
          const int d0 )
     {
-        return array<I>(idx)[idx.i()][d0];
+        return array<I>(idx.s())[idx.i()][d0];
     }
 
     template<std::size_t I>
@@ -209,7 +212,7 @@ class AoSoA
     get( const Index& idx,
          const int d0 ) const
     {
-        return array<I>(idx)[idx.i()][d0];
+        return array<I>(idx.s())[idx.i()][d0];
     }
 
     // Access the data value at a given struct member index and offset index
@@ -222,7 +225,7 @@ class AoSoA
          const int d0,
          const int d1 )
     {
-        return array<I>(idx)[idx.i()][d0][d1];
+        return array<I>(idx.s())[idx.i()][d0][d1];
     }
 
     template<std::size_t I>
@@ -233,7 +236,7 @@ class AoSoA
          const int d0,
          const int d1 ) const
     {
-        return array<I>(idx)[idx.i()][d0][d1];
+        return array<I>(idx.s())[idx.i()][d0][d1];
     }
 
     // Access the data value at a given struct member index and offset index
@@ -247,7 +250,7 @@ class AoSoA
          const int d1,
          const int d2 )
     {
-        return array<I>(idx)[idx.i()][d0][d1][d2];
+        return array<I>(idx.s())[idx.i()][d0][d1][d2];
     }
 
     template<std::size_t I>
@@ -259,7 +262,7 @@ class AoSoA
          const int d1,
          const int d2 ) const
     {
-        return array<I>(idx)[idx.i()][d0][d1][d2];
+        return array<I>(idx.s())[idx.i()][d0][d1][d2];
     }
 
   private:
@@ -271,7 +274,7 @@ class AoSoA
     std::size_t _num_struct;
 
     // Structs-of-Arrays.
-    struct_type* _structs;
+    soa_type* _data;
 };
 
 //---------------------------------------------------------------------------//