Coverage for src/km3pipe/tests/test_physics.py: 100%

108 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-19 03:14 +0000

1# Filename: test_physics.py 

2# pylint: disable=locally-disabled 

3 

4import km3io as ki 

5import numpy as np 

6import awkward as ak 

7 

8from km3pipe.dataclasses import Table 

9from km3pipe.testing import TestCase, data_path 

10from km3pipe.hardware import Detector 

11from km3pipe.calib import Calibration 

12from km3pipe.physics import cherenkov, get_closest, cut4d 

13import km3pipe.extras 

14 

15 

16__author__ = "Zineb ALY" 

17__copyright__ = "Copyright 2017, Tamas Gal and the KM3NeT collaboration." 

18__credits__ = [] 

19__license__ = "MIT" 

20__maintainer__ = "Zineb ALY" 

21__email__ = "zaly@km3net.de" 

22__status__ = "Development" 

23 

24 

25class TestGetCherenkov(TestCase): 

26 def setUp(self): 

27 self.calib_hits = { 

28 "pos_x": [ 

29 19.522, 

30 19.435, 

31 19.496, 

32 19.544, 

33 9.48, 

34 9.644, 

35 9.441, 

36 9.739, 

37 9.714, 

38 9.72, 

39 ], 

40 "pos_y": [ 

41 -12.053, 

42 -12.194, 

43 -12.196, 

44 -12.289, 

45 6.013, 

46 6.086, 

47 5.852, 

48 5.769, 

49 6.02, 

50 5.786, 

51 ], 

52 "pos_z": [ 

53 76.662, 

54 76.662, 

55 76.381, 

56 76.381, 

57 116.249, 

58 116.079, 

59 116.249, 

60 66.762, 

61 66.71, 

62 66.932, 

63 ], 

64 "dir_x": [ 

65 -0.39, 

66 -0.831, 

67 -0.527, 

68 -0.279, 

69 -0.606, 

70 0.22, 

71 -0.797, 

72 0.696, 

73 0.57, 

74 0.606, 

75 ], 

76 "dir_y": [ 

77 0.735, 

78 0.03, 

79 0.018, 

80 -0.447, 

81 0.57, 

82 0.93, 

83 -0.239, 

84 -0.655, 

85 0.606, 

86 -0.57, 

87 ], 

88 "dir_z": [ 

89 0.555, 

90 0.555, 

91 -0.85, 

92 -0.85, 

93 0.555, 

94 -0.295, 

95 0.555, 

96 -0.295, 

97 -0.555, 

98 0.555, 

99 ], 

100 } 

101 

102 self.track = { 

103 "pos_x": 1.7524502152598151, 

104 "pos_y": 39.06202405657308, 

105 "pos_z": 130.44049806891948, 

106 "dir_x": 0.028617421257374293, 

107 "dir_y": -0.489704257367248, 

108 "dir_z": -0.8714188335794505, 

109 "t": 70311441.452294, 

110 } 

111 

112 def test_cherenkov_from_dict(self): 

113 

114 arr = cherenkov(self.calib_hits, self.track) 

115 

116 self.assertAlmostEqual(arr["d_photon_closest"][0], 24.049593557846112) 

117 self.assertAlmostEqual(arr["d_photon_closest"][1], 24.085065395206847) 

118 

119 self.assertAlmostEqual(arr["d_photon"][0], 35.80244420413484) 

120 self.assertAlmostEqual(arr["d_photon"][1], 35.855250854478896) 

121 

122 self.assertAlmostEqual(arr["d_track"][0], 45.88106599210481) 

123 self.assertAlmostEqual(arr["d_track"][1], 45.90850564175342) 

124 

125 self.assertAlmostEqual(arr["t_photon"][0], 70311759.26448613) 

126 self.assertAlmostEqual(arr["t_photon"][1], 70311759.59904088) 

127 

128 self.assertAlmostEqual(arr["cos_photon_PMT"][0], -0.98123942583677) 

129 self.assertAlmostEqual(arr["cos_photon_PMT"][1], -0.6166369315726149) 

130 

131 self.assertAlmostEqual(arr["dir_x_photon"][0], 0.45964884122649263) 

132 self.assertAlmostEqual(arr["dir_x_photon"][1], 0.45652355929477095) 

133 

134 self.assertAlmostEqual(arr["dir_y_photon"][0], -0.8001372907490844) 

135 self.assertAlmostEqual(arr["dir_y_photon"][1], -0.8025165828910586) 

136 

137 self.assertAlmostEqual(arr["dir_z_photon"][0], -0.3853612055096594) 

138 self.assertAlmostEqual(arr["dir_z_photon"][1], -0.38412676812960095) 

139 

140 def test_cherenkov_from_Table(self): 

141 

142 arr = cherenkov(Table(self.calib_hits), Table(self.track)) 

143 

144 self.assertAlmostEqual(arr["d_photon_closest"][0], 24.049593557846112) 

145 self.assertAlmostEqual(arr["d_photon"][0], 35.80244420413484) 

146 self.assertAlmostEqual(arr["d_track"][0], 45.88106599210481) 

147 self.assertAlmostEqual(arr["t_photon"][0], 70311759.26448613) 

148 self.assertAlmostEqual(arr["cos_photon_PMT"][0], -0.98123942583677) 

149 self.assertAlmostEqual(arr["dir_x_photon"][0], 0.45964884122649263) 

150 self.assertAlmostEqual(arr["dir_y_photon"][0], -0.8001372907490844) 

151 self.assertAlmostEqual(arr["dir_z_photon"][0], -0.3853612055096594) 

152 

153 def test_cherenkov_from_DataFrame(self): 

154 

155 pd = km3pipe.extras.pandas() 

156 

157 arr = cherenkov(pd.DataFrame(self.calib_hits), pd.Series(self.track)) 

158 

159 self.assertAlmostEqual(arr["d_photon_closest"][0], 24.049593557846112) 

160 self.assertAlmostEqual(arr["d_photon"][0], 35.80244420413484) 

161 self.assertAlmostEqual(arr["d_track"][0], 45.88106599210481) 

162 self.assertAlmostEqual(arr["t_photon"][0], 70311759.26448613) 

163 self.assertAlmostEqual(arr["cos_photon_PMT"][0], -0.98123942583677) 

164 self.assertAlmostEqual(arr["dir_x_photon"][0], 0.45964884122649263) 

165 self.assertAlmostEqual(arr["dir_y_photon"][0], -0.8001372907490844) 

166 self.assertAlmostEqual(arr["dir_z_photon"][0], -0.3853612055096594) 

167 

168 def test_cherenkov_from_best_track_which_returns_awkward_Record(self): 

169 pd = km3pipe.extras.pandas() 

170 

171 best_track = ak.Record(self.track) 

172 

173 arr = cherenkov(pd.DataFrame(self.calib_hits), best_track) 

174 

175 self.assertAlmostEqual(arr["d_photon_closest"][0], 24.049593557846112) 

176 self.assertAlmostEqual(arr["d_photon"][0], 35.80244420413484) 

177 self.assertAlmostEqual(arr["d_track"][0], 45.88106599210481) 

178 self.assertAlmostEqual(arr["t_photon"][0], 70311759.26448613) 

179 self.assertAlmostEqual(arr["cos_photon_PMT"][0], -0.98123942583677) 

180 self.assertAlmostEqual(arr["dir_x_photon"][0], 0.45964884122649263) 

181 self.assertAlmostEqual(arr["dir_y_photon"][0], -0.8001372907490844) 

182 self.assertAlmostEqual(arr["dir_z_photon"][0], -0.3853612055096594) 

183 

184 

185class TestGetClosest(TestCase): 

186 def setUp(self): 

187 self.track = { 

188 "pos_x": 1.7524502152598151, 

189 "pos_y": 39.06202405657308, 

190 "pos_z": 130.44049806891948, 

191 "dir_x": 0.028617421257374293, 

192 "dir_y": -0.489704257367248, 

193 "dir_z": -0.8714188335794505, 

194 } 

195 

196 self.det = Detector(data_path("detx/detx_v3.detx")) 

197 

198 pd = km3pipe.extras.pandas() 

199 

200 self.DU = pd.DataFrame(self.det.dom_table).mean() 

201 

202 def test_get_closest(self): 

203 pd = km3pipe.extras.pandas() 

204 

205 DU = pd.DataFrame(self.det.dom_table).mean() 

206 d_closest, z_closest = get_closest(self.track, DU) 

207 

208 self.assertAlmostEqual(d_closest, 9.073491762564467) 

209 self.assertAlmostEqual(z_closest, 82.24928115091757) 

210 

211 def test_get_closest_from_DataFrame(self): 

212 pd = km3pipe.extras.pandas() 

213 d_closest, z_closest = get_closest(pd.Series(self.track), pd.Series(self.DU)) 

214 

215 self.assertAlmostEqual(d_closest, 9.073491762564467) 

216 self.assertAlmostEqual(z_closest, 82.24928115091757) 

217 

218 def test_get_closest_from_Table(self): 

219 trk = { 

220 "pos_x": 1.7524502152598151, 

221 "pos_y": 39.06202405657308, 

222 "pos_z": 130.44049806891948, 

223 "dir_x": 0.028617421257374293, 

224 "dir_y": -0.489704257367248, 

225 "dir_z": -0.8714188335794505, 

226 } 

227 

228 pd = km3pipe.extras.pandas() 

229 mean = pd.DataFrame(self.det.dom_table).mean() 

230 

231 d_closest, z_closest = get_closest(Table(trk), mean) 

232 

233 self.assertAlmostEqual(d_closest, 9.073491762564467) 

234 self.assertAlmostEqual(z_closest, 82.24928115091757) 

235 

236 

237class TestCut4D(TestCase): 

238 def test_cut4d(self): 

239 point4d = Table({"pos_x": [0], "pos_y": [3], "pos_z": [0], "t": [20]}) 

240 

241 items = Table( 

242 { 

243 "pos_x": [0, 10, 0, 20, 0], 

244 "pos_y": [10, 0, 0, 0, 30], 

245 "pos_z": [0, 0, 10, 0, 0], 

246 "time": [60, 15, 40, 20, 100], 

247 } 

248 ) 

249 

250 tmin = -50.0 

251 tmax = 10.0 

252 rmin = 3.0 

253 rmax = 80.0 

254 

255 selected_items = cut4d(point4d, tmin, tmax, rmin, rmax, items) 

256 assert len(selected_items) == 3 

257 

258 self.assertListEqual( 

259 list(np.array([0, 10, 0, 60])), 

260 list(selected_items.T[0]), 

261 ) 

262 self.assertListEqual( 

263 list(np.array([0, 0, 10, 40])), 

264 list(selected_items.T[1]), 

265 ) 

266 self.assertListEqual( 

267 list(np.array([0, 30, 0, 100])), 

268 list(selected_items.T[2]), 

269 )