WIP Add particle-grid periodic migrate and gather

ECP-copa · Jul 9, 2020 · fc4eb4e · fc4eb4e
1 parent dabc430
commit fc4eb4e
Show file tree

Hide file tree

Showing 6 changed files with 750 additions and 2 deletions.
diff --git a/core/src/CMakeLists.txt b/core/src/CMakeLists.txt
@@ -32,6 +32,12 @@ if(MPI_FOUND)
     )
 endif()
 
+if(Cabana_ENABLE_CAJITA)
+  list(APPEND HEADERS_PUBLIC
+    Cabana_ParticleGridCommunication.hpp
+    )
+endif()
+
 set(HEADERS_IMPL
   impl/Cabana_CartesianGrid.hpp
   impl/Cabana_Index.hpp

diff --git a/core/src/Cabana_Core.hpp b/core/src/Cabana_Core.hpp
@@ -33,6 +33,10 @@
 #include <Cabana_Halo.hpp>
 #endif
 
+#ifdef Cabana_ENABLE_CAJITA
+#include <Cabana_ParticleGridCommunication.hpp>
+#endif
+
 #ifdef Cabana_ENABLE_ARBORX
 #include <Cabana_Experimental_NeighborList.hpp>
 #endif

diff --git a/core/src/Cabana_Halo.hpp b/core/src/Cabana_Halo.hpp
@@ -222,9 +222,11 @@ struct is_halo : public is_halo_impl<typename std::remove_cv<T>::type>::type
   elements are expected to appear first (i.e. in the first halo.numLocal()
   elements) and the ghosted elements are expected to appear second (i.e. in
   the next halo.numGhost() elements()).
+
+  \param functor A functor to modify particle properties in the send buffer.
 */
-template <class Halo_t, class AoSoA_t>
-void gather( const Halo_t &halo, AoSoA_t &aosoa,
+template <class Halo_t, class AoSoA_t, class Functor_t>
+void gather( const Halo_t &halo, AoSoA_t &aosoa, Functor_t modify_buffer,
              typename std::enable_if<( is_halo<Halo_t>::value &&
                                        is_aosoa<AoSoA_t>::value ),
                                      int>::type * = 0 )
@@ -246,6 +248,7 @@ void gather( const Halo_t &halo, AoSoA_t &aosoa,
     auto gather_send_buffer_func = KOKKOS_LAMBDA( const std::size_t i )
     {
         send_buffer( i ) = aosoa.getTuple( steering( i ) );
+        modify_buffer( i );
     };
     Kokkos::RangePolicy<typename Halo_t::execution_space>
         gather_send_buffer_policy( 0, halo.totalNumExport() );
@@ -320,6 +323,15 @@ void gather( const Halo_t &halo, AoSoA_t &aosoa,
     MPI_Barrier( halo.comm() );
 }
 
+template <class Halo_t, class AoSoA_t>
+void gather( const Halo_t &halo, AoSoA_t &aosoa,
+             typename std::enable_if<( is_halo<Halo_t>::value &&
+                                       is_aosoa<AoSoA_t>::value ),
+                                     int>::type * = 0 )
+{
+    gather( halo, aosoa, []( const int i ) {} );
+}
+
 //---------------------------------------------------------------------------//
 /*!
   \brief Synchronously gather data from the local decomposition to the ghosts