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
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 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 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 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
#########################################
with open('iconv.pkl', 'wb') as fid:
pickle.dump(iconv, fid)
#########################################
# Save some optics as dict #
#########################################
seq = 'lhcb1'
mad = Madx()
mad.options.echo = False
mad.options.warn = False
mad.options.info = False
mad.call('lhcwbb_fortracking.seq')
mad.use(seq)
twiss_table = mad.twiss()
optics_dict = {
'betx': twiss_table.betx[0],
'bety': twiss_table.bety[0],
'alfx': twiss_table.alfx[0],
'alfy': twiss_table.alfy[0],
'disp_x': twiss_table.dx[0] / mad.sequence[seq].beam.beta,
'disp_px': twiss_table.dpx[0] / mad.sequence[seq].beam.beta,
'disp_y': twiss_table.dy[0] / mad.sequence[seq].beam.beta,
'disp_py': twiss_table.dpy[0] / mad.sequence[seq].beam.beta,
'qx': mad.table['summ']['q1'][0],
'qy': mad.table['summ']['q2'][0],
'dqx': mad.table['summ']['dq1'][0],
from pysixtrack.be_beambeam.tools import shift_strong_beam_based_on_close_ip
from pysixtrack.be_beambeam.tools import setup_beam_beam_in_line
from pysixtrack import MadPoint
ip_names = [1, 2, 5, 8]
# Parameters to be cross-checked
n_slices = 11
# Use cpymad to compute the required beam parameters
mad = Madx()
mad.options.echo = False
mad.options.warn = False
mad.options.info = False
# Sequences (b1 and b2) with clean machine to compute beam-beam parameters
mad.call('mad/lhcwbb.seq')
# Disable mad beam-beam kicks (we want unperturbed beams)
mad.globals.on_bb_charge = 0.
# Get IP locations from the survey
mad.use('lhcb1')
mad.twiss()
mad.survey()
IP_xyz_b1 = {}
for ip in ip_names:
IP_xyz_b1[ip] = MadPoint('ip%d' % ip+':1', mad, add_CO=False)
mad.use('lhcb2')
mad.twiss()
mad.survey()
import numpy as np
import matplotlib.pyplot as pl
from cpymad.madx import Madx
import sixtracklib as pystlib
import pysixtrack
# prepare madx
mad = Madx()
mad.options.echo = False
mad.call(file="SPS_Q20_thin.seq")
mad.use(sequence='sps')
twiss = mad.twiss()
q1mad = twiss.summary['q1']
q2mad = twiss.summary['q2']
print(q1mad, q2mad)
# Build elements for SixTrackLib
elements = pystlib.Elements.from_line(
pysixtrack.Line.from_madx_sequence(mad.sequence.sps)[0])
nturns = 2**14
ps_line, _ = pysixtrack.Line.from_madx_sequence(mad.sequence.sps)
elements = pystlib.Elements()
elements.append_line(ps_line)
bpm = elements.BeamMonitor(num_stores=nturns)
# Track one turn
npart = 10
particles = pystlib.Particles.from_ref(npart, p0c=26e6)
particles.x += np.linspace(0, 1e-6, npart)
job = pystlib.TrackJob(elements, particles, until_turn_elem_by_elem=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')
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)
neps_x = 1.63e-6
neps_y = 0.86e-6
delta_rms = 1.0e-3
V_RF_MV = 3
lag_RF_deg = 0.0
n_SCkicks = 250
length_fuzzy = 1.5
seq_name = "sps"
mad = Madx()
mad.options.echo = False
mad.options.info = False
mad.warn = False
mad.chdir("madx")
mad.call("sps_thin.madx")
mad.use(seq_name)
# Determine space charge locations
temp_line = pysixtrack.Line.from_madx_sequence(mad.sequence.sps)
sc_locations, sc_lengths = bt.determine_sc_locations(temp_line, n_SCkicks,
length_fuzzy)
# Install spacecharge place holders
sc_names = ["sc%d" % number for number in range(len(sc_locations))]
bt.install_sc_placeholders(mad, seq_name, sc_names, sc_locations, mode=sc_mode)
# twiss
twtable = mad.twiss()
# Generate line with spacecharge
class Madout:
def __init__(self):
self.out = []
def write(self, ll):
self.out.append(ll)
def __repr__(self):
return "".join(ll.decode() for ll in self.out)
madout = Madout()
mad = Madx(stdout=madout)
mad.call("h7ba_n8.seq")
mad.beam(energy=6, sequence="ring", particle="electron", radiate=False)
mad.use("ring")
mad.twiss()
print(mad.table.summ.q1, mad.table.summ.q2)
nslices = 4
mad.select(flag="makethin", clear=True)
mad.select(flag="makethin", class_="sbend", slice=nslices)
mad.select(flag="makethin", class_="quadrupole", slice=nslices)
mad.makethin(sequence="ring")
mad.use(sequence="ring")
print(mad.table.summ.q1, mad.table.summ.q2)
twiss = mad.twiss()
print(mad.table.summ.q1, mad.table.summ.q2)
class OrbitResponse:
"""
Computes the measured and the modeled orbit response matrix at a
single monitor with the available kickers in the transfer line
@param dataFile is the file path where the measured data is. It is
expected to be a measurement as produced by madgui in yaml format.
@param madxModelFile is the file path to the MAD-X model file. The model
should run in MAD-X.
"""
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 readData(self, dataFile):
"""
Reads data from the yaml file as written by madgui
and returns an ordered dictionary
"""
with open(dataFile) as f:
data = safe_load(f)
knobs = data['knobs']
# Normal ordering for kickers.
# Namely, according to their position s
# and in increasing order
data['records'] = sorted(data['records'],
key=lambda record: -1
if not record['optics'] else knobs.index(
list(record['optics'])[0]))
return data
def setData(self, dataFile):
self.dataFile = dataFile
self.data = self.readData(dataFile)
self.monitor = self.getMonitor()
self.kickers, self.kicks = self.getKicks()
def getMonitor(self):
records = self.data['records']
mess0 = records[0]['shots']
monitor = list(mess0[0].keys())
return monitor[0]
def getKicks(self):
"""
Returns a list of kickers names and corrector kick angles
that were used in the measurement
"""
records = self.data['records']
kickers, kicks = [], []
for messung in records:
kicker = list(messung['optics'].keys())
if (len(kicker) != 0):
kick = messung['optics']
kickers.append(kicker[0])
kicks.append(kick[kicker[0]])
return kickers, kicks
def getSequence(self):
sequences = ['hht1', 'hht2', 'hht3', 'hht4', 'hht5']
try:
for seqName in sequences:
if seqName in self.madxModelFile:
return seqName
except:
print('Sequence not found!')
#####################################################################
# Set here plotShots True to show the grid Reads with fitted peaks #
# Set showProfiles to True to show each measurement #
#####################################################################
def ormMeasured(self, plotShots=True):
"""
Computes the measured orbit responses at an specifical
monitor, returns the orbit response entries and their errors
as two arrays. The first entry of the arrays are the horizontal
response and the second is vertical response respectively.
"""
madx = self.madx
self.madx.call(file=self.madxModelFile, chdir=True)
self.madx.globals.update(self.data['model'])
if (self.monitor == 'h1dg1g' or self.monitor == 'h1dg2g'
or self.monitor == 'g3dg3g' or self.monitor == 'g3dg5g'):
profilePath = '../ormData/ormMessdata/2019-11-17/ORM_profile/'
profAnalizer = ProfileAnalyzer(self.data, profilePath)
else:
profAnalizer = ProfileAnalyzer(self.data, self.profilePath)
profAnalizer.fitProfiles(self.monitor.upper(),
showProfiles=False,
skipShots=1,
plot=plotShots)
pOrmx, pOrmy = profAnalizer.messDatax, profAnalizer.messDatay
gridProfiles = (len(pOrmx) != 0)
print(' GridProfiles: ', gridProfiles)
records = self.data['records']
beamMess = []
beamMessErr = []
for messung in records:
shots = messung['shots']
dataBeam = []
for shot in shots:
dataBeam.append(shot[self.monitor])
mean = np.mean(dataBeam, axis=0)
stdDev = np.std(dataBeam, axis=0)
# Last two entries correspond to the beam envelope measurement.
# (That's why we just take the first two)
beamMess.append(np.array(mean[:2]))
beamMessErr.append(np.array(stdDev[:2]))
orbitResponse = []
orbitResponseErr = []
# Note that len(kicks) + 1 == len(beamMess) must hold
# this condition could be implemented as control
for k in range(len(self.kicks)):
k0 = madx.globals[self.kickers[k]]
kickDiff = (self.kicks[k] - k0)
# First measurement is always the reference measurement
orbR_k = (beamMess[k + 1] - beamMess[0]) / kickDiff
orbRErr_k = np.sqrt(beamMessErr[k+1]**2 + beamMessErr[0]**2) \
/ kickDiff
if (gridProfiles):
# Grid horizontal axis is inverted
orbR_kx = -(pOrmx[self.kickers[k]][0] -
pOrmx[''][0]) / kickDiff
orbR_ky = (pOrmy[self.kickers[k]][0] - pOrmy[''][0]) / kickDiff
# Gaussian error propagation
if self.monitor == 'h1dg1g':
print(' Using fit for H1DG1G')
sigx = 0.41 * 10**-3
sigy = 0.059 * 10**-3
orbRErr_kx = np.sqrt(2 * sigx**2) / kickDiff
orbRErr_ky = np.sqrt(2 * sigy**2) / kickDiff
else:
orbRErr_kx = np.sqrt(pOrmx[self.kickers[k]][1]**2 + pOrmx[''][1]**2) \
/ kickDiff
orbRErr_ky = np.sqrt(pOrmy[self.kickers[k]][1]**2 + pOrmy[''][1]**2) \
/ kickDiff
orbitResponse.append([orbR_kx, orbR_ky])
orbitResponseErr.append([orbRErr_kx, orbRErr_ky])
else:
orbitResponse.append(orbR_k)
orbitResponseErr.append(orbRErr_k)
orbitResponse = np.transpose(orbitResponse)
orbitResponseErr = np.transpose(orbitResponseErr)
return orbitResponse, orbitResponseErr
def ormModel(self, kick=2e-4):
"""
Computes the orbit response according to the MADX model
with the same parameters as in the measurement
"""
madx = self.madx
madx.call(file=self.madxModelFile, chdir=True)
elems = madx.sequence[self.sequence].expanded_elements
iMonitor = elems.index(self.monitor)
madx.globals.update(self.data['model'])
twiss0 = madx.twiss(sequence=self.sequence,
RMatrix=True,
alfx=self.alfax,
alfy=self.alfay,
betx=self.betax,
bety=self.betay,
x=self.dx,
y=self.dy,
px=self.dpx,
py=self.dpy)
x0 = twiss0.x[iMonitor]
y0 = twiss0.y[iMonitor]
orbitResponse_x = []
orbitResponse_y = []
for k in self.kickers:
madx.globals.update(self.data['model'])
madx.globals[k] += kick
twiss1 = madx.twiss(sequence=self.sequence,
RMatrix=True,
alfx=self.alfax,
alfy=self.alfay,
betx=self.betax,
bety=self.betay,
x=self.dx,
y=self.dy,
px=self.dpx,
py=self.dpy)
x1 = twiss1.x[iMonitor]
y1 = twiss1.y[iMonitor]
# Horizontal and vertical response
cx = (x1 - x0) / kick
cy = (y1 - y0) / kick
orbitResponse_x.append(cx)
orbitResponse_y.append(cy)
return orbitResponse_x, orbitResponse_y
def ormModelpList(self,
pList,
dpList,
kickers,
dkickers,
dmonx,
dmony,
kick=2e-4):
"""
Computes the ORM and allows a user-defined list of
parameters to be changed during the computation
@param pList is a list containing the names of the
parameters
@param dpList is the absolute value by which the parameter
has to be changed in the computation from the nominal
value given in the measurement (self.data['model'])
"""
madx = self.madx
elems = madx.sequence[self.sequence].expanded_elements
iMonitor = elems.index(self.monitor)
orbitResponse_x = []
orbitResponse_y = []
madx.globals.update(self.data['model'])
# We update the model given the parameter list
# before we compute the Orbit Response
for p in range(len(pList)):
madx.globals[pList[p]] += dpList[p]
for k in range(len(kickers)):
madx.globals[kickers[k]] /= dkickers[k]
twiss1 = madx.twiss(sequence=self.sequence,
RMatrix=True,
alfx=self.alfax,
alfy=self.alfay,
betx=self.betax,
bety=self.betay,
x=self.dx,
y=self.dy,
px=self.dpx,
py=self.dpy)
x1 = twiss1.x[iMonitor]
y1 = twiss1.y[iMonitor]
for k_i in range(len(self.kickers)):
k = self.kickers[k_i]
madx.globals.update(self.data['model'])
# We update the model again. (To reset the Twiss command)
for p in range(len(pList)):
madx.globals[pList[p]] += dpList[p]
for kname in range(len(kickers)):
madx.globals[kickers[kname]] /= (dkickers[kname])
if len(dkickers):
madx.globals[k] = madx.globals[k] + kick / dkickers[k_i]
else:
madx.globals[k] += kick
twiss2 = madx.twiss(sequence=self.sequence,
RMatrix=True,
alfx=self.alfax,
alfy=self.alfay,
betx=self.betax,
bety=self.betay,
x=self.dx,
y=self.dy,
px=self.dpx,
py=self.dpy)
x2 = twiss2.x[iMonitor]
y2 = twiss2.y[iMonitor]
# Horizontal and vertical response
cx = (x2 - x1) / (kick * dmonx)
cy = (y2 - y1) / (kick * dmony)
orbitResponse_x.append(cx)
orbitResponse_y.append(cy)
return np.array(orbitResponse_x), np.array(orbitResponse_y)
def centralDorm(self,
pList,
dpList,
kickers=[],
dkickers=[],
dmonx=1,
dmony=1,
kick=2e-4,
dp=1e-5):
"""
Central finite difference
"""
dpList = np.array(dpList)
dCx = []
dCy = []
for param in range(len(pList)):
# Working copy of the parameter list
dpList_back = np.array(dpList)
dpList_forw = np.array(dpList)
dpList_back[param] -= dp
dpList_forw[param] += dp
ormMxback, ormMyback = self.ormModelpList(pList, dpList_back,
kickers, dkickers, dmonx,
dmony, kick)
ormMxforw, ormMyforw = self.ormModelpList(pList, dpList_forw,
kickers, dkickers, dmonx,
dmony, kick)
dCxdP = (ormMxforw - ormMxback) / (2 * dp)
dCydP = (ormMyforw - ormMyback) / (2 * dp)
dCx.append(dCxdP)
dCy.append(dCydP)
return np.array(dCx), np.array(dCy)
def visualizeData(self, messFile2='', saveFigs=False):
"""
This functions plots the data of a given file, together
with the modeled (expected) orbit response from MADX.
If a second measurement file is given, it will be assumed the
same conditions hold (same Monitor, same Kickers).
@param messFile2 is the second measurement
@param saveFigs(boolean) saves the plots if True
"""
ormG, dormG = self.ormMeasured()
ormMx, ormMy = self.ormModel()
# Horizontal response
y1 = ormG[0]
dy1 = dormG[0]
# Vertical response
y2 = ormG[1]
dy2 = dormG[1]
if messFile2 != '':
data2 = self.readData(messFile2)
ormG2, dormG2 = self.ormMeasured(data2)
y12 = ormG2[0]
dy12 = dormG2[0]
y22 = ormG2[1]
dy22 = dormG2[1]
self._plotData(y1, dy1, y2, dy2, ormMx, ormMy)
if (messFile2 != ''):
self._plotData(y12, dy12, y22, dy22, ormMx, ormMy, plotModel=False)
plt.show()
plt.clf()
plt.cla()
plt.close()
def _plotData(self,
y1,
dy1,
y2,
dy2,
ormMx,
ormMy,
save=False,
plotModel=True):
"""
Just for inner functionality
"""
x = np.linspace(0, len(y1), len(y1))
plt.figure(1, figsize=(8, 8))
plt.errorbar(x,
y1,
yerr=dy1,
label="Measurement",
marker="x",
linestyle="",
capsize=5)
if (plotModel):
plt.plot(x, ormMx, label="Model", marker=".", linestyle="-")
plt.xlabel('Kicker')
plt.ylabel(r'Horizontal Orbit Response [mm mrad$^{-1}$]')
locs, labels = plt.xticks()
plt.xticks(x, self.kickers, rotation='vertical')
plt.title("Monitor: {}".format(self.monitor))
plt.legend(loc=0)
if (save):
plt.savefig('Results/{}h'.format(self.monitor))
plt.close()
plt.figure(2, figsize=(8, 8))
plt.errorbar(x,
y2,
yerr=dy2,
label="Measurement",
marker="x",
linestyle="",
capsize=5)
if (plotModel):
plt.plot(x, ormMy, label="Model", marker=".", linestyle="-")
plt.xlabel('Kicker')
plt.ylabel(r'Vertical Orbit Response [mm mrad$^{-1}$]')
locs, labels = plt.xticks()
plt.xticks(x, self.kickers, rotation='vertical')
plt.title("Monitor: {}".format(self.monitor))
plt.legend(loc=0)
if (save):
plt.savefig('Results/{}v'.format(self.monitor))
plt.close()
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
from cpymad.madx import Madx
import pysixtrack
mad = Madx()
mad.options.echo = False
mad.options.warn = False
mad.options.info = False
mad.call("mad/lhcwbb.seq")
line, other = pysixtrack.Line.from_madx_sequence(mad.sequence.lhcb1)
print(line)
import os
import shutil
from cpymad.madx import Madx
os.system('gfortran headonslice.f -o headonslice')
mad=Madx()
mad.globals.mylhcbeam = 1
mad.globals.on_bb_switch = 1
mad.call('ts_collisions_ats30_en20_IMO380_C7_X160_I1.2_62.3100_60.3200.mask')
# mad.input('save, sequence=lhcb1,lhcb2, beam=true, file=lhcwbb.seq;')
try:
os.mkdir('../sixtrack')
except FileExistsError:
pass
shutil.copy('fc.2', '../sixtrack/fort.2')
with open('../sixtrack/fort.3', 'w') as fout:
with open('fort_beginning.3', 'r') as fid_fort3b:
fout.write(fid_fort3b.read())
with open('fc.3', 'r') as fid_fc3:
fout.write(fid_fc3.read())
with open('fort_end.3', 'r') as fid_fort3e:
fout.write(fid_fort3e.read())
import os
import shutil
from cpymad.madx import Madx
if not os.path.exists('beambeam_macros/headonslice'):
os.system('./init_beambeam_macros.sh')
mask_fname = 'ts_ats30_newMacro_en25_IMO550_C15_X160_I1.2_62.31_60.32.mask'
mad = Madx()
mad.globals.mylhcbeam = 1
mad.globals.on_bb_switch = 1
mad.call(mask_fname)
# mad.input('save, sequence=lhcb1,lhcb2, beam=true, file=lhcwbb.seq;')
try:
os.mkdir('sixtrack')
except FileExistsError:
pass
shutil.copy('fc.2', 'sixtrack/fort.2')
with open('sixtrack/fort.3', 'w') as fout:
with open('fort_beginning.3', 'r') as fid_fort3b:
fout.write(fid_fort3b.read())
with open('fc.3', 'r') as fid_fc3:
fout.write(fid_fc3.read())
with open('fort_end.3', 'r') as fid_fort3e:
fout.write(fid_fort3e.read())
def test_madx_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
seq_name = "psb1"
use_aperture = True
n_SCkicks = 120
length_fuzzy = 0.0
p0c = 0.571e6
particle = pysixtrack.Particles(p0c=p0c)
betagamma = particle.beta0 * particle.gamma0
# mass = pysixtrack.Particles.pmass
delta_rms = 1e-3
neps_x = 1.5e-6
neps_y = 1.5e-6
# for space charge bunched
number_of_particles = 1e11
bunchlength_rms = 1.0
# for space charge coasting
line_density = 1e11
for sc_mode in ["Bunched", "Coasting"]:
mad = Madx()
mad.options.echo = False
mad.options.info = False
mad.warn = False
file_path = os.path.realpath(__file__)
path = os.path.dirname(file_path) + "/psb/"
mad.call(path + "psb_fb_lhc.madx", chdir=True)
# Determine space charge locations
temp_line = pysixtrack.Line.from_madx_sequence(mad.sequence[seq_name])
sc_locations, sc_lengths = bt.determine_sc_locations(
temp_line, n_SCkicks, length_fuzzy
)
# Install spacecharge place holders
sc_names = ["sc%d" % number for number in range(len(sc_locations))]
bt.install_sc_placeholders(
mad, seq_name, sc_names, sc_locations, mode=sc_mode
)
# Generate line with spacecharge
line = pysixtrack.Line.from_madx_sequence(
mad.sequence[seq_name], install_apertures=use_aperture
)
# Get sc info from optics
mad_sc_names, sc_twdata = bt.get_spacecharge_names_twdata(
mad, seq_name, mode=sc_mode
)
# Check consistency
if sc_mode == "Bunched":
sc_elements, sc_names = line.get_elements_of_type(
pysixtrack.elements.SpaceChargeBunched
)
elif sc_mode == "Coasting":
sc_elements, sc_names = line.get_elements_of_type(
pysixtrack.elements.SpaceChargeCoasting
)
else:
raise ValueError("mode not understood")
bt.check_spacecharge_consistency(
sc_elements, sc_names, sc_lengths, mad_sc_names
)
# Setup spacecharge in the line
if sc_mode == "Bunched":
bt.setup_spacecharge_bunched_in_line(
sc_elements,
sc_lengths,
sc_twdata,
betagamma,
number_of_particles,
bunchlength_rms,
delta_rms,
neps_x,
neps_y,
)
elif sc_mode == "Coasting":
bt.setup_spacecharge_coasting_in_line(
sc_elements,
sc_lengths,
sc_twdata,
betagamma,
line_density,
delta_rms,
neps_x,
neps_y,
)
else:
raise ValueError("mode not understood")
from cpymad.madx import Madx
import pysixtrack
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
from cpymad.madx import Madx
# 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')
import pyht_beamsize
#plt.style.use('kostas')
plt.close('all')
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--noblock', dest='plot_block', action='store_false')
args = parser.parse_args()
mad = Madx()
mad.options.echo = False
mad.options.warn = False
#mad.options.info = False
mad.call('lhc_injection_fortracking.seq')
mad.use('lhcb1')
twiss = mad.twiss()
option = 'every-mb.b'
#option = 'between-mq'
fig1 = plt.figure(1)
ax1 = fig1.add_subplot(111)
ax1.plot(twiss.s, twiss.x)
ax1.set_xlabel('$\mathbf{s [m]}$')
ax1.set_ylabel('$\mathbf{x_{CO},y_{CO} [m]}$')
ax1.plot(twiss.s, twiss.y)
ax1.set_xlabel('$\mathbf{s [m]}$')
ax1.set_ylabel('$\mathbf{y_{CO} [m]}$')
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]
# %% Unmask the mask
beam = 1
i_octupoles = 100.
emittance_in_um = 2.3
n_particles = 2.25e11
chromaticity = 15
xing_angle_urad = 245.
seedran = 1
fname_mask = 'hl14_template_split.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}')
# %% Dump the unmasked mask on file
unmask_fname = fname_mask.split('.mask')[0] + '_unmask.mask'
with open(unmask_fname, 'w') as fid:
fid.write(mask_content)
#os.system('madx '+unmask_fname)
mad = Madx()
mad.call(unmask_fname)
# %%
from cpymad.madx import Madx
import pyblep
mad = Madx()
mad.options.echo = False
mad.options.warn = False
mad.options.info = False
mad.call('mad/lhcwbb.seq')
line, other = pyblep.from_madx_sequence(mad.sequence.lhcb1)
)
from pysixtrack.be_beamfields.tools import setup_beam_beam_in_line
from pysixtrack import MadPoint
ip_names = [1, 2, 5, 8]
# Parameters to be cross-checked
n_slices = 11
# Use cpymad to compute the required beam parameters
mad = Madx()
mad.options.echo = False
mad.options.warn = False
mad.options.info = False
# Sequences (b1 and b2) with clean machine to compute beam-beam parameters
mad.call("mad/lhcwbb.seq")
# Disable mad beam-beam kicks (we want unperturbed beams)
mad.globals.on_bb_charge = 0.0
# Get IP locations from the survey
mad.use("lhcb1")
mad.twiss()
mad.survey()
IP_xyz_b1 = {}
for ip in ip_names:
IP_xyz_b1[ip] = MadPoint.from_survey("ip%d" % ip + ":1", mad)
mad.use("lhcb2")
mad.twiss()
mad.survey()
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')
part.partid = i_part
part.state = 1
part.elemid = 0
part.turn = 0
p.from_pysixtrack(part, i_part)
return ps
# In[3]:
from cpymad.madx import Madx
madx = Madx(stdout=False)
madx.call("lhc_colin_track.seq")
madx.use("lhcb1") # if not called, no beam present in sequence???
_ = madx.twiss() # if not called, cavities don't get a frequency???
# In[4]:
# Line and sixtracklib Elements
line = pysixtrack.Line.from_madx_sequence(madx.sequence["lhcb1"])
# Only the coordinates at the end of tracking. To keep coordinates at each
# turn, give "num_stores=nturns" when creating the BeamMonitor element
_ = line.append_element(
pysixtrack.elements.BeamMonitor(num_stores=1, is_rolling=True),
"turn_monitor")
elements = sixtracklib.Elements.from_line(line)
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'):
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
i_ip = list(mad.table.twiss.name).index('ip5:1')
from cpymad.madx import Madx import sixtracklib as pystlib import pysixtrack 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 ps_line, _ = pysixtrack.Line.from_madx_sequence(sis18) elements = pystlib.Elements()
import pysixtrack
import matplotlib.pyplot as plt
import os
import shutil
from cpymad.madx import Madx
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--noblock', dest='plot_block', action='store_false')
args = parser.parse_args()
mad = Madx()
mad.call('lhc2018_injection.mask')
#mad.command.esave(file='lattice_errors.err')
os.remove('db4')
os.remove('db5')
os.remove('fidel')
os.remove('slhc')
os.remove('wise')
os.remove('twiss.b1')
os.remove('twiss.b2')
shutil.rmtree('temp')
mad.use('lhcb1')
twiss = mad.twiss(sequence='lhcb1')
from cpymad.madx import Madx
mad = Madx()
#mad.chdir('./madx')
#mad.call('sps/beams/lhc_beam_injection.beamx')
mad.call('sps_thin_crabcavity.madx')
#call, file = 'sps/beams/lhc_beam_injection.beamx';
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;')
madx.input('use, sequence = PS;')
madx.input('seqedit,sequence = PS;flatten;endedit;')
madx.input('seqedit,sequence = PS;remove, element=SELECTED;endedit;')
