diff --git a/km3io/tools.py b/km3io/tools.py
index c42841919e0a5265aee6d41276f1ed65a7358e03..5d0ab67596a374f98d6ef0d41da73adb61771fed 100644
--- a/km3io/tools.py
+++ b/km3io/tools.py
@@ -263,222 +263,75 @@ def best_track(tracks, startend=None, minmax=None, stages=None):
return _max_lik_track(_longest_tracks(selected_tracks))
-def _longest_tracks(tracks):
- """Select the longest reconstructed track"""
- if tracks.is_single:
- stages_nesting_level = 1
- tracks_nesting_level = 0
-
- else:
- stages_nesting_level = 2
- tracks_nesting_level = 1
-
- len_stages = count_nested(tracks.rec_stages, axis=stages_nesting_level)
- longest = tracks[len_stages == ak.max(len_stages, axis=tracks_nesting_level)]
-
- return longest
-
-
-def _max_lik_track(tracks):
- """Select the track with the highest likelihood """
- if tracks.is_single:
- tracks_nesting_level = 0
- else:
- tracks_nesting_level = 1
-
- return tracks[tracks.lik == ak.max(tracks.lik, axis=tracks_nesting_level)]
-
-
-def mask(tracks, stages=None, startend=None, minmax=None):
- """Create a mask for tracks.rec_stages.
+def mask(arr, sequence=None, startend=None, minmax=None, atleast=None):
+ """Return a boolean mask which check each nested sub-array for a condition.
Parameters
----------
- tracks : km3io.offline.OfflineBranch
- tracks, or one track, or slice of tracks, or slice of one track.
- stages : list or set
- reconstruction stages of interest:
- - list: the order of rec_stages in respected.
- - set: the order of rec_stages in irrelevant.
+ arr : awkward.Array with ndim>=2
+ The array to mask.
startend: tuple(int, int), optional
- The required first and last stage in tracks.rec_stages.
+ True for entries where the first and last element are matching the tuple.
minmax: tuple(int, int), optional
- The range (minimum and maximum) value of rec_stages to take into account.
-
- Returns
- -------
- awkward1.Array(bool)
- an awkward1 Array mask where True corresponds to the positions
- where stages were found. False otherwise.
-
- Raises
- ------
- ValueError
- - too many inputs specified.
- - no inputs are specified.
+ True for entries where each element is within the min-max-range.
+ sequence : list(int), optional
+ True for entries which contain the exact same elements (in that specific
+ order)
+ atleast : list(int), optional
+ True for entries where at least the provided elements are present.
"""
- inputs = (stages, startend, minmax)
-
- if all(v is None for v in inputs):
- raise ValueError("either stages, startend or minmax must be specified.")
-
- if stages is not None and (startend is not None or minmax is not None):
- raise ValueError("Please specify either a range or a set of rec stages.")
-
- if stages is not None and startend is None and minmax is None:
- if isinstance(stages, list):
- # order of stages is conserved
- return _mask_explicit_rec_stages(tracks, stages)
- if isinstance(stages, set):
- # order of stages is no longer conserved
- return _mask_rec_stages_in_set(tracks, stages)
-
- if startend is not None and minmax is None and stages is None:
- return _mask_rec_stages_between_start_end(tracks, *startend)
-
- if minmax is not None and startend is None and stages is None:
- return _mask_rec_stages_in_range_min_max(tracks, *minmax)
-
-
-def _mask_rec_stages_between_start_end(tracks, start, end):
- """Mask tracks.rec_stages that start exactly with start and end exactly
- with end. ie [start, a, b ...,z , end]"""
builder = ak.ArrayBuilder()
- if tracks.is_single:
- _find_between_single(tracks.rec_stages, start, end, builder)
- return (builder.snapshot() == 1)[0]
- else:
- _find_between(tracks.rec_stages, start, end, builder)
- return builder.snapshot() == 1
-
-
-@nb.jit(nopython=True)
-def _find_between(rec_stages, start, end, builder):
- """Find tracks.rec_stages where rec_stages[0] == start and rec_stages[-1] == end."""
-
- for s in rec_stages:
- builder.begin_list()
- for i in s:
- num_stages = len(i)
- if num_stages != 0:
- if (i[0] == start) and (i[-1] == end):
- builder.append(1)
+ _mask(arr, builder, sequence, startend, minmax, atleast)
+ return builder.snapshot()
+
+#nb.njit # TODO: not supported in awkward yet
+# see https://github.com/scikit-hep/awkward-1.0/issues/572
+def _mask(arr, builder, sequence=None, startend=None, minmax=None, atleast=None):
+ if arr.ndim == 2: # recursion stop
+ if startend is not None:
+ start, end = startend
+ for els in arr:
+ if ak.count(els) > 0 and els[0] == start and els[-1] == end:
+ builder.boolean(True)
else:
- builder.append(0)
- else:
- builder.append(0)
- builder.end_list()
-
-
-@nb.jit(nopython=True)
-def _find_between_single(rec_stages, start, end, builder):
- """Find tracks.rec_stages where rec_stages[0] == start and
- rec_stages[-1] == end in a single track."""
-
- builder.begin_list()
- for s in rec_stages:
- num_stages = len(s)
- if num_stages != 0:
- if (s[0] == start) and (s[-1] == end):
- builder.append(1)
- else:
- builder.append(0)
- else:
- builder.append(0)
- builder.end_list()
-
-
-def _mask_explicit_rec_stages(tracks, stages):
- """Mask explicit rec_stages .
-
- Parameters
- ----------
- tracks : km3io.offline.OfflineBranch
- tracks or one track, or slice of tracks.
- stages : list
- reconstruction stages of interest. The order of stages is conserved.
-
- Returns
- -------
- awkward1.Array
- an awkward1 Array mask where True corresponds to the positions
- where stages were found. False otherwise.
- """
-
- builder = ak.ArrayBuilder()
- if tracks.is_single:
- _find_single(tracks.rec_stages, ak.Array(stages), builder)
- return (builder.snapshot() == 1)[0]
- else:
- _find(tracks.rec_stages, ak.Array(stages), builder)
- return builder.snapshot() == 1
-
-
-@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.
+ builder.boolean(False)
+ elif minmax is not None:
+ min, max = minmax
+ for els in arr:
+ for el in els:
+ if el < min or el > max:
+ builder.boolean(False)
+ break
+ else:
+ builder.boolean(True)
+ elif sequence is not None:
+ n = len(sequence)
+ for els in arr:
+ if len(els) != n:
+ builder.boolean(False)
+ else:
+ for i in range(n):
+ if els[i] != sequence[i]:
+ builder.boolean(False)
+ break
+ else:
+ builder.boolean(True)
+ elif atleast is not None:
+ for els in arr:
+ for e in atleast:
+ if e not in els:
+ builder.boolean(False)
+ break
+ else:
+ builder.boolean(True)
+ return
- Parameters
- ----------
- rec_stages : awkward1.Array
- tracks.rec_stages from multiple events.
- stages : awkward1.Array
- reconstruction stages of interest.
- builder : awkward1.highlevel.ArrayBuilder
- awkward1 Array builder.
- """
- for s in rec_stages:
+ for subarray in arr:
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)
+ _mask(subarray, builder, sequence, startend, minmax, atleast)
builder.end_list()
-@nb.jit(nopython=True)
-def _find_single(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 from a SINGLE event.
- stages : awkward1.Array
- reconstruction stages of interest.
- builder : awkward1.highlevel.ArrayBuilder
- awkward1 Array builder.
- """
- builder.begin_list()
- for s in rec_stages:
- num_stages = len(s)
- if num_stages == len(stages):
- found = 0
- for j in range(num_stages):
- if s[j] == stages[j]:
- found += 1
- if found == num_stages:
- builder.append(1)
- else:
- builder.append(0)
- else:
- builder.append(0)
- builder.end_list()
-
def best_jmuon(tracks):
"""Select the best JMUON track.
@@ -560,211 +413,6 @@ def best_dusjshower(tracks):
return _max_lik_track(_longest_tracks(tracks[mask]))
-def _mask_rec_stages_in_range_min_max(tracks, min_stage=None, max_stage=None):
- """Mask tracks where rec_stages are withing the range(min, max).
-
- Parameters
- ----------
- tracks : km3io.offline.OfflineBranch
- tracks, or one track, or slice of tracks, or slices of tracks.
- min_stage : int
- minimum value of rec_stages.
- max_stage : int
- maximum value of rec_stages.
-
-
- Returns
- -------
- awkward1.Array
- an awkward1 Array mask where True corresponds to the positions
- where stages were found. False otherwise.
- """
- if (min_stage is not None) and (max_stage is not None):
-
- builder = ak.ArrayBuilder()
- if tracks.is_single:
- _find_in_range_single(tracks.rec_stages, min_stage, max_stage, builder)
- return (builder.snapshot() == 1)[0]
- else:
- _find_in_range(tracks.rec_stages, min_stage, max_stage, builder)
- return builder.snapshot() == 1
-
- else:
- raise ValueError("please provide min_stage and max_stage.")
-
-
-@nb.jit(nopython=True)
-def _find_in_range(rec_stages, min_stage, max_stage, builder):
- """Construct an awkward1 array with the same structure as tracks.rec_stages.
-
- When stages are within the range(min, max), the Array is filled with
- value 1, otherwise it is filled with value 0.
-
- Parameters
- ----------
- rec_stages : awkward1.Array
- tracks.rec_stages of MULTILPLE events.
- min_stage: int
- minimum value of rec_stages.
- max_stage: int
- minimum value of rec_stages.
- 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 != 0:
- found = 0
- for j in i:
- if min_stage <= j <= max_stage:
- found += 1
- if found == num_stages:
- builder.append(1)
- else:
- builder.append(0)
- else:
- builder.append(0)
- builder.end_list()
-
-
-@nb.jit(nopython=True)
-def _find_in_range_single(rec_stages, min_stage, max_stage, builder):
- """Construct an awkward1 array with the same structure as tracks.rec_stages.
-
- When stages are within the range(min, max), the Array is filled with
- value 1, otherwise it is filled with value 0.
-
- Parameters
- ----------
- rec_stages : awkward1.Array
- tracks.rec_stages of a SINGLE event.
- min_stage: int
- minimum value of rec_stages.
- max_stage: int
- minimum value of rec_stages.
- builder : awkward1.highlevel.ArrayBuilder
- awkward1 Array builder.
- """
- builder.begin_list()
- for s in rec_stages:
- num_stages = len(s)
- if num_stages != 0:
- found = 0
- for i in s:
- if min_stage <= i <= max_stage:
- found += 1
- if found == num_stages:
- builder.append(1)
- else:
- builder.append(0)
- else:
- builder.append(0)
- builder.end_list()
-
-
-def _mask_rec_stages_in_set(tracks, stages):
- """Mask tracks where rec_stages are withing the range(min, max).
-
- Parameters
- ----------
- tracks : km3io.offline.OfflineBranch
- tracks, or one track, or slice of tracks, or slices of tracks.
- stages : set
- set of stages to look for in tracks.rec_stages.
-
- Returns
- -------
- awkward1.Array
- an awkward1 Array mask where True corresponds to the positions
- where stages were found. False otherwise.
- """
- if isinstance(stages, set):
-
- builder = ak.ArrayBuilder()
- if tracks.is_single:
- _find_in_set_single(tracks.rec_stages, stages, builder)
- return (builder.snapshot() == 1)[0]
- else:
- _find_in_set(tracks.rec_stages, stages, builder)
- return builder.snapshot() == 1
-
- else:
- raise ValueError("stages must be a set")
-
-
-@nb.jit(nopython=True)
-def _find_in_set(rec_stages, stages, builder):
- """Construct an awkward1 array with the same structure as tracks.rec_stages.
-
- When all stages are found in rec_stages, the Array is filled with
- value 1, otherwise it is filled with value 0.
-
- Parameters
- ----------
- rec_stages : awkward1.Array
- tracks.rec_stages of MULTILPLE events.
- stages : set
- set of stages.
- builder : awkward1.highlevel.ArrayBuilder
- awkward1 Array builder.
-
- """
- n = len(stages)
- for s in rec_stages:
- builder.begin_list()
- for i in s:
- num_stages = len(i)
- if num_stages != 0:
- found = 0
- for j in i:
- if j in stages:
- found += 1
- if found == n:
- builder.append(1)
- else:
- builder.append(0)
- else:
- builder.append(0)
- builder.end_list()
-
-
-@nb.jit(nopython=True)
-def _find_in_set_single(rec_stages, stages, builder):
- """Construct an awkward1 array with the same structure as tracks.rec_stages.
-
- When all stages are found in rec_stages, the Array is filled with
- value 1, otherwise it is filled with value 0.
-
- Parameters
- ----------
- rec_stages : awkward1.Array
- tracks.rec_stages of a SINGLE event.
- stages : set
- set of stages.
- builder : awkward1.highlevel.ArrayBuilder
- awkward1 Array builder.
- """
- n = len(stages)
- builder.begin_list()
- for s in rec_stages:
- num_stages = len(s)
- if num_stages != 0:
- found = 0
- for j in s:
- if j in stages:
- found += 1
- if found == n:
- builder.append(1)
- else:
- builder.append(0)
- else:
- builder.append(0)
- builder.end_list()
-
-
def is_cc(fobj):
"""Determin if events are a result of a Charged Curent interaction (CC)
or a Neutral Curent interaction (NC).