Python Support, Cleaning

ECP-WarpX · Jan 4, 2024 · 286459d · 286459d
1 parent 55e4c43
commit 286459d
Show file tree

Hide file tree

Showing 5 changed files with 41 additions and 20 deletions.
diff --git a/src/initialization/InitDistribution.H b/src/initialization/InitDistribution.H
@@ -12,10 +12,11 @@
 
 #include "particles/ImpactXParticleContainer.H"
 
-#include <AMReX_Extension.H> // for AMREX_RESTRICT
+#include <AMReX_Extension.H>  // for AMREX_RESTRICT
 #include <AMReX_REAL.H>
 
-#include <type_traits>
+#include <utility>  // for std::move
+
 
 namespace impactx::initialization
 {
@@ -77,9 +78,15 @@ namespace impactx::initialization
             amrex::RandomEngine const & engine
         ) const
         {
-            m_distribution(m_part_x[i], m_part_y[i], m_part_t[i],
-                           m_part_px[i], m_part_py[i], m_part_pt[i],
-                           engine);
+            m_distribution(
+                m_part_x[i],
+                m_part_y[i],
+                m_part_t[i],
+                m_part_px[i],
+                m_part_py[i],
+                m_part_pt[i],
+                engine
+            );
         }
 
     private:

diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp
@@ -7,8 +7,9 @@
  * Authors: Axel Huebl, Chad Mitchell, Ji Qiang
  * License: BSD-3-Clause-LBNL
  */
-#include "ImpactX.H"
 #include "initialization/InitDistribution.H"
+
+#include "ImpactX.H"
 #include "particles/ImpactXParticleContainer.H"
 #include "particles/distribution/All.H"
 

diff --git a/src/particles/ImpactXParticleContainer.H b/src/particles/ImpactXParticleContainer.H
@@ -151,9 +151,9 @@ namespace impactx
 
         /** Add new particles to the container for fixed s.
          *
-         *  In this version of AddNParticles, when running on GPUs, the input
-         *  vectors of particle data live on the host, and will be copied to
-         *  the device as part of the method call.
+         * In this version of AddNParticles, when running on GPUs, the input
+         * vectors of particle data live on the host, and will be copied to
+         * the device as part of the method call.
          *
          * Note: This can only be used *after* the initialization (grids) have
          *       been created, meaning after the call to AmrCore::InitFromScratch
@@ -178,13 +178,13 @@ namespace impactx
                        amrex::Vector<amrex::ParticleReal> const & px,
                        amrex::Vector<amrex::ParticleReal> const & py,
                        amrex::Vector<amrex::ParticleReal> const & pt,
-                       amrex::ParticleReal const & qm,
-                       amrex::ParticleReal const & bchchg);
+                       amrex::ParticleReal qm,
+                       amrex::ParticleReal bchchg);
 
         /** Add new particles to the container for fixed s.
          *
-         *  In this version of AddNParticles, when running on GPUs, the input
-         *  vectors of particle data have already been copied to the device.
+         * In this version of AddNParticles, when running on GPUs, the input
+         * vectors of particle data have already been copied to the device.
          *
          * Note: This can only be used *after* the initialization (grids) have
          *       been created, meaning after the call to AmrCore::InitFromScratch
@@ -209,8 +209,8 @@ namespace impactx
                        amrex::Gpu::DeviceVector<amrex::ParticleReal> const & px,
                        amrex::Gpu::DeviceVector<amrex::ParticleReal> const & py,
                        amrex::Gpu::DeviceVector<amrex::ParticleReal> const & pt,
-                       amrex::ParticleReal const & qm,
-                       amrex::ParticleReal const & bchchg);
+                       amrex::ParticleReal qm,
+                       amrex::ParticleReal bchchg);
 
         /** Register storage for lost particles
          *

diff --git a/src/particles/ImpactXParticleContainer.cpp b/src/particles/ImpactXParticleContainer.cpp
@@ -107,8 +107,8 @@ namespace impactx
                                              amrex::Vector<amrex::ParticleReal> const & px,
                                              amrex::Vector<amrex::ParticleReal> const & py,
                                              amrex::Vector<amrex::ParticleReal> const & pt,
-                                             amrex::ParticleReal const & qm,
-                                             amrex::ParticleReal const & bchchg)
+                                             amrex::ParticleReal qm,
+                                             amrex::ParticleReal bchchg)
     {
         BL_PROFILE("ImpactX::AddNParticles");
 
@@ -175,8 +175,8 @@ namespace impactx
                                              amrex::Gpu::DeviceVector<amrex::ParticleReal> const & px,
                                              amrex::Gpu::DeviceVector<amrex::ParticleReal> const & py,
                                              amrex::Gpu::DeviceVector<amrex::ParticleReal> const & pt,
-                                             amrex::ParticleReal const & qm,
-                                             amrex::ParticleReal const & bchchg)
+                                             amrex::ParticleReal qm,
+                                             amrex::ParticleReal bchchg)
     {
         BL_PROFILE("ImpactX::AddNParticles");
 

diff --git a/src/python/ImpactXParticleContainer.cpp b/src/python/ImpactXParticleContainer.cpp
@@ -58,12 +58,25 @@ void init_impactxparticlecontainer(py::module& m)
         )
 
         .def("add_n_particles",
-             &ImpactXParticleContainer::AddNParticles,
+             py::overload_cast<
+                 int,
+                 amrex::Vector<amrex::ParticleReal> const &,
+                 amrex::Vector<amrex::ParticleReal> const &,
+                 amrex::Vector<amrex::ParticleReal> const &,
+                 amrex::Vector<amrex::ParticleReal> const &,
+                 amrex::Vector<amrex::ParticleReal> const &,
+                 amrex::Vector<amrex::ParticleReal> const &,
+                 amrex::ParticleReal,
+                 amrex::ParticleReal
+             >(&ImpactXParticleContainer::AddNParticles),
              py::arg("lev"),
              py::arg("x"), py::arg("y"), py::arg("t"),
              py::arg("px"), py::arg("py"), py::arg("pt"),
              py::arg("qm"), py::arg("bchchg"),
              "Add new particles to the container for fixed s.\n\n"
+             "In this version of AddNParticles, when running on GPUs, the input\n"
+             "vectors of particle data live on the host, and will be copied to\n"
+             "the device as part of the method call.\n\n"
              "Note: This can only be used *after* the initialization (grids) have\n"
              "      been created, meaning after the call to ImpactX.init_grids\n"
              "      has been made in the ImpactX class.\n\n"