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_attr_alias(self): aliases = { 'xcen_all': 'x0', 'ycen_all': 'y0', 'xpcen_all': 'xp0', 'ypcen_all': 'yp0', 'phicen_all': 'phi0', 'dEkcen_all': 'dEk0', 'xrms': 'x0_rms', 'yrms': 'y0_rms', 'xprms': 'xp0_rms', 'yprms': 'yp0_rms', 'phirms': 'phi0_rms', 'dEkrms': 'dEk0_rms', 'xcen': 'x0_env', 'ycen': 'y0_env', 'xpcen': 'xp0_env', 'ypcen': 'yp0_env', 'phicen': 'phi0_env', 'dEkcen': 'dEk0_env', 'cenvector': 'moment0_env', 'cenvector_all': 'moment0', 'rmsvector': 'moment0_rms', 'beammatrix_all': 'moment1', 'beammatrix': 'moment1_env', } ms = BeamState(latfile=self.latfile) for k,v in aliases.items(): left_val, right_val = getattr(ms, k), getattr(ms, v) if isinstance(left_val, np.ndarray): self.assertTrue(((left_val == right_val) | (np.isnan(left_val) & np.isnan(right_val))).all()) else: self.assertAlmostEqual(left_val, right_val)
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_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_transmat(self): with open(self.latfile, 'rb') as f: m = Machine(f) s = m.allocState({}) m.propagate(s, 0, 10) ms = BeamState(s) left_val = ms.transfer_matrix right_val = s.transmat self.assertTrue((left_val == right_val).all())
def test_twiss_parameter(self): ms = BeamState(latfile=self.latfile) self.assertAlmostEqual(ms.xtwsb, ms.xrms*ms.xrms/ms.xeps) self.assertAlmostEqual(ms.ytwsb, ms.yrms*ms.yrms/ms.yeps) self.assertAlmostEqual(ms.ztwsb, ms.phirms*ms.phirms/ms.zeps) self.assertAlmostEqual(ms.xtwsa, -ms.moment1_env[0, 1]/ms.xeps*1e3) self.assertAlmostEqual(ms.ytwsa, -ms.moment1_env[2, 3]/ms.yeps*1e3) self.assertAlmostEqual(ms.ztwsa, -ms.moment1_env[4, 5]/ms.zeps) ms.set_twiss('x', alpha=0.2, beta=3.0, emittance=5.0, cs=0) ms.set_twiss('y', alpha=0.2, beta=3.0, emittance=5.0, cs=0) ms.set_twiss('z', alpha=0.2, beta=3.0, emittance=5.0, cs=0) self.assertAlmostEqual(ms.xtwsa_all[0], 0.2) self.assertAlmostEqual(ms.ytwsa_all[0], 0.2) self.assertAlmostEqual(ms.ztwsa_all[0], 0.2) self.assertAlmostEqual(ms.xtwsb_all[0], 3.0) self.assertAlmostEqual(ms.ytwsb_all[0], 3.0) self.assertAlmostEqual(ms.ztwsb_all[0], 3.0) self.assertAlmostEqual(ms.xeps_all[0], 5.0) self.assertAlmostEqual(ms.yeps_all[0], 5.0) self.assertAlmostEqual(ms.zeps_all[0], 5.0) mat = twiss_to_matrix('x', 1, 2, 3) mat = twiss_to_matrix('y', 1, 2, 3, matrix=mat) mat = twiss_to_matrix('z', 1, 2, 3, matrix=mat) for i in [0, 2]: self.assertAlmostEqual(mat[i, i], 6.0) self.assertAlmostEqual(mat[i, i+1], -3e-3) self.assertAlmostEqual(mat[i+1, i], -3e-3) self.assertAlmostEqual(mat[i+1, i+1], 3e-06) self.assertAlmostEqual(mat[4, 4], 6.0) self.assertAlmostEqual(mat[4, 5], -3e0) self.assertAlmostEqual(mat[5, 4], -3e0) self.assertAlmostEqual(mat[5, 5], 3e0)
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_rms_size(self): ms = BeamState(latfile=self.latfile) for k in ('xrms', 'yrms', 'xprms', 'yprms', 'phirms', 'dEkrms'): self.assertEqual(getattr(ms, k), getattr(ms, k + '_all')[0])