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;')