remove events from best_track

km3io/tools.py

@@ -271,13 +271,13 @@ def mask(rec_stages, stages):
     return builder.snapshot() == 1
 
 
-def best_track(events, strategy="default", rec_type=None):
+def best_track(tracks, strategy="default", rec_type=None):
     """best track selection based on different strategies
 
     Parameters
     ----------
-    events : class km3io.offline.OfflineBranch
-        a subset of reconstructed events where `events.n_tracks > 0` is always true. 
+    tracks : class km3io.offline.OfflineBranch
+        a subset of reconstructed tracks where `events.n_tracks > 0` is always true. 
     strategy : str
         the trategy desired to select the best tracks. It is either: 
             - "first" : to select the first tracks.
@@ -303,13 +303,12 @@ def best_track(events, strategy="default", rec_type=None):
     if strategy not in options:
         raise ValueError("{} not in {}".format(strategy, options))
 
-    n_events = len(events)
-    if n_events > 1 and any(events.n_tracks == 0):
+    n_events = len(tracks)
+    if n_events > 1 and any(count_nested(tracks.lik, axis=1) == 0):
         raise ValueError(
             "'events' should not contain empty tracks. Consider applying the mask: events.n_tracks>0"
         )
 
-    tracks = events.tracks
     if strategy == "first":
         if n_events == 1:
             out = tracks[0]

tests/test_tools.py

@@ -60,8 +60,8 @@ class TestBestTrack(unittest.TestCase):
     def test_best_tracks(self):
         # test selection from multiple events
         events = self.events[self.events.n_tracks > 0]
-        first_tracks = best_track(events, strategy="first")
-        default_best = best_track(events,
+        first_tracks = best_track(events.tracks, strategy="first")
+        default_best = best_track(events.tracks,
                                   strategy="default",
                                   rec_type="JPP_RECONSTRUCTION_TYPE")
 
@@ -73,8 +73,8 @@ class TestBestTrack(unittest.TestCase):
         assert default_best.rec_type[0] == 4000
 
         # test selection from one event
-        first_track = best_track(self.one_event, strategy="first")
-        best = best_track(self.one_event,
+        first_track = best_track(self.one_event.tracks, strategy="first")
+        best = best_track(self.one_event.tracks,
                           strategy="default",
                           rec_type="JPP_RECONSTRUCTION_TYPE")
 
@@ -86,7 +86,7 @@ class TestBestTrack(unittest.TestCase):
 
         # test raising ValueError
         with self.assertRaises(ValueError):
-            best_track(events, strategy="Zineb")
+            best_track(events.tracks, strategy="Zineb")
 
 
 class TestGetMultiplicity(unittest.TestCase):