def test_get_all_types(self):
fm = ModelFlame(self.testfile)
etypes = {
'quadrupole', 'bpm', 'drift', 'source', 'rfcavity', 'sbend',
'orbtrim', 'solenoid', 'stripper'
}
self.assertEqual(set(fm.get_all_types()), etypes)
def test_set_machine(self):
fm_none = ModelFlame()
self.assertIsNone(fm_none.machine)
with open(self.testfile, 'rb') as f:
m = Machine(f)
fm_none.machine = m
self.assertEqual(fm_none.machine, m)
def test_get_index_by_type(self):
fm = ModelFlame(self.testfile)
m = fm.machine
all_types = fm.get_all_types()
for n in range(2, len(all_types)):
etyps = [random.choice(all_types) for _ in range(n)]
e = fm.get_index_by_type(type=etyps)
e0 = {t: m.find(type=t) for t in etyps}
self.assertEqual(e, e0)
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_configure_source1(self):
"""Update source, as well as Ion_Z (and others)
"""
latfile = self.testfile
fm = ModelFlame(lat_file=latfile)
s = generate_source(fm.bmstate)
s['properties']['IonChargeStates'] = np.asarray([0.1, ])
fm.configure(econf=s)
self.assertEqual(fm.bmstate.IonZ, np.asarray([0.1, ]))
def test_get_index_by_name(self):
fm = ModelFlame(self.testfile)
m = fm.machine
all_names = fm.get_all_names()
for n in range(2, 20):
enames = [random.choice(all_names) for _ in range(n)]
e = fm.get_index_by_name(name=enames)
e0 = {n: m.find(name=n) for n in enames}
self.assertEqual(e, e0)
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_find(self):
fm = ModelFlame(self.testfile)
m = fm.machine
all_types = fm.get_all_types()
for i in range(2, 20):
e = fm.find(m.conf(i)['name'])[0]
self.assertEqual(i, e)
for ntype in all_types:
e0 = m.find(type=ntype)
e = fm.find(type=ntype)
self.assertEqual(e, e0)
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 plot_orbit(*args, **kwargs):
"""Plot orbit (x,y) from FLAME lattice file or model instance,
first passing interested data columns as tuples, see the following
examples.
Keyword Arguments
-----------------
latfile : str
Name of FLAME lattice file.
flame_model :
ModelFlame object.
Returns
-------
r : dict
Orignal data.
Examples
--------
>>> fm = ModelFlame(latfile)
>>> # plot two curves: s v.s. x,y
>>> plot_orbit(('pos', 'xcen'), ('pos', 'ycen'), flame_model=fm)
>>>
See Also
--------
ModelFlame
"""
latfile = kwargs.get('latfile', None)
flame_model = kwargs.get('flame_model', None)
if latfile is not None:
fm = ModelFlame(lat_file=latfile)
elif flame_model is not None:
fm = flame_model
else:
print("Please define 'latfile' or 'flame_model'.")
return None
ks = flatten(args)
monitors = range(len(fm.machine))
r, s = fm.run(monitor=monitors)
data = fm.collect_data(r, **dict(zip(ks, [True]*len(ks))))
fig = plt.figure()
ax = fig.add_subplot(111)
for xlbl, ylbl in args:
ax.plot(data[xlbl], data[ylbl], label='${}$'.format(ylbl.upper()))
ax.set_xlabel('${}$'.format(xlbl.upper()))
ax.legend(loc='best')
plt.show()
return data
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_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_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_collect_data(self):
""" test_collect_data: get pos, x0, IonEk
"""
latfile = self.testfile
with open(latfile, 'rb') as f:
m0 = Machine(f)
s0 = m0.allocState({})
r0 = m0.propagate(s0, 0, 100, observe=range(100))
data0 = collect_data(r0, pos=True, x0=True, IonEk=True)
fm = ModelFlame(latfile)
r, s = fm.run(from_element=1, to_element=99, monitor=range(100))
data = fm.collect_data(r, pos=True, x0=True, IonEk=True)
self.assertEqual(data0['pos'][1:].tolist(), data['pos'].tolist())
self.assertEqual(data0['x0'][1:].tolist(), data['x0'].tolist())
self.assertEqual(data0['IonEk'][1:].tolist(), data['IonEk'].tolist())
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_transfer_matrix(self):
"""Calculate transfer matrix from A to B.
"""
cs = 0
latfile = self.testfile
fm = ModelFlame(latfile)
r, s = fm.run(from_element=1, to_element=3, monitor=[1,2,3])
s1 = r[0][-1]
s2 = r[1][-1]
s3 = r[2][-1]
m21 = fm.get_transfer_matrix(from_element=1, to_element=2,
charge_state_index=cs)
m31 = fm.get_transfer_matrix(from_element=1, to_element=3,
charge_state_index=cs)
self.assertEqual(m21.tolist(),
s2.transfer_matrix[:, :, cs].tolist())
self.assertEqual(m31.tolist(),
np.dot(
s3.transfer_matrix[:, :, cs],
s2.transfer_matrix[:, :, cs]).tolist())
for i in range(7):
self.assertAlmostEqual(np.dot(m31, s1.moment0[:,0])[i],
s3.moment0[:,0][i])
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_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 setUp(self):
testfile = os.path.join(curdir, 'lattice/test_0.lat')
self.testfile = make_latfile(testfile)
self.fm = ModelFlame(self.testfile)
import numpy as np
from math import sqrt
from flame import Machine
from flame_utils import generate_latfile
from flame_utils import BeamState
from flame_utils import ModelFlame
import logging
logging.getLogger('flame.machine').disabled = True
import ray
latfile = 'fitSolscan_20180430.lat'
fm = ModelFlame(latfile)
sols = fm.get_element(type='solenoid')
r, s = fm.run(monitor=-1)
refbg = s.bg[0]
refbg
[
s.moment0_rms[0], s.moment0_rms[2],
s.moment1[0, 2, 0] / s.moment0_rms[0] / s.moment0_rms[2]
]
def get_m1_mat(emitx, betax, alphax, emity, betay, alphay, cxy, cxyp, cyxp,
cxpyp):
emitx = emitx / 1e6
emity = emity / 1e6
m1_mat = np.eye(4)
def test_set_latfile(self):
fm_none = ModelFlame()
self.assertIsNone(fm_none.latfile)
fm_none.latfile = self.testfile
self.assertEqual(fm_none.latfile, self.testfile)
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)