diff --git a/km3io/offline.py b/km3io/offline.py index b3ed691ed3b73f6503d2a5420756ed1472989b3e..e51f319d33220b03f2a7b304069044ee128967c7 100644 --- a/km3io/offline.py +++ b/km3io/offline.py @@ -4,7 +4,7 @@ import warnings import numba as nb import awkward1 as ak1 -from .definitions import mc_header, fitparameters +from .definitions import mc_header, fitparameters, reconstruction from .tools import cached_property, to_num, unfold_indices from .rootio import Branch, BranchMapper @@ -21,6 +21,18 @@ def _nested_mapper(key): return '_'.join(key.split('.')[1:]) +def rec_types(): + """name of the reconstruction type as defined in the official + KM3NeT-Dataformat. + + Returns + ------- + dict_keys + reconstruction types. + """ + return reconstruction.keys() + + def fitinf(fitparam, tracks): """Access fit parameters in tracks.fitinf. @@ -140,7 +152,7 @@ def mask(rec_stages, stages): return builder.snapshot() == 1 -def best_track(tracks, strategy="first", rec_stages=None): +def best_track(tracks, strategy="first", rec_type=None, rec_stages=None): """best track selection based on different strategies Parameters @@ -154,6 +166,9 @@ def best_track(tracks, strategy="first", rec_stages=None): return tracks[:, 0] if strategy == "rec_stages" and rec_stages is not None: return tracks[mask(tracks.rec_stages, rec_stages)] + if strategy == "default" and rec_type is not None: + return tracks[tracks.rec_type == reconstruction[rec_type]][ + tracks.lik == ak1.max(tracks.lik, axis=1)][:, 0] EVENTS_MAP = BranchMapper(name="events",