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
2 files
+ 143
− 551
Compare changes
  • Side-by-side
  • Inline
Files
2
km3io/tools.py
+ 96
− 478
#!/usr/bin/env python3
from collections import namedtuple
import numba as nb
import numpy as np
import awkward as ak
@@ -206,12 +207,14 @@ def get_multiplicity(tracks, rec_stages):
awkward1.Array
tracks multiplicty.
"""
masked_tracks = tracks[mask(tracks, stages=rec_stages)]
masked_tracks = tracks[mask(tracks.rec_stages, sequence=rec_stages)]
if tracks.is_single:
out = count_nested(masked_tracks.rec_stages, axis=0)
else:
out = count_nested(masked_tracks.rec_stages, axis=1)
try:
axis = tracks.ndim
except AttributeError:
axis = 0
out = count_nested(masked_tracks.rec_stages, axis=axis)
return out
@@ -252,517 +255,132 @@ def best_track(tracks, startend=None, minmax=None, stages=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:
selected_tracks = tracks[mask(tracks, stages=stages)]
if isinstance(stages, list):
m1 = mask(tracks.rec_stages, sequence=stages)
elif isinstance(stages, set):
m1 = mask(tracks.rec_stages, atleast=list(stages))
else:
raise ValueError("stages must be a list or a set of integers")
if startend is not None and minmax is None and stages is None:
selected_tracks = tracks[mask(tracks, startend=startend)]
m1 = mask(tracks.rec_stages, startend=startend)
if minmax is not None and startend is None and stages is None:
selected_tracks = tracks[mask(tracks, minmax=minmax)]
return _max_lik_track(_longest_tracks(selected_tracks))
m1 = mask(tracks.rec_stages, minmax=minmax)
def _longest_tracks(tracks):
"""Select the longest reconstructed track"""
if tracks.is_single:
stages_nesting_level = 1
tracks_nesting_level = 0
try:
axis = tracks.ndim
except AttributeError:
axis = 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
tracks = tracks[m1]
rec_stage_lengths = ak.num(tracks.rec_stages, axis=axis+1)
max_rec_stage_length = ak.max(rec_stage_lengths, axis=axis)
m2 = rec_stage_lengths == max_rec_stage_length
tracks = tracks[m2]
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
m3 = ak.argmax(tracks.lik, axis=axis, keepdims=True)
return tracks[tracks.lik == ak.max(tracks.lik, axis=tracks_nesting_level)]
out = tracks[m3]
if isinstance(out, ak.highlevel.Record):
return namedtuple("BestTrack", out.fields)(*[getattr(out, a)[0] for a in out.fields])
return out
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)
inputs = (sequence, startend, minmax, atleast)
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)
raise ValueError("either sequence, startend, minmax or atleast must be specified.")
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.
Parameters
----------
tracks : km3io.offline.OfflineBranch
tracks, or one track, or slice of tracks, or slices of tracks.
Returns
-------
km3io.offline.OfflineBranch
the longest + highest likelihood track reconstructed with JMUON.
"""
mask = _mask_rec_stages_in_range_min_max(
tracks, min_stage=krec.JMUONBEGIN, max_stage=krec.JMUONEND
)
return _max_lik_track(_longest_tracks(tracks[mask]))
"""Select the best JMUON track."""
return best_track(tracks, minmax=(krec.JMUONBEGIN, krec.JMUONEND))
def best_jshower(tracks):
"""Select the best JSHOWER track.
Parameters
----------
tracks : km3io.offline.OfflineBranch
tracks, or one track, or slice of tracks, or slices of tracks.
Returns
-------
km3io.offline.OfflineBranch
the longest + highest likelihood track reconstructed with JSHOWER.
"""
mask = _mask_rec_stages_in_range_min_max(
tracks, min_stage=krec.JSHOWERBEGIN, max_stage=krec.JSHOWEREND
)
return _max_lik_track(_longest_tracks(tracks[mask]))
"""Select the best JSHOWER track."""
return best_track(tracks, minmax=(krec.JSHOWERBEGIN, krec.JSHOWEREND))
def best_aashower(tracks):
"""Select the best AASHOWER track.
Parameters
----------
tracks : km3io.offline.OfflineBranch
tracks, or one track, or slice of tracks, or slices of tracks.
Returns
-------
km3io.offline.OfflineBranch
the longest + highest likelihood track reconstructed with AASHOWER.
"""
mask = _mask_rec_stages_in_range_min_max(
tracks, min_stage=krec.AASHOWERBEGIN, max_stage=krec.AASHOWEREND
)
return _max_lik_track(_longest_tracks(tracks[mask]))
"""Select the best AASHOWER track. """
return best_track(tracks, minmax=(krec.AASHOWERBEGIN, krec.AASHOWEREND))
def best_dusjshower(tracks):
"""Select the best DISJSHOWER track.
Parameters
----------
tracks : km3io.offline.OfflineBranch
tracks, or one track, or slice of tracks, or slices of tracks.
Returns
-------
km3io.offline.OfflineBranch
the longest + highest likelihood track reconstructed with DUSJSHOWER.
"""
mask = _mask_rec_stages_in_range_min_max(
tracks, min_stage=krec.DUSJSHOWERBEGIN, max_stage=krec.DUSJSHOWEREND
)
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()
"""Select the best DISJSHOWER track."""
return best_track(tracks, minmax=(krec.DUSJSHOWERBEGIN, krec.DUSJSHOWEREND))
def is_cc(fobj):