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)