From 9ee30500fe4ba22452d5ce467bba2b345e699e48 Mon Sep 17 00:00:00 2001 From: Marco Garten Date: Tue, 23 Aug 2022 12:04:22 -0700 Subject: [PATCH 01/19] Add MADX Input Parser Using Andreas Adelmann's pymadxparser --- examples/CMakeLists.txt | 10 + examples/chicane/input_chicane.madx | 15 + examples/chicane/run_chicane_madx.py | 87 ++++++ examples/fodo/input_fodo.madx | 10 + examples/fodo/run_fodo_madx.py | 88 ++++++ src/python/impactx/MADXParser.py | 377 ++++++++++++++++++++++++++ src/python/impactx/__init__.py | 2 + src/python/impactx/madx_to_impactx.py | 48 ++++ 8 files changed, 637 insertions(+) create mode 100644 examples/chicane/input_chicane.madx create mode 100755 examples/chicane/run_chicane_madx.py create mode 100644 examples/fodo/input_fodo.madx create mode 100755 examples/fodo/run_fodo_madx.py create mode 100644 src/python/impactx/MADXParser.py create mode 100644 src/python/impactx/madx_to_impactx.py diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 39bc90a7c..7eec85435 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -104,6 +104,16 @@ add_impactx_test(FODO.py examples/fodo/plot_fodo.py ) +# Python MADX: FODO Cell ###################################################### +# +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 --save-png + ) + # Python: MPI-parallel FODO Cell ############################################## # add_impactx_test(FODO.py.MPI diff --git a/examples/chicane/input_chicane.madx b/examples/chicane/input_chicane.madx new file mode 100644 index 000000000..342f1cbbb --- /dev/null +++ b/examples/chicane/input_chicane.madx @@ -0,0 +1,15 @@ +beam, particle=electron, energy=40.e-3; + +DLEFT: drift, L=0.7578; +DMIDDLE: drift, L=0.9575; +DRIGHT: drift, L=0.7578; +! TODO make this work with inline calculations +! theta=pi/10.; +! TODO put theta again everywhere after `angle=` with its right sign +D114: sbend, L=0.264687, angle= 0.31415926535, e1= 0.000, e2= 0.31415926535, k1= 0.00; +D115: sbend, L=0.264687, angle=-0.31415926535, e1=-0.31415926535, e2= 0.000, k1= 0.00; +D116: sbend, L=0.264687, angle=-0.31415926535, e1= 0.000, e2=-0.31415926535, k1= 0.00; +D117: sbend, L=0.264687, angle= 0.31415926535, e1= 0.31415926535, e2= 0.000, k1= 0.00; + +BC1: Line=(D114,DLEFT,D115,DMIDDLE,D116,DRIGHT,D117); +USE, SEQUENCE = BC1; diff --git a/examples/chicane/run_chicane_madx.py b/examples/chicane/run_chicane_madx.py new file mode 100755 index 000000000..b571e8ef3 --- /dev/null +++ b/examples/chicane/run_chicane_madx.py @@ -0,0 +1,87 @@ +#!/usr/bin/env python3 +# +# Copyright 2022 ImpactX contributors +# Authors: Marco Garten, Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +import sys +import os + +import amrex +from impactx import ImpactX, RefPart, distribution, elements, MADXParser, madx2impactx + +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() + +# @TODO make this better ... shouldn't need that. But somehow the working directory is +dir_path = os.path.dirname(os.path.realpath(__file__)) + +try: + madx = MADXParser() + + madx.parse('{}/input_chicane.madx'.format(dir_path)) + +except Exception as e: + print(f"Unexpected {e = }, {type(e) = }") + raise + +# load a 40 MeV 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 +# @TODO get from dictionary b/c MADX knows P +mass_MeV = 0.510998950 +qm_qeeV = -1.0e-6/mass_MeV # charge/mass +npart = 10000 # number of macro particles + +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(qm_qeeV, charge_C, distr, npart) + +# design the accelerator lattice +ns = 25 # number of slices per ds in the element + + +# print summary +print(madx) + +# MADX default energy is in GeV +energy_MeV = float(madx.getEtot()) * 1e3 +beamline = madx.getBeamline() + +# set the energy in the reference particle +sim.particle_container().ref_particle() \ + .set_energy_MeV(energy_MeV, mass_MeV) + +chicane = madx2impactx(beamline) +# assign a fodo segment +sim.lattice.extend(chicane) + +# run simulation +sim.evolve() + +# clean shutdown +del sim +amrex.finalize() + + + diff --git a/examples/fodo/input_fodo.madx b/examples/fodo/input_fodo.madx new file mode 100644 index 000000000..1fbee02a6 --- /dev/null +++ b/examples/fodo/input_fodo.madx @@ -0,0 +1,10 @@ +! TODO implement handling of title +! TITLE,'FODO.MADX'; +BEAM, PARTICLE=ELECTRON,ENERGY=3.0; + +D: DRIFT,L=0.25; +QF: QUADRUPOLE,L=0.5,K1=1.0; +QD: QUADRUPOLE,L=0.5,K1=-1.0; + +FODO: LINE=(D,QF,D,QD,D); +USE, SEQUENCE = FODO; diff --git a/examples/fodo/run_fodo_madx.py b/examples/fodo/run_fodo_madx.py new file mode 100755 index 000000000..ee0d946e1 --- /dev/null +++ b/examples/fodo/run_fodo_madx.py @@ -0,0 +1,88 @@ +#!/usr/bin/env python3 +# +# Copyright 2022 ImpactX contributors +# Authors: Marco Garten, Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +import sys +import os + +import amrex +from impactx import ImpactX, RefPart, distribution, elements, MADXParser, madx2impactx + +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() + +# @TODO make this better ... shouldn't need that +dir_path = os.path.dirname(os.path.realpath(__file__)) + +try: + madx = MADXParser() + + madx.parse('{}/input_fodo.madx'.format(dir_path)) + +except Exception as e: + print(f"Unexpected {e = }, {type(e) = }") + raise + +# 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 +# @TODO get from dictionary b/c MADX knows P +mass_MeV = 0.510998950 +qm_qeeV = -1.0e-6/mass_MeV # charge/mass +npart = 10000 # number of macro particles + +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(qm_qeeV, charge_C, distr, npart) + +# design the accelerator lattice +ns = 25 # number of slices per ds in the element + + +# print summary +print(madx) + +# MADX default energy is in GeV +energy_MeV = float(madx.getEtot()) * 1e3 +beamline = madx.getBeamline() + +# set the energy in the reference particle +sim.particle_container().ref_particle() \ + .set_energy_MeV(energy_MeV, mass_MeV) + +fodo = madx2impactx(beamline) +# assign a fodo segment +sim.lattice.extend(fodo) + +# run simulation +sim.evolve() + +# clean shutdown +del sim +amrex.finalize() + + + + diff --git a/src/python/impactx/MADXParser.py b/src/python/impactx/MADXParser.py new file mode 100644 index 000000000..8c779ffa7 --- /dev/null +++ b/src/python/impactx/MADXParser.py @@ -0,0 +1,377 @@ +# Author: Matthias Frey, Andreas Adelmann +# Date: 13. Oct. 2017 +# +# This file parses MAD-X file into a beamline +# composed of pyAcceLEGOrator elements. +# +# 10.8.2022 Adapted for standalone use +# +import re + +class MADXParserError(Exception): + pass + +class MADXInputError(MADXParserError): + def __init__(self, args, with_traceback): + self.args = args + self.with_traceback = with_traceback + +class MADXParser: + """ + Simple MADX parser. + It expects a single line per element. + """ + + def __init__(self): + + self.__drift = { + 'name': '', + 'l': 0.0, + 'type': 'drift' + } + + self.__drift_pattern = '(.*):drift,(.*)=(.*);' + + self.__quadrupole = { + 'name': '', + 'l': 0.0, + 'k1': 0.0, + 'type': 'quad' + } + + # don't count name and type --> len - 2 + self.__nQuad = 2 * (len(self.__quadrupole) - 2) + + self.quad_pattern = '(.*):quadrupole,(.*)=(.*),(.*)=(.*);' + + self.__sbend = { + 'name': '', + 'l': 0.0, + 'angle': 0.0, + 'k1': 0.0, + 'e1': 0.0, + 'e2': 0.0, + 'type': 'sbend' + } + + self.__sbend_pattern = '(.*):sbend,(.*)=(.*),(.*)=(.*),(.*)=(.*),(.*)=(.*),(.*)=(.*);' + + # don't count name and type --> len - 2 + # TODO add rbend + self.__nDipole = 2 * (len(self.__sbend) - 2) + + self.beam = { + 'energy': 0.0, + # TODO extend by 'PC' + 'particle': '' + } + + self.__nBeam = 2 * len(self.beam) + + self.beam_pattern = 'beam,(.*)=(.*),(.*)=(.*);' + + self.__line = { + 'name': '', + 'elem': [], + } + + self.__line_pattern = '(.*):line=\(+(.*)\);' + + self.sequence = { + 'name': '' + } + + self.seq_pattern = 'use,sequence=(.*);' + + self.__elements = [] + self.__lines = [] + + self.__lattice = [] + + + # nonblank_lines inspired by + # https://stackoverflow.com/questions/4842057/easiest-way-to-ignore-blank-lines-when-reading-a-file-in-python + def nonblank_lines_to_lowercase(self, f): + for l in f: + line = l.rstrip().casefold() + if line: + yield line + + def parse(self, fn): + """ + fn (str) filename + + """ + + nLine = 0 + + with open(fn, 'r') as f: + for line in self.nonblank_lines_to_lowercase(f): + nLine += 1 + line = self._noWhitespace(line) + + if line[0] == '!': + # this is a comment + pass + elif 'drift' in line: + + obj = re.match(self.__drift_pattern, line) + + # first tag is name + self.__drift['name'] = obj.group(1) + + if obj.group(2) in self.__drift: + self.__drift[obj.group(2)] = float(obj.group(3)) + else: + raise MADXInputError('Drift', + 'Line ' + str(nLine) + ': Parameter ' + "'" + obj.group(2) + "'" + + ' does not exist for drift.') + + self.__elements.append( self.__drift.copy() ) + + elif 'quadrupole' in line: + + obj = re.match(self.quad_pattern, line) + + # first tag is name + self.__quadrupole['name'] = obj.group(1) + + for i in range(2, self.__nQuad+2, 2): + if obj.group(i) in self.__quadrupole: + self.__quadrupole[obj.group(i)] = float(obj.group(i+1)) + else: + raise MADXInputError('Quadrupole', + 'Line ' + str(nLine) + ': Parameter ' + "'" + obj.group(i) + "'" + + ' does not exist for quadrupole.') + + self.__elements.append( self.__quadrupole.copy() ) + + elif 'sbend' in line: + + obj = re.match(self.__sbend_pattern, line) + if obj: # TODO remove DEBUG line + print("OBJ: ", obj) # TODO remove DEBUG line + else: + print("Oops") + + # first tag is name + self.__sbend['name'] = obj.group(1) + + for i in range(2, self.__nDipole + 2, 2): + + if obj.group(i) in self.__sbend: + self.__sbend[obj.group(i)] = float(obj.group(i + 1)) + else: + raise MADXInputError('Dipole', + 'Line ' + str(nLine) + ': Parameter ' + + "'" + obj.group(i) + "'" + + ' does not exist for dipole.') + + self.__elements.append(self.__sbend.copy()) + elif 'marker' in line: + pass + + elif 'beam' in line: + + obj = re.match(self.beam_pattern, line) + + for i in range(1, self.__nBeam, 2): + if obj.group(i) in self.beam: + if obj.group(i+1).isdigit(): + self.beam[obj.group(i)] = float(obj.group(i+1)) + else: + self.beam[obj.group(i)] = obj.group(i+1) + + else: + raise MADXInputError('Beam', + 'Line ' + str(nLine) + ': Parameter ' + + "'" + obj.group(i) + "'" + + ' does not exist for beam.') + + elif 'line' in line: + + obj = re.match(self.__line_pattern, line) + + self.__line['name'] = obj.group(1) + + lines = obj.group(2).split(',') + newlines = [] + + # check multiplication and insert that many lines + for l in lines: + if '*' in l: + tmp = l.split('*') + + n = 0 + ll = '' + + if tmp[0].isdigit(): + n = int(tmp[0]) + ll = tmp[1] + else: + n = int(tmp[1]) + ll = tmp[0] + + newlines += [ll] * n + + else: + newlines.append(l) + + self.__line['elem'] = newlines + + self.__lines.append( self.__line.copy() ) + + elif 'use' in line and 'sequence' in line: + + obj = re.match(self.seq_pattern, line) + + self.sequence['name'] = obj.group(1) + else: + raise MADXInputError('', 'Error at line ' + str(nLine)) + + + # 14. Oct. 2017, + # https://stackoverflow.com/questions/7900882/extract-item-from-list-of-dictionaries + for l in self.__lines: # TODO remove this DEBUG line + print(l) + print(self.sequence) + start = [l for l in self.__lines if l['name'] == self.sequence['name']][0] + + self._flatten(start) + + # we need to add start line + self.__lattice = [start] + self.__lattice + + self.__lattice = self._combine(self.__lattice) + + + def _flatten(self, line): + """ + Find sublines. + + """ + + name = line['name'] + + for l in self.__lines: + if name in l['name']: + for ll in l['elem']: + for lll in self.__lines: + if lll['name'] == ll: + self.__lattice.append(lll) + self._flatten(lll) + + + def _combine(self, lattice): + """ + Combine to one list of all basic + elements. + + return a list of of element dictionaries + """ + + + l1 = self.__lattice[0] + + for i in range(1, len(self.__lattice)): + l2 = self.__lattice[i] + + for e in l1['elem']: + if l2['name'] == e: + idx = l1['elem'].index(e) + l1['elem'].remove(e) + l1['elem'] = l1['elem'][0:idx] + l2['elem'] + l1['elem'][idx:] + + return l1 + + + def _noWhitespace(self, string): + """ + Remove white space from a string. + + 14. Oct. 2017, + https://stackoverflow.com/questions/3739909/how-to-strip-all-whitespace-from-string + + """ + return string.replace(' ', '') + + + def __str__(self): + + if self.__lattice: + length = 0.0 + + # drift, dipole, quadrupole + nTypes = [0, 0, 0] + + for elem in self.__lattice['elem']: + for e in self.__elements: + if elem == e['name']: + length += e['l'] + + if e['type'] == 'drift': + nTypes[0] += 1 + elif e['type'] == 'sbend': + nTypes[1] += 1 + elif e['type'] == 'quad': + nTypes[2] += 1 + break + + sign = '*' * 70 + info = sign + '\n'\ + + 'MADX-Parser information:\n' \ + + ' length:\t' + str(length) + ' [m]\n'\ + + ' #elements:\t' + str(len(self.__lattice['elem'])) + '\n'\ + + ' * #drift:\t' + str(nTypes[0]) + '\n'\ + + ' * #dipole:\t' + str(nTypes[1]) + '\n'\ + + ' * #quadrupole:\t' + str(nTypes[2]) + '\n'\ + + ' beam:\t\n'\ + + ' * particle:\t' + self.beam['particle'] + '\n'\ + + ' * total energy:\t' + str(self.beam['energy']) + ' [GeV]\n'\ + + sign + + return info + else: + return 'No information available.' + + def getBeamline(self): + if self.__lattice: + + beamline = [] + + for elem in self.__lattice['elem']: + for e in self.__elements: + if elem == e['name']: + if e['type'] == 'drift': + print('Drift L= '+str(e['l'])) + beamline.append(e) + elif e['type'] == 'sbend': + print("Sbend L= ",e['l'],' angle = ',e['angle']) + beamline.append(e) + elif e['type'] == 'quad': + print("Quadrupole L= ",e['l'],' k1 = ',e['k1']) + beamline.append(e) + else: + print ( 'Skipping element type ' + "'" + + e['type'] + "'" ) + return beamline + + else: + return [] + + + def getParticle(self): + particle = None + if self.beam['particle'] == 'electron': + particle = 'Electron' + elif self.beam['particle'] == 'proton': + particle = 'Proton' + else: + raise MADXInputError('', 'No particle type ' + "'" + + self.beam['particle'] + "' available.") + + return particle + + + def getEtot(self): + return self.beam['energy'] diff --git a/src/python/impactx/__init__.py b/src/python/impactx/__init__.py index a9ac835fb..04eb5e19e 100644 --- a/src/python/impactx/__init__.py +++ b/src/python/impactx/__init__.py @@ -1,5 +1,7 @@ from . import impactx_pybind from .impactx_pybind import * # noqa +from .madx_to_impactx import madx2impactx # noqa +from .MADXParser import * # noqa __version__ = impactx_pybind.__version__ __doc__ = impactx_pybind.__doc__ diff --git a/src/python/impactx/madx_to_impactx.py b/src/python/impactx/madx_to_impactx.py new file mode 100644 index 000000000..86ff3b02f --- /dev/null +++ b/src/python/impactx/madx_to_impactx.py @@ -0,0 +1,48 @@ +from impactx import elements + +def madx2impactx(parsed_beamline,nslice=25): + """ + Function that converts a list of elements in the MADXParser former into ImpactX format + :param parsed_beamline: list of dictionaries + :return: list of translated dictionaries + """ + + madx_to_impactx_dict = { + "MARKER": "None", + "DRIFT": "Drift", + "SBEND": "Sbend", # Sector Bending Magnet + "QUAD": "Quad", # Quadrupole + "DIPEDGE": "DipEdge", + "MULTIPOLE": "Multipole", + "NLLENS": "NonlinearLens", + # TODO Figure out how to identify these + "ShortRF": "ShortRF", + "ConstF": "ConstF", + } + + impactx_beamline = [] + + for d in parsed_beamline: + + if d['type'] in [k.casefold() for k in list(madx_to_impactx_dict.keys())]: + if d['name'] == "drift": + impactx_beamline.append( + elements.Drift(ds=d['l'], nslice=nslice) + ) + elif d['name'] == "quadrupole": + impactx_beamline.append( + elements.Quad(ds=d['l'], k=d['k1'], nslice=nslice) + ) + else: + raise NotImplementedError( + "The beamline element named ", + d['name'], + "of type ", + d['type'], + "is not implemented in impactx.elements.", + "Available elements are:", + list(madx_to_impactx_dict.keys()) + ) + return impactx_beamline + + From 3eedb86296ce5909ec3e4ccf4d49a61eda1ea799 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 25 Aug 2022 01:21:28 +0000 Subject: [PATCH 02/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- examples/chicane/input_chicane.madx | 2 +- examples/chicane/run_chicane_madx.py | 7 ++----- examples/fodo/run_fodo_madx.py | 8 ++------ src/python/impactx/MADXParser.py | 5 +++-- src/python/impactx/__init__.py | 2 +- src/python/impactx/madx_to_impactx.py | 3 +-- 6 files changed, 10 insertions(+), 17 deletions(-) diff --git a/examples/chicane/input_chicane.madx b/examples/chicane/input_chicane.madx index 342f1cbbb..19b27766e 100644 --- a/examples/chicane/input_chicane.madx +++ b/examples/chicane/input_chicane.madx @@ -10,6 +10,6 @@ D114: sbend, L=0.264687, angle= 0.31415926535, e1= 0.000, e2= 0.31415926535, k1= D115: sbend, L=0.264687, angle=-0.31415926535, e1=-0.31415926535, e2= 0.000, k1= 0.00; D116: sbend, L=0.264687, angle=-0.31415926535, e1= 0.000, e2=-0.31415926535, k1= 0.00; D117: sbend, L=0.264687, angle= 0.31415926535, e1= 0.31415926535, e2= 0.000, k1= 0.00; - + BC1: Line=(D114,DLEFT,D115,DMIDDLE,D116,DRIGHT,D117); USE, SEQUENCE = BC1; diff --git a/examples/chicane/run_chicane_madx.py b/examples/chicane/run_chicane_madx.py index b571e8ef3..1d8f60506 100755 --- a/examples/chicane/run_chicane_madx.py +++ b/examples/chicane/run_chicane_madx.py @@ -6,11 +6,11 @@ # # -*- coding: utf-8 -*- -import sys import os import amrex -from impactx import ImpactX, RefPart, distribution, elements, MADXParser, madx2impactx +from impactx import (ImpactX, MADXParser, RefPart, distribution, elements, + madx2impactx) sim = ImpactX() @@ -82,6 +82,3 @@ # clean shutdown del sim amrex.finalize() - - - diff --git a/examples/fodo/run_fodo_madx.py b/examples/fodo/run_fodo_madx.py index ee0d946e1..d26a87af3 100755 --- a/examples/fodo/run_fodo_madx.py +++ b/examples/fodo/run_fodo_madx.py @@ -6,11 +6,11 @@ # # -*- coding: utf-8 -*- -import sys import os import amrex -from impactx import ImpactX, RefPart, distribution, elements, MADXParser, madx2impactx +from impactx import (ImpactX, MADXParser, RefPart, distribution, elements, + madx2impactx) sim = ImpactX() @@ -82,7 +82,3 @@ # clean shutdown del sim amrex.finalize() - - - - diff --git a/src/python/impactx/MADXParser.py b/src/python/impactx/MADXParser.py index 8c779ffa7..2e106a2c8 100644 --- a/src/python/impactx/MADXParser.py +++ b/src/python/impactx/MADXParser.py @@ -1,5 +1,5 @@ # Author: Matthias Frey, Andreas Adelmann -# Date: 13. Oct. 2017 +# Date: 13. Oct. 2017 # # This file parses MAD-X file into a beamline # composed of pyAcceLEGOrator elements. @@ -8,6 +8,7 @@ # import re + class MADXParserError(Exception): pass @@ -338,7 +339,7 @@ def getBeamline(self): if self.__lattice: beamline = [] - + for elem in self.__lattice['elem']: for e in self.__elements: if elem == e['name']: diff --git a/src/python/impactx/__init__.py b/src/python/impactx/__init__.py index 04eb5e19e..481d188bf 100644 --- a/src/python/impactx/__init__.py +++ b/src/python/impactx/__init__.py @@ -1,7 +1,7 @@ from . import impactx_pybind +from .MADXParser import * # noqa from .impactx_pybind import * # noqa from .madx_to_impactx import madx2impactx # noqa -from .MADXParser import * # noqa __version__ = impactx_pybind.__version__ __doc__ = impactx_pybind.__doc__ diff --git a/src/python/impactx/madx_to_impactx.py b/src/python/impactx/madx_to_impactx.py index 86ff3b02f..74337f676 100644 --- a/src/python/impactx/madx_to_impactx.py +++ b/src/python/impactx/madx_to_impactx.py @@ -1,5 +1,6 @@ from impactx import elements + def madx2impactx(parsed_beamline,nslice=25): """ Function that converts a list of elements in the MADXParser former into ImpactX format @@ -44,5 +45,3 @@ def madx2impactx(parsed_beamline,nslice=25): list(madx_to_impactx_dict.keys()) ) return impactx_beamline - - From c55422230a0cf66efb73a05f7b0dabfef5147be0 Mon Sep 17 00:00:00 2001 From: Marco Garten Date: Thu, 25 Aug 2022 17:38:26 -0700 Subject: [PATCH 03/19] Copy MADX input via CMakeLists.txt Co-authored-by: Axel Huebl --- examples/CMakeLists.txt | 3 +++ 1 file changed, 3 insertions(+) diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 7eec85435..91a2d2636 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -106,6 +106,9 @@ add_impactx_test(FODO.py # Python MADX: FODO Cell ###################################################### # +# copy MAD-X lattice file +file(COPY ${ImpactX_SOURCE_DIR}/examples/fodo/input_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 From 74c730b590e4e132580cd11a48f49ec7619124c1 Mon Sep 17 00:00:00 2001 From: Marco Garten Date: Fri, 26 Aug 2022 16:19:22 -0700 Subject: [PATCH 04/19] Add beam function to MADX translation --- examples/CMakeLists.txt | 13 ++++++ examples/chicane/run_chicane_madx.py | 33 ++++++++------ examples/fodo/input_fodo.madx | 2 +- examples/fodo/run_fodo_madx.py | 31 +++++++------ src/python/impactx/MADXParser.py | 18 +++++--- src/python/impactx/__init__.py | 2 +- src/python/impactx/madx_to_impactx.py | 64 ++++++++++++++++++++++++++- 7 files changed, 124 insertions(+), 39 deletions(-) diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 91a2d2636..1948e7239 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -147,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/input_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 --save-png + ) + # Constant Focusing Channel ################################################### # add_impactx_test(cfchannel diff --git a/examples/chicane/run_chicane_madx.py b/examples/chicane/run_chicane_madx.py index 1d8f60506..09c899e3b 100755 --- a/examples/chicane/run_chicane_madx.py +++ b/examples/chicane/run_chicane_madx.py @@ -10,7 +10,7 @@ import amrex from impactx import (ImpactX, MADXParser, RefPart, distribution, elements, - madx2impactx) + madx2impactx_lattice, madx2impactx_beam) sim = ImpactX() @@ -35,14 +35,27 @@ print(f"Unexpected {e = }, {type(e) = }") raise + +# print summary +print(madx) + +beamline = madx.getBeamline() + +ref_particle_dict = madx2impactx_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 # 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 -# @TODO get from dictionary b/c MADX knows P -mass_MeV = 0.510998950 -qm_qeeV = -1.0e-6/mass_MeV # charge/mass + +qm_qeeV = ref_particle_dict['mass']/ref_particle_dict['charge'] # charge/mass npart = 10000 # number of macro particles distr = distribution.Waterbag( @@ -60,19 +73,11 @@ # design the accelerator lattice ns = 25 # number of slices per ds in the element - -# print summary -print(madx) - -# MADX default energy is in GeV -energy_MeV = float(madx.getEtot()) * 1e3 -beamline = madx.getBeamline() - # set the energy in the reference particle sim.particle_container().ref_particle() \ - .set_energy_MeV(energy_MeV, mass_MeV) + .set_energy_MeV(ref_particle_dict['energy'], ref_particle_dict['mass']) -chicane = madx2impactx(beamline) +chicane = madx2impactx_lattice(beamline) # assign a fodo segment sim.lattice.extend(chicane) diff --git a/examples/fodo/input_fodo.madx b/examples/fodo/input_fodo.madx index 1fbee02a6..bdb706b98 100644 --- a/examples/fodo/input_fodo.madx +++ b/examples/fodo/input_fodo.madx @@ -1,6 +1,6 @@ ! TODO implement handling of title ! TITLE,'FODO.MADX'; -BEAM, PARTICLE=ELECTRON,ENERGY=3.0; +BEAM, PARTICLE=ELECTRON,ENERGY=2.0; D: DRIFT,L=0.25; QF: QUADRUPOLE,L=0.5,K1=1.0; diff --git a/examples/fodo/run_fodo_madx.py b/examples/fodo/run_fodo_madx.py index d26a87af3..2a141eb40 100755 --- a/examples/fodo/run_fodo_madx.py +++ b/examples/fodo/run_fodo_madx.py @@ -10,7 +10,7 @@ import amrex from impactx import (ImpactX, MADXParser, RefPart, distribution, elements, - madx2impactx) + madx2impactx_lattice, madx2impactx_beam) sim = ImpactX() @@ -35,14 +35,25 @@ print(f"Unexpected {e = }, {type(e) = }") raise +# print summary +print(madx) + +beamline = madx.getBeamline() + +ref_particle_dict = madx2impactx_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 -# @TODO get from dictionary b/c MADX knows P -mass_MeV = 0.510998950 -qm_qeeV = -1.0e-6/mass_MeV # charge/mass + +qm_qeeV = ref_particle_dict['mass']/ref_particle_dict['charge'] # charge/mass npart = 10000 # number of macro particles distr = distribution.Waterbag( @@ -60,19 +71,11 @@ # design the accelerator lattice ns = 25 # number of slices per ds in the element - -# print summary -print(madx) - -# MADX default energy is in GeV -energy_MeV = float(madx.getEtot()) * 1e3 -beamline = madx.getBeamline() - # set the energy in the reference particle sim.particle_container().ref_particle() \ - .set_energy_MeV(energy_MeV, mass_MeV) + .set_energy_MeV(ref_particle_dict['energy'], ref_particle_dict['mass']) -fodo = madx2impactx(beamline) +fodo = madx2impactx_lattice(beamline) # assign a fodo segment sim.lattice.extend(fodo) diff --git a/src/python/impactx/MADXParser.py b/src/python/impactx/MADXParser.py index 2e106a2c8..d1930532c 100644 --- a/src/python/impactx/MADXParser.py +++ b/src/python/impactx/MADXParser.py @@ -7,6 +7,7 @@ # 10.8.2022 Adapted for standalone use # import re +import warnings class MADXParserError(Exception): @@ -17,6 +18,9 @@ def __init__(self, args, with_traceback): self.args = args self.with_traceback = with_traceback +class MADXInputWarning(UserWarning): + pass + class MADXParser: """ Simple MADX parser. @@ -362,14 +366,14 @@ def getBeamline(self): def getParticle(self): - particle = None - if self.beam['particle'] == 'electron': - particle = 'Electron' - elif self.beam['particle'] == 'proton': - particle = 'Proton' + particle = self.beam['particle'] + known_particles = ['positron', 'electron', 'proton', 'antiproton', 'posmuon', 'negmuon', 'ion'] + + if particle not in known_particles: + warning_message = 'No particle type ' + "'" + self.beam['particle'] + "' available." + warnings.warn(warning_message, MADXInputWarning) else: - raise MADXInputError('', 'No particle type ' + "'" + - self.beam['particle'] + "' available.") + pass return particle diff --git a/src/python/impactx/__init__.py b/src/python/impactx/__init__.py index 481d188bf..a2af79478 100644 --- a/src/python/impactx/__init__.py +++ b/src/python/impactx/__init__.py @@ -1,7 +1,7 @@ from . import impactx_pybind from .MADXParser import * # noqa from .impactx_pybind import * # noqa -from .madx_to_impactx import madx2impactx # noqa +from .madx_to_impactx import madx2impactx_lattice, madx2impactx_beam # noqa __version__ = impactx_pybind.__version__ __doc__ = impactx_pybind.__doc__ diff --git a/src/python/impactx/madx_to_impactx.py b/src/python/impactx/madx_to_impactx.py index 74337f676..821cebe9f 100644 --- a/src/python/impactx/madx_to_impactx.py +++ b/src/python/impactx/madx_to_impactx.py @@ -1,9 +1,22 @@ +#!/usr/bin/env python3 +# +# Copyright 2022 ImpactX contributors +# Authors: Marco Garten +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +import warnings + +import scipy.constants as sc + from impactx import elements -def madx2impactx(parsed_beamline,nslice=25): + +def madx2impactx_lattice(parsed_beamline,nslice=25): """ - Function that converts a list of elements in the MADXParser former into ImpactX format + Function that converts a list of elements in the MADXParser format into ImpactX format :param parsed_beamline: list of dictionaries :return: list of translated dictionaries """ @@ -45,3 +58,50 @@ def madx2impactx(parsed_beamline,nslice=25): list(madx_to_impactx_dict.keys()) ) return impactx_beamline + +def madx2impactx_beam(particle, charge=None, mass=None, energy=None): + """ + Function that converts a list of beam parameter dictionaries in the MADXParser format into ImpactX format + + Rules following https://mad.web.cern.ch/mad/releases/5.02.08/madxuguide.pdf pages 55f. + + :param str particle: reference particle name + :param float charge: particle charge (proton charge units) + :param float mass: particle mass (electron masses) + :param float energy: particle energy (GeV) + - MAD-X default: 1 GeV + :return dict: dictionary containing particle and beam attributes in ImpactX units + """ + + GeV2MeV = 1e-3 + kg2MeV = sc.c ** 2 / sc.electron_volt * 1e-6 + muon_mass = sc.physical_constants['electron-muon mass ratio'][0] / sc.m_e + if energy is None: + energy_MeV = 1e3 # MAD-X default is 1 GeV particle energy + else: + energy_MeV = energy * GeV2MeV + + impactx_beam = { + 'positron': {'mass': sc.m_e * kg2MeV, 'charge': 1}, + 'electron': {'mass': sc.m_e * kg2MeV, 'charge': -1}, + 'proton': {'mass': sc.m_p * kg2MeV, 'charge': 1}, + 'antiproton': {'mass': sc.m_p * kg2MeV, 'charge': -1}, + 'posmuon': {'mass': muon_mass * kg2MeV, 'charge': 1}, # positively charged muon (anti-muon) + 'negmuon': {'mass': muon_mass * kg2MeV, 'charge': -1}, # negatively charged muon + 'ion': {'mass': sc.m_u * kg2MeV, 'charge': 1}, + 'generic': {'mass': mass, 'charge': charge} + } + + + if particle not in impactx_beam.keys(): + warnings.warn('Particle species name "' + particle + '" not in ' + impactx_beam.keys(), UserWarning) + print('Choosing generic particle species, using provided `charge`, `mass` and `energy`.') + _particle = 'generic' + else: + _particle = particle + print('Choosing provided particle species "', particle, '", ignoring potentially provided `charge` and `mass` and setting defaults.') + + reference_particle = impactx_beam[_particle] + reference_particle['energy'] = energy_MeV + + return reference_particle From 649f68786253ee2db8b78a09ebf9d1c595c59262 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 26 Aug 2022 23:20:05 +0000 Subject: [PATCH 05/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- examples/chicane/run_chicane_madx.py | 2 +- examples/fodo/run_fodo_madx.py | 2 +- src/python/impactx/__init__.py | 2 +- src/python/impactx/madx_to_impactx.py | 1 - 4 files changed, 3 insertions(+), 4 deletions(-) diff --git a/examples/chicane/run_chicane_madx.py b/examples/chicane/run_chicane_madx.py index 09c899e3b..9398bad51 100755 --- a/examples/chicane/run_chicane_madx.py +++ b/examples/chicane/run_chicane_madx.py @@ -10,7 +10,7 @@ import amrex from impactx import (ImpactX, MADXParser, RefPart, distribution, elements, - madx2impactx_lattice, madx2impactx_beam) + madx2impactx_beam, madx2impactx_lattice) sim = ImpactX() diff --git a/examples/fodo/run_fodo_madx.py b/examples/fodo/run_fodo_madx.py index 2a141eb40..2f35f4c2b 100755 --- a/examples/fodo/run_fodo_madx.py +++ b/examples/fodo/run_fodo_madx.py @@ -10,7 +10,7 @@ import amrex from impactx import (ImpactX, MADXParser, RefPart, distribution, elements, - madx2impactx_lattice, madx2impactx_beam) + madx2impactx_beam, madx2impactx_lattice) sim = ImpactX() diff --git a/src/python/impactx/__init__.py b/src/python/impactx/__init__.py index a2af79478..2e27cfaa8 100644 --- a/src/python/impactx/__init__.py +++ b/src/python/impactx/__init__.py @@ -1,7 +1,7 @@ from . import impactx_pybind from .MADXParser import * # noqa from .impactx_pybind import * # noqa -from .madx_to_impactx import madx2impactx_lattice, madx2impactx_beam # noqa +from .madx_to_impactx import madx2impactx_beam, madx2impactx_lattice # noqa __version__ = impactx_pybind.__version__ __doc__ = impactx_pybind.__doc__ diff --git a/src/python/impactx/madx_to_impactx.py b/src/python/impactx/madx_to_impactx.py index 821cebe9f..91b9e23d8 100644 --- a/src/python/impactx/madx_to_impactx.py +++ b/src/python/impactx/madx_to_impactx.py @@ -13,7 +13,6 @@ from impactx import elements - def madx2impactx_lattice(parsed_beamline,nslice=25): """ Function that converts a list of elements in the MADXParser format into ImpactX format From 6d8722dd2973cd1e3c862b6ac4c01dea4e310624 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 30 Aug 2022 09:47:08 -0700 Subject: [PATCH 06/19] Clean up: Test & Input File Names --- examples/CMakeLists.txt | 32 +++++++++---------- .../{input_chicane.madx => chicane.madx} | 0 examples/chicane/run_chicane_madx.py | 2 +- examples/fodo/{input_fodo.madx => fodo.madx} | 0 examples/fodo/run_fodo_madx.py | 2 +- 5 files changed, 18 insertions(+), 18 deletions(-) rename examples/chicane/{input_chicane.madx => chicane.madx} (100%) rename examples/fodo/{input_fodo.madx => fodo.madx} (100%) diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 1948e7239..1cade33ba 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -107,15 +107,15 @@ add_impactx_test(FODO.py # Python MADX: FODO Cell ###################################################### # # copy MAD-X lattice file -file(COPY ${ImpactX_SOURCE_DIR}/examples/fodo/input_fodo.madx DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/FODO_MADX.py) +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 --save-png - ) +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 --save-png +) # Python: MPI-parallel FODO Cell ############################################## # @@ -150,15 +150,15 @@ add_impactx_test(chicane.py # Python MADX: Chicane ###################################################### # # copy MAD-X lattice file -file(COPY ${ImpactX_SOURCE_DIR}/examples/chicane/input_chicane.madx DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/chicane_MADX.py) +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 --save-png - ) +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 --save-png +) # Constant Focusing Channel ################################################### # diff --git a/examples/chicane/input_chicane.madx b/examples/chicane/chicane.madx similarity index 100% rename from examples/chicane/input_chicane.madx rename to examples/chicane/chicane.madx diff --git a/examples/chicane/run_chicane_madx.py b/examples/chicane/run_chicane_madx.py index 9398bad51..207ab813d 100755 --- a/examples/chicane/run_chicane_madx.py +++ b/examples/chicane/run_chicane_madx.py @@ -29,7 +29,7 @@ try: madx = MADXParser() - madx.parse('{}/input_chicane.madx'.format(dir_path)) + madx.parse(f"{dir_path}/chicane.madx") except Exception as e: print(f"Unexpected {e = }, {type(e) = }") diff --git a/examples/fodo/input_fodo.madx b/examples/fodo/fodo.madx similarity index 100% rename from examples/fodo/input_fodo.madx rename to examples/fodo/fodo.madx diff --git a/examples/fodo/run_fodo_madx.py b/examples/fodo/run_fodo_madx.py index 2f35f4c2b..55a2d3ff7 100755 --- a/examples/fodo/run_fodo_madx.py +++ b/examples/fodo/run_fodo_madx.py @@ -29,7 +29,7 @@ try: madx = MADXParser() - madx.parse('{}/input_fodo.madx'.format(dir_path)) + madx.parse(f"{dir_path}/fodo.madx") except Exception as e: print(f"Unexpected {e = }, {type(e) = }") From 95af5e332c8f51a338f58590d7b046d154407887 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 30 Aug 2022 16:57:13 +0000 Subject: [PATCH 07/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- examples/chicane/run_chicane_madx.py | 41 ++-- examples/fodo/run_fodo_madx.py | 41 ++-- src/python/impactx/MADXParser.py | 291 +++++++++++++++----------- src/python/impactx/madx_to_impactx.py | 67 +++--- 4 files changed, 254 insertions(+), 186 deletions(-) diff --git a/examples/chicane/run_chicane_madx.py b/examples/chicane/run_chicane_madx.py index 207ab813d..a6ac56afc 100755 --- a/examples/chicane/run_chicane_madx.py +++ b/examples/chicane/run_chicane_madx.py @@ -9,15 +9,22 @@ import os import amrex -from impactx import (ImpactX, MADXParser, RefPart, distribution, elements, - madx2impactx_beam, madx2impactx_lattice) +from impactx import ( + ImpactX, + MADXParser, + RefPart, + distribution, + elements, + madx2impactx_beam, + madx2impactx_lattice, +) 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_diagnostics(False) # benchmarking sim.set_slice_step_diagnostics(True) # domain decomposition & space charge mesh @@ -45,7 +52,7 @@ 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 + energy=float(madx.getEtot()), # MADX default energy is in GeV ) @@ -55,27 +62,29 @@ # @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 +qm_qeeV = ref_particle_dict["mass"] / ref_particle_dict["charge"] # charge/mass npart = 10000 # number of macro particles 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) + 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(qm_qeeV, 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']) +sim.particle_container().ref_particle().set_energy_MeV( + ref_particle_dict["energy"], ref_particle_dict["mass"] +) chicane = madx2impactx_lattice(beamline) # assign a fodo segment diff --git a/examples/fodo/run_fodo_madx.py b/examples/fodo/run_fodo_madx.py index 55a2d3ff7..b484f8806 100755 --- a/examples/fodo/run_fodo_madx.py +++ b/examples/fodo/run_fodo_madx.py @@ -9,15 +9,22 @@ import os import amrex -from impactx import (ImpactX, MADXParser, RefPart, distribution, elements, - madx2impactx_beam, madx2impactx_lattice) +from impactx import ( + ImpactX, + MADXParser, + RefPart, + distribution, + elements, + madx2impactx_beam, + madx2impactx_lattice, +) 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_diagnostics(False) # benchmarking sim.set_slice_step_diagnostics(True) # domain decomposition & space charge mesh @@ -44,7 +51,7 @@ 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 + energy=float(madx.getEtot()), # MADX default energy is in GeV ) # load a 2 GeV electron beam with an initial @@ -53,27 +60,29 @@ # @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 +qm_qeeV = ref_particle_dict["mass"] / ref_particle_dict["charge"] # charge/mass npart = 10000 # number of macro particles 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) + 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(qm_qeeV, 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']) +sim.particle_container().ref_particle().set_energy_MeV( + ref_particle_dict["energy"], ref_particle_dict["mass"] +) fodo = madx2impactx_lattice(beamline) # assign a fodo segment diff --git a/src/python/impactx/MADXParser.py b/src/python/impactx/MADXParser.py index d1930532c..2aad8f586 100644 --- a/src/python/impactx/MADXParser.py +++ b/src/python/impactx/MADXParser.py @@ -13,14 +13,17 @@ class MADXParserError(Exception): pass + class MADXInputError(MADXParserError): def __init__(self, args, with_traceback): self.args = args self.with_traceback = with_traceback + class MADXInputWarning(UserWarning): pass + class MADXParser: """ Simple MADX parser. @@ -29,71 +32,61 @@ class MADXParser: def __init__(self): - self.__drift = { - 'name': '', - 'l': 0.0, - 'type': 'drift' - } + self.__drift = {"name": "", "l": 0.0, "type": "drift"} - self.__drift_pattern = '(.*):drift,(.*)=(.*);' + self.__drift_pattern = "(.*):drift,(.*)=(.*);" - self.__quadrupole = { - 'name': '', - 'l': 0.0, - 'k1': 0.0, - 'type': 'quad' - } + self.__quadrupole = {"name": "", "l": 0.0, "k1": 0.0, "type": "quad"} # don't count name and type --> len - 2 self.__nQuad = 2 * (len(self.__quadrupole) - 2) - self.quad_pattern = '(.*):quadrupole,(.*)=(.*),(.*)=(.*);' + self.quad_pattern = "(.*):quadrupole,(.*)=(.*),(.*)=(.*);" self.__sbend = { - 'name': '', - 'l': 0.0, - 'angle': 0.0, - 'k1': 0.0, - 'e1': 0.0, - 'e2': 0.0, - 'type': 'sbend' + "name": "", + "l": 0.0, + "angle": 0.0, + "k1": 0.0, + "e1": 0.0, + "e2": 0.0, + "type": "sbend", } - self.__sbend_pattern = '(.*):sbend,(.*)=(.*),(.*)=(.*),(.*)=(.*),(.*)=(.*),(.*)=(.*);' + self.__sbend_pattern = ( + "(.*):sbend,(.*)=(.*),(.*)=(.*),(.*)=(.*),(.*)=(.*),(.*)=(.*);" + ) # don't count name and type --> len - 2 # TODO add rbend self.__nDipole = 2 * (len(self.__sbend) - 2) self.beam = { - 'energy': 0.0, + "energy": 0.0, # TODO extend by 'PC' - 'particle': '' + "particle": "", } self.__nBeam = 2 * len(self.beam) - self.beam_pattern = 'beam,(.*)=(.*),(.*)=(.*);' + self.beam_pattern = "beam,(.*)=(.*),(.*)=(.*);" self.__line = { - 'name': '', - 'elem': [], + "name": "", + "elem": [], } - self.__line_pattern = '(.*):line=\(+(.*)\);' + self.__line_pattern = "(.*):line=\(+(.*)\);" - self.sequence = { - 'name': '' - } + self.sequence = {"name": ""} - self.seq_pattern = 'use,sequence=(.*);' + self.seq_pattern = "use,sequence=(.*);" self.__elements = [] self.__lines = [] self.__lattice = [] - # nonblank_lines inspired by # https://stackoverflow.com/questions/4842057/easiest-way-to-ignore-blank-lines-when-reading-a-file-in-python def nonblank_lines_to_lowercase(self, f): @@ -110,105 +103,131 @@ def parse(self, fn): nLine = 0 - with open(fn, 'r') as f: + with open(fn, "r") as f: for line in self.nonblank_lines_to_lowercase(f): nLine += 1 line = self._noWhitespace(line) - if line[0] == '!': + if line[0] == "!": # this is a comment pass - elif 'drift' in line: + elif "drift" in line: obj = re.match(self.__drift_pattern, line) # first tag is name - self.__drift['name'] = obj.group(1) + self.__drift["name"] = obj.group(1) if obj.group(2) in self.__drift: self.__drift[obj.group(2)] = float(obj.group(3)) else: - raise MADXInputError('Drift', - 'Line ' + str(nLine) + ': Parameter ' + "'" + obj.group(2) + "'" + - ' does not exist for drift.') + raise MADXInputError( + "Drift", + "Line " + + str(nLine) + + ": Parameter " + + "'" + + obj.group(2) + + "'" + + " does not exist for drift.", + ) - self.__elements.append( self.__drift.copy() ) + self.__elements.append(self.__drift.copy()) - elif 'quadrupole' in line: + elif "quadrupole" in line: obj = re.match(self.quad_pattern, line) # first tag is name - self.__quadrupole['name'] = obj.group(1) + self.__quadrupole["name"] = obj.group(1) - for i in range(2, self.__nQuad+2, 2): + for i in range(2, self.__nQuad + 2, 2): if obj.group(i) in self.__quadrupole: - self.__quadrupole[obj.group(i)] = float(obj.group(i+1)) + self.__quadrupole[obj.group(i)] = float(obj.group(i + 1)) else: - raise MADXInputError('Quadrupole', - 'Line ' + str(nLine) + ': Parameter ' + "'" + obj.group(i) + "'" + - ' does not exist for quadrupole.') + raise MADXInputError( + "Quadrupole", + "Line " + + str(nLine) + + ": Parameter " + + "'" + + obj.group(i) + + "'" + + " does not exist for quadrupole.", + ) - self.__elements.append( self.__quadrupole.copy() ) + self.__elements.append(self.__quadrupole.copy()) - elif 'sbend' in line: + elif "sbend" in line: obj = re.match(self.__sbend_pattern, line) - if obj: # TODO remove DEBUG line + if obj: # TODO remove DEBUG line print("OBJ: ", obj) # TODO remove DEBUG line else: print("Oops") # first tag is name - self.__sbend['name'] = obj.group(1) + self.__sbend["name"] = obj.group(1) for i in range(2, self.__nDipole + 2, 2): if obj.group(i) in self.__sbend: self.__sbend[obj.group(i)] = float(obj.group(i + 1)) else: - raise MADXInputError('Dipole', - 'Line ' + str(nLine) + ': Parameter ' + - "'" + obj.group(i) + "'" + - ' does not exist for dipole.') + raise MADXInputError( + "Dipole", + "Line " + + str(nLine) + + ": Parameter " + + "'" + + obj.group(i) + + "'" + + " does not exist for dipole.", + ) self.__elements.append(self.__sbend.copy()) - elif 'marker' in line: + elif "marker" in line: pass - elif 'beam' in line: + elif "beam" in line: obj = re.match(self.beam_pattern, line) for i in range(1, self.__nBeam, 2): if obj.group(i) in self.beam: - if obj.group(i+1).isdigit(): - self.beam[obj.group(i)] = float(obj.group(i+1)) + if obj.group(i + 1).isdigit(): + self.beam[obj.group(i)] = float(obj.group(i + 1)) else: - self.beam[obj.group(i)] = obj.group(i+1) + self.beam[obj.group(i)] = obj.group(i + 1) else: - raise MADXInputError('Beam', - 'Line ' + str(nLine) + ': Parameter ' + - "'" + obj.group(i) + "'" + - ' does not exist for beam.') - - elif 'line' in line: + raise MADXInputError( + "Beam", + "Line " + + str(nLine) + + ": Parameter " + + "'" + + obj.group(i) + + "'" + + " does not exist for beam.", + ) + + elif "line" in line: obj = re.match(self.__line_pattern, line) - self.__line['name'] = obj.group(1) + self.__line["name"] = obj.group(1) - lines = obj.group(2).split(',') + lines = obj.group(2).split(",") newlines = [] # check multiplication and insert that many lines for l in lines: - if '*' in l: - tmp = l.split('*') + if "*" in l: + tmp = l.split("*") n = 0 - ll = '' + ll = "" if tmp[0].isdigit(): n = int(tmp[0]) @@ -222,25 +241,24 @@ def parse(self, fn): else: newlines.append(l) - self.__line['elem'] = newlines + self.__line["elem"] = newlines - self.__lines.append( self.__line.copy() ) + self.__lines.append(self.__line.copy()) - elif 'use' in line and 'sequence' in line: + elif "use" in line and "sequence" in line: obj = re.match(self.seq_pattern, line) - self.sequence['name'] = obj.group(1) + self.sequence["name"] = obj.group(1) else: - raise MADXInputError('', 'Error at line ' + str(nLine)) - + raise MADXInputError("", "Error at line " + str(nLine)) # 14. Oct. 2017, # https://stackoverflow.com/questions/7900882/extract-item-from-list-of-dictionaries for l in self.__lines: # TODO remove this DEBUG line print(l) print(self.sequence) - start = [l for l in self.__lines if l['name'] == self.sequence['name']][0] + start = [l for l in self.__lines if l["name"] == self.sequence["name"]][0] self._flatten(start) @@ -249,24 +267,22 @@ def parse(self, fn): self.__lattice = self._combine(self.__lattice) - def _flatten(self, line): """ Find sublines. """ - name = line['name'] + name = line["name"] for l in self.__lines: - if name in l['name']: - for ll in l['elem']: + if name in l["name"]: + for ll in l["elem"]: for lll in self.__lines: - if lll['name'] == ll: + if lll["name"] == ll: self.__lattice.append(lll) self._flatten(lll) - def _combine(self, lattice): """ Combine to one list of all basic @@ -275,21 +291,19 @@ def _combine(self, lattice): return a list of of element dictionaries """ - l1 = self.__lattice[0] for i in range(1, len(self.__lattice)): l2 = self.__lattice[i] - for e in l1['elem']: - if l2['name'] == e: - idx = l1['elem'].index(e) - l1['elem'].remove(e) - l1['elem'] = l1['elem'][0:idx] + l2['elem'] + l1['elem'][idx:] + for e in l1["elem"]: + if l2["name"] == e: + idx = l1["elem"].index(e) + l1["elem"].remove(e) + l1["elem"] = l1["elem"][0:idx] + l2["elem"] + l1["elem"][idx:] return l1 - def _noWhitespace(self, string): """ Remove white space from a string. @@ -298,8 +312,7 @@ def _noWhitespace(self, string): https://stackoverflow.com/questions/3739909/how-to-strip-all-whitespace-from-string """ - return string.replace(' ', '') - + return string.replace(" ", "") def __str__(self): @@ -309,74 +322,98 @@ def __str__(self): # drift, dipole, quadrupole nTypes = [0, 0, 0] - for elem in self.__lattice['elem']: + for elem in self.__lattice["elem"]: for e in self.__elements: - if elem == e['name']: - length += e['l'] + if elem == e["name"]: + length += e["l"] - if e['type'] == 'drift': + if e["type"] == "drift": nTypes[0] += 1 - elif e['type'] == 'sbend': + elif e["type"] == "sbend": nTypes[1] += 1 - elif e['type'] == 'quad': + elif e["type"] == "quad": nTypes[2] += 1 break - sign = '*' * 70 - info = sign + '\n'\ - + 'MADX-Parser information:\n' \ - + ' length:\t' + str(length) + ' [m]\n'\ - + ' #elements:\t' + str(len(self.__lattice['elem'])) + '\n'\ - + ' * #drift:\t' + str(nTypes[0]) + '\n'\ - + ' * #dipole:\t' + str(nTypes[1]) + '\n'\ - + ' * #quadrupole:\t' + str(nTypes[2]) + '\n'\ - + ' beam:\t\n'\ - + ' * particle:\t' + self.beam['particle'] + '\n'\ - + ' * total energy:\t' + str(self.beam['energy']) + ' [GeV]\n'\ + sign = "*" * 70 + info = ( + sign + + "\n" + + "MADX-Parser information:\n" + + " length:\t" + + str(length) + + " [m]\n" + + " #elements:\t" + + str(len(self.__lattice["elem"])) + + "\n" + + " * #drift:\t" + + str(nTypes[0]) + + "\n" + + " * #dipole:\t" + + str(nTypes[1]) + + "\n" + + " * #quadrupole:\t" + + str(nTypes[2]) + + "\n" + + " beam:\t\n" + + " * particle:\t" + + self.beam["particle"] + + "\n" + + " * total energy:\t" + + str(self.beam["energy"]) + + " [GeV]\n" + sign + ) return info else: - return 'No information available.' + return "No information available." def getBeamline(self): if self.__lattice: beamline = [] - for elem in self.__lattice['elem']: + for elem in self.__lattice["elem"]: for e in self.__elements: - if elem == e['name']: - if e['type'] == 'drift': - print('Drift L= '+str(e['l'])) + if elem == e["name"]: + if e["type"] == "drift": + print("Drift L= " + str(e["l"])) beamline.append(e) - elif e['type'] == 'sbend': - print("Sbend L= ",e['l'],' angle = ',e['angle']) + elif e["type"] == "sbend": + print("Sbend L= ", e["l"], " angle = ", e["angle"]) beamline.append(e) - elif e['type'] == 'quad': - print("Quadrupole L= ",e['l'],' k1 = ',e['k1']) + elif e["type"] == "quad": + print("Quadrupole L= ", e["l"], " k1 = ", e["k1"]) beamline.append(e) else: - print ( 'Skipping element type ' + "'" + - e['type'] + "'" ) + print("Skipping element type " + "'" + e["type"] + "'") return beamline else: return [] - def getParticle(self): - particle = self.beam['particle'] - known_particles = ['positron', 'electron', 'proton', 'antiproton', 'posmuon', 'negmuon', 'ion'] + particle = self.beam["particle"] + known_particles = [ + "positron", + "electron", + "proton", + "antiproton", + "posmuon", + "negmuon", + "ion", + ] if particle not in known_particles: - warning_message = 'No particle type ' + "'" + self.beam['particle'] + "' available." + warning_message = ( + "No particle type " + "'" + self.beam["particle"] + "' available." + ) warnings.warn(warning_message, MADXInputWarning) else: pass return particle - def getEtot(self): - return self.beam['energy'] + return self.beam["energy"] diff --git a/src/python/impactx/madx_to_impactx.py b/src/python/impactx/madx_to_impactx.py index 91b9e23d8..1264e5937 100644 --- a/src/python/impactx/madx_to_impactx.py +++ b/src/python/impactx/madx_to_impactx.py @@ -13,7 +13,7 @@ from impactx import elements -def madx2impactx_lattice(parsed_beamline,nslice=25): +def madx2impactx_lattice(parsed_beamline, nslice=25): """ Function that converts a list of elements in the MADXParser format into ImpactX format :param parsed_beamline: list of dictionaries @@ -37,27 +37,26 @@ def madx2impactx_lattice(parsed_beamline,nslice=25): for d in parsed_beamline: - if d['type'] in [k.casefold() for k in list(madx_to_impactx_dict.keys())]: - if d['name'] == "drift": + if d["type"] in [k.casefold() for k in list(madx_to_impactx_dict.keys())]: + if d["name"] == "drift": + impactx_beamline.append(elements.Drift(ds=d["l"], nslice=nslice)) + elif d["name"] == "quadrupole": impactx_beamline.append( - elements.Drift(ds=d['l'], nslice=nslice) - ) - elif d['name'] == "quadrupole": - impactx_beamline.append( - elements.Quad(ds=d['l'], k=d['k1'], nslice=nslice) + elements.Quad(ds=d["l"], k=d["k1"], nslice=nslice) ) else: raise NotImplementedError( "The beamline element named ", - d['name'], + d["name"], "of type ", - d['type'], + d["type"], "is not implemented in impactx.elements.", "Available elements are:", - list(madx_to_impactx_dict.keys()) + list(madx_to_impactx_dict.keys()), ) return impactx_beamline + def madx2impactx_beam(particle, charge=None, mass=None, energy=None): """ Function that converts a list of beam parameter dictionaries in the MADXParser format into ImpactX format @@ -73,34 +72,48 @@ def madx2impactx_beam(particle, charge=None, mass=None, energy=None): """ GeV2MeV = 1e-3 - kg2MeV = sc.c ** 2 / sc.electron_volt * 1e-6 - muon_mass = sc.physical_constants['electron-muon mass ratio'][0] / sc.m_e + kg2MeV = sc.c**2 / sc.electron_volt * 1e-6 + muon_mass = sc.physical_constants["electron-muon mass ratio"][0] / sc.m_e if energy is None: energy_MeV = 1e3 # MAD-X default is 1 GeV particle energy else: energy_MeV = energy * GeV2MeV impactx_beam = { - 'positron': {'mass': sc.m_e * kg2MeV, 'charge': 1}, - 'electron': {'mass': sc.m_e * kg2MeV, 'charge': -1}, - 'proton': {'mass': sc.m_p * kg2MeV, 'charge': 1}, - 'antiproton': {'mass': sc.m_p * kg2MeV, 'charge': -1}, - 'posmuon': {'mass': muon_mass * kg2MeV, 'charge': 1}, # positively charged muon (anti-muon) - 'negmuon': {'mass': muon_mass * kg2MeV, 'charge': -1}, # negatively charged muon - 'ion': {'mass': sc.m_u * kg2MeV, 'charge': 1}, - 'generic': {'mass': mass, 'charge': charge} + "positron": {"mass": sc.m_e * kg2MeV, "charge": 1}, + "electron": {"mass": sc.m_e * kg2MeV, "charge": -1}, + "proton": {"mass": sc.m_p * kg2MeV, "charge": 1}, + "antiproton": {"mass": sc.m_p * kg2MeV, "charge": -1}, + "posmuon": { + "mass": muon_mass * kg2MeV, + "charge": 1, + }, # positively charged muon (anti-muon) + "negmuon": { + "mass": muon_mass * kg2MeV, + "charge": -1, + }, # negatively charged muon + "ion": {"mass": sc.m_u * kg2MeV, "charge": 1}, + "generic": {"mass": mass, "charge": charge}, } - if particle not in impactx_beam.keys(): - warnings.warn('Particle species name "' + particle + '" not in ' + impactx_beam.keys(), UserWarning) - print('Choosing generic particle species, using provided `charge`, `mass` and `energy`.') - _particle = 'generic' + warnings.warn( + 'Particle species name "' + particle + '" not in ' + impactx_beam.keys(), + UserWarning, + ) + print( + "Choosing generic particle species, using provided `charge`, `mass` and `energy`." + ) + _particle = "generic" else: _particle = particle - print('Choosing provided particle species "', particle, '", ignoring potentially provided `charge` and `mass` and setting defaults.') + print( + 'Choosing provided particle species "', + particle, + '", ignoring potentially provided `charge` and `mass` and setting defaults.', + ) reference_particle = impactx_beam[_particle] - reference_particle['energy'] = energy_MeV + reference_particle["energy"] = energy_MeV return reference_particle From 64181d124f70cf9983256839f2e88f607418ff7a Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 30 Aug 2022 13:30:03 -0700 Subject: [PATCH 08/19] MADXParser: Regex Literals Use proper regex literals. Avoids: ``` DeprecationWarning: invalid escape sequence \( ``` et al. --- src/python/impactx/MADXParser.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/python/impactx/MADXParser.py b/src/python/impactx/MADXParser.py index 2aad8f586..804de11da 100644 --- a/src/python/impactx/MADXParser.py +++ b/src/python/impactx/MADXParser.py @@ -34,14 +34,14 @@ def __init__(self): self.__drift = {"name": "", "l": 0.0, "type": "drift"} - self.__drift_pattern = "(.*):drift,(.*)=(.*);" + self.__drift_pattern = r"(.*):drift,(.*)=(.*);" self.__quadrupole = {"name": "", "l": 0.0, "k1": 0.0, "type": "quad"} # don't count name and type --> len - 2 self.__nQuad = 2 * (len(self.__quadrupole) - 2) - self.quad_pattern = "(.*):quadrupole,(.*)=(.*),(.*)=(.*);" + self.quad_pattern = r"(.*):quadrupole,(.*)=(.*),(.*)=(.*);" self.__sbend = { "name": "", @@ -54,7 +54,7 @@ def __init__(self): } self.__sbend_pattern = ( - "(.*):sbend,(.*)=(.*),(.*)=(.*),(.*)=(.*),(.*)=(.*),(.*)=(.*);" + r"(.*):sbend,(.*)=(.*),(.*)=(.*),(.*)=(.*),(.*)=(.*),(.*)=(.*);" ) # don't count name and type --> len - 2 @@ -69,18 +69,18 @@ def __init__(self): self.__nBeam = 2 * len(self.beam) - self.beam_pattern = "beam,(.*)=(.*),(.*)=(.*);" + self.beam_pattern = r"beam,(.*)=(.*),(.*)=(.*);" self.__line = { "name": "", "elem": [], } - self.__line_pattern = "(.*):line=\(+(.*)\);" + self.__line_pattern = r"(.*):line=\(+(.*)\);" self.sequence = {"name": ""} - self.seq_pattern = "use,sequence=(.*);" + self.seq_pattern = r"use,sequence=(.*);" self.__elements = [] self.__lines = [] From f05c93fb9d95acdb83ac2ce57a0fafe2cc5f3820 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 30 Aug 2022 14:13:28 -0700 Subject: [PATCH 09/19] Lattice: Cleaner User API & Docs --- docs/source/usage/python.rst | 19 +++++++++ examples/chicane/run_chicane_madx.py | 20 ++-------- examples/fodo/run_fodo_madx.py | 20 ++-------- src/python/elements.cpp | 47 +++++++++++++--------- src/python/impactx/__init__.py | 18 +++++---- src/python/impactx/madx_to_impactx.py | 56 +++++++++++++++++++-------- 6 files changed, 105 insertions(+), 75 deletions(-) diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index de23e0507..835df7b96 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -305,6 +305,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/chicane/run_chicane_madx.py b/examples/chicane/run_chicane_madx.py index a6ac56afc..7da2a3d59 100755 --- a/examples/chicane/run_chicane_madx.py +++ b/examples/chicane/run_chicane_madx.py @@ -6,18 +6,10 @@ # # -*- coding: utf-8 -*- -import os import amrex -from impactx import ( - ImpactX, - MADXParser, - RefPart, - distribution, - elements, - madx2impactx_beam, - madx2impactx_lattice, -) +from impactx import (ImpactX, MADXParser, RefPart, distribution, elements, + madx2impactx_beam, madx2impactx_lattice) sim = ImpactX() @@ -30,13 +22,10 @@ # domain decomposition & space charge mesh sim.init_grids() -# @TODO make this better ... shouldn't need that. But somehow the working directory is -dir_path = os.path.dirname(os.path.realpath(__file__)) - try: madx = MADXParser() - madx.parse(f"{dir_path}/chicane.madx") + madx.parse("chicane.madx") except Exception as e: print(f"Unexpected {e = }, {type(e) = }") @@ -86,9 +75,8 @@ ref_particle_dict["energy"], ref_particle_dict["mass"] ) -chicane = madx2impactx_lattice(beamline) # assign a fodo segment -sim.lattice.extend(chicane) +sim.lattice.load_file("chicane.madx", nslice=25) # run simulation sim.evolve() diff --git a/examples/fodo/run_fodo_madx.py b/examples/fodo/run_fodo_madx.py index b484f8806..66a9a5ce4 100755 --- a/examples/fodo/run_fodo_madx.py +++ b/examples/fodo/run_fodo_madx.py @@ -6,18 +6,10 @@ # # -*- coding: utf-8 -*- -import os import amrex -from impactx import ( - ImpactX, - MADXParser, - RefPart, - distribution, - elements, - madx2impactx_beam, - madx2impactx_lattice, -) +from impactx import (ImpactX, MADXParser, RefPart, distribution, elements, + madx2impactx_beam, madx2impactx_lattice) sim = ImpactX() @@ -30,13 +22,10 @@ # domain decomposition & space charge mesh sim.init_grids() -# @TODO make this better ... shouldn't need that -dir_path = os.path.dirname(os.path.realpath(__file__)) - try: madx = MADXParser() - madx.parse(f"{dir_path}/fodo.madx") + madx.parse("fodo.madx") except Exception as e: print(f"Unexpected {e = }, {type(e) = }") @@ -84,9 +73,8 @@ ref_particle_dict["energy"], ref_particle_dict["mass"] ) -fodo = madx2impactx_lattice(beamline) # assign a fodo segment -sim.lattice.extend(fodo) +sim.lattice.load_file("fodo.madx", nslice=25) # run simulation sim.evolve() diff --git a/src/python/elements.cpp b/src/python/elements.cpp index 8c821a4f8..28961bd82 100644 --- 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(); - 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(); + 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 */ diff --git a/src/python/impactx/__init__.py b/src/python/impactx/__init__.py index 2e27cfaa8..14ea8dbe0 100644 --- a/src/python/impactx/__init__.py +++ b/src/python/impactx/__init__.py @@ -1,13 +1,15 @@ -from . import impactx_pybind -from .MADXParser import * # noqa +from . import impactx_pybind as cxx from .impactx_pybind import * # noqa -from .madx_to_impactx import madx2impactx_beam, madx2impactx_lattice # noqa +from .madx_to_impactx import read_beam, read_lattice # noqa -__version__ = impactx_pybind.__version__ -__doc__ = impactx_pybind.__doc__ -__license__ = impactx_pybind.__license__ -__author__ = impactx_pybind.__author__ +__version__ = cxx.__version__ +__doc__ = cxx.__doc__ +__license__ = cxx.__license__ +__author__ = cxx.__author__ # 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 diff --git a/src/python/impactx/madx_to_impactx.py b/src/python/impactx/madx_to_impactx.py index 1264e5937..00eee51b3 100644 --- a/src/python/impactx/madx_to_impactx.py +++ b/src/python/impactx/madx_to_impactx.py @@ -12,11 +12,14 @@ from impactx import elements +from .MADXParser import MADXParser -def madx2impactx_lattice(parsed_beamline, nslice=25): + +def lattice(parsed_beamline, nslice=1): """ Function that converts a list of elements in the MADXParser format into ImpactX format :param parsed_beamline: list of dictionaries + :param nslice: number of ds slices per element :return: list of translated dictionaries """ @@ -57,7 +60,19 @@ def madx2impactx_lattice(parsed_beamline, nslice=25): return impactx_beamline -def madx2impactx_beam(particle, charge=None, mass=None, energy=None): +def read_lattice(madx_file, nslice=1): + """ + 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 + :param nslice: number of ds slices per element + :return: list of ImpactX.KnownElements + """ + madx = MADXParser() + beamline = madx.parse(madx_file).getBeamline() + return lattice(beamline, nslice) + + +def beam(particle, charge=None, mass=None, energy=None): """ Function that converts a list of beam parameter dictionaries in the MADXParser format into ImpactX format @@ -71,34 +86,34 @@ def madx2impactx_beam(particle, charge=None, mass=None, energy=None): :return dict: dictionary containing particle and beam attributes in ImpactX units """ - GeV2MeV = 1e-3 - kg2MeV = sc.c**2 / sc.electron_volt * 1e-6 + GeV2MeV = 1.0e-3 + kg2MeV = sc.c**2 / sc.electron_volt * 1.0e-6 muon_mass = sc.physical_constants["electron-muon mass ratio"][0] / sc.m_e if energy is None: - energy_MeV = 1e3 # MAD-X default is 1 GeV particle energy + energy_MeV = 1.0e3 # MAD-X default is 1 GeV particle energy else: energy_MeV = energy * GeV2MeV impactx_beam = { - "positron": {"mass": sc.m_e * kg2MeV, "charge": 1}, - "electron": {"mass": sc.m_e * kg2MeV, "charge": -1}, - "proton": {"mass": sc.m_p * kg2MeV, "charge": 1}, - "antiproton": {"mass": sc.m_p * kg2MeV, "charge": -1}, + "positron": {"mass": sc.m_e * kg2MeV, "charge": 1.0}, + "electron": {"mass": sc.m_e * kg2MeV, "charge": -1.0}, + "proton": {"mass": sc.m_p * kg2MeV, "charge": 1.0}, + "antiproton": {"mass": sc.m_p * kg2MeV, "charge": -1.0}, "posmuon": { "mass": muon_mass * kg2MeV, - "charge": 1, + "charge": 1.0, }, # positively charged muon (anti-muon) "negmuon": { "mass": muon_mass * kg2MeV, - "charge": -1, + "charge": -1.0, }, # negatively charged muon - "ion": {"mass": sc.m_u * kg2MeV, "charge": 1}, + "ion": {"mass": sc.m_u * kg2MeV, "charge": 1.0}, "generic": {"mass": mass, "charge": charge}, } if particle not in impactx_beam.keys(): warnings.warn( - 'Particle species name "' + particle + '" not in ' + impactx_beam.keys(), + f"Particle species name '{particle}' not in '{impactx_beam.keys()}'", UserWarning, ) print( @@ -108,12 +123,21 @@ def madx2impactx_beam(particle, charge=None, mass=None, energy=None): else: _particle = particle print( - 'Choosing provided particle species "', - particle, - '", ignoring potentially provided `charge` and `mass` and setting defaults.', + f"Choosing provided particle species '{particle}', ignoring potentially provided `charge` and `mass` and setting defaults.", ) reference_particle = impactx_beam[_particle] reference_particle["energy"] = energy_MeV return reference_particle + + +def read_beam(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) From e585bfe6e6ef6ada60d08439276344cceaede7f5 Mon Sep 17 00:00:00 2001 From: Marco Garten Date: Tue, 30 Aug 2022 14:31:16 -0700 Subject: [PATCH 10/19] Fix chicane MADX example and add parser elements --- examples/chicane/chicane.madx | 23 +++++------ src/python/impactx/MADXParser.py | 56 ++++++++++++++++++++++----- src/python/impactx/madx_to_impactx.py | 8 ++++ 3 files changed, 66 insertions(+), 21 deletions(-) diff --git a/examples/chicane/chicane.madx b/examples/chicane/chicane.madx index 19b27766e..49c935787 100644 --- a/examples/chicane/chicane.madx +++ b/examples/chicane/chicane.madx @@ -1,15 +1,16 @@ -beam, particle=electron, energy=40.e-3; +beam, particle=electron, energy=5.0; -DLEFT: drift, L=0.7578; -DMIDDLE: drift, L=0.9575; -DRIGHT: drift, L=0.7578; +DRIFT1: drift, L=5.0058489435; +DRIFT2: drift, L=1.0; +DRIFT3: drift, L=2.0; ! TODO make this work with inline calculations -! theta=pi/10.; +! theta=0.50037/10.35; +! inv_rho=1.0/10.35 ! TODO put theta again everywhere after `angle=` with its right sign -D114: sbend, L=0.264687, angle= 0.31415926535, e1= 0.000, e2= 0.31415926535, k1= 0.00; -D115: sbend, L=0.264687, angle=-0.31415926535, e1=-0.31415926535, e2= 0.000, k1= 0.00; -D116: sbend, L=0.264687, angle=-0.31415926535, e1= 0.000, e2=-0.31415926535, k1= 0.00; -D117: sbend, L=0.264687, angle= 0.31415926535, e1= 0.31415926535, e2= 0.000, k1= 0.00; +SBEND1: sbend, L=0.50037, angle=0.048345, e1=0.000, e2=0.000, k1=0.00; +SBEND2: sbend, L=0.50037, angle=-0.048345, e1=0.000, e2=0.000, k1=0.00; +DIPEDGE1: dipedge, H=0.096618357, e1=-0.048345620280243, fint=0.000, hgap=0.000, tilt=0.00; +DIPEDGE2: dipedge, H=0.096618357, e1=0.048345620280243, fint=0.000, hgap=0.000, tilt=0.00;; -BC1: Line=(D114,DLEFT,D115,DMIDDLE,D116,DRIGHT,D117); -USE, SEQUENCE = BC1; +CHICANE: Line=(SBEND1,DIPEDGE1,DRIFT1,DIPEDGE2,SBEND2,DRIFT2,SBEND2,DIPEDGE2,DRIFT1,DIPEDGE1,SBEND1,DRIFT3); +USE, SEQUENCE = CHICANE; diff --git a/src/python/impactx/MADXParser.py b/src/python/impactx/MADXParser.py index 804de11da..392a88bce 100644 --- a/src/python/impactx/MADXParser.py +++ b/src/python/impactx/MADXParser.py @@ -61,6 +61,21 @@ def __init__(self): # TODO add rbend self.__nDipole = 2 * (len(self.__sbend) - 2) + self.__dipedge = { + 'name': '', + 'h': 0.0, + 'e1': 0.0, + 'fint': 0.0, + 'hgap': 0.0, + 'tilt': 0.0, + 'type': 'dipedge' + } + + self.__dipedge_pattern = '(.*):dipedge,(.*)=(.*),(.*)=(.*),(.*)=(.*),(.*)=(.*),(.*)=(.*);' + + # don't count name and type --> len - 2 + self.__nDipedge = 2 * (len(self.__dipedge) - 2) + self.beam = { "energy": 0.0, # TODO extend by 'PC' @@ -161,10 +176,6 @@ def parse(self, fn): elif "sbend" in line: obj = re.match(self.__sbend_pattern, line) - if obj: # TODO remove DEBUG line - print("OBJ: ", obj) # TODO remove DEBUG line - else: - print("Oops") # first tag is name self.__sbend["name"] = obj.group(1) @@ -186,7 +197,27 @@ def parse(self, fn): ) self.__elements.append(self.__sbend.copy()) - elif "marker" in line: + + elif 'dipedge' in line: + + obj = re.match(self.__dipedge_pattern, line) + + # first tag is name + self.__dipedge['name'] = obj.group(1) + + 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('DipEgde', + 'Line ' + str(nLine) + ': Parameter ' + + "'" + obj.group(i) + "'" + + ' does not exist for dipole edge.') + + self.__elements.append(self.__dipedge.copy()) + + elif 'marker' in line: pass elif "beam" in line: @@ -255,9 +286,6 @@ def parse(self, fn): # 14. Oct. 2017, # https://stackoverflow.com/questions/7900882/extract-item-from-list-of-dictionaries - for l in self.__lines: # TODO remove this DEBUG line - print(l) - print(self.sequence) start = [l for l in self.__lines if l["name"] == self.sequence["name"]][0] self._flatten(start) @@ -319,8 +347,8 @@ def __str__(self): if self.__lattice: length = 0.0 - # drift, dipole, quadrupole - nTypes = [0, 0, 0] + # drift, dipole, quadrupole, dipedge + nTypes = [0, 0, 0, 0] for elem in self.__lattice["elem"]: for e in self.__elements: @@ -333,6 +361,8 @@ def __str__(self): nTypes[1] += 1 elif e["type"] == "quad": nTypes[2] += 1 + elif e['type'] == 'dipedge': + nTypes[3] += 1 break sign = "*" * 70 @@ -355,6 +385,9 @@ def __str__(self): + " * #quadrupole:\t" + str(nTypes[2]) + "\n" + + " * #dipedge:\t" + + str(nTypes[3]) + + "\n" + " beam:\t\n" + " * particle:\t" + self.beam["particle"] @@ -386,6 +419,9 @@ def getBeamline(self): elif e["type"] == "quad": print("Quadrupole L= ", e["l"], " k1 = ", e["k1"]) beamline.append(e) + elif e['type'] == 'dipedge': + print("Dipedge H= ",e['h'],' E1 = ',e['e1']) + 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 00eee51b3..0b8b3b3d4 100644 --- a/src/python/impactx/madx_to_impactx.py +++ b/src/python/impactx/madx_to_impactx.py @@ -47,6 +47,14 @@ def lattice(parsed_beamline, nslice=1): impactx_beamline.append( elements.Quad(ds=d["l"], k=d["k1"], nslice=nslice) ) + elif d['name'] == "sbend": + impactx_beamline.append( + elements.Sbend(ds=d['l'], rc=d['l']/d['angle'], nslice=nslice) + ) + elif d['name'] == "dipedge": + impactx_beamline.append( + elements.DipEdge(psi=d['e1'], rc=1/d['h'], nslice=nslice) + ) else: raise NotImplementedError( "The beamline element named ", From d63a1d04b40196d7012ba7dbb53df0388747b76a Mon Sep 17 00:00:00 2001 From: Marco Garten Date: Tue, 30 Aug 2022 15:30:07 -0700 Subject: [PATCH 11/19] Add additional parameters to DipEdge and adjust formatting --- examples/chicane/chicane.madx | 1 + src/python/impactx/MADXParser.py | 42 ++++++++++++++++----------- src/python/impactx/madx_to_impactx.py | 10 +++++-- 3 files changed, 34 insertions(+), 19 deletions(-) diff --git a/examples/chicane/chicane.madx b/examples/chicane/chicane.madx index 49c935787..5c11b2464 100644 --- a/examples/chicane/chicane.madx +++ b/examples/chicane/chicane.madx @@ -9,6 +9,7 @@ DRIFT3: drift, L=2.0; ! TODO put theta again everywhere after `angle=` with its right sign SBEND1: sbend, L=0.50037, angle=0.048345, e1=0.000, e2=0.000, k1=0.00; SBEND2: sbend, L=0.50037, angle=-0.048345, e1=0.000, e2=0.000, k1=0.00; +! dipole edge elements DIPEDGE1: dipedge, H=0.096618357, e1=-0.048345620280243, fint=0.000, hgap=0.000, tilt=0.00; DIPEDGE2: dipedge, H=0.096618357, e1=0.048345620280243, fint=0.000, hgap=0.000, tilt=0.00;; diff --git a/src/python/impactx/MADXParser.py b/src/python/impactx/MADXParser.py index 392a88bce..a6520ad8b 100644 --- a/src/python/impactx/MADXParser.py +++ b/src/python/impactx/MADXParser.py @@ -62,16 +62,18 @@ def __init__(self): self.__nDipole = 2 * (len(self.__sbend) - 2) self.__dipedge = { - 'name': '', - 'h': 0.0, - 'e1': 0.0, - 'fint': 0.0, - 'hgap': 0.0, - 'tilt': 0.0, - 'type': 'dipedge' + "name": "", + "h": 0.0, + "e1": 0.0, + "fint": 0.0, + "hgap": 0.0, + "tilt": 0.0, + "type": "dipedge" } - self.__dipedge_pattern = '(.*):dipedge,(.*)=(.*),(.*)=(.*),(.*)=(.*),(.*)=(.*),(.*)=(.*);' + self.__dipedge_pattern = ( + r"(.*):dipedge,(.*)=(.*),(.*)=(.*),(.*)=(.*),(.*)=(.*),(.*)=(.*);" + ) # don't count name and type --> len - 2 self.__nDipedge = 2 * (len(self.__dipedge) - 2) @@ -198,26 +200,32 @@ def parse(self, fn): self.__elements.append(self.__sbend.copy()) - elif 'dipedge' in line: + elif "dipedge" in line: obj = re.match(self.__dipedge_pattern, line) # first tag is name - self.__dipedge['name'] = obj.group(1) + self.__dipedge["name"] = obj.group(1) 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('DipEgde', - 'Line ' + str(nLine) + ': Parameter ' + - "'" + obj.group(i) + "'" + - ' does not exist for dipole edge.') + raise MADXInputError( + "DipEgde", + "Line " + + str(nLine) + + ": Parameter " + + "'" + + obj.group(i) + + "'" + + " does not exist for dipole edge." + ) self.__elements.append(self.__dipedge.copy()) - elif 'marker' in line: + elif "marker" in line: pass elif "beam" in line: @@ -419,8 +427,8 @@ def getBeamline(self): elif e["type"] == "quad": print("Quadrupole L= ", e["l"], " k1 = ", e["k1"]) beamline.append(e) - elif e['type'] == 'dipedge': - print("Dipedge H= ",e['h'],' E1 = ',e['e1']) + elif e["type"] == "dipedge": + print("Dipedge H= ",e["h"]," E1 = ",e["e1"]) beamline.append(e) else: print("Skipping element type " + "'" + e["type"] + "'") diff --git a/src/python/impactx/madx_to_impactx.py b/src/python/impactx/madx_to_impactx.py index 0b8b3b3d4..d0d0d7f88 100644 --- a/src/python/impactx/madx_to_impactx.py +++ b/src/python/impactx/madx_to_impactx.py @@ -49,11 +49,17 @@ def lattice(parsed_beamline, nslice=1): ) 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": impactx_beamline.append( - elements.DipEdge(psi=d['e1'], rc=1/d['h'], nslice=nslice) + 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) ) else: raise NotImplementedError( From ece6babb88ff8fadb89d6f8af5954ba99a4e3e3e Mon Sep 17 00:00:00 2001 From: Marco Garten Date: Tue, 30 Aug 2022 16:58:49 -0700 Subject: [PATCH 12/19] Update examples according to API changes Caution: Examples still not working again --- examples/chicane/chicane.madx | 2 +- examples/chicane/run_chicane_madx.py | 4 ++-- examples/fodo/run_fodo_madx.py | 4 ++-- src/python/impactx/MADXParser.py | 4 ++-- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/examples/chicane/chicane.madx b/examples/chicane/chicane.madx index 5c11b2464..21c8c29f4 100644 --- a/examples/chicane/chicane.madx +++ b/examples/chicane/chicane.madx @@ -11,7 +11,7 @@ SBEND1: sbend, L=0.50037, angle=0.048345, e1=0.000, e2=0.000, k1=0.00; SBEND2: sbend, L=0.50037, angle=-0.048345, e1=0.000, e2=0.000, k1=0.00; ! dipole edge elements DIPEDGE1: dipedge, H=0.096618357, e1=-0.048345620280243, fint=0.000, hgap=0.000, tilt=0.00; -DIPEDGE2: dipedge, H=0.096618357, e1=0.048345620280243, fint=0.000, hgap=0.000, tilt=0.00;; +DIPEDGE2: dipedge, H=0.096618357, e1=0.048345620280243, fint=0.000, hgap=0.000, tilt=0.00; CHICANE: Line=(SBEND1,DIPEDGE1,DRIFT1,DIPEDGE2,SBEND2,DRIFT2,SBEND2,DIPEDGE2,DRIFT1,DIPEDGE1,SBEND1,DRIFT3); USE, SEQUENCE = CHICANE; diff --git a/examples/chicane/run_chicane_madx.py b/examples/chicane/run_chicane_madx.py index 7da2a3d59..bab8ed69f 100755 --- a/examples/chicane/run_chicane_madx.py +++ b/examples/chicane/run_chicane_madx.py @@ -9,7 +9,7 @@ import amrex from impactx import (ImpactX, MADXParser, RefPart, distribution, elements, - madx2impactx_beam, madx2impactx_lattice) + read_beam, read_lattice) sim = ImpactX() @@ -37,7 +37,7 @@ beamline = madx.getBeamline() -ref_particle_dict = madx2impactx_beam( +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, diff --git a/examples/fodo/run_fodo_madx.py b/examples/fodo/run_fodo_madx.py index 66a9a5ce4..afd68530e 100755 --- a/examples/fodo/run_fodo_madx.py +++ b/examples/fodo/run_fodo_madx.py @@ -9,7 +9,7 @@ import amrex from impactx import (ImpactX, MADXParser, RefPart, distribution, elements, - madx2impactx_beam, madx2impactx_lattice) + read_beam, read_lattice) sim = ImpactX() @@ -36,7 +36,7 @@ beamline = madx.getBeamline() -ref_particle_dict = madx2impactx_beam( +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, diff --git a/src/python/impactx/MADXParser.py b/src/python/impactx/MADXParser.py index a6520ad8b..f6c6d815d 100644 --- a/src/python/impactx/MADXParser.py +++ b/src/python/impactx/MADXParser.py @@ -213,7 +213,7 @@ def parse(self, fn): self.__dipedge[obj.group(i)] = float(obj.group(i + 1)) else: raise MADXInputError( - "DipEgde", + "DipEdge", "Line " + str(nLine) + ": Parameter " @@ -393,7 +393,7 @@ def __str__(self): + " * #quadrupole:\t" + str(nTypes[2]) + "\n" - + " * #dipedge:\t" + + " * #dipedge:\t" + str(nTypes[3]) + "\n" + " beam:\t\n" From f88d6851fc72c504bf3e9c86ded6d3d0027225d4 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 30 Aug 2022 23:33:32 -0700 Subject: [PATCH 13/19] Cleaning --- docs/source/usage/python.rst | 6 ++++ examples/CMakeLists.txt | 4 +-- examples/chicane/run_chicane_madx.py | 49 +++++--------------------- examples/fodo/run_fodo_madx.py | 45 +++++------------------- src/python/impactx/MADXParser.py | 50 +++++++++++++++------------ src/python/impactx/__init__.py | 10 ++++-- src/python/impactx/madx_to_impactx.py | 33 ++++++++++++------ 7 files changed, 82 insertions(+), 115 deletions(-) diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index 835df7b96..9449c6b44 100644 --- 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 index 1cade33ba..98ca987fd 100644 --- 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 index bab8ed69f..4f8cb1008 100755 --- 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 index afd68530e..85ca93a75 100755 --- 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 index f6c6d815d..d0992451a 100644 --- 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 index 14ea8dbe0..40cc06cb5 100644 --- 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 index d0d0d7f88..8edd61348 100644 --- 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"]) From 1369a80dfddfb4288f5350432e76bb441a68995d Mon Sep 17 00:00:00 2001 From: Marco Garten Date: Wed, 31 Aug 2022 11:13:27 -0700 Subject: [PATCH 14/19] Rename elements in chicane.madx to avoid matching of `drift` in `use, sequence` line with parser. --- examples/chicane/chicane.madx | 16 ++++++++-------- src/python/impactx/madx_to_impactx.py | 1 + 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/examples/chicane/chicane.madx b/examples/chicane/chicane.madx index 21c8c29f4..418bf1803 100644 --- a/examples/chicane/chicane.madx +++ b/examples/chicane/chicane.madx @@ -1,17 +1,17 @@ beam, particle=electron, energy=5.0; -DRIFT1: drift, L=5.0058489435; -DRIFT2: drift, L=1.0; -DRIFT3: drift, L=2.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 theta again everywhere after `angle=` with its right sign -SBEND1: sbend, L=0.50037, angle=0.048345, e1=0.000, e2=0.000, k1=0.00; -SBEND2: sbend, L=0.50037, angle=-0.048345, e1=0.000, e2=0.000, k1=0.00; +SB1: sbend, L=0.50037, angle=0.048345, e1=0.000, e2=0.000, k1=0.00; +SB2: sbend, L=0.50037, angle=-0.048345, e1=0.000, e2=0.000, k1=0.00; ! dipole edge elements -DIPEDGE1: dipedge, H=0.096618357, e1=-0.048345620280243, fint=0.000, hgap=0.000, tilt=0.00; -DIPEDGE2: dipedge, H=0.096618357, e1=0.048345620280243, fint=0.000, hgap=0.000, tilt=0.00; +DIPE1: dipedge, H=0.096618357, e1=-0.048345620280243, fint=0.000, hgap=0.000, tilt=0.00; +DIPE2: dipedge, H=0.096618357, e1=0.048345620280243, fint=0.000, hgap=0.000, tilt=0.00; -CHICANE: Line=(SBEND1,DIPEDGE1,DRIFT1,DIPEDGE2,SBEND2,DRIFT2,SBEND2,DIPEDGE2,DRIFT1,DIPEDGE1,SBEND1,DRIFT3); +CHICANE: Line=(SB1,DIP1,D1,DIP2,SB2,D2,SB2,DIPE2,D1,DIPE1,SB1,D3); USE, SEQUENCE = CHICANE; diff --git a/src/python/impactx/madx_to_impactx.py b/src/python/impactx/madx_to_impactx.py index 8edd61348..61669ae27 100644 --- a/src/python/impactx/madx_to_impactx.py +++ b/src/python/impactx/madx_to_impactx.py @@ -152,6 +152,7 @@ def beam(particle, charge=None, mass=None, energy=None): def read_beam(ref: RefPart, madx_file): """ Function that reads elements from a MAD-X file into a list of ImpactX.KnownElements + :param RefPart ref: ImpactX reference particle (pybind_11 object) :param madx_file: file name to MAD-X file with beamline elements :return: list of ImpactX.KnownElements """ From 5f80f3713e5effdf3e2a2d9f4ca798b9abd66516 Mon Sep 17 00:00:00 2001 From: Marco Garten Date: Wed, 31 Aug 2022 18:36:18 -0700 Subject: [PATCH 15/19] Fix: element type instead of name change MADXParser Quad type to Quadrupole to stay consistent with MADX input --- examples/fodo/fodo.madx | 9 +++++---- src/python/impactx/MADXParser.py | 6 +++--- src/python/impactx/madx_to_impactx.py | 13 +++++++------ 3 files changed, 15 insertions(+), 13 deletions(-) diff --git a/examples/fodo/fodo.madx b/examples/fodo/fodo.madx index bdb706b98..4b51d4a7b 100644 --- a/examples/fodo/fodo.madx +++ b/examples/fodo/fodo.madx @@ -2,9 +2,10 @@ ! TITLE,'FODO.MADX'; BEAM, PARTICLE=ELECTRON,ENERGY=2.0; -D: DRIFT,L=0.25; -QF: QUADRUPOLE,L=0.5,K1=1.0; -QD: QUADRUPOLE,L=0.5,K1=-1.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=(D,QF,D,QD,D); +FODO: LINE=(D1,QF,D2,QD,D1); USE, SEQUENCE = FODO; diff --git a/src/python/impactx/MADXParser.py b/src/python/impactx/MADXParser.py index d0992451a..4095b6333 100644 --- a/src/python/impactx/MADXParser.py +++ b/src/python/impactx/MADXParser.py @@ -37,7 +37,7 @@ def __init__(self): self.__drift_pattern = r"(.*):drift,(.*)=(.*);" - self.__quadrupole = {"name": "", "l": 0.0, "k1": 0.0, "type": "quad"} + self.__quadrupole = {"name": "", "l": 0.0, "k1": 0.0, "type": "quadrupole"} # don't count name and type --> len - 2 self.__nQuad = 2 * (len(self.__quadrupole) - 2) @@ -371,7 +371,7 @@ def __str__(self): nTypes[0] += 1 elif e["type"] == "sbend": nTypes[1] += 1 - elif e["type"] == "quad": + elif e["type"] == "quadrupole": nTypes[2] += 1 elif e["type"] == "dipedge": nTypes[3] += 1 @@ -428,7 +428,7 @@ def getBeamline(self): elif e["type"] == "sbend": print("Sbend L= ", e["l"], " angle = ", e["angle"]) beamline.append(e) - elif e["type"] == "quad": + elif e["type"] == "quadrupole": print("Quadrupole L= ", e["l"], " k1 = ", e["k1"]) beamline.append(e) elif e["type"] == "dipedge": diff --git a/src/python/impactx/madx_to_impactx.py b/src/python/impactx/madx_to_impactx.py index 61669ae27..27088094e 100644 --- a/src/python/impactx/madx_to_impactx.py +++ b/src/python/impactx/madx_to_impactx.py @@ -27,7 +27,7 @@ def lattice(parsed_beamline, nslice=1): "MARKER": "None", "DRIFT": "Drift", "SBEND": "Sbend", # Sector Bending Magnet - "QUAD": "Quad", # Quadrupole + "QUADRUPOLE": "Quad", # Quadrupole "DIPEDGE": "DipEdge", "MULTIPOLE": "Multipole", "NLLENS": "NonlinearLens", @@ -41,17 +41,18 @@ def lattice(parsed_beamline, nslice=1): for d in parsed_beamline: if d["type"] in [k.casefold() for k in list(madx_to_impactx_dict.keys())]: - if d["name"] == "drift": + + if d["type"] == "drift": impactx_beamline.append(elements.Drift(ds=d["l"], nslice=nslice)) - elif d["name"] == "quadrupole": + elif d["type"] == "quadrupole": impactx_beamline.append( elements.Quad(ds=d["l"], k=d["k1"], nslice=nslice) ) - elif d["name"] == "sbend": + elif d["type"] == "sbend": impactx_beamline.append( elements.Sbend(ds=d["l"], rc=d["l"] / d["angle"], nslice=nslice) ) - elif d["name"] == "dipedge": + elif d["type"] == "dipedge": impactx_beamline.append( elements.DipEdge( psi=d["e1"], @@ -152,7 +153,7 @@ def beam(particle, charge=None, mass=None, energy=None): def read_beam(ref: RefPart, madx_file): """ Function that reads elements from a MAD-X file into a list of ImpactX.KnownElements - :param RefPart ref: ImpactX reference particle (pybind_11 object) + :param RefPart ref: ImpactX reference particle (passed by reference) :param madx_file: file name to MAD-X file with beamline elements :return: list of ImpactX.KnownElements """ From ccb7cd7b289efb63627a6de0f0aab7345dc1cccd Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Wed, 31 Aug 2022 18:46:30 -0700 Subject: [PATCH 16/19] Update Copyright MAD-X Parser Script Confirmed today via email :) --- src/python/impactx/MADXParser.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/src/python/impactx/MADXParser.py b/src/python/impactx/MADXParser.py index 4095b6333..c460d8fd2 100644 --- a/src/python/impactx/MADXParser.py +++ b/src/python/impactx/MADXParser.py @@ -1,11 +1,14 @@ -# Author: Matthias Frey, Andreas Adelmann -# Date: 13. Oct. 2017 +#!/usr/bin/env python3 # -# This file parses MAD-X file into a beamline -# composed of pyAcceLEGOrator elements. +# Copyright 2022 ImpactX contributors +# Authors: Matthias Frey, Andreas Adelmann, Marco Garten +# License: BSD-3-Clause # -# 10.8.2022 Adapted for standalone use +# Changelog: +# Oct 13th, 2017: original version for pyAcceLEGOrator elements +# Aug 10th, 2022: adapted for standalone use # +# -*- coding: utf-8 -*- import os import re import warnings From 191be5d0f521676d7b2883c91cfbc7b0f13afd70 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Thu, 1 Sep 2022 07:08:28 -0700 Subject: [PATCH 17/19] Fix: DipEdge is a thin element --- src/python/impactx/madx_to_impactx.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/python/impactx/madx_to_impactx.py b/src/python/impactx/madx_to_impactx.py index 27088094e..0275b44ff 100644 --- a/src/python/impactx/madx_to_impactx.py +++ b/src/python/impactx/madx_to_impactx.py @@ -60,7 +60,6 @@ def lattice(parsed_beamline, nslice=1): # MAD-X is using half the gap height g=2.0 * d["hgap"], K2=d["fint"], - nslice=nslice, ) ) else: From 370ec24f232626ba394e5b9a773abf5d0a211442 Mon Sep 17 00:00:00 2001 From: Marco Garten Date: Thu, 1 Sep 2022 15:15:26 -0700 Subject: [PATCH 18/19] Fix Chicane input to make test pass also fix energy conversion in madx_to_impactx --- examples/chicane/chicane.madx | 12 ++++++------ examples/chicane/run_chicane_madx.py | 1 - src/python/impactx/madx_to_impactx.py | 4 +++- 3 files changed, 9 insertions(+), 8 deletions(-) diff --git a/examples/chicane/chicane.madx b/examples/chicane/chicane.madx index 418bf1803..ec60c591e 100644 --- a/examples/chicane/chicane.madx +++ b/examples/chicane/chicane.madx @@ -6,12 +6,12 @@ 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 theta again everywhere after `angle=` with its right sign -SB1: sbend, L=0.50037, angle=0.048345, e1=0.000, e2=0.000, k1=0.00; -SB2: sbend, L=0.50037, angle=-0.048345, e1=0.000, e2=0.000, k1=0.00; +! 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.096618357, e1=-0.048345620280243, fint=0.000, hgap=0.000, tilt=0.00; -DIPE2: dipedge, H=0.096618357, e1=0.048345620280243, fint=0.000, hgap=0.000, tilt=0.00; +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,DIP1,D1,DIP2,SB2,D2,SB2,DIPE2,D1,DIPE1,SB1,D3); +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 index 4f8cb1008..ce34c5d2e 100755 --- a/examples/chicane/run_chicane_madx.py +++ b/examples/chicane/run_chicane_madx.py @@ -23,7 +23,6 @@ # load a 5 GeV electron beam with an initial # normalized transverse rms emittance of 1 um -energy_MeV = 5.0e3 # reference energy bunch_charge_C = 1.0e-9 # used with space charge npart = 10000 # number of macro particles diff --git a/src/python/impactx/madx_to_impactx.py b/src/python/impactx/madx_to_impactx.py index 0275b44ff..5faa10665 100644 --- a/src/python/impactx/madx_to_impactx.py +++ b/src/python/impactx/madx_to_impactx.py @@ -103,7 +103,7 @@ def beam(particle, charge=None, mass=None, energy=None): :return dict: dictionary containing particle and beam attributes in ImpactX units """ - GeV2MeV = 1.0e-3 + GeV2MeV = 1.0e3 kg2MeV = sc.c**2 / sc.electron_volt * 1.0e-6 muon_mass = sc.physical_constants["electron-muon mass ratio"][0] / sc.m_e if energy is None: @@ -169,3 +169,5 @@ def read_beam(ref: RefPart, madx_file): ref.set_charge_qe(ref_particle_dict["charge"]) ref.set_mass_MeV(ref_particle_dict["mass"]) ref.set_energy_MeV(ref_particle_dict["energy"]) + + return ref From a5fb7dd42007060212367ae82e0a2dd73065cf02 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Thu, 1 Sep 2022 18:18:51 -0700 Subject: [PATCH 19/19] Comment Debug Printing --- src/python/impactx/MADXParser.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/python/impactx/MADXParser.py b/src/python/impactx/MADXParser.py index c460d8fd2..250cc6232 100644 --- a/src/python/impactx/MADXParser.py +++ b/src/python/impactx/MADXParser.py @@ -426,16 +426,16 @@ def getBeamline(self): for e in self.__elements: if elem == e["name"]: if e["type"] == "drift": - print("Drift L= " + str(e["l"])) + # print("Drift L= " + str(e["l"])) beamline.append(e) elif e["type"] == "sbend": - print("Sbend L= ", e["l"], " angle = ", e["angle"]) + # print("Sbend L= ", e["l"], " angle = ", e["angle"]) beamline.append(e) elif e["type"] == "quadrupole": - print("Quadrupole L= ", e["l"], " k1 = ", e["k1"]) + # 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"] + "'")