diff --git a/CollectObject.py b/CollectObject.py new file mode 100644 index 0000000..0b834a4 --- /dev/null +++ b/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] + diff --git a/analyzer.py b/analyzer.py index 165beec..f2f939e 100644 --- a/analyzer.py +++ b/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)