Python: Expose Manual Push (#468)

* Python: Expose Manual Push Enable to push a particle container through an element class without being in the simulation's evolve loop. * Simplify Overload Address MSVC link issue. Will keep option simple if operator is implemented as `const` method.
ECP-WarpX · Nov 15, 2023 · 5c0ddc8 · 5c0ddc8
1 parent c8a165a
commit 5c0ddc8
Show file tree

Hide file tree

Showing 5 changed files with 167 additions and 29 deletions.
diff --git a/src/particles/Push.cpp b/src/particles/Push.cpp
@@ -11,7 +11,6 @@
 
 #include <AMReX_BLProfiler.H>
 
-#include <string>
 #include <variant>
 
 
@@ -24,12 +23,7 @@ namespace impactx
         // here we just access the element by its respective type
         std::visit([&pc, step](auto&& element)
         {
-            // performance profiling per element
-            std::string element_name;
-            element_name = element.name;
-            std::string const profile_name = "impactx::Push::" + element_name;
             BL_PROFILE("impactx::Push");
-            BL_PROFILE(profile_name);
 
             // push reference particle & all particles
             element(pc, step);

diff --git a/src/particles/PushAll.H b/src/particles/PushAll.H
@@ -36,6 +36,10 @@ namespace impactx
             [[maybe_unused]] bool omp_parallel = true
     )
     {
+        // performance profiling per element
+        std::string const profile_name = "impactx::Push::" + std::string(T_Element::name);
+        BL_PROFILE(profile_name);
+
         // preparing to access reference particle data: RefPart
         RefPart & ref_part = pc.GetRefParticle();
 

diff --git a/src/particles/elements/diagnostics/openPMD.cpp b/src/particles/elements/diagnostics/openPMD.cpp
@@ -15,6 +15,7 @@
 #include <ablastr/particles/IndexHandling.H>
 
 #include <AMReX.H>
+#include <AMReX_BLProfiler.H>
 #include <AMReX_REAL.H>
 #include <AMReX_ParmParse.H>
 
@@ -23,6 +24,7 @@
 namespace io = openPMD;
 #endif
 
+#include <string>
 #include <utility>
 
 
@@ -310,6 +312,9 @@ namespace detail
         int step
     )
     {
+        std::string profile_name = "impactx::Push::" + std::string(BeamMonitor::name);
+        BL_PROFILE(profile_name);
+
         // preparing to access reference particle data: RefPart
         RefPart & ref_part = pc.GetRefParticle();
 

diff --git a/src/python/elements.cpp b/src/python/elements.cpp
@@ -5,16 +5,37 @@
  */
 #include "pyImpactX.H"
 
+#include <particles/Push.H>
 #include <particles/elements/All.H>
 #include <AMReX.H>
 
+#include <type_traits>
 #include <utility>
 #include <vector>
 
 namespace py = pybind11;
 using namespace impactx;
 
 
+namespace
+{
+    /** Registers the mixin BeamOptics::operator methods
+     */
+    template<typename T_PyClass>
+    void register_beamoptics_push(T_PyClass & cl)
+    {
+        using Element = typename T_PyClass::type;  // py::class<T, options...>
+
+        cl.def("push",
+            [](Element & el, ImpactXParticleContainer & pc, int step) {
+                el(pc, step);
+            },
+            py::arg("pc"), py::arg("step")=0,
+            "Push first the reference particle, then all other particles."
+        );
+    }
+}
+
 void init_elements(py::module& m)
 {
     py::module_ const me = m.def_submodule(
@@ -91,25 +112,30 @@ void init_elements(py::module& m)
 
     // diagnostics
 
-    py::class_<diagnostics::BeamMonitor, elements::Thin>(me, "BeamMonitor")
+    py::class_<diagnostics::BeamMonitor, elements::Thin> py_BeamMonitor(me, "BeamMonitor");
+    py_BeamMonitor
         .def(py::init<std::string, std::string, std::string>(),
              py::arg("name"), py::arg("backend")="default", py::arg("encoding")="g",
              "This element writes the particle beam out to openPMD data."
         )
     ;
+    register_beamoptics_push(py_BeamMonitor);
 
     // beam optics
 
-    py::class_<ChrDrift, elements::Thick>(me, "ChrDrift")
+    py::class_<ChrDrift, elements::Thick> py_ChrDrift(me, "ChrDrift");
+    py_ChrDrift
         .def(py::init<
                 amrex::ParticleReal const,
                 int const >(),
              py::arg("ds"), py::arg("nslice") = 1,
              "A Drift with chromatic effects included."
         )
     ;
+    register_beamoptics_push(py_ChrDrift);
 
-    py::class_<ChrQuad, elements::Thick>(me, "ChrQuad")
+    py::class_<ChrQuad, elements::Thick> py_ChrQuad(me, "ChrQuad");
+    py_ChrQuad
         .def(py::init<
                 amrex::ParticleReal const,
                 amrex::ParticleReal const,
@@ -119,8 +145,10 @@ void init_elements(py::module& m)
              "A Quadrupole magnet with chromatic effects included."
         )
     ;
+    register_beamoptics_push(py_ChrQuad);
 
-    py::class_<ChrAcc, elements::Thick>(me, "ChrAcc")
+    py::class_<ChrAcc, elements::Thick> py_ChrAcc(me, "ChrAcc");
+    py_ChrAcc
         .def(py::init<
                 amrex::ParticleReal const,
                 amrex::ParticleReal const,
@@ -130,8 +158,10 @@ void init_elements(py::module& m)
              "A region of Uniform Acceleration, with chromatic effects included."
         )
     ;
+    register_beamoptics_push(py_ChrAcc);
 
-    py::class_<ConstF, elements::Thick>(me, "ConstF")
+    py::class_<ConstF, elements::Thick> py_ConstF(me, "ConstF");
+    py_ConstF
         .def(py::init<
                 amrex::ParticleReal const,
                 amrex::ParticleReal const,
@@ -142,8 +172,10 @@ void init_elements(py::module& m)
              "A linear Constant Focusing element."
         )
     ;
+    register_beamoptics_push(py_ConstF);
 
-    py::class_<DipEdge, elements::Thin>(me, "DipEdge")
+    py::class_<DipEdge, elements::Thin> py_DipEdge(me, "DipEdge");
+    py_DipEdge
         .def(py::init<
                 amrex::ParticleReal const,
                 amrex::ParticleReal const,
@@ -153,26 +185,32 @@ void init_elements(py::module& m)
              "Edge focusing associated with bend entry or exit."
         )
     ;
+    register_beamoptics_push(py_DipEdge);
 
-    py::class_<Drift, elements::Thick>(me, "Drift")
+    py::class_<Drift, elements::Thick> py_Drift(me, "Drift");
+    py_Drift
         .def(py::init<
                 amrex::ParticleReal const,
                 int const >(),
              py::arg("ds"), py::arg("nslice") = 1,
              "A drift."
         )
     ;
+    register_beamoptics_push(py_Drift);
 
-    py::class_<ExactDrift, elements::Thick>(me, "ExactDrift")
+    py::class_<ExactDrift, elements::Thick> py_ExactDrift(me, "ExactDrift");
+    py_ExactDrift
         .def(py::init<
                 amrex::ParticleReal const,
                 int const >(),
              py::arg("ds"), py::arg("nslice") = 1,
              "A Drift using the exact nonlinear map."
         )
     ;
+    register_beamoptics_push(py_ExactDrift);
 
-    py::class_<ExactSbend, elements::Thick>(me, "ExactSbend")
+    py::class_<ExactSbend, elements::Thick> py_ExactSbend(me, "ExactSbend");
+    py_ExactSbend
         .def(py::init<
                 amrex::ParticleReal const,
                 amrex::ParticleReal const,
@@ -182,8 +220,10 @@ void init_elements(py::module& m)
              "An ideal sector bend using the exact nonlinear map.  When B = 0, the reference bending radius is defined by r0 = length / (angle in rad), corresponding to a magnetic field of B = rigidity / r0; otherwise the reference bending radius is defined by r0 = rigidity / B."
         )
     ;
+    register_beamoptics_push(py_ExactSbend);
 
-    py::class_<Kicker, elements::Thin>(me, "Kicker")
+    py::class_<Kicker, elements::Thin> py_Kicker(me, "Kicker");
+    py_Kicker
         .def(py::init([](
                 amrex::ParticleReal xkick,
                 amrex::ParticleReal ykick,
@@ -201,8 +241,10 @@ void init_elements(py::module& m)
              R"(A thin transverse kicker element. Kicks are for units "dimensionless" or in "T-m".)"
         )
     ;
+    register_beamoptics_push(py_Kicker);
 
-    py::class_<Multipole, elements::Thin>(me, "Multipole")
+    py::class_<Multipole, elements::Thin> py_Multipole(me, "Multipole");
+    py_Multipole
         .def(py::init<
                 int const,
                 amrex::ParticleReal const,
@@ -211,21 +253,26 @@ void init_elements(py::module& m)
              "A general thin multipole element."
         )
     ;
+    register_beamoptics_push(py_Multipole);
 
-    py::class_<None, elements::Thin>(me, "None")
+    py::class_<None, elements::Thin> py_None(me, "None");
+    py_None
         .def(py::init<>(),
              "This element does nothing."
         )
     ;
+    register_beamoptics_push(py_None);
 
-    py::class_<NonlinearLens, elements::Thin>(me, "NonlinearLens")
+    py::class_<NonlinearLens, elements::Thin> py_NonlinearLens(me, "NonlinearLens");
+    py_NonlinearLens
         .def(py::init<
                 amrex::ParticleReal const,
                 amrex::ParticleReal const>(),
              py::arg("knll"), py::arg("cnll"),
              "Single short segment of the nonlinear magnetic insert element."
         )
     ;
+    register_beamoptics_push(py_NonlinearLens);
 
     py::class_<Programmable>(me, "Programmable", py::dynamic_attr())
         .def(py::init<
@@ -270,7 +317,8 @@ void init_elements(py::module& m)
         )
     ;
 
-    py::class_<Quad, elements::Thick>(me, "Quad")
+    py::class_<Quad, elements::Thick> py_Quad(me, "Quad");
+    py_Quad
         .def(py::init<
                 amrex::ParticleReal const,
                 amrex::ParticleReal const,
@@ -279,8 +327,10 @@ void init_elements(py::module& m)
              "A Quadrupole magnet."
         )
     ;
+    register_beamoptics_push(py_Quad);
 
-    py::class_<RFCavity, elements::Thick>(me, "RFCavity")
+    py::class_<RFCavity, elements::Thick> py_RFCavity(me, "RFCavity");
+    py_RFCavity
         .def(py::init<
                 amrex::ParticleReal,
                 amrex::ParticleReal,
@@ -297,8 +347,10 @@ void init_elements(py::module& m)
              "An RF cavity."
         )
     ;
+    register_beamoptics_push(py_RFCavity);
 
-    py::class_<Sbend, elements::Thick>(me, "Sbend")
+    py::class_<Sbend, elements::Thick> py_Sbend(me, "Sbend");
+    py_Sbend
         .def(py::init<
                 amrex::ParticleReal const,
                 amrex::ParticleReal const,
@@ -307,8 +359,10 @@ void init_elements(py::module& m)
              "An ideal sector bend."
         )
     ;
+    register_beamoptics_push(py_Sbend);
 
-    py::class_<CFbend, elements::Thick>(me, "CFbend")
+    py::class_<CFbend, elements::Thick> py_CFbend(me, "CFbend");
+    py_CFbend
         .def(py::init<
                 amrex::ParticleReal const,
                 amrex::ParticleReal const,
@@ -318,17 +372,21 @@ void init_elements(py::module& m)
              "An ideal combined function bend (sector bend with quadrupole component)."
         )
     ;
+    register_beamoptics_push(py_CFbend);
 
-    py::class_<Buncher, elements::Thin>(me, "Buncher")
+    py::class_<Buncher, elements::Thin> py_Buncher(me, "Buncher");
+    py_Buncher
         .def(py::init<
                 amrex::ParticleReal const,
                 amrex::ParticleReal const>(),
              py::arg("V"), py::arg("k"),
              "A short linear RF cavity element at zero-crossing for bunching."
         )
     ;
+    register_beamoptics_push(py_Buncher);
 
-    py::class_<ShortRF, elements::Thin>(me, "ShortRF")
+    py::class_<ShortRF, elements::Thin> py_ShortRF(me, "ShortRF");
+    py_ShortRF
         .def(py::init<
                 amrex::ParticleReal const,
                 amrex::ParticleReal const,
@@ -337,8 +395,10 @@ void init_elements(py::module& m)
              "A short RF cavity element."
         )
     ;
+    register_beamoptics_push(py_ShortRF);
 
-    py::class_<SoftSolenoid, elements::Thick>(me, "SoftSolenoid")
+    py::class_<SoftSolenoid, elements::Thick> py_SoftSolenoid(me, "SoftSolenoid");
+    py_SoftSolenoid
         .def(py::init<
                  amrex::ParticleReal const,
                  amrex::ParticleReal const,
@@ -353,8 +413,10 @@ void init_elements(py::module& m)
              "A soft-edge solenoid."
         )
     ;
+    register_beamoptics_push(py_SoftSolenoid);
 
-    py::class_<Sol, elements::Thick>(me, "Sol")
+    py::class_<Sol, elements::Thick> py_Sol(me, "Sol");
+    py_Sol
         .def(py::init<
                 amrex::ParticleReal const,
                 amrex::ParticleReal const,
@@ -363,17 +425,21 @@ void init_elements(py::module& m)
              "An ideal hard-edge Solenoid magnet."
         )
     ;
+    register_beamoptics_push(py_Sol);
 
-    py::class_<PRot, elements::Thin>(me, "PRot")
+    py::class_<PRot, elements::Thin> py_PRot(me, "PRot");
+    py_PRot
         .def(py::init<
                 amrex::ParticleReal const,
                 amrex::ParticleReal const>(),
              py::arg("phi_in"), py::arg("phi_out"),
              "An exact pole-face rotation in the x-z plane. Both angles are in degrees."
         )
     ;
+    register_beamoptics_push(py_PRot);
 
-    py::class_<SoftQuadrupole, elements::Thick>(me, "SoftQuadrupole")
+    py::class_<SoftQuadrupole, elements::Thick> py_SoftQuadrupole(me, "SoftQuadrupole");
+    py_SoftQuadrupole
         .def(py::init<
                  amrex::ParticleReal,
                  amrex::ParticleReal,
@@ -388,5 +454,11 @@ void init_elements(py::module& m)
              "A soft-edge quadrupole."
         )
     ;
+    register_beamoptics_push(py_SoftQuadrupole);
 
+    // free-standing push function
+    m.def("push", &Push,
+        py::arg("pc"), py::arg("element"), py::arg("step")=0,
+        "Push particles through an element"
+    );
 }