Calculate x+u automatically inside slice

ORNL · Sep 27, 2023 · 8c26f83 · 8c26f83
1 parent b0f56f1
commit 8c26f83
Showing 1 changed file with 21 additions and 0 deletions.
diff --git a/src/CabanaPD_Particles.hpp b/src/CabanaPD_Particles.hpp
@@ -88,6 +88,7 @@ class Particles<DeviceType, PMB, Dimension>
     using self_type = Particles<DeviceType, PMB, Dimension>;
     using device_type = DeviceType;
     using memory_space = typename device_type::memory_space;
+    using execution_space = typename memory_space::execution_space;
     static constexpr int dim = Dimension;
 
     // Per particle.
@@ -305,10 +306,14 @@ class Particles<DeviceType, PMB, Dimension>
     }
     auto sliceCurrentPosition()
     {
+        // Update before returning data.
+        updateCurrentPosition();
         return Cabana::slice<0>( _aosoa_y, "current_positions" );
     }
     auto sliceCurrentPosition() const
     {
+        // Update before returning data.
+        updateCurrentPosition();
         return Cabana::slice<0>( _aosoa_y, "current_positions" );
     }
     auto sliceDisplacement()
@@ -370,6 +375,22 @@ class Particles<DeviceType, PMB, Dimension>
         return Cabana::slice<0>( _aosoa_nofail, "no_fail_region" );
     }
 
+    void updateCurrentPosition()
+    {
+        // Not using slice function because this is called inside.
+        auto y = Cabana::slice<0>( _aosoa_y, "current_positions" );
+        auto x = sliceRefPosition();
+        auto u = sliceDisplacement();
+        Kokkos::RangePolicy<execution_space> policy( 0, n_ghost );
+        auto sum_x_u = KOKKOS_LAMBDA( const std::size_t pid )
+        {
+            for ( int d = 0; d < 3; d++ )
+                y( pid, d ) = x( pid, d ) + u( pid, d );
+        };
+        Kokkos::parallel_for( "CabanaPD::CalculateCurrentPositions", policy,
+                              sum_x_u );
+    }
+
     void resize( int new_local, int new_ghost )
     {
         n_local = new_local;