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 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 make_lhc_beams(madx: Madx,
energy: float = 6500,
emittance: float = 3.75e-6,
**kwargs) -> None:
"""
Define beams with default configuratons for `LHCB1` and `LHCB2` sequences.
Args:
madx (cpymad.madx.Madx): an instanciated cpymad Madx object.
energy (float): beam energy in GeV. Defaults to 6500.
emittance (float): emittance in meters, which will be used to calculate geometric emittance,
then fed to the BEAM command.
Keyword Args:
Any keyword argument that can be given to the MAD-X BEAM command.
"""
logger.info("Making default beams for 'lhcb1' and 'lhbc2' sequences")
madx.globals["NRJ"] = energy
madx.globals["brho"] = energy * 1e9 / madx.globals.clight
geometric_emit = madx.globals["geometric_emit"] = emittance / (energy /
0.938)
for beam in (1, 2):
logger.trace(f"Defining beam for sequence 'lhcb{beam:d}'")
madx.command.beam(
sequence=f"lhcb{beam:d}",
particle="proton",
bv=1 if beam == 1 else -1,
energy=energy,
npart=1.0e10,
ex=geometric_emit,
ey=geometric_emit,
**kwargs,
)
def apply_lhc_colinearity_knob(madx: Madx,
colinearity_knob_value: float = 0,
ir: int = None) -> None:
"""
Applies the LHC colinearity knob. If you don't know what this is, you should not be using this
function.
Args:
madx (cpymad.madx.Madx): an instanciated cpymad Madx object.
colinearity_knob_value (float): Units of the colinearity knob to apply. Defaults to 0 so users
don't mess up local coupling by mistake. This should be a positive integer, normally between 1
and 10.
ir (int): The Interaction Region to apply the knob to, should be one of [1, 2, 5, 8].
Classically 1 or 5.
"""
logger.info(
f"Applying Colinearity knob with a unit setting of {colinearity_knob_value}"
)
logger.warning("You should re-match tunes & chromaticities after this")
knob_variables = (f"KQSX3.R{ir:d}", f"KQSX3.L{ir:d}"
) # MQSX IP coupling correctors
right_knob, left_knob = knob_variables
madx.globals[right_knob] = colinearity_knob_value * 1e-4
logger.trace(f"Set '{right_knob}' to {madx.globals[right_knob]}")
madx.globals[left_knob] = -1 * colinearity_knob_value * 1e-4
logger.trace(f"Set '{left_knob}' to {madx.globals[left_knob]}")
def reset(self):
"""
If MAD-X fails, re-spawn process
"""
self.madx = Madx(stdout=False)
self.madx.option(echo=False,
warn=False,
info=False,
debug=False,
verbose=False)
self.madx.input('BEAM, PARTICLE=PROTON, PC = 2.14')
self.madx.input('BRHO := BEAM->PC * 3.3356;')
self.madx.call(file='ps_mu.seq')
self.madx.call(file='ps_ss_mod.seq')
self.madx.call(file='ps_50LeQ.str')
self.madx.call(file='ps_pro_bare_machine.str')
self.madx.call(file='remove_elements.seq')
self.madx.input('seqedit, sequence = PS;')
self.madx.input('select, flag=seqedit, class = MQNAAIAP;')
self.madx.input('select, flag=seqedit, class = MQNABIAP;')
self.madx.input('select, flag=seqedit, class = MQSAAIAP;')
self.madx.input('select, flag=seqedit, class = QNSD;')
self.madx.input('select, flag=seqedit, class = QNSF;')
self.madx.input('use, sequence = PS;')
self.madx.input('seqedit,sequence = PS;flatten;endedit;')
self.madx.input(
'seqedit,sequence = PS;remove, element=SELECTED;endedit;')
self.madx.input('endedit;')
self.madx.input('use, sequence = PS;')
def madxp(inputFile, pythonData=None, outputDF='output.pkl', command_log_file='log.madx', stdout_file='stdout.madx', verbose=False):
'''
It runs the MADX dataframe using the MADX extended syntax.
inputFile: the MADX input file.
outputDF: the file to dump the output DF. The MADX variable and the pythonData will be available for all
code sections.
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)
aux=df2run(madx, madx2df(inputFile), pythonData=pythonData, verbose=verbose)
if outputDF!=None:
aux.to_pickle(outputDF)
print('Profiling DF saved.')
def test_independent_instances():
# Check independence by defining a variable differently in each
# instance:
with Madx(prompt='X1:> ') as mad1, Madx(prompt='X2:> ') as mad2:
mad1.input('ANSWER=42;')
mad2.input('ANSWER=43;')
assert mad1.eval('ANSWER') == 42
assert mad2.eval('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 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 read_tfsfile(filename):
"""
Read TFS table.
"""
from cpymad.madx import Madx
global TFS_READER
if not TFS_READER:
TFS_READER = Madx()
TFS_READER.command.readmytable(file=abspath(filename), table='user')
return TFS_READER.get_table('user')
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 apply_lhc_rigidity_waist_shift_knob(madx: Madx,
rigidty_waist_shift_value: float = 0,
ir: int = None,
side: str = "left") -> None:
"""
Applies the LHC rigidity waist shift knob, moving the waist left or right of IP. If you don't know what
this is, you should not be using this function. The waist shift is done by unbalancing the
triplet powering knob of the left and right-hand sides of the IP.
Warning: Applying the shift will modify your tunes and most likely flip them, making a subsequent
matching impossible if your lattice has coupling. To avoid this, match to tunes split further apart
before applying the waist shift knob, and then match to the desired working point. For instance for
the LHC, matching to (62.27, 60.36) before applying and afterwards rematching to (62.31, 60.32) usually
works well.
Args:
madx (cpymad.madx.Madx): an instanciated cpymad Madx object.
rigidty_waist_shift_value (float): Units of the rigidity waist shift knob (positive values only).
ir (int): The Interaction Region to apply the knob to, should be one of [1, 2, 5, 8].
Classically 1 or 5.
side (str): Which side of the IP to move the waist to, determines a sign in the calculation.
Defaults to 'left', which means s_waist < s_ip (and setting it to 'right' would move the waist
to s_waist > s_ip).
"""
logger.info(
f"Applying Rigidity Waist Shift knob with a unit setting of {rigidty_waist_shift_value}"
)
logger.warning("You should re-match tunes & chromaticities after this")
right_knob, left_knob = f"kqx.r{ir:d}", f"kqx.l{ir:d}" # IP triplet default knob (no trims)
current_right_knob = madx.globals[right_knob]
current_left_knob = madx.globals[left_knob]
if side.lower() == "left":
madx.globals[right_knob] = (
1 - rigidty_waist_shift_value * 0.005) * current_right_knob
madx.globals[left_knob] = (
1 + rigidty_waist_shift_value * 0.005) * current_left_knob
elif side.lower() == "right":
madx.globals[right_knob] = (
1 + rigidty_waist_shift_value * 0.005) * current_right_knob
madx.globals[left_knob] = (
1 - rigidty_waist_shift_value * 0.005) * current_left_knob
else:
logger.error(
f"Given side '{side}' invalid, only 'left' and 'right' are accepted values."
)
raise ValueError("Invalid value for parameter 'side'.")
logger.trace(f"Set '{right_knob}' to {madx.globals[right_knob]}")
logger.trace(f"Set '{left_knob}' to {madx.globals[left_knob]}")
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_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 main():
np.set_printoptions(
**{
'precision': 5,
'suppress': True, # no scientific notation
'linewidth': 120,
})
# load gantry model
m = Madx(stdout=False)
m.call('../hit_models/hht3/run.madx', True)
d_kick = 0.1e-3
# define elements
els = m.sequence.hht3.expanded_elements
H = els.index('h1ms2h')
V = els.index('h1ms1v')
M = [
el.index for el in els
if el.index > max(H, V) and el.base_name == 'monitor'
]
orm_tab_1 = calc_orms(m, H, V, M, d_kick)
m.globals.kl_b3qd12 += 1e-7
orm_tab_2 = calc_orms(m, H, V, M, d_kick)
m.globals.kl_b3qd12 -= 1e-7
orm_tab = (orm_tab_2 - orm_tab_1) / 1e-7
orm_tab = orm_tab_2
xlabel = [els[m].name for m in M]
plt.plot(xlabel, orm_tab[:, 0], 'o', label="sin x")
plt.plot(xlabel, orm_tab[:, 2], label="var x")
plt.plot(xlabel, orm_tab[:, 4], label="sec x")
plt.legend()
plt.setp(plt.xticks()[1], rotation=50)
plt.show()
plt.clf()
plt.plot(xlabel, orm_tab[:, 1], 'o', label="sin y")
plt.plot(xlabel, orm_tab[:, 3], label="var y")
plt.plot(xlabel, orm_tab[:, 5], label="sec y")
plt.legend()
plt.setp(plt.xticks()[1], rotation=50)
plt.show()
def main():
filenames = glob.glob(
'../data/orm/2019-01-20_quadscan/M8-E108-F1-I9-G1/*/*_X.yml')
raw_data = [yaml.safe_load(read_file(filename)) for filename in filenames]
m = Madx(stdout=False)
m.verbose()
m.call('../hit_models/hht3/sequence.madx', chdir=True)
m.command.beam()
m.use('hht3')
#m.call('../hit_models/hht3/strengths0.madx', chdir=True)
for data in raw_data:
mean = mean_values(data)
err = stddevs(data)
kl = np.array([
sum(optics.values())
for optics, group in itertools.groupby(data['records'],
key=lambda r: r['optics'])
])
mon = data['monitors'][0]
knob, = data['optics'][0].keys()
quad = [
elem.name for elem in m.sequence.hht3.expanded_elements
if elem.base_name == 'quadrupole'
and knob in m.expr_vars(elem.defs.k1)
][0]
m.globals.update(data['base_optics'])
for i, x in enumerate('xy'):
#plt.subplot(2, 1, i+1)
plt.title(f"pos{x}_{mon}({knob})")
plt.ylabel("{x} [m]")
plt.xlabel(f"{knob} [$m^{{-1}}$]")
plt.errorbar(kl, mean[:, 0, i], err[:, 0, i], label=x)
for pos in np.linspace(-0.002, 0.002, 5):
modelled = np.array([
track(m, optics, range=f'{quad}/{mon}', **{x: pos})[i]
for optics, group in itertools.groupby(
data['records'], key=lambda r: r['optics'])
])
i0 = np.argmin(kl)
modelled += mean[i0, 0, i] - modelled[i0]
plt.plot(kl, modelled, label=f'${x}_0$={pos}')
plt.legend()
plt.show()
def switch_magnetic_errors(madx: Madx, **kwargs) -> None:
"""
INITIAL IMPLEMENTATION CREDITS GO TO JOSCHUA DILLY (@JoschD).
Applies magnetic field orders. This will only work for LHC and HLLHC machines.
Args:
madx (cpymad.madx.Madx): an instanciated cpymad Madx object.
Keyword Args:
default: sets global default to this value. Defaults to `False`.
AB#: sets the default for all of that order, the order being the # number.
A# or B#: sets the default for systematic and random of this id.
A#s, B#r etc.: sets the specific value.
In all kwargs, the order # should be in the range (1...15), where 1 == dipolar field.
"""
logger.debug("Setting magnetic errors")
global_default = kwargs.get("default", False)
for order in range(1, 16):
logger.trace(f"Setting up for order {order}")
order_default = kwargs.get(f"AB{order:d}", global_default)
for ab in "AB":
ab_default = kwargs.get(f"{ab}{order:d}", order_default)
for sr in "sr":
name = f"{ab}{order:d}{sr}"
error_value = int(kwargs.get(name, ab_default))
logger.trace(
f"Setting global for 'ON_{name}' to {error_value}")
madx.globals[f"ON_{name}"] = error_value
def test_command_log():
"""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'
try:
# feed some input lines and compare with history file:
lines = """
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.strip() for line in lines if line.strip()]
with Madx(command_log=history_filename) as mad:
for line in lines:
mad.input(line)
with open(history_filename) as history_file:
history = history_file.read()
assert history.strip() == '\n'.join(lines).strip()
finally:
# remove history file
os.remove(history_filename)
def _test_transfer_map(seq, range_, doc, rtol=1e-7, atol=1e-15):
with Madx() as mad:
mad.input(doc)
mad.use(seq)
par = ['x', 'px', 'y', 'py', 't', 'pt']
val = [+0.0010, -0.0015, -0.0020, +0.0025, +0.0000, +0.0000]
twiss = {'betx': 0.0012, 'alfx': 0.0018,
'bety': 0.0023, 'alfy': 0.0027}
twiss.update(zip(par, val))
elems = range_.split('/')
smap = mad.sectormap(elems, sequence=seq, **twiss)[-1]
tw = mad.twiss(sequence=seq, range=range_, **twiss)
# transport of coordinate vector:
x_init = np.array(val)
x_final_tw = np.array([tw[p][-1] for p in par])
x_final_sm = np.dot(smap, np.hstack((x_init, 1)))
assert_allclose(x_final_tw[:4], x_final_sm[:4],
rtol=rtol, atol=atol)
# transport of beam matrix:
tm = smap[0:6, 0:6]
tab_len = len(tw['sig11'])
sig_init = tw.sigmat(0)
sig_final_tw = tw.sigmat(tab_len-1)
sig_final_sm = np.dot(tm, np.dot(sig_init, tm.T))
assert_allclose(sig_final_tw[0:4, 0:4], sig_final_sm[0:4, 0:4],
rtol=rtol, atol=atol)
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 test_format_command(self):
m = self.madx = Madx()
self.assertEqual(util.format_command('twiss', sequence='lhc'),
'twiss, sequence="lhc";')
self.assertEqual(util.format_command('option', echo=True),
'option, echo=true;')
self.assertEqual(
util.format_command('constraint', betx=Constraint(max=3.13)),
'constraint, betx<3.13;')
self.assertEqual(util.format_command(m.command.quadrupole, k1='pi/2'),
'quadrupole, k1:=pi/2;')
self.assertEqual(util.format_command('multipole', knl=[0.0, 1.0, 2.0]),
'multipole, knl={0.0,1.0,2.0};')
self.assertEqual(util.format_command('twiss', range=Range('#s', '#e')),
'twiss, range=#s/#e;')
self.assertEqual(util.format_command(
'select',
class_='quadrupole',
), 'select, class="quadrupole";')
self.assertEqual(
util.format_command(m.elements.quadrupole, k1="hello + world"),
'quadrupole, k1:=hello + world;')
self.assertEqual(
util.format_command(
# match->sequence parameter is list in MAD-X!
m.command.match,
sequence="foo"),
"match, sequence=foo;")
self.assertEqual(
util.format_command(
m.command.select,
class_='quadrupole',
), 'select, class=quadrupole;')
def load_file(cls,
filename,
madx=None,
*,
undo_stack=None,
interpolate=0,
**madx_kwargs):
"""Load model from .madx or .yml file and pass additional arguments
to the Model constructor."""
madx = madx or Madx(**madx_kwargs)
madx.option(echo=False)
filename = os.path.abspath(filename)
path, name = os.path.split(filename)
ext = os.path.splitext(name)[1].lower()
if ext in ('.yml', '.yaml'):
data = yaml.load_file(filename)
path = os.path.join(path, data.get('path', '.'))
_load_params(data, 'beam', path)
_load_params(data, 'twiss', path)
for fname in data.get('init-files', []):
_call(madx, path, fname)
else:
_call(madx, path, filename)
seqname = _guess_main_sequence(madx)
data = _get_seq_model(madx, seqname)
data['init-files'] = [filename]
return cls(madx,
data,
undo_stack=undo_stack,
filename=filename,
interpolate=interpolate)
def apply_lhc_coupling_knob(madx: Madx,
coupling_knob: float = 0,
beam: int = 1,
telescopic_squeeze: bool = False) -> None:
"""
Applies the LHC coupling knob to reach the desired C- value.
Args:
madx (cpymad.madx.Madx): an instanciated cpymad Madx object.
coupling_knob (float): Desired value for the Cminus, typically a few units of 1E-3. Defaults to 0
so users don't mess up coupling by mistake
beam (int): beam to apply the knob to, defaults to beam 1.
telescopic_squeeze (bool): if set to True, uses the knobs for Telescopic Squeeze configuration.
Defaults to False.
"""
logger.info("Applying coupling knob")
logger.warning("You should re-match tunes & chromaticities after this")
suffix = "_sq" if telescopic_squeeze else ""
knob_name = f"CMRS.b{beam:d}{suffix}"
logger.trace(
f"Knob '{knob_name}' is {madx.globals[knob_name]} before implementation"
)
madx.globals[knob_name] = coupling_knob
logger.trace(f"Set '{knob_name}' to {madx.globals[knob_name]}")
def __init__(self, dataFile, madxModelFile, profilePath):
self.madxModelFile = madxModelFile
self.dataFile = dataFile
self.profilePath = profilePath
self.data = self.readData(dataFile)
self.monitor = self.getMonitor()
self.kickers, self.kicks = self.getKicks()
self.sequence = self.getSequence()
self.madx = Madx(stdout=False)
self.madx.call(file=self.madxModelFile, chdir=True)
# This are the initial conditions for the Twiss Module of MAD-X
# there doesn't seem to be a strong dependence on them
self.dx = 1.0e-4
self.dpx = 1.0e-6
self.dy = 1.0e-4
self.dpy = 1.0e-6
def test_context_manager():
output = []
with Madx(stdout=output.append) as m:
m.input('foo = 3;')
assert m.globals.foo == 3
assert b'+ MAD-X finished normally ' in output[-1]
assert not bool(m)
with raises(RuntimeError):
m.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 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 test_append_semicolon():
"""Check that semicolon is automatically appended to input() text."""
# Regression test for #73
log = []
with Madx(command_log=log.append) as mad:
mad.input('a = 0')
mad.input('b = 1')
assert log == ['a = 0;', 'b = 1;']
assert mad.globals.a == 0
assert mad.globals.b == 1
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()
def correct_lhc_orbit(madx: Madx,
orbit_tolerance: float = 1e-14,
iterations: int = 3,
mode: str = "svd",
**kwargs) -> None:
"""
Routine for orbit correction using 'MCB.*' elements in the LHC.
Args:
madx (cpymad.madx.Madx): an instanciated cpymad Madx object.
orbit_tolerance (float): the tolerance for the correction. Defaults to 1e-14.
iterations (int): the number of iterations of the correction to perform. Defaults to 3.
mode (str): the method to use for the correction. Defaults to 'svd'.
Keyword Args:
Any keyword argument that can be given to the MAD-X CORRECT command, such as `mode`, `ncorr`, etc.
"""
logger.info("Starting orbit correction")
for default_kicker in ("kicker", "hkicker", "vkicker", "virtualcorrector"):
logger.trace(f"Disabling default corrector class '{default_kicker}'")
madx.command.usekick(status="off", class_=default_kicker)
logger.debug("Selecting '^MCB.*' correctors")
madx.command.usekick(status="on", pattern="^MCB.*")
madx.command.usemonitor(status="on", class_="monitor")
for _ in range(iterations):
logger.trace("Doing orbit correction for Y then X plane")
madx.twiss(chrom=True)
madx.command.correct(plane="y",
flag="ring",
error=orbit_tolerance,
mode=mode,
**kwargs)
madx.command.correct(plane="x",
flag="ring",
error=orbit_tolerance,
mode=mode,
**kwargs)
def __init__(self, dataFile, madxModelFile, profilePath):
self.madxModelFile = madxModelFile
self.dataFile = dataFile
self.profilePath = profilePath
self.data = self.readData(dataFile)
self.monitor = self.getMonitor()
self.kickers, self.kicks = self.getKicks()
self.sequence = self.getSequence()
self.madx = Madx(stdout=False)
self.madx.call(file=self.madxModelFile, chdir=True)
# This are the initial conditions for the Twiss Module of MAD-X
# there doesn't seem to be a strong dependence on them
self.dx = 1.0e-4
self.dpx = 1.0e-6
self.dy = 1.0e-4
self.dpy = 1.0e-6
# Initial twiss parameters
self.alfax = -3.036572962956109
self.alfay = 0.24718095605022355
self.betax = 20.549681146151855
self.betay = 2.671346837756801
def get_rdts(madx: Madx, order: int = 4, file: Union[Path, str] = None) -> tfs.TfsDataFrame:
"""
INITIAL IMPLEMENTATION CREDITS GO TO JOSCHUA DILLY (@JoschD).
Calculate the RDTs via PTC_TWISS.
Args:
madx (cpymad.madx.Madx): an instanciated cpymad Madx object.
order (int): maximum derivative order (only 0, 1 or 2 implemented in PTC). Defaults to `2`.
file (Union[Path, str]): path to output file. Default `None`
Returns:
A TfsDataframe with results.
"""
logger.info(f"Entering PTC to calculate RDTs up to order {order}")
madx.ptc_create_universe()
logger.trace("Creating PTC layout")
madx.ptc_create_layout(model=3, method=4, nst=3, exact=True)
# madx.ptc_create_layout(model=3, method=6, nst=1) # from Michi
logger.trace("Incorporating MAD-X alignment errors")
madx.ptc_align() # use madx alignment errors
# madx.ptc_setswitch(fringe=True) # include fringe effects
logger.debug("Executing PTC Twiss")
madx.ptc_twiss(icase=6, no=order, normal=True, trackrdts=True)
madx.ptc_end()
logger.debug("Extracting results to TfsDataFrame")
dframe = tfs.TfsDataFrame(madx.table.twissrdt.dframe())
dframe.columns = dframe.columns.str.upper()
dframe.NAME = dframe.NAME.str.upper()
if file:
logger.debug(f"Exporting results to disk at '{Path(file).absolute()}'")
tfs.write(file, dframe)
return dframe
def make_sixtrack_output(madx: Madx, energy: int) -> None:
"""
INITIAL IMPLEMENTATION CREDITS GO TO JOSCHUA DILLY (@JoschD).
Prepare output for sixtrack run.
Args:
madx (cpymad.madx.Madx): an instanciated cpymad Madx object.
energy (float): beam energy in GeV.
"""
logger.info("Preparing outputs for SixTrack")
logger.debug("Powering RF cavities")
madx.globals[
"VRF400"] = 8 if energy < 5000 else 16 # is 6 at injection for protons iirc?
madx.globals[
"LAGRF400.B1"] = 0.5 # cavity phase difference in units of 2pi
madx.globals["LAGRF400.B2"] = 0.0
logger.debug("Executing TWISS and SIXTRACK commands")
madx.twiss() # used by sixtrack
madx.sixtrack(
cavall=True,
radius=0.017) # this value is only ok for HL(LHC) magnet radius
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 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])
from cpymad.madx import Madx import pysixtracklib as pyst import time from scipy.constants import e, m_p, c import numpy as np p0c = 6 * 1e9 # in eV Etot = np.sqrt(p0c**2 + (m_p/e)**2 * c**4) * 1e-9 # in GeV mad = Madx() mad.options.echo = False mad.call(file="fodo.madx") mad.command.beam(particle='proton', energy=str(Etot)) mad.use(sequence="FODO") mad.twiss() mad.command.select(flag="makethin", class_="quadrupole", slice='8') mad.command.select(flag="makethin", class_="sbend", slice='8') mad.command.makethin(makedipedge=False, style="teapot", sequence="fodo") mad.twiss() sis18 = mad.sequence.FODO nturns = 1 elements = pyst.Elements.from_mad(sis18) def prepare(npart=int(1e6), p0c=p0c, elements=elements, device='cpu'):
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)
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))
