def test_solenoid(self):
"Propagate a state matrix through a solenoid"
T = self.expect = numpy.asfarray([[0.1, 0, 0, 0, 0, 0, 0],
[0, 0.2, 0, 0, 0, 0, 0],
[0, 0, 0.3, 0, 0, 0, 0],
[0, 0, 0, 0.4, 0, 0, 0],
[0, 0, 0, 0, 0.5, 0, 0],
[0, 0, 0, 0, 0, 0.6, 0],
[0, 0, 0, 0, 0, 0, 0]])
self.M = Machine({
'sim_type':
'MomentMatrix',
'elements': [
{
'name': 'elem0',
'type': 'source',
'initial': T
},
{
'name': 'elem1',
'type': 'solenoid',
'L': 1.123,
'K': -1.0e3,
'B': 0.0
},
]
})
S = self.M.allocState({})
self.M.propagate(S)
assert_aequal(S.state,
[[
3.95745247e-01, 1.18242831e-02, -2.36485661e-02,
0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0
],
[
1.18242831e-02, 2.98581749e-01, -1.73472348e-18,
-2.36485661e-02, 0.00000000e+00, 0.00000000e+00, 0
],
[
-2.36485661e-02, 0.00000000e+00, 2.01418251e-01,
1.18242831e-02, 0.00000000e+00, 0.00000000e+00, 0
],
[
8.67361738e-19, -2.36485661e-02, 1.18242831e-02,
1.04254753e-01, 0.00000000e+00, 0.00000000e+00, 0
],
[
0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
0.00000000e+00, 5.00000000e-01, 0.00000000e+00, 0
],
[
0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
0.00000000e+00, 0.00000000e+00, 6.00000000e-01, 0
], [0, 0, 0, 0, 0, 0, 0]],
decimal=6)
def test_quad2(self):
"Propagate a state matrix through a defocusing quadrupole"
T = self.expect = numpy.asfarray([[0.1, 0, 0, 0, 0, 0, 0],
[0, 0.2, 0, 0, 0, 0, 0],
[0, 0, 0.3, 0, 0, 0, 0],
[0, 0, 0, 0.4, 0, 0, 0],
[0, 0, 0, 0, 0.5, 0, 0],
[0, 0, 0, 0, 0, 0.6, 0],
[0, 0, 0, 0, 0, 0, 0]])
self.M = Machine({
'sim_type':
'MomentMatrix',
'elements': [
{
'name': 'elem0',
'type': 'source',
'initial': T
},
{
'name': 'elem1',
'type': 'quadrupole',
'L': 2.0e-3,
'K': -1.1e6,
'B2': 0.0
},
]
})
S = self.M.allocState({})
self.M.propagate(S)
assert_aequal(S.state,
[[
4.63617503e+00, 4.90314048e+00, 0.00000000e+00,
0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0
],
[
4.90314048e+00, 5.18979253e+00, 0.00000000e+00,
0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0
],
[
0.00000000e+00, 0.00000000e+00, 3.47549344e-01,
-2.90072396e-02, 0.00000000e+00, 0.00000000e+00, 0
],
[
0.00000000e+00, 0.00000000e+00, -2.90072396e-02,
3.47695722e-01, 0.00000000e+00, 0.00000000e+00, 0
],
[
0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
0.00000000e+00, 5.00000000e-01, 0.00000000e+00, 0
],
[
0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
0.00000000e+00, 0.00000000e+00, 6.00000000e-01, 0
], [0, 0, 0, 0, 0, 0, 0]],
decimal=6)
def test_generic(self):
"Propagate a state matrix through a generic section"
T = self.expect = numpy.asfarray([[1, 0, 0, 0, 0, 0, 0],
[0, 1, 0, 0, 0, 0, 0],
[0, 0, 1, 0, 0, 0, 0],
[0, 0, 0, 1, 0, 0, 0],
[0, 0, 0, 0, 1, 0, 0],
[0, 0, 0, 0, 0, 1, 0],
[0, 0, 0, 0, 0, 0, 1]])
self.M = Machine({
'sim_type':
'MomentMatrix',
'elements': [{
'name': 'elem0',
'type': 'source',
'initial': T
}, {
'name': 'elem1',
'type': 'generic',
'transfer': numpy.identity(7)
}]
})
S = self.M.allocState({})
self.M.propagate(S)
assert_aequal(S.state, [[1, 0, 0, 0, 0, 0, 0], [0, 1, 0, 0, 0, 0, 0],
[0, 0, 1, 0, 0, 0, 0], [0, 0, 0, 1, 0, 0, 0],
[0, 0, 0, 0, 1, 0, 0], [0, 0, 0, 0, 0, 1, 0],
[0, 0, 0, 0, 0, 0, 1]])
def test_quad1(self):
"Propagate a state matrix through a focusing quadrupole"
T = self.expect = numpy.asfarray([[0.1, 0, 0, 0, 0, 0, 0],
[0, 0.2, 0, 0, 0, 0, 0],
[0, 0, 0.3, 0, 0, 0, 0],
[0, 0, 0, 0.4, 0, 0, 0],
[0, 0, 0, 0, 0.5, 0, 0],
[0, 0, 0, 0, 0, 0.6, 0],
[0, 0, 0, 0, 0, 0, 0]])
self.M = Machine({
'sim_type':
'MomentMatrix',
'elements': [
{
'name': 'elem0',
'type': 'source',
'initial': T
},
{
'name': 'elem1',
'type': 'quadrupole',
'L': 2.0e-3,
'K': 1.1e6,
'B2': 0.0
},
]
})
S = self.M.allocState({})
self.M.propagate(S)
assert_aequal(S.state,
[[
1.61134871e-01, -3.72950224e-02, 0.00000000e+00,
0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0
],
[
-3.72950224e-02, 1.32751642e-01, 0.00000000e+00,
0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0
],
[
0.00000000e+00, 0.00000000e+00, 1.09819606e+01,
1.15461050e+01, 0.00000000e+00, 0.00000000e+00, 0
],
[
0.00000000e+00, 0.00000000e+00, 1.15461050e+01,
1.21501566e+01, 0.00000000e+00, 0.00000000e+00, 0
],
[
0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
0.00000000e+00, 5.00000000e-01, 0.00000000e+00, 0
],
[
0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
0.00000000e+00, 0.00000000e+00, 6.00000000e-01, 0
], [0, 0, 0, 0, 0, 0, 0]],
decimal=6)
def testCopySubStruct(self):
A = _Value(_Type([
('x', ('S', None, [
('y', 'i'),
('z', 'ai'),
('q', 'v'),
('m', 'av'),
('a', ('S', None, [
('A', 'i'),
])),
])),
]), {
'x.y':42,
'x.z':range(3),
'x.q':'hello',
'x.m':['hello', 52],
'x.a.A':100,
})
B = _Value(A.type(), {
'x.y':43,
'x.z':range(4),
'x.q':15,
'x.m':['world', 62],
'x.a.A':101,
})
print(A.x, B.x)
B.x = A.x
self.assertEqual(B.x.y, 42)
assert_aequal(B.x.z, [0,1,2])
self.assertEqual(B.x.q, 'hello')
self.assertEqual(B.x.m, ['hello', 52])
self.assertEqual(B.x.a.A, 100)
def test_sbend1(self):
"Propagate a state matrix through a focusing sector bend"
T = self.expect = numpy.asfarray([[0.1, 0, 0, 0, 0, 0, 0],
[0, 0.2, 0, 0, 0, 0, 0],
[0, 0, 0.3, 0, 0, 0, 0],
[0, 0, 0, 0.4, 0, 0, 0],
[0, 0, 0, 0, 0.5, 0, 0],
[0, 0, 0, 0, 0, 0.6, 0],
[0, 0, 0, 0, 0, 0, 0]])
self.M = Machine({
'sim_type':
'MomentMatrix',
'elements': [
{
'name': 'elem0',
'type': 'source',
'initial': T
},
{
'name': 'elem1',
'type': 'sbend',
'L': 2.0e-3,
'phi': math.radians(25.0),
'K': 1.0e6
},
]
})
S = self.M.allocState({})
self.M.propagate(S)
assert_aequal(S.state,
[[
1.71805650e-01, -3.79121906e-02, 0.00000000e+00,
0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0
],
[
-3.79121906e-02, 1.24776654e-01, 0.00000000e+00,
0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0
],
[
0.00000000e+00, 0.00000000e+00, 9.50788149e+00,
9.55147102e+00, 0.00000000e+00, 0.00000000e+00, 0
],
[
0.00000000e+00, 0.00000000e+00, 9.55147102e+00,
9.60788149e+00, 0.00000000e+00, 0.00000000e+00, 0
],
[
0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
0.00000000e+00, 5.00000000e-01, 0.00000000e+00, 0
],
[
0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
0.00000000e+00, 0.00000000e+00, 6.00000000e-01, 0
], [0, 0, 0, 0, 0, 0, 0]],
decimal=6)
def test_sbend2(self):
"Propagate a state matrix through a defocusing sector bend"
T = self.expect = numpy.asfarray([[0.1, 0, 0, 0, 0, 0, 0],
[0, 0.2, 0, 0, 0, 0, 0],
[0, 0, 0.3, 0, 0, 0, 0],
[0, 0, 0, 0.4, 0, 0, 0],
[0, 0, 0, 0, 0.5, 0, 0],
[0, 0, 0, 0, 0, 0.6, 0],
[0, 0, 0, 0, 0, 0, 0]])
self.M = Machine({
'sim_type':
'MomentMatrix',
'elements': [
{
'name': 'elem0',
'type': 'source',
'initial': T
},
{
'name': 'elem1',
'type': 'sbend',
'L': 2.0e-3,
'phi': math.radians(25.0),
'K': -1.0e6
},
]
})
S = self.M.allocState({})
self.M.propagate(S)
assert_aequal(S.state,
[[
3.78918063e+00, 3.74852747e+00, 0.00000000e+00,
0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0
],
[
3.74852747e+00, 3.71358865e+00, 0.00000000e+00,
0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0
],
[
0.00000000e+00, 0.00000000e+00, 3.82682181e-01,
-3.78401248e-02, 0.00000000e+00, 0.00000000e+00, 0
],
[
0.00000000e+00, 0.00000000e+00, -3.78401248e-02,
3.17317819e-01, 0.00000000e+00, 0.00000000e+00, 0
],
[
0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
0.00000000e+00, 5.00000000e-01, 0.00000000e+00, 0
],
[
0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
0.00000000e+00, 0.00000000e+00, 6.00000000e-01, 0
], [0, 0, 0, 0, 0, 0, 0]],
decimal=6)
def test_unwrap_3d(self):
pixels = numpy.array([[[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11]],
[[12, 13, 14, 15],
[16, 17, 18, 19],
[20, 21, 22, 23]]])
V = Value(nt.NTNDArray.buildType(), {
'value': numpy.arange(24),
'dimension': [
{'size': 4}, # X, columns
{'size': 3}, # Y, rows
{'size': 2}, # "color"
],
'attribute': [
{'name': 'ColorMode', 'value': 4},
],
})
img = nt.NTNDArray.unwrap(V)
self.assertEqual(img.shape, (2, 3, 4))
assert_aequal(img, pixels)
V2 = nt.NTNDArray().wrap(img)
assert_aequal(V.value, V2.value)
self.assertEqual(V.dimension[0].size, V2.dimension[0].size)
self.assertEqual(V.dimension[1].size, V2.dimension[1].size)
self.assertEqual(V.dimension[2].size, V2.dimension[2].size)
def test_solenoid(self):
"Propagate a state matrix through a solenoid"
T = self.expect = numpy.asfarray([
[0.1, 0, 0, 0, 0, 0, 0],
[0, 0.2, 0, 0, 0, 0, 0],
[0, 0, 0.3, 0, 0, 0, 0],
[0, 0, 0, 0.4, 0, 0, 0],
[0, 0, 0, 0, 0.5, 0, 0],
[0, 0, 0, 0, 0, 0.6, 0],
[0, 0, 0, 0, 0, 0, 0]
])
self.M = Machine({
'sim_type':'MomentMatrix',
'elements':[
{'name':'elem0', 'type':'source', 'initial':T},
{'name':'elem1', 'type':'solenoid', 'L':1.123, 'K':-1.0e3, 'B':0.0},
]
})
S = self.M.allocState({})
self.M.propagate(S)
assert_aequal(S.state, [
[ 3.95745247e-01, 1.18242831e-02, -2.36485661e-02, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0],
[ 1.18242831e-02, 2.98581749e-01, -1.73472348e-18, -2.36485661e-02, 0.00000000e+00, 0.00000000e+00, 0],
[-2.36485661e-02, 0.00000000e+00, 2.01418251e-01, 1.18242831e-02, 0.00000000e+00, 0.00000000e+00, 0],
[ 8.67361738e-19, -2.36485661e-02, 1.18242831e-02, 1.04254753e-01, 0.00000000e+00, 0.00000000e+00, 0],
[ 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 5.00000000e-01, 0.00000000e+00, 0],
[ 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 6.00000000e-01, 0],
[ 0, 0, 0, 0, 0, 0, 0]
], decimal=6)
def test_generic(self):
"Propagate a state matrix through a generic section"
T = self.expect = numpy.asfarray([
[1, 0, 0, 0, 0, 0, 0],
[0, 1, 0, 0, 0, 0, 0],
[0, 0, 1, 0, 0, 0, 0],
[0, 0, 0, 1, 0, 0, 0],
[0, 0, 0, 0, 1, 0, 0],
[0, 0, 0, 0, 0, 1, 0],
[0, 0, 0, 0, 0, 0, 1]
])
self.M = Machine({
'sim_type':'MomentMatrix',
'elements':[
{'name':'elem0', 'type':'source', 'initial':T},
{'name':'elem1', 'type':'generic', 'transfer':numpy.identity(7)}
]
})
S = self.M.allocState({})
self.M.propagate(S)
assert_aequal(S.state, [
[1, 0, 0, 0, 0, 0, 0],
[0, 1, 0, 0, 0, 0, 0],
[0, 0, 1, 0, 0, 0, 0],
[0, 0, 0, 1, 0, 0, 0],
[0, 0, 0, 0, 1, 0, 0],
[0, 0, 0, 0, 0, 1, 0],
[0, 0, 0, 0, 0, 0, 1]
])
def test_unwrap_None(self):
V = Value(nt.NTNDArray.buildType(), {})
img = nt.NTNDArray.unwrap(V)
self.assertIsInstance(img, numpy.ndarray)
assert_aequal(img.shape, (0, ))
def test_drift(self):
"""Propogate an identity matrix to find
the cumulative transfer matrix of two drifts
"""
S = self.M.allocState({})
assert_aequal(S.state, [
[1, 0, 0, 0, 0, 0],
[0, 1, 0, 0, 0, 0],
[0, 0, 1, 0, 0, 0],
[0, 0, 0, 1, 0, 0],
[0, 0, 0, 0, 1, 0],
[0, 0, 0, 0, 0, 1],
])
self.M.propagate(S)
assert_aequal(S.state, [
[1, 2, 0, 0, 0, 0],
[0, 1, 0, 0, 0, 0],
[0, 0, 1, 2, 0, 0],
[0, 0, 0, 1, 0, 0],
[0, 0, 0, 0, 1, 0],
[0, 0, 0, 0, 0, 1],
])
def test_drift(self):
"Propagate a state matrix through a drift"
T = self.expect = numpy.asfarray([
[0.1, 0, 0, 0, 0, 0, 0],
[0, 0.2, 0, 0, 0, 0, 0],
[0, 0, 0.3, 0, 0, 0, 0],
[0, 0, 0, 0.4, 0, 0, 0],
[0, 0, 0, 0, 0.5, 0, 0],
[0, 0, 0, 0, 0, 0.6, 0],
[0, 0, 0, 0, 0, 0, 0]
])
self.M = Machine({
'sim_type':'MomentMatrix',
'elements':[
{'name':'elem0', 'type':'source', 'initial':T},
{'name':'elem1', 'type':'drift', 'L':1.234},
]
})
S = self.M.allocState({})
self.M.propagate(S)
assert_aequal(S.state, [
[ 3.045513e+05, 2.468000e+02, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0],
[ 2.468000e+02, 2.000000e-01, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0],
[ 0.000000e+00, 0.000000e+00, 6.091027e+05, 4.936000e+02, 0.000000e+00, 0.000000e+00, 0],
[ 0.000000e+00, 0.000000e+00, 4.936000e+02, 4.000000e-01, 0.000000e+00, 0.000000e+00, 0],
[ 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 5.000000e-01, 0.000000e+00, 0],
[ 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 6.000000e-01, 0],
[ 0, 0, 0, 0, 0, 0, 0]
])
def test_matrix(self):
T = numpy.asfarray([
[1, 0, 1, 0, 1, 0],
[0, 1, 0, 1, 0, 1],
[1, 0, 1, 0, 1, 0],
[0, 1, 0, 1, 0, 1],
[1, 0, 1, 0, 1, 0],
[0, 1, 0, 1, 0, 1],
])
M = Machine({
'sim_type':
'TransferMatrix',
'elements': [
{
'name': 'elem0',
'type': 'source',
'initial': T
},
],
})
S = M.allocState({})
M.propagate(S)
assert_aequal(S.state, T)
def test_array_wrap(self):
NT = nt.NTScalar('ad') # array of double
A = numpy.asarray([1.0, 5.0])
V = NT.wrap(A)
assert_aequal(V.value, A)
self.assertEqual(V.alarm.severity, 0)
def testVariantUnion(self):
V = _Value(_Type([
('x', 'v'),
]))
self.assertIsNone(V.x)
V.x = 5
self.assertEqual(V.x, 5)
V.x = 1.1
self.assertEqual(V.x, 1.1)
V.x = 'test'
self.assertEqual(V.x, u'test')
V.x = u'test'
self.assertEqual(V.x, u'test')
V.x = np.asarray([1, 2])
assert_aequal(V.x, np.asarray([1,2]))
V.x = np.asfarray([1, 2])
assert_aequal(V.x, np.asfarray([1,2]))
# clearing unions is broken prior to 7.0.0
if pvdVersion()>=(7,0,0,0):
V.x = None
self.assertIsNone(V.x)
def test_quad2(self):
"Propagate a state matrix through a defocusing quadrupole"
T = self.expect = numpy.asfarray([
[0.1, 0, 0, 0, 0, 0, 0],
[0, 0.2, 0, 0, 0, 0, 0],
[0, 0, 0.3, 0, 0, 0, 0],
[0, 0, 0, 0.4, 0, 0, 0],
[0, 0, 0, 0, 0.5, 0, 0],
[0, 0, 0, 0, 0, 0.6, 0],
[0, 0, 0, 0, 0, 0, 0]
])
self.M = Machine({
'sim_type':'MomentMatrix',
'elements':[
{'name':'elem0', 'type':'source', 'initial':T},
{'name':'elem1', 'type':'quadrupole', 'L':2.0e-3, 'K':-1.1e6, 'B2':0.0},
]
})
S = self.M.allocState({})
self.M.propagate(S)
assert_aequal(S.state, [
[ 4.63617503e+00, 4.90314048e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0],
[ 4.90314048e+00, 5.18979253e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0],
[ 0.00000000e+00, 0.00000000e+00, 3.47549344e-01, -2.90072396e-02, 0.00000000e+00, 0.00000000e+00, 0],
[ 0.00000000e+00, 0.00000000e+00, -2.90072396e-02, 3.47695722e-01, 0.00000000e+00, 0.00000000e+00, 0],
[ 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 5.00000000e-01, 0.00000000e+00, 0],
[ 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 6.00000000e-01, 0],
[ 0, 0, 0, 0, 0, 0, 0]
], decimal=6)
def test_drift(self):
"""Propogate an identity matrix to find
the cumulative transfer matrix of two drifts
"""
S = self.M.allocState({})
assert_aequal(
S.state,
[
[1, 0, 0, 0, 0, 0],
[0, 1, 0, 0, 0, 0],
[0, 0, 1, 0, 0, 0],
[0, 0, 0, 1, 0, 0],
[0, 0, 0, 0, 1, 0],
[0, 0, 0, 0, 0, 1],
],
)
self.M.propagate(S)
assert_aequal(
S.state,
[
[1, 2, 0, 0, 0, 0],
[0, 1, 0, 0, 0, 0],
[0, 0, 1, 2, 0, 0],
[0, 0, 0, 1, 0, 0],
[0, 0, 0, 0, 1, 0],
[0, 0, 0, 0, 0, 1],
],
)
def test_zero_length(self):
V = Value(
nt.NTNDArray.buildType(),
{
'value': numpy.arange(0),
'dimension': [
{
'size': 3
}, # X, columns
{
'size': 0
}, # Y, rows
],
'attribute': [
{
'name': 'ColorMode',
'value': 0
},
],
})
img = nt.NTNDArray.unwrap(V)
self.assertIsInstance(img, numpy.ndarray)
assert_aequal(img.shape, (0, 3))
def test_sbend1(self):
"Propagate a state matrix through a focusing sector bend"
T = self.expect = numpy.asfarray([
[0.1, 0, 0, 0, 0, 0, 0],
[0, 0.2, 0, 0, 0, 0, 0],
[0, 0, 0.3, 0, 0, 0, 0],
[0, 0, 0, 0.4, 0, 0, 0],
[0, 0, 0, 0, 0.5, 0, 0],
[0, 0, 0, 0, 0, 0.6, 0],
[0, 0, 0, 0, 0, 0, 0]
])
self.M = Machine({
'sim_type':'MomentMatrix',
'elements':[
{'name':'elem0', 'type':'source', 'initial':T},
{'name':'elem1', 'type':'sbend', 'L':2.0e-3, 'phi':math.radians(25.0),
'K':1.0e6},
]
})
S = self.M.allocState({})
self.M.propagate(S)
assert_aequal(S.state, [
[ 1.71805650e-01, -3.79121906e-02, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0],
[-3.79121906e-02, 1.24776654e-01, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0],
[ 0.00000000e+00, 0.00000000e+00, 9.50788149e+00, 9.55147102e+00, 0.00000000e+00, 0.00000000e+00, 0],
[ 0.00000000e+00, 0.00000000e+00, 9.55147102e+00, 9.60788149e+00, 0.00000000e+00, 0.00000000e+00, 0],
[ 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 5.00000000e-01, 0.00000000e+00, 0],
[ 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 6.00000000e-01, 0],
[ 0, 0, 0, 0, 0, 0, 0]
], decimal=6)
def test_sbend2(self):
"Propagate a state matrix through a defocusing sector bend"
T = self.expect = numpy.asfarray([
[0.1, 0, 0, 0, 0, 0, 0],
[0, 0.2, 0, 0, 0, 0, 0],
[0, 0, 0.3, 0, 0, 0, 0],
[0, 0, 0, 0.4, 0, 0, 0],
[0, 0, 0, 0, 0.5, 0, 0],
[0, 0, 0, 0, 0, 0.6, 0],
[0, 0, 0, 0, 0, 0, 0]
])
self.M = Machine({
'sim_type':'MomentMatrix',
'elements':[
{'name':'elem0', 'type':'source', 'initial':T},
{'name':'elem1', 'type':'sbend', 'L':2.0e-3, 'phi':math.radians(25.0),
'K':-1.0e6},
]
})
S = self.M.allocState({})
self.M.propagate(S)
assert_aequal(S.state, [
[ 3.78918063e+00, 3.74852747e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0],
[ 3.74852747e+00, 3.71358865e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0],
[ 0.00000000e+00, 0.00000000e+00, 3.82682181e-01, -3.78401248e-02, 0.00000000e+00, 0.00000000e+00, 0],
[ 0.00000000e+00, 0.00000000e+00, -3.78401248e-02, 3.17317819e-01, 0.00000000e+00, 0.00000000e+00, 0],
[ 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 5.00000000e-01, 0.00000000e+00, 0],
[ 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 6.00000000e-01, 0],
[ 0, 0, 0, 0, 0, 0, 0]
], decimal=6)
def test_quad1(self):
"Propagate a state matrix through a focusing quadrupole"
T = self.expect = numpy.asfarray([
[0.1, 0, 0, 0, 0, 0, 0],
[0, 0.2, 0, 0, 0, 0, 0],
[0, 0, 0.3, 0, 0, 0, 0],
[0, 0, 0, 0.4, 0, 0, 0],
[0, 0, 0, 0, 0.5, 0, 0],
[0, 0, 0, 0, 0, 0.6, 0],
[0, 0, 0, 0, 0, 0, 0]
])
self.M = Machine({
'sim_type':'MomentMatrix',
'elements':[
{'name':'elem0', 'type':'source', 'initial':T},
{'name':'elem1', 'type':'quadrupole', 'L':2.0e-3, 'K':1.1e6, 'B2':0.0},
]
})
S = self.M.allocState({})
self.M.propagate(S)
assert_aequal(S.state, [
[ 1.61134871e-01, -3.72950224e-02, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0],
[-3.72950224e-02, 1.32751642e-01, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0],
[ 0.00000000e+00, 0.00000000e+00, 1.09819606e+01, 1.15461050e+01, 0.00000000e+00, 0.00000000e+00, 0],
[ 0.00000000e+00, 0.00000000e+00, 1.15461050e+01, 1.21501566e+01, 0.00000000e+00, 0.00000000e+00, 0],
[ 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 5.00000000e-01, 0.00000000e+00, 0],
[ 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 6.00000000e-01, 0],
[ 0, 0, 0, 0, 0, 0, 0]
], decimal=6)
def test_array_unwrap(self):
NT = nt.NTScalar('ad') # array of double
A = numpy.asarray(range(10))[2:5]
V = NT.wrap(A)
P = nt.NTScalar.unwrap(V)
self.assertIsNot(P, A)
assert_aequal(P, A)
self.assertEqual(P.severity, 0)
def test_generic(self):
"Propogate a state matrix through a generic section"
S = self.M.allocState({})
self.M.propagate(S)
assert_aequal(S.moment0, self.expect0)
assert_aequal(S.state, self.expect)
def test_matrix(self):
T = numpy.asfarray([1, 0, 1, 0, 1, 0])
M = Machine({"sim_type": "Vector", "elements": [{"name": "elem0", "type": "source", "initial": T}]})
S = M.allocState({})
M.propagate(S)
assert_aequal(S.state, T)
def test_drift(self):
"Propagate a state matrix through a drift"
T = self.expect = numpy.asfarray([[0.1, 0, 0, 0, 0, 0, 0],
[0, 0.2, 0, 0, 0, 0, 0],
[0, 0, 0.3, 0, 0, 0, 0],
[0, 0, 0, 0.4, 0, 0, 0],
[0, 0, 0, 0, 0.5, 0, 0],
[0, 0, 0, 0, 0, 0.6, 0],
[0, 0, 0, 0, 0, 0, 0]])
self.M = Machine({
'sim_type':
'MomentMatrix',
'elements': [
{
'name': 'elem0',
'type': 'source',
'initial': T
},
{
'name': 'elem1',
'type': 'drift',
'L': 1.234
},
]
})
S = self.M.allocState({})
self.M.propagate(S)
assert_aequal(S.state, [[
3.045513e+05, 2.468000e+02, 0.000000e+00, 0.000000e+00,
0.000000e+00, 0.000000e+00, 0
],
[
2.468000e+02, 2.000000e-01, 0.000000e+00,
0.000000e+00, 0.000000e+00, 0.000000e+00, 0
],
[
0.000000e+00, 0.000000e+00, 6.091027e+05,
4.936000e+02, 0.000000e+00, 0.000000e+00, 0
],
[
0.000000e+00, 0.000000e+00, 4.936000e+02,
4.000000e-01, 0.000000e+00, 0.000000e+00, 0
],
[
0.000000e+00, 0.000000e+00, 0.000000e+00,
0.000000e+00, 5.000000e-01, 0.000000e+00, 0
],
[
0.000000e+00, 0.000000e+00, 0.000000e+00,
0.000000e+00, 0.000000e+00, 6.000000e-01, 0
], [0, 0, 0, 0, 0, 0, 0]])
def test_expected(self):
"""Test that the expected strength actually results in the expected output state
"""
S = self.M.allocState({})
S.state[:] = [1, 1e-3, 1, 1e-3, 1, 1e-3]
self.M.reconfigure(1, {"L": 1.0e-1, "phi": 1.0e-6, "K": self._expect_K})
self.M.propagate(S)
assert_aequal(S.state, self._expected, 1.0e8)
def test_wrap(self):
NT = nt.NTTable(columns=[
('a', 'i'),
('b', 's'),
])
V = NT.wrap([
{'a': 5, 'b': 'one'},
{'a': 6, 'b': 'two'},
])
assert_aequal(V.value.a, [5, 6])
self.assertEqual(V.value.b, ['one', 'two'])
def testArray(self):
V = _Value(_Type([
('ival', 'ai'),
('dval', 'ad'),
('sval', 'as'),
]), {
'ival': [1, 2, 3],
'dval': np.asfarray([1.1, 2.2]),
'sval': ['a', u'b'],
})
assert_aequal(V.ival, np.asarray([1, 2, 3]))
assert_aequal(V.dval, np.asfarray([1.1, 2.2]))
self.assertListEqual(V.sval, [u'a', u'b'])
def test_matrix(self):
T = numpy.asfarray([1, 0, 1, 0, 1, 0])
M = Machine({
'sim_type':'Vector',
'elements':[
{'name':'elem0', 'type':'source', 'initial':T},
],
})
S = M.allocState({})
M.propagate(S)
assert_aequal(S.state, T)
def test_expected(self):
"""Test that the expected strength actually results in the expected output state
"""
S = self.M.allocState({})
S.state[:] = [1, 1e-3, 1, 1e-3, 1, 1e-3]
self.M.reconfigure(1, {
'L': 1.0e-1,
'phi': 1.0e-6,
'K': self._expect_K,
})
self.M.propagate(S)
assert_aequal(S.state, self._expected, 1.0e8)
def test_expected(self):
"""Test that the expected strength actually results in the expected output state
"""
S = self.M.allocState({})
S.state = [1, 1e-3, 1, 1e-3, 1, 1e-3]
self.M.reconfigure(1, {
'L':1.0e-1,
'phi':1.0e-6,
'K':self._expect_K,
})
self.M.propagate(S)
assert_aequal(S.state, self._expected, decimal=8)
def testVariantUnionRoundTrip(self):
"Ensure type is preserved"
T = Type([
('x', 'v'),
])
V = T()
for X in (
None,
True,
42,
4.2,
u"hello",
[u"hello", u"world"],
):
V.x = X
self.assertEqual(X, V.x)
self.assertIs(type(X), type(V.x))
V2 = T()
V2[None] = V
self.assertEqual(X, V2.x)
self.assertIs(type(X), type(V2.x))
for X in (
np.asarray([True, False]),
np.asarray([-0x80, 0x7f], dtype='i1'),
np.asarray([-0x8000, 0x7fff], dtype='i2'),
np.asarray([-0x80000000, 0x7fffffff], dtype='i4'),
np.asarray([-0x8000000000000000, 0x7fffffffffffffff],
dtype='i8'),
np.asarray([0xff], dtype='u1'),
np.asarray([0xffff], dtype='u2'),
np.asarray([0xffffffff], dtype='u4'),
np.asarray([0xffffffffffffffff], dtype='u8'),
# float test should not be sensitive to roundoff
np.asarray([1.23456789], dtype='f4'),
np.asarray([1.23456789e100], dtype='f8'),
):
try:
V.x = X
assert_aequal(V.x, X)
self.assertEqual(V.x.dtype, X.dtype)
V2 = T()
V2[None] = V
assert_aequal(V2.x, X)
self.assertEqual(V2.x.dtype, X.dtype)
except RuntimeError as e:
raise RuntimeError("With %s : %s" % (X, e))
def test_drift(self):
"Propogate a state vector through a drift section"
S = self.M.allocState({})
self.assertEqual(S.next_elem, 0)
S.state[:] = [1, 1e-3, 0, 0, 0, 0]
assert_aequal(S.state, [1.000, 1e-3, 0, 0, 0, 0])
self.M.propagate(S)
self.assertEqual(S.next_elem, 1)
assert_aequal(S.state, [1.001, 1e-3, 0, 0, 0, 0])
S.next_elem = 2
self.assertEqual(S.next_elem, 2)
def test_drift(self):
"Propogate a state vector through a drift section"
S = self.M.allocState({})
self.assertEqual(S.next_elem, 0)
S.state[:] = [1, 1e-3, 0, 0, 0, 0]
assert_aequal(S.state, [1.000, 1e-3, 0, 0, 0, 0])
self.M.propagate(S)
self.assertEqual(S.next_elem, 1)
assert_aequal(S.state, [1.001, 1e-3, 0, 0, 0, 0])
S.next_elem = 2
self.assertEqual(S.next_elem, 2)
def testVariantUnionExplicit(self):
V = Value(Type([
('x', 'v'),
]))
V.x = ('B', 0xffff)
self.assertEqual(V.x, 0xff)
V.x = ('d', 42)
self.assertEqual(V.x, 42.0)
self.assertEqual(type(V.x), float)
V.x = ('ad', np.asarray([1, 2]))
assert_aequal(V.x, np.asfarray([1, 2]))
V.x = ('as', ['one', 'two'])
self.assertEqual(V.x, ['one', 'two'])
def test_all(self):
S = self.M.allocState({})
S.state[:] = [0, 0, 1, 1e-3, 0, 0]
results = self.M.propagate(S, observe=range(5))
self.assertIsNot(results, None)
self.assertEqual(results[0][0], 0)
self.assertEqual(results[1][0], 1)
self.assertEqual(results[2][0], 2)
self.assertEqual(results[3][0], 3)
self.assertEqual(results[4][0], 4)
assert_aequal(results[0][1].state, [0, 0, 1.002, 1e-3, 0, 0])
assert_aequal(results[1][1].state, [0, 0, 1.004, 1e-3, 0, 0])
assert_aequal(results[2][1].state, [0, 0, 1.006, 1e-3, 0, 0])
assert_aequal(results[3][1].state, [0, 0, 1.008, 1e-3, 0, 0])
assert_aequal(results[4][1].state, [0, 0, 1.010, 1e-3, 0, 0])
def test_generic(self):
T = numpy.asfarray(
[
[1, 0, 1, 0, 1, 0],
[0, 1, 0, 1, 0, 1],
[1, 0, 1, 0, 1, 0],
[0, 1, 0, 1, 0, 1],
[1, 0, 1, 0, 1, 0],
[0, 1, 0, 1, 0, 1],
]
)
M = Machine({"sim_type": "TransferMatrix", "elements": [{"name": "elem0", "type": "generic", "transfer": T}]})
S = M.allocState({})
M.propagate(S)
assert_aequal(S.state, T)
def test_all(self):
S = self.M.allocState({})
S.state[:] = [0, 0, 1, 1e-3, 0, 0]
results = self.M.propagate(S, observe=range(5))
self.assertIsNot(results, None)
self.assertEqual(results[0][0], 0)
self.assertEqual(results[1][0], 1)
self.assertEqual(results[2][0], 2)
self.assertEqual(results[3][0], 3)
self.assertEqual(results[4][0], 4)
assert_aequal(results[0][1].state, [0, 0, 1.002, 1e-3, 0, 0])
assert_aequal(results[1][1].state, [0, 0, 1.004, 1e-3, 0, 0])
assert_aequal(results[2][1].state, [0, 0, 1.006, 1e-3, 0, 0])
assert_aequal(results[3][1].state, [0, 0, 1.008, 1e-3, 0, 0])
assert_aequal(results[4][1].state, [0, 0, 1.010, 1e-3, 0, 0])
def test_source(self):
"Initialize w/ all zeros, then propagate through source element to overwrite"
C = self.M.conf()
S = self.M.allocState({}, inherit=False)
self.M.propagate(S, max=1) # propagate through source element
self.assertEqual(S.pos, 0.0)
self.assertEqual(S.IonEs[0], C['IonEs'])
self.assertEqual(S.IonW[0], C['IonEk']+C['IonEs'])
self.assertEqual(S.IonEk[0], C['IonEk'])
self.assertEqual(S.gamma[0], (C['IonEk']+C['IonEs'])/C['IonEs'])
self.assertAlmostEqual(S.beta[0], sqrt(1-1/(S.gamma[0]**2)))
print("moment0", S.moment0_env, C['IV'])
assert_aequal(S.moment0_env, C['IV'])
print("state", S.moment1_env, C['IM'])
assert_aequal(S.moment1_env, C['IM'].reshape((7,7)), 1e10)
def test_parse(self):
"Test global scope when parsing"
M = Machine(b"""
sim_type = "Vector";
L = 2.0e-3;
elem0: drift;
elem1: drift;
elem2: drift;
foo: LINE = (elem0, elem1, elem2);
""")
S = M.allocState({})
S.state = [0, 0, 1, 1e-3, 0, 0]
assert_aequal(S.state, [0, 0, 1.000, 1e-3, 0, 0])
M.propagate(S)
assert_aequal(S.state, [0, 0, 1.006, 1e-3, 0, 0])
def test_parse(self):
"Test global scope when parsing"
M = Machine(b"""
sim_type = "Vector";
L = 2.0e-3;
elem0: drift;
elem1: drift;
elem2: drift;
foo: LINE = (elem0, elem1, elem2);
""")
S = M.allocState({})
S.state[:] = [0, 0, 1, 1e-3, 0, 0]
assert_aequal(S.state, [0, 0, 1.000, 1e-3, 0, 0])
M.propagate(S)
assert_aequal(S.state, [0, 0, 1.006, 1e-3, 0, 0])
def test_lat(self):
# Propagate a state matrix through a lattice defined by a lattice file.
P = GLPSParser()
with open(os.path.join(datadir, 'moment_jb.lat'), 'rb') as inf:
self.M = Machine(inf.read())
S = self.M.allocState({})
self.M.propagate(S)
assert_aequal(S.state, [
[ 3.15034165e+00, 9.27334719e-03, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0],
[ 9.27334719e-03, 6.37581482e-03, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0],
[ 0.00000000e+00, 0.00000000e+00, 1.59854625e+01, 4.96594365e+00, 0.00000000e+00, 0.00000000e+00, 0],
[ 0.00000000e+00, 0.00000000e+00, 4.96594365e+00, 1.55019577e+00, 0.00000000e+00, 0.00000000e+00, 0],
[ 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 5.00000000e-01, 0.00000000e+00, 0],
[ 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 6.00000000e-01, 0],
[ 0, 0, 0, 0, 0, 0, 0]
], decimal=6)
def test_generic(self):
T = numpy.asfarray([
[1, 0, 1, 0, 1, 0],
[0, 1, 0, 1, 0, 1],
[1, 0, 1, 0, 1, 0],
[0, 1, 0, 1, 0, 1],
[1, 0, 1, 0, 1, 0],
[0, 1, 0, 1, 0, 1],
])
M = Machine({
'sim_type':'TransferMatrix',
'elements':[
{'name':'elem0', 'type':'generic', 'transfer':T},
],
})
S = M.allocState({})
M.propagate(S)
assert_aequal(S.state, T)
def test_parse(self):
"Test global scope when parsing"
# Propagate a state matrix through a lattice defined by a lattice file.
P = GLPSParser()
with open(os.path.join(datadir, 'moment_jb_2.lat'), 'rb') as inf:
self.M = Machine(inf.read())
S = self.M.allocState({})
self.M.propagate(S)
assert_aequal(S.state, [
[ 7.50527918e+00, 4.29398123e-03, 1.26662933e-02, -1.18523706e-06, 3.17367476e-04, -2.34554323e-05, 0],
[ 4.29398123e-03, 3.84946607e-06, -2.41285106e-05, -1.85409878e-08, 1.06777952e-07, 5.28564016e-09, 0],
[ 1.26662933e-02, -2.41285106e-05, 8.08978916e+00, 5.33808178e-03, -1.19544245e-03, 7.51043870e-05, 0],
[-1.18523706e-06, -1.85409878e-08, 5.33808178e-03, 4.89711389e-06, -5.01614882e-07, 5.57484222e-08, 0],
[ 3.17367476e-04, 1.06777952e-07, -1.19544245e-03, -5.01614882e-07, 6.71687185e-04, -1.23223342e-05, 0],
[-2.34554323e-05, 5.28564016e-09, 7.51043870e-05, 5.57484222e-08, -1.23223342e-05, 1.99524669e-06, 0],
[ 0, 0, 0, 0, 0, 0, 0]
], decimal=6)
def test_unwrap_mono(self):
pixels = numpy.asarray([ # 2x3
[0, 1, 2],
[3, 4, 5],
])
# check my understanding of numpy
self.assertTupleEqual(pixels.shape, (2, 3)) # inner-most right (in a pixel loop)
assert_aequal(pixels.flatten(), [0, 1, 2, 3, 4, 5]) # row major
V = Value(nt.NTNDArray.buildType(), {
'value': numpy.arange(6),
'dimension': [
{'size': 3}, # X, columns
{'size': 2}, # Y, rows
],
'attribute': [
{'name': 'ColorMode', 'value': 0},
],
})
img = nt.NTNDArray.unwrap(V)
self.assertEqual(img.shape, (2, 3))
assert_aequal(img, pixels)
V2 = nt.NTNDArray().wrap(img)
assert_aequal(V.value, V2.value)
self.assertEqual(V.dimension[0].size, V2.dimension[0].size)
self.assertEqual(V.dimension[1].size, V2.dimension[1].size)
def test_source_single(self):
"""See that source element initializes correctly for a single charge state
Use cstate=1 to select
S.IonZ = IonChargeStates[1]
S.moment0_env = IV1
S.moment1_env = IM1
"""
M = Machine(self.lattice, extra={"cstate":1})
C = M.conf()
S = M.allocState({}, inherit=False) # defaults
M.propagate(S, max=1) # propagate through source element
self.assertEqual(S.pos, 0.0)
self.assertEqual(S.ref_IonEk, 1.0)
self.assertEqual(S.ref_IonEs, 1.0)
self.assertEqual(S.ref_gamma, 2.0)
self.assertAlmostEqual(S.ref_beta, 0.8660254037844386)
assert_aequal(S.IonQ, C['NCharge'][1:])
print("moment0", S.moment0_env, C['IV1'])
assert_aequal(S.moment0_env, C['IV1'])
print("state", S.moment1_env, C['IM1'])
assert_aequal(S.moment1_env, C['IM1'].reshape((7,7)), 1e10)
def test_parse(self):
"Test global scope when parsing"
# Propagate a state matrix through a lattice defined by a lattice file.
P = GLPSParser()
with open(os.path.join(datadir, 'moment_jb_2.lat'), 'rb') as inf:
self.M = Machine(inf.read())
S = self.M.allocState({})
self.M.propagate(S)
assert_aequal(
S.state,
[[
7.50527918e+00, 4.29398123e-03, 1.26662933e-02,
-1.18523706e-06, 3.17367476e-04, -2.34554323e-05, 0
],
[
4.29398123e-03, 3.84946607e-06, -2.41285106e-05,
-1.85409878e-08, 1.06777952e-07, 5.28564016e-09, 0
],
[
1.26662933e-02, -2.41285106e-05, 8.08978916e+00,
5.33808178e-03, -1.19544245e-03, 7.51043870e-05, 0
],
[
-1.18523706e-06, -1.85409878e-08, 5.33808178e-03,
4.89711389e-06, -5.01614882e-07, 5.57484222e-08, 0
],
[
3.17367476e-04, 1.06777952e-07, -1.19544245e-03,
-5.01614882e-07, 6.71687185e-04, -1.23223342e-05, 0
],
[
-2.34554323e-05, 5.28564016e-09, 7.51043870e-05,
5.57484222e-08, -1.23223342e-05, 1.99524669e-06, 0
], [0, 0, 0, 0, 0, 0, 0]],
decimal=6)
def test_lat(self):
# Propagate a state matrix through a lattice defined by a lattice file.
P = GLPSParser()
with open(os.path.join(datadir, 'moment_jb.lat'), 'rb') as inf:
self.M = Machine(inf.read())
S = self.M.allocState({})
self.M.propagate(S)
assert_aequal(
S.state, [[
3.15034165e+00, 9.27334719e-03, 0.00000000e+00,
0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0
],
[
9.27334719e-03, 6.37581482e-03, 0.00000000e+00,
0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0
],
[
0.00000000e+00, 0.00000000e+00, 1.59854625e+01,
4.96594365e+00, 0.00000000e+00, 0.00000000e+00, 0
],
[
0.00000000e+00, 0.00000000e+00, 4.96594365e+00,
1.55019577e+00, 0.00000000e+00, 0.00000000e+00, 0
],
[
0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
0.00000000e+00, 5.00000000e-01, 0.00000000e+00, 0
],
[
0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
0.00000000e+00, 0.00000000e+00, 6.00000000e-01, 0
], [0, 0, 0, 0, 0, 0, 0]],
decimal=6)
def test_reconfig(self):
self.M.reconfigure(1, {'transfer':numpy.identity(7)*2.0})
S = self.M.allocState({})
self.M.propagate(S)
assert_aequal(S.moment0, self.expect0*2)
assert_aequal(S.state, self.expect*4)
self.M.reconfigure(1, {'transfer':numpy.identity(7)*5.0})
S = self.M.allocState({})
self.M.propagate(S)
assert_aequal(S.moment0, self.expect0*5)
assert_aequal(S.state, self.expect*25)
def test_modify_state(self):
S = self.M.allocState({}, inherit=False)
# assign scalar
S.pos = 42.2
self.assertEqual(S.pos, 42.2)
# assign 1d
S.moment0_env = numpy.asfarray([1, 2, 3, 4, 5, 6, 7])
assert_aequal(S.moment0_env, numpy.asfarray([1, 2, 3, 4, 5, 6, 7]))
# assign 2d
S.moment1_env = self.expect*2.0
assert_aequal(S.moment1_env, self.expect*2.0)
# assign 3d
X = self.expect.reshape((7,7,1))
S.moment1 = X*2.0
assert_aequal(S.moment1, X*2.0)
def test_source_multi(self):
"""See that source element initializes correctly for many (two) charge states
Use cstate=1 to select
S.IonZ = IonChargeStates[1]
S.moment0_env = IV1
S.moment1_env = IM1
"""
M = Machine(self.lattice)
C = M.conf()
S = M.allocState({}, inherit=False) # defaults
M.propagate(S, max=1) # propagate through source element
self.assertEqual(S.pos, 0.0)
self.assertEqual(S.ref_IonEk, 1.0)
self.assertEqual(S.ref_IonEs, 1.0)
self.assertEqual(S.ref_gamma, 2.0)
self.assertAlmostEqual(S.ref_beta, 0.8660254037844386)
# check access to array attributes
assert_aequal(S.IonQ, C['NCharge'])
assert_aequal(S.moment0[...,0], C['IV0'])
assert_aequal(S.moment0[...,1], C['IV1'])
assert_aequal(S.moment0.shape, numpy.asarray([7, 2]))
IM0 = C['IM0'].reshape((7,7))
IM1 = C['IM1'].reshape((7,7))
assert_aequal(S.moment1[...,0], IM0, 1e10)
assert_aequal(S.moment1[...,1], IM1, 1e10)
assert_aequal(S.moment1.shape, numpy.asarray([7, 7, 2]))
# check consistency of *_env and *_rms stats
# moment0_env is average of moment0 over charge state weighted by charge in each state
W = (S.IonQ/S.IonQ.sum()).reshape((1,2)).repeat(7,axis=0)
M0env = (S.moment0*W).sum(axis=1)
print("moment0_env", S.moment0_env, M0env)
assert_aequal(S.moment0_env, M0env)
# moment1_env is ... more complex
IM = numpy.zeros(IM0.shape)
Qs, M0, M0env, M1 = S.IonQ, S.moment0, S.moment0_env, S.moment1
# using a loop because outer() doesn't understand axis=
for i in range(len(Qs)):
Q = Qs[i]
m0diff = M0[:,i]-M0env
IM[:7,:7] += Q*(M1[...,i]+numpy.outer(m0diff, m0diff))
IM /= Qs.sum()
print("moment1_env", S.moment1_env, IM)
assert_aequal(S.moment1_env, IM, 1e10)
# moment0_rms is derived from moment1_env
M0rms = numpy.sqrt(numpy.diagonal(S.moment1_env))
print("moment0_rms", S.moment0_rms, M0rms)
assert_aequal(S.moment0_rms, M0rms)
def test_reconfig(self):
"Change the length after construction"
S = self.M.allocState({})
S.state[:] = [1, 1e-3, 0, 0, 0, 0]
assert_aequal(S.state, [1.000, 1e-3, 0, 0, 0, 0])
self.M.propagate(S)
assert_aequal(S.state, [1.001, 1e-3, 0, 0, 0, 0])
S.state[:] = [1, 1e-3, 0, 0, 0, 0]
assert_aequal(S.state, [1.000, 1e-3, 0, 0, 0, 0])
self.M.reconfigure(0, {"L": 2.0e-3})
self.M.propagate(S)
assert_aequal(S.state, [1.002, 1e-3, 0, 0, 0, 0])
S.state[:] = [1, 1e-3, 0, 0, 0, 0]
assert_aequal(S.state, [1.000, 1e-3, 0, 0, 0, 0])
self.M.reconfigure(0, {"L": 5.0e-3})
self.M.propagate(S)
assert_aequal(S.state, [1.005, 1e-3, 0, 0, 0, 0])
def test_config(self):
C = self.M.conf()
self.assertTrue("elements" in C)
self.assertEqual(C['IonEk'], 500e3)
assert_aequal(C['IV'], numpy.asfarray([1,1,0,0,0,0,0]))
