def test_zero_errors():
# check that zero-errors are loaded without error
cpymad_spec = util.find_spec("cpymad")
if cpymad_spec is None:
print("cpymad is not available - abort test")
sys.exit(0)
from cpymad.madx import Madx
madx = Madx()
madx.input('''
qd: multipole, knl={0,-0.3};
qf: multipole, knl={0, 0.3};
testseq: sequence, l = 1;
qd, at = 0.3;
qf, at = 0.6;
endsequence;
''')
madx.beam()
madx.use('testseq')
madx.select(flag='error', pattern='qf')
madx.command.efcomp(dkn=[0, 0, 0, 0, 0.0, 0.0, 0.0], dks=[0.0, 0.0, 0, 0])
madx.command.ealign(dx=0.0,
dy=0.0,
ds=0.0,
DPHI=0.0,
DTHETA=0.0,
DPSI=0.0,
MREX=0.0,
MREY=0.0,
MSCALX=0.0,
MSCALY=0.0,
AREX=0.0,
AREY=0.0)
def test_command_log(self):
"""Check that the command log contains all input commands."""
# create a new Madx instance that uses the history feature:
history_filename = '_test_madx.madx.tmp'
mad = Madx(command_log=history_filename)
try:
# feed some input lines and compare with history file:
lines = dedent("""
l = 5;
f = 200;
fodo: sequence, refer=entry, l=100;
QF: quadrupole, l=5, at= 0, k1= 1/(f*l);
QD: quadrupole, l=5, at=50, k1=-1/(f*l);
endsequence;
beam, particle=proton, energy=2;
use, sequence=fodo;
""").splitlines()
lines = [line for line in lines if line.strip()]
for line in lines:
mad.input(line)
with open(history_filename) as history_file:
history = history_file.read()
self.assertEqual(history.strip(), '\n'.join(lines).strip())
finally:
# remove history file
mad.quit()
del mad
os.remove(history_filename)
def test_independent_instances(self):
# create a second Madx instance (1st one is created in setUp)
madxness = Madx()
# Check independence by defining a variable differently in each
# instance:
self.mad.input('ANSWER=42;')
madxness.input('ANSWER=43;')
self.assertEqual(self.mad.evaluate('ANSWER'), 42)
self.assertEqual(madxness.evaluate('ANSWER'), 43)
def test_independent_instances(self):
# create a second Madx instance (1st one is created in setUp)
madxness = Madx()
# Check independence by defining a variable differently in each
# instance:
self.mad.input('ANSWER=42;')
madxness.input('ANSWER=43;')
self.assertEqual(self.mad.evaluate('ANSWER'), 42);
self.assertEqual(madxness.evaluate('ANSWER'), 43);
def df2run(myDF,
pythonData=None,
command_log_file='log.madx',
stdout_file='stdout.madx',
verbose=False):
'''
It runs the MADX dataframe using the MADX extended syntax.
myDF: the MADX DF to run.
command_log_file: the filename of the logging file. Use the None variable not to log.
stdout_file: the filename of the file to redirect the stdout. Use the None variable not to log.
verbose: boolean flag to have verbose output during the execution.
'''
if command_log_file == None:
if stdout_file == None:
madx = Madx()
else:
with open(stdout_file, 'w') as f:
madx = Madx(stdout=f)
else:
if stdout_file == None:
madx = Madx(command_log=command_log_file)
else:
with open(stdout_file, 'w') as f:
madx = Madx(stdout=f, command_log=command_log_file)
myGlobals = []
for section in myDF.iterrows():
print(section[0])
start_time = time.time()
codeSubSections = section[1]['Code subsections']
pythonDictionary = {}
with madx.batch():
myCheck = [list(code.keys())[0] for code in codeSubSections]
if verbose: print(codeSubSections)
if ('madx' in myCheck) & ('python' in myCheck):
raise Exception(
'Do not put madx and python code in the same section!')
for code in codeSubSections:
myType = list(code.keys())[0]
if myType == 'markdown':
pass
elif myType == 'python':
exec(code['python']) # local variables
elif myType == 'madx':
madx.input(code['madx'])
else:
assert (0)
execution_time_s = time.time() - start_time
myDict = {}
myDict = dict(madx.globals)
myDict['execution time [s]'] = execution_time_s
myDict['pythonDictionary'] = pythonDictionary
myDict['Code subsections'] = section[1]['Code subsections']
myDict['Code section'] = section[1]['Code section']
myGlobals.append(myDict)
profileDF = pd.DataFrame(myGlobals, index=myDF.index)
return profileDF
def test_command_log(self):
"""Check that the command log contains all input commands."""
# create a new Madx instance that uses the history feature:
history_filename = '_test_madx.madx.tmp'
mad = Madx(command_log=history_filename)
# feed some input and compare with history file:
for line in self.doc.splitlines():
mad.input(line)
with open(history_filename) as history_file:
history = history_file.read()
self.assertEqual(history.strip(), self.doc.strip())
# remove history file
del mad
os.remove(history_filename)
def test_append_semicolon(self):
"""Check that semicolon is automatically appended to input() text."""
# Regression test for #73
log = []
mad = Madx(command_log=log.append)
try:
mad.input('a = 0')
mad.input('b = 1')
self.assertEqual(log, ['a = 0;', 'b = 1;'])
self.assertEqual(mad.globals.a, 0)
self.assertEqual(mad.globals.b, 1)
finally:
mad.quit()
del mad
def test_command_log(self):
"""Check that the command log contains all input commands."""
# create a new Madx instance that uses the history feature:
history_filename = '_test_madx.madx.tmp'
mad = Madx(command_log=history_filename)
# feed some input and compare with history file:
for line in self.doc.splitlines():
mad.input(line)
with open(history_filename) as history_file:
history = history_file.read()
self.assertEqual(history.strip(), self.doc.strip())
# remove history file
del mad
os.remove(history_filename)
def test_streamreader(self):
output = []
m = Madx(stdout=output.append)
self.assertEqual(len(output), 1)
self.assertIn(b'++++++++++++++++++++++++++++++++++++++++++++', output[0])
self.assertIn(b'+ Support: [email protected],', output[0])
self.assertIn(b'+ Release date: ', output[0])
self.assertIn(b'+ Execution date: ', output[0])
# self.assertIn(b'+ Support: [email protected], ', output[1])
m.input('foo = 3;')
self.assertEqual(len(output), 1)
m.input('foo = 3;')
self.assertEqual(len(output), 2)
self.assertEqual(output[1], b'++++++ info: foo redefined\n')
m.quit()
self.assertEqual(len(output), 3)
self.assertIn(b'+ MAD-X finished normally ', output[2])
def test_streamreader(self):
output = []
m = Madx(stdout=output.append)
self.assertEqual(len(output), 1)
self.assertIn(b'++++++++++++++++++++++++++++++++++++++++++++', output[0])
self.assertIn(b'+ Support: [email protected],', output[0])
self.assertIn(b'+ Release date: ', output[0])
self.assertIn(b'+ Execution date: ', output[0])
# self.assertIn(b'+ Support: [email protected], ', output[1])
m.input('foo = 3;')
self.assertEqual(len(output), 1)
m.input('foo = 3;')
self.assertEqual(len(output), 2)
self.assertEqual(output[1], b'++++++ info: foo redefined\n')
m.quit()
self.assertEqual(len(output), 3)
self.assertIn(b'+ MAD-X finished normally ', output[2])
def mad_benchmark(mtype, attrs, pc=0.2, x=0, px=0, y=0, py=0, t=0, pt=0):
import pysixtrack
from cpymad.madx import Madx
mad = Madx(stdout=False)
madtype = mad.command[mtype]
mad.beam(particle="proton", pc=pc)
madtype.clone("mm", **attrs)
mad.input("bench: line=(mm)")
mad.use(sequence="bench")
mad.track(onepass=True, dump=False)
mad.start(x=x, px=px, y=y, py=py, t=t, pt=pt)
mad.run()
mad.endtrack()
p_mad = pysixtrack.Particles.from_madx_track(mad)
p_six = p_mad.copy(0)
line = pysixtrack.Line.from_madx_sequence(mad.sequence.bench,
exact_drift=True)
line.track(p_six)
p_mad.copy(-1).compare(p_six, rel_tol=0)
return mad, line, p_mad, p_six
def match_no_coupling_through_ripkens(
madx: Madx,
sequence: str = None,
location: str = None,
vary_knobs: Sequence[str] = None) -> None:
"""
Matching routine to get cross-term Ripken parameters beta_12 and beta_21 to be 0 at a given location.
Args:
madx (cpymad.madx.Madx): an instanciated cpymad Madx object.
sequence (str): name of the sequence to activate for the matching.
location (str): the name of the element at which one wants the cross-term Ripkens to be 0.
vary_knobs (Sequence[str]): the variables names to 'vary' in the MADX routine.
"""
logger.info(
f"Matching Ripken parameters for no coupling at location {location}")
logger.debug("Creating macro tu update Ripkens")
madx.input("do_ripken: macro = {twiss, ripken=True;}"
) # cpymad needs .input for macros
logger.debug("Matching Parameters")
madx.command.match(sequence=sequence, use_macro=True, chrom=True)
for knob in vary_knobs:
madx.command.vary(name=knob)
madx.command.use_macro(name="do_ripken")
madx.input(f"constraint, expr=table(twiss, {location}, beta12)=0"
) # need input else includes " and fails
madx.input(f"constraint, expr=table(twiss, {location}, beta21)=0"
) # need input else includes " and fails
madx.command.lmdif(calls=500, tolerance=1e-21)
madx.command.endmatch()
from pysixtrack.particles import Particles
import pysixtrack.be_beamfields.tools as bt
import simulation_parameters as pp
os.makedirs(pp.input_dir, exist_ok=True)
mad = Madx()
mad.options.echo = False
mad.options.info = False
mad.warn = False
mad.chdir('madx')
mad.call('sps_thin_crabcavity.madx')
for parameter in pp.madx_settings:
setting = pp.madx_settings[parameter]
mad.input(f'{parameter} = {setting};')
mad.use(pp.seq_name)
# Include b3b5b7 in MBA and MBB
mad.call('./sps/cmd/sps_setMultipoles_upto7.cmd')
mad.input('exec, set_Multipoles_270GeV;')
mad.call('./sps/cmd/sps_assignMultipoles_upto7.cmd')
mad.input('exec, AssignMultipoles;')
mad.command.readtable(file='err.out', table='errors')
errors = mad.table.errors
# Tune and Chromaticity matching
mad.call('./sps/cmd/sps_matching.cmd')
mad.input('exec, SPS_matchtunes(QH, QV);')
# %% Run MADX
madx = Madx()
# %%
import time
myGlobals = []
mySmallDF = myDF
myString = ''
for block in mySmallDF.iterrows():
print(block[0])
start_time = time.time()
myString = myString + '! %%' + block[0] + '\n' + block[1]['Code string'][
0:-1]
with madx.batch():
madx.input('! %%' + block[0] + '\n' + block[1]['Code string'][0:-1])
execution_time_s = time.time() - start_time
myDict = {}
myDict = dict(madx.globals)
myDict['execution time [s]'] = execution_time_s
myGlobals.append(myDict)
profileDF = pd.DataFrame(myGlobals, index=mySmallDF.index)
# %%
with open('inverse.mask', 'w') as fid:
fid.write(myString)
# %%
# %%
madx.input(aux[2])
# %%
# MAD-X parameters dictionary
madx_settings = {'QH': 26.13, 'QV': 26.18, 'QPH': 0.0, 'QPV': 0.0}
seq_name = 'sps'
harmonic_number = 4620
mad = Madx()
mad.options.echo = False
mad.options.info = False
mad.warn = False
mad.chdir('./madx')
mad.call('sps_thin_crabcavity.madx')
for parameter in madx_settings:
setting = madx_settings[parameter]
mad.input(f'{parameter} = {setting};')
mad.use(seq_name)
# Include b3b5b7 in MBA and MBB
mad.call('./sps/cmd/sps_setMultipoles_upto7.cmd')
mad.input('exec, set_Multipoles_26GeV;')
mad.call('./sps/cmd/sps_assignMultipoles_upto7.cmd')
mad.input('exec, AssignMultipoles;')
# Tune matching. After the latest update the names of the quadrupoles changed. Therfore, to performe the tune matching correctly the following macro from the toolkit repository is needed.
# https://gitlab.cern.ch/acc-models/acc-models-sps/-/tree/2021/
mad.call('./sps/toolkit/macro.madx')
mad.input('exec, sps_match_tunes(QH,QV);')
#mad.call('./sps/cmd/sps_matching.cmd')
#mad.input('exec, SPS_matchtunes(QH, QV);')
class TestMadx(unittest.TestCase, _TestCaseCompat):
"""
Test methods for the Madx class.
The tests are directly based on the specifics of the sequence in
test/testseq.madx
Please compare this file for reference.
"""
def setUp(self):
self.mad = Madx(command_log=CommandLog(sys.stdout, 'X:> '))
here = os.path.dirname(__file__)
there = os.path.join(here, 'testseq.madx')
self.doc = open(there).read()
for line in self.doc.splitlines():
self.mad._libmadx.input(line)
def tearDown(self):
del self.mad
def test_version(self):
"""Check that the Madx.version attribute can be used as expected."""
version = self.mad.version
# check format:
major, minor, mini = map(int, version.release.split('.'))
# We need at least MAD-X 5.02.03:
self.assertGreaterEqual(major, 5)
self.assertGreaterEqual(minor, 2)
self.assertGreaterEqual(mini, 3)
# check format:
year, month, day = map(int, version.date.split('.'))
self.assertGreaterEqual(year, 2014)
self.assertGreaterEqual(month, 1)
self.assertGreaterEqual(day, 1)
self.assertLessEqual(month, 12)
self.assertLessEqual(day, 31)
def test_independent_instances(self):
# create a second Madx instance (1st one is created in setUp)
madxness = Madx()
# Check independence by defining a variable differently in each
# instance:
self.mad.input('ANSWER=42;')
madxness.input('ANSWER=43;')
self.assertEqual(self.mad.evaluate('ANSWER'), 42)
self.assertEqual(madxness.evaluate('ANSWER'), 43)
def test_command_log(self):
"""Check that the command log contains all input commands."""
# create a new Madx instance that uses the history feature:
history_filename = '_test_madx.madx.tmp'
mad = Madx(command_log=history_filename)
# feed some input and compare with history file:
for line in self.doc.splitlines():
mad.input(line)
with open(history_filename) as history_file:
history = history_file.read()
self.assertEqual(history.strip(), self.doc.strip())
# remove history file
del mad
os.remove(history_filename)
# TODO:
# def test_command(self):
# def test_help(self):
# def test_call(self):
def _check_twiss(self, seq_name):
beam = 'beam, ex=1, ey=2, particle=electron, sequence={0};'.format(
seq_name)
self.mad.command(beam)
initial = dict(alfx=0.5, alfy=1.5, betx=2.5, bety=3.5)
# by explicitly specifying the 'columns' parameter a persistent copy
# is returned. We check that this copy contains the data we want and
# that it has a 'summary' attribute:
twiss = self.mad.twiss(sequence=seq_name,
columns=['betx', 'bety', 'alfx', 'alfy'],
**initial)
betx, bety = twiss['betx'], twiss['bety']
alfx, alfy = twiss['alfx'], twiss['alfy']
# Check initial values:
self.assertAlmostEqual(twiss['alfx'][0], initial['alfx'])
self.assertAlmostEqual(twiss['alfy'][0], initial['alfy'])
self.assertAlmostEqual(twiss['betx'][0], initial['betx'])
self.assertAlmostEqual(twiss['bety'][0], initial['bety'])
self.assertAlmostEqual(twiss.summary['ex'], 1)
self.assertAlmostEqual(twiss.summary['ey'], 2)
# Check that keys are all lowercase:
for k in twiss:
self.assertEqual(k, k.lower())
for k in twiss.summary:
self.assertEqual(k, k.lower())
def test_twiss_1(self):
self._check_twiss('s1') # s1 can be computed at start
self._check_twiss('s1') # s1 can be computed multiple times
self._check_twiss('s2') # s2 can be computed after s1
def test_twiss_2(self):
self._check_twiss('s2') # s2 can be computed at start
self._check_twiss('s1') # s1 can be computed after s2
def test_twiss_with_range(self):
beam = 'beam, ex=1, ey=2, particle=electron, sequence=s1;'
self.mad.command(beam)
params = dict(alfx=0.5,
alfy=1.5,
betx=2.5,
bety=3.5,
columns=['betx', 'bety'],
sequence='s1')
# Compute TWISS on full sequence, then on a sub-range, then again on
# the full sequence. This checks that none of the range selections
# have side-effects on each other:
betx_full1 = self.mad.twiss(**params)['betx']
betx_range = self.mad.twiss(range=('dr[2]', 'sb'), **params)['betx']
betx_full2 = self.mad.twiss(**params)['betx']
# Check that the results have the expected lengths:
self.assertEqual(len(betx_full1), 9)
self.assertEqual(len(betx_range), 4)
self.assertEqual(len(betx_full2), 9)
# Check numeric results. Since the first 3 elements of range and full
# sequence are identical, equal results are expected. And non-equal
# results afterwards.
self.assertAlmostEqual(betx_range[0], betx_full1[1]) # dr:2, dr:1
self.assertAlmostEqual(betx_range[1], betx_full1[2]) # qp:2, qp:1
self.assertAlmostEqual(betx_range[2], betx_full1[3]) # dr:3, dr:2
self.assertNotAlmostEqual(betx_range[3], betx_full1[4]) # sb, qp:2
def test_range_row_api(self):
beam = 'beam, ex=1, ey=2, particle=electron, sequence=s1;'
self.mad.command(beam)
params = dict(alfx=0.5, alfy=1.5, betx=2.5, bety=3.5, sequence='s1')
tab = self.mad.twiss(range=('dr[2]', 'sb'), **params)
self.assertEqual(tab.range, ('dr[2]', 'sb'))
self.assertIn('betx', tab)
# def test_survey(self):
# def test_aperture(self):
# def test_use(self):
# def test_match(self):
# def test_verbose(self):
def test_active_sequence(self):
self.mad.command('beam, ex=1, ey=2, particle=electron, sequence=s1;')
self.mad.active_sequence = 's1'
self.assertEqual(self.mad.active_sequence.name, 's1')
def test_get_sequence(self):
with self.assertRaises(KeyError):
self.mad.sequences['sN']
s1 = self.mad.sequences['s1']
self.assertEqual(s1.name, 's1')
def test_get_sequences(self):
seqs = self.mad.sequences
self.assertItemsEqual(seqs, ['s1', 's2'])
def test_evaluate(self):
val = self.mad.evaluate("1/QP_K1")
self.assertAlmostEqual(val, 0.5)
def test_set_value(self):
self.mad.set_value('FOO', 1)
self.mad.set_value('BAR', 'FOO')
self.mad.set_value('FOO', 2)
self.assertEqual(self.mad.evaluate('FOO'), 2)
self.assertEqual(self.mad.evaluate('BAR'), 1)
def test_set_expression(self):
self.mad.set_expression('FOO', 'BAR')
self.mad.set_value('BAR', 42)
self.mad.evaluate('FOO')
self.assertEqual(self.mad.evaluate('FOO'), 42)
self.mad.set_value('BAR', 43)
self.assertEqual(self.mad.evaluate('FOO'), 43)
def test_globals(self):
g = self.mad.globals
# Membership:
self.assertNotIn('FOO', g)
# Setting values:
g['FOO'] = 2
self.assertIn('FOO', g)
self.assertEqual(g['FOO'], 2)
self.assertEqual(self.mad.evaluate('FOO'), 2)
# Re-setting values:
g['FOO'] = 3
self.assertEqual(self.mad.evaluate('FOO'), 3)
# Setting expressions:
g['BAR'] = '3*foo'
self.assertEqual(self.mad.evaluate('BAR'), 9)
g['FOO'] = 4
self.assertEqual(self.mad.evaluate('BAR'), 12)
def test_elements(self):
self.assertIn('sb', self.mad.elements)
self.assertIn('sb', list(self.mad.elements))
self.assertNotIn('foobar', self.mad.elements)
self.assertAlmostEqual(self.mad.elements['sb']['angle'], 3.14 / 4)
idx = self.mad.elements.index('qp1')
elem = self.mad.elements[idx]
self.assertEqual(elem['k1'], 3)
# def test_sequence_beam(self):
# def test_sequence_twiss(self):
# def test_sequence_twissname(self):
def _get_elems(self, seq_name):
elems = self.mad.sequences[seq_name].elements
elem_idx = dict((el['name'], i) for i, el in enumerate(elems))
return elems, elem_idx
def test_sequence_get_elements_s1(self):
s1, idx = self._get_elems('s1')
qp1 = s1['qp[1]']
qp2 = s1['qp[2]']
sb1 = s1['sb[1]']
self.assertLess(idx['qp'], idx['qp[2]'])
self.assertLess(idx['qp[2]'], idx['sb'])
self.assertAlmostEqual(qp1['at'], 1)
self.assertAlmostEqual(qp2['at'], 3)
self.assertAlmostEqual(sb1['at'], 5)
self.assertAlmostEqual(qp1['l'], 1)
self.assertAlmostEqual(qp2['l'], 1)
self.assertAlmostEqual(sb1['l'], 2)
self.assertAlmostEqual(float(qp1['k1']), 2)
self.assertAlmostEqual(float(qp2['k1']), 2)
self.assertAlmostEqual(float(sb1['angle']), 3.14 / 4)
self.assertEqual(str(qp1['k1']).lower(), "qp_k1")
def test_sequence_get_elements_s2(self):
s2, idx = self._get_elems('s2')
qp1 = s2['qp1[1]']
qp2 = s2['qp2[1]']
self.assertLess(idx['qp1'], idx['qp2'])
self.assertAlmostEqual(qp1['at'], 0)
self.assertAlmostEqual(qp2['at'], 1)
self.assertAlmostEqual(qp1['l'], 1)
self.assertAlmostEqual(qp2['l'], 2)
self.assertAlmostEqual(float(qp1['k1']), 3)
self.assertAlmostEqual(float(qp2['k1']), 2)
# def test_sequence_get_expanded_elements(self):
def test_crash(self):
"""Check that a RuntimeError is raised in case MAD-X crashes."""
# a.t.m. MAD-X crashes on this input, because the L (length)
# parametere is missing:
self.assertRaises(RuntimeError, self.mad.input, 'XXX: sequence;')
def test_sequence_elements(self):
elements = self.mad.sequences['s1'].elements
iqp2 = elements.index('qp[2]')
qp1 = elements['qp[1]']
qp2 = elements[iqp2]
self.assertAlmostEqual(qp1['at'], 1)
self.assertAlmostEqual(qp2['at'], 3)
self.assertEqual(iqp2, elements.at(3.1))
def test_sequence_expanded_elements(self):
beam = 'beam, ex=1, ey=2, particle=electron, sequence=s1;'
self.mad.command(beam)
self.mad.use('s1')
elements = self.mad.sequences['s1'].expanded_elements
iqp2 = elements.index('qp[2]')
qp1 = elements['qp[1]']
qp2 = elements[iqp2]
self.assertAlmostEqual(qp1['at'], 1)
self.assertAlmostEqual(qp2['at'], 3)
self.assertEqual(iqp2, elements.at(3.1))
# MAD-X parameters dictionary
madx_settings = {'QH': 26.13, 'QV': 26.18, 'QPH': 0.0, 'QPV': 0.0}
seq_name = 'sps'
harmonic_number = 4620
mad = Madx()
mad.options.echo = False
mad.options.info = False
mad.warn = False
mad.chdir('./madx')
mad.call('sps_thin_crabcavity.madx')
for parameter in madx_settings:
setting = madx_settings[parameter]
mad.input(f'{parameter} = {setting};')
mad.use(seq_name)
# Include b3b5b7 in MBA and MBB
mad.call('./sps/cmd/sps_setMultipoles_upto7.cmd')
mad.input('exec, set_Multipoles_26GeV;')
mad.call('./sps/cmd/sps_assignMultipoles_upto7.cmd')
mad.input('exec, AssignMultipoles;')
# Tune matching. After the latest update the names of the quadrupoles changed. Therfore, to performe the tune matching correctly the following macro from the toolkit repository is needed.
# https://gitlab.cern.ch/acc-models/acc-models-sps/-/tree/2021/
mad.call('./sps/toolkit/macro.madx')
mad.input('exec, sps_match_tunes(QH,QV);')
#mad.call('./sps/cmd/sps_matching.cmd')
#mad.input('exec, SPS_matchtunes(QH, QV);')
def test_error_import():
cpymad_spec = util.find_spec("cpymad")
if cpymad_spec is None:
print("cpymad is not available - abort test")
sys.exit(0)
from cpymad.madx import Madx
madx = Madx()
madx.input('''
MQ1: Quadrupole, K1:=KQ1, L=1.0, apertype=CIRCLE, aperture={0.04};
MQ2: Quadrupole, K1:=KQ2, L=1.0, apertype=CIRCLE, aperture={0.04};
MQ3: Quadrupole, K1:=0.0, L=1.0, apertype=CIRCLE, aperture={0.04};
KQ1 = 0.02;
KQ2 = -0.02;
testseq: SEQUENCE, l = 20.0;
MQ1, at = 5;
MQ2, at = 12;
MQ3, at = 18;
ENDSEQUENCE;
!---the usual stuff
BEAM, PARTICLE=PROTON, ENERGY=7000.0, EXN=2.2e-6, EYN=2.2e-6;
USE, SEQUENCE=testseq;
Select, flag=makethin, pattern="MQ1", slice=2;
makethin, sequence=testseq;
use, sequence=testseq;
!---assign misalignments and field errors
select, flag = error, clear;
select, flag = error, pattern = "MQ1";
ealign, dx = 0.01, dy = 0.01, arex = 0.02, arey = 0.02;
select, flag = error, clear;
select, flag = error, pattern = "MQ2";
ealign, dx = 0.04, dy = 0.04, dpsi = 0.1;
select, flag = error, clear;
select, flag = error, pattern = "MQ3";
ealign, dx = 0.00, dy = 0.00, arex = 0.00, arey = 0.00, dpsi = 0.00;
efcomp, DKN = {0.0, 0.0, 0.001, 0.002}, DKS = {0.0, 0.0, 0.003, 0.004, 0.005};
select, flag = error, full;
''')
seq = madx.sequence.testseq
pysixtrack_line = pysixtrack.Line.from_madx_sequence(
seq,
install_apertures=True,
apply_madx_errors=True,
)
madx.input('stop;')
expected_element_num = (
2 # start and end marker
+ 6 # drifts (including drift between MQ1 slices)
+ 3 + 2 # quadrupoles + MQ1 slices
+ 3 + 2 # corresponding aperture elements
+ 2*(3+1) # dx/y in/out for MQ1 slices and MQ2
+ 2 # tilt in/out for MQ2
+ 2*3 # arex/y in/out for MQ1 slices
)
assert len(pysixtrack_line) == expected_element_num
expected_element_order = [
pysixtrack.elements.Drift, # start marker
pysixtrack.elements.Drift,
pysixtrack.elements.XYShift, # dx/y in of MQ1 1st slice
pysixtrack.elements.Multipole, # MQ1 1st slice
pysixtrack.elements.XYShift, # arex/y in for MQ1 1st slice
pysixtrack.elements.LimitEllipse, # MQ1 1st slice aperture
pysixtrack.elements.XYShift, # arex/y out for MQ1 1st slice
pysixtrack.elements.XYShift, # dx/y out for MQ1 1st slice
pysixtrack.elements.Drift,
pysixtrack.elements.XYShift, # dx/y in for MQ1 marker
pysixtrack.elements.Drift, # MQ1 marker
pysixtrack.elements.XYShift, # arex/y in for MQ1 marker
pysixtrack.elements.LimitEllipse, # MQ1 marker aperture
pysixtrack.elements.XYShift, # arex/y out for MQ1 marker
pysixtrack.elements.XYShift, # dx/y out for MQ1 marker
pysixtrack.elements.Drift,
pysixtrack.elements.XYShift, # dx/y in for MQ1 2nd slice
pysixtrack.elements.Multipole, # MQ1 2nd slice
pysixtrack.elements.XYShift, # arex/y in for MQ1 2nd slice
pysixtrack.elements.LimitEllipse, # MQ1 2nd slice aperture
pysixtrack.elements.XYShift, # arex/y out for MQ1 2nd slice
pysixtrack.elements.XYShift, # dx/y out for MQ1 2nd slice
pysixtrack.elements.Drift,
pysixtrack.elements.XYShift, # dx/y in for MQ2
pysixtrack.elements.SRotation, # tilt in for MQ2
pysixtrack.elements.Multipole, # MQ2
pysixtrack.elements.LimitEllipse, # MQ2 aperture
pysixtrack.elements.SRotation, # tilt out for MQ2
pysixtrack.elements.XYShift, # dx/y out for MQ2
pysixtrack.elements.Drift,
pysixtrack.elements.Multipole, # MQ3
pysixtrack.elements.LimitEllipse, # MQ3 aperture
pysixtrack.elements.Drift,
pysixtrack.elements.Drift # end marker
]
for element, expected_element in zip(pysixtrack_line.elements,
expected_element_order):
assert isinstance(element, expected_element)
idx_MQ3 = pysixtrack_line.element_names.index('mq3')
MQ3 = pysixtrack_line.elements[idx_MQ3]
assert abs(MQ3.knl[2] - 0.001) < 1e-14
assert abs(MQ3.knl[3] - 0.002) < 1e-14
assert abs(MQ3.ksl[2] - 0.003) < 1e-14
assert abs(MQ3.ksl[3] - 0.004) < 1e-14
assert abs(MQ3.ksl[4] - 0.005) < 1e-14
sequence += 'mb_arc%dl_cell%d_%d: mb, at=%e;\n' % (
i_arc, i_cell, i_bend + n_dip_half_cell, s_start_cell +
(i_bend + 1) * L_halfcell / (n_dip_half_cell + 1) + L_halfcell)
s_start_cell += L_halfcell * 2
sequence += '''
end_machine: marker at=circum;
endsequence;
beam, particle = proton, sequence=toyring, energy = 6500., NPART=1.05E11, sige= 2.5e;
'''
with open('sequence.seq', 'w') as fid:
fid.write(sequence)
mad.input(sequence)
mad.use('toyring')
str_factor_list = np.linspace(.7, 2, 50)
plt.close('all')
fig1 = plt.figure()
ax1 = fig1.add_subplot(111)
betax_max = np.zeros((len(str_factor_list), len(str_factor_list)))
betax_min = np.zeros((len(str_factor_list), len(str_factor_list)))
for ii1, ff1 in enumerate(str_factor_list):
mad.input('kqf := %e' % (ff1 * k1l_quad))
for ii2, ff2 in enumerate(str_factor_list):
mad.input('kqd := %e' % (-ff2 * k1l_quad))
def build_mad_instance_with_bb(sequences_file_name, bb_data_frames,
beam_names, sequence_names,
mad_echo=False, mad_warn=False, mad_info=False):
mad = Madx()
mad.options.echo = mad_echo
mad.options.warn = mad_warn
mad.options.info = mad_info
mad.call(sequences_file_name)# assuming a sequence rotated in IR3
for bb_df in bb_data_frames:
mad.input(bb_df['elementDefinition'].str.cat(sep='\n'))
# %% seqedit
for beam, bb_df, seq in zip(beam_names, bb_data_frames, sequence_names):
myBBDFFiltered=bb_df[bb_df['beam']==beam]
mad.input(f'seqedit, sequence={"lhc"+beam};')
mad.input('flatten;')
mad.input(myBBDFFiltered['elementInstallation'].str.cat(sep='\n'))
mad.input('flatten;')
mad.input(f'endedit;')
return mad
class TestMadx(unittest.TestCase, _TestCaseCompat):
"""
Test methods for the Madx class.
The tests are directly based on the specifics of the sequence in
test/testseq.madx
Please compare this file for reference.
"""
def setUp(self):
self.mad = Madx(command_log=CommandLog(sys.stdout, 'X:> '))
here = os.path.dirname(__file__)
there = os.path.join(here, 'testseq.madx')
with open(there) as f:
self.doc = f.read()
for line in self.doc.splitlines():
line = line.split('!')[0].strip()
if line:
self.mad._libmadx.input(line)
def tearDown(self):
self.mad.quit()
del self.mad
def test_copyright(self):
import cpymad
notice = cpymad.get_copyright_notice()
self.assertIsInstance(notice, type(u""))
def test_version(self):
"""Check that the Madx.version attribute can be used as expected."""
version = self.mad.version
# check format:
major, minor, micro = map(int, version.release.split('.'))
# We need at least MAD-X 5.04.02:
self.assertGreaterEqual((major, minor, micro), (5, 4, 2))
# check format:
year, month, day = map(int, version.date.split('.'))
self.assertGreaterEqual((year, month, day), (2018, 10, 3))
self.assertLessEqual(month, 12)
self.assertLessEqual(day, 31)
self.assertTrue(str(version).startswith(
'MAD-X {}'.format(version.release)))
def test_metadata(self):
version = self.mad.version
self.assertEqual(metadata.__version__, version.release)
self.assertIsInstance(metadata.get_copyright_notice(), type(u""))
def test_independent_instances(self):
# create a second Madx instance (1st one is created in setUp)
madxness = Madx()
# Check independence by defining a variable differently in each
# instance:
self.mad.input('ANSWER=42;')
madxness.input('ANSWER=43;')
self.assertEqual(self.mad.eval('ANSWER'), 42)
self.assertEqual(madxness.eval('ANSWER'), 43)
madxness.quit()
def test_streamreader(self):
output = []
m = Madx(stdout=output.append)
self.assertEqual(len(output), 1)
self.assertIn(b'++++++++++++++++++++++++++++++++++++++++++++', output[0])
self.assertIn(b'+ Support: [email protected],', output[0])
self.assertIn(b'+ Release date: ', output[0])
self.assertIn(b'+ Execution date: ', output[0])
# self.assertIn(b'+ Support: [email protected], ', output[1])
m.input('foo = 3;')
self.assertEqual(len(output), 1)
m.input('foo = 3;')
self.assertEqual(len(output), 2)
self.assertEqual(output[1], b'++++++ info: foo redefined\n')
m.quit()
self.assertEqual(len(output), 3)
self.assertIn(b'+ MAD-X finished normally ', output[2])
def test_quit(self):
self.mad.quit()
self.assertIsNot(self.mad._process.returncode, None)
self.assertFalse(bool(self.mad))
with self.assertRaises(RuntimeError):
self.mad.input(';')
def test_context_manager(self):
output = []
with Madx(stdout=output.append) as m:
m.input('foo = 3;')
self.assertEqual(m.globals.foo, 3)
self.assertIn(b'+ MAD-X finished normally ', output[-1])
self.assertFalse(bool(m))
with self.assertRaises(RuntimeError):
m.input(';')
def test_command_log(self):
"""Check that the command log contains all input commands."""
# create a new Madx instance that uses the history feature:
history_filename = '_test_madx.madx.tmp'
mad = Madx(command_log=history_filename)
# feed some input and compare with history file:
for line in self.doc.splitlines():
mad.input(line)
with open(history_filename) as history_file:
history = history_file.read()
self.assertEqual(history.strip(), self.doc.strip())
# remove history file
mad.quit()
del mad
os.remove(history_filename)
def test_call_and_chdir(self):
folder = os.path.abspath(os.path.dirname(__file__))
parent = os.path.dirname(folder)
getcwd = self.mad._libmadx.getcwd
g = self.mad.globals
self.mad.chdir(folder)
self.assertEqual(getcwd(), folder)
self.mad.call('answer_42.madx')
self.assertEqual(g.answer, 42)
with self.mad.chdir('..'):
self.assertEqual(getcwd(), parent)
self.mad.call('test/answer_43.madx')
self.assertEqual(g.answer, 43)
self.mad.call('test/answer_call42.madx', True)
self.assertEqual(g.answer, 42)
self.assertEqual(getcwd(), folder)
self.mad.call('answer_43.madx')
self.assertEqual(g.answer, 43)
self.mad.chdir('..')
self.assertEqual(getcwd(), parent)
def _check_twiss(self, seq_name):
beam = 'ex=1, ey=2, particle=electron, sequence={0};'.format(seq_name)
self.mad.command.beam(beam)
self.mad.use(seq_name)
initial = dict(alfx=0.5, alfy=1.5,
betx=2.5, bety=3.5)
twiss = self.mad.twiss(sequence=seq_name, **initial)
# Check initial values:
self.assertAlmostEqual(twiss['alfx'][0], initial['alfx'])
self.assertAlmostEqual(twiss['alfy'][0], initial['alfy'])
self.assertAlmostEqual(twiss['betx'][0], initial['betx'])
self.assertAlmostEqual(twiss['bety'][0], initial['bety'])
self.assertAlmostEqual(twiss.summary['ex'], 1)
self.assertAlmostEqual(twiss.summary['ey'], 2)
# Check that keys are all lowercase:
for k in twiss:
self.assertEqual(k, k.lower())
for k in twiss.summary:
self.assertEqual(k, k.lower())
def test_error(self):
self.mad.input("""
seq: sequence, l=1;
endsequence;
beam;
use, sequence=seq;
""")
# Errors in MAD-X must not crash, but return False instead:
self.assertFalse(self.mad.input('twiss;'))
self.assertTrue(self.mad.input('twiss, betx=1, bety=1;'))
def test_twiss_1(self):
self._check_twiss('s1') # s1 can be computed at start
self._check_twiss('s1') # s1 can be computed multiple times
self._check_twiss('s2') # s2 can be computed after s1
def test_twiss_2(self):
self._check_twiss('s2') # s2 can be computed at start
self._check_twiss('s1') # s1 can be computed after s2
def test_twiss_with_range(self):
beam = 'ex=1, ey=2, particle=electron, sequence=s1;'
self.mad.command.beam(beam)
self.mad.use('s1')
params = dict(alfx=0.5, alfy=1.5,
betx=2.5, bety=3.5,
sequence='s1')
# Compute TWISS on full sequence, then on a sub-range, then again on
# the full sequence. This checks that none of the range selections
# have side-effects on each other:
betx_full1 = self.mad.twiss(**params)['betx']
betx_range = self.mad.twiss(range=('dr[2]', 'sb'), **params)['betx']
betx_full2 = self.mad.twiss(**params)['betx']
# Check that the results have the expected lengths:
self.assertEqual(len(betx_full1), 9)
self.assertEqual(len(betx_range), 4)
self.assertEqual(len(betx_full2), 9)
# Check numeric results. Since the first 3 elements of range and full
# sequence are identical, equal results are expected. And non-equal
# results afterwards.
self.assertAlmostEqual(betx_range[0], betx_full1[1]) # dr:2, dr:1
self.assertAlmostEqual(betx_range[1], betx_full1[2]) # qp:2, qp:1
self.assertAlmostEqual(betx_range[2], betx_full1[3]) # dr:3, dr:2
self.assertNotAlmostEqual(betx_range[3], betx_full1[4]) # sb, qp:2
def test_range_row_api(self):
beam = 'ex=1, ey=2, particle=electron, sequence=s1;'
self.mad.command.beam(beam)
self.mad.use('s1')
params = dict(alfx=0.5, alfy=1.5,
betx=2.5, bety=3.5,
sequence='s1')
tab = self.mad.twiss(range=('dr[2]', 'sb'), **params)
self.assertEqual(tab.range, ('dr[2]', 'sb'))
self.assertIn('betx', tab)
def test_survey(self):
self.mad.beam()
self.mad.use('s1')
tab = self.mad.survey()
self.assertEqual(tab._name, 'survey')
self.assertIn('x', tab)
self.assertIn('y', tab)
self.assertIn('z', tab)
self.assertIn('theta', tab)
self.assertIn('phi', tab)
self.assertIn('psi', tab)
self.assertLess(tab.x[-1], -1)
assert_allclose(tab.y, 0)
self.assertGreater(tab.z[-1], 7)
def test_match(self):
beam = 'ex=1, ey=2, particle=electron, sequence=s2;'
self.mad.command.beam(beam)
self.mad.use('s2')
params = dict(alfx=0.5, alfy=1.5,
betx=2.5, bety=3.5,
sequence='s2')
self.mad.match(constraints=[dict(range='s1$end', betx=2.0)],
weight={'betx': 2},
vary=['qp2->k1'],
**params)
twiss = self.mad.twiss(**params)
val = twiss.betx[-1]
self.assertAlmostEqual(val, 2.0, places=2)
def test_verbose(self):
self.mad.verbose(False)
self.assertEqual(self.mad.options.echo, False)
self.assertEqual(self.mad.options.info, False)
self.mad.verbose(True)
self.assertEqual(self.mad.options.echo, True)
self.assertEqual(self.mad.options.info, True)
def test_active_sequence(self):
self.mad.command.beam('ex=1, ey=2, particle=electron, sequence=s1;')
self.mad.use('s1')
self.assertEqual(self.mad.sequence(), 's1')
self.mad.beam()
self.mad.use('s2')
self.assertEqual(self.mad.sequence().name, 's2')
def test_get_sequence(self):
with self.assertRaises(KeyError):
self.mad.sequence['sN']
s1 = self.mad.sequence['s1']
self.assertEqual(s1.name, 's1')
seqs = self.mad.sequence
self.assertItemsEqual(seqs, ['s1', 's2'])
def test_eval(self):
self.assertEqual(self.mad.eval(True), True)
self.assertEqual(self.mad.eval(13), 13)
self.assertEqual(self.mad.eval(1.3), 1.3)
self.assertEqual(self.mad.eval([2, True, 'QP_K1']), [2, True, 2.0])
self.assertAlmostEqual(self.mad.eval("1/QP_K1"), 0.5)
def test_globals(self):
g = self.mad.globals
# Membership:
self.assertNotIn('FOO', g)
# Setting values:
g['FOO'] = 2
self.assertIn('FOO', g)
self.assertEqual(g['FOO'], 2)
self.assertEqual(self.mad.eval('FOO'), 2)
# Re-setting values:
g['FOO'] = 3
self.assertEqual(self.mad.eval('FOO'), 3)
# Setting expressions:
g['BAR'] = '3*foo'
self.assertEqual(self.mad.eval('BAR'), 9)
g['FOO'] = 4
self.assertEqual(self.mad.eval('BAR'), 12)
self.assertEqual(g.defs.bar, "3*foo")
self.assertEqual(g.cmdpar.bar.definition, "3*foo")
# attribute access:
g.bar = 42
self.assertEqual(g.defs.bar, 42)
self.assertEqual(g.cmdpar.bar.definition, 42)
self.assertEqual(g.BAR, 42)
# repr
self.assertIn("'bar': 42.0", str(g))
with self.assertRaises(NotImplementedError):
del g['bar']
with self.assertRaises(NotImplementedError):
del g.bar
self.assertEqual(g.bar, 42) # still there
self.assertIn('bar', list(g))
self.assertIn('foo', list(g))
# self.assertEqual(list(g), list(g.defs))
# self.assertEqual(list(g), list(g.cmdpar))
self.assertEqual(len(g), len(list(g)))
self.assertEqual(len(g.defs), len(list(g.defs)))
self.assertEqual(len(g.cmdpar), len(list(g.cmdpar)))
def test_elements(self):
self.assertIn('sb', self.mad.elements)
self.assertIn('sb', list(self.mad.elements))
self.assertNotIn('foobar', self.mad.elements)
self.assertAlmostEqual(self.mad.elements['sb']['angle'], 3.14/4)
idx = self.mad.elements.index('qp1')
elem = self.mad.elements[idx]
self.assertEqual(elem['k1'], 3)
def test_sequence_map(self):
seq = self.mad.sequence
self.assertEqual(len(seq), 2)
self.assertEqual(set(seq), {'s1', 's2'})
self.assertIn('s1', seq)
self.assertNotIn('s3', seq)
self.assertTrue(hasattr(seq, 's1'))
self.assertFalse(hasattr(seq, 's3'))
self.assertEqual(seq.s1.name, 's1')
self.assertEqual(seq.s2.name, 's2')
with self.assertRaises(AttributeError):
seq.s3
def test_table_map(self):
self.mad.beam()
self.mad.use('s2')
self.mad.survey(sequence='s2')
tab = self.mad.table
self.assertIn('survey', list(tab))
self.assertIn('survey', tab)
self.assertNotIn('foobar', tab)
self.assertEqual(len(tab), len(list(tab)))
with self.assertRaises(AttributeError):
tab.foobar
def test_sequence(self):
s1 = self.mad.sequence.s1
self.assertEqual(str(s1), '<Sequence: s1>')
self.assertEqual(s1, self.mad.sequence.s1)
self.assertEqual(s1, 's1')
self.assertNotEqual(s1, self.mad.sequence.s2)
self.assertNotEqual(s1, 's2')
with self.assertRaises(RuntimeError):
s1.beam
with self.assertRaises(RuntimeError):
s1.twiss_table
with self.assertRaises(RuntimeError):
s1.twiss_table_name
self.assertFalse(s1.has_beam)
self.assertFalse(s1.is_expanded)
s1.expand()
self.assertTrue(s1.has_beam)
self.assertTrue(s1.is_expanded)
s1.expand() # idempotent
self.assertTrue(s1.has_beam)
self.assertTrue(s1.is_expanded)
initial = dict(alfx=0.5, alfy=1.5,
betx=2.5, bety=3.5)
self.mad.twiss(sequence='s1', sectormap=True,
table='my_twiss', **initial)
# Now works:
self.assertEqual(s1.beam.particle, 'positron')
self.assertEqual(s1.twiss_table_name, 'my_twiss')
self.assertEqual(s1.twiss_table.betx[0], 2.5)
self.assertEqual(s1.element_names(), [
's1$start',
'dr', 'qp', 'dr[2]', 'qp[2]', 'dr[3]', 'sb', 'dr[4]',
's1$end',
])
self.assertEqual(s1.expanded_element_names(), s1.element_names())
self.assertEqual(len(s1.element_names()), len(s1.element_positions()))
self.assertEqual(s1.element_positions(), [
0.0, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 7.0, 8.0])
self.assertEqual(s1.expanded_element_positions(), [
0.0, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 7.0, 8.0])
self.assertEqual(s1.elements[0].name, 's1$start')
self.assertEqual(s1.elements[-1].name, 's1$end')
self.assertEqual(s1.elements[-1].index, len(s1.elements)-1)
self.assertEqual(s1.elements[3].index, 3)
self.assertEqual(s1.elements.index('#s'), 0)
self.assertEqual(s1.elements.index('#e'), len(s1.elements)-1)
self.assertEqual(s1.elements.index('sb'), 6)
def _get_elems(self, seq_name):
elems = self.mad.sequence[seq_name].elements
elem_idx = dict((el.node_name, i) for i, el in enumerate(elems))
return elems, elem_idx
def test_sequence_get_elements_s1(self):
s1, idx = self._get_elems('s1')
qp1 = s1['qp[1]']
qp2 = s1['qp[2]']
sb1 = s1['sb[1]']
self.assertLess(idx['qp'], idx['qp[2]'])
self.assertLess(idx['qp[2]'], idx['sb'])
self.assertAlmostEqual(qp1['at'], 1.5)
self.assertAlmostEqual(qp2['at'], 3.5)
self.assertAlmostEqual(sb1['at'], 6)
self.assertAlmostEqual(qp1.position, 1)
self.assertAlmostEqual(qp2.position, 3)
self.assertAlmostEqual(sb1.position, 5)
self.assertAlmostEqual(qp1['l'], 1)
self.assertAlmostEqual(qp2['l'], 1)
self.assertAlmostEqual(sb1['l'], 2)
self.assertAlmostEqual(float(qp1['k1']), 2)
self.assertAlmostEqual(float(qp2['k1']), 2)
self.assertAlmostEqual(float(sb1['angle']), 3.14/4)
self.assertEqual(qp1.cmdpar.k1.expr.lower(), "qp_k1")
def test_sequence_get_elements_s2(self):
s2, idx = self._get_elems('s2')
qp1 = s2['qp1[1]']
qp2 = s2['qp2[1]']
self.assertLess(idx['qp1'], idx['qp2'])
self.assertAlmostEqual(qp1['at'], 0)
self.assertAlmostEqual(qp2['at'], 1)
self.assertAlmostEqual(qp1['l'], 1)
self.assertAlmostEqual(qp2['l'], 2)
self.assertAlmostEqual(float(qp1['k1']), 3)
self.assertAlmostEqual(float(qp2['k1']), 2)
# def test_sequence_get_expanded_elements(self):
def test_crash(self):
"""Check that a RuntimeError is raised in case MAD-X crashes."""
self.assertTrue(self.mad)
# a.t.m. MAD-X crashes on this input, because the L (length)
# parametere is missing:
self.assertRaises(RuntimeError, self.mad.input, 'XXX: sequence;')
self.assertFalse(self.mad)
def test_sequence_elements(self):
elements = self.mad.sequence['s1'].elements
iqp2 = elements.index('qp[2]')
qp1 = elements['qp[1]']
qp2 = elements[iqp2]
self.assertAlmostEqual(qp1['at'], 1.5)
self.assertAlmostEqual(qp2['at'], 3.5)
self.assertAlmostEqual(qp1.position, 1)
self.assertAlmostEqual(qp2.position, 3)
self.assertEqual(iqp2, elements.at(3.1))
def test_sequence_expanded_elements(self):
beam = 'ex=1, ey=2, particle=electron, sequence=s1;'
self.mad.command.beam(beam)
self.mad.use('s1')
elements = self.mad.sequence['s1'].expanded_elements
iqp2 = elements.index('qp[2]')
qp1 = elements['qp[1]']
qp2 = elements[iqp2]
self.assertAlmostEqual(qp1['at'], 1.5)
self.assertAlmostEqual(qp2['at'], 3.5)
self.assertAlmostEqual(qp1.position, 1)
self.assertAlmostEqual(qp2.position, 3)
self.assertEqual(iqp2, elements.at(3.1))
def test_element_inform(self):
beam = 'ex=1, ey=2, particle=electron, sequence=s1;'
self.mad.command.beam(beam)
self.mad.use('s1')
elem = self.mad.sequence.s1.expanded_elements['qp']
self.assertSetEqual({'k1', 'l', 'at'}, {
name for name in elem
if elem.cmdpar[name].inform
})
def test_table(self):
beam = 'ex=1, ey=2, particle=electron, sequence=s1;'
self.mad.command.beam(beam)
self.mad.use('s1')
initial = dict(alfx=0.5, alfy=1.5,
betx=2.5, bety=3.5)
twiss = self.mad.twiss(sequence='s1', sectormap=True, **initial)
sector = self.mad.table.sectortable
self.assertTrue(str(twiss).startswith("<Table 'twiss': "))
self.assertTrue(str(sector).startswith("<Table 'sectortable': "))
self.assertIn('betx', twiss)
self.assertIn('t111', sector)
self.assertNotIn('t111', twiss)
self.assertNotIn('betx', sector)
self.assertEqual(len(twiss), len(list(twiss)))
self.assertEqual(set(twiss), set(twiss[0]))
self.assertEqual(twiss.s[5], twiss[5].s)
self.assertEqual(twiss.s[-1], twiss[-1].s)
copy = twiss.copy()
assert_allclose(copy['betx'], twiss.betx)
self.assertEqual(set(copy), set(twiss))
copy = twiss.copy(['betx'])
self.assertEqual(set(copy), {'betx'})
ALL = slice(None)
self.assertEqual(sector.tmat(0).shape, (6, 6, 6))
assert_allclose(sector.tmat(ALL)[0, 0, 0, :], sector.t111)
assert_allclose(sector.tmat(ALL)[1, 5, 3, :], sector.t264)
assert_allclose(sector.tmat(ALL)[3, 0, 3, :], sector.t414)
assert_allclose(sector.tmat(ALL)[4, 4, 4, :], sector.t555)
assert_allclose(sector.rmat(ALL)[0, 0, :], sector.r11)
assert_allclose(sector.rmat(ALL)[1, 5, :], sector.r26)
assert_allclose(sector.rmat(ALL)[3, 0, :], sector.r41)
assert_allclose(sector.rmat(ALL)[4, 4, :], sector.r55)
assert_allclose(sector.kvec(ALL)[0, :], sector.k1)
assert_allclose(sector.kvec(ALL)[1, :], sector.k2)
assert_allclose(sector.kvec(ALL)[3, :], sector.k4)
assert_allclose(sector.kvec(ALL)[4, :], sector.k5)
r = self.mad.sectortable()[:, :6, :6]
k = self.mad.sectortable()[:, 6, :6]
t = self.mad.sectortable2()
num_elems = len(self.mad.sequence.s1.elements)
self.assertEqual(t.shape, (num_elems, 6, 6, 6))
self.assertEqual(r.shape, (num_elems, 6, 6))
self.assertEqual(k.shape, (num_elems, 6))
assert_allclose(t[:, 0, 0, 0], sector.t111)
assert_allclose(t[:, 1, 5, 3], sector.t264)
assert_allclose(t[:, 3, 0, 3], sector.t414)
assert_allclose(t[:, 4, 4, 4], sector.t555)
assert_allclose(r[:, 0, 0], sector.r11)
assert_allclose(r[:, 1, 5], sector.r26)
assert_allclose(r[:, 3, 0], sector.r41)
assert_allclose(r[:, 4, 4], sector.r55)
assert_allclose(k[:, 0], sector.k1)
assert_allclose(k[:, 1], sector.k2)
assert_allclose(k[:, 3], sector.k4)
assert_allclose(k[:, 4], sector.k5)
def test_attr(self):
self.assertTrue(hasattr(self.mad, 'constraint'))
self.assertTrue(hasattr(self.mad, 'constraint_'))
self.assertTrue(hasattr(self.mad, 'global_'))
self.assertFalse(hasattr(self.mad, 'foobar'))
self.assertFalse(hasattr(self.mad, '_constraint'))
def test_expr(self):
g = self.mad.globals
vars = self.mad.expr_vars
g.foo = 1
g.bar = 2
self.assertEqual(set(vars('foo')), {'foo'})
self.assertEqual(set(vars('(foo) * sin(2*pi*bar)')), {'foo', 'bar'})
def test_command(self):
twiss = self.mad.command.twiss
sbend = self.mad.elements.sb
clone = sbend.clone('foobar', angle="pi/5", l=1)
self.assertIn('betx=0', str(twiss))
self.assertIn('angle=', str(sbend))
self.assertIn('tilt', sbend)
self.assertEqual(sbend.tilt, 0)
self.assertEqual(len(sbend), len(list(sbend)))
self.assertIn('tilt', list(sbend))
self.assertEqual(clone.name, 'foobar')
self.assertEqual(clone.base_type.name, 'sbend')
self.assertEqual(clone.parent.name, 'sb')
self.assertEqual(clone.defs.angle, 'pi / 5')
self.assertAlmostEqual(clone.angle, 0.6283185307179586)
self.assertEqual(len(clone), len(sbend))
self.assertIn('angle=0.628', str(clone))
self.assertNotIn('tilt', str(clone))
clone.angle = 0.125
clone = self.mad.elements.foobar # need to update cache
self.assertEqual(clone.angle, 0.125)
self.assertEqual(len(twiss), len(list(twiss)))
self.assertIn('betx', list(twiss))
self.assertNotEqual(clone.angle, clone.parent.angle)
del clone.angle
clone = self.mad.elements.foobar # need to update cache
self.assertEqual(clone.angle, clone.parent.angle)
with self.assertRaises(AttributeError):
clone.missing_attribute
with self.assertRaises(NotImplementedError):
del twiss['betx']
with self.assertRaises(NotImplementedError):
del clone.base_type.angle
def test_array_attribute(self):
self.mad.globals.nine = 9
clone = self.mad.elements.multipole.clone('foo', knl=[0, 'nine/3', 4])
knl = clone.knl
self.assertEqual(knl[0], 0)
self.assertEqual(knl[1], 3)
self.assertEqual(knl[2], 4)
self.assertEqual(len(knl), 3)
self.assertEqual(list(knl), [0.0, 3.0, 4.0])
self.assertEqual(str(knl), '[0.0, 3.0, 4.0]')
knl[1] = '3*nine'
self.assertEqual(self.mad.elements.foo.defs.knl[1], '3 * nine')
self.assertEqual(self.mad.elements.foo.knl[1], 27)
def test_command_map(self):
command = self.mad.command
self.assertIn('match', command)
self.assertIn('sbend', command)
self.assertNotIn('foooo', command)
self.assertIn('match', list(command))
self.assertEqual(len(command), len(list(command)))
self.assertIn('match', str(command))
self.assertIn('sbend', str(command))
self.assertIn('sbend', self.mad.base_types)
self.assertNotIn('match', self.mad.base_types)
def test_comments(self):
mad = self.mad
var = mad.globals
mad('x = 1; ! x = 2;')
self.assertEqual(var.x, 1)
mad('x = 2; // x = 3;')
self.assertEqual(var.x, 2)
mad('x = 3; /* x = 4; */')
self.assertEqual(var.x, 3)
mad('/* x = 3; */ x = 4;')
self.assertEqual(var.x, 4)
mad('x = 5; ! /* */ x = 6;')
self.assertEqual(var.x, 5)
mad('x = 5; /* ! */ x = 6;')
self.assertEqual(var.x, 6)
def test_multiline_input(self):
mad = self.mad
var = mad.globals
mad('''
x = 1;
y = 2;
''')
self.assertEqual(var.x, 1)
self.assertEqual(var.y, 2)
mad('''
x = /* 3;
y =*/ 4;
''')
self.assertEqual(var.x, 4)
self.assertEqual(var.y, 2)
mad('''
x = 1; /* */ x = 2;
*/ if (x == 1) {
x = 3;
}
''')
self.assertEqual(var.x, 2)
mad('''
x = 1; /* x = 2;
*/ if (x == 1) {
x = 3;
}
''')
self.assertEqual(var.x, 3)
large_plot_vrange = [-12, 12]
zoom_vrange = [-1, 1]
zoom_srange = [-60, 60]
mad = Madx()
mad.options.echo = False
mad.options.warn = False
mad.options.info = False
gamma = 7000.
nemitt = 2.5e-6
epsx = nemitt / gamma
epsy = nemitt / gamma
mad.input(f"""Beam,particle=proton,sequence=lhcb1,energy=6500.0,NPART=1.2E11,
sige=1.1e-4,sigt=0.075,ex={epsx},ey={epsy};""")
mad.call(optics_repo + "/lhc_as-built.seq")
betastar_vect = []
sigmax_mat = []
k1l_mat = []
for i_opt in range(1, 31):
if i_opt == 29:
continue
mad.call(optics_repo + '/PROTON/opticsfile.%d' % i_opt)
mad.use('lhcb1')
mad.twiss()
# find_ip
def test_neutral_errors():
# make sure that some misaligned drifts do not influence particle
cpymad_spec = util.find_spec("cpymad")
if cpymad_spec is None:
print("cpymad is not available - abort test")
sys.exit(0)
from cpymad.madx import Madx
madx = Madx()
madx.input('''
T1: Collimator, L=1.0, apertype=CIRCLE, aperture={0.5};
T2: Collimator, L=1.0, apertype=CIRCLE, aperture={0.5};
T3: Collimator, L=1.0, apertype=CIRCLE, aperture={0.5};
KQ1 = 0.02;
KQ2 = -0.02;
testseq: SEQUENCE, l = 20.0;
T1, at = 5;
T2, at = 12;
T3, at = 18;
ENDSEQUENCE;
!---the usual stuff
BEAM, PARTICLE=PROTON, ENERGY=7000.0, EXN=2.2e-6, EYN=2.2e-6;
USE, SEQUENCE=testseq;
Select, flag=makethin, pattern="T1", slice=2;
makethin, sequence=testseq;
use, sequence=testseq;
!---misalign collimators
select, flag = error, clear;
select, flag = error, pattern = "T1";
ealign, dx = 0.01, dy = 0.01, arex = 0.02, arey = 0.02;
select, flag = error, clear;
select, flag = error, pattern = "T2";
ealign, dx = 0.04, dy = 0.04, dpsi = 0.1;
select, flag = error, clear;
select, flag = error, pattern = "T3";
ealign, dx = 0.02, dy = 0.01, arex = 0.03, arey = 0.02, dpsi = 0.1;
select, flag = error, full;
''')
seq = madx.sequence.testseq
pysixtrack_line = pysixtrack.Line.from_madx_sequence(
seq,
install_apertures=True,
apply_madx_errors=True,
)
madx.input('stop;')
initial_x = 0.025
initial_y = -0.015
particle = pysixtrack.Particles()
particle.x = initial_x
particle.y = initial_y
particle.state = 1
pysixtrack_line.track(particle)
assert abs(particle.x-initial_x) < 1e-14
assert abs(particle.y-initial_y) < 1e-14
def _mad(self, doc):
mad = Madx(command_log=CommandLog(sys.stdout, 'X:> '))
mad.input(doc)
return mad
def test_error_functionality():
# check if errors are actually working as intended
cpymad_spec = util.find_spec("cpymad")
if cpymad_spec is None:
print("cpymad is not available - abort test")
sys.exit(0)
from cpymad.madx import Madx
import numpy as np
madx = Madx()
madx.input('''
T1: Collimator, L=0.0, apertype=CIRCLE, aperture={0.5};
T2: Marker;
T3: Collimator, L=0.0, apertype=CIRCLE, aperture={0.5};
testseq: SEQUENCE, l = 20.0;
T1, at = 5;
T2, at = 10;
T3, at = 15;
ENDSEQUENCE;
!---the usual stuff
BEAM, PARTICLE=PROTON, ENERGY=7000.0, EXN=2.2e-6, EYN=2.2e-6;
USE, SEQUENCE=testseq;
!---assign misalignments and field errors
select, flag = error, clear;
select, flag = error, pattern = "T1";
ealign, dx = 0.01, dy = 0.02, arex = 0.03, arey = 0.04;
select, flag = error, clear;
select, flag = error, pattern = "T3";
ealign, dx = 0.07, dy = 0.08, dpsi = 0.7, arex = 0.08, arey = 0.09;
select, flag = error, full;
''')
seq = madx.sequence.testseq
pysixtrack_line = pysixtrack.Line.from_madx_sequence(
seq,
install_apertures=True,
apply_madx_errors=True,
)
madx.input('stop;')
x_init = 0.1*np.random.rand(10)
y_init = 0.1*np.random.rand(10)
particles = pysixtrack.Particles(
x=x_init.copy(),
y=y_init.copy()
)
T1_checked = False
T1_aper_checked = False
T2_checked = False
T3_checked = False
T3_aper_checked = False
for element, element_name in zip(pysixtrack_line.elements,
pysixtrack_line.element_names):
ret = element.track(particles)
if element_name == 't1':
T1_checked = True
assert np.all(abs(particles.x - (x_init - 0.01)) < 1e-14)
assert np.all(abs(particles.y - (y_init - 0.02)) < 1e-14)
if element_name == 't1_aperture':
T1_aper_checked = True
assert np.all(abs(particles.x - (x_init - 0.01 - 0.03)) < 1e-14)
assert np.all(abs(particles.y - (y_init - 0.02 - 0.04)) < 1e-14)
if element_name == 't2':
T2_checked = True
assert np.all(abs(particles.x - x_init) < 1e-14)
assert np.all(abs(particles.y - y_init) < 1e-14)
cospsi = np.cos(0.7)
sinpsi = np.sin(0.7)
if element_name == 't3':
T3_checked = True
assert np.all(abs(
particles.x
- (x_init - 0.07)*cospsi
- (y_init - 0.08)*sinpsi
) < 1e-14)
assert np.all(abs(
particles.y
+ (x_init - 0.07)*sinpsi
- (y_init - 0.08)*cospsi
) < 1e-14)
if element_name == 't3_aperture':
T3_aper_checked = True
assert np.all(abs(
particles.x
- (x_init - 0.07)*cospsi
- (y_init - 0.08)*sinpsi
- (-0.08)
) < 1e-14)
assert np.all(abs(
particles.y
+ (x_init - 0.07)*sinpsi
- (y_init - 0.08)*cospsi
- (-0.09)
) < 1e-14)
if ret is not None:
break
assert not ret
assert np.all([T1_checked, T1_aper_checked,
T2_checked, T3_checked, T3_aper_checked])
class Runner(object):
A = 238
Q = 28
Ekin_per_nucleon = 0.2e9 # in eV
epsx_rms_fin = 35e-6 / 4 # geometrical emittances
epsy_rms_fin = 15e-6 / 4
limit_n_rms_x = 2
limit_n_rms_y = 2
limit_n_rms_z = 3.4
sig_z = 58 / 4. # in m
sig_dp = 0.5e-3
def __init__(self, nturns=20000, npart=1000):
self.nturns = nturns
self.npart = npart
mass = self.A * nmass * 1e9 * e / c**2 # in kg
charge = self.Q * e # in Coul
Ekin = self.Ekin_per_nucleon * self.A
p0c = np.sqrt(Ekin**2 + 2 * Ekin * mass / e * c**2) # in eV
Etot = np.sqrt(p0c**2 + (mass / e)**2 * c**4) * 1e-9 # in GeV
p0 = p0c / c * e # in SI units
gamma = np.sqrt(1 + (p0 / (mass * c))**2)
beta = np.sqrt(1 - gamma**-2)
self.beta = beta
self.gamma = gamma
self.p0 = p0
self.Etot = Etot
self.p0c = p0c
self.charge = charge
self.mass = mass
epsx_gauss = self.epsx_rms_fin * 1.43
epsy_gauss = self.epsy_rms_fin * 1.41
self.epsn_x = epsx_gauss * beta * gamma
self.epsn_y = epsy_gauss * beta * gamma
self.sig_z = self.sig_z * 1.22
self.sig_dp = self.sig_dp * 1.22
self.beta_z = self.sig_z / self.sig_dp
self.madx = Madx()
self.madx.options.echo = False
self.madx.options.warn = False
self.madx.options.info = False
def prepareMaxSimple(self):
# prepare madx
self.madx.command.beam(particle='ion',
mass=self.A * nmass,
charge=self.Q,
energy=self.Etot)
self.madx.call(file="SIS100_RF_220618_9slices.thin.seq")
self.madx.use(sequence='sis100ring')
twiss = self.madx.twiss()
return twiss
def prepareMax(self, qx, qy):
# prepare madx
self.madx.command.beam(particle='ion',
mass=self.A * nmass,
charge=self.Q,
energy=self.Etot)
self.madx.call(file="SIS100_RF_220618_9slices.thin.seq")
self.madx.use(sequence='sis100ring')
self.madx.input('''
match, sequence=SIS100RING;
global, sequence=SIS100RING, q1={qx}, q2={qy};
vary, name=kqf, step=0.00001;
vary, name=kqd, step=0.00001;
lmdif, calls=500, tolerance=1.0e-10;
endmatch;
'''.format(qx=qx, qy=qy))
twiss = self.madx.twiss()
self.madx.input('cavity_voltage = 58.2/1000/number_cavities;')
return twiss
def setup_pysixtrack(self):
seqname = 'sis100ring'
sis100 = getattr(self.madx.sequence, seqname)
pysixtrack_elements = pysixtrack.Line.from_madx_sequence(
self.madx.sequence.sis100ring,
exact_drift=True,
install_apertures=True)
pysixtrack_elements.remove_zero_length_drifts(inplace=True)
pysixtrack_elements.merge_consecutive_drifts(inplace=True)
return pysixtrack_elements
def prepareDis(self, twiss, closed_orbit):
if closed_orbit is not None:
x_co = twiss[0]['x']
y_co = twiss[0]['y']
else:
x_co = 0
y_co = 0
np.random.seed(0)
D_x_0 = twiss[0]['dx'] * self.beta
D_y_0 = twiss[0]['dy'] * self.beta
Dp_x_0 = twiss[0]['dpx'] * self.beta
Dp_y_0 = twiss[0]['dpy'] * self.beta
bx_0 = twiss[0]['betx']
by_0 = twiss[0]['bety']
s0 = twiss[-1]['s']
circumference = s0
alfx_0 = twiss[0]['alfx']
alfy_0 = twiss[0]['alfy']
pyht_beam = generators.generate_Gaussian6DTwiss(
self.npart,
1,
self.charge,
self.mass,
s0,
self.gamma,
alfx_0,
alfy_0,
bx_0,
by_0,
1,
self.epsn_x,
self.epsn_y,
1,
dispersion_x=None,
dispersion_y=None,
limit_n_rms_x=self.limit_n_rms_x**2,
limit_n_rms_y=self.limit_n_rms_y**2,
limit_n_rms_z=self.limit_n_rms_z**2,
)
distribution_z_uncut = generators.gaussian2D(self.sig_z**2)
is_accepted = generators.make_is_accepted_within_n_sigma(
epsn_rms=self.sig_z,
limit_n_rms=2.5,
)
distribution_z_cut = generators.cut_distribution(
distribution_z_uncut, is_accepted)
z, dp = distribution_z_cut(self.npart)
pyht_beam.z, pyht_beam.dp = z, dp / self.beta_z
# recentre on 0 to avoid dipolar motion:
pyht_beam.x -= pyht_beam.mean_x()
pyht_beam.xp -= pyht_beam.mean_xp()
pyht_beam.y -= pyht_beam.mean_y()
pyht_beam.yp -= pyht_beam.mean_yp()
pyht_beam.z -= pyht_beam.mean_z()
pyht_beam.dp -= pyht_beam.mean_dp()
# PyHT generates around 0, need to offset with closed orbit:
pyht_beam.x += x_co
pyht_beam.y += y_co
# add dispersive contribution to coordinates:
pyht_beam.x += D_x_0 * pyht_beam.dp
pyht_beam.y += D_y_0 * pyht_beam.dp
# also need to add D'_{x,y} to momenta:
pyht_beam.xp += Dp_x_0 * pyht_beam.dp
pyht_beam.yp += Dp_y_0 * pyht_beam.dp
return pyht_beam
def setup_sixtracklib(self, pysixtrack_elements, pyht_beam):
elements = pystlib.Elements.from_line(pysixtrack_elements)
elements.BeamMonitor(num_stores=self.nturns)
particles = pystlib.Particles.from_ref(self.npart,
p0c=self.p0c,
mass0=self.A * nmass * 1e9,
q0=self.Q)
particles.x[:] = pyht_beam.x
particles.px[:] = pyht_beam.xp
particles.y[:] = pyht_beam.y
particles.py[:] = pyht_beam.yp
particles.zeta[:] = pyht_beam.z
particles.delta[:] = pyht_beam.dp
particles.rpp[:] = 1. / (pyht_beam.dp + 1)
restmass = self.mass * c**2
restmass_sq = restmass**2
E0 = np.sqrt((self.p0 * c)**2 + restmass_sq)
p = self.p0 * (1 + pyht_beam.dp)
E = np.sqrt((p * c)**2 + restmass_sq)
particles.psigma[:] = (E - E0) / (self.beta * self.p0 * c)
gammai = E / restmass
betai = np.sqrt(1 - 1. / (gammai * gammai))
particles.rvv[:] = betai / self.beta
### prepare trackjob in SixTrackLib
job = pystlib.TrackJob(elements, particles)
return job
def plotDis(self, pyht_beam):
fig, ax = plt.subplots(1, 3, figsize=(15, 5))
plt.sca(ax[0])
plt.title('horizontal')
plt.scatter(pyht_beam.x[::100] * 1e3,
pyht_beam.xp[::100] * 1e3,
s=10,
marker='.')
plt.xlim(1.1 * pyht_beam.x.min() * 1e3, 1.1 * pyht_beam.x.max() * 1e3)
plt.ylim(1.1 * pyht_beam.xp.min() * 1e3,
1.1 * pyht_beam.xp.max() * 1e3)
plt.xlabel('$x$ [mm]')
plt.ylabel("$x'$ [mrad]")
plt.sca(ax[1])
plt.title('vertical')
plt.scatter(pyht_beam.y[::100] * 1e3,
pyht_beam.yp[::100] * 1e3,
s=10,
marker='.')
plt.xlim(1.1 * pyht_beam.y.min() * 1e3, 1.1 * pyht_beam.y.max() * 1e3)
plt.ylim(1.1 * pyht_beam.yp.min() * 1e3,
1.1 * pyht_beam.yp.max() * 1e3)
plt.xlabel('$y$ [mm]')
plt.ylabel("$y'$ [mrad]")
plt.sca(ax[2])
plt.title('longitudinal')
plt.scatter(pyht_beam.z[::100],
pyht_beam.dp[::100] * 1e3,
s=10,
marker='.')
plt.xlabel('$z$ [m]')
plt.ylabel(r"$\Delta p/p_0'$ [$10^{-3}$]")
plt.tight_layout()
plt.savefig('phasespace.png')
def tuneFFT(self, x, y, twiss):
q1mad = twiss.summary['q1']
q2mad = twiss.summary['q2']
ff = np.linspace(0, 0.5, self.nturns // 2 + 1)
xf = abs(np.fft.rfft(x))
q1st = ff[xf.argmax()]
yf = abs(np.fft.rfft(y))
q2st = ff[yf.argmax()]
#
q1madFrac = q1mad % 1
q1madFrac = q1madFrac if q1madFrac < 0.5 else 1 - q1madFrac
q2madFrac = q2mad % 1
q2madFrac = q2madFrac if q2madFrac < 0.5 else 1 - q2madFrac
#
#
print('horizontal:', q1mad, round(1 - q1st, 2), q1madFrac - q1st)
print('vertical:', q2mad, round(1 - q2st, 2), q2madFrac - q2st)
return 1 - q1st, 1 - q2st
def setup_sixtracklib_fft(self, pysixtrack_elements):
elements = pystlib.Elements.from_line(pysixtrack_elements)
elements.BeamMonitor(num_stores=self.nturns)
particles = pystlib.Particles.from_ref(self.npart, p0c=self.p0c)
particles.x += np.linspace(0, 1e-6, self.npart)
particles.y += np.linspace(0, 1e-6, self.npart)
job = pystlib.TrackJob(elements, particles)
return job
def getData(self, job):
x = job.output.particles[0].x[1::self.npart]
y = job.output.particles[0].y[1::self.npart]
match_string = '''
match, sequence=PS;
vary, name= kd, step= 0.00001;
vary, name= kf, step= 0.00001;
global,sequence=PS,Q1= 6.20;
global,sequence=PS,Q2= 6.20;
jacobian, calls = 100, tolerance=1.0e-21;
endmatch;
'''
############## The cell where we first remove all LeQ's and then install the correct ones ##################
## Madx definition ----------------------------------------------------------------------------------------------------------------------
madx = Madx()
#madx = Madx()
madx.input('BEAM, PARTICLE=PROTON, PC = 2.14')
madx.input('BRHO := BEAM->PC * 3.3356;')
# using the latest model directly from EOS (if you are on SWAN)
madx.call(file='/eos/project-a/acc-models/public/ps/2021/ps_mu.seq')
madx.call(file='/eos/project-a/acc-models/public/ps/2021/ps_ss.seq')
madx.call(file='/eos/project-a/acc-models/public/ps/2021/scenarios/bare_machine/0_proton_injection_energy/ps_pro_bare_machine.str')
madx.call(file='/eos/project-a/acc-models/public/ps/supplementary/space_charge_simulations/remove_elements.seq')
# removing the LEQ in addition to all other elements mentioned in the file remove_elements.seq
madx.input('select, flag=seqedit, class = MQNAAIAP;')
madx.input('select, flag=seqedit, class = MQNABIAP;')
madx.input('select, flag=seqedit, class = MQSAAIAP;')
madx.input('select, flag=seqedit, class = QNSD;')
madx.input('select, flag=seqedit, class = QNSF;')
def place_quads(pos, ssz):
madx = Madx(stdout=False)
madx.option(echo=False, warn=False, info=False, debug=False, verbose=False)
madx.input('BEAM, PARTICLE=PROTON, PC = 2.14')
madx.input('BRHO := BEAM->PC * 3.3356;')
madx.call(file='ps_mu.seq')
madx.call(file='ps_ss_mod.seq')
madx.call(file='ps_50LeQ.str')
madx.call(file='ps_pro_bare_machine.str')
madx.call(file='remove_elements.seq')
madx.input('seqedit, sequence = PS;')
madx.input('select, flag=seqedit, class = MQNAAIAP;')
madx.input('select, flag=seqedit, class = MQNABIAP;')
madx.input('select, flag=seqedit, class = MQSAAIAP;')
madx.input('select, flag=seqedit, class = QNSD;')
madx.input('select, flag=seqedit, class = QNSF;')
madx.input('use, sequence = PS;')
madx.input('seqedit,sequence = PS;flatten;endedit;')
madx.input('seqedit,sequence = PS;remove, element=SELECTED;endedit;')
madx.input('endedit;')
madx.input('use, sequence = PS;')
madx.input('seqedit, sequence = PS;')
for s_idx in pos:
if s_idx == 99:
madx.input('PR.QDN00: MULTIPOLE, KNL:={0,kd};')
madx.input('install, element=PR.QDN00, at=' + str(623.834315) +
';')
elif (s_idx % 2) == 1:
madx.input('PR.QDN%02d: MULTIPOLE, KNL:={0,kd};' % (s_idx + 1))
madx.input('install, element=PR.QDN%02d, at=' % (s_idx + 1) +
str(ssz[s_idx]) + ';')
else:
madx.input('PR.QFN%02d: MULTIPOLE, KNL:={0,kf};' % (s_idx + 1))
madx.input('install, element=PR.QFN%02d, at=' % (s_idx + 1) +
str(ssz[s_idx]) + ';')
madx.input('endedit;')
madx.input('use, sequence=PS;')
madx.input('''
match, sequence=PS;
vary, name= kd, step= 0.00001;
vary, name= kf, step= 0.00001;
global,sequence=PS,Q1= 6.10;
global,sequence=PS,Q2= 6.10;
jacobian, calls = 50000, tolerance=1.0e-15;
endmatch;
''')
madx.twiss()
return madx
class TestMadx(unittest.TestCase, _compat.TestCase):
"""
Test methods for the Madx class.
The tests are directly based on the specifics of the sequence in
test/data/lebt/init.madx
Please compare this file for reference.
"""
def setUp(self):
self.mad = Madx(command_log=CommandLog(sys.stdout, 'X:> '))
here = os.path.dirname(__file__)
there = os.path.join(here, 'data', 'lebt', 'init.madx')
self.doc = open(there).read()
for line in self.doc.splitlines():
self.mad._libmadx.input(line)
def tearDown(self):
del self.mad
def test_version(self):
"""Check that the Madx.version attribute can be used as expected."""
version = self.mad.version
# check format:
major, minor, mini = map(int, version.release.split('.'))
# We need at least MAD-X 5.02.03:
self.assertGreaterEqual(major, 5)
self.assertGreaterEqual(minor, 2)
self.assertGreaterEqual(mini, 3)
# check format:
year, month, day = map(int, version.date.split('.'))
self.assertGreaterEqual(year, 2014)
self.assertGreaterEqual(month, 1)
self.assertGreaterEqual(day, 1)
self.assertLessEqual(month, 12)
self.assertLessEqual(day, 31)
def test_independent_instances(self):
# create a second Madx instance (1st one is created in setUp)
madxness = Madx()
# Check independence by defining a variable differently in each
# instance:
self.mad.input('ANSWER=42;')
madxness.input('ANSWER=43;')
self.assertEqual(self.mad.evaluate('ANSWER'), 42);
self.assertEqual(madxness.evaluate('ANSWER'), 43);
def test_command_log(self):
"""Check that the command log contains all input commands."""
# create a new Madx instance that uses the history feature:
history_filename = '_test_madx.madx.tmp'
mad = Madx(command_log=history_filename)
# feed some input and compare with history file:
for line in self.doc.splitlines():
mad.input(line)
with open(history_filename) as history_file:
history = history_file.read()
self.assertEqual(history.strip(), self.doc.strip())
# remove history file
del mad
os.remove(history_filename)
# TODO:
# def test_command(self):
# def test_help(self):
# def test_call(self):
def _check_twiss(self, seq_name):
beam = 'beam, ex=1, ey=2, particle=electron, sequence={0};'.format(seq_name)
self.mad.command(beam)
initial = dict(alfx=0.5, alfy=1.5,
betx=2.5, bety=3.5)
# by explicitly specifying the 'columns' parameter a persistent copy
# is returned. We check that this copy contains the data we want and
# that it has a 'summary' attribute:
twiss = self.mad.twiss(sequence=seq_name,
columns=['betx', 'bety', 'alfx', 'alfy'],
**initial)
betx, bety = twiss['betx'], twiss['bety']
alfx, alfy = twiss['alfx'], twiss['alfy']
# Check initial values:
self.assertAlmostEqual(twiss['alfx'][0], initial['alfx'])
self.assertAlmostEqual(twiss['alfy'][0], initial['alfy'])
self.assertAlmostEqual(twiss['betx'][0], initial['betx'])
self.assertAlmostEqual(twiss['bety'][0], initial['bety'])
self.assertAlmostEqual(twiss.summary['ex'], 1)
self.assertAlmostEqual(twiss.summary['ey'], 2)
# Check that keys are all lowercase:
for k in twiss:
self.assertEqual(k, k.lower())
for k in twiss.summary:
self.assertEqual(k, k.lower())
def test_twiss_1(self):
self._check_twiss('s1') # s1 can be computed at start
self._check_twiss('s1') # s1 can be computed multiple times
self._check_twiss('s2') # s2 can be computed after s1
def test_twiss_2(self):
self._check_twiss('s2') # s2 can be computed at start
self._check_twiss('s1') # s1 can be computed after s2
def test_twiss_with_range(self):
beam = 'beam, ex=1, ey=2, particle=electron, sequence=s1;'
self.mad.command(beam)
params = dict(alfx=0.5, alfy=1.5,
betx=2.5, bety=3.5,
columns=['betx', 'bety'],
sequence='s1')
# Compute TWISS on full sequence, then on a sub-range, then again on
# the full sequence. This checks that none of the range selections
# have side-effects on each other:
betx_full1 = self.mad.twiss(**params)['betx']
betx_range = self.mad.twiss(range=('dr[2]', 'sb'), **params)['betx']
betx_full2 = self.mad.twiss(**params)['betx']
# Check that the results have the expected lengths:
self.assertEqual(len(betx_full1), 9)
self.assertEqual(len(betx_range), 4)
self.assertEqual(len(betx_full2), 9)
# Check numeric results. Since the first 3 elements of range and full
# sequence are identical, equal results are expected. And non-equal
# results afterwards.
self.assertAlmostEqual(betx_range[0], betx_full1[1]) # dr:2, dr:1
self.assertAlmostEqual(betx_range[1], betx_full1[2]) # qp:2, qp:1
self.assertAlmostEqual(betx_range[2], betx_full1[3]) # dr:3, dr:2
self.assertNotAlmostEqual(betx_range[3], betx_full1[4]) # sb, qp:2
def test_range_row_api(self):
beam = 'beam, ex=1, ey=2, particle=electron, sequence=s1;'
self.mad.command(beam)
params = dict(alfx=0.5, alfy=1.5,
betx=2.5, bety=3.5,
sequence='s1')
tab = self.mad.twiss(range=('dr[2]', 'sb'), **params)
self.assertEqual(tab.range, ('dr[2]', 'sb'))
self.assertIn('betx', tab)
# def test_survey(self):
# def test_aperture(self):
# def test_use(self):
# def test_match(self):
# def test_verbose(self):
def test_active_sequence(self):
self.mad.command('beam, ex=1, ey=2, particle=electron, sequence=s1;')
self.mad.active_sequence = 's1'
self.assertEqual(self.mad.active_sequence.name, 's1')
def test_get_sequence(self):
with self.assertRaises(KeyError):
self.mad.sequences['sN']
s1 = self.mad.sequences['s1']
self.assertEqual(s1.name, 's1')
def test_get_sequences(self):
seqs = self.mad.sequences
self.assertItemsEqual(seqs, ['s1', 's2'])
def test_evaluate(self):
val = self.mad.evaluate("1/QP_K1")
self.assertAlmostEqual(val, 0.5)
def test_set_value(self):
self.mad.set_value('FOO', 1)
self.mad.set_value('BAR', 'FOO')
self.mad.set_value('FOO', 2)
self.assertEqual(self.mad.evaluate('FOO'), 2)
self.assertEqual(self.mad.evaluate('BAR'), 1)
def test_set_expression(self):
self.mad.set_expression('FOO', 'BAR')
self.mad.set_value('BAR', 42)
self.mad.evaluate('FOO')
self.assertEqual(self.mad.evaluate('FOO'), 42)
self.mad.set_value('BAR', 43)
self.assertEqual(self.mad.evaluate('FOO'), 43)
def test_globals(self):
g = self.mad.globals
# Membership:
self.assertNotIn('FOO', g)
# Setting values:
g['FOO'] = 2
self.assertIn('FOO', g)
self.assertEqual(g['FOO'], 2)
self.assertEqual(self.mad.evaluate('FOO'), 2)
# Re-setting values:
g['FOO'] = 3
self.assertEqual(self.mad.evaluate('FOO'), 3)
# Setting expressions:
g['BAR'] = '3*foo'
self.assertEqual(self.mad.evaluate('BAR'), 9)
g['FOO'] = 4
self.assertEqual(self.mad.evaluate('BAR'), 12)
def test_elements(self):
self.assertIn('sb', self.mad.elements)
self.assertIn('sb', list(self.mad.elements))
self.assertNotIn('foobar', self.mad.elements)
self.assertAlmostEqual(self.mad.elements['sb']['angle'], 3.14/4)
idx = self.mad.elements.index('qp1')
elem = self.mad.elements[idx]
self.assertEqual(elem['k1'], 3)
# def test_sequence_beam(self):
# def test_sequence_twiss(self):
# def test_sequence_twissname(self):
def _get_elems(self, seq_name):
elems = self.mad.sequences[seq_name].elements
elem_idx = dict((el['name'], i) for i, el in enumerate(elems))
return elems, elem_idx
def test_sequence_get_elements_s1(self):
s1, idx = self._get_elems('s1')
qp1 = s1['qp[1]']
qp2 = s1['qp[2]']
sb1 = s1['sb[1]']
self.assertLess(idx['qp'], idx['qp[2]'])
self.assertLess(idx['qp[2]'], idx['sb'])
self.assertAlmostEqual(qp1['at'], 1)
self.assertAlmostEqual(qp2['at'], 3)
self.assertAlmostEqual(sb1['at'], 5)
self.assertAlmostEqual(qp1['l'], 1)
self.assertAlmostEqual(qp2['l'], 1)
self.assertAlmostEqual(sb1['l'], 2)
self.assertAlmostEqual(float(qp1['k1']), 2)
self.assertAlmostEqual(float(qp2['k1']), 2)
self.assertAlmostEqual(float(sb1['angle']), 3.14/4)
self.assertEqual(str(qp1['k1']).lower(), "qp_k1")
def test_sequence_get_elements_s2(self):
s2, idx = self._get_elems('s2')
qp1 = s2['qp1[1]']
qp2 = s2['qp2[1]']
self.assertLess(idx['qp1'], idx['qp2'])
self.assertAlmostEqual(qp1['at'], 0)
self.assertAlmostEqual(qp2['at'], 1)
self.assertAlmostEqual(qp1['l'], 1)
self.assertAlmostEqual(qp2['l'], 2)
self.assertAlmostEqual(float(qp1['k1']), 3)
self.assertAlmostEqual(float(qp2['k1']), 2)
# def test_sequence_get_expanded_elements(self):
def test_crash(self):
"""Check that a RuntimeError is raised in case MAD-X crashes."""
# a.t.m. MAD-X crashes on this input, because the L (length)
# parametere is missing:
self.assertRaises(RuntimeError, self.mad.input, 'XXX: sequence;')
def test_sequence_elements(self):
elements = self.mad.sequences['s1'].elements
iqp2 = elements.index('qp[2]')
qp1 = elements['qp[1]']
qp2 = elements[iqp2]
self.assertAlmostEqual(qp1['at'], 1)
self.assertAlmostEqual(qp2['at'], 3)
self.assertEqual(iqp2, elements.at(3.1))
def test_sequence_expanded_elements(self):
beam = 'beam, ex=1, ey=2, particle=electron, sequence=s1;'
self.mad.command(beam)
self.mad.use('s1')
elements = self.mad.sequences['s1'].expanded_elements
iqp2 = elements.index('qp[2]')
qp1 = elements['qp[1]']
qp2 = elements[iqp2]
self.assertAlmostEqual(qp1['at'], 1)
self.assertAlmostEqual(qp2['at'], 3)
self.assertEqual(iqp2, elements.at(3.1))
class TestMadx(unittest.TestCase, _TestCaseCompat):
"""
Test methods for the Madx class.
The tests are directly based on the specifics of the sequence in
test/testseq.madx
Please compare this file for reference.
"""
def setUp(self):
self.mad = Madx(command_log=CommandLog(sys.stdout, 'X:> '))
here = os.path.dirname(__file__)
there = os.path.join(here, 'testseq.madx')
self.mad.call(there)
def tearDown(self):
self.mad.quit()
del self.mad
def test_copyright(self):
import cpymad
notice = cpymad.get_copyright_notice()
self.assertIsInstance(notice, type(u""))
def test_version(self):
"""Check that the Madx.version attribute can be used as expected."""
version = self.mad.version
# check format:
major, minor, micro = map(int, version.release.split('.'))
# We need at least MAD-X 5.05.00:
self.assertGreaterEqual((major, minor, micro), (5, 5, 0))
# check format:
year, month, day = map(int, version.date.split('.'))
self.assertGreaterEqual((year, month, day), (2019, 5, 10))
self.assertLessEqual(month, 12)
self.assertLessEqual(day, 31)
self.assertTrue(
str(version).startswith('MAD-X {}'.format(version.release)))
def test_metadata(self):
version = self.mad.version
self.assertEqual(metadata.__version__, version.release)
self.assertIsInstance(metadata.get_copyright_notice(), type(u""))
def test_independent_instances(self):
# create a second Madx instance (1st one is created in setUp)
madxness = Madx()
try:
# Check independence by defining a variable differently in each
# instance:
self.mad.input('ANSWER=42;')
madxness.input('ANSWER=43;')
self.assertEqual(self.mad.eval('ANSWER'), 42)
self.assertEqual(madxness.eval('ANSWER'), 43)
finally:
madxness.quit()
# TODO: We need to fix this on windows, but for now, I just need it to
# pass so that the CI builds the release...
@unittest.skipIf(sys.platform == 'win32', 'Known to be broken on win32!')
def test_streamreader(self):
output = []
m = Madx(stdout=output.append)
try:
self.assertEqual(len(output), 1)
self.assertIn(b'+++++++++++++++++++++++++++++++++', output[0])
self.assertIn(b'+ Support: [email protected],', output[0])
self.assertIn(b'+ Release date: ', output[0])
self.assertIn(b'+ Execution date: ', output[0])
# self.assertIn(b'+ Support: [email protected], ', output[1])
m.input('foo = 3;')
self.assertEqual(len(output), 1)
m.input('foo = 3;')
self.assertEqual(len(output), 2)
self.assertEqual(output[1], b'++++++ info: foo redefined\n')
finally:
m.quit()
self.assertEqual(len(output), 3)
self.assertIn(b'+ MAD-X finished normally ', output[2])
def test_quit(self):
self.mad.quit()
self.assertIsNot(self.mad._process.returncode, None)
self.assertFalse(bool(self.mad))
with self.assertRaises(RuntimeError):
self.mad.input(';')
@unittest.skipIf(sys.platform == 'win32', 'Known to be broken on win32!')
def test_context_manager(self):
output = []
with Madx(stdout=output.append) as m:
m.input('foo = 3;')
self.assertEqual(m.globals.foo, 3)
self.assertIn(b'+ MAD-X finished normally ', output[-1])
self.assertFalse(bool(m))
with self.assertRaises(RuntimeError):
m.input(';')
def test_command_log(self):
"""Check that the command log contains all input commands."""
# create a new Madx instance that uses the history feature:
history_filename = '_test_madx.madx.tmp'
mad = Madx(command_log=history_filename)
try:
# feed some input lines and compare with history file:
lines = dedent("""
l = 5;
f = 200;
fodo: sequence, refer=entry, l=100;
QF: quadrupole, l=5, at= 0, k1= 1/(f*l);
QD: quadrupole, l=5, at=50, k1=-1/(f*l);
endsequence;
beam, particle=proton, energy=2;
use, sequence=fodo;
""").splitlines()
lines = [line for line in lines if line.strip()]
for line in lines:
mad.input(line)
with open(history_filename) as history_file:
history = history_file.read()
self.assertEqual(history.strip(), '\n'.join(lines).strip())
finally:
# remove history file
mad.quit()
del mad
os.remove(history_filename)
def test_append_semicolon(self):
"""Check that semicolon is automatically appended to input() text."""
# Regression test for #73
log = []
mad = Madx(command_log=log.append)
try:
mad.input('a = 0')
mad.input('b = 1')
self.assertEqual(log, ['a = 0;', 'b = 1;'])
self.assertEqual(mad.globals.a, 0)
self.assertEqual(mad.globals.b, 1)
finally:
mad.quit()
del mad
def test_call_and_chdir(self):
folder = os.path.abspath(os.path.dirname(__file__))
parent = os.path.dirname(folder)
getcwd = self.mad._libmadx.getcwd
g = self.mad.globals
self.mad.chdir(folder)
self.assertEqual(normalize(getcwd()), normalize(folder))
self.mad.call('answer_42.madx')
self.assertEqual(g.answer, 42)
with self.mad.chdir('..'):
self.assertEqual(normalize(getcwd()), normalize(parent))
self.mad.call('test/answer_43.madx')
self.assertEqual(g.answer, 43)
self.mad.call('test/answer_call42.madx', True)
self.assertEqual(g.answer, 42)
self.assertEqual(normalize(getcwd()), normalize(folder))
self.mad.call('answer_43.madx')
self.assertEqual(g.answer, 43)
self.mad.chdir('..')
self.assertEqual(normalize(getcwd()), normalize(parent))
def _check_twiss(self, seq_name):
beam = 'ex=1, ey=2, particle=electron, sequence={0};'.format(seq_name)
self.mad.command.beam(beam)
self.mad.use(seq_name)
initial = dict(alfx=0.5, alfy=1.5, betx=2.5, bety=3.5)
twiss = self.mad.twiss(sequence=seq_name, **initial)
# Check initial values:
self.assertAlmostEqual(twiss['alfx'][0], initial['alfx'])
self.assertAlmostEqual(twiss['alfy'][0], initial['alfy'])
self.assertAlmostEqual(twiss['betx'][0], initial['betx'])
self.assertAlmostEqual(twiss['bety'][0], initial['bety'])
self.assertAlmostEqual(twiss.summary['ex'], 1)
self.assertAlmostEqual(twiss.summary['ey'], 2)
# Check that keys are all lowercase:
for k in twiss:
self.assertEqual(k, k.lower())
for k in twiss.summary:
self.assertEqual(k, k.lower())
def test_error(self):
self.mad.input("""
seq: sequence, l=1;
endsequence;
beam;
use, sequence=seq;
""")
# Errors in MAD-X must not crash, but return False instead:
self.assertFalse(self.mad.input('twiss;'))
self.assertTrue(self.mad.input('twiss, betx=1, bety=1;'))
def test_twiss_1(self):
self._check_twiss('s1') # s1 can be computed at start
self._check_twiss('s1') # s1 can be computed multiple times
self._check_twiss('s2') # s2 can be computed after s1
def test_twiss_2(self):
self._check_twiss('s2') # s2 can be computed at start
self._check_twiss('s1') # s1 can be computed after s2
def test_twiss_with_range(self):
beam = 'ex=1, ey=2, particle=electron, sequence=s1;'
self.mad.command.beam(beam)
self.mad.use('s1')
params = dict(alfx=0.5, alfy=1.5, betx=2.5, bety=3.5, sequence='s1')
# Compute TWISS on full sequence, then on a sub-range, then again on
# the full sequence. This checks that none of the range selections
# have side-effects on each other:
betx_full1 = self.mad.twiss(**params)['betx']
betx_range = self.mad.twiss(range=('dr[2]', 'sb'), **params)['betx']
betx_full2 = self.mad.twiss(**params)['betx']
# Check that the results have the expected lengths:
self.assertEqual(len(betx_full1), 9)
self.assertEqual(len(betx_range), 4)
self.assertEqual(len(betx_full2), 9)
# Check numeric results. Since the first 3 elements of range and full
# sequence are identical, equal results are expected. And non-equal
# results afterwards.
self.assertAlmostEqual(betx_range[0], betx_full1[1]) # dr:2, dr:1
self.assertAlmostEqual(betx_range[1], betx_full1[2]) # qp:2, qp:1
self.assertAlmostEqual(betx_range[2], betx_full1[3]) # dr:3, dr:2
self.assertNotAlmostEqual(betx_range[3], betx_full1[4]) # sb, qp:2
def test_range_row_api(self):
beam = 'ex=1, ey=2, particle=electron, sequence=s1;'
self.mad.command.beam(beam)
self.mad.use('s1')
params = dict(alfx=0.5, alfy=1.5, betx=2.5, bety=3.5, sequence='s1')
tab = self.mad.twiss(range=('dr[2]', 'sb'), **params)
self.assertEqual(tab.range, ('dr[2]', 'sb'))
self.assertIn('betx', tab)
def test_survey(self):
self.mad.beam()
self.mad.use('s1')
tab = self.mad.survey()
self.assertEqual(tab._name, 'survey')
self.assertIn('x', tab)
self.assertIn('y', tab)
self.assertIn('z', tab)
self.assertIn('theta', tab)
self.assertIn('phi', tab)
self.assertIn('psi', tab)
self.assertLess(tab.x[-1], -1)
assert_allclose(tab.y, 0)
self.assertGreater(tab.z[-1], 7)
def test_match(self):
beam = 'ex=1, ey=2, particle=electron, sequence=s2;'
self.mad.command.beam(beam)
self.mad.use('s2')
params = dict(alfx=0.5, alfy=1.5, betx=2.5, bety=3.5, sequence='s2')
self.mad.match(constraints=[dict(range='s1$end', betx=2.0)],
weight={'betx': 2},
vary=['qp2->k1'],
**params)
twiss = self.mad.twiss(**params)
val = twiss.betx[-1]
self.assertAlmostEqual(val, 2.0, places=2)
def test_verbose(self):
self.mad.verbose(False)
self.assertEqual(self.mad.options.echo, False)
self.assertEqual(self.mad.options.info, False)
self.mad.verbose(True)
self.assertEqual(self.mad.options.echo, True)
self.assertEqual(self.mad.options.info, True)
def test_active_sequence(self):
self.mad.command.beam('ex=1, ey=2, particle=electron, sequence=s1;')
self.mad.use('s1')
self.assertEqual(self.mad.sequence(), 's1')
self.mad.beam()
self.mad.use('s2')
self.assertEqual(self.mad.sequence().name, 's2')
def test_get_sequence(self):
with self.assertRaises(KeyError):
self.mad.sequence['sN']
s1 = self.mad.sequence['s1']
self.assertEqual(s1.name, 's1')
seqs = self.mad.sequence
self.assertItemsEqual(seqs, ['s1', 's2'])
def test_eval(self):
self.assertEqual(self.mad.eval(True), True)
self.assertEqual(self.mad.eval(13), 13)
self.assertEqual(self.mad.eval(1.3), 1.3)
self.assertEqual(self.mad.eval([2, True, 'QP_K1']), [2, True, 2.0])
self.assertAlmostEqual(self.mad.eval("1/QP_K1"), 0.5)
def test_globals(self):
g = self.mad.globals
# Membership:
self.assertNotIn('FOO', g)
# Setting values:
g['FOO'] = 2
self.assertIn('FOO', g)
self.assertEqual(g['FOO'], 2)
self.assertEqual(self.mad.eval('FOO'), 2)
# Re-setting values:
g['FOO'] = 3
self.assertEqual(self.mad.eval('FOO'), 3)
# Setting expressions:
g['BAR'] = '3*foo'
self.assertEqual(self.mad.eval('BAR'), 9)
g['FOO'] = 4
self.assertEqual(self.mad.eval('BAR'), 12)
self.assertEqual(g.defs.bar, "3*foo")
self.assertEqual(g.cmdpar.bar.definition, "3*foo")
# attribute access:
g.bar = 42
self.assertEqual(g.defs.bar, 42)
self.assertEqual(g.cmdpar.bar.definition, 42)
self.assertEqual(g.BAR, 42)
# repr
self.assertIn("'bar': 42.0", str(g))
with self.assertRaises(NotImplementedError):
del g['bar']
with self.assertRaises(NotImplementedError):
del g.bar
self.assertEqual(g.bar, 42) # still there
self.assertIn('bar', list(g))
self.assertIn('foo', list(g))
# self.assertEqual(list(g), list(g.defs))
# self.assertEqual(list(g), list(g.cmdpar))
self.assertEqual(len(g), len(list(g)))
self.assertEqual(len(g.defs), len(list(g.defs)))
self.assertEqual(len(g.cmdpar), len(list(g.cmdpar)))
def test_elements(self):
self.assertIn('sb', self.mad.elements)
self.assertIn('sb', list(self.mad.elements))
self.assertNotIn('foobar', self.mad.elements)
self.assertAlmostEqual(self.mad.elements['sb']['angle'], 3.14 / 4)
idx = self.mad.elements.index('qp1')
elem = self.mad.elements[idx]
self.assertEqual(elem['k1'], 3)
def test_sequence_map(self):
seq = self.mad.sequence
self.assertEqual(len(seq), 2)
self.assertEqual(set(seq), {'s1', 's2'})
self.assertIn('s1', seq)
self.assertNotIn('s3', seq)
self.assertTrue(hasattr(seq, 's1'))
self.assertFalse(hasattr(seq, 's3'))
self.assertEqual(seq.s1.name, 's1')
self.assertEqual(seq.s2.name, 's2')
with self.assertRaises(AttributeError):
seq.s3
def test_table_map(self):
self.mad.beam()
self.mad.use('s2')
self.mad.survey(sequence='s2')
tab = self.mad.table
self.assertIn('survey', list(tab))
self.assertIn('survey', tab)
self.assertNotIn('foobar', tab)
self.assertEqual(len(tab), len(list(tab)))
with self.assertRaises(AttributeError):
tab.foobar
def test_sequence(self):
s1 = self.mad.sequence.s1
self.assertEqual(str(s1), '<Sequence: s1>')
self.assertEqual(s1, self.mad.sequence.s1)
self.assertEqual(s1, 's1')
self.assertNotEqual(s1, self.mad.sequence.s2)
self.assertNotEqual(s1, 's2')
with self.assertRaises(RuntimeError):
s1.beam
with self.assertRaises(RuntimeError):
s1.twiss_table
with self.assertRaises(RuntimeError):
s1.twiss_table_name
self.assertFalse(s1.has_beam)
self.assertFalse(s1.is_expanded)
s1.expand()
self.assertTrue(s1.has_beam)
self.assertTrue(s1.is_expanded)
s1.expand() # idempotent
self.assertTrue(s1.has_beam)
self.assertTrue(s1.is_expanded)
initial = dict(alfx=0.5, alfy=1.5, betx=2.5, bety=3.5)
self.mad.twiss(sequence='s1',
sectormap=True,
table='my_twiss',
**initial)
# Now works:
self.assertEqual(s1.beam.particle, 'positron')
self.assertEqual(s1.twiss_table_name, 'my_twiss')
self.assertEqual(s1.twiss_table.betx[0], 2.5)
self.assertEqual(s1.element_names(), [
's1$start',
'dr',
'qp',
'dr[2]',
'qp[2]',
'dr[3]',
'sb',
'dr[4]',
's1$end',
])
self.assertEqual(s1.expanded_element_names(), s1.element_names())
self.assertEqual(len(s1.element_names()), len(s1.element_positions()))
self.assertEqual(s1.element_positions(),
[0.0, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 7.0, 8.0])
self.assertEqual(s1.expanded_element_positions(),
[0.0, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 7.0, 8.0])
self.assertEqual(s1.elements[0].name, 's1$start')
self.assertEqual(s1.elements[-1].name, 's1$end')
self.assertEqual(s1.elements[-1].index, len(s1.elements) - 1)
self.assertEqual(s1.elements[3].index, 3)
self.assertEqual(s1.elements.index('#s'), 0)
self.assertEqual(s1.elements.index('#e'), len(s1.elements) - 1)
self.assertEqual(s1.elements.index('sb'), 6)
self.assertEqual(s1.length, 8.0)
def _get_elems(self, seq_name):
elems = self.mad.sequence[seq_name].elements
elem_idx = dict((el.node_name, i) for i, el in enumerate(elems))
return elems, elem_idx
def test_sequence_get_elements_s1(self):
s1, idx = self._get_elems('s1')
qp1 = s1['qp[1]']
qp2 = s1['qp[2]']
sb1 = s1['sb[1]']
self.assertLess(idx['qp'], idx['qp[2]'])
self.assertLess(idx['qp[2]'], idx['sb'])
self.assertAlmostEqual(qp1['at'], 1.5)
self.assertAlmostEqual(qp2['at'], 3.5)
self.assertAlmostEqual(sb1['at'], 6)
self.assertAlmostEqual(qp1.position, 1)
self.assertAlmostEqual(qp2.position, 3)
self.assertAlmostEqual(sb1.position, 5)
self.assertAlmostEqual(qp1['l'], 1)
self.assertAlmostEqual(qp2['l'], 1)
self.assertAlmostEqual(sb1['l'], 2)
self.assertAlmostEqual(float(qp1['k1']), 2)
self.assertAlmostEqual(float(qp2['k1']), 2)
self.assertAlmostEqual(float(sb1['angle']), 3.14 / 4)
self.assertEqual(qp1.cmdpar.k1.expr.lower(), "qp_k1")
def test_sequence_get_elements_s2(self):
s2, idx = self._get_elems('s2')
qp1 = s2['qp1[1]']
qp2 = s2['qp2[1]']
self.assertLess(idx['qp1'], idx['qp2'])
self.assertAlmostEqual(qp1['at'], 0)
self.assertAlmostEqual(qp2['at'], 1)
self.assertAlmostEqual(qp1['l'], 1)
self.assertAlmostEqual(qp2['l'], 2)
self.assertAlmostEqual(float(qp1['k1']), 3)
self.assertAlmostEqual(float(qp2['k1']), 2)
# def test_sequence_get_expanded_elements(self):
def test_crash(self):
"""Check that a RuntimeError is raised in case MAD-X crashes."""
self.assertTrue(self.mad)
# a.t.m. MAD-X crashes on this input, because the L (length)
# parametere is missing:
self.assertRaises(RuntimeError, self.mad.input, 'XXX: sequence;')
self.assertFalse(self.mad)
def test_sequence_elements(self):
elements = self.mad.sequence['s1'].elements
iqp2 = elements.index('qp[2]')
qp1 = elements['qp[1]']
qp2 = elements[iqp2]
self.assertAlmostEqual(qp1['at'], 1.5)
self.assertAlmostEqual(qp2['at'], 3.5)
self.assertAlmostEqual(qp1.position, 1)
self.assertAlmostEqual(qp2.position, 3)
self.assertEqual(iqp2, elements.at(3.1))
def test_sequence_expanded_elements(self):
beam = 'ex=1, ey=2, particle=electron, sequence=s1;'
self.mad.command.beam(beam)
self.mad.use('s1')
elements = self.mad.sequence['s1'].expanded_elements
iqp2 = elements.index('qp[2]')
qp1 = elements['qp[1]']
qp2 = elements[iqp2]
self.assertAlmostEqual(qp1['at'], 1.5)
self.assertAlmostEqual(qp2['at'], 3.5)
self.assertAlmostEqual(qp1.position, 1)
self.assertAlmostEqual(qp2.position, 3)
self.assertEqual(iqp2, elements.at(3.1))
def test_element_inform(self):
beam = 'ex=1, ey=2, particle=electron, sequence=s1;'
self.mad.command.beam(beam)
self.mad.use('s1')
elem = self.mad.sequence.s1.expanded_elements['qp']
self.assertSetEqual(
{'k1', 'l', 'at'},
{name
for name in elem if elem.cmdpar[name].inform})
def test_table(self):
beam = 'ex=1, ey=2, particle=electron, sequence=s1;'
self.mad.command.beam(beam)
self.mad.use('s1')
initial = dict(alfx=0.5, alfy=1.5, betx=2.5, bety=3.5)
twiss = self.mad.twiss(sequence='s1', sectormap=True, **initial)
sector = self.mad.table.sectortable
self.assertTrue(str(twiss).startswith("<Table 'twiss': "))
self.assertTrue(str(sector).startswith("<Table 'sectortable': "))
self.assertIn('betx', twiss)
self.assertIn('t111', sector)
self.assertNotIn('t111', twiss)
self.assertNotIn('betx', sector)
self.assertEqual(len(twiss), len(list(twiss)))
self.assertEqual(set(twiss), set(twiss[0]))
self.assertEqual(twiss.s[5], twiss[5].s)
self.assertEqual(twiss.s[-1], twiss[-1].s)
copy = twiss.copy()
assert_allclose(copy['betx'], twiss.betx)
self.assertEqual(set(copy), set(twiss))
copy = twiss.copy(['betx'])
self.assertEqual(set(copy), {'betx'})
ALL = slice(None)
self.assertEqual(sector.tmat(0).shape, (6, 6, 6))
assert_allclose(sector.tmat(ALL)[0, 0, 0, :], sector.t111)
assert_allclose(sector.tmat(ALL)[1, 5, 3, :], sector.t264)
assert_allclose(sector.tmat(ALL)[3, 0, 3, :], sector.t414)
assert_allclose(sector.tmat(ALL)[4, 4, 4, :], sector.t555)
assert_allclose(sector.rmat(ALL)[0, 0, :], sector.r11)
assert_allclose(sector.rmat(ALL)[1, 5, :], sector.r26)
assert_allclose(sector.rmat(ALL)[3, 0, :], sector.r41)
assert_allclose(sector.rmat(ALL)[4, 4, :], sector.r55)
assert_allclose(sector.kvec(ALL)[0, :], sector.k1)
assert_allclose(sector.kvec(ALL)[1, :], sector.k2)
assert_allclose(sector.kvec(ALL)[3, :], sector.k4)
assert_allclose(sector.kvec(ALL)[4, :], sector.k5)
r = self.mad.sectortable()[:, :6, :6]
k = self.mad.sectortable()[:, 6, :6]
t = self.mad.sectortable2()
num_elems = len(self.mad.sequence.s1.elements)
self.assertEqual(t.shape, (num_elems, 6, 6, 6))
self.assertEqual(r.shape, (num_elems, 6, 6))
self.assertEqual(k.shape, (num_elems, 6))
assert_allclose(t[:, 0, 0, 0], sector.t111)
assert_allclose(t[:, 1, 5, 3], sector.t264)
assert_allclose(t[:, 3, 0, 3], sector.t414)
assert_allclose(t[:, 4, 4, 4], sector.t555)
assert_allclose(r[:, 0, 0], sector.r11)
assert_allclose(r[:, 1, 5], sector.r26)
assert_allclose(r[:, 3, 0], sector.r41)
assert_allclose(r[:, 4, 4], sector.r55)
assert_allclose(k[:, 0], sector.k1)
assert_allclose(k[:, 1], sector.k2)
assert_allclose(k[:, 3], sector.k4)
assert_allclose(k[:, 4], sector.k5)
def test_attr(self):
self.assertTrue(hasattr(self.mad, 'constraint'))
self.assertTrue(hasattr(self.mad, 'constraint_'))
self.assertTrue(hasattr(self.mad, 'global_'))
self.assertFalse(hasattr(self.mad, 'foobar'))
self.assertFalse(hasattr(self.mad, '_constraint'))
def test_expr(self):
g = self.mad.globals
vars = self.mad.expr_vars
g.foo = 1
g.bar = 2
self.assertEqual(set(vars('foo')), {'foo'})
self.assertEqual(set(vars('(foo) * sin(2*pi*bar)')), {'foo', 'bar'})
def test_command(self):
twiss = self.mad.command.twiss
sbend = self.mad.elements.sb
clone = sbend.clone('foobar', angle="pi/5", l=1)
self.assertIn('betx=0', str(twiss))
self.assertIn('angle=', str(sbend))
self.assertIn('tilt', sbend)
self.assertEqual(sbend.tilt, 0)
self.assertEqual(len(sbend), len(list(sbend)))
self.assertIn('tilt', list(sbend))
self.assertEqual(clone.name, 'foobar')
self.assertEqual(clone.base_type.name, 'sbend')
self.assertEqual(clone.parent.name, 'sb')
self.assertEqual(clone.defs.angle, 'pi / 5')
self.assertAlmostEqual(clone.angle, 0.6283185307179586)
self.assertEqual(len(clone), len(sbend))
self.assertIn('angle=0.628', str(clone))
self.assertNotIn('tilt', str(clone))
clone.angle = 0.125
clone = self.mad.elements.foobar # need to update cache
self.assertEqual(clone.angle, 0.125)
self.assertEqual(len(twiss), len(list(twiss)))
self.assertIn('betx', list(twiss))
self.assertNotEqual(clone.angle, clone.parent.angle)
del clone.angle
clone = self.mad.elements.foobar # need to update cache
self.assertEqual(clone.angle, clone.parent.angle)
with self.assertRaises(AttributeError):
clone.missing_attribute
with self.assertRaises(NotImplementedError):
del twiss['betx']
with self.assertRaises(NotImplementedError):
del clone.base_type.angle
def test_array_attribute(self):
self.mad.globals.nine = 9
clone = self.mad.elements.multipole.clone('foo', knl=[0, 'nine/3', 4])
knl = clone.knl
self.assertEqual(knl[0], 0)
self.assertEqual(knl[1], 3)
self.assertEqual(knl[2], 4)
self.assertEqual(len(knl), 3)
self.assertEqual(list(knl), [0.0, 3.0, 4.0])
self.assertEqual(str(knl), '[0.0, 3.0, 4.0]')
knl[1] = '3*nine'
self.assertEqual(self.mad.elements.foo.defs.knl[1], '3 * nine')
self.assertEqual(self.mad.elements.foo.knl[1], 27)
def test_command_map(self):
command = self.mad.command
self.assertIn('match', command)
self.assertIn('sbend', command)
self.assertNotIn('foooo', command)
self.assertIn('match', list(command))
self.assertEqual(len(command), len(list(command)))
self.assertIn('match', str(command))
self.assertIn('sbend', str(command))
self.assertIn('sbend', self.mad.base_types)
self.assertNotIn('match', self.mad.base_types)
def test_comments(self):
mad = self.mad
var = mad.globals
mad('x = 1; ! x = 2;')
self.assertEqual(var.x, 1)
mad('x = 2; // x = 3;')
self.assertEqual(var.x, 2)
mad('x = 3; /* x = 4; */')
self.assertEqual(var.x, 3)
mad('/* x = 3; */ x = 4;')
self.assertEqual(var.x, 4)
mad('x = 5; ! /* */ x = 6;')
self.assertEqual(var.x, 5)
mad('x = 5; /* ! */ x = 6;')
self.assertEqual(var.x, 6)
def test_multiline_input(self):
mad = self.mad
var = mad.globals
mad('''
x = 1;
y = 2;
''')
self.assertEqual(var.x, 1)
self.assertEqual(var.y, 2)
mad('''
x = /* 3;
y =*/ 4;
''')
self.assertEqual(var.x, 4)
self.assertEqual(var.y, 2)
mad('''
x = 1; /* */ x = 2;
*/ if (x == 1) {
x = 3;
}
''')
self.assertEqual(var.x, 2)
mad('''
x = 1; /* x = 2;
*/ if (x == 1) {
x = 3;
}
''')
self.assertEqual(var.x, 3)
def test_errors(self):
mad = self.mad
mad.beam()
mad.use(sequence='s1')
mad.select(flag='error', range='qp')
dkn = [1e-6, 2e-6, 3e-6]
dks = [4e-6, 5e-6, 6e-6]
mad.efcomp(dkn=dkn, dks=dks)
mad.ealign(dx=1e-3, dy=-4e-3)
fd = mad.sequence['s1'].expanded_elements['qp'].field_errors
al = mad.sequence['s1'].expanded_elements['qp'].align_errors
expected_dkn = np.hstack((dkn, np.zeros(len(fd.dkn) - len(dkn))))
expected_dks = np.hstack((dks, np.zeros(len(fd.dks) - len(dks))))
assert_allclose(fd.dkn, expected_dkn)
assert_allclose(fd.dks, expected_dks)
assert_allclose(al.dx, 1e-3)
assert_allclose(al.dy, -4e-3)
def test_makethin(self):
# regression test for segfault, see #67:
self.mad.chdir(os.path.dirname(__file__))
self.mad.call('test_makethin.madx')
def place_quads_wmarkers(int_steps, pos, s_pos, ssz):
madx = Madx(stdout=False)
madx.option(echo=False, warn=False, info=False, debug=False, verbose=False)
madx.input('BEAM, PARTICLE=PROTON, PC = 2.14')
madx.input('BRHO := BEAM->PC * 3.3356;')
madx.call(file='ps_mu.seq')
madx.call(file='ps_ss_mod.seq')
madx.call(file='ps_50LeQ.str')
madx.call(file='ps_pro_bare_machine.str')
madx.call(file='remove_elements.seq')
madx.input('seqedit, sequence = PS;')
madx.input('select, flag=seqedit, class = MQNAAIAP;')
madx.input('select, flag=seqedit, class = MQNABIAP;')
madx.input('select, flag=seqedit, class = MQSAAIAP;')
madx.input('select, flag=seqedit, class = QNSD;')
madx.input('select, flag=seqedit, class = QNSF;')
madx.input('use, sequence = PS;')
madx.input('seqedit,sequence = PS;flatten;endedit;')
madx.input('seqedit,sequence = PS;remove, element=SELECTED;endedit;')
madx.input('endedit;')
madx.input('use, sequence = PS;')
madx.input('seqedit, sequence = PS;')
for i in range(100):
madx.input('MARK%02d: MARKER;' % (i + 1))
madx.input('install, element= MARK%02d, at=' % (i + 1) + str(ssz[i]) +
';')
for i in s_pos:
madx.input('MARK%02d_2: MARKER;' % (i + 1))
madx.input('install, element= MARK%02d_2, at=' % (i + 1) +
str(ssz[i] + 0.01) + ';')
madx.input('endedit;')
madx.input('use, sequence=PS;')
madx.input('select, flag=makethin, CLASS=SBEND, THICK= false, SLICE =' +
str(int_steps) + ';')
madx.input('makethin, sequence=PS;')
madx.input('use, sequence=PS;')
madx.twiss()
posz = [
i for i, elem in enumerate(madx.table.twiss.name)
if elem.startswith('mark') and not (elem.endswith('_2:1'))
]
madx.input('seqedit, sequence = PS;')
for s_idx in pos:
if s_idx == 99:
madx.input('PR.QDN00: MULTIPOLE, KNL:={0,kd};')
madx.input('replace, element=MARK100, by=PR.QDN00;')
elif (s_idx % 2) == 1:
madx.input('PR.QDN%02d: MULTIPOLE, KNL:={0,kd};' % (s_idx + 1))
madx.input('replace, element=MARK%02d, by=PR.QDN%02d;' %
(s_idx + 1, s_idx + 1))
else:
madx.input('PR.QFN%02d: MULTIPOLE, KNL:={0,kf};' % (s_idx + 1))
madx.input('replace, element=MARK%02d, by=PR.QFN%02d;' %
(s_idx + 1, s_idx + 1))
madx.input('endedit;')
madx.input('use, sequence=PS;')
madx.input('''
match, sequence=PS;
vary, name= kd, step= 0.00001;
vary, name= kf, step= 0.00001;
global,sequence=PS,Q1= 6.10;
global,sequence=PS,Q2= 6.10;
jacobian, calls = 50000, tolerance=1.0e-15;
endmatch;
''')
madx.twiss()
return madx, posz
mbb_s = np.array(twiss.s[mbb_list])
ecloud_s = mbb_s.copy()
ecloud_l = 3 * 14.3 * np.ones_like(ecloud_s)
ecloud_arc.append(list(ecloud_s))
ecloud_lengths.append(list(ecloud_l))
elif option == 'between-mq':
for mq_list in mq_arc:
mq_s = np.array(twiss.s[mq_list])
ecloud_s = (mq_s[:-1] + mq_s[1:]) / 2.
ecloud_l = (mq_s[1:] - mq_s[:-1])
ecloud_arc.append(list(ecloud_s))
ecloud_lengths.append(list(ecloud_l))
ecloud_lengths_dict = {}
print('Number of eclouds per arc: %d' % len(ecloud_arc[0]))
mad.input('seqedit, sequence=lhcb1;flatten;')
for i, eclouds in enumerate(ecloud_arc):
ecloud_l = ecloud_lengths[i]
for j, ec in enumerate(eclouds):
ecloud_length = ecloud_l[j]
ecloud_name = f'ecloud.{name_arcs[i]}.{j}'
ecloud_lengths_dict[ecloud_name] = ecloud_length
mad.input((f'install, element={ecloud_name}, class=marker, at={ec};'))
#mad.input(f'install, ecloud.{i%d}.{j%d}
mad.input('flatten;endedit;')
mad.use('lhcb1')
line = pysixtrack.Line.from_madx_sequence(mad.sequence['lhcb1'])
#for el,name in zip(line.elements, line.element_names):
# if name[0:6] == 'ecloud':
def place_quads_wmarkers_FODO(int_steps):
madx = Madx(stdout=False)
madx.option(echo=False, warn=False, info=False, debug=False, verbose=False)
#madx.input('beam, particle=proton, energy=7000;')
madx.input('BEAM, PARTICLE=PROTON, PC = 2.14')
madx.input('BRHO := BEAM->PC * 3.3356;')
madx.call(file='fodo_ring.seq')
madx.input('seqedit,sequence = FODO_ring;flatten;endedit;')
madx.input('use, sequence=FODO_ring;')
madx.input('select, flag=makethin, CLASS=SBEND, THICK= false, SLICE =' +
str(int_steps - 2) + ';')
madx.input('makethin, sequence=FODO_ring;')
madx.input('use, sequence=FODO_ring;')
madx.twiss()
posz = [
i for i, elem in enumerate(madx.table.twiss.name)
if elem.startswith('q')
]
return madx, posz
def table_interpolation_df(my_s_list, my_table):
'''
Thanks to H. Bartosik for sharing his code and for the discussion.
This function will interpolate in a list of s-positions the MAD-X table passed as argument.
This table has to be a full MAD-X twiss table (e.g., 'e1', 'fint',... columns have to be present).
For each element in my_s_list there will be a twiss-command of a short sequence.
This can be time consuming for long lists. It this case please look to MAD-X interpolate command at
http://mad.web.cern.ch/mad/webguide/manual.html#Ch19.S10.
The rationale is to make a new instance of MAD-X in the body of the funciton and to build-and-twiss mini sequences.
The twiss will be performed with given initial values http://mad.web.cern.ch/mad/webguide/manual.html#Ch19.S3.
Args:
my_s_list: list of s-position to be evaluated
my_table: a MAD-X twiss table.
Returns:
The pandas DF with the interpolated values.
See madxp/examples/variablesExamples/000_run.py
'''
# my_s_list=[2.8], my_table=mt.table_df(mad.table.twiss)
my_list = []
madx = Madx(stdout=False)
for my_s in my_s_list:
try:
start_condition = my_table[my_table['s'] < my_s].index[-1]
except:
start_condition = my_table.index[0]
try:
my_element = my_table[my_table['s'] >= my_s].index[0]
except:
my_element = my_table.index[-1]
my_element_row = my_table.loc[my_element]
start_condition_row = my_table.loc[start_condition]
if my_element_row.s == my_s:
interpolation = pd.DataFrame(my_element_row).transpose()
else:
assert my_element_row.l > 0 # The elements need to be thick
if my_element_row.keyword in [
'quadrupole',
'drift',
'sextupole',
'octupole',
'placeholder',
'hmonitor',
'vmonitor',
'monitor',
]:
my_string = f'''
my_special_element: quadrupole,
l= {my_s-start_condition_row.s},
k1={my_element_row.k1l/my_element_row.l},
k1s={my_element_row.k1sl/my_element_row.l},
tilt={my_element_row.tilt};
'''
elif my_element_row.keyword in ['sbend']:
#
my_string = f'''
my_special_element: sbend,
l={my_s-start_condition_row.s},
angle={my_element_row.angle/my_element_row.l*{my_s-start_condition_row.s}},
tilt={my_element_row.tilt},
k1={my_element_row.k1l/my_element_row.l},
e1={my_element_row.e1},
fint={my_element_row.fint},
fintx=0,
hgap={my_element_row.hgap},
k1s={my_element_row.k1sl/my_element_row.l},
h1={my_element_row.h1},
h2=0,
kill_exi_fringe=0;
'''
elif my_element_row.keyword in ['rbend']:
my_string = f'''
option, rbarc=false;
my_special_element: rbend,
l= {my_s-start_condition_row.s},
angle={my_element_row.angle/my_element_row.l*{my_s-start_condition_row.s}},
k1={my_element_row.k1l/my_element_row.l},
e1={my_element_row.e1},
fint={my_element_row.fint},
fintx=0,
hgap={my_element_row.hgap},
k1s={my_element_row.k1sl/my_element_row.l},
h1={my_element_row.h1},
h2=0,
kill_exi_fringe=0;
'''
elif my_element_row.keyword in [
'matrix', 'collimator', 'elseparator', 'hacdipole',
'vacdipole', 'twcavity', 'tkicker', 'hkicker', 'vkicker',
'kicker', 'solenoid'
'rfcavity'
]:
print(
f'The element keyword {my_element_row.keyword} has not been implemented.'
)
print(
f'Consider to remove the interpolating position at {my_s} m.'
)
return
my_string = my_string + f'''
my_special_sequence: sequence, l={my_s-start_condition_row.s}, refer=entry;
at_{my_s:}:my_special_element, at=0;
endsequence;
beam, sequence=my_special_sequence; ! we assume normalized elements
use, sequence=my_special_sequence;
twiss,
betx = {start_condition_row.betx},
alfx = {start_condition_row.alfx},
mux = {start_condition_row.mux},
bety = {start_condition_row.bety},
alfy = {start_condition_row.alfy},
muy = {start_condition_row.muy},
dx = {start_condition_row.dx},
dpx = {start_condition_row.dpx},
dy = {start_condition_row.dy},
dpy = {start_condition_row.dpy},
x = {start_condition_row.x},
px = {start_condition_row.px},
y = {start_condition_row.y},
py = {start_condition_row.py},
t = {start_condition_row.t},
pt = {start_condition_row.pt},
wx = {start_condition_row.wx},
phix = {start_condition_row.phix},
dmux = {start_condition_row.dmux},
wy = {start_condition_row.wy},
phiy = {start_condition_row.phiy},
dmuy = {start_condition_row.dmuy},
ddx = {start_condition_row.ddx},
ddy = {start_condition_row.ddy},
ddpx = {start_condition_row.ddpx},
ddpy = {start_condition_row.ddpy},
r11 = {start_condition_row.r11},
r12 = {start_condition_row.r12},
r21 = {start_condition_row.r21},
r22 = {start_condition_row.r22},
table=special_twiss;
'''
madx.input(my_string)
madx.input(
'delete, sequence=my_special_sequence;my_special_element=0;')
interpolation = table_df(madx.table.special_twiss).iloc[[1]]
interpolation.keyword = 'interpolation'
interpolation.s = my_s
my_list.append(interpolation)
gc.collect()
return pd.concat(my_list)
from cpymad.madx import Madx
beam = 1
i_octupoles = 100.
emittance_in_um = 2.3
n_particles = 2.25e11
chromaticity = 15
xing_angle_urad = 245.
seedran = 1
madx = Madx()
fname_mask = 'hl14_template.mask'
with open(fname_mask) as fid:
mask_content = fid.read()
mask_content = mask_content.replace(r'%BEAM%', str(1))
mask_content = mask_content.replace(r'%OCT%', f'{i_octupoles:e}')
mask_content = mask_content.replace(r'%EMIT_BEAM', f'{emittance_in_um:e}')
mask_content = mask_content.replace(r'%NPART', f'{n_particles:e}')
mask_content = mask_content.replace(r'%CHROM%', f'{chromaticity:e}')
mask_content = mask_content.replace(r'%XING', f'{xing_angle_urad:e}')
mask_content = mask_content.replace(r'%SEEDRAN', f'{seedran:d}')
with open(fname_mask.split('.mask')[0] + '.unmask', 'w') as fid:
fid.write(mask_content)
madx.input(mask_content)
