From d89f17ec8865f8365a28113785abb7ac04c3dc95 Mon Sep 17 00:00:00 2001 From: Marco Garten Date: Tue, 22 Aug 2023 12:04:09 -0700 Subject: [PATCH] Add KICKER to accepted MAD-X input (#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 --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Axel Huebl --- examples/CMakeLists.txt | 10 +++++ examples/kicker/README.rst | 2 +- examples/kicker/kicker.madx | 8 ++++ examples/kicker/run_kicker_madx.py | 54 +++++++++++++++++++++++++++ src/python/impactx/MADXParser.py | 49 ++++++++++++++++++++++-- src/python/impactx/madx_to_impactx.py | 11 ++++++ 6 files changed, 129 insertions(+), 5 deletions(-) create mode 100644 examples/kicker/kicker.madx create mode 100644 examples/kicker/run_kicker_madx.py diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index c57fc639e..9a4bdc373 100644 --- a/examples/CMakeLists.txt +++ b/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 +) diff --git a/examples/kicker/README.rst b/examples/kicker/README.rst index e4a3a0ca7..be43fd2dd 100644 --- a/examples/kicker/README.rst +++ b/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. diff --git a/examples/kicker/kicker.madx b/examples/kicker/kicker.madx new file mode 100644 index 000000000..4f5056943 --- /dev/null +++ b/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; diff --git a/examples/kicker/run_kicker_madx.py b/examples/kicker/run_kicker_madx.py new file mode 100644 index 000000000..88147c770 --- /dev/null +++ b/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() diff --git a/src/python/impactx/MADXParser.py b/src/python/impactx/MADXParser.py index c25c1f725..8a5d348f2 100644 --- a/src/python/impactx/MADXParser.py +++ b/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 diff --git a/src/python/impactx/madx_to_impactx.py b/src/python/impactx/madx_to_impactx.py index 4e8993ab4..1ad9cee26 100644 --- a/src/python/impactx/madx_to_impactx.py +++ b/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))