Add KICKER to accepted MAD-X input (ECP-WarpX#419)

* Accept KICKER in MAD-X input

Added the logic to read lines from MAD-X input files that contain
kicker elements. Also added a MAD-X example input and a Python file that
reads it. The output and analysis script results match with the original
kicker example.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add kicker test with MAD-X input to ctests

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update license in examples/kicker/run_kicker_madx.py

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

examples/CMakeLists.txt

@@ -619,3 +619,13 @@ add_impactx_test(kicker.py
     examples/kicker/analysis_kicker.py
     OFF  # no plot script yet
 )
+# copy MAD-X lattice file
+file(COPY ${ImpactX_SOURCE_DIR}/examples/kicker/kicker.madx
+        DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kicker_madx.py)
+add_impactx_test(kicker_madx.py
+        examples/kicker/run_kicker_madx.py
+        OFF   # ImpactX MPI-parallel
+        ON   # ImpactX Python interface
+        examples/kicker/analysis_kicker.py
+        OFF  # no plot script yet
+)

examples/kicker/README.rst

@@ -7,7 +7,7 @@ This test applies two transverse momentum kicks, first in the horizontal directi
 
 We use a 2 GeV electron beam.
 
-The second beam moments should be unchanged, but the first beam moments corresponding to :math:`\p_x` and :math:`\p_y` should change according to the size of the kick.
+The second beam moments should be unchanged, but the first beam moments corresponding to :math:`p_x` and :math:`p_y` should change according to the size of the kick.
 
 In this test, the initial and final values of :math:`\sigma_x`, :math:`\sigma_y`, :math:`\sigma_t`, :math:`\epsilon_x`, :math:`\epsilon_y`, and :math:`\epsilon_t` must agree with nominal values.
 

examples/kicker/kicker.madx

@@ -0,0 +1,8 @@
+beam, particle=electron, energy=2.0;
+
+M1: MONITOR, L=0.0;
+HK1: KICKER, hkick=2.0e-3, vkick=0.0;
+VK1: KICKER, hkick=0.0, vkick=3.0e-3;
+
+KICKLATTICE: Line=(M1,HK1,VK1,M1);
+USE, SEQUENCE = KICKLATTICE;

examples/kicker/run_kicker_madx.py

@@ -0,0 +1,54 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2023 ImpactX contributors
+# Authors: Chad Mitchell, Axel Huebl, Marco Garten
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+
+import amrex.space3d as amr
+from impactx import ImpactX, RefPart, distribution, elements
+
+sim = ImpactX()
+
+# set numerical parameters and IO control
+sim.particle_shape = 2  # B-spline order
+sim.space_charge = False
+# sim.diagnostics = False  # benchmarking
+sim.slice_step_diagnostics = True
+
+# domain decomposition & space charge mesh
+sim.init_grids()
+
+# load a 2 GeV electron beam with an initial
+# unnormalized rms emittance of  nm
+bunch_charge_C = 1.0e-9  # used without space charge
+npart = 10000  # number of macro particles
+
+#   reference particle
+ref = sim.particle_container().ref_particle().load_file("kicker.madx")
+
+#   particle bunch
+distr = distribution.Waterbag(
+    sigmaX=4.0e-3,
+    sigmaY=4.0e-3,
+    sigmaT=1.0e-3,
+    sigmaPx=3.0e-4,
+    sigmaPy=3.0e-4,
+    sigmaPt=2.0e-3,
+)
+sim.add_particles(bunch_charge_C, distr, npart)
+
+# add beam diagnostics
+monitor = elements.BeamMonitor("monitor", backend="h5")
+
+# design the accelerator lattice
+sim.lattice.load_file("kicker.madx", nslice=1)
+
+# run simulation
+sim.evolve()
+
+# clean shutdown
+del sim
+amr.finalize()

src/python/impactx/MADXParser.py

@@ -99,6 +99,13 @@ def __init__(self):
         # don't count name and type --> len - 2
         self.__nDipedge = 2 * (len(self.__dipedge) - 2)
 
+        self.__kicker = {"name": "", "hkick": 0.0, "vkick": 0.0, "type": "kicker"}
+
+        self.__kicker_pattern = r"(.*):kicker,(.*)=(.*),(.*)=(.*);"
+
+        # don't count name and type --> len - 2
+        self.__nKicker = 2 * (len(self.__kicker) - 2)
+
         self.beam = {
             "energy": 0.0,
             # TODO extend by 'PC'
@@ -290,6 +297,29 @@ def parse(self, fn):
 
                     self.__elements.append(self.__dipedge.copy())
 
+                elif "kicker" in line:
+                    obj = re.match(self.__kicker_pattern, line)
+
+                    # first tag is name
+                    self.__kicker["name"] = obj.group(1)
+
+                    for i in range(2, self.__nKicker + 2, 2):
+                        if obj.group(i) in self.__kicker:
+                            self.__kicker[obj.group(i)] = float(obj.group(i + 1))
+                        else:
+                            raise MADXInputError(
+                                "Kicker",
+                                "Line "
+                                + str(nLine)
+                                + ": Parameter "
+                                + "'"
+                                + obj.group(i)
+                                + "'"
+                                + " does not exist for kicker.",
+                            )
+
+                    self.__elements.append(self.__kicker.copy())
+
                 elif "marker" in line:
                     pass
 
@@ -352,7 +382,10 @@ def parse(self, fn):
 
                     self.sequence["name"] = obj.group(1)
                 else:
-                    raise MADXInputError("", "Error at line " + str(nLine))
+                    raise MADXInputError(
+                        ("Error at line " + str(nLine), "Parsed line: " + str(line)),
+                        with_traceback=True,
+                    )
 
         # 14. Oct. 2017,
         # https://stackoverflow.com/questions/7900882/extract-item-from-list-of-dictionaries
@@ -410,14 +443,14 @@ def _noWhitespace(self, string):
         https://stackoverflow.com/questions/3739909/how-to-strip-all-whitespace-from-string
 
         """
-        return string.replace(" ", "")
+        return "".join(string.split())
 
     def __str__(self):
         if self.__lattice:
             length = 0.0
 
-            # drift, monitor, dipole, solenoid, quadrupole, dipedge
-            nTypes = [0, 0, 0, 0, 0, 0]
+            # drift, monitor, dipole, solenoid, quadrupole, dipedge, kicker
+            nTypes = [0, 0, 0, 0, 0, 0, 0]
 
             for elem in self.__lattice["elem"]:
                 for e in self.__elements:
@@ -437,6 +470,8 @@ def __str__(self):
                             nTypes[4] += 1
                         elif e["type"] == "dipedge":
                             nTypes[5] += 1
+                        elif e["type"] == "kicker":
+                            nTypes[6] += 1
                         break
 
             sign = "*" * 70
@@ -468,6 +503,9 @@ def __str__(self):
                 + "            *     #dipedge:\t"
                 + str(nTypes[5])
                 + "\n"
+                + "            *     #kicker:\t"
+                + str(nTypes[6])
+                + "\n"
                 + "           beam:\t\n"
                 + "            *     particle:\t"
                 + self.beam["particle"]
@@ -507,6 +545,9 @@ def getBeamline(self):
                         elif e["type"] == "dipedge":
                             # print("Dipedge H= ", e["h"], " E1 = ", e["e1"])
                             beamline.append(e)
+                        elif e["type"] == "kicker":
+                            # print("Kicker hkick= ", e["hkick"], " vkick = ", e["vkick"])
+                            beamline.append(e)
                         else:
                             print("Skipping element type " + "'" + e["type"] + "'")
             return beamline

src/python/impactx/madx_to_impactx.py

@@ -31,6 +31,10 @@ def lattice(parsed_beamline, nslice=1):
         "SOLENOID": "Sol",  # Ideal, thick Solenoid: MAD-X user guide 10.9 p78
         "QUADRUPOLE": "Quad",  # Quadrupole
         "DIPEDGE": "DipEdge",
+        # Kicker, idealized thin element,
+        # MADX also defines length "L" and a roll angle around the longitudinal axis "TILT"
+        # https://mad.web.cern.ch/mad/webguide/manual.html#Ch11.S11
+        "KICKER": "Kicker",
         # note: in MAD-X, this keeps track only of the beam centroid,
         # "In addition it serves to record the beam position for closed orbit correction."
         "MONITOR": "BeamMonitor",  # drift + output diagnostics
@@ -70,6 +74,13 @@ def lattice(parsed_beamline, nslice=1):
                         K2=d["fint"],
                     )
                 )
+            elif d["type"] == "kicker":
+                impactx_beamline.append(
+                    elements.Kicker(
+                        xkick=d["hkick"],
+                        ykick=d["vkick"],
+                    )
+                )
             elif d["type"] == "monitor":
                 if d["l"] > 0:
                     impactx_beamline.append(elements.Drift(ds=d["l"], nslice=nslice))