Merge pull request #50 from CoffeaTeam/topic_awkward_numba

Update fnal-column-analysis-tools to opportunistically use awkward.numba

.travis.yml

@@ -7,7 +7,7 @@ python:
   - 3.6
 env:
   matrix:
-  - AWKWARD="awkward>=0.7.0"
+  - AWKWARD="awkward>=0.8.1"
 matrix:
 addons:
   apt:
@@ -17,6 +17,7 @@ script: python setup.py pytest
 install:
   - pip install --upgrade setuptools_scm
   - pip install awkward
+  - pip install awkward-numba
   - pip install numba
   - pip install scipy
   - pip install numpy --upgrade

fnal_column_analysis_tools/__init__.py

@@ -36,3 +36,5 @@
 import fnal_column_analysis_tools.jetmet_tools
 
 from fnal_column_analysis_tools.version import __version__
+
+import fnal_column_analysis_tools.util

fnal_column_analysis_tools/analysis_objects/JaggedCandidateArray.py

@@ -1,6 +1,6 @@
-import awkward
 import uproot_methods
-import numpy as np
+from fnal_column_analysis_tools.util import awkward
+from fnal_column_analysis_tools.util import numpy as np
 
 JaggedTLorentzVectorArray = awkward.Methods.mixin(uproot_methods.classes.TLorentzVector.ArrayMethods, awkward.JaggedArray)
 
@@ -66,7 +66,7 @@ def candidatesfromcounts(cls,counts,**kwargs):
                                                               mass=column4,
                                                               ...)
         """
-        offsets = awkward.array.jagged.counts2offsets(counts)
+        offsets = awkward.JaggedArray.counts2offsets(counts)
         return cls.candidatesfromoffsets(offsets,**kwargs)
 
     @classmethod
@@ -322,7 +322,7 @@ def cross(self, other, nested=False):
         keys = outs.columns
         reorder = False
         for key in keys:
-            if not isinstance(outs[key].content,awkward.array.table.Table):
+            if not isinstance(outs[key].content,awkward.Table):
                 del outs[key]
                 reorder = True
         if reorder:

fnal_column_analysis_tools/hist/hist_tools.py

@@ -1,10 +1,10 @@
 from __future__ import division
 from collections import namedtuple
+from fnal_column_analysis_tools.util import numpy as np
 import copy
 import functools
 import math
 import numbers
-import numpy as np
 import re
 import warnings
 

fnal_column_analysis_tools/hist/plot.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 from __future__ import division
-import numpy as np
+from fnal_column_analysis_tools.util import numpy as np
 import scipy.stats
 import copy
 import warnings

fnal_column_analysis_tools/lookup_tools/csv_converters.py

@@ -1,4 +1,5 @@
-import numpy as np
+from fnal_column_analysis_tools.util import awkward
+from fnal_column_analysis_tools.util import numpy as np
 
 #pt except for reshaping, then discriminant
 btag_feval_dims = {0:[1],1:[1],2:[1],3:[2]}

fnal_column_analysis_tools/lookup_tools/dense_evaluated_lookup.py

@@ -1,7 +1,7 @@
 from fnal_column_analysis_tools.lookup_tools.lookup_base import lookup_base
 
-import numpy as np
-from awkward.array.jagged import JaggedArray
+from fnal_column_analysis_tools.util import numpy as np
+from fnal_column_analysis_tools.util import awkward
 from copy import deepcopy
 import numba
 
@@ -61,7 +61,7 @@ def __init__(self,values,dims,feval_dim=None):
     def _evaluate(self,*args):        
         indices = [] 
         for arg in args: 
-            if type(arg) == JaggedArray: raise Exception('JaggedArray in inputs')
+            if type(arg) == awkward.JaggedArray: raise Exception('JaggedArray in inputs')
         if self._dimension == 1:
             indices.append(np.clip(np.searchsorted(self._axes, args[0], side='right')-1,0,self._values.shape[0]-1))
         else:

fnal_column_analysis_tools/lookup_tools/dense_lookup.py

@@ -1,7 +1,7 @@
 from fnal_column_analysis_tools.lookup_tools.lookup_base import lookup_base
 
-import numpy as np
-from awkward.array.jagged import JaggedArray
+from fnal_column_analysis_tools.util import awkward
+from fnal_column_analysis_tools.util import numpy as np
 from copy import deepcopy
 
 class dense_lookup(lookup_base):
@@ -30,7 +30,7 @@ def __init__(self,values,dims,feval_dim=None):
     def _evaluate(self,*args):        
         indices = [] 
         for arg in args: 
-            if type(arg) == JaggedArray: raise Exception('JaggedArray in inputs')
+            if type(arg) == awkward.JaggedArray: raise Exception('JaggedArray in inputs')
         if self._dimension == 1:
             indices.append(np.clip(np.searchsorted(self._axes, args[0], side='right')-1,0,self._values.shape[0]-1))
         else:

fnal_column_analysis_tools/lookup_tools/extractor.py

@@ -1,7 +1,7 @@
 from __future__ import print_function
 import os
 
-import numpy as np
+from fnal_column_analysis_tools.util import numpy as np
 from .evaluator import evaluator
 
 from .root_converters import convert_histo_root_file

fnal_column_analysis_tools/lookup_tools/jec_uncertainty_lookup.py

@@ -1,7 +1,7 @@
 from fnal_column_analysis_tools.lookup_tools.lookup_base import lookup_base
 
-import numpy as np
-from awkward.array.jagged import JaggedArray
+from fnal_column_analysis_tools.util import awkward
+from fnal_column_analysis_tools.util import numpy as np
 from copy import deepcopy
 
 from scipy.interpolate import interp1d

fnal_column_analysis_tools/lookup_tools/jersf_lookup.py

@@ -1,7 +1,7 @@
 from fnal_column_analysis_tools.lookup_tools.lookup_base import lookup_base
 
-import numpy as np
-from awkward.array.jagged import JaggedArray
+from fnal_column_analysis_tools.util import awkward
+from fnal_column_analysis_tools.util import numpy as np
 from copy import deepcopy
 
 def masked_bin_eval(dim1_indices, dimN_bins, dimN_vals):

fnal_column_analysis_tools/lookup_tools/jme_standard_function.py

@@ -1,7 +1,7 @@
 from fnal_column_analysis_tools.lookup_tools.lookup_base import lookup_base
 
-import numpy as np
-from awkward.array.jagged import JaggedArray
+from fnal_column_analysis_tools.util import awkward
+from fnal_column_analysis_tools.util import numpy as np
 from copy import deepcopy
 
 from numpy import sqrt,log,exp,abs
@@ -137,7 +137,7 @@ def _evaluate(self,*args):
             else:
                 idxs = flatten_idxs(bin_tuple,self._eval_clamp_mins[eval_name])
                 clamp_mins = self._eval_clamp_mins[eval_name].content[idxs]
-                if isinstance(clamp_mins,JaggedArray):
+                if isinstance(clamp_mins,awkward.JaggedArray):
                     if clamp_mins.content.size == 1:
                         clamp_mins = clamp_mins.content[0]
                     else:
@@ -148,7 +148,7 @@ def _evaluate(self,*args):
             else:
                 idxs = flatten_idxs(bin_tuple,self._eval_clamp_maxs[eval_name])
                 clamp_maxs = self._eval_clamp_maxs[eval_name].content[idxs]
-                if isinstance(clamp_maxs,JaggedArray):
+                if isinstance(clamp_maxs,awkward.JaggedArray):
                     if clamp_maxs.content.size == 1:
                         clamp_maxs = clamp_maxs.content[0]
                     else:

fnal_column_analysis_tools/lookup_tools/json_converters.py

@@ -1,4 +1,4 @@
-import numpy as np
+from fnal_column_analysis_tools.util import numpy as np
 import json
 
 def extract_json_histo_structure(parselevel,axis_names,axes):

fnal_column_analysis_tools/lookup_tools/lookup_base.py

@@ -1,5 +1,5 @@
-import numpy as np
-from awkward import JaggedArray
+from fnal_column_analysis_tools.util import awkward
+from fnal_column_analysis_tools.util import numpy as np
 
 #base class for all objects that do some sort of value or function lookup
 class lookup_base(object):
@@ -12,7 +12,7 @@ def __call__(self,*args):
         # TODO: check can use offsets (this should always be true for striped)
         # Alternatively we can just use starts and stops
         for i in range(len(inputs)):
-            if isinstance(inputs[i], JaggedArray):
+            if isinstance(inputs[i], awkward.JaggedArray):
                 if offsets is not None and offsets.base is not inputs[i].offsets.base:
                     if type(offsets) is int:
                         raise Exception('Do not mix JaggedArrays and numpy arrays when calling derived class of lookup_base')
@@ -27,7 +27,7 @@ def __call__(self,*args):
                 offsets = -1
         retval = self._evaluate(*tuple(inputs))
         if offsets is not None and type(offsets) is not int:
-            retval = JaggedArray.fromoffsets(offsets,retval)
+            retval = awkward.JaggedArray.fromoffsets(offsets,retval)
         return retval
 
     def _evaluate(self,*args):

fnal_column_analysis_tools/lookup_tools/root_converters.py

@@ -1,4 +1,4 @@
-import numpy as np
+from fnal_column_analysis_tools.util import numpy as np
 import uproot
 import re
 

fnal_column_analysis_tools/lookup_tools/txt_converters.py

@@ -1,6 +1,5 @@
-import numpy as np
-import awkward
-from awkward import JaggedArray
+from fnal_column_analysis_tools.util import awkward
+from fnal_column_analysis_tools.util import numpy as np
 
 #for later
 #func = numbaize(formula,['p%i'%i for i in range(nParms)]+[varnames[i] for i in range(nEvalVars)])
@@ -112,7 +111,7 @@ def _build_standard_jme_lookup(name,layout,pars,nBinnedVars,nBinColumns,
                 theBins = np.union1d(binMins,binMaxs)
                 allBins = np.append(allBins,theBins)
                 counts  = np.append(counts,theBins.size)
-            bins[layout[i+offset_name]] = JaggedArray.fromcounts(counts,allBins)
+            bins[layout[i+offset_name]] = awkward.JaggedArray.fromcounts(counts,allBins)
         bin_order.append(layout[i+offset_name])
         offset_col += 1
 
@@ -133,8 +132,8 @@ def _build_standard_jme_lookup(name,layout,pars,nBinnedVars,nBinColumns,
     for i in range(nEvalVars):
         var_order.append(layout[i+offset_name])
         if not interpolatedFunc:
-            clamp_mins[layout[i+offset_name]] = JaggedArray.fromcounts(jagged_counts,np.atleast_1d(pars[columns[i+offset_col]]))
-            clamp_maxs[layout[i+offset_name]] = JaggedArray.fromcounts(jagged_counts,np.atleast_1d(pars[columns[i+offset_col+1]]))
+            clamp_mins[layout[i+offset_name]] = awkward.JaggedArray.fromcounts(jagged_counts,np.atleast_1d(pars[columns[i+offset_col]]))
+            clamp_maxs[layout[i+offset_name]] = awkward.JaggedArray.fromcounts(jagged_counts,np.atleast_1d(pars[columns[i+offset_col+1]]))
             offset_col += 1
 
 
@@ -143,7 +142,7 @@ def _build_standard_jme_lookup(name,layout,pars,nBinnedVars,nBinColumns,
     parm_order = []
     offset_col = 2*nBinnedVars+1 + (interpolatedFunc==False)*2*nEvalVars
     for i in range(nParms):
-        parms.append(JaggedArray.fromcounts(jagged_counts,pars[columns[i+offset_col]]))
+        parms.append(awkward.JaggedArray.fromcounts(jagged_counts,pars[columns[i+offset_col]]))
         parm_order.append('p%i'%(i))
 
     wrapped_up = {}
@@ -262,7 +261,7 @@ def convert_effective_area_file(eaFilePath):
                 theBins = np.union1d(binMins,binMaxs)
                 allBins = np.append(allBins,theBins)
                 counts  = np.append(counts,theBins.size)
-            bins[layout[i+offset_name]] = JaggedArray.fromcounts(counts,allBins)
+            bins[layout[i+offset_name]] = awkward.JaggedArray.fromcounts(counts,allBins)
         bin_order.append(layout[i+offset_name])
         offset_col += 1
 

fnal_column_analysis_tools/lumi_tools/__init__.py

@@ -1,4 +1,4 @@
-import numpy as np
+from fnal_column_analysis_tools.util import numpy as np
 import json
 
 

fnal_column_analysis_tools/striped/ColumnGroup2JaggedTable.py

@@ -1,6 +1,6 @@
 from fnal_column_analysis_tools.analysis_objects.JaggedCandidateArray import JaggedCandidateArray
 from fnal_column_analysis_tools.striped.StripedColumnTransformer import PhysicalColumnGroup
-import awkward
+from fnal_column_analysis_tools.util import awkward
 
 def jaggedFromColumnGroup(cgroup):
     if isinstance(cgroup,PhysicalColumnGroup):

fnal_column_analysis_tools/util.py

@@ -0,0 +1,6 @@
+try:
+    import awkward.numba as awkward
+except :
+    import awkward as awkward
+
+from awkward.util import numpy

fnal_column_analysis_tools/version.py

@@ -30,7 +30,7 @@
 
 import re
 
-__version__ = "0.1.3"
+__version__ = "0.2.0"
 version = __version__
 version_info = tuple(re.split(r"[-\.]", __version__))
 

tests/dummy_distributions.py

@@ -1,5 +1,5 @@
-import numpy as np
-import awkward
+from fnal_column_analysis_tools.util import awkward
+from fnal_column_analysis_tools.util import numpy as np
 import uproot_methods
 
 
@@ -48,14 +48,14 @@ def __init__(self):
     return events()
 
 def gen_reco_TLV():
-    gen_pt = awkward.fromiter([[10.0, 20.0, 30.0], [], [40.0, 50.0]])
-    reco_pt = awkward.fromiter([[20.2, 10.1, 30.3, 50.5], [50.5], [60]])
+    gen_pt = awkward.JaggedArray.fromiter([[10.0, 20.0, 30.0], [], [40.0, 50.0]])
+    reco_pt = awkward.JaggedArray.fromiter([[20.2, 10.1, 30.3, 50.5], [50.5], [60]])
 
-    gen_eta = awkward.fromiter([[-3.0, -2.0, 2.0], [], [-1.0, 1.0]])
-    reco_eta = awkward.fromiter([[-2.2, -3.3, 2.2, 0.0], [0.0], [1.1]])
+    gen_eta = awkward.JaggedArray.fromiter([[-3.0, -2.0, 2.0], [], [-1.0, 1.0]])
+    reco_eta = awkward.JaggedArray.fromiter([[-2.2, -3.3, 2.2, 0.0], [0.0], [1.1]])
 
-    gen_phi = awkward.fromiter([[-1.5, 0.0, 1.5], [], [0.78, -0.78]])
-    reco_phi = awkward.fromiter([[ 0.1, -1.4, 1.4, 0.78], [0.78], [-0.77]])
+    gen_phi = awkward.JaggedArray.fromiter([[-1.5, 0.0, 1.5], [], [0.78, -0.78]])
+    reco_phi = awkward.JaggedArray.fromiter([[ 0.1, -1.4, 1.4, 0.78], [0.78], [-0.77]])
 
     gen = uproot_methods.TLorentzVectorArray.from_ptetaphim(gen_pt, gen_eta, gen_phi, 0.2)
     reco = uproot_methods.TLorentzVectorArray.from_ptetaphim(reco_pt, reco_eta, reco_phi, 0.2)

tests/test_analysis_objects.py

@@ -4,8 +4,8 @@
 from fnal_column_analysis_tools.analysis_objects import JaggedCandidateArray, JaggedTLorentzVectorArray
 import uproot
 import uproot_methods
-import awkward
-import numpy as np
+from fnal_column_analysis_tools.util import awkward
+from fnal_column_analysis_tools.util import numpy as np
 
 from dummy_distributions import dummy_four_momenta, gen_reco_TLV
 

tests/test_hist_tools.py

@@ -1,7 +1,7 @@
 from __future__ import print_function, division
 
 from fnal_column_analysis_tools import hist
-import numpy as np
+from fnal_column_analysis_tools.util import numpy as np
 
 from dummy_distributions import dummy_jagged_eta_pt
 

tests/test_lookup_tools.py

@@ -2,8 +2,8 @@
 
 from fnal_column_analysis_tools import lookup_tools
 import uproot
-import awkward
-import numpy as np
+from fnal_column_analysis_tools.util import awkward
+from fnal_column_analysis_tools.util import numpy as np
 
 from dummy_distributions import dummy_jagged_eta_pt
 

tests/test_striped.py

@@ -3,8 +3,8 @@
 from fnal_column_analysis_tools.striped import ColumnGroup, PhysicalColumnGroup, jaggedFromColumnGroup
 import uproot
 import uproot_methods
-import awkward
-import numpy as np
+from fnal_column_analysis_tools.util import awkward
+from fnal_column_analysis_tools.util import numpy as np
 
 from dummy_distributions import dummy_events