diff --git a/km3io/offline.py b/km3io/offline.py index 5acf53c7905718e5e88b942f5f1259611fd9437d..8a8175d9b7d000c714cbc8d928bb931682ccc66a 100644 --- a/km3io/offline.py +++ b/km3io/offline.py @@ -582,16 +582,13 @@ class OfflineReader: "The stages {} are not found in your file.".format( str(stages))) else: - fit_data = np.array([ - i[k] - for i, k in zip(fitinf[mask], rec_stages[:,1][mask]) - ]) + fit_data = np.array( + [i[k] for i, k in zip(fitinf[mask], rec_stages[:, 1][mask])]) rec_array = np.core.records.fromarrays(fit_data.transpose(), names=keys) return rec_array def get_reco_hits(self, stages, keys, mc=False): - """construct a dictionary of hits class data based on the reconstruction stages of interest. For example, if the reconstruction stages of interest are [1, 2, 3, 4, 5], then get_reco_hits method will select the hits data diff --git a/tests/test_offline.py b/tests/test_offline.py index a37f69ef179734bc83415f1fd7b6e658b3a6ff71..19253494b882b8cdebc41a5e29faa57308fa0427 100644 --- a/tests/test_offline.py +++ b/tests/test_offline.py @@ -221,7 +221,6 @@ class TestOfflineReader(unittest.TestCase): pos = self.nu.get_reco_tracks([1, 2, 3, 4, 5], ["pos_x"])["pos_x"] mc_pos = self.nu.get_reco_tracks([], ["pos_x"], mc=True)["pos_x"] - self.assertEqual(pos.size, 9) self.assertEqual(mc_pos.size, 10)