Python: Particle Init

src/initialization/InitDistribution.H

@@ -0,0 +1,79 @@
+/* Copyright 2022 The Regents of the University of California, through Lawrence
+ *           Berkeley National Laboratory (subject to receipt of any required
+ *           approvals from the U.S. Dept. of Energy). All rights reserved.
+ *
+ * This file is part of ImpactX.
+ *
+ * Authors: Axel Huebl
+ * License: BSD-3-Clause-LBNL
+ */
+#include "particles/ImpactXParticleContainer.H"
+
+#include <AMReX.H>
+#include <AMReX_REAL.H>
+#include <AMReX_ParmParse.H>
+#include <AMReX_Print.H>
+#include <AMReX_Random.H>
+#include <AMReX_Vector.H>
+
+#include <string>
+
+
+namespace impactx
+{
+    /** Generate and add n particles to a particle container
+     *
+     * @tparam T_Distribution distribution function to draw from (type)
+     * @param pc a particle container to add particles to
+     * @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
+     */
+    template <typename T_Distribution>
+    void
+    generate_add_particles (
+        impactx::ImpactXParticleContainer & pc,
+        amrex::ParticleReal qm,
+        amrex::ParticleReal bunch_charge,
+        T_Distribution & distr,
+        int npart
+    )
+    {
+        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;
+
+        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) {
+            distr(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);
+        }
+
+        int const lev = 0;
+        pc.AddNParticles(lev, x, y, t, px, py, pt,
+                         qm, bunch_charge);
+    }
+} // namespace impactx

src/initialization/InitDistribution.cpp

@@ -7,6 +7,7 @@
  * 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"
@@ -25,65 +26,6 @@
 #include <string>
 
 
-namespace
-{
-    /** Generate and add n particles to a particle container
-     *
-     * @tparam T_Distribution distribution function to draw from (type)
-     * @param pc a particle container to add particles to
-     * @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
-     */
-    template <typename T_Distribution>
-    void
-    generate_add_particles (
-        impactx::ImpactXParticleContainer & pc,
-        amrex::ParticleReal qm,
-        amrex::ParticleReal bunch_charge,
-        T_Distribution & distr,
-        int npart
-    )
-    {
-        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;
-
-        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) {
-            distr(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);
-        }
-
-        int const lev = 0;
-        pc.AddNParticles(lev, x, y, t, px, py, pt,
-                         qm, bunch_charge);
-    }
-}
-
 namespace impactx
 {
     void ImpactX::initBeamDistributionFromInputs ()
@@ -267,6 +209,7 @@ namespace impactx
         refPart.z = 0.0;
         refPart.px = 0.0;
         refPart.py = 0.0;
+        // make the next two lines a helper function?
         refPart.pt = -energy/massE - 1.0_prt;
         refPart.pz = sqrt(pow(refPart.pt,2) - 1.0_prt);
         m_particle_container->SetRefParticle(refPart);

src/python/CMakeLists.txt

@@ -4,6 +4,9 @@
 #
 target_sources(pyImpactX
   PRIVATE
+    distribution.cpp
     elements.cpp
     ImpactX.cpp
+    ImpactXParticleContainer.cpp
+    ReferenceParticle.cpp
 )

src/python/ImpactX.cpp

@@ -6,6 +6,14 @@
 #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>
+#include <particles/distribution/KVdist.H>
+#include <particles/distribution/Semigaussian.H>
+#include <particles/distribution/Waterbag.H>
+
 #include <AMReX.H>
 #include <AMReX_ParmParse.H>
 
@@ -76,8 +84,12 @@ void init_ImpactX(py::module& m)
         .def("init_beam_distribution_from_inputs", &ImpactX::initBeamDistributionFromInputs)
         .def("init_lattice_elements_from_inputs", &ImpactX::initLatticeElementsFromInputs)
         .def("evolve", &ImpactX::evolve)
-
-        //.def_property("particle_container", &ImpactX::m_particle_container)
+        .def("particle_container",
+             [](ImpactX & ix) -> ImpactXParticleContainer & {
+                return *ix.m_particle_container;
+             },
+             py::return_value_policy::reference_internal
+        )
         //.def_readwrite("rho", &ImpactX::m_rho)
         .def_readwrite("lattice", &ImpactX::m_lattice)
         //.def_readwrite("lattice", &ImpactX::m_lattice_test)

src/python/ImpactXParticleContainer.cpp

@@ -0,0 +1,33 @@
+/* Copyright 2021-2022 The ImpactX Community
+ *
+ * Authors: Axel Huebl
+ * License: BSD-3-Clause-LBNL
+ */
+#include "pyImpactX.H"
+
+#include <particles/ImpactXParticleContainer.H>
+#include <AMReX.H>
+#include <AMReX_ParticleContainer.H>
+
+namespace py = pybind11;
+using namespace impactx;
+
+
+void init_impactxparticlecontainer(py::module& m)
+{
+    py::class_<
+        ImpactXParticleContainer,
+        amrex::ParticleContainer<0, 0, RealSoA::nattribs, IntSoA::nattribs>
+    >(m, "ImpactXParticleContainer")
+        //.def(py::init<>())
+        .def("add_n_particles", &ImpactXParticleContainer::AddNParticles)
+        .def("get_ref_particle",
+            py::overload_cast<>(&ImpactXParticleContainer::GetRefParticle),
+            py::return_value_policy::reference_internal
+        )
+        .def("set_ref_particle", &ImpactXParticleContainer::SetRefParticle)
+        .def("set_particle_shape",
+            py::overload_cast<int const>(&ImpactXParticleContainer::SetParticleShape)
+        )
+    ;
+}

src/python/ReferenceParticle.cpp

@@ -0,0 +1,28 @@
+/* Copyright 2021-2022 The ImpactX Community
+ *
+ * Authors: Axel Huebl
+ * License: BSD-3-Clause-LBNL
+ */
+#include "pyImpactX.H"
+
+#include <particles/ImpactXParticleContainer.H>
+#include <AMReX.H>
+
+namespace py = pybind11;
+using namespace impactx;
+
+
+void init_refparticle(py::module& m)
+{
+    py::class_<RefPart>(m, "RefPart")
+        .def(py::init<>())
+        .def_readwrite("x", &RefPart::x, "horizontal position x, in meters")
+        .def_readwrite("y", &RefPart::y, "vertical position y, in meters")
+        .def_readwrite("z", &RefPart::z, "longitudinal position y, in meters")
+        .def_readwrite("t", &RefPart::t, "clock time * c in meters")
+        .def_readwrite("px", &RefPart::px, "momentum in x, normalized to proper velocity")
+        .def_readwrite("py", &RefPart::py, "momentum in y, normalized to proper velocity")
+        .def_readwrite("pz", &RefPart::pz, "momentum in z, normalized to proper velocity")
+        .def_readwrite("pt", &RefPart::pt, "energy deviation, normalized by rest energy")
+    ;
+}

src/python/distribution.cpp

@@ -0,0 +1,113 @@
+/* Copyright 2021-2022 The ImpactX Community
+ *
+ * Authors: Axel Huebl
+ * License: BSD-3-Clause-LBNL
+ */
+#include "pyImpactX.H"
+
+#include <initialization/InitDistribution.H>
+#include <particles/distribution/All.H>
+
+#include <AMReX.H>
+#include <AMReX_REAL.H>
+
+#include <variant>
+
+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>);
+
+        if constexpr (I < std::variant_size_v<V> - 1)
+            register_generate_add_particles<I + 1>(md);
+    }
+}
+
+void init_distribution(py::module& m)
+{
+    py::module_ md = m.def_submodule(
+        "distribution",
+        "Particle beam distributions in ImpactX"
+    );
+
+    py::class_<distribution::Gaussian>(md, "Gaussian")
+        .def(py::init<
+                amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const,
+                amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const,
+                amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const
+             >(),
+             py::arg("sigmaX"), py::arg("sigmaY"), py::arg("sigmaT"),
+             py::arg("sigmaPx"), py::arg("sigmaPy"), py::arg("sigmaPt"),
+             py::arg("muxpx"), py::arg("muypy"), py::arg("mutpt")
+        );
+
+    py::class_<distribution::Kurth4D>(md, "Kurth4D")
+        .def(py::init<
+                 amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const,
+                 amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const,
+                 amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const
+         >(),
+             py::arg("sigmaX"), py::arg("sigmaY"), py::arg("sigmaT"),
+             py::arg("sigmaPx"), py::arg("sigmaPy"), py::arg("sigmaPt"),
+             py::arg("muxpx"), py::arg("muypy"), py::arg("mutpt")
+        );
+
+    py::class_<distribution::Kurth6D>(md, "Kurth6D")
+        .def(py::init<
+                 amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const,
+                 amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const,
+                 amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const
+             >(),
+             py::arg("sigmaX"), py::arg("sigmaY"), py::arg("sigmaT"),
+             py::arg("sigmaPx"), py::arg("sigmaPy"), py::arg("sigmaPt"),
+             py::arg("muxpx"), py::arg("muypy"), py::arg("mutpt")
+        );
+
+    py::class_<distribution::KVdist>(md, "KVdist")
+        .def(py::init<
+                 amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const,
+                 amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const,
+                 amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const
+             >(),
+             py::arg("sigmaX"), py::arg("sigmaY"), py::arg("sigmaT"),
+             py::arg("sigmaPx"), py::arg("sigmaPy"), py::arg("sigmaPt"),
+             py::arg("muxpx"), py::arg("muypy"), py::arg("mutpt")
+        );
+
+    py::class_<distribution::Semigaussian>(md, "Semigaussian")
+        .def(py::init<
+                 amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const,
+                 amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const,
+                 amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const
+             >(),
+             py::arg("sigmaX"), py::arg("sigmaY"), py::arg("sigmaT"),
+             py::arg("sigmaPx"), py::arg("sigmaPy"), py::arg("sigmaPt"),
+             py::arg("muxpx"), py::arg("muypy"), py::arg("mutpt")
+        );
+
+    py::class_<distribution::Waterbag>(md, "Waterbag")
+        .def(py::init<
+                 amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const,
+                 amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const,
+                 amrex::ParticleReal const, amrex::ParticleReal const, amrex::ParticleReal const
+             >(),
+             py::arg("sigmaX"), py::arg("sigmaY"), py::arg("sigmaT"),
+             py::arg("sigmaPx"), py::arg("sigmaPy"), py::arg("sigmaPt"),
+             py::arg("muxpx"), py::arg("muypy"), py::arg("mutpt")
+        );
+
+    register_generate_add_particles(md);
+}