def test_init(self):
dirrec = direction_reconstruction.CoincidenceDirectionReconstruction(
sentinel.cluster)
self.assertEqual(dirrec.direct,
direction_reconstruction.DirectAlgorithmCartesian3D)
self.assertEqual(dirrec.fit,
direction_reconstruction.RegressionAlgorithm3D)
self.assertEqual(dirrec.cluster, sentinel.cluster)
def test_set_cluster_timestamp(self):
cluster = Mock()
dirrec = direction_reconstruction.CoincidenceDirectionReconstruction(
cluster)
coincidence = [[sentinel.station_number, {
'timestamp': 1
}], [0, 0], [0, 0]]
theta, phi, nums = dirrec.reconstruct_coincidence(coincidence,
station_numbers=[])
cluster.set_timestamp.assert_called_with(1)
self.assertTrue(isnan(theta))
self.assertTrue(isnan(phi))
def test_reconstruct_coincidences(self, mock_reconstruct_coincidence):
mock_reconstruct_coincidence.return_value = [
sentinel.theta, sentinel.phi, sentinel.nums
]
dirrec = direction_reconstruction.CoincidenceDirectionReconstruction(
sentinel.cluster)
self.assertEqual(
dirrec.reconstruct_coincidences(
[sentinel.coincidence, sentinel.coincidence],
sentinel.station_numbers,
sentinel.offsets,
progress=False),
((sentinel.theta, sentinel.theta), (sentinel.phi, sentinel.phi),
(sentinel.nums, sentinel.nums)))
self.assertEqual(mock_reconstruct_coincidence.call_count, 2)
mock_reconstruct_coincidence.assert_called_with(
sentinel.coincidence, sentinel.station_numbers, sentinel.offsets)
self.assertEqual(
dirrec.reconstruct_coincidences([],
sentinel.station_numbers,
sentinel.offsets,
progress=False), ((), (), ()))
self.assertEqual(mock_reconstruct_coincidence.call_count, 2)
def test_reconstruct_coincidence(self, mock_station_arrival_time):
mock_station_arrival_time.return_value = 0.
cluster = MagicMock()
station = MagicMock()
cluster.get_station.return_value = station
station.calc_center_of_mass_coordinates.return_value = [
sentinel.x, sentinel.y, sentinel.z
]
dirrec = direction_reconstruction.CoincidenceDirectionReconstruction(
cluster)
dirrec.direct = Mock()
dirrec.fit = Mock()
dirrec.direct.reconstruct_common.return_value = (sentinel.theta,
sentinel.phi)
dirrec.fit.reconstruct_common.return_value = (sentinel.theta,
sentinel.phi)
coincidence_2 = [[sentinel.station_number, {
'timestamp': 1
}], [1, sentinel.event]]
coincidence_3 = [[sentinel.station_number, {
'timestamp': 1
}], [1, sentinel.event], [2, sentinel.event]]
coincidence_4 = [[sentinel.station_number, {
'timestamp': 1
}], [1, sentinel.event], [2, sentinel.event], [3, sentinel.event]]
theta, phi, nums = dirrec.reconstruct_coincidence(coincidence_2)
self.assertEqual(dirrec.direct.reconstruct_common.call_count, 0)
self.assertTrue(isnan(theta))
self.assertTrue(isnan(phi))
self.assertEqual(len(nums), 0)
theta, phi, nums = dirrec.reconstruct_coincidence(
coincidence_3, station_numbers=[1, 2])
self.assertEqual(dirrec.direct.reconstruct_common.call_count, 0)
self.assertTrue(isnan(theta))
self.assertTrue(isnan(phi))
self.assertEqual(len(nums), 2)
theta, phi, nums = dirrec.reconstruct_coincidence(coincidence_3)
cluster.set_timestamp.assert_called_with(1)
dirrec.direct.reconstruct_common.assert_called_once_with(
[0.] * 3, [sentinel.x] * 3, [sentinel.y] * 3, [sentinel.z] * 3)
self.assertEqual(dirrec.fit.reconstruct_common.call_count, 0)
self.assertEqual(theta, sentinel.theta)
self.assertEqual(phi, sentinel.phi)
self.assertEqual(len(nums), 3)
theta, phi, nums = dirrec.reconstruct_coincidence(coincidence_4)
cluster.set_timestamp.assert_called_with(1)
dirrec.fit.reconstruct_common.assert_called_once_with([0.] * 4,
[sentinel.x] * 4,
[sentinel.y] * 4,
[sentinel.z] * 4)
self.assertEqual(dirrec.direct.reconstruct_common.call_count, 1)
self.assertEqual(theta, sentinel.theta)
self.assertEqual(phi, sentinel.phi)
self.assertEqual(len(nums), 4)
mock_station_arrival_time.return_value = nan
theta, phi, nums = dirrec.reconstruct_coincidence(coincidence_4)
cluster.set_timestamp.assert_called_with(1)
self.assertEqual(dirrec.fit.reconstruct_common.call_count, 1)
self.assertEqual(dirrec.direct.reconstruct_common.call_count, 1)
self.assertTrue(isnan(theta))
self.assertTrue(isnan(phi))
self.assertEqual(len(nums), 0)
def setUp(self):
self.dirrec = direction_reconstruction.CoincidenceDirectionReconstruction(
sentinel.cluster)