def process(self, events):
# Initialize accumulator
out = self.accumulator.identity()
dataset = sample_name
# events.metadata['dataset']
# Stop processing if there is no event remain
if len(events) == 0:
return out
# Cut flow
cut0 = np.zeros(len(events))
# <----- Helper functions ------>#
# Sort by PT helper function
def sort_by_pt(ele, pho, jet):
ele = ele[ak.argsort(ele.pt, ascending=False, axis=1)]
pho = pho[ak.argsort(pho.pt, ascending=False, axis=1)]
jet = jet[ak.argsort(jet.pt, ascending=False, axis=1)]
return ele, pho, jet
# Lorentz vectors
from coffea.nanoevents.methods import vector
ak.behavior.update(vector.behavior)
def TLorentz_vector(vec):
vec = ak.zip(
{
"x": vec.x,
"y": vec.y,
"z": vec.z,
"t": vec.t
},
with_name="LorentzVector",
)
return vec
def TLorentz_vector_cylinder(vec):
vec = ak.zip(
{
"pt": vec.pt,
"eta": vec.eta,
"phi": vec.phi,
"mass": vec.mass,
},
with_name="PtEtaPhiMLorentzVector",
)
return vec
# Cut-based ID modification
@numba.njit
def PhotonVID(vid, idBit):
rBit = 0
for x in range(0, 7):
rBit |= (1 << x) if ((vid >> (x * 2)) & 0b11 >= idBit) else 0
return rBit
# Inverse Sieie and upper limit
@numba.njit
def make_fake_obj_mask(Pho, builder):
# for eventIdx,pho in enumerate(tqdm(Pho)): # --Event Loop
for eventIdx, pho in enumerate(Pho):
builder.begin_list()
if len(pho) < 1:
continue
for phoIdx, _ in enumerate(pho): # --Photon Loop
vid = Pho[eventIdx][phoIdx].vidNestedWPBitmap
vid_cuts1 = PhotonVID(vid, 1) # Loose photon
vid_cuts2 = PhotonVID(vid, 2) # Medium photon
vid_cuts3 = PhotonVID(vid, 3) # Tight photon
# Field name
# |0|0|0|0|0|0|0|
# |IsoPho|IsoNeu|IsoChg|Sieie|hoe|scEta|PT|
# 1. Turn off cut (ex turn off Sieie
# |1|1|1|0|1|1|1| = |1|1|1|0|1|1|1|
# 2. Inverse cut (ex inverse Sieie)
# |1|1|1|1|1|1|1| = |1|1|1|0|1|1|1|
# if (vid_cuts2 & 0b1111111 == 0b1111111): # Cut applied
# if vid_cuts2 & 0b1111111 == 0b1110111: # Inverse Sieie
if (vid_cuts2 & 0b1110111 == 0b1110111): # Without Sieie
builder.boolean(True)
else:
builder.boolean(False)
builder.end_list()
return builder
# Golden Json file
if self._year == "2018":
injson = "/x5/cms/jwkim/gitdir/JWCorp/JW_analysis/Coffea_WZG/Corrections/Cert_314472-325175_13TeV_Legacy2018_Collisions18_JSON.txt.RunABCD"
if self._year == "2017":
injson = "/x5/cms/jwkim/gitdir/JWCorp/JW_analysis/Coffea_WZG/Corrections/Cert_294927-306462_13TeV_UL2017_Collisions17_GoldenJSON.txt"
# --- Selection
Initial_events = events
# Good Run ( Golden Json files )
from coffea import lumi_tools
lumi_mask_builder = lumi_tools.LumiMask(injson)
lumimask = ak.Array(
lumi_mask_builder.__call__(events.run, events.luminosityBlock))
events = events[lumimask]
# print("{0}% of files pass good-run conditions".format(len(events)/ len(Initial_events)))
# Stop processing if there is no event remain
if len(events) == 0:
return out
##----------- Cut flow1: Passing Triggers
# double lepton trigger
is_double_ele_trigger = True
if not is_double_ele_trigger:
double_ele_triggers_arr = np.ones(len(events), dtype=np.bool)
else:
double_ele_triggers_arr = np.zeros(len(events), dtype=np.bool)
for path in self._doubleelectron_triggers[self._year]:
if path not in events.HLT.fields:
continue
double_ele_triggers_arr = double_ele_triggers_arr | events.HLT[
path]
# single lepton trigger
is_single_ele_trigger = True
if not is_single_ele_trigger:
single_ele_triggers_arr = np.ones(len(events), dtype=np.bool)
else:
single_ele_triggers_arr = np.zeros(len(events), dtype=np.bool)
for path in self._singleelectron_triggers[self._year]:
if path not in events.HLT.fields:
continue
single_ele_triggers_arr = single_ele_triggers_arr | events.HLT[
path]
events.Electron, events.Photon, events.Jet = sort_by_pt(
events.Electron, events.Photon, events.Jet)
# Apply cut1
Initial_events = events
# events = events[single_ele_triggers_arr | double_ele_triggers_arr]
events = events[double_ele_triggers_arr]
cut1 = np.ones(len(events))
# Set Particles
Electron = events.Electron
Muon = events.Muon
Photon = events.Photon
MET = events.MET
Jet = events.Jet
# Stop processing if there is no event remain
if len(Electron) == 0:
return out
# --Muon ( only used to calculate dR )
MuSelmask = ((Muon.pt >= 10)
& (abs(Muon.eta) <= 2.5)
& (Muon.tightId)
& (Muon.pfRelIso04_all < 0.15))
# Muon = ak.mask(Muon,MuSelmask)
Muon = Muon[MuSelmask]
##----------- Cut flow2: Electron Selection
EleSelmask = ((Electron.pt >= 20)
& (np.abs(Electron.eta + Electron.deltaEtaSC) < 1.479)
& (Electron.cutBased > 2)
& (abs(Electron.dxy) < 0.05)
& (abs(Electron.dz) < 0.1)) | (
(Electron.pt >= 20)
&
(np.abs(Electron.eta + Electron.deltaEtaSC) > 1.479)
& (np.abs(Electron.eta + Electron.deltaEtaSC) <= 2.5)
& (Electron.cutBased > 2)
& (abs(Electron.dxy) < 0.1)
& (abs(Electron.dz) < 0.2))
Electron = Electron[EleSelmask]
# apply cut 2
Tri_electron_mask = ak.num(Electron) >= 2
Electron = Electron[Tri_electron_mask]
Photon = Photon[Tri_electron_mask]
Jet = Jet[Tri_electron_mask]
MET = MET[Tri_electron_mask]
Muon = Muon[Tri_electron_mask]
events = events[Tri_electron_mask]
# Stop processing if there is no event remain
if len(Electron) == 0:
return out
cut2 = np.ones(len(Photon)) * 2
##----------- Cut flow3: Photon Selection
# Basic photon selection
isgap_mask = (abs(Photon.eta) < 1.442) | ((abs(Photon.eta) > 1.566) &
(abs(Photon.eta) < 2.5))
Pixel_seed_mask = ~Photon.pixelSeed
PT_mask = Photon.pt >= 20
# dR cut with selected Muon and Electrons
dr_pho_ele_mask = ak.all(Photon.metric_table(Electron) >= 0.5,
axis=-1) # default metric table: delta_r
dr_pho_mu_mask = ak.all(Photon.metric_table(Muon) >= 0.5, axis=-1)
# Photon basic mask
PhoSelmask = (PT_mask
& isgap_mask
& Pixel_seed_mask
& dr_pho_ele_mask
& dr_pho_mu_mask)
Photon = Photon[PhoSelmask]
# Apply cut 3
A_photon_mask = ak.num(Photon) > 0
Electron = Electron[A_photon_mask]
Photon = Photon[A_photon_mask]
Jet = Jet[A_photon_mask]
Muon = Muon[A_photon_mask]
MET = MET[A_photon_mask]
events = events[A_photon_mask]
# ID for fake photon
Photon_template_mask = make_fake_obj_mask(
Photon, ak.ArrayBuilder()).snapshot()
Photon = Photon[Photon_template_mask]
# Apply cut -Fake Photon -
A_photon_mask = ak.num(Photon) > 0
Electron = Electron[A_photon_mask]
Photon = Photon[A_photon_mask]
Jet = Jet[A_photon_mask]
Muon = Muon[A_photon_mask]
MET = MET[A_photon_mask]
events = events[A_photon_mask]
# Stop processing if there is no event remain
if len(Electron) == 0:
return out
cut3 = np.ones(len(Photon)) * 3
##----------- Cut flow4: Select 2 OSSF electrons from Z
@numba.njit
def find_2lep(events_leptons, builder):
for leptons in events_leptons:
builder.begin_list()
nlep = len(leptons)
for i0 in range(nlep):
for i1 in range(i0 + 1, nlep):
if leptons[i0].charge + leptons[i1].charge != 0:
continue
if nlep == 2:
builder.begin_tuple(2)
builder.index(0).integer(i0)
builder.index(1).integer(i1)
builder.end_tuple()
else:
for i2 in range(nlep):
if len({i0, i1, i2}) < 3:
continue
builder.begin_tuple(3)
builder.index(0).integer(i0)
builder.index(1).integer(i1)
builder.index(2).integer(i2)
builder.end_tuple()
builder.end_list()
return builder
ossf_idx = find_2lep(Electron, ak.ArrayBuilder()).snapshot()
# OSSF cut
ossf_mask = ak.num(ossf_idx) >= 1
ossf_idx = ossf_idx[ossf_mask]
Electron = Electron[ossf_mask]
Photon = Photon[ossf_mask]
Jet = Jet[ossf_mask]
MET = MET[ossf_mask]
Double_electron = [Electron[ossf_idx[idx]] for idx in "01"]
from coffea.nanoevents.methods import vector
ak.behavior.update(vector.behavior)
Diele = ak.zip({
"lep1":
Double_electron[0],
"lep2":
Double_electron[1],
"p4":
TLorentz_vector(Double_electron[0] + Double_electron[1]),
})
bestZ_idx = ak.singletons(
ak.argmin(abs(Diele.p4.mass - 91.1876), axis=1))
Diele = Diele[bestZ_idx]
cut4 = np.ones(len(Electron)) * 4
##----------- Cut flow 5: Event Selection
def make_leading_pair(target, base):
return target[ak.argmax(base.pt, axis=1, keepdims=True)]
leading_pho = make_leading_pair(Photon, Photon)
# Mee cut
Mee_cut_mask = ak.firsts(Diele.p4.mass) > 4
# Electron PT cuts
Elept_mask = ak.firsts((Diele.lep1.pt >= 25) & (Diele.lep2.pt >= 20))
# MET cuts
MET_mask = MET.pt > 20
# --------Mask -------#
Event_sel_mask = Mee_cut_mask & Elept_mask & MET_mask
Diele_sel = Diele[Event_sel_mask]
leading_pho_sel = leading_pho[Event_sel_mask]
Jet_sel = Jet[Event_sel_mask]
MET_sel = MET[Event_sel_mask]
cut5 = np.ones(len(Diele)) * 5
# Photon EE and EB
isEE_mask = leading_pho.isScEtaEE
isEB_mask = leading_pho.isScEtaEB
Pho_EE = leading_pho[isEE_mask & Event_sel_mask]
Pho_EB = leading_pho[isEB_mask & Event_sel_mask]
# -------------------- Flatten variables ---------------------------#
# -- Ele1 --#
Ele1_PT = ak.flatten(Diele_sel.lep1.pt)
Ele1_Eta = ak.flatten(Diele_sel.lep1.eta)
Ele1_Phi = ak.flatten(Diele_sel.lep1.phi)
# -- Ele2 --#
Ele2_PT = ak.flatten(Diele_sel.lep2.pt)
Ele2_Eta = ak.flatten(Diele_sel.lep2.eta)
Ele2_Phi = ak.flatten(Diele_sel.lep2.phi)
# -- Pho -- #
Pho_PT = ak.flatten(leading_pho_sel.pt)
Pho_Eta = ak.flatten(leading_pho_sel.eta)
Pho_Phi = ak.flatten(leading_pho_sel.phi)
# -- Pho EB --#
Pho_EB_PT = ak.flatten(Pho_EB.pt)
Pho_EB_Eta = ak.flatten(Pho_EB.eta)
Pho_EB_Phi = ak.flatten(Pho_EB.phi)
Pho_EB_Isochg = ak.flatten(Pho_EE.pfRelIso03_chg)
Pho_EB_Sieie = ak.flatten(Pho_EE.sieie)
# -- Pho EE --#
Pho_EE_PT = ak.flatten(Pho_EE.pt)
Pho_EE_Eta = ak.flatten(Pho_EE.eta)
Pho_EE_Phi = ak.flatten(Pho_EE.phi)
Pho_EE_Isochg = ak.flatten(Pho_EE.pfRelIso03_chg)
Pho_EE_Sieie = ak.flatten(Pho_EE.sieie)
# --Kinematics --#
Diele_mass = ak.flatten(Diele_sel.p4.mass)
eeg_vec = Diele_sel.p4 + leading_pho_sel
eeg_mass = ak.flatten(eeg_vec.mass)
leading_ele, subleading_ele = ak.flatten(
TLorentz_vector_cylinder(Diele_sel.lep1)), ak.flatten(
TLorentz_vector_cylinder(Diele_sel.lep2))
dR_e1pho = ak.flatten(
leading_ele.delta_r(leading_pho_sel)) # dR pho,ele1
dR_e2pho = ak.flatten(
subleading_ele.delta_r(leading_pho_sel)) # dR pho,ele2
dR_jpho = ak.flatten(Jet_sel[:, 0].delta_r(leading_pho_sel))
MET_PT = ak.to_numpy(MET_sel.pt)
# -------------------- Sieie bins---------------------------#
def make_bins(pt, eta, sieie, bin_range_str):
bin_dict = {
"PT_1_eta_1": (pt > 20) & (pt < 30) & (eta < 1),
"PT_1_eta_2": (pt > 20) & (pt < 30) & (eta > 1) & (eta < 1.5),
"PT_1_eta_3": (pt > 20) & (pt < 30) & (eta > 1.5) & (eta < 2),
"PT_1_eta_4": (pt > 20) & (pt < 30) & (eta > 2) & (eta < 2.5),
"PT_2_eta_1": (pt > 30) & (pt < 40) & (eta < 1),
"PT_2_eta_2": (pt > 30) & (pt < 40) & (eta > 1) & (eta < 1.5),
"PT_2_eta_3": (pt > 30) & (pt < 40) & (eta > 1.5) & (eta < 2),
"PT_2_eta_4": (pt > 30) & (pt < 40) & (eta > 2) & (eta < 2.5),
"PT_3_eta_1": (pt > 40) & (pt < 50) & (eta < 1),
"PT_3_eta_2": (pt > 40) & (pt < 50) & (eta > 1) & (eta < 1.5),
"PT_3_eta_3": (pt > 40) & (pt < 50) & (eta > 1.5) & (eta < 2),
"PT_3_eta_4": (pt > 40) & (pt < 50) & (eta > 2) & (eta < 2.5),
"PT_4_eta_1": (pt > 50) & (eta < 1),
"PT_4_eta_2": (pt > 50) & (eta > 1) & (eta < 1.5),
"PT_4_eta_3": (pt > 50) & (eta > 1.5) & (eta < 2),
"PT_4_eta_4": (pt > 50) & (eta > 2) & (eta < 2.5),
}
binmask = bin_dict[bin_range_str]
return ak.to_numpy(sieie[binmask])
bin_name_list = [
"PT_1_eta_1",
"PT_1_eta_2",
"PT_1_eta_3",
"PT_1_eta_4",
"PT_2_eta_1",
"PT_2_eta_2",
"PT_2_eta_3",
"PT_2_eta_4",
"PT_3_eta_1",
"PT_3_eta_2",
"PT_3_eta_3",
"PT_3_eta_4",
"PT_4_eta_1",
"PT_4_eta_2",
"PT_4_eta_3",
"PT_4_eta_4",
]
binned_sieie_hist = {}
for name in bin_name_list:
binned_sieie_hist[name] = make_bins(
ak.flatten(leading_pho_sel.pt),
ak.flatten(abs(leading_pho_sel.eta)),
ak.flatten(leading_pho_sel.sieie),
name,
)
# -------------------- Fill hist ---------------------------#
# Initial events
out["sumw"][dataset] += len(Initial_events)
# print("cut5: ",len(cut5))
# Cut flow loop
for cut in [cut0, cut1, cut2, cut3, cut4, cut5]:
out["cutflow"].fill(dataset=dataset, cutflow=cut)
# --Ele1 -- #
out["ele1pt"].fill(dataset=dataset, ele1pt=Ele1_PT)
out["ele1eta"].fill(dataset=dataset, ele1eta=Ele1_Eta)
out["ele1phi"].fill(dataset=dataset, ele1phi=Ele1_Phi)
# --Ele2 -- #
out["ele2pt"].fill(dataset=dataset, ele2pt=Ele2_PT)
out["ele2eta"].fill(dataset=dataset, ele2eta=Ele2_Eta)
out["ele2phi"].fill(dataset=dataset, ele2phi=Ele2_Phi)
# --Photon-- #
out["phopt"].fill(dataset=dataset, phopt=Pho_PT)
out["phoeta"].fill(dataset=dataset, phoeta=Pho_Eta)
out["phophi"].fill(dataset=dataset, phophi=Pho_Phi)
# --Photon EB --#
out["pho_EB_pt"].fill(
dataset=dataset,
pho_EB_pt=Pho_EB_PT,
)
out["pho_EB_eta"].fill(
dataset=dataset,
pho_EB_eta=Pho_EB_Eta,
)
out["pho_EB_phi"].fill(
dataset=dataset,
pho_EB_phi=Pho_EB_Phi,
)
out["pho_EB_sieie"].fill(
dataset=dataset,
pho_EB_sieie=Pho_EB_Sieie,
)
out["pho_EB_Iso_chg"].fill(dataset=dataset,
pho_EB_Iso_chg=Pho_EB_Isochg)
# --Photon EE --#
out["pho_EE_pt"].fill(
dataset=dataset,
pho_EE_pt=Pho_EE_PT,
)
out["pho_EE_eta"].fill(
dataset=dataset,
pho_EE_eta=Pho_EE_Eta,
)
out["pho_EE_phi"].fill(
dataset=dataset,
pho_EE_phi=Pho_EE_Phi,
)
out["pho_EE_sieie"].fill(
dataset=dataset,
pho_EE_sieie=Pho_EE_Sieie,
)
out["pho_EE_Iso_chg"].fill(dataset=dataset,
pho_EE_Iso_chg=Pho_EE_Isochg)
# -- Kinematic variables -- #
out["mass"].fill(dataset=dataset, Mee=Diele_mass)
out["mass_eea"].fill(dataset=dataset, mass_eea=eeg_mass)
out["met"].fill(dataset=dataset, met=MET_PT)
out["dR_ae1"].fill(dataset=dataset, dR_ae1=dR_e1pho)
out["dR_ae2"].fill(dataset=dataset, dR_ae2=dR_e2pho)
out["dR_aj"].fill(dataset=dataset, dR_aj=dR_jpho)
# test_target = binned_sieie_hist['PT_1_eta_1']
# print("CheckAAA: ",test_target[test_target > 0.05])
# print("CheckBBB: ",test_target[test_target > 0])
# -- Binned sieie hist -- #
if len(binned_sieie_hist["PT_1_eta_1"] > 0):
out["PT_1_eta_1"].fill(dataset=dataset,
PT_1_eta_1=binned_sieie_hist["PT_1_eta_1"])
if len(binned_sieie_hist["PT_1_eta_2"] > 0):
out["PT_1_eta_2"].fill(dataset=dataset,
PT_1_eta_2=binned_sieie_hist["PT_1_eta_2"])
if len(binned_sieie_hist["PT_1_eta_3"] > 0):
out["PT_1_eta_3"].fill(dataset=dataset,
PT_1_eta_3=binned_sieie_hist["PT_1_eta_3"])
if len(binned_sieie_hist["PT_1_eta_4"] > 0):
out["PT_1_eta_4"].fill(dataset=dataset,
PT_1_eta_4=binned_sieie_hist["PT_1_eta_4"])
if len(binned_sieie_hist["PT_2_eta_1"] > 0):
out["PT_2_eta_1"].fill(dataset=dataset,
PT_2_eta_1=binned_sieie_hist["PT_2_eta_1"])
if len(binned_sieie_hist["PT_2_eta_2"] > 0):
out["PT_2_eta_2"].fill(dataset=dataset,
PT_2_eta_2=binned_sieie_hist["PT_2_eta_2"])
if len(binned_sieie_hist["PT_2_eta_3"] > 0):
out["PT_2_eta_3"].fill(dataset=dataset,
PT_2_eta_3=binned_sieie_hist["PT_2_eta_3"])
if len(binned_sieie_hist["PT_2_eta_4"] > 0):
out["PT_2_eta_4"].fill(dataset=dataset,
PT_2_eta_4=binned_sieie_hist["PT_2_eta_4"])
if len(binned_sieie_hist["PT_3_eta_1"] > 0):
out["PT_3_eta_1"].fill(dataset=dataset,
PT_3_eta_1=binned_sieie_hist["PT_3_eta_1"])
if len(binned_sieie_hist["PT_3_eta_2"] > 0):
out["PT_3_eta_2"].fill(dataset=dataset,
PT_3_eta_2=binned_sieie_hist["PT_3_eta_2"])
if len(binned_sieie_hist["PT_3_eta_3"] > 0):
out["PT_3_eta_3"].fill(dataset=dataset,
PT_3_eta_3=binned_sieie_hist["PT_3_eta_3"])
if len(binned_sieie_hist["PT_3_eta_4"] > 0):
out["PT_3_eta_4"].fill(dataset=dataset,
PT_3_eta_4=binned_sieie_hist["PT_3_eta_4"])
if len(binned_sieie_hist["PT_4_eta_1"] > 0):
out["PT_4_eta_1"].fill(dataset=dataset,
PT_4_eta_1=binned_sieie_hist["PT_4_eta_1"])
if len(binned_sieie_hist["PT_4_eta_2"] > 0):
out["PT_4_eta_2"].fill(dataset=dataset,
PT_4_eta_2=binned_sieie_hist["PT_4_eta_2"])
if len(binned_sieie_hist["PT_4_eta_3"] > 0):
out["PT_4_eta_3"].fill(dataset=dataset,
PT_4_eta_3=binned_sieie_hist["PT_4_eta_3"])
if len(binned_sieie_hist["PT_4_eta_4"] > 0):
print("## show me the last bin: ", binned_sieie_hist["PT_4_eta_4"])
out["PT_4_eta_4"].fill(dataset=dataset,
PT_4_eta_4=binned_sieie_hist["PT_4_eta_4"])
return out
def process(self, events):
# Initialize accumulator
out = self.accumulator.identity()
dataset = sample_name
#events.metadata['dataset']
# Data or MC
isData = 'genWeight' not in events.fields
#Stop processing if there is no event remain
if len(events) == 0:
return out
# Golden Json file
if (self._year == "2018") and isData:
injson = "/x5/cms/jwkim/gitdir/JWCorp/JW_analysis/Coffea_WZG/Corrections/Cert_314472-325175_13TeV_Legacy2018_Collisions18_JSON.txt.RunABCD"
if (self._year == "2017") and isData:
injson = "/x5/cms/jwkim/gitdir/JWCorp/JW_analysis/Coffea_WZG/Corrections/Cert_294927-306462_13TeV_UL2017_Collisions17_GoldenJSON.txt"
# <----- Get Scale factors ------>#
if not isData:
# Egamma reco ID
get_ele_reco_above20_sf = self._corrections[
'get_ele_reco_above20_sf'][self._year]
get_ele_medium_id_sf = self._corrections['get_ele_medium_id_sf'][
self._year]
get_pho_medium_id_sf = self._corrections['get_pho_medium_id_sf'][
self._year]
# DoubleEG trigger # 2016, 2017 are not applied yet
if self._year == "2018":
get_ele_trig_leg1_SF = self._corrections[
'get_ele_trig_leg1_SF'][self._year]
get_ele_trig_leg1_data_Eff = self._corrections[
'get_ele_trig_leg1_data_Eff'][self._year]
get_ele_trig_leg1_mc_Eff = self._corrections[
'get_ele_trig_leg1_mc_Eff'][self._year]
get_ele_trig_leg2_SF = self._corrections[
'get_ele_trig_leg2_SF'][self._year]
get_ele_trig_leg2_data_Eff = self._corrections[
'get_ele_trig_leg2_data_Eff'][self._year]
get_ele_trig_leg2_mc_Eff = self._corrections[
'get_ele_trig_leg2_mc_Eff'][self._year]
# PU weight with custom made npy and multi-indexing
pu_weight_idx = ak.values_astype(events.Pileup.nTrueInt, "int64")
pu = self._puweight_arr[pu_weight_idx]
selection = processor.PackedSelection()
# Cut flow
cut0 = np.zeros(len(events))
# <----- Helper functions ------>#
# Sort by PT helper function
def sort_by_pt(ele, pho, jet):
ele = ele[ak.argsort(ele.pt, ascending=False, axis=1)]
pho = pho[ak.argsort(pho.pt, ascending=False, axis=1)]
jet = jet[ak.argsort(jet.pt, ascending=False, axis=1)]
return ele, pho, jet
# Lorentz vectors
from coffea.nanoevents.methods import vector
ak.behavior.update(vector.behavior)
def TLorentz_vector(vec):
vec = ak.zip({
"x": vec.x,
"y": vec.y,
"z": vec.z,
"t": vec.t
},
with_name="LorentzVector")
return vec
def TLorentz_vector_cylinder(vec):
vec = ak.zip(
{
"pt": vec.pt,
"eta": vec.eta,
"phi": vec.phi,
"mass": vec.mass,
},
with_name="PtEtaPhiMLorentzVector",
)
return vec
# Cut-based ID modification
@numba.njit
def PhotonVID(vid, idBit):
rBit = 0
for x in range(0, 7):
rBit |= (1 << x) if ((vid >> (x * 2)) & 0b11 >= idBit) else 0
return rBit
# Inverse Sieie and upper limit
@numba.njit
def make_fake_obj_mask(Pho, builder):
#for eventIdx,pho in enumerate(tqdm(Pho)): # --Event Loop
for eventIdx, pho in enumerate(Pho):
builder.begin_list()
if len(pho) < 1: continue
for phoIdx, _ in enumerate(pho): # --Photon Loop
vid = Pho[eventIdx][phoIdx].vidNestedWPBitmap
vid_cuts1 = PhotonVID(vid, 1) # Loose photon
vid_cuts2 = PhotonVID(vid, 2) # Medium photon
vid_cuts3 = PhotonVID(vid, 3) # Tight photon
# Field name
# |0|0|0|0|0|0|0|
# |IsoPho|IsoNeu|IsoChg|Sieie|hoe|scEta|PT|
# 1. Turn off cut (ex turn off Sieie
# |1|1|1|0|1|1|1| = |1|1|1|0|1|1|1|
# 2. Inverse cut (ex inverse Sieie)
# |1|1|1|1|1|1|1| = |1|1|1|0|1|1|1|
#if (vid_cuts2 & 0b1111111 == 0b1111111): # Cut applied
#if (vid_cuts2 & 0b1111111 == 0b1110111): # Inverse Sieie
if (vid_cuts2 & 0b1110111 == 0b1110111): # Without Sieie
builder.boolean(True)
else:
builder.boolean(False)
builder.end_list()
return builder
# <----- Selection ------>#
Initial_events = events
# Good Run ( Golden Json files )
from coffea import lumi_tools
if isData:
lumi_mask_builder = lumi_tools.LumiMask(injson)
lumimask = ak.Array(
lumi_mask_builder.__call__(events.run, events.luminosityBlock))
events = events[lumimask]
#print("{0}% of files pass good-run conditions".format(len(events)/ len(Initial_events)))
# Stop processing if there is no event remain
if len(events) == 0:
return out
##----------- Cut flow1: Passing Triggers
# double lepton trigger
is_double_ele_trigger = True
if not is_double_ele_trigger:
double_ele_triggers_arr = np.ones(len(events), dtype=np.bool)
else:
double_ele_triggers_arr = np.zeros(len(events), dtype=np.bool)
for path in self._doubleelectron_triggers[self._year]:
if path not in events.HLT.fields: continue
double_ele_triggers_arr = double_ele_triggers_arr | events.HLT[
path]
# single lepton trigger
is_single_ele_trigger = True
if not is_single_ele_trigger:
single_ele_triggers_arr = np.ones(len(events), dtype=np.bool)
else:
single_ele_triggers_arr = np.zeros(len(events), dtype=np.bool)
for path in self._singleelectron_triggers[self._year]:
if path not in events.HLT.fields: continue
single_ele_triggers_arr = single_ele_triggers_arr | events.HLT[
path]
events.Electron, events.Photon, events.Jet = sort_by_pt(
events.Electron, events.Photon, events.Jet)
# Good Primary vertex
nPV = events.PV.npvsGood
if not isData: nPV = nPV * pu
nPV_nw = nPV
# Apply cut1
events = events[double_ele_triggers_arr]
if not isData: pu = pu[double_ele_triggers_arr]
cut1 = np.ones(len(events))
# Set Particles
Electron = events.Electron
Muon = events.Muon
Photon = events.Photon
MET = events.MET
Jet = events.Jet
# Stop processing if there is no event remain
if len(Electron) == 0:
return out
# --Gen Photon for dR
genparts = events.GenPart
pdgID_mask = (genparts.pdgId == 22)
# mask2: isPrompt | fromHardProcess | isLastCopy
mask2 = (1 << 0) | (1 << 8) | (1 << 13)
# https://github.com/PKUHEPEWK/WGamma/blob/master/2018/wgRealPhotonTemplateModule.py
status_mask = ((genparts.statusFlags & mask2) == mask2)
gen_photons = genparts[pdgID_mask & status_mask]
assert (ak.all(ak.num(gen_photons) == 1)
) # Raise error if len(gen_photon) != 1
# --Muon ( only used to calculate dR )
MuSelmask = (Muon.pt >= 10) & (abs(
Muon.eta) <= 2.5) & (Muon.tightId) & (Muon.pfRelIso04_all < 0.15)
Muon = Muon[MuSelmask]
##----------- Cut flow2: Electron Selection
EleSelmask = ((Electron.pt >= 20) & (np.abs(Electron.eta + Electron.deltaEtaSC) < 1.479) & (Electron.cutBased > 2) & (abs(Electron.dxy) < 0.05) & (abs(Electron.dz) < 0.1)) | \
((Electron.pt >= 20) & (np.abs(Electron.eta + Electron.deltaEtaSC) > 1.479) & (np.abs(Electron.eta + Electron.deltaEtaSC) <= 2.5) & (Electron.cutBased > 2) & (abs(Electron.dxy) < 0.1) & (abs(Electron.dz) < 0.2))
Electron = Electron[EleSelmask]
# apply cut 2
Tri_electron_mask = ak.num(Electron) >= 2
Electron = Electron[Tri_electron_mask]
Photon = Photon[Tri_electron_mask]
Jet = Jet[Tri_electron_mask]
MET = MET[Tri_electron_mask]
Muon = Muon[Tri_electron_mask]
if not isData: pu = pu[Tri_electron_mask]
events = events[Tri_electron_mask]
gen_photons = gen_photons[Tri_electron_mask]
# Stop processing if there is no event remain
if len(Electron) == 0:
return out
cut2 = np.ones(len(Photon)) * 2
##----------- Cut flow3: Photon Selection
# Basic photon selection
isgap_mask = (abs(Photon.eta) < 1.442) | ((abs(Photon.eta) > 1.566) &
(abs(Photon.eta) < 2.5))
Pixel_seed_mask = ~Photon.pixelSeed
PT_mask = Photon.pt >= 20
# dR cut with selected Muon and Electrons
dr_pho_ele_mask = ak.all(Photon.metric_table(Electron) >= 0.5,
axis=-1) # default metric table: delta_r
dr_pho_mu_mask = ak.all(Photon.metric_table(Muon) >= 0.5, axis=-1)
PhoSelmask = PT_mask & isgap_mask & Pixel_seed_mask & dr_pho_ele_mask & dr_pho_mu_mask
Photon = Photon[PhoSelmask]
# Apply cut 3
A_photon_mask = ak.num(Photon) > 0
Electron = Electron[A_photon_mask]
Photon = Photon[A_photon_mask]
Jet = Jet[A_photon_mask]
Muon = Muon[A_photon_mask]
MET = MET[A_photon_mask]
if not isData: pu = pu[A_photon_mask]
events = events[A_photon_mask]
gen_photons = gen_photons[A_photon_mask]
Photon_template_mask = make_fake_obj_mask(
Photon, ak.ArrayBuilder()).snapshot()
Photon = Photon[Photon_template_mask]
# Apply cut 3
A_photon_mask = ak.num(Photon) > 0
Electron = Electron[A_photon_mask]
Photon = Photon[A_photon_mask]
Jet = Jet[A_photon_mask]
Muon = Muon[A_photon_mask]
MET = MET[A_photon_mask]
if not isData: pu = pu[A_photon_mask]
events = events[A_photon_mask]
gen_photons = gen_photons[A_photon_mask]
# Stop processing if there is no event remain
if len(Electron) == 0:
return out
cut3 = np.ones(len(Photon)) * 3
## -- Additional photon selection: Photon gen-matching
# Choose Photons that dR(genPhoton,Photon) <= 0.1
gen_match_photon_mask = ak.all(Photon.metric_table(gen_photons) <= 0.1,
axis=-1)
# Apply cut
Photon = Photon[gen_match_photon_mask]
gen_match_photon_evt_mask = ak.num(Photon) >= 1
Electron = Electron[gen_match_photon_evt_mask]
Photon = Photon[gen_match_photon_evt_mask]
Jet = Jet[gen_match_photon_evt_mask]
MET = MET[gen_match_photon_evt_mask]
gen_photons = gen_photons[gen_match_photon_evt_mask]
if not isData: pu = pu[gen_match_photon_evt_mask]
events = events[gen_match_photon_evt_mask]
##----------- Cut flow4: Select 2 OSSF electrons from Z
@numba.njit
def find_2lep(events_leptons, builder):
for leptons in events_leptons:
builder.begin_list()
nlep = len(leptons)
for i0 in range(nlep):
for i1 in range(i0 + 1, nlep):
if leptons[i0].charge + leptons[i1].charge != 0:
continue
if nlep == 2:
builder.begin_tuple(2)
builder.index(0).integer(i0)
builder.index(1).integer(i1)
builder.end_tuple()
else:
for i2 in range(nlep):
if len({i0, i1, i2}) < 3: continue
builder.begin_tuple(3)
builder.index(0).integer(i0)
builder.index(1).integer(i1)
builder.index(2).integer(i2)
builder.end_tuple()
builder.end_list()
return builder
ossf_idx = find_2lep(Electron, ak.ArrayBuilder()).snapshot()
# OSSF cut
ossf_mask = ak.num(ossf_idx) >= 1
ossf_idx = ossf_idx[ossf_mask]
Electron = Electron[ossf_mask]
Photon = Photon[ossf_mask]
Jet = Jet[ossf_mask]
MET = MET[ossf_mask]
events = events[ossf_mask]
if not isData: pu = pu[ossf_mask]
Double_electron = [Electron[ossf_idx[idx]] for idx in "01"]
from coffea.nanoevents.methods import vector
ak.behavior.update(vector.behavior)
Diele = ak.zip({
"lep1":
Double_electron[0],
"lep2":
Double_electron[1],
"p4":
TLorentz_vector(Double_electron[0] + Double_electron[1])
})
bestZ_idx = ak.singletons(
ak.argmin(abs(Diele.p4.mass - 91.1876), axis=1))
Diele = Diele[bestZ_idx]
# Stop processing if there is no event remain
if len(Electron) == 0:
return out
cut4 = np.ones(len(Electron)) * 4
leading_ele = Diele.lep1
subleading_ele = Diele.lep2
def make_leading_pair(target, base):
return target[ak.argmax(base.pt, axis=1, keepdims=True)]
leading_pho = make_leading_pair(Photon, Photon)
# -- Scale Factor for each electron
# Trigger weight helper function
def Trigger_Weight(eta1, pt1, eta2, pt2):
per_ev_MC =\
get_ele_trig_leg1_mc_Eff(eta1,pt1) * get_ele_trig_leg2_mc_Eff(eta2,pt2) +\
get_ele_trig_leg1_mc_Eff(eta2,pt2) * get_ele_trig_leg2_mc_Eff(eta1,pt1) -\
get_ele_trig_leg1_mc_Eff(eta1,pt1) * get_ele_trig_leg1_mc_Eff(eta2,pt2)
per_ev_data =\
get_ele_trig_leg1_data_Eff(eta1,pt1) * get_ele_trig_leg1_SF(eta1,pt1) * get_ele_trig_leg2_data_Eff(eta2,pt2) * get_ele_trig_leg2_SF(eta2,pt2) +\
get_ele_trig_leg1_data_Eff(eta2,pt2) * get_ele_trig_leg1_SF(eta2,pt2) * get_ele_trig_leg2_data_Eff(eta1,pt1) * get_ele_trig_leg2_SF(eta1,pt1) -\
get_ele_trig_leg1_data_Eff(eta1,pt1) * get_ele_trig_leg1_SF(eta1,pt1) * get_ele_trig_leg1_data_Eff(eta2,pt2) * get_ele_trig_leg1_SF(eta2,pt2)
return per_ev_data / per_ev_MC
if not isData:
## -------------< Egamma ID and Reco Scale factor > -----------------##
get_pho_medium_id_sf = get_pho_medium_id_sf(
ak.flatten(leading_pho.eta), ak.flatten(leading_pho.pt))
ele_reco_sf = get_ele_reco_above20_sf(
ak.flatten(leading_ele.deltaEtaSC + leading_ele.eta),
ak.flatten(leading_ele.pt)) * get_ele_reco_above20_sf(
ak.flatten(subleading_ele.deltaEtaSC + subleading_ele.eta),
ak.flatten(subleading_ele.pt))
ele_medium_id_sf = get_ele_medium_id_sf(
ak.flatten(leading_ele.deltaEtaSC + leading_ele.eta),
ak.flatten(leading_ele.pt)) * get_ele_medium_id_sf(
ak.flatten(subleading_ele.deltaEtaSC + subleading_ele.eta),
ak.flatten(subleading_ele.pt))
## -------------< Double Electron Trigger Scale factor > -----------------##
eta1 = ak.flatten(leading_ele.deltaEtaSC + leading_ele.eta)
eta2 = ak.flatten(subleading_ele.deltaEtaSC + subleading_ele.eta)
pt1 = ak.flatten(leading_ele.pt)
pt2 = ak.flatten(subleading_ele.pt)
# -- 2017,2016 are not applied yet
if self._year == '2018':
ele_trig_weight = Trigger_Weight(eta1, pt1, eta2, pt2)
##----------- Cut flow5: Event selection
# Mee cut
Mee_cut_mask = ak.firsts(Diele.p4.mass) > 4
# Electron PT cuts
Elept_mask = ak.firsts((Diele.lep1.pt >= 25) & (Diele.lep2.pt >= 20))
# MET cuts
MET_mask = MET.pt > 20
# --------Mask -------#
Event_sel_mask = Mee_cut_mask & Elept_mask & MET_mask
Diele_sel = Diele[Event_sel_mask]
leading_pho_sel = leading_pho[Event_sel_mask]
Jet_sel = Jet[Event_sel_mask]
MET_sel = MET[Event_sel_mask]
# Photon EE and EB
isEE_mask = leading_pho.isScEtaEE
isEB_mask = leading_pho.isScEtaEB
Pho_EE = leading_pho[isEE_mask & Event_sel_mask]
Pho_EB = leading_pho[isEB_mask & Event_sel_mask]
#Stop processing if there is no event remain
if len(leading_pho_sel) == 0:
return out
cut5 = np.ones(len(Diele)) * 5
# -------------------- Flatten variables ---------------------------#
# -- Ele1 --#
Ele1_PT = ak.flatten(Diele_sel.lep1.pt)
Ele1_Eta = ak.flatten(Diele_sel.lep1.eta)
Ele1_Phi = ak.flatten(Diele_sel.lep1.phi)
# -- Ele2 --#
Ele2_PT = ak.flatten(Diele_sel.lep2.pt)
Ele2_Eta = ak.flatten(Diele_sel.lep2.eta)
Ele2_Phi = ak.flatten(Diele_sel.lep2.phi)
# -- Pho -- #
Pho_PT = ak.flatten(leading_pho_sel.pt)
Pho_Eta = ak.flatten(leading_pho_sel.eta)
Pho_Phi = ak.flatten(leading_pho_sel.phi)
# -- Pho EB --#
Pho_EB_PT = ak.flatten(Pho_EB.pt)
Pho_EB_Eta = ak.flatten(Pho_EB.eta)
Pho_EB_Phi = ak.flatten(Pho_EB.phi)
Pho_EB_Isochg = ak.flatten(Pho_EE.pfRelIso03_chg)
Pho_EB_Sieie = ak.flatten(Pho_EE.sieie)
# -- Pho EE --#
Pho_EE_PT = ak.flatten(Pho_EE.pt)
Pho_EE_Eta = ak.flatten(Pho_EE.eta)
Pho_EE_Phi = ak.flatten(Pho_EE.phi)
Pho_EE_Isochg = ak.flatten(Pho_EE.pfRelIso03_chg)
Pho_EE_Sieie = ak.flatten(Pho_EE.sieie)
# --Kinematics --#
Diele_mass = ak.flatten(Diele_sel.p4.mass)
leading_ele, subleading_ele = ak.flatten(
TLorentz_vector_cylinder(Diele_sel.lep1)), ak.flatten(
TLorentz_vector_cylinder(Diele_sel.lep2))
dR_e1pho = ak.flatten(
leading_ele.delta_r(leading_pho_sel)) # dR pho,ele1
dR_e2pho = ak.flatten(
subleading_ele.delta_r(leading_pho_sel)) # dR pho,ele2
dR_jpho = ak.flatten(Jet_sel[:, 0].delta_r(leading_pho_sel))
MET_PT = ak.to_numpy(MET_sel.pt)
# -------------------- Sieie bins---------------------------#
def make_bins(pt, eta, sieie, bin_range_str):
bin_dict = {
'PT_1_eta_1': (pt > 20) & (pt < 30) & (eta < 1),
'PT_1_eta_2': (pt > 20) & (pt < 30) & (eta > 1) & (eta < 1.5),
'PT_1_eta_3': (pt > 20) & (pt < 30) & (eta > 1.5) & (eta < 2),
'PT_1_eta_4': (pt > 20) & (pt < 30) & (eta > 2) & (eta < 2.5),
'PT_2_eta_1': (pt > 30) & (pt < 40) & (eta < 1),
'PT_2_eta_2': (pt > 30) & (pt < 40) & (eta > 1) & (eta < 1.5),
'PT_2_eta_3': (pt > 30) & (pt < 40) & (eta > 1.5) & (eta < 2),
'PT_2_eta_4': (pt > 30) & (pt < 40) & (eta > 2) & (eta < 2.5),
'PT_3_eta_1': (pt > 40) & (pt < 50) & (eta < 1),
'PT_3_eta_2': (pt > 40) & (pt < 50) & (eta > 1) & (eta < 1.5),
'PT_3_eta_3': (pt > 40) & (pt < 50) & (eta > 1.5) & (eta < 2),
'PT_3_eta_4': (pt > 40) & (pt < 50) & (eta > 2) & (eta < 2.5),
'PT_4_eta_1': (pt > 50) & (eta < 1),
'PT_4_eta_2': (pt > 50) & (eta > 1) & (eta < 1.5),
'PT_4_eta_3': (pt > 50) & (eta > 1.5) & (eta < 2),
'PT_4_eta_4': (pt > 50) & (eta > 2) & (eta < 2.5)
}
binmask = bin_dict[bin_range_str]
return ak.to_numpy(sieie[binmask]), binmask
bin_name_list = [
'PT_1_eta_1', 'PT_1_eta_2', 'PT_1_eta_3', 'PT_1_eta_4',
'PT_2_eta_1', 'PT_2_eta_2', 'PT_2_eta_3', 'PT_2_eta_4',
'PT_3_eta_1', 'PT_3_eta_2', 'PT_3_eta_3', 'PT_3_eta_4',
'PT_4_eta_1', 'PT_4_eta_2', 'PT_4_eta_3', 'PT_4_eta_4'
]
binned_sieie_hist = {}
binmask_dict = {}
for name in bin_name_list:
binned_sieie_hist[name], _ = make_bins(
ak.flatten(leading_pho_sel.pt),
ak.flatten(abs(leading_pho_sel.eta)),
ak.flatten(leading_pho_sel.sieie), name)
_, binmask_dict[name] = make_bins(ak.flatten(leading_pho.pt),
ak.flatten(abs(leading_pho.eta)),
ak.flatten(leading_pho.sieie),
name)
print("Show me the last bin: ", binned_sieie_hist['PT_4_eta_4'])
# --- Apply weight and hist
weights = processor.Weights(len(cut4))
# --- skim cut-weight
def skim_weight(arr):
mask1 = ~ak.is_none(arr)
subarr = arr[mask1]
mask2 = subarr != 0
return ak.to_numpy(subarr[mask2])
cuts = Event_sel_mask
cuts_pho_EE = ak.flatten(isEE_mask)
cuts_pho_EB = ak.flatten(isEB_mask)
print(
"cut0: {0}, cut1: {1}, cut2: {2}, cut3: {3}, cut4: {4} ,cut5 {5} ".
format(len(Initial_events), len(cut1), len(cut2), len(cut3),
len(cut4), len(cut5)))
# Weight and SF here
if not isData:
weights.add('pileup', pu)
weights.add('ele_id', ele_medium_id_sf)
weights.add('pho_id', get_pho_medium_id_sf)
weights.add('ele_reco', ele_reco_sf)
# 2016,2017 are not applied yet
if self._year == "2018":
weights.add('ele_trigger', ele_trig_weight)
# ---------------------------- Fill hist --------------------------------------#
# Initial events
out["sumw"][dataset] += len(Initial_events)
# Cut flow loop
for cut in [cut0, cut1, cut2, cut3, cut4, cut5]:
out["cutflow"].fill(dataset=dataset, cutflow=cut)
# Primary vertex
out['nPV'].fill(
dataset=dataset,
nPV=nPV,
)
out['nPV_nw'].fill(dataset=dataset, nPV_nw=nPV_nw)
# Fill hist
# -- met -- #
out["met"].fill(dataset=dataset,
met=MET_PT,
weight=skim_weight(weights.weight() * cuts))
# --mass -- #
out["mass"].fill(dataset=dataset,
mass=Diele_mass,
weight=skim_weight(weights.weight() * cuts))
# -- Ele1 -- #
out["ele1pt"].fill(dataset=dataset,
ele1pt=Ele1_PT,
weight=skim_weight(weights.weight() * cuts))
out["ele1eta"].fill(dataset=dataset,
ele1eta=Ele1_Eta,
weight=skim_weight(weights.weight() * cuts))
out["ele1phi"].fill(dataset=dataset,
ele1phi=Ele1_Phi,
weight=skim_weight(weights.weight() * cuts))
# --Ele2 --#
out["ele2pt"].fill(dataset=dataset,
ele2pt=Ele2_PT,
weight=skim_weight(weights.weight() * cuts))
out["ele2eta"].fill(dataset=dataset,
ele2eta=Ele2_Eta,
weight=skim_weight(weights.weight() * cuts))
out["ele2phi"].fill(dataset=dataset,
ele2phi=Ele2_Phi,
weight=skim_weight(weights.weight() * cuts))
# -- Photon -- #
out["phopt"].fill(dataset=dataset,
phopt=Pho_PT,
weight=skim_weight(weights.weight() * cuts))
out["phoeta"].fill(dataset=dataset,
phoeta=Pho_Eta,
weight=skim_weight(weights.weight() * cuts))
out["phophi"].fill(dataset=dataset,
phophi=Pho_Phi,
weight=skim_weight(weights.weight() * cuts))
# -- Binned sieie hist -- #
if len(binned_sieie_hist['PT_1_eta_1'] > 0):
out['PT_1_eta_1'].fill(dataset=dataset,
PT_1_eta_1=binned_sieie_hist['PT_1_eta_1'])
if len(binned_sieie_hist['PT_1_eta_2'] > 0):
out['PT_1_eta_2'].fill(dataset=dataset,
PT_1_eta_2=binned_sieie_hist['PT_1_eta_2'])
if len(binned_sieie_hist['PT_1_eta_3'] > 0):
out['PT_1_eta_3'].fill(dataset=dataset,
PT_1_eta_3=binned_sieie_hist['PT_1_eta_3'])
if len(binned_sieie_hist['PT_1_eta_4'] > 0):
out['PT_1_eta_4'].fill(dataset=dataset,
PT_1_eta_4=binned_sieie_hist['PT_1_eta_4'])
if len(binned_sieie_hist['PT_2_eta_1'] > 0):
out['PT_2_eta_1'].fill(dataset=dataset,
PT_2_eta_1=binned_sieie_hist['PT_2_eta_1'])
if len(binned_sieie_hist['PT_2_eta_2'] > 0):
out['PT_2_eta_2'].fill(dataset=dataset,
PT_2_eta_2=binned_sieie_hist['PT_2_eta_2'])
if len(binned_sieie_hist['PT_2_eta_3'] > 0):
out['PT_2_eta_3'].fill(dataset=dataset,
PT_2_eta_3=binned_sieie_hist['PT_2_eta_3'])
if len(binned_sieie_hist['PT_2_eta_4'] > 0):
out['PT_2_eta_4'].fill(dataset=dataset,
PT_2_eta_4=binned_sieie_hist['PT_2_eta_4'])
if len(binned_sieie_hist['PT_3_eta_1'] > 0):
out['PT_3_eta_1'].fill(dataset=dataset,
PT_3_eta_1=binned_sieie_hist['PT_3_eta_1'])
if len(binned_sieie_hist['PT_3_eta_2'] > 0):
out['PT_3_eta_2'].fill(dataset=dataset,
PT_3_eta_2=binned_sieie_hist['PT_3_eta_2'])
if len(binned_sieie_hist['PT_3_eta_3'] > 0):
out['PT_3_eta_3'].fill(dataset=dataset,
PT_3_eta_3=binned_sieie_hist['PT_3_eta_3'])
if len(binned_sieie_hist['PT_3_eta_4'] > 0):
out['PT_3_eta_4'].fill(dataset=dataset,
PT_3_eta_4=binned_sieie_hist['PT_3_eta_4'])
if len(binned_sieie_hist['PT_4_eta_1'] > 0):
out['PT_4_eta_1'].fill(dataset=dataset,
PT_4_eta_1=binned_sieie_hist['PT_4_eta_1'])
if len(binned_sieie_hist['PT_4_eta_2'] > 0):
out['PT_4_eta_2'].fill(dataset=dataset,
PT_4_eta_2=binned_sieie_hist['PT_4_eta_2'])
if len(binned_sieie_hist['PT_4_eta_3'] > 0):
out['PT_4_eta_3'].fill(dataset=dataset,
PT_4_eta_3=binned_sieie_hist['PT_4_eta_3'])
if len(binned_sieie_hist['PT_4_eta_4'] > 0):
out['PT_4_eta_4'].fill(dataset=dataset,
PT_4_eta_4=binned_sieie_hist['PT_4_eta_4'])
return out
def process(self, events):
# Initialize accumulator
out = self.accumulator.identity()
dataset = sample_name
# events.metadata['dataset']
# Data or MC
isData = "genWeight" not in events.fields
isFake = self._isFake
# Stop processing if there is no event remain
if len(events) == 0:
return out
# Golden Json file
if (self._year == "2018") and isData:
injson = "/x5/cms/jwkim/gitdir/JWCorp/JW_analysis/Coffea_WZG/Corrections/Cert_314472-325175_13TeV_Legacy2018_Collisions18_JSON.txt.RunABD"
if (self._year == "2017") and isData:
injson = "/x5/cms/jwkim/gitdir/JWCorp/JW_analysis/Coffea_WZG/Corrections/Cert_294927-306462_13TeV_UL2017_Collisions17_GoldenJSON.txt"
# <----- Get Scale factors ------>#
if not isData:
# Egamma reco ID
get_ele_reco_above20_sf = self._corrections[
"get_ele_reco_above20_sf"][self._year]
get_ele_medium_id_sf = self._corrections["get_ele_medium_id_sf"][
self._year]
get_pho_medium_id_sf = self._corrections["get_pho_medium_id_sf"][
self._year]
# DoubleEG trigger # 2016, 2017 are not applied yet
if self._year == "2018":
get_ele_trig_leg1_SF = self._corrections[
"get_ele_trig_leg1_SF"][self._year]
get_ele_trig_leg1_data_Eff = self._corrections[
"get_ele_trig_leg1_data_Eff"][self._year]
get_ele_trig_leg1_mc_Eff = self._corrections[
"get_ele_trig_leg1_mc_Eff"][self._year]
get_ele_trig_leg2_SF = self._corrections[
"get_ele_trig_leg2_SF"][self._year]
get_ele_trig_leg2_data_Eff = self._corrections[
"get_ele_trig_leg2_data_Eff"][self._year]
get_ele_trig_leg2_mc_Eff = self._corrections[
"get_ele_trig_leg2_mc_Eff"][self._year]
# Muon ID, Iso
get_mu_tight_id_sf = self._corrections["get_mu_tight_id_sf"][
self._year]
get_mu_tight_iso_sf = self._corrections["get_mu_tight_iso_sf"][
self._year]
# PU weight with custom made npy and multi-indexing
pu_weight_idx = ak.values_astype(events.Pileup.nTrueInt, "int64")
pu = self._puweight_arr[pu_weight_idx]
# <----- Helper functions ------>#
# Sort by PT helper function
def sort_by_pt(ele, pho, jet):
ele = ele[ak.argsort(ele.pt, ascending=False, axis=1)]
pho = pho[ak.argsort(pho.pt, ascending=False, axis=1)]
jet = jet[ak.argsort(jet.pt, ascending=False, axis=1)]
return ele, pho, jet
# Lorentz vectors
from coffea.nanoevents.methods import vector
ak.behavior.update(vector.behavior)
def TLorentz_vector(vec):
vec = ak.zip(
{
"x": vec.x,
"y": vec.y,
"z": vec.z,
"t": vec.t
},
with_name="LorentzVector",
)
return vec
def TLorentz_vector_cylinder(vec):
vec = ak.zip(
{
"pt": vec.pt,
"eta": vec.eta,
"phi": vec.phi,
"mass": vec.mass,
},
with_name="PtEtaPhiMLorentzVector",
)
return vec
# <----- Selection ------>#
Initial_events = events
# Good Run ( Golden Json files )
from coffea import lumi_tools
if isData:
lumi_mask_builder = lumi_tools.LumiMask(injson)
lumimask = ak.Array(
lumi_mask_builder.__call__(events.run, events.luminosityBlock))
events = events[lumimask]
# print("{0}% of files pass good-run conditions".format(len(events)/ len(Initial_events)))
# Stop processing if there is no event remain
if len(events) == 0:
return out
# Cut flow
cut0 = np.zeros(len(events))
##----------- Cut flow1: Passing Triggers
# double lepton trigger
is_double_ele_trigger = True
if not is_double_ele_trigger:
double_ele_triggers_arr = np.ones(len(events), dtype=np.bool)
else:
double_ele_triggers_arr = np.zeros(len(events), dtype=np.bool)
for path in self._doubleelectron_triggers[self._year]:
if path not in events.HLT.fields:
continue
double_ele_triggers_arr = double_ele_triggers_arr | events.HLT[
path]
# single lepton trigger
is_single_ele_trigger = True
if not is_single_ele_trigger:
single_ele_triggers_arr = np.ones(len(events), dtype=np.bool)
else:
single_ele_triggers_arr = np.zeros(len(events), dtype=np.bool)
for path in self._singleelectron_triggers[self._year]:
if path not in events.HLT.fields:
continue
single_ele_triggers_arr = single_ele_triggers_arr | events.HLT[
path]
events.Electron, events.Photon, events.Jet = sort_by_pt(
events.Electron, events.Photon, events.Jet)
# Good Primary vertex
nPV = events.PV.npvsGood
nPV_nw = events.PV.npvsGood
if not isData:
nPV = nPV * pu
print(pu)
# Apply cut1
events = events[double_ele_triggers_arr]
if not isData:
pu = pu[double_ele_triggers_arr]
# Stop processing if there is no event remain
if len(events) == 0:
return out
cut1 = np.ones(len(events))
# Set Particles
Electron = events.Electron
Muon = events.Muon
Photon = events.Photon
MET = events.MET
Jet = events.Jet
##----------- Cut flow2: Muon Selection
MuSelmask = ((Muon.pt >= 10)
& (abs(Muon.eta) <= 2.5)
& (Muon.tightId)
& (Muon.pfRelIso04_all < 0.15))
Muon = Muon[MuSelmask]
# Exatly one muon
Muon_sel_mask = ak.num(Muon) == 1
Electron = Electron[Muon_sel_mask]
Photon = Photon[Muon_sel_mask]
Jet = Jet[Muon_sel_mask]
MET = MET[Muon_sel_mask]
Muon = Muon[Muon_sel_mask]
events = events[Muon_sel_mask]
if not isData:
pu = pu[Muon_sel_mask]
# Stop processing if there is no event remain
if len(Electron) == 0:
return out
cut2 = np.ones(len(Photon)) * 2
##----------- Cut flow3: Electron Selection
EleSelmask = ((Electron.pt >= 10)
& (np.abs(Electron.eta + Electron.deltaEtaSC) < 1.479)
& (Electron.cutBased > 2)
& (abs(Electron.dxy) < 0.05)
& (abs(Electron.dz) < 0.1)) | (
(Electron.pt >= 10)
&
(np.abs(Electron.eta + Electron.deltaEtaSC) > 1.479)
& (np.abs(Electron.eta + Electron.deltaEtaSC) <= 2.5)
& (Electron.cutBased > 2)
& (abs(Electron.dxy) < 0.1)
& (abs(Electron.dz) < 0.2))
Electron = Electron[EleSelmask]
# Exactly two electrons
ee_mask = ak.num(Electron) == 2
Electron = Electron[ee_mask]
Photon = Photon[ee_mask]
Jet = Jet[ee_mask]
MET = MET[ee_mask]
Muon = Muon[ee_mask]
if not isData:
pu = pu[ee_mask]
events = events[ee_mask]
# Stop processing if there is no event remain
if len(Electron) == 0:
return out
cut3 = np.ones(len(Photon)) * 3
##----------- Cut flow4: Photon Selection
# Basic photon selection
isgap_mask = (abs(Photon.eta) < 1.442) | ((abs(Photon.eta) > 1.566) &
(abs(Photon.eta) < 2.5))
Pixel_seed_mask = ~Photon.pixelSeed
if (dataset == "ZZ") and (self._year == "2017"):
PT_ID_mask = (Photon.pt >= 20) & (
Photon.cutBasedBitmap >= 3
) # 2^0(Loose) + 2^1(Medium) + 2^2(Tights)
else:
PT_ID_mask = (Photon.pt >= 20) & (Photon.cutBased > 1)
# dR cut with selected Muon and Electrons
dr_pho_ele_mask = ak.all(Photon.metric_table(Electron) >= 0.5,
axis=-1) # default metric table: delta_r
dr_pho_mu_mask = ak.all(Photon.metric_table(Muon) >= 0.5, axis=-1)
# genPartFlav cut
"""
if dataset == "WZG":
isPrompt = (Photon.genPartFlav == 1) | (Photon.genPartFlav == 11)
PhoSelmask = PT_ID_mask & isgap_mask & Pixel_seed_mask & isPrompt & dr_pho_ele_mask & dr_pho_mu_mask
elif dataset == "WZ":
isPrompt = (Photon.genPartFlav == 1)
PhoSelmask = PT_ID_mask & isgap_mask & Pixel_seed_mask & ~isPrompt & dr_pho_ele_mask & dr_pho_mu_mask
else:
PhoSelmask = PT_ID_mask & isgap_mask & Pixel_seed_mask & dr_pho_ele_mask & dr_pho_mu_mask
"""
# Add genPartFlav to remove Fake Photon in MC samples ( They are already considered by data driven method )
if not isData:
genPartFlav_mask = (Photon.genPartFlav == 1)
PhoSelmask = (genPartFlav_mask & PT_ID_mask & isgap_mask
& Pixel_seed_mask & dr_pho_ele_mask & dr_pho_mu_mask)
else:
PhoSelmask = (PT_ID_mask & isgap_mask & Pixel_seed_mask
& dr_pho_ele_mask & dr_pho_mu_mask)
Photon = Photon[PhoSelmask]
# Apply cut 4
A_photon_mask = ak.num(Photon) > 0
Electron = Electron[A_photon_mask]
Photon = Photon[A_photon_mask]
Jet = Jet[A_photon_mask]
Muon = Muon[A_photon_mask]
MET = MET[A_photon_mask]
if not isData:
pu = pu[A_photon_mask]
events = events[A_photon_mask]
# Stop processing if there is no event remain
if len(Electron) == 0:
return out
def make_leading_pair(target, base):
return target[ak.argmax(base.pt, axis=1, keepdims=True)]
leading_pho = make_leading_pair(Photon, Photon)
# -------------------- Make Fake Photon BKGs---------------------------#
def make_bins(pt, eta, bin_range_str):
bin_dict = {
"PT_1_eta_1": (pt > 20) & (pt < 30) & (eta < 1),
"PT_1_eta_2": (pt > 20) & (pt < 30) & (eta > 1) & (eta < 1.5),
"PT_1_eta_3": (pt > 20) & (pt < 30) & (eta > 1.5) & (eta < 2),
"PT_1_eta_4": (pt > 20) & (pt < 30) & (eta > 2) & (eta < 2.5),
"PT_2_eta_1": (pt > 30) & (pt < 40) & (eta < 1),
"PT_2_eta_2": (pt > 30) & (pt < 40) & (eta > 1) & (eta < 1.5),
"PT_2_eta_3": (pt > 30) & (pt < 40) & (eta > 1.5) & (eta < 2),
"PT_2_eta_4": (pt > 30) & (pt < 40) & (eta > 2) & (eta < 2.5),
"PT_3_eta_1": (pt > 40) & (pt < 50) & (eta < 1),
"PT_3_eta_2": (pt > 40) & (pt < 50) & (eta > 1) & (eta < 1.5),
"PT_3_eta_3": (pt > 40) & (pt < 50) & (eta > 1.5) & (eta < 2),
"PT_3_eta_4": (pt > 40) & (pt < 50) & (eta > 2) & (eta < 2.5),
"PT_4_eta_1": (pt > 50) & (eta < 1),
"PT_4_eta_2": (pt > 50) & (eta > 1) & (eta < 1.5),
"PT_4_eta_3": (pt > 50) & (eta > 1.5) & (eta < 2),
"PT_4_eta_4": (pt > 50) & (eta > 2) & (eta < 2.5),
}
binmask = bin_dict[bin_range_str]
return binmask
bin_name_list = [
"PT_1_eta_1",
"PT_1_eta_2",
"PT_1_eta_3",
"PT_1_eta_4",
"PT_2_eta_1",
"PT_2_eta_2",
"PT_2_eta_3",
"PT_2_eta_4",
"PT_3_eta_1",
"PT_3_eta_2",
"PT_3_eta_3",
"PT_3_eta_4",
"PT_4_eta_1",
"PT_4_eta_2",
"PT_4_eta_3",
"PT_4_eta_4",
]
## -- Fake-fraction Lookup table --##
if isFake:
# Make Bin-range mask
binned_pteta_mask = {}
for name in bin_name_list:
binned_pteta_mask[name] = make_bins(
ak.flatten(leading_pho.pt),
ak.flatten(abs(leading_pho.eta)),
name,
)
# Read Fake fraction --> Mapping bin name to int()
if self._year == "2018":
in_dict = np.load("Fitting_2018/Fit_results.npy",
allow_pickle="True")[()]
if self._year == "2017":
in_dict = np.load("Fitting_2017/Fit_results.npy",
allow_pickle="True")[()]
idx = 0
fake_dict = {}
for i, j in in_dict.items():
fake_dict[idx] = j
idx += 1
# Reconstruct Fake_weight
fw = 0
for i, j in binned_pteta_mask.items():
fw = fw + j * fake_dict[bin_name_list.index(i)]
# Process 0 weight to 1
@numba.njit
def zero_one(x):
if x == 0:
x = 1
return x
vec_zero_one = np.vectorize(zero_one)
fw = vec_zero_one(fw)
else:
fw = np.ones(len(events))
cut4 = np.ones(len(Photon)) * 4
print("Fake fraction weight: ", len(fw), len(cut4), fw)
##----------- Cut flow5: OSSF
ossf_mask = Electron.charge[:, 0] + Electron.charge[:, 1] == 0
# Apply cut 5
Electron = Electron[ossf_mask]
Photon = Photon[ossf_mask]
fw = fw[ossf_mask]
Jet = Jet[ossf_mask]
MET = MET[ossf_mask]
Muon = Muon[ossf_mask]
if not isData:
pu = pu[ossf_mask]
events = events[ossf_mask]
# Stop processing if there is no event remain
if len(Electron) == 0:
return out
cut5 = np.ones(ak.sum(ak.num(Electron) > 0)) * 5
# Define Electron Triplet
Diele = ak.zip({
"lep1": Electron[:, 0],
"lep2": Electron[:, 1],
"p4": TLorentz_vector(Electron[:, 0] + Electron[:, 1]),
})
leading_ele = Diele.lep1
subleading_ele = Diele.lep2
def make_leading_pair(target, base):
return target[ak.argmax(base.pt, axis=1, keepdims=True)]
leading_pho = make_leading_pair(Photon, Photon)
# -- Scale Factor for each electron
# Trigger weight helper function
def Trigger_Weight(eta1, pt1, eta2, pt2):
per_ev_MC = (get_ele_trig_leg1_mc_Eff(eta1, pt1) *
get_ele_trig_leg2_mc_Eff(eta2, pt2) +
get_ele_trig_leg1_mc_Eff(eta2, pt2) *
get_ele_trig_leg2_mc_Eff(eta1, pt1) -
get_ele_trig_leg1_mc_Eff(eta1, pt1) *
get_ele_trig_leg1_mc_Eff(eta2, pt2))
per_ev_data = (
get_ele_trig_leg1_data_Eff(eta1, pt1) * get_ele_trig_leg1_SF(
eta1, pt1) * get_ele_trig_leg2_data_Eff(eta2, pt2) *
get_ele_trig_leg2_SF(eta2, pt2) +
get_ele_trig_leg1_data_Eff(eta2, pt2) * get_ele_trig_leg1_SF(
eta2, pt2) * get_ele_trig_leg2_data_Eff(eta1, pt1) *
get_ele_trig_leg2_SF(eta1, pt1) -
get_ele_trig_leg1_data_Eff(eta1, pt1) * get_ele_trig_leg1_SF(
eta1, pt1) * get_ele_trig_leg1_data_Eff(eta2, pt2) *
get_ele_trig_leg1_SF(eta2, pt2))
return per_ev_data / per_ev_MC
if not isData:
## -------------< Egamma ID and Reco Scale factor > -----------------##
get_pho_medium_id_sf = get_pho_medium_id_sf(
ak.flatten(leading_pho.eta), ak.flatten(leading_pho.pt))
ele_reco_sf = get_ele_reco_above20_sf(
leading_ele.deltaEtaSC + leading_ele.eta,
leading_ele.pt,
) * get_ele_reco_above20_sf(
subleading_ele.deltaEtaSC + subleading_ele.eta,
subleading_ele.pt,
)
ele_medium_id_sf = get_ele_medium_id_sf(
leading_ele.deltaEtaSC + leading_ele.eta,
leading_ele.pt,
) * get_ele_medium_id_sf(
subleading_ele.deltaEtaSC + subleading_ele.eta,
subleading_ele.pt,
)
## -------------< Muon ID and Iso Scale factor > -----------------##
get_mu_tight_id_sf = get_mu_tight_id_sf(ak.flatten(abs(Muon.eta)),
ak.flatten(Muon.pt))
get_mu_tight_iso_sf = get_mu_tight_iso_sf(
ak.flatten(abs(Muon.eta)), ak.flatten(Muon.pt))
## -------------< Double Electron Trigger Scale factor > -----------------##
eta1 = leading_ele.deltaEtaSC + leading_ele.eta
eta2 = subleading_ele.deltaEtaSC + subleading_ele.eta
pt1 = leading_ele.pt
pt2 = subleading_ele.pt
# -- 2017,2016 are not applied yet
if self._year == "2018":
ele_trig_weight = Trigger_Weight(eta1, pt1, eta2, pt2)
##----------- Cut flow6: Baseline selection
# Mee cut
Mee_cut_mask = Diele.p4.mass > 4
# Lepton PT cuts
Leppt_mask = ak.firsts((Diele.lep1.pt >= 25) & (Diele.lep2.pt >= 20)
& (Muon.pt >= 25))
# MET cuts
MET_mask = MET.pt > 20 # Baseline
# Assemble!!
Baseline_mask = Leppt_mask & MET_mask & Mee_cut_mask # SR,CR
# Apply cut6
Diele_base = Diele[Baseline_mask]
leading_pho_base = leading_pho[Baseline_mask]
Jet_base = Jet[Baseline_mask]
MET_base = MET[Baseline_mask]
Muon_base = Muon[Baseline_mask]
events_base = events[Baseline_mask]
# Photon EE and EB
isEE_mask = leading_pho.isScEtaEE
isEB_mask = leading_pho.isScEtaEB
Pho_EE_base = leading_pho[isEE_mask & Baseline_mask]
Pho_EB_base = leading_pho[isEB_mask & Baseline_mask]
# Stop processing if there is no event remain
if len(leading_pho_base) == 0:
return out
cut6 = np.ones(ak.sum(ak.num(leading_pho_base) > 0)) * 6
base_arr_dict = {
"Diele_sel": Diele_base,
"leading_pho_sel": leading_pho_base,
"Jet_sel": Jet_base,
"MET_sel": MET_base,
"Muon_sel": Muon_base,
"Pho_EE_sel": Pho_EE_base,
"Pho_EB_sel": Pho_EB_base,
}
##----------- << SR >>
Zmass_window_mask = abs(Diele.p4.mass - 91.1876) < 15
MET_mask = MET.pt > 30
bjet_veto = ak.firsts(Jet.btagDeepB > 0.7665) == 0
Mlll_mask = ((Diele.p4 + Muon[:, 0]).mass) > 100
SR_mask = Zmass_window_mask & MET_mask & bjet_veto & Mlll_mask
SR_mask = Baseline_mask & SR_mask
Diele_SR = Diele[SR_mask]
leading_pho_SR = leading_pho[SR_mask]
Muon_SR = Muon[SR_mask]
MET_SR = MET[SR_mask]
Jet_SR = Jet[SR_mask]
events_SR = events[SR_mask]
Pho_EE_SR = leading_pho[isEE_mask & SR_mask]
Pho_EB_SR = leading_pho[isEB_mask & SR_mask]
SR_arr_dict = {
"Diele_sel": Diele_SR,
"leading_pho_sel": leading_pho_SR,
"Jet_sel": Jet_SR,
"MET_sel": MET_SR,
"Muon_sel": Muon_SR,
"Pho_EE_sel": Pho_EE_SR,
"Pho_EB_sel": Pho_EB_SR,
}
##----------- << CR-Z+Jets >>
Zmass_window_mask = abs(Diele.p4.mass - 91.1876) < 15
MET_mask = MET.pt <= 30
bjet_veto = ak.firsts(Jet.btagDeepB > 0.7665) == 0
Mlll_mask = ((Diele.p4 + Muon[:, 0]).mass) > 100
CR_ZJets_mask = Zmass_window_mask & MET_mask & bjet_veto & Mlll_mask
CR_ZJets_mask = Baseline_mask & CR_ZJets_mask
Diele_CR_ZJets = Diele[CR_ZJets_mask]
leading_pho_CR_ZJets = leading_pho[CR_ZJets_mask]
Muon_CR_ZJets = Muon[CR_ZJets_mask]
MET_CR_ZJets = MET[CR_ZJets_mask]
Jet_CR_ZJets = Jet[CR_ZJets_mask]
events_CR_ZJets = events[CR_ZJets_mask]
Pho_EE_CR_ZJets = leading_pho[isEE_mask & CR_ZJets_mask]
Pho_EB_CR_ZJets = leading_pho[isEB_mask & CR_ZJets_mask]
CR_ZJets_arr_dict = {
"Diele_sel": Diele_CR_ZJets,
"leading_pho_sel": leading_pho_CR_ZJets,
"Jet_sel": Jet_CR_ZJets,
"MET_sel": MET_CR_ZJets,
"Muon_sel": Muon_CR_ZJets,
"Pho_EE_sel": Pho_EE_CR_ZJets,
"Pho_EB_sel": Pho_EB_CR_ZJets,
}
##----------- << CR-T-enriched >>
Zmass_window_mask = abs(Diele.p4.mass - 91.1876) > 5
MET_mask = MET.pt > 30
bjet_veto = ak.firsts(Jet.btagDeepB > 0.7665) > 0
Mlll_mask = ((Diele.p4 + Muon[:, 0]).mass) > 100
CR_Tenri_mask = Zmass_window_mask & MET_mask & bjet_veto & Mlll_mask
CR_Tenri_mask = Baseline_mask & CR_Tenri_mask
Diele_CR_t = Diele[CR_Tenri_mask]
leading_pho_CR_t = leading_pho[CR_Tenri_mask]
Muon_CR_t = Muon[CR_Tenri_mask]
MET_CR_t = MET[CR_Tenri_mask]
Jet_CR_t = Jet[CR_Tenri_mask]
events_CR_t = events[CR_Tenri_mask]
Pho_EE_CR_t = leading_pho[isEE_mask & CR_Tenri_mask]
Pho_EB_CR_t = leading_pho[isEB_mask & CR_Tenri_mask]
CR_tEnriched_arr_dict = {
"Diele_sel": Diele_CR_t,
"leading_pho_sel": leading_pho_CR_t,
"Jet_sel": Jet_CR_t,
"MET_sel": MET_CR_t,
"Muon_sel": Muon_CR_t,
"Pho_EE_sel": Pho_EE_CR_t,
"Pho_EB_sel": Pho_EB_CR_t,
}
##----------- << CR-Conversion >>
Zmass_window_mask = abs(Diele.p4.mass - 91.1876) > 15
MET_mask = MET.pt <= 30
bjet_veto = ak.firsts(Jet.btagDeepB > 0.7665) == 0
Mlll_mask = ((Diele.p4 + Muon[:, 0]).mass) <= 100
CR_conv_mask = Zmass_window_mask & MET_mask & bjet_veto & Mlll_mask
CR_conv_mask = Baseline_mask & CR_conv_mask
Diele_CR_conv = Diele[CR_conv_mask]
leading_pho_CR_conv = leading_pho[CR_conv_mask]
Muon_CR_conv = Muon[CR_conv_mask]
MET_CR_conv = MET[CR_conv_mask]
Jet_CR_conv = Jet[CR_conv_mask]
events_CR_conv = events[CR_conv_mask]
Pho_EE_CR_conv = leading_pho[isEE_mask & CR_conv_mask]
Pho_EB_CR_conv = leading_pho[isEB_mask & CR_conv_mask]
CR_Conversion_dict = {
"Diele_sel": Diele_CR_conv,
"leading_pho_sel": leading_pho_CR_conv,
"Jet_sel": Jet_CR_conv,
"MET_sel": MET_CR_conv,
"Muon_sel": Muon_CR_conv,
"Pho_EE_sel": Pho_EE_CR_conv,
"Pho_EB_sel": Pho_EB_CR_conv,
}
## -------------------- Prepare making hist --------------#
regions = {
"Baseline": base_arr_dict,
"Signal": SR_arr_dict,
"CR_ZJets": CR_ZJets_arr_dict,
"CR_tEnriched": CR_tEnriched_arr_dict,
"CR_conversion": CR_Conversion_dict,
}
mask_dict = {
"Baseline": Baseline_mask,
"Signal": SR_mask,
"CR_ZJets": CR_ZJets_mask,
"CR_tEnriched": CR_Tenri_mask,
"CR_conversion": CR_conv_mask,
}
for region, arr_dict in regions.items():
# Photon
phoPT = ak.flatten(arr_dict["leading_pho_sel"].pt)
phoEta = ak.flatten(arr_dict["leading_pho_sel"].eta)
phoPhi = ak.flatten(arr_dict["leading_pho_sel"].phi)
# Photon EE
if len(arr_dict["Pho_EE_sel"].pt) != 0:
Pho_EE_PT = ak.flatten(arr_dict["Pho_EE_sel"].pt)
Pho_EE_Eta = ak.flatten(arr_dict["Pho_EE_sel"].eta)
Pho_EE_Phi = ak.flatten(arr_dict["Pho_EE_sel"].phi)
Pho_EE_sieie = ak.flatten(arr_dict["Pho_EE_sel"].sieie)
Pho_EE_Iso_charge = ak.flatten(
arr_dict["Pho_EE_sel"].pfRelIso03_chg)
# Photon EB
if len(arr_dict["Pho_EB_sel"].pt) != 0:
Pho_EB_PT = ak.flatten(arr_dict["Pho_EB_sel"].pt)
Pho_EB_Eta = ak.flatten(arr_dict["Pho_EB_sel"].eta)
Pho_EB_Phi = ak.flatten(arr_dict["Pho_EB_sel"].phi)
Pho_EB_sieie = ak.flatten(arr_dict["Pho_EB_sel"].sieie)
Pho_EB_Iso_charge = ak.flatten(
arr_dict["Pho_EB_sel"].pfRelIso03_chg)
# Electrons
ele1PT = arr_dict["Diele_sel"].lep1.pt
ele1Eta = arr_dict["Diele_sel"].lep1.eta
ele1Phi = arr_dict["Diele_sel"].lep1.phi
ele2PT = arr_dict["Diele_sel"].lep2.pt
ele2Eta = arr_dict["Diele_sel"].lep2.eta
ele2Phi = arr_dict["Diele_sel"].lep2.phi
# Muon
muPT = ak.flatten(arr_dict["Muon_sel"].pt)
muEta = ak.flatten(arr_dict["Muon_sel"].eta)
muPhi = ak.flatten(arr_dict["Muon_sel"].phi)
# MET
met = ak.to_numpy(arr_dict["MET_sel"].pt)
# M(eea) M(ee)
diele = arr_dict["Diele_sel"].p4
lll_vec = diele + arr_dict["Muon_sel"][:, 0]
Mlll = lll_vec.mass
Mee = diele.mass
# W MT (--> beta)
MT = np.sqrt(
2 * arr_dict["Muon_sel"].pt * arr_dict["MET_sel"].pt *
(1 - np.cos(
abs(arr_dict["MET_sel"].delta_phi(arr_dict["Muon_sel"])))))
MT = np.array(ak.firsts(MT))
# --- Apply weight and hist
weights = processor.Weights(len(cut5))
# --- skim cut-weight
def skim_weight(arr):
mask1 = ~ak.is_none(arr)
subarr = arr[mask1]
mask2 = subarr != 0
return ak.to_numpy(subarr[mask2])
cuts = mask_dict[region]
cuts_pho_EE = ak.flatten(isEE_mask)
cuts_pho_EB = ak.flatten(isEB_mask)
if isFake:
weights.add("fake_fraction", fw)
# Weight and SF here
if not (isData | isFake):
weights.add("pileup", pu)
weights.add("ele_id", ele_medium_id_sf)
weights.add("ele_reco", ele_reco_sf)
weights.add("pho_id", get_pho_medium_id_sf)
weights.add("mu_id", get_mu_tight_id_sf)
weights.add("mu_iso", get_mu_tight_id_sf)
# 2016,2017 are not applied yet
if self._year == "2018":
weights.add("ele_trigger", ele_trig_weight)
# ---------------------------- Fill hist --------------------------------------#
# Initial events
out["sumw"][dataset] += len(Initial_events)
print(
"region: {0} ### cut0: {1},cut1: {2}, cut2: {3},cut3: {4},cut4: {5},cut5: {6},cut6: {7}, cut7: {8}"
.format(region, len(cut0), len(cut1), len(cut2), len(cut3),
len(cut4), len(cut5), len(cut6), len(met)))
# Fill hist
# -- met -- #
out["met"].fill(
dataset=dataset,
region=region,
met=met,
weight=skim_weight(weights.weight() * cuts),
)
# --mass -- #
out["MT"].fill(
dataset=dataset,
region=region,
MT=MT,
weight=skim_weight(weights.weight() * cuts),
)
out["mass"].fill(
dataset=dataset,
region=region,
mass=Mee,
weight=skim_weight(weights.weight() * cuts),
)
out["mass_lll"].fill(
dataset=dataset,
region=region,
mass_lll=Mlll,
weight=skim_weight(weights.weight() * cuts),
)
# -- Muon -- #
out["mupt"].fill(
dataset=dataset,
region=region,
mupt=muPT,
weight=skim_weight(weights.weight() * cuts),
)
out["mueta"].fill(
dataset=dataset,
region=region,
mueta=muEta,
weight=skim_weight(weights.weight() * cuts),
)
out["muphi"].fill(
dataset=dataset,
region=region,
muphi=muPhi,
weight=skim_weight(weights.weight() * cuts),
)
# -- Electron -- #
out["ele1pt"].fill(
dataset=dataset,
region=region,
ele1pt=ele1PT,
weight=skim_weight(weights.weight() * cuts),
)
out["ele1eta"].fill(
dataset=dataset,
region=region,
ele1eta=ele1Eta,
weight=skim_weight(weights.weight() * cuts),
)
out["ele1phi"].fill(
dataset=dataset,
region=region,
ele1phi=ele1Phi,
weight=skim_weight(weights.weight() * cuts),
)
out["ele2pt"].fill(
dataset=dataset,
region=region,
ele2pt=ele2PT,
weight=skim_weight(weights.weight() * cuts),
)
out["ele2eta"].fill(
dataset=dataset,
region=region,
ele2eta=ele2Eta,
weight=skim_weight(weights.weight() * cuts),
)
out["ele2phi"].fill(
dataset=dataset,
region=region,
ele2phi=ele2Phi,
weight=skim_weight(weights.weight() * cuts),
)
# -- Photon -- #
out["phopt"].fill(
dataset=dataset,
region=region,
phopt=phoPT,
weight=skim_weight(weights.weight() * cuts),
)
out["phoeta"].fill(
dataset=dataset,
region=region,
phoeta=phoEta,
weight=skim_weight(weights.weight() * cuts),
)
out["phophi"].fill(
dataset=dataset,
region=region,
phophi=phoPhi,
weight=skim_weight(weights.weight() * cuts),
)
if len(arr_dict["Pho_EE_sel"].pt) != 0:
out["pho_EE_pt"].fill(
dataset=dataset,
region=region,
pho_EE_pt=Pho_EE_PT,
weight=skim_weight(weights.weight() * cuts * cuts_pho_EE),
)
out["pho_EE_eta"].fill(
dataset=dataset,
region=region,
pho_EE_eta=Pho_EE_Eta,
weight=skim_weight(weights.weight() * cuts * cuts_pho_EE),
)
out["pho_EE_phi"].fill(
dataset=dataset,
region=region,
pho_EE_phi=Pho_EE_Phi,
weight=skim_weight(weights.weight() * cuts * cuts_pho_EE),
)
out["pho_EE_sieie"].fill(
dataset=dataset,
region=region,
pho_EE_sieie=Pho_EE_sieie,
weight=skim_weight(weights.weight() * cuts * cuts_pho_EE),
)
out["pho_EE_Iso_chg"].fill(
dataset=dataset,
region=region,
pho_EE_Iso_chg=Pho_EE_Iso_charge,
weight=skim_weight(weights.weight() * cuts * cuts_pho_EE),
)
if len(arr_dict["Pho_EB_sel"].pt) != 0:
out["pho_EB_pt"].fill(
dataset=dataset,
region=region,
pho_EB_pt=Pho_EB_PT,
weight=skim_weight(weights.weight() * cuts * cuts_pho_EB),
)
out["pho_EB_eta"].fill(
dataset=dataset,
region=region,
pho_EB_eta=Pho_EB_Eta,
weight=skim_weight(weights.weight() * cuts * cuts_pho_EB),
)
out["pho_EB_phi"].fill(
dataset=dataset,
region=region,
pho_EB_phi=Pho_EB_Phi,
weight=skim_weight(weights.weight() * cuts * cuts_pho_EB),
)
out["pho_EB_sieie"].fill(
dataset=dataset,
region=region,
pho_EB_sieie=Pho_EB_sieie,
weight=skim_weight(weights.weight() * cuts * cuts_pho_EB),
)
out["pho_EB_Iso_chg"].fill(
dataset=dataset,
region=region,
pho_EB_Iso_chg=Pho_EB_Iso_charge,
weight=skim_weight(weights.weight() * cuts * cuts_pho_EB),
)
return out
def process(self, events):
# Initialize accumulator
out = self.accumulator.identity()
dataset = sample_name
# events.metadata['dataset']
# Data or MC
isData = "genWeight" not in events.fields
isFake = self._isFake
# Stop processing if there is no event remain
if len(events) == 0:
return out
# Golden Json file
if (self._year == "2018") and isData:
injson = "/x5/cms/jwkim/gitdir/JWCorp/JW_analysis/Coffea_WZG/Corrections/Cert_314472-325175_13TeV_Legacy2018_Collisions18_JSON.txt.RunABD"
if (self._year == "2017") and isData:
injson = "/x5/cms/jwkim/gitdir/JWCorp/JW_analysis/Coffea_WZG/Corrections/Cert_294927-306462_13TeV_UL2017_Collisions17_GoldenJSON.txt"
# <----- Get Scale factors ------>#
if not isData:
# Egamma reco ID
get_ele_reco_above20_sf = self._corrections["get_ele_reco_above20_sf"][
self._year
]
get_ele_medium_id_sf = self._corrections["get_ele_medium_id_sf"][self._year]
get_pho_medium_id_sf = self._corrections["get_pho_medium_id_sf"][self._year]
# DoubleEG trigger # 2016, 2017 are not applied yet
if self._year == "2018":
get_ele_trig_leg1_SF = self._corrections["get_ele_trig_leg1_SF"][
self._year
]
get_ele_trig_leg1_data_Eff = self._corrections[
"get_ele_trig_leg1_data_Eff"
][self._year]
get_ele_trig_leg1_mc_Eff = self._corrections[
"get_ele_trig_leg1_mc_Eff"
][self._year]
get_ele_trig_leg2_SF = self._corrections["get_ele_trig_leg2_SF"][
self._year
]
get_ele_trig_leg2_data_Eff = self._corrections[
"get_ele_trig_leg2_data_Eff"
][self._year]
get_ele_trig_leg2_mc_Eff = self._corrections[
"get_ele_trig_leg2_mc_Eff"
][self._year]
# PU weight with custom made npy and multi-indexing
pu_weight_idx = ak.values_astype(events.Pileup.nTrueInt, "int64")
pu = self._puweight_arr[pu_weight_idx]
print("## pu_idx: ",len(pu_weight_idx),pu_weight_idx)
print("## pu_arr: ",len(self._puweight_arr),self._puweight_arr)
print("## pu:",len(pu),pu)
selection = processor.PackedSelection()
# Cut flow
cut0 = np.zeros(len(events))
out["cutflow"].fill(dataset=dataset, cutflow=cut0)
# <----- Helper functions ------>#
# Sort by PT helper function
def sort_by_pt(ele, pho, jet):
ele = ele[ak.argsort(ele.pt, ascending=False, axis=1)]
pho = pho[ak.argsort(pho.pt, ascending=False, axis=1)]
jet = jet[ak.argsort(jet.pt, ascending=False, axis=1)]
return ele, pho, jet
# Lorentz vectors
from coffea.nanoevents.methods import vector
ak.behavior.update(vector.behavior)
def TLorentz_vector(vec):
vec = ak.zip(
{"x": vec.x, "y": vec.y, "z": vec.z, "t": vec.t},
with_name="LorentzVector",
)
return vec
def TLorentz_vector_cylinder(vec):
vec = ak.zip(
{
"pt": vec.pt,
"eta": vec.eta,
"phi": vec.phi,
"mass": vec.mass,
},
with_name="PtEtaPhiMLorentzVector",
)
return vec
# <----- Selection ------>#
Initial_events = events
# Good Run ( Golden Json files )
from coffea import lumi_tools
if isData:
lumi_mask_builder = lumi_tools.LumiMask(injson)
lumimask = ak.Array(
lumi_mask_builder.__call__(events.run, events.luminosityBlock)
)
events = events[lumimask]
# print("{0}% of files pass good-run conditions".format(len(events)/ len(Initial_events)))
# Stop processing if there is no event remain
if len(events) == 0:
return out
##----------- Cut flow1: Passing Triggers
# double lepton trigger
is_double_ele_trigger = True
if not is_double_ele_trigger:
double_ele_triggers_arr = np.ones(len(events), dtype=np.bool)
else:
double_ele_triggers_arr = np.zeros(len(events), dtype=np.bool)
for path in self._doubleelectron_triggers[self._year]:
if path not in events.HLT.fields:
continue
double_ele_triggers_arr = double_ele_triggers_arr | events.HLT[path]
# single lepton trigger
is_single_ele_trigger = True
if not is_single_ele_trigger:
single_ele_triggers_arr = np.ones(len(events), dtype=np.bool)
else:
single_ele_triggers_arr = np.zeros(len(events), dtype=np.bool)
for path in self._singleelectron_triggers[self._year]:
if path not in events.HLT.fields:
continue
single_ele_triggers_arr = single_ele_triggers_arr | events.HLT[path]
events.Electron, events.Photon, events.Jet = sort_by_pt(
events.Electron, events.Photon, events.Jet
)
# Good Primary vertex
nPV = events.PV.npvsGood
nPV_nw = events.PV.npvsGood
if not isData:
nPV = nPV * pu
print(pu)
# Apply cut1
events = events[double_ele_triggers_arr]
if not isData:
pu = pu[double_ele_triggers_arr]
# Stop processing if there is no event remain
if len(events) == 0:
return out
cut1 = np.ones(len(events))
out["cutflow"].fill(dataset=dataset, cutflow=cut1)
# Set Particles
Electron = events.Electron
Muon = events.Muon
Photon = events.Photon
MET = events.MET
Jet = events.Jet
# --Muon ( only used to calculate dR )
MuSelmask = (
(Muon.pt >= 10)
& (abs(Muon.eta) <= 2.5)
& (Muon.tightId)
& (Muon.pfRelIso04_all < 0.15)
)
Muon = Muon[MuSelmask]
# --Loose Muon ( For Loose Muon veto )
LoooseMuSelmask = (
(Muon.pt > 20)
& (abs(Muon.eta) < 2.4)
& (Muon.isPFcand)
& (Muon.isGlobal | Muon.isTracker)
& (Muon.pfRelIso03_all < 0.25)
)
# Reference: VBS Zgamma+2jets
VetoMuon = Muon[LoooseMuSelmask]
##----------- Cut flow2: Electron Selection
EleSelmask = (
(Electron.pt >= 10)
& (np.abs(Electron.eta + Electron.deltaEtaSC) < 1.479)
& (Electron.cutBased > 2)
& (abs(Electron.dxy) < 0.05)
& (abs(Electron.dz) < 0.1)
) | (
(Electron.pt >= 10)
& (np.abs(Electron.eta + Electron.deltaEtaSC) > 1.479)
& (np.abs(Electron.eta + Electron.deltaEtaSC) <= 2.5)
& (Electron.cutBased > 2)
& (abs(Electron.dxy) < 0.1)
& (abs(Electron.dz) < 0.2)
)
Electron = Electron[EleSelmask]
# Event with 3 Electrons
# apply cut 2
Tri_electron_mask = ak.num(Electron) == 3
Electron = Electron[Tri_electron_mask]
Photon = Photon[Tri_electron_mask]
Jet = Jet[Tri_electron_mask]
MET = MET[Tri_electron_mask]
Muon = Muon[Tri_electron_mask]
VetoMuon = VetoMuon[Tri_electron_mask]
if not isData:
pu = pu[Tri_electron_mask]
events = events[Tri_electron_mask]
# Stop processing if there is no event remain
if len(Electron) == 0:
return out
cut2 = np.ones(len(Photon)) * 2
out["cutflow"].fill(dataset=dataset, cutflow=cut2)
##----------- Cut flow3: 4th lepton veto (Loose Muon)
# Veto 4th Loose muon
# apply cut 3
fourth_lepton_veto = ak.num(VetoMuon) < 1
Electron = Electron[fourth_lepton_veto]
Photon = Photon[fourth_lepton_veto]
Jet = Jet[fourth_lepton_veto]
MET = MET[fourth_lepton_veto]
Muon = Muon[fourth_lepton_veto]
if not isData:
pu = pu[fourth_lepton_veto]
events = events[fourth_lepton_veto]
# Stop processing if there is no event remain
if len(Electron) == 0:
return out
cut3 = np.ones(len(Photon)) * 3
out["cutflow"].fill(dataset=dataset, cutflow=cut3)
##----------- Cut flow4: Photon Selection
# Basic photon selection
isgap_mask = (abs(Photon.eta) < 1.442) | (
(abs(Photon.eta) > 1.566) & (abs(Photon.eta) < 2.5)
)
Pixel_seed_mask = ~Photon.pixelSeed
if (dataset == "ZZ") and (self._year == "2017"):
PT_ID_mask = (Photon.pt >= 20) & (
Photon.cutBasedBitmap >= 3
) # 2^0(Loose) + 2^1(Medium) + 2^2(Tights)
else:
PT_ID_mask = (Photon.pt >= 20) & (Photon.cutBased > 1)
# dR cut with selected Muon and Electrons
dr_pho_ele_mask = ak.all(
Photon.metric_table(Electron) >= 0.5, axis=-1
) # default metric table: delta_r
dr_pho_mu_mask = ak.all(Photon.metric_table(Muon) >= 0.5, axis=-1)
# genPartFlav cut
"""
if dataset == "WZG":
isPrompt = (Photon.genPartFlav == 1) | (Photon.genPartFlav == 11)
PhoSelmask = PT_ID_mask & isgap_mask & Pixel_seed_mask & isPrompt & dr_pho_ele_mask & dr_pho_mu_mask
elif dataset == "WZ":
isPrompt = (Photon.genPartFlav == 1)
PhoSelmask = PT_ID_mask & isgap_mask & Pixel_seed_mask & ~isPrompt & dr_pho_ele_mask & dr_pho_mu_mask
else:
PhoSelmask = PT_ID_mask & isgap_mask & Pixel_seed_mask & dr_pho_ele_mask & dr_pho_mu_mask
"""
PhoSelmask = (
PT_ID_mask & isgap_mask & Pixel_seed_mask & dr_pho_ele_mask & dr_pho_mu_mask
)
Photon = Photon[PhoSelmask]
# Apply cut 4
A_photon_mask = ak.num(Photon) > 0
Electron = Electron[A_photon_mask]
Photon = Photon[A_photon_mask]
Jet = Jet[A_photon_mask]
Muon = Muon[A_photon_mask]
MET = MET[A_photon_mask]
if not isData:
pu = pu[A_photon_mask]
events = events[A_photon_mask]
# Stop processing if there is no event remain
if len(Electron) == 0:
return out
cut4 = np.ones(len(Photon)) * 4
out["cutflow"].fill(dataset=dataset, cutflow=cut4)
##----------- Cut flow5: OSSF
# OSSF index maker
@numba.njit
def find_3lep(events_leptons, builder):
for leptons in events_leptons:
builder.begin_list()
nlep = len(leptons)
for i0 in range(nlep):
for i1 in range(i0 + 1, nlep):
if leptons[i0].charge + leptons[i1].charge != 0:
continue
for i2 in range(nlep):
if len({i0, i1, i2}) < 3:
continue
builder.begin_tuple(3)
builder.index(0).integer(i0)
builder.index(1).integer(i1)
builder.index(2).integer(i2)
builder.end_tuple()
builder.end_list()
return builder
eee_triplet_idx = find_3lep(Electron, ak.ArrayBuilder()).snapshot()
ossf_mask = ak.num(eee_triplet_idx) == 2
# Apply cut 5
eee_triplet_idx = eee_triplet_idx[ossf_mask]
Electron = Electron[ossf_mask]
Photon = Photon[ossf_mask]
Jet = Jet[ossf_mask]
MET = MET[ossf_mask]
if not isData:
pu = pu[ossf_mask]
events = events[ossf_mask]
# Stop processing if there is no event remain
if len(Electron) == 0:
return out
cut5 = np.ones(ak.sum(ak.num(Electron) > 0)) * 5
out["cutflow"].fill(dataset=dataset, cutflow=cut5)
# Define Electron Triplet
Triple_electron = [Electron[eee_triplet_idx[idx]] for idx in "012"]
Triple_eee = ak.zip(
{
"lep1": Triple_electron[0],
"lep2": Triple_electron[1],
"lep3": Triple_electron[2],
"p4": TLorentz_vector(Triple_electron[0] + Triple_electron[1]),
}
)
# Ele pair selector --> Close to Z mass
bestZ_idx = ak.singletons(ak.argmin(abs(Triple_eee.p4.mass - 91.1876), axis=1))
Triple_eee = Triple_eee[bestZ_idx]
leading_ele = Triple_eee.lep1
subleading_ele = Triple_eee.lep2
third_ele = Triple_eee.lep3
def make_leading_pair(target, base):
return target[ak.argmax(base.pt, axis=1, keepdims=True)]
leading_pho = make_leading_pair(Photon, Photon)
# -- Scale Factor for each electron
# Trigger weight helper function
def Trigger_Weight(eta1, pt1, eta2, pt2):
per_ev_MC = (
get_ele_trig_leg1_mc_Eff(eta1, pt1)
* get_ele_trig_leg2_mc_Eff(eta2, pt2)
+ get_ele_trig_leg1_mc_Eff(eta2, pt2)
* get_ele_trig_leg2_mc_Eff(eta1, pt1)
- get_ele_trig_leg1_mc_Eff(eta1, pt1)
* get_ele_trig_leg1_mc_Eff(eta2, pt2)
)
per_ev_data = (
get_ele_trig_leg1_data_Eff(eta1, pt1)
* get_ele_trig_leg1_SF(eta1, pt1)
* get_ele_trig_leg2_data_Eff(eta2, pt2)
* get_ele_trig_leg2_SF(eta2, pt2)
+ get_ele_trig_leg1_data_Eff(eta2, pt2)
* get_ele_trig_leg1_SF(eta2, pt2)
* get_ele_trig_leg2_data_Eff(eta1, pt1)
* get_ele_trig_leg2_SF(eta1, pt1)
- get_ele_trig_leg1_data_Eff(eta1, pt1)
* get_ele_trig_leg1_SF(eta1, pt1)
* get_ele_trig_leg1_data_Eff(eta2, pt2)
* get_ele_trig_leg1_SF(eta2, pt2)
)
return per_ev_data / per_ev_MC
if not isData:
## -------------< Egamma ID and Reco Scale factor > -----------------##
get_pho_medium_id_sf = get_pho_medium_id_sf(
ak.flatten(leading_pho.eta), ak.flatten(leading_pho.pt)
)
ele_reco_sf = (
get_ele_reco_above20_sf(
ak.flatten(leading_ele.deltaEtaSC + leading_ele.eta),
ak.flatten(leading_ele.pt),
)
* get_ele_reco_above20_sf(
ak.flatten(subleading_ele.deltaEtaSC + subleading_ele.eta),
ak.flatten(subleading_ele.pt),
)
* get_ele_reco_above20_sf(
ak.flatten(third_ele.deltaEtaSC + third_ele.eta),
ak.flatten(third_ele.pt),
)
)
ele_medium_id_sf = (
get_ele_medium_id_sf(
ak.flatten(leading_ele.deltaEtaSC + leading_ele.eta),
ak.flatten(leading_ele.pt),
)
* get_ele_medium_id_sf(
ak.flatten(subleading_ele.deltaEtaSC + subleading_ele.eta),
ak.flatten(subleading_ele.pt),
)
* get_ele_medium_id_sf(
ak.flatten(third_ele.deltaEtaSC + third_ele.eta),
ak.flatten(third_ele.pt),
)
)
## -------------< Double Electron Trigger Scale factor > -----------------##
eta1 = ak.flatten(leading_ele.deltaEtaSC + leading_ele.eta)
eta2 = ak.flatten(subleading_ele.deltaEtaSC + subleading_ele.eta)
pt1 = ak.flatten(leading_ele.pt)
pt2 = ak.flatten(subleading_ele.pt)
# -- 2017,2016 are not applied yet
if self._year == "2018":
ele_trig_weight = Trigger_Weight(eta1, pt1, eta2, pt2)
##----------- Cut flow6: Event selection
# Mee cut
diele = Triple_eee.p4
Mee_cut_mask = ak.firsts(diele.mass) > 4
# Z mass window
# zmass_window_mask = ak.firsts(abs(diele.mass - 91.1876)) < 15 # SR, CR_ZZA, CR_Z+jets, CR_Conversion
# zmass_window_mask = ak.firsts(abs(diele.mass - 91.1876)) > 5 # CR_t-enriched
# zmass_window_mask = ak.firsts(abs(diele.mass - 91.1876)) > 15 # CR_Conversion
# M(eee) cut SR, CR_ZZA, CR_Z+jets, CR_t enriched
# eee = Triple_eee.lep1 + Triple_eee.lep2 + Triple_eee.lep3
# Meee_cut_mask = ak.firsts(eee.mass > 100)
# Meee_cut_mask = ak.firsts(eee.mass <= 100)
# b-Jet veto cut #SR, CR_ZZA, CR_Z+jets, CR_Conversion
# bjet_mask = (Jet.btagCSVV2 > 0.4184) & (Jet.pt > 30)
# bjet_veto_mask = ak.num(Jet[bjet_mask]) == 0
# bjet_veto_mask = ak.num(Jet[bjet_mask]) > 0 # CR_t-enriched
# Electron PT cuts
Elept_mask = ak.firsts(
(leading_ele.pt >= 25) & (subleading_ele.pt >= 10) & (third_ele.pt >= 25)
)
# MET cuts
MET_mask = MET > 20 # Baseline
# MET_mask = MET.pt > 30 # SR, CR-ZZE, CR-t-entirched
# MET_mask = MET.pt <= 30 # CR-Z+jets. CR-Conversion
# Mask
# Event_sel_mask = Elept_mask & MET_mask & bjet_veto_mask & Mee_cut_mask & zmass_window_mask & Meee_cut_mask # SR,CR
Event_sel_mask = Elept_mask & MET_mask & Mee_cut_mask # SR,CR
# Apply cut6
Triple_eee_sel = Triple_eee[Event_sel_mask]
leading_pho_sel = leading_pho[Event_sel_mask]
MET_sel = MET[Event_sel_mask]
events = events[Event_sel_mask]
# Photon EE and EB
isEE_mask = leading_pho.isScEtaEE
isEB_mask = leading_pho.isScEtaEB
Pho_EE = leading_pho[isEE_mask & Event_sel_mask]
Pho_EB = leading_pho[isEB_mask & Event_sel_mask]
# Stop processing if there is no event remain
if len(leading_pho_sel) == 0:
return out
cut6 = np.ones(ak.sum(ak.num(leading_pho_sel) > 0)) * 6
out["cutflow"].fill(dataset=dataset, cutflow=cut6)
## -------------------- Prepare making hist --------------#
# Photon
phoPT = ak.flatten(leading_pho_sel.pt)
phoEta = ak.flatten(leading_pho_sel.eta)
phoPhi = ak.flatten(leading_pho_sel.phi)
# Photon EE
if len(Pho_EE.pt) != 0:
Pho_EE_PT = ak.flatten(Pho_EE.pt)
Pho_EE_Eta = ak.flatten(Pho_EE.eta)
Pho_EE_Phi = ak.flatten(Pho_EE.phi)
Pho_EE_sieie = ak.flatten(Pho_EE.sieie)
Pho_EE_hoe = ak.flatten(Pho_EE.hoe)
Pho_EE_Iso_charge = ak.flatten(Pho_EE.pfRelIso03_chg)
# Photon EB
if len(Pho_EB.pt) != 0:
Pho_EB_PT = ak.flatten(Pho_EB.pt)
Pho_EB_Eta = ak.flatten(Pho_EB.eta)
Pho_EB_Phi = ak.flatten(Pho_EB.phi)
Pho_EB_sieie = ak.flatten(Pho_EB.sieie)
Pho_EB_hoe = ak.flatten(Pho_EB.hoe)
Pho_EB_Iso_charge = ak.flatten(Pho_EB.pfRelIso03_chg)
# Electrons
ele1PT = ak.flatten(Triple_eee_sel.lep1.pt)
ele1Eta = ak.flatten(Triple_eee_sel.lep1.eta)
ele1Phi = ak.flatten(Triple_eee_sel.lep1.phi)
ele2PT = ak.flatten(Triple_eee_sel.lep2.pt)
ele2Eta = ak.flatten(Triple_eee_sel.lep2.eta)
ele2Phi = ak.flatten(Triple_eee_sel.lep2.phi)
ele3PT = ak.flatten(Triple_eee_sel.lep3.pt)
ele3Eta = ak.flatten(Triple_eee_sel.lep3.eta)
ele3Phi = ak.flatten(Triple_eee_sel.lep3.phi)
charge = ak.flatten(Triple_eee.lep1.charge + Triple_eee.lep2.charge)
# MET
met = ak.to_numpy(MET_sel)
# M(eea) M(ee)
diele = Triple_eee_sel.p4
eeg_vec = diele + leading_pho_sel
Meea = ak.flatten(eeg_vec.mass)
Mee = ak.flatten(Triple_eee_sel.p4.mass)
# --- Apply weight and hist
if isFake:
weights = processor.Weights(len(cut6))
else:
weights = processor.Weights(len(cut5))
# -------------------- Sieie bins---------------------------#
def make_bins(pt, eta, bin_range_str):
bin_dict = {
"PT_1_eta_1": (pt > 20) & (pt < 30) & (eta < 1),
"PT_1_eta_2": (pt > 20) & (pt < 30) & (eta > 1) & (eta < 1.5),
"PT_1_eta_3": (pt > 20) & (pt < 30) & (eta > 1.5) & (eta < 2),
"PT_1_eta_4": (pt > 20) & (pt < 30) & (eta > 2) & (eta < 2.5),
"PT_2_eta_1": (pt > 30) & (pt < 40) & (eta < 1),
"PT_2_eta_2": (pt > 30) & (pt < 40) & (eta > 1) & (eta < 1.5),
"PT_2_eta_3": (pt > 30) & (pt < 40) & (eta > 1.5) & (eta < 2),
"PT_2_eta_4": (pt > 30) & (pt < 40) & (eta > 2) & (eta < 2.5),
"PT_3_eta_1": (pt > 40) & (pt < 50) & (eta < 1),
"PT_3_eta_2": (pt > 40) & (pt < 50) & (eta > 1) & (eta < 1.5),
"PT_3_eta_3": (pt > 40) & (pt < 50) & (eta > 1.5) & (eta < 2),
"PT_3_eta_4": (pt > 40) & (pt < 50) & (eta > 2) & (eta < 2.5),
"PT_4_eta_1": (pt > 50) & (eta < 1),
"PT_4_eta_2": (pt > 50) & (eta > 1) & (eta < 1.5),
"PT_4_eta_3": (pt > 50) & (eta > 1.5) & (eta < 2),
"PT_4_eta_4": (pt > 50) & (eta > 2) & (eta < 2.5),
}
binmask = bin_dict[bin_range_str]
return binmask
bin_name_list = [
"PT_1_eta_1",
"PT_1_eta_2",
"PT_1_eta_3",
"PT_1_eta_4",
"PT_2_eta_1",
"PT_2_eta_2",
"PT_2_eta_3",
"PT_2_eta_4",
"PT_3_eta_1",
"PT_3_eta_2",
"PT_3_eta_3",
"PT_3_eta_4",
"PT_4_eta_1",
"PT_4_eta_2",
"PT_4_eta_3",
"PT_4_eta_4",
]
## -- Fake-fraction Lookup table --##
if isFake:
# Make Bin-range mask
binned_pteta_mask = {}
for name in bin_name_list:
binned_pteta_mask[name] = make_bins(
ak.flatten(leading_pho_sel.pt),
ak.flatten(abs(leading_pho_sel.eta)),
name,
)
# Read Fake fraction --> Mapping bin name to int()
in_dict = np.load('Fitting_v2/results_210517.npy',allow_pickle="True")[()]
idx=0
fake_dict ={}
for i,j in in_dict.items():
fake_dict[idx] = j
idx+=1
# Reconstruct Fake_weight
fw= 0
for i,j in binned_pteta_mask.items():
fw = fw + j*fake_dict[bin_name_list.index(i)]
# Process 0 weight to 1
@numba.njit
def zero_one(x):
if x == 0:
x = 1
return x
vec_zero_one = np.vectorize(zero_one)
fw = vec_zero_one(fw)
# --- skim cut-weight
if not isFake:
def skim_weight(arr):
mask1 = ~ak.is_none(arr)
subarr = arr[mask1]
mask2 = subarr != 0
return ak.to_numpy(subarr[mask2])
else:
def skim_weight(arr):
return arr
if not isFake:
cuts = Event_sel_mask
cuts_pho_EE = ak.flatten(isEE_mask)
cuts_pho_EB = ak.flatten(isEB_mask)
if isFake:
cuts = np.ones(len(Event_sel_mask))
cuts_pho_EE = ak.flatten(isEE_mask & Event_sel_mask)
cuts_pho_EB = ak.flatten(isEB_mask & Event_sel_mask)
if isFake:
weights.add("fake_fraction", fw)
# Weight and SF here
if not (isData | isFake):
weights.add("pileup", pu)
weights.add("ele_id", ele_medium_id_sf)
weights.add("pho_id", get_pho_medium_id_sf)
weights.add("ele_reco", ele_reco_sf)
# 2016,2017 are not applied yet
if self._year == "2018":
weights.add("ele_trigger", ele_trig_weight)
# ---------------------------- Fill hist --------------------------------------#
# Initial events
out["sumw"][dataset] += len(Initial_events)
print("cut1: {0},cut2: {1},cut3: {2},cut4: {3},cut5: {4},cut6: {5},cut7: {6}".format(len(cut0), len(cut1), len(cut2), len(cut3), len(cut4), len(cut5),len(cut6)))
## Cut flow loop
#for cut in [cut0, cut1, cut2, cut3, cut4, cut5,cut6]:
# out["cutflow"].fill(dataset=dataset, cutflow=cut)
# Primary vertex
out["nPV"].fill(
dataset=dataset,
nPV=nPV,
)
out["nPV_nw"].fill(dataset=dataset, nPV_nw=nPV_nw)
# Fill hist
# -- met -- #
out["met"].fill(
dataset=dataset, met=met, weight=skim_weight(weights.weight() * cuts)
)
# --mass -- #
out["mass"].fill(
dataset=dataset, mass=Mee, weight=skim_weight(weights.weight() * cuts)
)
out["mass_eea"].fill(
dataset=dataset, mass_eea=Meea, weight=skim_weight(weights.weight() * cuts)
)
# -- Electron -- #
out["ele1pt"].fill(
dataset=dataset, ele1pt=ele1PT, weight=skim_weight(weights.weight() * cuts)
)
out["ele1eta"].fill(
dataset=dataset,
ele1eta=ele1Eta,
weight=skim_weight(weights.weight() * cuts),
)
out["ele1phi"].fill(
dataset=dataset,
ele1phi=ele1Phi,
weight=skim_weight(weights.weight() * cuts),
)
out["ele2pt"].fill(
dataset=dataset, ele2pt=ele2PT, weight=skim_weight(weights.weight() * cuts)
)
out["ele2eta"].fill(
dataset=dataset,
ele2eta=ele2Eta,
weight=skim_weight(weights.weight() * cuts),
)
out["ele2phi"].fill(
dataset=dataset,
ele2phi=ele2Phi,
weight=skim_weight(weights.weight() * cuts),
)
out["ele3pt"].fill(
dataset=dataset, ele3pt=ele3PT, weight=skim_weight(weights.weight() * cuts)
)
# -- Photon -- #
out["phopt"].fill(
dataset=dataset, phopt=phoPT, weight=skim_weight(weights.weight() * cuts)
)
out["phoeta"].fill(
dataset=dataset, phoeta=phoEta, weight=skim_weight(weights.weight() * cuts)
)
out["phophi"].fill(
dataset=dataset, phophi=phoPhi, weight=skim_weight(weights.weight() * cuts)
)
if len(Pho_EE.pt) != 0:
out["pho_EE_pt"].fill(
dataset=dataset,
pho_EE_pt=Pho_EE_PT,
weight=skim_weight(weights.weight() * cuts * cuts_pho_EE),
)
out["pho_EE_eta"].fill(
dataset=dataset,
pho_EE_eta=Pho_EE_Eta,
weight=skim_weight(weights.weight() * cuts * cuts_pho_EE),
)
out["pho_EE_phi"].fill(
dataset=dataset,
pho_EE_phi=Pho_EE_Phi,
weight=skim_weight(weights.weight() * cuts * cuts_pho_EE),
)
out["pho_EE_hoe"].fill(
dataset=dataset,
pho_EE_hoe=Pho_EE_hoe,
weight=skim_weight(weights.weight() * cuts * cuts_pho_EE),
)
out["pho_EE_sieie"].fill(
dataset=dataset,
pho_EE_sieie=Pho_EE_sieie,
weight=skim_weight(weights.weight() * cuts * cuts_pho_EE),
)
out["pho_EE_Iso_chg"].fill(
dataset=dataset,
pho_EE_Iso_chg=Pho_EE_Iso_charge,
weight=skim_weight(weights.weight() * cuts * cuts_pho_EE),
)
if len(Pho_EB.pt) != 0:
out["pho_EB_pt"].fill(
dataset=dataset,
pho_EB_pt=Pho_EB_PT,
weight=skim_weight(weights.weight() * cuts * cuts_pho_EB),
)
out["pho_EB_eta"].fill(
dataset=dataset,
pho_EB_eta=Pho_EB_Eta,
weight=skim_weight(weights.weight() * cuts * cuts_pho_EB),
)
out["pho_EB_phi"].fill(
dataset=dataset,
pho_EB_phi=Pho_EB_Phi,
weight=skim_weight(weights.weight() * cuts * cuts_pho_EB),
)
out["pho_EB_hoe"].fill(
dataset=dataset,
pho_EB_hoe=Pho_EB_hoe,
weight=skim_weight(weights.weight() * cuts * cuts_pho_EB),
)
out["pho_EB_sieie"].fill(
dataset=dataset,
pho_EB_sieie=Pho_EB_sieie,
weight=skim_weight(weights.weight() * cuts * cuts_pho_EB),
)
out["pho_EB_Iso_chg"].fill(
dataset=dataset,
pho_EB_Iso_chg=Pho_EB_Iso_charge,
weight=skim_weight(weights.weight() * cuts * cuts_pho_EB),
)
return out