diff --git a/km3io/rootio.py b/km3io/rootio.py index 4668dae30a0a5a2e3d91efdd18e2f643007094d5..fe7cb3a5684e1b0fd6d698d25fd2dbd1b1fbadf2 100644 --- a/km3io/rootio.py +++ b/km3io/rootio.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 import numpy as np -import awkward1 as ak +import awkward as ak import uproot3 as uproot from .tools import unfold_indices diff --git a/km3io/tools.py b/km3io/tools.py index 2f62aeb889fb63893c3eddf402e5810099f3acec..11b7b419d31b0477a3cd32408d961da7777eb071 100644 --- a/km3io/tools.py +++ b/km3io/tools.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 import numba as nb import numpy as np -import awkward1 as ak1 +import awkward as ak import uproot3 as uproot from km3io.definitions import reconstruction as krec @@ -160,7 +160,7 @@ def fitinf(fitparam, tracks): out = params[:, index] else: params = fit[count_nested(fit, axis=2) > index] - out = ak1.Array([i[:, index] for i in params]) + out = ak.Array([i[:, index] for i in params]) return out @@ -183,11 +183,11 @@ def count_nested(arr, axis=0): counts of elements found in a nested awkward1 Array. """ if axis == 0: - return ak1.num(arr, axis=0) + return ak.num(arr, axis=0) if axis == 1: - return ak1.num(arr, axis=1) + return ak.num(arr, axis=1) if axis == 2: - return ak1.count(arr, axis=2) + return ak.count(arr, axis=2) def get_multiplicity(tracks, rec_stages): @@ -276,7 +276,7 @@ def _longest_tracks(tracks): tracks_nesting_level = 1 len_stages = count_nested(tracks.rec_stages, axis=stages_nesting_level) - longest = tracks[len_stages == ak1.max(len_stages, axis=tracks_nesting_level)] + longest = tracks[len_stages == ak.max(len_stages, axis=tracks_nesting_level)] return longest @@ -288,7 +288,7 @@ def _max_lik_track(tracks): else: tracks_nesting_level = 1 - return tracks[tracks.lik == ak1.max(tracks.lik, axis=tracks_nesting_level)] + return tracks[tracks.lik == ak.max(tracks.lik, axis=tracks_nesting_level)] def mask(tracks, stages=None, startend=None, minmax=None): @@ -345,7 +345,7 @@ def mask(tracks, stages=None, startend=None, minmax=None): 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 = ak1.ArrayBuilder() + builder = ak.ArrayBuilder() if tracks.is_single: _find_between_single(tracks.rec_stages, start, end, builder) return (builder.snapshot() == 1)[0] @@ -407,12 +407,12 @@ def _mask_explicit_rec_stages(tracks, stages): where stages were found. False otherwise. """ - builder = ak1.ArrayBuilder() + builder = ak.ArrayBuilder() if tracks.is_single: - _find_single(tracks.rec_stages, ak1.Array(stages), builder) + _find_single(tracks.rec_stages, ak.Array(stages), builder) return (builder.snapshot() == 1)[0] else: - _find(tracks.rec_stages, ak1.Array(stages), builder) + _find(tracks.rec_stages, ak.Array(stages), builder) return builder.snapshot() == 1 @@ -583,7 +583,7 @@ def _mask_rec_stages_in_range_min_max(tracks, min_stage=None, max_stage=None): """ if (min_stage is not None) and (max_stage is not None): - builder = ak1.ArrayBuilder() + builder = ak.ArrayBuilder() if tracks.is_single: _find_in_range_single(tracks.rec_stages, min_stage, max_stage, builder) return (builder.snapshot() == 1)[0] @@ -685,7 +685,7 @@ def _mask_rec_stages_in_set(tracks, stages): """ if isinstance(stages, set): - builder = ak1.ArrayBuilder() + builder = ak.ArrayBuilder() if tracks.is_single: _find_in_set_single(tracks.rec_stages, stages, builder) return (builder.snapshot() == 1)[0] @@ -788,14 +788,14 @@ def is_cc(fobj): """ program = fobj.header.simul.program w2list = fobj.events.w2list - len_w2lists = ak1.num(w2list, axis=1) + len_w2lists = ak.num(w2list, axis=1) if all(len_w2lists <= 7): # old nu file have w2list of len 7. usr_names = fobj.events.mc_tracks.usr_names usr_data = fobj.events.mc_tracks.usr mask_cc_flag = usr_names[:, 0] == b"cc" inter_ID = usr_data[:, 0][mask_cc_flag] - out = ak1.flatten(inter_ID == 2) # 2 is the interaction ID for CC. + out = ak.flatten(inter_ID == 2) # 2 is the interaction ID for CC. else: if "gseagen" in program.lower(): diff --git a/km3io/utils/kprinttree.py b/km3io/utils/kprinttree.py index 47d38c5ba6d52eb82c2134d6c6f29c861ee80063..7b5af9da0147ab07305cf52987ff72446c08ed70 100644 --- a/km3io/utils/kprinttree.py +++ b/km3io/utils/kprinttree.py @@ -14,7 +14,7 @@ Options: -h --help Show this screen. """ -import uproot4 as uproot +import uproot def print_tree(filename): diff --git a/requirements/install.txt b/requirements/install.txt index aa82996e87d817de8569229d5b6bd031214f5d7c..761911d0512f358850bcd09d5759efc9fbdcdea2 100644 --- a/requirements/install.txt +++ b/requirements/install.txt @@ -1,6 +1,6 @@ docopt numba>=0.50 -awkward1>=0.3.1 -uproot4 +awkward>=1.0.0rc1 +uproot>=4.0.0rc3 uproot3 setuptools_scm