Cleaning

ECP-WarpX · Aug 31, 2022 · f88d685 · f88d685
1 parent ece6bab
commit f88d685
Show file tree

Hide file tree

Showing 7 changed files with 82 additions and 115 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
 --------------------------

diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt
@@ -114,7 +114,7 @@ add_impactx_test(FODO.MADX.py
       OFF  # ImpactX MPI-parallel
       ON   # ImpactX Python interface
     examples/fodo/analysis_fodo.py
-    examples/fodo/plot_fodo.py --save-png
+    examples/fodo/plot_fodo.py
 )
 
 # Python: MPI-parallel FODO Cell ##############################################
@@ -157,7 +157,7 @@ add_impactx_test(chicane.MADX.py
       OFF  # ImpactX MPI-parallel
       ON   # ImpactX Python interface
     examples/chicane/analysis_chicane.py
-    examples/chicane/plot_chicane.py --save-png
+    examples/chicane/plot_chicane.py
 )
 
 # Constant Focusing Channel ###################################################

diff --git a/examples/chicane/run_chicane_madx.py b/examples/chicane/run_chicane_madx.py
@@ -8,8 +8,7 @@
 
 
 import amrex
-from impactx import (ImpactX, MADXParser, RefPart, distribution, elements,
-                     read_beam, read_lattice)
+from impactx import ImpactX, RefPart, distribution, elements
 
 sim = ImpactX()
 
@@ -22,38 +21,16 @@
 # domain decomposition & space charge mesh
 sim.init_grids()
 
-try:
-    madx = MADXParser()
-
-    madx.parse("chicane.madx")
-
-except Exception as e:
-    print(f"Unexpected {e = }, {type(e) = }")
-    raise
-
-
-# print summary
-print(madx)
-
-beamline = madx.getBeamline()
-
-ref_particle_dict = read_beam(
-    madx.getParticle(),  # if particle species is known, mass, charge, and potentially energy are set to default
-    # TODO MADX parser needs to extract charge if it's given,
-    # TODO MADX parser needs to extract mass if it's given,
-    energy=float(madx.getEtot()),  # MADX default energy is in GeV
-)
-
-
-# load a 40 MeV electron beam with an initial
+# load a 5 GeV electron beam with an initial
 # normalized transverse rms emittance of 1 um
-
-# @TODO read CHARGE (single particle charge) from MADX as well
-charge_C = 1.0e-9  # used with space charge
-
-qm_qeeV = ref_particle_dict["mass"] / ref_particle_dict["charge"]  # charge/mass
+energy_MeV = 5.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("chicane.madx")
+
+#   particle bunch
 distr = distribution.Waterbag(
     sigmaX=2.2951017632e-5,
     sigmaY=1.3084093142e-5,
@@ -65,17 +42,9 @@
     muypy=0.933345606203060,
     mutpt=0.999999961419755,
 )
-sim.add_particles(qm_qeeV, charge_C, distr, npart)
+sim.add_particles(bunch_charge_C, distr, npart)
 
 # design the accelerator lattice
-ns = 25  # number of slices per ds in the element
-
-# set the energy in the reference particle
-sim.particle_container().ref_particle().set_energy_MeV(
-    ref_particle_dict["energy"], ref_particle_dict["mass"]
-)
-
-# assign a fodo segment
 sim.lattice.load_file("chicane.madx", nslice=25)
 
 # run simulation

diff --git a/examples/fodo/run_fodo_madx.py b/examples/fodo/run_fodo_madx.py
@@ -8,8 +8,7 @@
 
 
 import amrex
-from impactx import (ImpactX, MADXParser, RefPart, distribution, elements,
-                     read_beam, read_lattice)
+from impactx import ImpactX, RefPart, distribution, elements
 
 sim = ImpactX()
 
@@ -22,36 +21,16 @@
 # domain decomposition & space charge mesh
 sim.init_grids()
 
-try:
-    madx = MADXParser()
-
-    madx.parse("fodo.madx")
-
-except Exception as e:
-    print(f"Unexpected {e = }, {type(e) = }")
-    raise
-
-# print summary
-print(madx)
-
-beamline = madx.getBeamline()
-
-ref_particle_dict = read_beam(
-    madx.getParticle(),  # if particle species is known, mass, charge, and potentially energy are set to default
-    # TODO MADX parser needs to extract charge if it's given,
-    # TODO MADX parser needs to extract mass if it's given,
-    energy=float(madx.getEtot()),  # MADX default energy is in GeV
-)
-
 # load a 2 GeV electron beam with an initial
 # unnormalized rms emittance of 2 nm
-
-# @TODO read CHARGE (single particle charge) from MADX as well
-charge_C = 1.0e-9  # used with space charge
-
-qm_qeeV = ref_particle_dict["mass"] / ref_particle_dict["charge"]  # charge/mass
+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,
@@ -63,17 +42,9 @@
     muypy=0.846574929020762,
     mutpt=0.0,
 )
-sim.add_particles(qm_qeeV, charge_C, distr, npart)
+sim.add_particles(bunch_charge_C, distr, npart)
 
 # design the accelerator lattice
-ns = 25  # number of slices per ds in the element
-
-# set the energy in the reference particle
-sim.particle_container().ref_particle().set_energy_MeV(
-    ref_particle_dict["energy"], ref_particle_dict["mass"]
-)
-
-# assign a fodo segment
 sim.lattice.load_file("fodo.madx", nslice=25)
 
 # run simulation

diff --git a/src/python/impactx/MADXParser.py b/src/python/impactx/MADXParser.py
@@ -6,6 +6,7 @@
 #
 # 10.8.2022 Adapted for standalone use
 #
+import os
 import re
 import warnings
 
@@ -68,11 +69,11 @@ def __init__(self):
             "fint": 0.0,
             "hgap": 0.0,
             "tilt": 0.0,
-            "type": "dipedge"
+            "type": "dipedge",
         }
 
         self.__dipedge_pattern = (
-             r"(.*):dipedge,(.*)=(.*),(.*)=(.*),(.*)=(.*),(.*)=(.*),(.*)=(.*);"
+            r"(.*):dipedge,(.*)=(.*),(.*)=(.*),(.*)=(.*),(.*)=(.*),(.*)=(.*);"
         )
 
         # don't count name and type --> len - 2
@@ -118,6 +119,9 @@ def parse(self, fn):
 
         """
 
+        if not os.path.isfile(fn):
+            raise FileNotFoundError(f"File '{fn}' not found!")
+
         nLine = 0
 
         with open(fn, "r") as f:
@@ -202,28 +206,28 @@ def parse(self, fn):
 
                 elif "dipedge" in line:
 
-                        obj = re.match(self.__dipedge_pattern, line)
+                    obj = re.match(self.__dipedge_pattern, line)
 
-                        # first tag is name
-                        self.__dipedge["name"] = obj.group(1)
+                    # first tag is name
+                    self.__dipedge["name"] = obj.group(1)
 
-                        for i in range(2, self.__nDipedge + 2, 2):
+                    for i in range(2, self.__nDipedge + 2, 2):
 
-                            if obj.group(i) in self.__dipedge:
-                                self.__dipedge[obj.group(i)] = float(obj.group(i + 1))
-                            else:
-                                raise MADXInputError(
-                                    "DipEdge",
-                                    "Line "
-                                    + str(nLine)
-                                    + ": Parameter "
-                                    + "'"
-                                    + obj.group(i)
-                                    + "'"
-                                    + " does not exist for dipole edge."
-                                )
-
-                        self.__elements.append(self.__dipedge.copy())
+                        if obj.group(i) in self.__dipedge:
+                            self.__dipedge[obj.group(i)] = float(obj.group(i + 1))
+                        else:
+                            raise MADXInputError(
+                                "DipEdge",
+                                "Line "
+                                + str(nLine)
+                                + ": Parameter "
+                                + "'"
+                                + obj.group(i)
+                                + "'"
+                                + " does not exist for dipole edge.",
+                            )
+
+                    self.__elements.append(self.__dipedge.copy())
 
                 elif "marker" in line:
                     pass
@@ -369,7 +373,7 @@ def __str__(self):
                             nTypes[1] += 1
                         elif e["type"] == "quad":
                             nTypes[2] += 1
-                        elif e['type'] == 'dipedge':
+                        elif e["type"] == "dipedge":
                             nTypes[3] += 1
                         break
 
@@ -428,7 +432,7 @@ def getBeamline(self):
                             print("Quadrupole L= ", e["l"], " k1 = ", e["k1"])
                             beamline.append(e)
                         elif e["type"] == "dipedge":
-                            print("Dipedge H= ",e["h"]," E1 = ",e["e1"])
+                            print("Dipedge H= ", e["h"], " E1 = ", e["e1"])
                             beamline.append(e)
                         else:
                             print("Skipping element type " + "'" + e["type"] + "'")

diff --git a/src/python/impactx/__init__.py b/src/python/impactx/__init__.py
@@ -10,6 +10,10 @@
 # at this place we can enhance Python classes with additional methods written
 # in pure Python or add some other Python logic
 
-# MAD-X reader for beamline lattice elements
-#   adds an overload to existing methods
-elements.KnownElementsList.load_file = lambda madx_file, nslice=1 : read_lattice(madx_file, nslice)  # noqa
+# MAD-X file reader for beamline lattice elements
+elements.KnownElementsList.load_file = lambda self, madx_file, nslice=1: self.extend(
+    read_lattice(madx_file, nslice)
+)  # noqa
+
+# MAD-X file reader for reference particle
+RefPart.load_file = lambda self, madx_file: read_beam(self, madx_file)  # noqa
diff --git a/src/python/impactx/madx_to_impactx.py b/src/python/impactx/madx_to_impactx.py
@@ -10,7 +10,7 @@
 
 import scipy.constants as sc
 
-from impactx import elements
+from impactx import RefPart, elements
 
 from .MADXParser import MADXParser
 
@@ -47,19 +47,20 @@ def lattice(parsed_beamline, nslice=1):
                 impactx_beamline.append(
                     elements.Quad(ds=d["l"], k=d["k1"], nslice=nslice)
                 )
-            elif d['name'] == "sbend":
+            elif d["name"] == "sbend":
                 impactx_beamline.append(
-                    elements.Sbend(ds=d["l"], rc=d["l"]/d["angle"], nslice=nslice)
+                    elements.Sbend(ds=d["l"], rc=d["l"] / d["angle"], nslice=nslice)
                 )
-            elif d['name'] == "dipedge":
+            elif d["name"] == "dipedge":
                 impactx_beamline.append(
                     elements.DipEdge(
                         psi=d["e1"],
                         rc=1.0 / d["h"],
                         # MAD-X is using half the gap height
                         g=2.0 * d["hgap"],
-                        K2=d['fint'],
-                        nslice=nslice)
+                        K2=d["fint"],
+                        nslice=nslice,
+                    )
                 )
         else:
             raise NotImplementedError(
@@ -82,7 +83,9 @@ def read_lattice(madx_file, nslice=1):
     :return: list of ImpactX.KnownElements
     """
     madx = MADXParser()
-    beamline = madx.parse(madx_file).getBeamline()
+    madx.parse(madx_file)
+    beamline = madx.getBeamline()
+
     return lattice(beamline, nslice)
 
 
@@ -146,12 +149,22 @@ def beam(particle, charge=None, mass=None, energy=None):
     return reference_particle
 
 
-def read_beam(madx_file):
+def read_beam(ref: RefPart, madx_file):
     """
     Function that reads elements from a MAD-X file into a list of ImpactX.KnownElements
     :param madx_file: file name to MAD-X file with beamline elements
     :return: list of ImpactX.KnownElements
     """
     madx = MADXParser()
-    beamline = madx.parse(madx_file).getBeamline()
-    return beam(beamline)
+    madx.parse(madx_file)
+
+    ref_particle_dict = beam(
+        madx.getParticle(),  # if particle species is known, mass, charge, and potentially energy are set to default
+        # TODO MADX parser needs to extract charge if it's given,
+        # TODO MADX parser needs to extract mass if it's given,
+        energy=float(madx.getEtot()),  # MADX default energy is in GeV
+    )
+
+    ref.set_charge_qe(ref_particle_dict["charge"])
+    ref.set_mass_MeV(ref_particle_dict["mass"])
+    ref.set_energy_MeV(ref_particle_dict["energy"])