def test_track_many(self):
"""
Test MatrixTracking.track_many
"""
tracking = MatrixTracking(self.matrix_list, self.offset_list,
self.offset_in)
hit_in = Hit.new_from_dict(
{
"mass": Common.pdg_pid_to_mass[2212],
"charge": 1.,
"pid": 2212,
"pz": 1.,
"x": 1.,
"py": 1.,
"energy": 1001.
}, "pz")
hit_in_2 = Hit.new_from_dict(
{
"mass": Common.pdg_pid_to_mass[2212],
"charge": 1.,
"pid": 2212,
"pz": 1.,
"x": 1.,
"py": 1.,
"energy": 1002.
}, "pz")
hit_list_of_lists = tracking.track_many([hit_in] * 3 + [hit_in_2])
for hit_list in hit_list_of_lists[1:-1]:
self.assertEqual(hit_list, hit_list_of_lists[0])
self.assertNotEqual(hit_list_of_lists[-1], hit_list_of_lists[0])
self.assertEqual(tracking.last, hit_list_of_lists)
def __init__(self, list_of_transfer_matrices, list_of_offsets, offset_in,
timeout):
"""
As per matrix tracking but with additional timeout parameter
- timeout: causes track_many to pause for timeout seconds before return
"""
MatrixTracking.__init__(self, list_of_transfer_matrices,
list_of_offsets, offset_in)
self.timeout = timeout
def setUp(self):
"""Setup the EllipseClosedOrbitFinder"""
self.num_turns = 5
# x is conventional ellipse
# y is phase advance = pi
# t is phase advance = complex
data = [[1.0, 0.5, 0.0, 0.0, 0.0,
0.0], [-0.5, 0.75, 0.0, 0.0, 0.0, 0.0],
[0.0, 0.0, 0.0, -1.0, 0.0, 0.0],
[0.0, 0.0, 1.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 2.0, 2.0],
[0.0, 0.0, 0.0, 0.0, 2.0, 2.5]]
matrix = numpy.matrix(data)
self.matrix_list = [matrix**(i + 1) for i in range(self.num_turns)]
offset = numpy.matrix([10., 7., 0., 0., 0., 1000.])
self.offset_list = [offset] * self.num_turns
self.offset_in = numpy.matrix([10., 7., 0., 0., 0., 1000.])
self.tracking = MatrixTracking(self.matrix_list, self.offset_list,
self.offset_in)
self.co = Hit.new_from_dict({
'x': 10.,
'px': 7.,
'energy': 1000.,
'pid': 2212,
'mass': common.pdg_pid_to_mass[2212]
})
if abs(numpy.linalg.det(matrix) - 1.0) > 1e-9:
raise ValueError("TM determinant should be 1, got "+\
str(numpy.linalg.det(matrix)))
def get_tracking_object():
"""
Create a tracking object. This tracking object uses a list of matrices to
generate tracking. The tracking goes like
- x is a conventional phase space ellipse
- y has phase advance of exactly pi
- t has phase advance is complex i.e. on a resonance.
Each cell or turn we advance through another transfer matrix. The closed
orbit is generated by a hard coded offset.
Return value is a tuple consisting of a hit object on the closed_orbit and
a tracking object that will propagate hit objects through the matrices.
"""
number_of_turns = 5
# x is conventional ellipse
# y is phase advance = pi
# t is phase advance = complex
data = [[1.0, 0.5, 0.0, 0.0, 0.0, 0.0], [-0.5, 0.75, 0.0, 0.0, 0.0, 0.0],
[0.0, 0.0, 0.0, -1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 0.0, 2.0, 2.0], [0.0, 0.0, 0.0, 0.0, 2.0, 2.5]]
matrix = numpy.matrix(data)
matrix_list = [matrix**(i + 1) for i in range(number_of_turns)]
offset = numpy.matrix([10., 7., 0., 0., 0., 1000.])
offset_list = [offset] * number_of_turns
offset_in = numpy.matrix([10., 7., 0., 0., 0., 1000.])
tracking = MatrixTracking(matrix_list, offset_list, offset_in)
closed_orbit_hit = Hit.new_from_dict(
{
'x': 10.,
'px': 7.,
'energy': 1000.,
'pid': 2212,
'mass': common.pdg_pid_to_mass[2212]
}, 'pz')
return closed_orbit_hit, tracking
def test_init(self):
"""
Test MatrixTracking.__init__
"""
try: # wrong lengths
tracking = MatrixTracking(self.matrix_list, self.offset_list[0:4],
self.offset_in)
self.assertTrue(False, msg="Should have raised")
except IndexError:
pass
try: # offset wrong type
tracking = MatrixTracking(self.matrix_list,
[1] * len(self.matrix_list),
self.offset_in)
self.assertTrue(False, msg="Should have raised")
except TypeError:
pass
try: # offset wrong shape
tracking = MatrixTracking(self.matrix_list, self.matrix_list,
self.offset_in)
self.assertTrue(False, msg="Should have raised")
except TypeError:
pass
try: # matrix wrong type
tracking = MatrixTracking([1] * len(self.matrix_list),
self.offset_list, self.offset_in)
self.assertTrue(False, msg="Should have raised")
except TypeError:
pass
try: # matrix wrong shape
tracking = MatrixTracking(self.offset_list, self.offset_list,
self.offset_in)
self.assertTrue(False, msg="Should have raised")
except TypeError:
pass
try: # offset_in wrong type
tracking = MatrixTracking(self.matrix_list, self.offset_list, 1)
self.assertTrue(False, msg="Should have raised")
except TypeError:
pass
try: # offset_in wrong shape
tracking = MatrixTracking(self.offset_list, self.offset_list,
self.matrix_list[0])
self.assertTrue(False, msg="Should have raised")
except TypeError:
pass
tracking = MatrixTracking(self.matrix_list, self.offset_list,
self.offset_in)
self.assertEqual(tracking.last, [])
def test_track_one(self):
"""
Test MatrixTracking.track_one
"""
tracking = MatrixTracking(self.matrix_list, self.offset_list,
self.offset_in)
hit_in = Hit.new_from_dict(
{
"mass": Common.pdg_pid_to_mass[2212],
"charge": 1.,
"pid": 2212,
"pz": 1.,
"x": 1.,
"py": 1.,
"energy": 1001.
}, "pz")
hit_list = tracking.track_one(hit_in)
self.assertEqual(len(hit_list), self.num_turns + 1)
self.assertEqual(hit_list[0], hit_in)
self.assertLess(abs(hit_list[-1]['x']), 10.) # should be contained
self.assertLess(abs(hit_list[-1]['px']), 10.) # should be contained
self.assertAlmostEqual(abs(hit_list[-1]['py']), hit_in['py']) # drift
self.assertEqual(tracking.last, [hit_list])
def matrix(phi_x, num_turns):
# x is conventional ellipse
# y is phase advance = pi
# t is phase advance = complex
f_x = 2 * math.pi - phi_x
cos = math.cos(f_x)
sin = math.sin(f_x)
b_x = 10.
a_x = 1.
g_x = (1. + a_x**2.) / b_x
data = [[cos + sin * a_x, sin * b_x, 0.0, 0.0, 0.0, 0.0],
[-sin * g_x, cos - sin * a_x, 0.0, 0.0, 0.0, 0.0],
[0.0, 0.0, 0.0, -1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 1.0]]
matrix = numpy.matrix(data)
matrix_list = [matrix]
for i in range(num_turns - 1):
matrix_list.append(matrix * matrix_list[-1])
for j in range(6):
if math.isnan(matrix_list[-1][j, j]):
print(("ARG", i, matrix_list[-1], '\n\n', matrix_list[-5]))
raise ValueError("ARG")
matrix_list = matrix_list
offset = numpy.matrix([10., 7., 0., 0., 0., 1000.])
offset_list = [offset] * num_turns
offset_in = numpy.matrix([10., 7., 0., 0., 0., 1000.])
tracking = MatrixTracking(matrix_list, offset_list, offset_in)
co = Hit.new_from_dict({
'x': 10.,
'px': 7.,
'energy': 1000.,
'pid': 2212,
'mass': common.pdg_pid_to_mass[2212]
})
if abs(numpy.linalg.det(matrix) - 1.0) > 1e-9:
raise ValueError("TM determinant should be 1, got "+\
str(numpy.linalg.det(matrix)))
return co, tracking
def track_many(self, hit_list):
"""Track the hits through the matrices"""
tracking_out = MatrixTracking.track_many(self, hit_list)
time.sleep(self.timeout)
return tracking_out