diff --git a/km3io/tools.py b/km3io/tools.py index 4b95dd42abb9591a00ff837e484ca76dac5de426..22462762068b4d9bc32a6174265d28d4352dc080 100644 --- a/km3io/tools.py +++ b/km3io/tools.py @@ -271,13 +271,13 @@ def mask(rec_stages, stages): return builder.snapshot() == 1 -def best_track(tracks, strategy="first", rec_type=None, rec_stages=None): +def best_track(events, strategy="first", rec_type=None, rec_stages=None): """best track selection based on different strategies Parameters ---------- - tracks : class km3io.offline.OfflineBranch - the tracks branch. + events : class km3io.offline.OfflineBranch + the events branch. strategy : str the trategy desired to select the best tracks. It is either: - "first" : to select the first tracks. @@ -297,6 +297,7 @@ def best_track(tracks, strategy="first", rec_type=None, rec_stages=None): class km3io.offline.OfflineBranch tracks class with the desired "best tracks" selection. """ + tracks = events.tracks[events.n_tracks > 0] if strategy == "first": return tracks[:, 0] diff --git a/tests/test_tools.py b/tests/test_tools.py index 325e6814c4e86da581d0d78a1e39187c6ff6ebd0..71c4a409c0c603d873e97e934fbb9bc3279f816c 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -8,11 +8,16 @@ from pathlib import Path from km3io import OfflineReader from km3io.tools import (to_num, cached_property, unfold_indices, unique, uniquecount, fitinf, fitparams, count_nested, _find, - mask, best_track, rec_types) + mask, best_track, rec_types, get_w2list_param) SAMPLES_DIR = Path(__file__).parent / 'samples' OFFLINE_FILE = OfflineReader(SAMPLES_DIR / 'km3net_offline.root') +# class TestGetw2listParam(unittest.TestCase): +# def test_get_w2list_param(self): +# xsec_mean = get_w2list_param(OFFLINE_FILE.events, "gseagen", "W2LIST_GSEAGEN_XSEC_MEAN") +# print(xsec_mean) + class TestFitinf(unittest.TestCase): def setUp(self): @@ -48,25 +53,25 @@ class TestRecoTypes(unittest.TestCase): class TestBestTrack(unittest.TestCase): def setUp(self): - self.tracks = OFFLINE_FILE.events.tracks + self.events = OFFLINE_FILE.events def test_best_tracks(self): - first_tracks = best_track(self.tracks, strategy="first") - rec_stages_tracks = best_track(self.tracks, + first_tracks = best_track(self.events, strategy="first") + rec_stages_tracks = best_track(self.events, strategy="rec_stages", rec_stages=[1, 3, 5, 4]) - default_best = best_track(self.tracks, + default_best = best_track(self.events, strategy="default", rec_type="JPP_RECONSTRUCTION_TYPE") - assert first_tracks.dir_z[0] == self.tracks.dir_z[0][0] - assert first_tracks.dir_x[1] == self.tracks.dir_x[1][0] + assert first_tracks.dir_z[0] == self.events.tracks.dir_z[0][0] + assert first_tracks.dir_x[1] == self.events.tracks.dir_x[1][0] assert rec_stages_tracks.rec_stages[0] == [1, 3, 5, 4] assert rec_stages_tracks.rec_stages[1] == [1, 3, 5, 4] - assert default_best.lik[0] == ak.max(self.tracks.lik[0]) - assert default_best.lik[1] == ak.max(self.tracks.lik[1]) + assert default_best.lik[0] == ak.max(self.events.tracks.lik[0]) + assert default_best.lik[1] == ak.max(self.events.tracks.lik[1]) assert default_best.rec_type[0] == 4000