def test_insert_in_modelflame(self):
latfile = self.testfile
fm = ModelFlame(latfile)
r0,s0 = fm.run(to_element=6)
econf_before_insertion = fm.get_element(index=5)[0]
total_before_insertion = len(fm.machine)
new_econf = {'index':5, 'properties':{'name':'test_drift', 'type':'drift', 'L':0.05588}}
fm.insert_element(econf=new_econf)
total_after_insertion = len(fm._mach_ins)
test_econf = fm.get_element(index=5)[0]
self.assertEqual(test_econf['index'], new_econf['index'])
self.assertEqual(test_econf['properties']['name'], new_econf['properties']['name'])
self.assertEqual(test_econf['properties']['type'], new_econf['properties']['type'])
self.assertEqual(test_econf['properties']['L'], new_econf['properties']['L'])
self.assertEqual(total_before_insertion+1, total_after_insertion)
test_econf2 = fm.get_element(index=6)[0]
self.assertEqual(test_econf2['index'], 6)
self.assertEqual(test_econf2['properties']['name'], econf_before_insertion['properties']['name'])
self.assertEqual(test_econf2['properties']['type'], econf_before_insertion['properties']['type'])
r1,s1 = fm.run(to_element=6)
compare_mstates(self, s0, s1)
def test_init_with_mix(self):
with open(self.latfile, 'rb') as f:
m = Machine(f)
ms = BeamState(machine=m, latfile=self.latfile)
s = m.allocState({})
m.propagate(s, 0, 1)
compare_mstates(self, ms, s)
def test_set_bmstate(self):
fm_none = ModelFlame()
self.assertIsNone(fm_none.bmstate)
with open(self.testfile, 'rb') as f:
m = Machine(f)
s = m.allocState({})
m.propagate(s, 0, 1)
fm_none.bmstate = s
compare_mstates(self, fm_none.bmstate, s)
def test_init_with_s2(self):
""" test_init_with_s2: s is None
"""
with open(self.latfile, 'rb') as f:
m = Machine(f)
s = m.allocState({})
m.propagate(s, 0, 1)
ms = BeamState()
ms.state = s
compare_mstates(self, ms, s)
def test_run_3(self):
""" test run_3: test initial states
"""
latfile = self.testfile
with open(latfile, 'rb') as f:
m0 = Machine(f)
s0 = m0.allocState({})
m0.propagate(s0, 0, 1)
fm = ModelFlame(latfile)
r, s = fm.run(from_element=0, to_element=0)
compare_mstates(self, s0, s)
def test_run_1(self):
""" test_run_1: propagate from the first to last, monitor None
"""
latfile = self.testfile
with open(latfile, 'rb') as f:
m0 = Machine(f)
s0 = m0.allocState({})
m0.propagate(s0, 0, len(m0))
fm = ModelFlame(latfile)
r,s = fm.run()
self.assertEqual(r, [])
compare_mstates(self, s, s0)
def test_run_6(self):
""" test_run_6: optional monitor setting 'all'
"""
latfile = self.testfile
with open(latfile, 'rb') as f:
m0 = Machine(f)
s0 = m0.allocState({})
fm = ModelFlame(latfile)
r0 = m0.propagate(s0, 0, len(m0), observe=range(len(m0)))
r,s = fm.run(monitor='all')
rs0 = [ts for (ti,ts) in r0]
rs = [ts for (ti,ts) in r]
for (is1, is2) in zip(rs0, rs):
compare_mstates(self, is1, is2)
compare_mstates(self, s, s0)
def test_run_8(self):
""" test_run_8: include_initial_state
"""
latfile = self.testfile
with open(latfile, 'rb') as f:
m0 = Machine(f)
s0 = m0.allocState({})
fm = ModelFlame(latfile)
m0.propagate(s0, 0, 1)
r0 = m0.propagate(s0, 1, len(m0), observe=range(len(m0)))
r,s = fm.run(monitor='all', include_initial_state=False)
rs0 = [ts for (ti,ts) in r0]
rs = [ts for (ti,ts) in r]
for (is1, is2) in zip(rs0, rs):
compare_mstates(self, is1, is2)
compare_mstates(self, s, s0)
def test_run_7(self):
""" test_run_7: optional monitor setting 'type'
"""
latfile = self.testfile
with open(latfile, 'rb') as f:
m0 = Machine(f)
s0 = m0.allocState({})
fm = ModelFlame(latfile)
obs = fm.get_index_by_type(type='bpm')['bpm']
r0 = m0.propagate(s0, 0, len(m0), observe=obs)
r,s = fm.run(monitor='bpm')
rs0 = [ts for (ti,ts) in r0]
rs = [ts for (ti,ts) in r]
for (is1, is2) in zip(rs0, rs):
compare_mstates(self, is1, is2)
compare_mstates(self, s, s0)
def test_run_2(self):
""" test_run_2: propagate from the first to last, monitor all BPMs
"""
latfile = self.testfile
with open(latfile, 'rb') as f:
m0 = Machine(f)
s0 = m0.allocState({})
fm = ModelFlame(latfile)
obs = fm.get_index_by_type(type='bpm')['bpm']
r0 = m0.propagate(s0, 0, len(m0), observe=obs)
r,s = fm.run(monitor=obs)
rs0 = [ts for (ti,ts) in r0]
rs = [ts for (ti,ts) in r]
for (is1, is2) in zip(rs0, rs):
compare_mstates(self, is1, is2)
compare_mstates(self, s, s0)
def test_generate_source(self):
latfile = self.testfile
fm = ModelFlame(latfile)
ms = fm.bmstate
sconf = generate_source(ms)
sconf0 = fm.get_element(type='source')[0]
compare_source_element(self, sconf, sconf0)
r0, s0 = fm.run(monitor=range(len(fm.machine)))
fm.configure(sconf)
r, s = fm.run(monitor=range(len(fm.machine)))
compare_mstates(self, s, s0)
rs0 = [ts for (ti,ts) in r0]
rs = [ts for (ti,ts) in r]
for (is1, is2) in zip(rs0, rs):
compare_mstates(self, is1, is2)
def test_reconfigure(self):
latfile = self.testfile
with open(latfile, 'rb') as f:
m0 = Machine(f)
s0 = m0.allocState({})
e_cor_idx = 10
e_name = m0.conf(e_cor_idx)['name']
m0.reconfigure(10, {'theta_x': 0.005})
r0 = m0.propagate(s0, 0, len(m0), range(len(m0)))
fm = ModelFlame(latfile)
fm.reconfigure(e_name, {'theta_x': 0.005})
r, s = fm.run(monitor=range(len(m0)))
rs0 = [ts for (ti,ts) in r0]
rs = [ts for (ti,ts) in r]
for (is1, is2) in zip(rs0, rs):
compare_mstates(self, is1, is2)
def test_run_4(self):
""" test_run_4: run and monitor from element index of 10 to 20
"""
latfile = self.testfile
with open(latfile, 'rb') as f:
m0 = Machine(f)
s0 = m0.allocState({})
m0.propagate(s0, 0, 10)
r0 = m0.propagate(s0, 10, 11, observe=range(10, 21))
fm = ModelFlame(latfile)
r, s = fm.run(from_element=10, to_element=20, monitor=range(10,21))
compare_mstates(self, s0, s)
rs0 = [ts for (ti,ts) in r0]
rs = [ts for (ti,ts) in r]
for (is1, is2) in zip(rs0, rs):
compare_mstates(self, is1, is2)
def test_configure(self):
latfile = self.testfile
with open(latfile, 'rb') as f:
m0 = Machine(f)
s0 = m0.allocState({})
e_cor_idx = 10
m0.reconfigure(10, {'theta_x': 0.005})
r0 = m0.propagate(s0, 0, len(m0), range(len(m0)))
fm = ModelFlame(latfile)
e = fm.get_element(index=10)[0]
e['properties']['theta_x'] = 0.005
fm.configure(e)
r, s = fm.run(monitor=range(len(m0)))
rs0 = [ts for (ti,ts) in r0]
rs = [ts for (ti,ts) in r]
for (is1, is2) in zip(rs0, rs):
compare_mstates(self, is1, is2)
def test_run_5(self):
""" test_run_5: using BeamState object
"""
latfile = self.testfile
with open(latfile, 'rb') as f:
m0 = Machine(f)
ms = BeamState(machine=m0)
fm = ModelFlame()
fm.bmstate = ms
fm.machine = m0
obs = fm.get_index_by_type(type='bpm')['bpm']
r,s = fm.run(monitor=obs)
s0 = m0.allocState({})
m0.propagate(s0, 0, 1)
r0 = m0.propagate(s0, 1, len(m0), observe=obs)
rs0 = [ts for (ti,ts) in r0]
rs = [ts for (ti,ts) in r]
for (is1, is2) in zip(rs0, rs):
compare_mstates(self, is1, is2)
compare_mstates(self, s, s0)
def test_init_with_s1(self):
""" test_init_with_s1: s is not None
"""
with open(self.latfile, 'rb') as f:
m = Machine(f)
s0 = m.allocState({})
s1 = s0.clone()
m.propagate(s1, 0, 1)
ms0 = BeamState(s0)
compare_mstates(self, ms0, s0)
ms1 = BeamState(s0, machine=m)
compare_mstates(self, ms1, s1)
ms1_1 = BeamState(s0, latfile=self.latfile)
compare_mstates(self, ms1_1, s1)
def test_init_machine(self):
fm_none = ModelFlame()
m, s = fm_none.init_machine(self.testfile)
fm_none.machine, fm_none.bmstate = m, s
self.assertEqual(fm_none.machine, m)
compare_mstates(self, fm_none.bmstate, s)
