From b1441e97e6d1ece6c47fde5eb3b9654dc2414a02 Mon Sep 17 00:00:00 2001
From: Sam Reeve <6740307+streeve@users.noreply.github.com>
Date: Mon, 9 Dec 2024 18:01:26 -0500
Subject: [PATCH] Add custom particle position/volume creation

---
 src/CabanaPD_Particles.hpp | 113 ++++++++++++++++++++++++++++++++++---
 1 file changed, 106 insertions(+), 7 deletions(-)

diff --git a/src/CabanaPD_Particles.hpp b/src/CabanaPD_Particles.hpp
index 84b87323..494facd7 100644
--- a/src/CabanaPD_Particles.hpp
+++ b/src/CabanaPD_Particles.hpp
@@ -160,6 +160,23 @@ class Particles<MemorySpace, PMB, TemperatureIndependent, BaseOutput, Dimension>
         resize( 0, 0 );
     }
 
+    // Constructor with existing particle data.
+    template <class ExecSpace, class PositionType, class VolumeType>
+    Particles( const ExecSpace& exec_space, const PositionType& x,
+               const VolumeType& vol, std::array<double, dim> low_corner,
+               std::array<double, dim> high_corner,
+               const std::array<int, dim> num_cells, const int max_halo_width )
+        : halo_width( max_halo_width )
+        , _plist_x( "positions" )
+        , _plist_f( "forces" )
+    {
+        createDomain( low_corner, high_corner, num_cells );
+
+        _init_timer.start();
+        implInitCustomParticles( exec_space, x, vol );
+        _init_timer.stop();
+    }
+
     // Constructor which initializes particles on regular grid.
     template <class ExecSpace>
     Particles( const ExecSpace& exec_space, std::array<double, dim> low_corner,
@@ -468,6 +485,44 @@ class Particles<MemorySpace, PMB, TemperatureIndependent, BaseOutput, Dimension>
     friend class Comm<self_type, PMB, TemperatureIndependent>;
     friend class Comm<self_type, PMB, TemperatureDependent>;
 
+    // This is necessary to avoid Cuda lambda capture issues in the constructor.
+    template <class ExecSpace, class PositionType, class VolumeType>
+    void implInitCustomParticles( const ExecSpace, const PositionType& x,
+                                  const VolumeType& vol )
+    {
+        resize( vol.size(), 0 );
+        auto p_x = sliceReferencePosition();
+        auto p_vol = sliceVolume();
+        auto v = sliceVelocity();
+        auto f = sliceForce();
+        auto type = sliceType();
+        auto rho = sliceDensity();
+        auto u = sliceDisplacement();
+        auto nofail = sliceNoFail();
+
+        static_assert(
+            Cabana::is_accessible_from<
+                memory_space, typename PositionType::execution_space>::value );
+
+        Kokkos::parallel_for(
+            "copy_to_particles", Kokkos::RangePolicy<ExecSpace>( 0, n_local ),
+            KOKKOS_LAMBDA( const int pid ) {
+                // Set the particle position and volume.
+                // Set everything else to zero.
+                p_vol( pid ) = vol( pid );
+                for ( int d = 0; d < 3; d++ )
+                {
+                    p_x( pid, d ) = x( pid, d );
+                    u( pid, d ) = 0.0;
+                    v( pid, d ) = 0.0;
+                    f( pid, d ) = 0.0;
+                }
+                type( pid ) = 0;
+                nofail( pid ) = 0;
+                rho( pid ) = 1.0;
+            } );
+    }
+
   protected:
     aosoa_u_type _aosoa_u;
     aosoa_y_type _aosoa_y;
@@ -798,13 +853,10 @@ class Particles<MemorySpace, ModelType, ThermalType, EnergyOutput, Dimension>
         _aosoa_output = aosoa_output_type( "Particle Output Fields", 0 );
     }
 
-    // Constructor which initializes particles on regular grid.
-    template <class ExecSpace>
-    Particles( const ExecSpace& exec_space, std::array<double, dim> low_corner,
-               std::array<double, dim> high_corner,
-               const std::array<int, dim> num_cells, const int max_halo_width )
-        : base_type( exec_space, low_corner, high_corner, num_cells,
-                     max_halo_width )
+    // Base constructor.
+    template <typename... Args>
+    Particles( Args&&... args )
+        : base_type( std::forward<Args>( args )... )
     {
         _aosoa_output = aosoa_output_type( "Particle Output Fields", n_local );
         init_output();
@@ -945,6 +997,53 @@ auto createParticles(
         exec_space, low_corner, high_corner, num_cells, max_halo_width,
         EnergyOutput{} );
 }
+
+template <typename MemorySpace, typename ModelType, typename ExecSpace,
+          typename PositionType, typename VolumeType, std::size_t Dim,
+          typename OutputType>
+auto createParticles( const ExecSpace& exec_space, const PositionType& x,
+                      const VolumeType& vol, std::array<double, Dim> low_corner,
+                      std::array<double, Dim> high_corner,
+                      const std::array<int, Dim> num_cells,
+                      const int max_halo_width, OutputType )
+{
+    return std::make_shared<
+        CabanaPD::Particles<MemorySpace, typename ModelType::base_model,
+                            typename ModelType::thermal_type, OutputType>>(
+        exec_space, x, vol, low_corner, high_corner, num_cells,
+        max_halo_width );
+}
+
+template <typename MemorySpace, typename ModelType, typename ThermalType,
+          typename ExecSpace, typename PositionType, typename VolumeType,
+          std::size_t Dim, typename OutputType>
+auto createParticles(
+    const ExecSpace& exec_space, const PositionType& x, const VolumeType& vol,
+    std::array<double, Dim> low_corner, std::array<double, Dim> high_corner,
+    const std::array<int, Dim> num_cells, const int max_halo_width, OutputType,
+    typename std::enable_if<( is_temperature_dependent<ThermalType>::value ),
+                            int>::type* = 0 )
+{
+    return std::make_shared<CabanaPD::Particles<
+        MemorySpace, ModelType, typename ThermalType::base_type, OutputType>>(
+        exec_space, x, vol, low_corner, high_corner, num_cells,
+        max_halo_width );
+}
+
+template <typename MemorySpace, typename ModelType, typename ExecSpace,
+          typename PositionType, typename VolumeType, std::size_t Dim>
+auto createParticles( const ExecSpace& exec_space, const PositionType& x,
+                      const VolumeType& vol, std::array<double, Dim> low_corner,
+                      std::array<double, Dim> high_corner,
+                      const std::array<int, Dim> num_cells,
+                      const int max_halo_width )
+{
+    return createParticles<MemorySpace, ModelType, ExecSpace, PositionType,
+                           VolumeType, Dim>( exec_space, x, vol, low_corner,
+                                             high_corner, num_cells,
+                                             max_halo_width, EnergyOutput{} );
+}
+
 } // namespace CabanaPD
 
 #endif