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