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_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())
