diff --git a/km3io/offline.py b/km3io/offline.py index 053c8b07252cce02b6636d3a0cee9d70b5a577e8..e39238c61ce523ebf0fc9be623871a8c1524d0bd 100644 --- a/km3io/offline.py +++ b/km3io/offline.py @@ -87,6 +87,19 @@ def count_nested(Array, axis=0): @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: @@ -106,6 +119,21 @@ def _find(rec_stages, stages, builder): 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