diff --git a/km3io/offline.py b/km3io/offline.py index 4677316320d9bdcdf130cc32ec921438204af46f..03b4dba36018af70a72229de332d5bdc62d68479 100644 --- a/km3io/offline.py +++ b/km3io/offline.py @@ -560,9 +560,9 @@ class OfflineReader: are not found in the file. """ keys = ", ".join(self.keys.fit_keys[:-1]) - fit_info = self.tracks.fitinf rec_stages = np.array( [match for match in self._find_rec_stages(stages)]) + mask = rec_stages[:, 1] != None if np.all(rec_stages[:, 1] == None): raise ValueError( "The stages {} are not found in your file.".format( @@ -570,9 +570,8 @@ class OfflineReader: else: fit_data = np.array([ i[k] - for i, j, k in zip(fit_info, rec_stages[:, 0], rec_stages[:, - 1]) - if k is not None + for i, k in zip(self.tracks.fitinf[mask], rec_stages[:, + 1][mask]) ]) rec_array = np.core.records.fromarrays(fit_data.transpose(), names=keys) @@ -709,7 +708,7 @@ class OfflineReader: list of reconstruction stages of interest. for example [1, 2, 3, 4, 5]. - Yieldsma + Yields ------ generator the track id and the index of the reconstruction stages of