diff --git a/km3io/offline.py b/km3io/offline.py index 9026e684b9cafe7f2d4dac233123545634927bee..fe9794d563fe84b0c567779e12c7a6d0b84d1768 100644 --- a/km3io/offline.py +++ b/km3io/offline.py @@ -21,158 +21,6 @@ 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. - - Parameters - ---------- - fitparam : str - the fit parameter name according to fitparameters defined in - KM3NeT-Dataformat. - tracks : class km3io.offline.OfflineBranch - the tracks class. both full tracks branch or a slice of the - tracks branch (example tracks[:, 0]) work. - - Returns - ------- - awkward array - awkward array of the values of the fit parameter requested. - """ - fit = tracks.fitinf - index = fitparameters[fitparam] - try: - params = fit[count_nested(fit, axis=2) > index] - return ak1.Array([i[:, index] for i in params]) - except ValueError: - # This is the case for tracks[:, 0] or any other selection. - params = fit[count_nested(fit, axis=1) > index] - return params[:, index] - - -def fitparams(): - """name of the fit parameters as defined in the official - KM3NeT-Dataformat. - - Returns - ------- - dict_keys - fit parameters keys. - """ - return fitparameters.keys() - - -def count_nested(Array, axis=0): - """count elements in a nested awkward Array. - - Parameters - ---------- - Array : Awkward1 Array - Array of data. Example tracks.fitinf or tracks.rec_stages. - axis : int, optional - axis = 0: to count elements in the outmost level of nesting. - axis = 1: to count elements in the first level of nesting. - axis = 2: to count elements in the second level of nesting. - - Returns - ------- - awkward1 Array or int - counts of elements found in a nested awkward1 Array. - """ - if axis == 0: - return ak1.num(Array, axis=0) - if axis == 1: - return ak1.num(Array, axis=1) - if axis == 2: - return ak1.count(Array, axis=2) - - -@nb.jit(nopython=True) -def _find(rec_stages, stages, builder): - """construct an awkward1 array with the same structure as tracks.rec_stages. - When stages are found, the Array is filled with value 1, otherwise it is filled - with value 0. - - Parameters - ---------- - rec_stages : awkward1 Array - tracks.rec_stages . - stages : awkward1 Array - reconstruction stages of interest. - builder : awkward1.highlevel.ArrayBuilder - awkward1 Array builder. - """ - for s in rec_stages: - builder.begin_list() - for i in s: - num_stages = len(i) - if num_stages == len(stages): - found = 0 - for j in range(num_stages): - if i[j] == stages[j]: - found += 1 - if found == num_stages: - builder.append(1) - else: - builder.append(0) - else: - builder.append(0) - builder.end_list() - - -def mask(rec_stages, stages): - """create a mask on tracks.rec_stages . - - Parameters - ---------- - rec_stages : awkward1 Array - tracks.rec_stages . - stages : list - reconstruction stages of interest. - - Returns - ------- - awkward1 Array - an awkward1 Array mask where True corresponds to the positions - where stages were found. False otherwise. - """ - builder = ak1.ArrayBuilder() - _find(rec_stages, ak1.Array(stages), builder) - return builder.snapshot() == 1 - - -def best_track(tracks, strategy="first", rec_type=None, rec_stages=None): - """best track selection based on different strategies - - Parameters - ---------- - tracks : class km3io.offline.OfflineBranch - the tracks branch. - strategy : str - the trategy desired to select the best tracks. - """ - if strategy == "first": - 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", key="Evt", extra={