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)