add_particles: Resize & Redistribute (#181)

* add_particles: Resize & Redistribute

Move to ImpactX class and call Resize and Redistribute afterwards.
Otherwise, we would need to manually call the latter two
from Python, which is verbose and unnecessary.

* work-around for nvcc

examples/fodo/run_fodo.py

@@ -32,8 +32,7 @@
     muxpx = -0.846574929020762,
     muypy = 0.846574929020762,
     mutpt = 0.0)
-distribution.generate_add_particles(
-    impactX.particle_container(), qm_qeeV, charge_C, distr, npart)
+impactX.add_particles(qm_qeeV, charge_C, distr, npart)
 
 # init reference particle
 refPart = RefPart()

src/ImpactX.H

@@ -10,6 +10,7 @@
 #ifndef IMPACT_X_H
 #define IMPACT_X_H
 
+#include "particles/distribution/All.H"
 #include "particles/elements/All.H"
 #include "particles/ImpactXParticleContainer.H"
 
@@ -63,6 +64,25 @@ namespace impactx
          */
         void initLatticeElementsFromInputs ();
 
+        /** Generate and add n particles to the particle container
+         *
+         * Will also resize the geometry based on the updated particle
+         * distribution's extent and then redistribute particles in according
+         * AMReX grid boxes.
+         *
+         * @param qm charge/mass ratio (q_e/eV)
+         * @param bunch_charge bunch charge (C)
+         * @param distr distribution function to draw from (object)
+         * @param npart number of particles to draw
+         */
+        void
+        add_particles (
+            amrex::ParticleReal qm,
+            amrex::ParticleReal bunch_charge,
+            distribution::KnownDistributions distr,
+            int npart
+        );
+
         /** Run the main simulation loop for a number of steps
          */
         void evolve ();

src/initialization/InitDistribution.cpp

@@ -7,15 +7,9 @@
  * Authors: Axel Huebl, Chad Mitchell, Ji Qiang
  * License: BSD-3-Clause-LBNL
  */
-#include "InitDistribution.H"
 #include "ImpactX.H"
 #include "particles/ImpactXParticleContainer.H"
-#include "particles/distribution/Waterbag.H"
-#include "particles/distribution/Kurth6D.H"
-#include "particles/distribution/Gaussian.H"
-#include "particles/distribution/KVdist.H"
-#include "particles/distribution/Kurth4D.H"
-#include "particles/distribution/Semigaussian.H"
+#include "particles/distribution/All.H"
 
 #include <AMReX.H>
 #include <AMReX_BLProfiler.H>
@@ -28,6 +22,61 @@
 
 namespace impactx
 {
+    void
+    ImpactX::add_particles (
+        amrex::ParticleReal qm,
+        amrex::ParticleReal bunch_charge,
+        distribution::KnownDistributions distr,
+        int npart
+    )
+    {
+        BL_PROFILE("ImpactX::add_particles");
+
+        amrex::Vector<amrex::ParticleReal> x, y, t;
+        amrex::Vector<amrex::ParticleReal> px, py, pt;
+        amrex::ParticleReal ix, iy, it, ipx, ipy, ipt;
+        amrex::RandomEngine rng;
+
+        // Logic: We initialize 1/Nth of particles, independent of their
+        // position, per MPI rank. We then measure the distribution's spatial
+        // extent, create a grid, resize it to fit the beam, and then
+        // redistribute particles so that they reside on the correct MPI rank.
+        int myproc = amrex::ParallelDescriptor::MyProc();
+        int nprocs = amrex::ParallelDescriptor::NProcs();
+        int navg = npart / nprocs;
+        int nleft = npart - navg * nprocs;
+        int npart_this_proc = (myproc < nleft) ? navg+1 : navg;
+
+        std::visit([&](auto&& distribution){
+            x.reserve(npart_this_proc);
+            y.reserve(npart_this_proc);
+            t.reserve(npart_this_proc);
+            px.reserve(npart_this_proc);
+            py.reserve(npart_this_proc);
+            pt.reserve(npart_this_proc);
+
+            for (amrex::Long i = 0; i < npart_this_proc; ++i) {
+                distribution(ix, iy, it, ipx, ipy, ipt, rng);
+                x.push_back(ix);
+                y.push_back(iy);
+                t.push_back(it);
+                px.push_back(ipx);
+                py.push_back(ipy);
+                pt.push_back(ipt);
+            }
+        }, distr);
+
+        int const lev = 0;
+        m_particle_container->AddNParticles(lev, x, y, t, px, py, pt,
+                                            qm, bunch_charge);
+
+        // Resize the mesh to fit the spatial extent of the beam and then
+        // redistribute particles, so they reside on the MPI rank that is
+        // responsible for the respective spatial particle position.
+        this->ResizeMesh();
+        m_particle_container->Redistribute();
+    }
+
     void ImpactX::initBeamDistributionFromInputs ()
     {
         BL_PROFILE("ImpactX::initBeamDistributionFromInputs");
@@ -78,12 +127,12 @@ namespace impactx
           pp_dist.query("muypy", muypy);
           pp_dist.query("mutpt", mutpt);
 
-          impactx::distribution::Waterbag waterbag(
+          distribution::KnownDistributions waterbag(distribution::Waterbag(
               sigx, sigy, sigt,
               sigpx, sigpy, sigpt,
-              muxpx, muypy, mutpt);
+              muxpx, muypy, mutpt));
 
-          generate_add_particles(*m_particle_container, qm, bunch_charge, waterbag, npart);
+          add_particles(qm, bunch_charge, waterbag, npart);
 
         } else if (distribution_type == "kurth6d") {
           amrex::ParticleReal sigx,sigy,sigt,sigpx,sigpy,sigpt;
@@ -98,12 +147,12 @@ namespace impactx
           pp_dist.query("muypy", muypy);
           pp_dist.query("mutpt", mutpt);
 
-          impactx::distribution::Kurth6D kurth6D(
+          distribution::KnownDistributions kurth6D(distribution::Kurth6D(
             sigx, sigy, sigt,
             sigpx, sigpy, sigpt,
-            muxpx, muypy, mutpt);
+            muxpx, muypy, mutpt));
 
-          generate_add_particles(*m_particle_container, qm, bunch_charge, kurth6D, npart);
+          add_particles(qm, bunch_charge, kurth6D, npart);
 
         } else if (distribution_type == "gaussian") {
           amrex::ParticleReal sigx,sigy,sigt,sigpx,sigpy,sigpt;
@@ -118,12 +167,12 @@ namespace impactx
           pp_dist.query("muypy", muypy);
           pp_dist.query("mutpt", mutpt);
 
-          impactx::distribution::Gaussian gaussian(
+          distribution::KnownDistributions gaussian(distribution::Gaussian(
             sigx, sigy, sigt,
             sigpx, sigpy, sigpt,
-            muxpx, muypy, mutpt);
+            muxpx, muypy, mutpt));
 
-          generate_add_particles(*m_particle_container, qm, bunch_charge, gaussian, npart);
+          add_particles(qm, bunch_charge, gaussian, npart);
 
         } else if (distribution_type == "kvdist") {
           amrex::ParticleReal sigx,sigy,sigt,sigpx,sigpy,sigpt;
@@ -138,12 +187,12 @@ namespace impactx
           pp_dist.query("muypy", muypy);
           pp_dist.query("mutpt", mutpt);
 
-          impactx::distribution::KVdist kvDist(
+          distribution::KnownDistributions kvDist(distribution::KVdist(
             sigx, sigy, sigt,
             sigpx, sigpy, sigpt,
-            muxpx, muypy, mutpt);
+            muxpx, muypy, mutpt));
 
-          generate_add_particles(*m_particle_container, qm, bunch_charge, kvDist, npart);
+          add_particles(qm, bunch_charge, kvDist, npart);
 
         } else if (distribution_type == "kurth4d") {
           amrex::ParticleReal sigx,sigy,sigt,sigpx,sigpy,sigpt;
@@ -158,12 +207,12 @@ namespace impactx
           pp_dist.query("muypy", muypy);
           pp_dist.query("mutpt", mutpt);
 
-          impactx::distribution::Kurth4D kurth4D(
+          distribution::KnownDistributions kurth4D(distribution::Kurth4D(
             sigx, sigy, sigt,
             sigpx, sigpy, sigpt,
-            muxpx, muypy, mutpt);
+            muxpx, muypy, mutpt));
 
-          generate_add_particles(*m_particle_container, qm, bunch_charge, kurth4D, npart);
+          add_particles(qm, bunch_charge, kurth4D, npart);
         } else if (distribution_type == "semigaussian") {
           amrex::ParticleReal sigx,sigy,sigt,sigpx,sigpy,sigpt;
           amrex::ParticleReal muxpx = 0.0, muypy = 0.0, mutpt = 0.0;
@@ -177,22 +226,16 @@ namespace impactx
           pp_dist.query("muypy", muypy);
           pp_dist.query("mutpt", mutpt);
 
-          impactx::distribution::Semigaussian semigaussian(
+          distribution::KnownDistributions semigaussian(distribution::Semigaussian(
             sigx, sigy, sigt,
             sigpx, sigpy, sigpt,
-            muxpx, muypy, mutpt);
+            muxpx, muypy, mutpt));
 
-          generate_add_particles(*m_particle_container, qm, bunch_charge, semigaussian, npart);
+          add_particles(qm, bunch_charge, semigaussian, npart);
         } else {
             amrex::Abort("Unknown distribution: " + distribution_type);
         }
 
-        // Resize the mesh to fit the spatial extent of the beam and then
-        // redistribute particles, so they reside on the MPI rank that is
-        // responsible for the respective spatial particle position.
-        this->ResizeMesh();
-        m_particle_container->Redistribute();
-
         // reference particle
         amrex::ParticleReal massE;  // MeV
         if (particle_type == "electron") {

src/python/ImpactX.cpp

@@ -6,7 +6,6 @@
 #include "pyImpactX.H"
 
 #include <ImpactX.H>
-#include <initialization/InitDistribution.H>
 #include <particles/distribution/Gaussian.H>
 #include <particles/distribution/Kurth4D.H>
 #include <particles/distribution/Kurth6D.H>
@@ -37,7 +36,8 @@ void init_ImpactX(py::module& m)
     >(m, "MultiFabPerLevel");
     */
 
-    py::class_<ImpactX>(m, "ImpactX")
+    py::class_<ImpactX> impactx(m, "ImpactX");
+    impactx
         .def(py::init<>())
 
         .def("load_inputs_file",
@@ -83,6 +83,10 @@ void init_ImpactX(py::module& m)
         .def("init_grids", &ImpactX::initGrids)
         .def("init_beam_distribution_from_inputs", &ImpactX::initBeamDistributionFromInputs)
         .def("init_lattice_elements_from_inputs", &ImpactX::initLatticeElementsFromInputs)
+        .def("add_particles", &ImpactX::add_particles,
+             py::arg("qm"), py::arg("bunch_charge"),
+             py::arg("distr"), py::arg("npart")
+        )
         .def("evolve", &ImpactX::evolve)
         .def("particle_container",
              [](ImpactX & ix) -> ImpactXParticleContainer & {

src/python/distribution.cpp

@@ -5,7 +5,6 @@
  */
 #include "pyImpactX.H"
 
-#include <initialization/InitDistribution.H>
 #include <particles/distribution/All.H>
 
 #include <AMReX.H>
@@ -16,28 +15,6 @@
 namespace py = pybind11;
 using namespace impactx;
 
-namespace
-{
-    /** Register impactx::generate_add_particles for each distribution
-     *
-     * @tparam I counter through all known distributions in impactx::distribution::KnownDistributions
-     * @param md the module to register the function in
-     */
-    template< std::size_t I = 0 >
-    void register_generate_add_particles(py::module& md)
-    {
-        using V = impactx::distribution::KnownDistributions;
-        using T = std::variant_alternative_t<I, V>;
-
-        md.def("generate_add_particles", &generate_add_particles<T>,
-            py::arg("pc"), py::arg("qm"), py::arg("bunch_charge"),
-            py::arg("distr"), py::arg("npart")
-        );
-
-        if constexpr (I < std::variant_size_v<V> - 1)
-            register_generate_add_particles<I + 1>(md);
-    }
-}
 
 void init_distribution(py::module& m)
 {
@@ -111,6 +88,4 @@ void init_distribution(py::module& m)
              py::arg("sigmaPx"), py::arg("sigmaPy"), py::arg("sigmaPt"),
              py::arg("muxpx")=0.0, py::arg("muypy")=0.0, py::arg("mutpt")=0.0
         );
-
-    register_generate_add_particles(md);
 }

tests/python/test_impactx.py

@@ -45,8 +45,7 @@ def test_impactx_nofile():
         muxpx = -0.846574929020762,
         muypy = 0.846574929020762,
         mutpt = 0.0)
-    distribution.generate_add_particles(
-        impactX.particle_container(), qm_qeeV, charge_C, distr, npart)
+    impactX.add_particles(qm_qeeV, charge_C, distr, npart)
 
     # init reference particle
     refPart = RefPart()