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