Feature/object collection

CollectObject.py

@@ -0,0 +1,62 @@
+import numpy as np
+import awkward as ak
+
+class CollectObject:
+
+    def __init__(self, events):
+        self.events = events
+
+    def CreateMask(self, variable, cut):
+        loaded = self.objects[variable]
+        sign = cut[0]
+        value = cut[1]
+        if sign == "==":
+            self.mask = self.mask & (loaded == value)
+        elif sign == ">=":
+            self.mask = self.mask & (loaded >= value)
+        elif sign == ">":
+            self.mask = self.mask & (loaded > value)
+        elif sign == "<=":
+            self.mask = self.mask & (loaded <= value)
+        elif sign == "<":
+            self.mask = self.mask & (loaded < value)
+        elif sign == "!=":
+            self.mask = self.mask & (loaded != value)
+        else:
+            pass
+
+    def Electron(self, cuts):
+        self.objects = self.events["Electron"]
+        self.mask = ak.ones_like(self.objects["pt"], dtype=bool)
+        for variable, cut in cuts.items():
+            self.CreateMask(variable, cut)
+        return self.objects[self.mask]
+
+    def Muon(self, cuts):
+        self.objects = self.events["Muon"]
+        self.mask = ak.ones_like(self.objects["pt"], dtype=bool)
+        for variable, cut in cuts.items():
+            self.CreateMask(variable, cut)
+        return self.objects[self.mask]
+
+    def Tau(self, cuts):
+        self.objects = self.events["Tau"]
+        self.mask = ak.ones_like(self.objects["pt"], dtype=bool)
+        for variable, cut in cuts.items():
+            self.CreateMask(variable, cut)
+        return self.objects[self.mask]
+
+    def Jet(self, cuts):
+        self.objects = self.events["Jet"]
+        self.mask = ak.ones_like(self.objects["pt"], dtype=bool)
+        for variable, cut in cuts.items():
+            self.CreateMask(variable, cut)
+        return self.objects[self.mask]
+
+    def FatJet(self, cuts):
+        self.objects = self.events["FatJet"]
+        self.mask = ak.ones_like(self.objects["pt"], dtype=bool)
+        for variable, cut in cuts.items():
+            self.CreateMask(variable, cut)
+        return self.objects[self.mask]
+

analyzer.py

@@ -9,6 +9,7 @@
 import dask.array as da
 import dask_awkward as dak
 from coffea.analysis_tools import PackedSelection
+from CollectObject import CollectObject
 import time
 import re
 
@@ -175,8 +176,15 @@ def process(self, events):
         # OBJECT SELECTION #
         ####################
 
-        # muon and electron selections are broken out into standalone functions
-        tightElectrons, looseElectrons = selectElectrons(events)
+        coll = CollectObject(events)
+        tightElectrons = coll.Electron({
+            "pt" : (">", 53),
+            "eta" : ("<", 2.4),
+        })
+        looseElectrons = coll.Electron({
+            "pt" : (">", 53),
+            "eta" : ("<", 2.4),
+        })
         nTightElectrons = ak.num(tightElectrons)
 
         tightMuons, looseMuons = selectMuons(events)