Element Push: Restructure

src/ImpactX.cpp

@@ -183,7 +183,7 @@ namespace impactx
                 // assuming that the distribution did not change
 
                 // push all particles with external maps
-                Push(*m_particle_container, element_variant);
+                Push(*m_particle_container, element_variant, global_step);
 
                 // just prints an empty newline at the end of the slice_step
                 amrex::Print() << "\n";

src/particles/Push.H

@@ -20,11 +20,13 @@ namespace impactx
 {
     /** Push particles
      *
-     * @param pc container of the particles to push
-     * @param element_variant a single element to push the particles through
+     * @param[inout] pc container of the particles to push
+     * @param[inout] element_variant a single element to push the particles through
+     * @param[in] step global step for diagnostics
      */
     void Push (ImpactXParticleContainer & pc,
-               KnownElements const & element_variant);
+               KnownElements & element_variant,
+               int step);
 
 } // namespace impactx
 

src/particles/Push.cpp

@@ -10,56 +10,30 @@
 #include "Push.H"
 
 #include <AMReX_BLProfiler.H>
-#include <AMReX_Extension.H>  // for AMREX_RESTRICT
-#include <AMReX_REAL.H>       // for ParticleReal
 
-#include <utility>
+#include <string>
+#include <variant>
 
 
 namespace impactx
 {
     void Push (ImpactXParticleContainer & pc,
-               KnownElements const & element_variant)
+               KnownElements & element_variant,
+               int step)
     {
-        // performance profiling per element
-        std::string element_name;
-        std::visit([&element_name](auto&& element){ element_name = element.name; }, element_variant);
-        std::string const profile_name = "impactx::Push::" + element_name;
-        BL_PROFILE("impactx::Push");
-        BL_PROFILE(profile_name);
-
-        using namespace amrex::literals; // for _rt and _prt
-
-        // preparing to access reference particle data: RefPart
-        RefPart & ref_part = pc.GetRefParticle();
-
-        // loop over refinement levels
-        int const nLevel = pc.finestLevel();
-        for (int lev = 0; lev <= nLevel; ++lev)
+        // here we just access the element by its respective type
+        std::visit([&pc, step](auto&& element)
         {
-            // get simulation geometry information
-            //const amrex::Geometry& gm = this->Geom(lev);
-            //const auto prob_lo = gm.ProbLo();
-
-            // loop over all particle boxes
-            using ParIt = ImpactXParticleContainer::iterator;
-#ifdef AMREX_USE_OMP
-#pragma omp parallel if (amrex::Gpu::notInLaunchRegion())
-#endif
-            for (ParIt pti(pc, lev); pti.isValid(); ++pti) {
-                // here we just access the element by its respective type
-                std::visit(
-                    [&pti, &ref_part](auto element) {
-                        // push reference particle in global coordinates
-                        element(ref_part);
-
-                        // push beam particles relative to reference particle
-                        element(pti, ref_part);
-                    },
-                    element_variant
-                );
-            } // end loop over all particle boxes
-        } // env mesh-refinement level loop
+            // 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);
+        }, element_variant);
     }
 
 } // namespace impactx

src/particles/PushAll.H

@@ -0,0 +1,59 @@
+/* Copyright 2022-2023 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
+ */
+#ifndef IMPACTX_PUSH_ALL_H
+#define IMPACTX_PUSH_ALL_H
+
+#include "particles/ImpactXParticleContainer.H"
+
+
+namespace impactx
+{
+    /** Push all particles in a particle container.
+     *
+     * This element pushes first the reference particle, then all other particles.
+     * All particles are pushed independently with the same logic.
+     * Particles are relative to the reference particle.
+     *
+     * @param[in,out] pc particle container to push
+     * @param[in,out] element the beamline element
+     * @param[in] step global step for diagnostics
+     */
+    template<typename T_Element>
+    void push_all (
+            ImpactXParticleContainer & pc,
+            T_Element & element,
+            [[maybe_unused]] int step
+    )
+    {
+        // preparing to access reference particle data: RefPart
+        RefPart & ref_part = pc.GetRefParticle();
+
+        // loop over refinement levels
+        int const nLevel = pc.finestLevel();
+        for (int lev = 0; lev <= nLevel; ++lev)
+        {
+            // loop over all particle boxes
+            using ParIt = ImpactXParticleContainer::iterator;
+#ifdef AMREX_USE_OMP
+#pragma omp parallel if (amrex::Gpu::notInLaunchRegion())
+#endif
+            for (ParIt pti(pc, lev); pti.isValid(); ++pti) {
+                // push reference particle in global coordinates
+                element(ref_part);
+
+                // push beam particles relative to reference particle
+                element(pti, ref_part);
+            } // end loop over all particle boxes
+        } // env mesh-refinement level loop
+    }
+
+} // namespace impactx
+
+#endif // IMPACTX_PUSH_ALL_H

src/particles/elements/None.H

@@ -31,6 +31,14 @@ namespace impactx
         {
         }
 
+        /** Push all particles - nothing to do here */
+        void operator() (
+                ImpactXParticleContainer & /* pc */,
+                int /* step */
+        ) {
+            // nothing to do
+        }
+
         /** Push all particles - nothing to do here */
         void operator() (
             ImpactXParticleContainer::iterator & /* pti */,

src/particles/elements/Programmable.H

@@ -26,9 +26,21 @@ namespace impactx
         static constexpr auto name = "Programmable";
         using PType = ImpactXParticleContainer::ParticleType;
 
-        /** This element writes the particle beam out to openPMD data.
+        /** This element can be programmed
          */
-        Programmable () {}
+        Programmable (amrex::ParticleReal ds=0.0, int nslice=1)
+            : m_ds(ds), m_nslice(nslice)
+        {}
+
+        /** Push all particles relative to the reference particle
+         *
+         * @param[in,out] pc particle container to push
+         * @param[in] step global step for diagnostics
+         */
+        void operator() (
+            ImpactXParticleContainer & pc,
+            int step
+        ) const;
 
         /** Push all particles relative to the reference particle */
         void operator() (
@@ -68,6 +80,7 @@ namespace impactx
         amrex::ParticleReal m_ds = 0.0; //! segment length in m
         int m_nslice = 1; //! number of slices used for the application of space charge
 
+        std::function<void(ImpactXParticleContainer *, int)> m_push; //! hook for push of whole container
         std::function<void(ImpactXParticleContainer::iterator *, RefPart &)> m_beam_particles; //! hook for beam particles
         std::function<void(RefPart &)> m_ref_particle; //! hook for reference particle
     };

src/particles/elements/Programmable.cpp

@@ -9,17 +9,33 @@
  */
 
 #include "Programmable.H"
+#include "particles/PushAll.H"
 
 
 namespace impactx
 {
+    void
+    Programmable::operator() (
+        ImpactXParticleContainer & pc,
+        int step
+    ) const
+    {
+        if (m_push == nullptr) {
+            push_all(pc, *this, step);
+        }
+        else {
+            m_push(&pc, step);
+        }
+    }
+
     void
     Programmable::operator() (
         ImpactXParticleContainer::iterator & pti,
         RefPart & ref_part
     ) const
     {
         if (m_beam_particles == nullptr)
+            // TODO: only if verbose mode is set
             amrex::AllPrint() << "Programmable element - all particles: NO HOOK\n";
         else
             m_beam_particles(&pti, ref_part);
@@ -29,6 +45,7 @@ namespace impactx
     Programmable::operator() (RefPart & ref_part) const
     {
         if (m_ref_particle == nullptr)
+            // TODO: only if verbose mode is set
             amrex::AllPrint() << "Programmable element - ref particles: NO HOOK\n";
         else
             m_ref_particle(ref_part);

src/particles/elements/mixin/beamoptic.H

@@ -11,10 +11,13 @@
 #define IMPACTX_ELEMENTS_MIXIN_BEAMOPTIC_H
 
 #include "particles/ImpactXParticleContainer.H"
+#include "particles/PushAll.H"
 
-#include <AMReX_Extension.H>
+#include <AMReX_Extension.H> // for AMREX_RESTRICT
 #include <AMReX_REAL.H>
 
+#include <type_traits>
+
 
 namespace impactx::elements
 {
@@ -130,6 +133,21 @@ namespace detail
     template<typename T_Element>
     struct BeamOptic
     {
+        /** Push first the reference particle, then all other particles */
+        void operator() (
+            ImpactXParticleContainer & pc,
+            int step
+        )
+        {
+            static_assert(
+                std::is_base_of_v<BeamOptic, T_Element>,
+                "BeamOptic can only be used as a mixin class!"
+            );
+
+            T_Element& element = *static_cast<T_Element*>(this);
+            push_all(pc, element, step);
+        }
+
         /** This pushes the particles on a particle iterator tile or box.
          *
          * Particles are relative to the reference particle.
@@ -140,7 +158,13 @@ namespace detail
          void operator() (
             ImpactXParticleContainer::iterator & pti,
             RefPart & AMREX_RESTRICT ref_part
-         ) {
+         )
+         {
+            static_assert(
+                std::is_base_of_v<BeamOptic, T_Element>,
+                "BeamOptic can only be used as a mixin class!"
+            );
+
             T_Element& element = *static_cast<T_Element*>(this);
             detail::push_all_particles<T_Element>(pti, ref_part, element);
          }

src/python/elements.cpp

@@ -88,7 +88,7 @@ void init_elements(py::module& m)
         .def_property_readonly("ds", &elements::Thin::ds)
     ;
 
-    // beam optics below
+    // beam optics
 
     py::class_<ConstF, elements::Thick>(me, "ConstF")
         .def(py::init<
@@ -148,7 +148,10 @@ void init_elements(py::module& m)
     ;
 
     py::class_<Programmable>(me, "Programmable")
-        .def(py::init<>(),
+        .def(py::init<
+                 amrex::ParticleReal,
+                 int>(),
+             py::arg("ds") = 0.0, py::arg("nslice") = 1,
              "A programmable beam optics element."
         )
         .def_property("nslice",
@@ -159,17 +162,26 @@ void init_elements(py::module& m)
               [](Programmable & p) { return p.ds(); },
               [](Programmable & p, amrex::ParticleReal ds) { p.m_ds = ds; }
         )
+        .def_property("push",
+              [](Programmable & p) { return p.m_push; },
+              [](Programmable & p,
+                 std::function<void(ImpactXParticleContainer *, int)> new_hook
+              ) { p.m_push = new_hook; },
+              "hook for push of whole container (pc, step)"
+        )
         .def_property("beam_particles",
               [](Programmable & p) { return p.m_beam_particles; },
               [](Programmable & p,
                  std::function<void(ImpactXParticleContainer::iterator *, RefPart &)> new_hook
-              ) { p.m_beam_particles = new_hook; }
+              ) { p.m_beam_particles = new_hook; },
+              "hook for beam particles (pti, RefPart)"
         )
         .def_property("ref_particle",
               [](Programmable & p) { return p.m_ref_particle; },
               [](Programmable & p,
                  std::function<void(RefPart &)> new_hook
-              ) { p.m_ref_particle = new_hook; }
+              ) { p.m_ref_particle = new_hook; },
+              "hook for reference particle (RefPart)"
         )
     ;