From abad0a5a6eaa769438a8a9d25ddfd4bc2e8335ed Mon Sep 17 00:00:00 2001
From: Zineb Aly <zaly@km3net.de>
Date: Tue, 29 Sep 2020 22:42:55 +0200
Subject: [PATCH] count rec stages lentgh in best_track
---
km3io/tools.py | 18 +++++++++++++-----
1 file changed, 13 insertions(+), 5 deletions(-)
diff --git a/km3io/tools.py b/km3io/tools.py
index e28acdb..d6cff15 100644
--- a/km3io/tools.py
+++ b/km3io/tools.py
@@ -315,13 +315,21 @@ def best_track(tracks, strategy="default", rec_type=None):
else:
out = tracks[:, 0]
- elif strategy == "default" and rec_type is not None:
+ if strategy == "default" and rec_type is None:
+ raise KeyError(
+ "rec_type must be provided when the default strategy is used.")
+
+ if strategy == "default" and rec_type is not None:
if n_events == 1:
- out = tracks[tracks.rec_type == reconstruction[rec_type]][
- tracks.lik == ak1.max(tracks.lik, axis=0)][0]
+ len_stages = count_nested(tracks.rec_stages, axis=1)
+ rec_types = tracks[tracks.rec_type == reconstruction[rec_type]]
+ longest = rec_types[len_stages == ak1.max(len_stages, axis=0)]
+ out = longest[longest.lik == ak1.max(longest.lik, axis=0)]
else:
- out = tracks[tracks.rec_type == reconstruction[rec_type]][
- tracks.lik == ak1.max(tracks.lik, axis=1)][:, 0]
+ len_stages = count_nested(tracks.rec_stages, axis=2)
+ rec_types = tracks[tracks.rec_type == reconstruction[rec_type]]
+ longest = rec_types[len_stages == ak1.max(len_stages, axis=1)]
+ out = longest[longest.lik == ak1.max(longest.lik, axis=1)]
return out
--
GitLab