From e1756439ea218fef88702aa91fb82f1f22c2bc72 Mon Sep 17 00:00:00 2001
From: zineb aly <aly.zineb.az@gmail.com>
Date: Tue, 26 May 2020 09:29:27 +0200
Subject: [PATCH] add best track
---
km3io/offline.py | 10 +++++-----
1 file changed, 5 insertions(+), 5 deletions(-)
diff --git a/km3io/offline.py b/km3io/offline.py
index b3fe535..b3ed691 100644
--- a/km3io/offline.py
+++ b/km3io/offline.py
@@ -1,8 +1,8 @@
from collections import namedtuple
import uproot
import warnings
-import awkward1 as ak1
import numba as nb
+import awkward1 as ak1
from .definitions import mc_header, fitparameters
from .tools import cached_property, to_num, unfold_indices
@@ -152,8 +152,8 @@ def best_track(tracks, strategy="first", rec_stages=None):
"""
if strategy == "first":
return tracks[:, 0]
- # if strategy == "rec_stages" and rec_stages is not None:
- # mask = tracks.rec_stages[]
+ if strategy == "rec_stages" and rec_stages is not None:
+ return tracks[mask(tracks.rec_stages, rec_stages)]
EVENTS_MAP = BranchMapper(name="events",
@@ -182,8 +182,8 @@ SUBBRANCH_MAPS = [
BranchMapper(name="mc_tracks",
key="mc_trks",
exclude=EXCLUDE_KEYS + [
- 'mc_trks.rec_stages',
- 'mc_trks.fitinf', 'mc_trks.fUniqueID', 'mc_trks.fBits'
+ 'mc_trks.rec_stages', 'mc_trks.fitinf',
+ 'mc_trks.fUniqueID', 'mc_trks.fBits'
],
attrparser=_nested_mapper,
toawkward=['usr', 'usr_names'],
--
GitLab