Add MADX Input Parser (#214)

* Add MADX Input Parser Using Andreas Adelmann's pymadxparser * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Copy MADX input via CMakeLists.txt Co-authored-by: Axel Huebl <[email protected]> * Add beam function to MADX translation * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Clean up: Test & Input File Names * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * MADXParser: Regex Literals Use proper regex literals. Avoids: ``` DeprecationWarning: invalid escape sequence \( ``` et al. * Lattice: Cleaner User API & Docs * Fix chicane MADX example and add parser elements * Add additional parameters to DipEdge and adjust formatting * Update examples according to API changes Caution: Examples still not working again * Cleaning * Rename elements in chicane.madx to avoid matching of `drift` in `use, sequence` line with parser. * Fix: element type instead of name change MADXParser Quad type to Quadrupole to stay consistent with MADX input * Update Copyright MAD-X Parser Script Confirmed today via email :) * Fix: DipEdge is a thin element * Fix Chicane input to make test pass also fix energy conversion in madx_to_impactx * Comment Debug Printing Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Axel Huebl <[email protected]>
ECP-WarpX · Sep 2, 2022 · e5dabe5 · e5dabe5
1 parent 55f89f4
commit e5dabe5
Show file tree

Hide file tree

Showing 10 changed files with 873 additions and 25 deletions.
diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst
@@ -240,6 +240,12 @@ Particles
 
       Write-only: Set reference particle energy.
 
+   .. py:method:: load_file(madx_file)
+
+      Load reference particle information from a MAD-X file.
+
+      :param madx_file: file name to MAD-X file with a ``BEAM`` entry
+
 
 Initial Beam Distributions
 --------------------------
@@ -305,6 +311,25 @@ This module provides elements for the accelerator lattice.
 
    An iterable, ``list``-like type of elements.
 
+   .. py:method:: clear()
+
+      Clear the list to become empty.
+
+   .. py:method:: extend(list)
+
+      Add a list of elements to the list.
+
+   .. py:method:: append(element)
+
+      Add a single element to the list.
+
+   .. py:method:: load_file(madx_file, nslice=1)
+
+      Load and append an accelerator lattice description from a MAD-X file.
+
+      :param madx_file: file name to MAD-X file with beamline elements
+      :param nslice: number of slices used for the application of space charge
+
 .. py:class:: impactx.elements.ConstF(ds, kx, ky, kt, nslice=1)
 
    A linear Constant Focusing element.

diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt
@@ -104,6 +104,19 @@ add_impactx_test(FODO.py
     examples/fodo/plot_fodo.py
 )
 
+# Python MADX: FODO Cell ######################################################
+#
+# copy MAD-X lattice file
+file(COPY ${ImpactX_SOURCE_DIR}/examples/fodo/fodo.madx DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/FODO.MADX.py)
+
+add_impactx_test(FODO.MADX.py
+    examples/fodo/run_fodo_madx.py
+      OFF  # ImpactX MPI-parallel
+      ON   # ImpactX Python interface
+    examples/fodo/analysis_fodo.py
+    examples/fodo/plot_fodo.py
+)
+
 # Python: MPI-parallel FODO Cell ##############################################
 #
 add_impactx_test(FODO.py.MPI
@@ -134,6 +147,19 @@ add_impactx_test(chicane.py
     examples/chicane/plot_chicane.py
 )
 
+# Python MADX: Chicane ######################################################
+#
+# copy MAD-X lattice file
+file(COPY ${ImpactX_SOURCE_DIR}/examples/chicane/chicane.madx DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/chicane.MADX.py)
+
+add_impactx_test(chicane.MADX.py
+    examples/chicane/run_chicane_madx.py
+      OFF  # ImpactX MPI-parallel
+      ON   # ImpactX Python interface
+    examples/chicane/analysis_chicane.py
+    examples/chicane/plot_chicane.py
+)
+
 # Constant Focusing Channel ###################################################
 #
 add_impactx_test(cfchannel

diff --git a/examples/chicane/chicane.madx b/examples/chicane/chicane.madx
@@ -0,0 +1,17 @@
+beam, particle=electron, energy=5.0;
+
+D1: drift, L=5.0058489435;
+D2: drift, L=1.0;
+D3: drift, L=2.0;
+! TODO make this work with inline calculations
+! theta=0.50037/10.35;
+! inv_rho=1.0/10.35
+! TODO put `angle=theta` for SBENDs and `e1=theta` for DIPEDGEs with their right sign
+SB1: sbend, L=0.50037, angle=-0.04834492753623188, e1=0.000, e2=0.000, k1=0.00;
+SB2: sbend, L=0.50037, angle=0.04834492753623188, e1=0.000, e2=0.000, k1=0.00;
+! dipole edge elements
+DIPE1: dipedge, H=-0.0966183574879227, e1=-0.048345620280243, fint=0.000, hgap=0.000, tilt=0.00;
+DIPE2: dipedge, H=0.0966183574879227, e1=0.048345620280243, fint=0.000, hgap=0.000, tilt=0.00;
+
+CHICANE:     Line=(SB1,DIPE1,D1,DIPE2,SB2,D2,SB2,DIPE2,D1,DIPE1,SB1,D3);
+USE, SEQUENCE = CHICANE;
diff --git a/examples/chicane/run_chicane_madx.py b/examples/chicane/run_chicane_madx.py
@@ -0,0 +1,54 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022 ImpactX contributors
+# Authors: Marco Garten, Axel Huebl, Chad Mitchell
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+
+import amrex
+from impactx import ImpactX, RefPart, distribution, elements
+
+sim = ImpactX()
+
+# set numerical parameters and IO control
+sim.set_particle_shape(2)  # B-spline order
+sim.set_space_charge(False)
+# sim.set_diagnostics(False)  # benchmarking
+sim.set_slice_step_diagnostics(True)
+
+# domain decomposition & space charge mesh
+sim.init_grids()
+
+# load a 5 GeV electron beam with an initial
+# normalized transverse rms emittance of 1 um
+bunch_charge_C = 1.0e-9  # used with space charge
+npart = 10000  # number of macro particles
+
+#   reference particle
+ref = sim.particle_container().ref_particle().load_file("chicane.madx")
+
+#   particle bunch
+distr = distribution.Waterbag(
+    sigmaX=2.2951017632e-5,
+    sigmaY=1.3084093142e-5,
+    sigmaT=5.5555553e-8,
+    sigmaPx=1.598353425e-6,
+    sigmaPy=2.803697378e-6,
+    sigmaPt=2.000000000e-6,
+    muxpx=0.933345606203060,
+    muypy=0.933345606203060,
+    mutpt=0.999999961419755,
+)
+sim.add_particles(bunch_charge_C, distr, npart)
+
+# design the accelerator lattice
+sim.lattice.load_file("chicane.madx", nslice=25)
+
+# run simulation
+sim.evolve()
+
+# clean shutdown
+del sim
+amrex.finalize()
diff --git a/examples/fodo/fodo.madx b/examples/fodo/fodo.madx
@@ -0,0 +1,11 @@
+! TODO implement handling of title
+! TITLE,'FODO.MADX';
+BEAM, PARTICLE=ELECTRON,ENERGY=2.0;
+
+D1: DRIFT,L=0.25;
+D2: DRIFT,L=0.50;
+QF: QUADRUPOLE,L=1.0,K1=1.0;
+QD: QUADRUPOLE,L=1.0,K1=-1.0;
+
+FODO: LINE=(D1,QF,D2,QD,D1);
+USE, SEQUENCE = FODO;
diff --git a/examples/fodo/run_fodo_madx.py b/examples/fodo/run_fodo_madx.py
@@ -0,0 +1,55 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022 ImpactX contributors
+# Authors: Marco Garten, Axel Huebl, Chad Mitchell
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+
+import amrex
+from impactx import ImpactX, RefPart, distribution, elements
+
+sim = ImpactX()
+
+# set numerical parameters and IO control
+sim.set_particle_shape(2)  # B-spline order
+sim.set_space_charge(False)
+# sim.set_diagnostics(False)  # benchmarking
+sim.set_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 2 nm
+energy_MeV = 2.0e3  # reference energy
+bunch_charge_C = 1.0e-9  # used with space charge
+npart = 10000  # number of macro particles
+
+#   reference particle
+ref = sim.particle_container().ref_particle().load_file("fodo.madx")
+
+#   particle bunch
+distr = distribution.Waterbag(
+    sigmaX=3.9984884770e-5,
+    sigmaY=3.9984884770e-5,
+    sigmaT=1.0e-3,
+    sigmaPx=2.6623538760e-5,
+    sigmaPy=2.6623538760e-5,
+    sigmaPt=2.0e-3,
+    muxpx=-0.846574929020762,
+    muypy=0.846574929020762,
+    mutpt=0.0,
+)
+sim.add_particles(bunch_charge_C, distr, npart)
+
+# design the accelerator lattice
+sim.lattice.load_file("fodo.madx", nslice=25)
+
+# run simulation
+sim.evolve()
+
+# clean shutdown
+del sim
+amrex.finalize()
diff --git a/src/python/elements.cpp b/src/python/elements.cpp
@@ -31,25 +31,34 @@ void init_elements(py::module& m)
             return v;
         }))
 
-        .def("append", [](KnownElementsList &v, KnownElements el) { v.emplace_back(el); })
-
-        .def("extend", [](KnownElementsList &v, KnownElementsList l) {
-            for (auto const & el : l)
-                v.push_back(el);
-            return v;
-        })
-        .def("extend", [](KnownElementsList &v, py::list l) {
-            for (auto const & handle : l)
-            {
-                auto el = handle.cast<KnownElements>();
-                v.push_back(el);
-            }
-            return v;
-        })
-
-        .def("clear", &KnownElementsList::clear)
-        .def("pop_back", &KnownElementsList::pop_back)
-        .def("__len__", [](const KnownElementsList &v) { return v.size(); })
+        .def("append", [](KnownElementsList &v, KnownElements el) { v.emplace_back(el); },
+             "Add a single element to the list.")
+
+        .def("extend",
+             [](KnownElementsList &v, KnownElementsList l) {
+                 for (auto const & el : l)
+                    v.push_back(el);
+                 return v;
+             },
+             "Add a list of elements to the list.")
+        .def("extend",
+             [](KnownElementsList &v, py::list l) {
+                 for (auto const & handle : l)
+                 {
+                    auto el = handle.cast<KnownElements>();
+                    v.push_back(el);
+                 }
+                 return v;
+             },
+             "Add a list of elements to the list."
+         )
+
+        .def("clear", &KnownElementsList::clear,
+             "Clear the list to become empty.")
+        .def("pop_back", &KnownElementsList::pop_back,
+             "Return and remove the last element of the list.")
+        .def("__len__", [](const KnownElementsList &v) { return v.size(); },
+             "The length of the list.")
         .def("__iter__", [](KnownElementsList &v) {
             return py::make_iterator(v.begin(), v.end());
         }, py::keep_alive<0, 1>()) /* Keep list alive while iterator is used */