Skip to content
Snippets Groups Projects

Resolve "uproot4 integration"

Merged Resolve "uproot4 integration"
58-uproot4-integration-2 into master
Merged Tamas Gal requested to merge 58-uproot4-integration-2 into master
Compare
and
Show latest version
1 file
+ 56
− 408
Compare changes
  • Side-by-side
  • Inline
km3io/tools.py
+ 56
− 408
@@ -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).