Fastpuppi v1 dask integration

python/boost_hist.py

@@ -1,7 +1,6 @@
-import awkward as ak
+import hist.dask as dah
 import hist
-from hist import Hist
-
+import awkward as ak
 
 def TH1F(name, title, nbins, bin_low, bin_high):
     b_axis_name = 'X'
@@ -10,12 +9,13 @@ def TH1F(name, title, nbins, bin_low, bin_high):
         b_axis_name = title_split[1]
     b_name = title_split[0]
     b_label = name
-    return Hist(
-        hist.axis.Regular(bins=nbins, start=bin_low, stop=bin_high, name=b_axis_name),
-        label=b_label,
-        name=b_name,
-        storage=hist.storage.Weight()
-        )
+
+    return hist.dask.Hist(
+         hist.axis.Regular(bins=nbins, start=bin_low, stop=bin_high, name=b_axis_name),
+         label=b_label,
+         name=b_name,
+         storage=hist.storage.Weight()
+         )
 
 def TH2F(name, title, x_nbins, x_bin_low, x_bin_high, y_nbins, y_bin_low, y_bin_high):
     b_x_axis_name = 'X'
@@ -27,43 +27,25 @@ def TH2F(name, title, x_nbins, x_bin_low, x_bin_high, y_nbins, y_bin_low, y_bin_
         b_y_axis_name = title_split[2]
     b_name = title_split[0]
     b_label = name
-    return Hist(
-        hist.axis.Regular(bins=x_nbins, start=x_bin_low, stop=x_bin_high, name=b_x_axis_name),
-        hist.axis.Regular(bins=y_nbins, start=y_bin_low, stop=y_bin_high, name=b_y_axis_name),
-        label=b_label,
+
+    return hist.dask.Hist(
+        hist.axis.Regular(bins=x_nbins, start=x_bin_low, stop=x_bin_high, name=b_x_axis_name), 
+        hist.axis.Regular(bins=y_nbins, start=y_bin_low, stop=y_bin_high, name=b_y_axis_name), 
+        label=b_label, 
         name=b_name,
         storage=hist.storage.Weight()
         )
 
-
-def TH2F_category(name, title, x_categories, y_nbins, y_bin_low, y_bin_high):
-    b_x_axis_name = 'X'
-    b_y_axis_name = 'Y'
-    title_split = title.split(';')
-    if len(title_split) > 1:
-        b_x_axis_name = title_split[1]
-    if len(title_split) > 2:
-        b_y_axis_name = title_split[2]
-    b_name = title_split[0]
-    b_label = name
-    return Hist(
-        hist.axis.StrCategory(x_categories, name=b_x_axis_name),
-        hist.axis.Regular(bins=y_nbins, start=y_bin_low, stop=y_bin_high, name=b_y_axis_name),
-        label=b_label,
-        name=b_name,
-        storage=hist.storage.Weight()
-        )
-
-
 def fill_1Dhist(hist, array, weights=None):
     flar = ak.drop_none(ak.flatten(array))
+
     if weights is None:
         hist.fill(flar, threads=None)
         # ROOT.fill_1Dhist(hist=hist, array=flar)
     else:
         hist.fill(flar, weights)
         # ROOT.fill_1Dhist(hist=hist, array=flar, weights=weights)
-
+        
 def fill_2Dhist(hist, arrayX, arrayY, weights=None):
     flar_x = ak.drop_none(ak.flatten(arrayX))
     flar_y = ak.drop_none(ak.flatten(arrayY))
@@ -73,5 +55,4 @@ def fill_2Dhist(hist, arrayX, arrayY, weights=None):
         hist.fill(flar_x, flar_y, threads=None)
     else:
         # ROOT.fill_2Dhist(hist=hist, arrayX=flar_x, arrayY=flar_y, weights=weights)
-        hist.fill(flar_x, flar_y, weights)
-
+        hist.fill(flar_x, flar_y, weights)

python/l1THistos.py

@@ -84,6 +84,10 @@ def write(self, upfile):
         for histo in [a for a in dir(self) if a.startswith('h_')]:
             writeable_hist = getattr(self, histo)
             # print (f"Writing {histo} class {writeable_hist.__class__.__name__}")
+
+            name = writeable_hist.label         
+            writeable_hist = writeable_hist.compute()
+
             if 'GraphBuilder' in writeable_hist.__class__.__name__ :
                 continue
             elif 'TH1' in writeable_hist.__class__.__name__ or 'TH2' in writeable_hist.__class__.__name__:
@@ -93,7 +97,7 @@ def write(self, upfile):
                 # print('ok')
             else:
                 up_writeable_hist = up.to_writable(writeable_hist)
-                upfile[f'{dir_name}/{writeable_hist.label}'] = up_writeable_hist
+                upfile[f'{dir_name}/{name}'] = up_writeable_hist
 
     # def normalize(self, norm):
     #     className = self.__class__.__name__
@@ -590,7 +594,7 @@ def __init__(self, name, root_file=None, debug=False):
             self.h_pfIsoPV = bh.TH1F(f'{name}_pfIsoPV', 'Iso; rel-iso^{PV}_{pf}', 100, 0, 2)
             self.h_n = bh.TH1F(f'{name}_n', '# objects per event', 100, 0, 100)
             self.h_compBdt = bh.TH1F(f'{name}_compBdt', 'BDT Score Comp ID', 50, 0, 1)
-
+            
         BaseHistos.__init__(self, name, root_file, debug)
 
     def fill(self, egs):

python/selections.py

@@ -612,7 +612,8 @@ def compare_selections(sel1, sel2):
 
 wps = working_points_histomax[version]
 labels = ['LE', 'HE']
-wls = zip(wps, labels, strict=False)
+
+wls = zip(wps, labels)
 # for i,
 tphgc_egbdt_sel = []
 

python/tree_reader.py

@@ -5,6 +5,9 @@
 import awkward as ak
 import vector
 
+import coffea
+from coffea.nanoevents import NanoEventsFactory, NanoAODSchema, BaseSchema
+
 vector.register_awkward()
 
 class TreeReader:
@@ -94,28 +97,40 @@ def getDataFrame(self, prefix, entry_block, fallback=None):
                     if br.startswith(f'{prefix}_') and
                     br != f'{prefix}_n']
         names = ['_'.join(br.split('_')[1:]) for br in branches]
-        name_map = dict(zip(names, branches, strict=False))
+        name_map = dict(zip(names, branches))
         if len(branches) == 0:
             if fallback is not None:
                 return self.getDataFrame(prefix=fallback, entry_block=entry_block)
             prefs = set([br.split('_')[0] for br in self._branches])
             print(f'stored branch prefixes are: {prefs}')
             raise ValueError(f'[TreeReader::getDataFrame] No branches with prefix: {prefix}')
 
-        akarray = self.tree.arrays(names,
-                                   library='ak',
-                                   aliases=name_map,
-                                   entry_start=self.file_entry,
-                                   entry_stop=self.file_entry+entry_block)
+        dask_akarray = NanoEventsFactory.from_root(
+                                   self.tree,
+                                   schemaclass=NanoAODSchema).events()
+
+        #akarray = self.tree.arrays(names, 
+        #library='ak', 
+        #aliases=name_map, 
+        #entry_start=self.file_entry, 
+        #entry_stop=self.file_entry+entry_block)
+        #print("[0] prefix to select: ", prefix)
+
+        dask_akarray = dask_akarray[prefix]     
+        #print("[1] Selected fields from prefix", dask_akarray.fields)   
+
+        dask_akarray = dask_akarray[names]
+
+        #print("[2] specific fields with names", dask_akarray.fields)
+        dask_akarray = dask_akarray[self.file_entry : self.file_entry + entry_block]
 
-        # print(akarray)
         records = {}
-        for field in akarray.fields:
-            records[field] = akarray[field]
+        for field in dask_akarray.fields:
+            records[field] = dask_akarray[field]
 
         if 'pt' in names and 'eta' in names and 'phi' in names:
             if 'mass' not in names and 'energy' not in names:
-                records['mass'] = 0.*akarray['pt']
+                records['mass'] = 0.*dask_akarray['pt']
             return vector.zip(records)
 
         return ak.zip(records)
@@ -124,5 +139,4 @@ def getDataFrame(self, prefix, entry_block, fallback=None):
         # ele_rec = ak.zip({'pt': tkele.pt, 'eta': tkele.eta, 'phi': tkele.phi}, with_name="pippo")
         # this would allow to handle the records and assign behaviours....
 
-        # return akarray
-
+        # return akarray