Kicker Update API: Enums & Strings

Update the C++ API to use an enum for the unit system and the Python API & inputs file syntax to accept a string. That makes the parameters at the call sites self-describing. Also clarifies in doc strings what dimensionless is referring to (units of magnetic rigidity of the reference particle).
ECP-WarpX · Aug 9, 2023 · fd5fe09 · fd5fe09
1 parent 5d9862c
commit fd5fe09
Show file tree

Hide file tree

Showing 8 changed files with 47 additions and 29 deletions.
diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst
@@ -489,7 +489,7 @@ Lattice Elements
 
             * ``<element_name>.ykick`` (``float``, dimensionless OR in T-m) the vertical kick strength
 
-            * ``<element_name>.units`` (``integer``, dimensionless) specification of units (default: ``0``)
+            * ``<element_name>.units`` (``string``) specification of units: ``dimensionless`` (default, in units of the magnetic rigidity of the reference particle) or ``T-m``
 
         * ``beam_monitor`` a beam monitor, writing all beam particles at fixed ``s`` to openPMD files.
           If the same element name is used multiple times, then an output series is created with multiple outputs.

diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst
@@ -477,7 +477,7 @@ This module provides elements for the accelerator lattice.
 
    :param xkick: horizontal kick strength (dimensionless OR T-m)
    :param ykick: vertical kick strength (dimensionless OR T-m)
-   :param units: specification of units (0 - dimensionless, 1 - T-m)
+   :param units: specification of units (``"dimensionless"`` in units of the magnetic rigidity of the reference particle or ``"T-m"``)
 
 .. py:class:: impactx.elements.Multipole(multipole, K_normal, K_skew)
 

diff --git a/examples/kicker/README.rst b/examples/kicker/README.rst
@@ -1,6 +1,6 @@
 .. _examples-kicker:
 
-Test of a transverse kicker
+Test of a Transverse Kicker
 ===========================
 
 This test applies two transverse momentum kicks, first in the horizontal direction (2 mrad) and then in the vertical direction (3 mrad).

diff --git a/examples/kicker/input_kicker.in b/examples/kicker/input_kicker.in
@@ -27,12 +27,12 @@ monitor.type = beam_monitor
 monitor.backend = h5
 
 hkick.type = kicker
-hkick.xkick = 2.0e-3      //2 mrad horizontal kick
+hkick.xkick = 2.0e-3      # 2 mrad horizontal kick
 hkick.ykick = 0.0
 
 vkick.type = kicker
 vkick.xkick = 0.0
-vkick.ykick = 3.0e-3     //3 mrad vertical kick
+vkick.ykick = 3.0e-3     # 3 mrad vertical kick
 
 ###############################################################################
 # Algorithms

diff --git a/examples/kicker/run_kicker.py b/examples/kicker/run_kicker.py
@@ -47,8 +47,8 @@
 # design the accelerator lattice
 kicklattice = [
     monitor,
-    elements.Kicker(xkick=2.0e-3, ykick=0.0, units=0),
-    elements.Kicker(xkick=0.0, ykick=3.0e-3, units=0),
+    elements.Kicker(xkick=2.0e-3, ykick=0.0, units="dimensionless"),
+    elements.Kicker(xkick=0.0, ykick=3.0e-3, units="dimensionless"),
     monitor,
 ]
 # assign a lattice

diff --git a/src/initialization/InitElement.cpp b/src/initialization/InitElement.cpp
@@ -233,10 +233,15 @@ namespace detail
             m_lattice.emplace_back( ChrAcc(ds, ez, bz, nslice) );
         } else if (element_type == "kicker") {
             amrex::Real xkick, ykick;
-            int units = 0;
+            std::string units_str = "dimensionless";
             pp_element.get("xkick", xkick);
             pp_element.get("ykick", ykick);
-            pp_element.queryAdd("units", units);
+            pp_element.queryAdd("units", units_str);
+            AMREX_ALWAYS_ASSERT_WITH_MESSAGE(units_str == "dimensionless" || units_str == "T-m",
+                                             element_name + ".units must be \"dimensionless\" or \"T-m\"");
+            Kicker::UnitSystem units = units_str == "dimensionless" ?
+                Kicker::UnitSystem::dimensionless :
+                Kicker::UnitSystem::Tm;
             m_lattice.emplace_back( Kicker(xkick, ykick, units) );
         } else if (element_type == "beam_monitor") {
             std::string openpmd_name = element_name;

diff --git a/src/particles/elements/Kicker.H b/src/particles/elements/Kicker.H
@@ -31,16 +31,22 @@ namespace impactx
         static constexpr auto name = "Kicker";
         using PType = ImpactXParticleContainer::ParticleType;
 
+        enum UnitSystem
+        {
+            dimensionless = 0, //! in units of the magnetic rigidity of the reference particle
+            Tm = 1  //! T-m
+        };
+
         /** A thin transverse kicker element, inducing a transverse
          *  momentum displacement from the reference orbit.
          *
          * @param xkick Strength of horizontal kick
          * @param ykick Strength of vertical kick
-         * @param unit (0 - dimensionless, 1 - T-m)
+         * @param unit units of xkick and ykick
          */
-        Kicker( amrex::ParticleReal const xkick,
-                amrex::ParticleReal const ykick,
-                int const unit )
+        Kicker (amrex::ParticleReal xkick,
+                amrex::ParticleReal ykick,
+                UnitSystem unit)
         : m_xkick(xkick), m_ykick(ykick), m_unit(unit)
         {
         }
@@ -59,12 +65,12 @@ namespace impactx
          */
         AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
         void operator() (
-                PType& AMREX_RESTRICT p,
-                amrex::ParticleReal & AMREX_RESTRICT px,
-                amrex::ParticleReal & AMREX_RESTRICT py,
-                amrex::ParticleReal & AMREX_RESTRICT pt,
-                [[maybe_unused]] RefPart const & refpart) const {
-
+            PType& AMREX_RESTRICT p,
+            amrex::ParticleReal & AMREX_RESTRICT px,
+            amrex::ParticleReal & AMREX_RESTRICT py,
+            amrex::ParticleReal & AMREX_RESTRICT pt,
+            [[maybe_unused]] RefPart const & refpart) const
+        {
             using namespace amrex::literals; // for _rt and _prt
 
             // access AoS data such as positions and cpu/id
@@ -75,7 +81,7 @@ namespace impactx
             // normalize quad units to MAD-X convention if needed
             amrex::ParticleReal dpx = m_xkick;
             amrex::ParticleReal dpy = m_ykick;
-            if (m_unit == 1) {
+            if (m_unit == UnitSystem::Tm) {
                   dpx /= refpart.rigidity_Tm();
                   dpy /= refpart.rigidity_Tm();
             }
@@ -99,7 +105,6 @@ namespace impactx
             px = pxout;
             py = pyout;
             pt = ptout;
-
         }
 
         /** This pushes the reference particle. */
@@ -108,8 +113,7 @@ namespace impactx
     private:
         amrex::ParticleReal m_xkick; //! horizontal kick strength
         amrex::ParticleReal m_ykick; //! vertical kick strength
-        int m_unit; //! Kicks are for 0 dimensionless, or for 1 in T-m."
-
+        UnitSystem m_unit; //! Kicks are for 0 dimensionless, or for 1 in T-m."
     };
 
 } // namespace impactx

diff --git a/src/python/elements.cpp b/src/python/elements.cpp
@@ -183,12 +183,21 @@ void init_elements(py::module& m)
     ;
 
     py::class_<Kicker, elements::Thin>(me, "Kicker")
-        .def(py::init<
-                amrex::ParticleReal const,
-                amrex::ParticleReal const,
-                int const>(),
-             py::arg("xkick"), py::arg("ykick"), py::arg("units") = 0,
-             "A thin transverse kicker element. Kicks are for unit 0 dimensionless, or for 1 in T-m."
+        .def(py::init([](
+                amrex::ParticleReal xkick,
+                amrex::ParticleReal ykick,
+                std::string units)
+             {
+                 if (units != "dimensionless" && units != "T-m")
+                     throw std::runtime_error("units must be \"dimensionless\" or \"T-m\"");
+
+                 Kicker::UnitSystem u = units == "dimensionless" ?
+                                            Kicker::UnitSystem::dimensionless :
+                                            Kicker::UnitSystem::Tm;
+                 return new Kicker(xkick, ykick, u);
+             }),
+             py::arg("xkick"), py::arg("ykick"), py::arg("units") = "dimensionless",
+             "A thin transverse kicker element. Kicks are for units \"dimensionless\" or in \"T-m\"."
         )
     ;