def test_firsts():
assert ak.to_list(
ak.firsts(
ak.singletons(ak.Array([1.1, 2.2, None, 3.3, None, None, 4.4, 5.5])), axis=1
)
) == [1.1, 2.2, None, 3.3, None, None, 4.4, 5.5]
assert (
ak.to_list(
ak.firsts(
ak.singletons(
ak.Array([[1.1, 2.2, None], [3.3, None], [None], [4.4, 5.5]])
),
axis=2,
)
)
== [[1.1, 2.2, None], [3.3, None], [None], [4.4, 5.5]]
)
assert (
ak.to_list(
ak.firsts(
ak.singletons(
ak.Array(
[[[1.1, 2.2, None]], [[3.3, None]], [[None]], [[4.4, 5.5]]]
)
),
axis=3,
)
)
== [[[1.1, 2.2, None]], [[3.3, None]], [[None]], [[4.4, 5.5]]]
)
def get_leading_pair(taus, Pt_thr):
tau_1 = taus[0]
tau_2 = taus[1]
# select only true taus that pass generator preselection
num_tau_mask_1 = num_mask_eff(tau_1, Pt_thr=Pt_thr)
num_tau_mask_2 = num_mask_eff(tau_2, Pt_thr=Pt_thr)
pair_mask = num_tau_mask_1 & num_tau_mask_2
ev_mask = ditau_selection(num_tau_mask_1, num_tau_mask_2)
tau_1_selected = (tau_1[pair_mask])[ev_mask]
tau_2_selected = (tau_2[pair_mask])[ev_mask]
# get arrays of leading and subleading taus
tau_leading = ak.firsts(tau_1_selected, axis=-1)
tau_subleading = ak.firsts(tau_2_selected, axis=-1)
return tau_leading, tau_subleading
def calculate_selection(self, syst_tag, events):
"""
"""
electrons = events.ele
electrons["label"] = -1 * awkward.ones_like(electrons.pt)
if not self.is_data:
electrons["label"] = awkward.where(
electrons.genPartFlav == 1, awkward.ones_like(electrons.label),
electrons.label)
electrons["label"] = awkward.where(
(electrons.genPartFlav == 3) | (electrons.genPartFlav == 4) |
(electrons.genPartFlav == 5),
awkward.zeros_like(electrons.label), electrons.label)
fields = [x for x in events.fields if x not in ["ele", "Electron"]]
for x in fields:
if x == "ZCand":
electrons[x] = awkward.firsts(events[x])
else:
electrons[x] = events[x]
electrons = awkward.flatten(electrons)
dummy_cut = electrons.pt >= 0
return dummy_cut, electrons
def set_genZ(events, genBranches, selection_options, debug):
if genBranches is None:
events["genZ_decayMode"] = ak.from_numpy(-1 * numpy.ones(len(events)))
else:
electron_idxs = abs(genBranches.pdgId) == 11
muon_idxs = abs(genBranches.pdgId) == 13
tau_idxs = abs(genBranches.pdgId) == 15
motherOfElectrons = genBranches.genPartIdxMother[electron_idxs]
motherOfMuons = genBranches.genPartIdxMother[muon_idxs]
motherOfTaus = genBranches.genPartIdxMother[tau_idxs]
ZToEleEvents = ak.sum(
(genBranches.pdgId[motherOfElectrons] == 23), axis=1) > 0
ZToMuEvents = ak.sum(
(genBranches.pdgId[motherOfMuons] == 23), axis=1) > 0
ZToTauEvents = ak.sum(
(genBranches.pdgId[motherOfTaus] == 23), axis=1) > 0
# W decay dudes
WToEEvents = ak.sum(
(abs(genBranches.pdgId[motherOfElectrons]) == 24), axis=1) > 0
WToMuEvents = ak.sum(
(abs(genBranches.pdgId[motherOfMuons]) == 24), axis=1) > 0
WToTauEvents = ak.sum(
(abs(genBranches.pdgId[motherOfTaus]) == 24), axis=1) > 0
events["genZ_decayMode"] = ak.from_numpy(
numpy.zeros(len(events))
) + 1 * ZToEleEvents + 2 * ZToMuEvents + 3 * ZToTauEvents + 4 * WToEEvents + 5 * WToMuEvents + 6 * WToTauEvents
events["tau_motherID"] = ak.fill_none(
ak.firsts(genBranches.pdgId[motherOfTaus]), 0)
return events
def VQQgenmatch(events):
dataset = events.metadata['dataset']
if 'WJetsToQQ' in dataset: motherId = 24
elif 'ZJetsToQQ' in dataset: motherId = 23
elif 'VectorDiJet' in dataset: motherId = 55
else:
return np.ones(len(events.FatJet.pt),
dtype=bool) #events.FatJet.pt.ones_like()
mother = ak.flatten(events.GenPart[
(abs(events.GenPart.pdgId) == motherId)
& events.GenPart.hasFlags(["isLastCopy", "fromHardProcess"])])
print(mother)
print(mother.children.pdgId)
print(mother.children)
try:
q0 = mother.children[:, 0]
q1 = mother.children[:, 1]
except:
q0 = mother.children[:, 0]
q1 = mother.children[:, 1]
print(q0.pdgId, q1.pdgId)
leading_jet = ak.firsts(events.FatJet)
return (leading_jet.delta_r2(q0) < 0.8 * 0.8) & (leading_jet.delta_r2(q1) <
0.8 * 0.8)
def test_cartesian():
muon = ak.Array([[{"pt": 1.0}], []], with_name="muon")
electron = ak.Array([[], [{"pt": 1.0}]], with_name="electron")
muon = muon[muon.pt > 5]
electron = electron[electron.pt > 5]
leptons = ak.concatenate([muon, electron], axis=1)
candidate = ak.firsts(leptons)
assert ak.to_list(ak.Array(candidate)) == [None, None]
result = ak.cartesian([candidate, candidate], axis=0)
assert ak.to_list(result) == [
(None, None),
(None, None),
(None, None),
(None, None),
]
result = ak.cartesian([candidate, ak.Array([[1, 2, 3], []])], axis=1)
assert ak.to_list(result) == [None, None]
one, two = ak.broadcast_arrays(candidate, ak.Array([[1, 2, 3], []]))
assert ak.to_list(one) == [None, None]
assert ak.to_list(two) == [None, None]
def test_firsts_singletons():
array = ak.Array([None, 1.1, 2.2, None, 3.3, None, None, 4.4, 5.5, None])
one = ak.singletons(array)
assert ak.to_list(one) == [[], [1.1], [2.2], [], [3.3], [], [], [4.4], [5.5], []]
two = ak.firsts(one)
assert ak.to_list(two) == [None, 1.1, 2.2, None, 3.3, None, None, 4.4, 5.5, None]
array = ak.repartition(array, 3)
assert isinstance(array.layout, ak.partition.PartitionedArray)
one = ak.singletons(array)
assert isinstance(one.layout, ak.partition.PartitionedArray)
assert ak.to_list(one) == [[], [1.1], [2.2], [], [3.3], [], [], [4.4], [5.5], []]
two = ak.firsts(one)
assert isinstance(two.layout, ak.partition.PartitionedArray)
assert ak.to_list(two) == [None, 1.1, 2.2, None, 3.3, None, None, 4.4, 5.5, None]
def nearest(
self,
other,
axis=1,
metric=lambda a, b: a.delta_r(b),
return_metric=False,
threshold=None,
):
"""Return nearest object to this one
Finds item in ``other`` satisfying ``min(metric(self, other))``.
The two arrays should be broadcast-compatible on all axes other than the specified
axis, which will be used to form a cartesian product. If axis=None, broadcast arrays directly.
The return shape will be that of ``self``.
Parameters
----------
other : awkward.Array
Another array with same shape in all but ``axis``
axis : int, optional
The axis to form the cartesian product (default 1). If None, the metric
is directly evaluated on the input arrays (i.e. they should broadcast)
metric : callable
A function of two arguments, returning a scalar. The default metric is `delta_r`.
return_metric : bool, optional
If true, return both the closest object and its metric (default false)
threshold : Number, optional
If set, any objects with ``metric > threshold`` will be masked from the result
"""
mval, (a, b) = self.metric_table(other,
axis,
metric,
return_combinations=True)
if axis is None:
# NotImplementedError: awkward.firsts with axis=-1
axis = other.layout.purelist_depth - 2
mmin = awkward.argmin(mval, axis=axis + 1, keepdims=True)
out = awkward.firsts(b[mmin], axis=axis + 1)
metric = awkward.firsts(mval[mmin], axis=axis + 1)
if threshold is not None:
out = out.mask[metric <= threshold]
if return_metric:
return out, metric
return out
def test_firsts():
array = ak.Array([[[0, 1, 2], []], [[3, 4]], [], [[5], [6, 7, 8, 9]]])
assert ak.to_list(ak.firsts(array, axis=0)) == [[0, 1, 2], []]
assert ak.to_list(ak.firsts(array, axis=1)) == [[0, 1, 2], [3, 4], None,
[5]]
assert ak.to_list(ak.firsts(array, axis=2)) == [[0, None], [3], [], [5, 6]]
assert ak.to_list(ak.firsts(array, axis=-1)) == [[0, None], [3], [],
[5, 6]]
assert ak.to_list(ak.firsts(array, axis=-2)) == [[0, 1, 2], [3, 4], None,
[5]]
assert ak.to_list(ak.firsts(array, axis=-3)) == [[0, 1, 2], []]
with pytest.raises(ValueError):
ak.firsts(array, axis=-4)
def get_vpt(check_offshell=False):
"""Only the leptonic samples have no resonance in the decay tree, and only
when M is beyond the configured Breit-Wigner cutoff (usually 15*width)
"""
boson = ak.firsts(
genpart[((genpart.pdgId == 23) | (abs(genpart.pdgId) == 24))
& genpart.hasFlags(["fromHardProcess", "isLastCopy"])])
if check_offshell:
offshell = genpart[
genpart.hasFlags(["fromHardProcess", "isLastCopy"])
& ak.is_none(boson)
& (abs(genpart.pdgId) >= 11) &
(abs(genpart.pdgId) <= 16)].sum()
return ak.where(ak.is_none(boson.pt), offshell.pt, boson.pt)
return np.array(ak.fill_none(boson.pt, 0.))
def TTsemileptonicmatch(events):
dataset = events.metadata['dataset']
if 'TTToSemiLeptonic' not in dataset:
return np.zeros(len(events.FatJet.pt), dtype=bool)
child = events.GenPart[
(abs(events.GenPart.pdgId) == 24)
& events.GenPart.hasFlags(["isLastCopy", "fromHardProcess"])].children
fatjet = ak.firsts(events.FatJet)
n_matched_quarks = np.zeros(len(fatjet))
for ii in [0, 1]:
for jj in [0, 1]:
n_matched_quarks = n_matched_quarks + ak.fill_none(
(fatjet.delta_r2(child[:, ii, jj]) < 0.8**2) &
(abs(child[:, ii, jj].pdgId) < 6), 0.)
print(n_matched_quarks)
return n_matched_quarks
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
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 get_hpt():
boson = ak.firsts(
genpart[(genpart.pdgId == 25)
& genpart.hasFlags(["fromHardProcess", "isLastCopy"])])
return np.array(ak.fill_none(boson.pt, 0.))
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]
# 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
inspected_event_weight = weight/abs(weight)
inspected_event_weight = inspected_event_weight.sum()
# 4. Particle Definition and Selections
print(".... Particle and Event selection")
# make Jet selection mask
Jet_sel_mask = (Jet.PT > 35 ) & (abs(Jet.Eta) < 2.4)
Jet = Jet[Jet_sel_mask] # selected jet
# Baseline selection
NJet_mask = ak.num(Jet) >= 9 # selected jet >=9 mask
HT_mask = ak.firsts(HT >= 700) # HT cut
Nbtag_mask = ak.sum(Jet.BTag,axis=1) >= 3 # at leat 3 b-tag
base_line_mask = NJet_mask & HT_mask & Nbtag_mask
# Apply cut using boolean indexing
Jet = Jet[base_line_mask]
HT = HT[base_line_mask]
inspected_event_weight = inspected_event_weight[base_line_mask]
print("After N jet cut: ",len(Jet),len(HT),len(inspected_event_weight))
# 5. Flatten the arrays ( 2D -> 1D )
print(".... Write file")
# Fltten and chage type Awkward array to numpy array
def process(self, events):
output = self.accumulator.identity()
dataset = events.metadata['dataset']
isRealData = 'genWeight' not in events.fields
if not isRealData:
output['sumw'][dataset] += sum(events.genWeight)
JECversion = JECversions[str(self.year)]['MC']
else:
output['nbtagmu'][dataset] += ak.count(events.event)
JECversion = JECversions[str(
self.year)]['Data'][dataset.split('BTagMu')[1]]
############
# Some corrections
weights = processor.Weights(len(events))
corrections = {}
if not isRealData:
weights.add('genWeight', events.genWeight)
weights.add(
'pileup_weight',
self.puReweight(self.puFile, self.nTrueFile,
dataset)(events.Pileup.nPU))
events.FatJet = self.applyJEC(events.FatJet,
events.fixedGridRhoFastjetAll,
events.caches[0], 'AK8PFPuppi',
isRealData, JECversion)
cuts = processor.PackedSelection()
############
# Trigger selection
if self.year == 2016:
if 'BTagMu_AK4Jet300_Mu5' not in events.HLT.fields:
self.triggers = [
trigger.replace('AK4', '') for trigger in self.triggers
]
elif self.year == 2018:
for (i, trigger) in enumerate(self.triggers):
if trigger.strip("HLT_") not in events.HLT.fields:
self.triggers[i] = trigger + "_noalgo"
trig_arrs = [
events.HLT[_trig.strip("HLT_")] for _trig in self.triggers
]
req_trig = np.zeros(len(events), dtype='bool')
for t in trig_arrs:
req_trig = req_trig | t
cuts.add('trigger', ak.to_numpy(req_trig))
############
# Basic cuts
## Muon cuts
# muon twiki: https://twiki.cern.ch/twiki/bin/view/CMS/SWGuideMuonIdRun2
events.Muon = events.Muon[(events.Muon.pt > 5)
& (abs(events.Muon.eta < 2.4)) &
(events.Muon.tightId != 1) &
(events.Muon.pfRelIso04_all > 0.15)]
events.Muon = ak.pad_none(events.Muon, 2, axis=1)
## Jet cuts (not used)
events.Jet = events.Jet[(events.Jet.pt > 25)
& (abs(events.Jet.eta) <= 2.5)]
#req_jets = (ak.count(events.Jet.pt, axis=1) >= 2)
## FatJet cuts
events.FatJet = events.FatJet[
(events.FatJet.pt > self._mask_fatjets['basic']['pt_cut']) &
(abs(events.FatJet.eta) <= self._mask_fatjets['basic']['eta_cut'])
& (events.FatJet.jetId > self._mask_fatjets['basic']['jetId_cut'])
& (ak.count(events.FatJet.subjets.pt, axis=2) >=
2)] ## subjet sel to crosscheck
#print(events['FatJetSVs'])
## Event level variables
eventVariables = {}
eventVariables['nfatjet'] = ak.num(events.FatJet)
## Leading jet variables
leadfatjet = ak.firsts(events.FatJet)
leadfatjet['tau21'] = leadfatjet.tau2 / leadfatjet.tau1
subjet1 = ak.pad_none(leadfatjet.subjets, 2)[:, 0]
subjet2 = ak.pad_none(leadfatjet.subjets, 2)[:, 1]
leadfatjet['nsv1'] = get_nsv(subjet1, events.SV)
leadfatjet['nsv2'] = get_nsv(subjet2, events.SV)
leadfatjet['nmusj1'] = ak.num(subjet1.delta_r(events.Muon) < 0.4)
leadfatjet['nmusj2'] = ak.num(subjet2.delta_r(events.Muon) < 0.4)
fatjet_mutag = (leadfatjet.nmusj1 >= 1) & (leadfatjet.nmusj2 >= 1)
cuts.add('fatjet_mutag', ak.to_numpy(fatjet_mutag))
for DDX in self._mask_DDX.keys():
for wp, cut in self._mask_DDX[DDX].items():
DDX_pass = (leadfatjet[f'btag{DDX}vLV2'] > cut)
DDX_fail = (leadfatjet[f'btag{DDX}vLV2'] < cut)
cuts.add(f'{DDX}_pass{wp}wp', ak.to_numpy(DDX_pass))
cuts.add(f'{DDX}_fail{wp}wp', ak.to_numpy(DDX_fail))
flavors = {}
if not isRealData:
flavors['b'] = (leadfatjet.hadronFlavour == 5)
flavors['c'] = (leadfatjet.hadronFlavour == 4)
flavors['l'] = (leadfatjet.hadronFlavour < 4)
flavors['bb'] = abs(leadfatjet.hadronFlavour == 5) & (
leadfatjet.nBHadrons >= 2) #& (leadfatjet.nCHadrons == 0)
flavors['cc'] = abs(leadfatjet.hadronFlavour == 4) & (
leadfatjet.nBHadrons == 0) & (leadfatjet.nCHadrons >= 2)
#flavors['ll'] = abs(leadfatjet.hadronFlavour < 4) & (leadfatjet.nBHadrons == 0) & (leadfatjet.nCHadrons == 0)
flavors['b'] = flavors['b'] & ~flavors['bb']
flavors['c'] = flavors['c'] & ~flavors['cc']
flavors['l'] = flavors['l'] & ~flavors['bb'] & ~flavors[
'cc'] & ~flavors['b'] & ~flavors['c']
#flavors['others'] = ~flavors['l'] & ~flavors['bb'] & ~flavors['cc'] & ~flavors['b'] & ~flavors['c']
else:
flavors['Data'] = np.ones(len(events), dtype='bool')
for selname, cut in self._mask_fatjets.items():
sel = (leadfatjet.pt > cut['pt_cut']) & \
(leadfatjet.msoftdrop > cut['mass_cut']) & \
(abs(leadfatjet.eta) < cut['eta_cut']) & \
(leadfatjet.jetId >= cut['jetId_cut']) & \
(leadfatjet.tau21 < cut['tau21_cut'])
#(leadfatjet.Hbb > cut['Hbb'])
cuts.add(selname, ak.to_numpy(sel))
selection = {}
selection['basic'] = {'trigger', 'basic'}
selection['pt350msd50'] = {'trigger', 'fatjet_mutag', 'pt350msd50'}
selection['msd100tau06'] = {'trigger', 'fatjet_mutag', 'msd100tau06'}
selection['pt400msd100tau06'] = {
'trigger', 'fatjet_mutag', 'pt400msd100tau06'
}
for mask_f in self._final_mask:
for DDX in self._mask_DDX.keys():
for wp, cut in self._mask_DDX[DDX].items():
selection[f'{mask_f}{DDX}pass{wp}wp'] = selection[
mask_f].copy()
selection[f'{mask_f}{DDX}pass{wp}wp'].add(
f'{DDX}_pass{wp}wp')
selection[f'{mask_f}{DDX}fail{wp}wp'] = selection[
mask_f].copy()
selection[f'{mask_f}{DDX}fail{wp}wp'].add(
f'{DDX}_fail{wp}wp')
for histname, h in output.items():
sel = [r for r in selection.keys() if r in histname.split('_')]
if ((histname in self.fatjet_hists) |
('hist2d_fatjet' in histname)):
for flav, mask in flavors.items():
weight = weights.weight() * cuts.all(
*selection[sel[0]]) * ak.to_numpy(mask)
fields = {
k: ak.fill_none(leadfatjet[k], -9999)
for k in h.fields if k in dir(leadfatjet)
}
h.fill(dataset=dataset,
flavor=flav,
**fields,
weight=weight)
if histname in self.event_hists:
for flav, mask in flavors.items():
weight = weights.weight() * cuts.all(
*selection[sel[0]]) * ak.to_numpy(mask)
fields = {
k: ak.fill_none(eventVariables[k], -9999)
for k in h.fields if k in eventVariables.keys()
}
h.fill(dataset=dataset,
flavor=flav,
**fields,
weight=weight)
return output
def process_shift(self, events, shift_name):
dataset = events.metadata['dataset']
isRealData = not hasattr(events, "genWeight")
selection = PackedSelection()
weights = Weights(len(events), storeIndividual=True)
output = self.make_output()
if shift_name is None and not isRealData:
output['sumw'] = ak.sum(events.genWeight)
if isRealData or self._newTrigger:
trigger = np.zeros(len(events), dtype='bool')
for t in self._triggers[self._year]:
if t in events.HLT.fields:
trigger = trigger | events.HLT[t]
selection.add('trigger', trigger)
del trigger
else:
selection.add('trigger', np.ones(len(events), dtype='bool'))
if isRealData:
selection.add(
'lumimask', lumiMasks[self._year](events.run,
events.luminosityBlock))
else:
selection.add('lumimask', np.ones(len(events), dtype='bool'))
if isRealData and self._skipRunB and self._year == '2017':
selection.add('dropB', events.run > 299329)
else:
selection.add('dropB', np.ones(len(events), dtype='bool'))
if isRealData:
trigger = np.zeros(len(events), dtype='bool')
for t in self._muontriggers[self._year]:
if t in events.HLT.fields:
trigger |= np.array(events.HLT[t])
selection.add('muontrigger', trigger)
del trigger
else:
selection.add('muontrigger', np.ones(len(events), dtype='bool'))
metfilter = np.ones(len(events), dtype='bool')
for flag in self._met_filters[
self._year]['data' if isRealData else 'mc']:
metfilter &= np.array(events.Flag[flag])
selection.add('metfilter', metfilter)
del metfilter
fatjets = events.FatJet
fatjets['msdcorr'] = corrected_msoftdrop(fatjets)
fatjets['qcdrho'] = 2 * np.log(fatjets.msdcorr / fatjets.pt)
fatjets['n2ddt'] = fatjets.n2b1 - n2ddt_shift(fatjets, year=self._year)
fatjets['msdcorr_full'] = fatjets['msdcorr'] * self._msdSF[self._year]
candidatejet = fatjets[
# https://github.com/DAZSLE/BaconAnalyzer/blob/master/Analyzer/src/VJetLoader.cc#L269
(fatjets.pt > 200)
& (abs(fatjets.eta) < 2.5)
& fatjets.isTight # this is loose in sampleContainer
]
candidatejet = candidatejet[:, :
2] # Only consider first two to match generators
if self._jet_arbitration == 'pt':
candidatejet = ak.firsts(candidatejet)
elif self._jet_arbitration == 'mass':
candidatejet = ak.firsts(candidatejet[ak.argmax(
candidatejet.msdcorr, axis=1, keepdims=True)])
elif self._jet_arbitration == 'n2':
candidatejet = ak.firsts(candidatejet[ak.argmin(candidatejet.n2ddt,
axis=1,
keepdims=True)])
elif self._jet_arbitration == 'ddb':
candidatejet = ak.firsts(candidatejet[ak.argmax(
candidatejet.btagDDBvLV2, axis=1, keepdims=True)])
elif self._jet_arbitration == 'ddc':
candidatejet = ak.firsts(candidatejet[ak.argmax(
candidatejet.btagDDCvLV2, axis=1, keepdims=True)])
else:
raise RuntimeError("Unknown candidate jet arbitration")
if self._tagger == 'v1':
bvl = candidatejet.btagDDBvL
cvl = candidatejet.btagDDCvL
cvb = candidatejet.btagDDCvB
elif self._tagger == 'v2':
bvl = candidatejet.btagDDBvLV2
cvl = candidatejet.btagDDCvLV2
cvb = candidatejet.btagDDCvBV2
elif self._tagger == 'v3':
bvl = candidatejet.particleNetMD_Xbb
cvl = candidatejet.particleNetMD_Xcc / (
1 - candidatejet.particleNetMD_Xbb)
cvb = candidatejet.particleNetMD_Xcc / (
candidatejet.particleNetMD_Xcc +
candidatejet.particleNetMD_Xbb)
elif self._tagger == 'v4':
bvl = candidatejet.particleNetMD_Xbb
cvl = candidatejet.btagDDCvLV2
cvb = candidatejet.particleNetMD_Xcc / (
candidatejet.particleNetMD_Xcc +
candidatejet.particleNetMD_Xbb)
else:
raise ValueError("Not an option")
selection.add('minjetkin', (candidatejet.pt >= 450)
& (candidatejet.pt < 1200)
& (candidatejet.msdcorr >= 40.)
& (candidatejet.msdcorr < 201.)
& (abs(candidatejet.eta) < 2.5))
selection.add('_strict_mass', (candidatejet.msdcorr > 85) &
(candidatejet.msdcorr < 130))
selection.add('_high_score', cvl > 0.8)
selection.add('minjetkinmu', (candidatejet.pt >= 400)
& (candidatejet.pt < 1200)
& (candidatejet.msdcorr >= 40.)
& (candidatejet.msdcorr < 201.)
& (abs(candidatejet.eta) < 2.5))
selection.add('minjetkinw', (candidatejet.pt >= 200)
& (candidatejet.pt < 1200)
& (candidatejet.msdcorr >= 40.)
& (candidatejet.msdcorr < 201.)
& (abs(candidatejet.eta) < 2.5))
selection.add('jetid', candidatejet.isTight)
selection.add('n2ddt', (candidatejet.n2ddt < 0.))
if not self._tagger == 'v2':
selection.add('ddbpass', (bvl >= 0.89))
selection.add('ddcpass', (cvl >= 0.83))
selection.add('ddcvbpass', (cvb >= 0.2))
else:
selection.add('ddbpass', (bvl >= 0.7))
selection.add('ddcpass', (cvl >= 0.45))
selection.add('ddcvbpass', (cvb >= 0.03))
jets = events.Jet
jets = jets[(jets.pt > 30.) & (abs(jets.eta) < 2.5) & jets.isTight]
# only consider first 4 jets to be consistent with old framework
jets = jets[:, :4]
dphi = abs(jets.delta_phi(candidatejet))
selection.add(
'antiak4btagMediumOppHem',
ak.max(jets[dphi > np.pi / 2][self._ak4tagBranch],
axis=1,
mask_identity=False) <
BTagEfficiency.btagWPs[self._ak4tagger][self._year]['medium'])
ak4_away = jets[dphi > 0.8]
selection.add(
'ak4btagMedium08',
ak.max(ak4_away[self._ak4tagBranch], axis=1, mask_identity=False) >
BTagEfficiency.btagWPs[self._ak4tagger][self._year]['medium'])
met = events.MET
selection.add('met', met.pt < 140.)
goodmuon = ((events.Muon.pt > 10)
& (abs(events.Muon.eta) < 2.4)
& (events.Muon.pfRelIso04_all < 0.25)
& events.Muon.looseId)
nmuons = ak.sum(goodmuon, axis=1)
leadingmuon = ak.firsts(events.Muon[goodmuon])
if self._looseTau:
goodelectron = ((events.Electron.pt > 10)
& (abs(events.Electron.eta) < 2.5)
&
(events.Electron.cutBased >= events.Electron.VETO))
nelectrons = ak.sum(goodelectron, axis=1)
ntaus = ak.sum(
((events.Tau.pt > 20)
& (abs(events.Tau.eta) < 2.3)
& events.Tau.idDecayMode
& ((events.Tau.idMVAoldDM2017v2 & 2) != 0)
& ak.all(events.Tau.metric_table(events.Muon[goodmuon]) > 0.4,
axis=2)
& ak.all(events.Tau.metric_table(
events.Electron[goodelectron]) > 0.4,
axis=2)),
axis=1,
)
else:
goodelectron = (
(events.Electron.pt > 10)
& (abs(events.Electron.eta) < 2.5)
& (events.Electron.cutBased >= events.Electron.LOOSE))
nelectrons = ak.sum(goodelectron, axis=1)
ntaus = ak.sum(
(events.Tau.pt > 20)
&
events.Tau.idDecayMode # bacon iso looser than Nano selection
& ak.all(events.Tau.metric_table(events.Muon[goodmuon]) > 0.4,
axis=2)
& ak.all(events.Tau.metric_table(events.Electron[goodelectron])
> 0.4,
axis=2),
axis=1,
)
selection.add('noleptons',
(nmuons == 0) & (nelectrons == 0) & (ntaus == 0))
selection.add('onemuon',
(nmuons == 1) & (nelectrons == 0) & (ntaus == 0))
selection.add('muonkin',
(leadingmuon.pt > 55.) & (abs(leadingmuon.eta) < 2.1))
selection.add('muonDphiAK8',
abs(leadingmuon.delta_phi(candidatejet)) > 2 * np.pi / 3)
# W-Tag (Tag and Probe)
# tag side
selection.add(
'ak4btagMediumOppHem',
ak.max(jets[dphi > np.pi / 2][self._ak4tagBranch],
axis=1,
mask_identity=False) >
BTagEfficiency.btagWPs[self._ak4tagger][self._year]['medium'])
selection.add('met40p', met.pt > 40.)
selection.add('tightMuon',
(leadingmuon.tightId) & (leadingmuon.pt > 53.))
# selection.add('ptrecoW', (leadingmuon + met).pt > 250.)
selection.add('ptrecoW200', (leadingmuon + met).pt > 200.)
selection.add(
'ak4btagNearMu',
leadingmuon.delta_r(leadingmuon.nearest(ak4_away, axis=None)) <
2.0)
_bjets = jets[self._ak4tagBranch] > BTagEfficiency.btagWPs[
self._ak4tagger][self._year]['medium']
# _nearAK8 = jets.delta_r(candidatejet) < 0.8
# _nearMu = jets.delta_r(ak.firsts(events.Muon)) < 0.3
# selection.add('ak4btagOld', ak.sum(_bjets & ~_nearAK8 & ~_nearMu, axis=1) >= 1)
_nearAK8 = jets.delta_r(candidatejet) < 0.8
_nearMu = jets.delta_r(leadingmuon) < 0.3
selection.add('ak4btagOld',
ak.sum(_bjets & ~_nearAK8 & ~_nearMu, axis=1) >= 1)
# _nearAK8 = jets.delta_r(candidatejet) < 0.8
# _nearMu = jets.delta_r(candidatejet.nearest(events.Muon[goodmuon], axis=None)) < 0.3
# selection.add('ak4btagNew', ak.sum(_bjets & ~_nearAK8 & ~_nearMu, axis=1) >= 1)
# probe side
selection.add('minWjetpteta',
(candidatejet.pt >= 200) & (abs(candidatejet.eta) < 2.4))
# selection.add('noNearMuon', candidatejet.delta_r(candidatejet.nearest(events.Muon[goodmuon], axis=None)) > 1.0)
selection.add('noNearMuon', candidatejet.delta_r(leadingmuon) > 1.0)
#####
if isRealData:
genflavor = ak.zeros_like(candidatejet.pt)
else:
if 'HToCC' in dataset or 'HToBB' in dataset:
if self._ewkHcorr:
add_HiggsEW_kFactors(weights, events.GenPart, dataset)
weights.add('genweight', events.genWeight)
if "PSWeight" in events.fields:
add_ps_weight(weights, events.PSWeight)
else:
add_ps_weight(weights, None)
if "LHEPdfWeight" in events.fields:
add_pdf_weight(weights, events.LHEPdfWeight)
else:
add_pdf_weight(weights, None)
if "LHEScaleWeight" in events.fields:
add_scalevar_7pt(weights, events.LHEScaleWeight)
add_scalevar_3pt(weights, events.LHEScaleWeight)
else:
add_scalevar_7pt(weights, [])
add_scalevar_3pt(weights, [])
add_pileup_weight(weights, events.Pileup.nPU, self._year, dataset)
bosons = getBosons(events.GenPart)
matchedBoson = candidatejet.nearest(bosons,
axis=None,
threshold=0.8)
if self._tightMatch:
match_mask = (
(candidatejet.pt - matchedBoson.pt) / matchedBoson.pt <
0.5) & ((candidatejet.msdcorr - matchedBoson.mass) /
matchedBoson.mass < 0.3)
selmatchedBoson = ak.mask(matchedBoson, match_mask)
genflavor = bosonFlavor(selmatchedBoson)
else:
genflavor = bosonFlavor(matchedBoson)
genBosonPt = ak.fill_none(ak.firsts(bosons.pt), 0)
if self._newVjetsKfactor:
add_VJets_kFactors(weights, events.GenPart, dataset)
else:
add_VJets_NLOkFactor(weights, genBosonPt, self._year, dataset)
if shift_name is None:
output['btagWeight'].fill(val=self._btagSF.addBtagWeight(
weights, ak4_away, self._ak4tagBranch))
if self._nnlops_rew and dataset in [
'GluGluHToCC_M125_13TeV_powheg_pythia8'
]:
weights.add('minlo_rew',
powheg_to_nnlops(ak.to_numpy(genBosonPt)))
if self._newTrigger:
add_jetTriggerSF(
weights, ak.firsts(fatjets),
self._year if not self._skipRunB else f'{self._year}CDEF',
selection)
else:
add_jetTriggerWeight(weights, candidatejet.msdcorr,
candidatejet.pt, self._year)
add_mutriggerSF(weights, leadingmuon, self._year, selection)
add_mucorrectionsSF(weights, leadingmuon, self._year, selection)
if self._year in ("2016", "2017"):
weights.add("L1Prefiring", events.L1PreFiringWeight.Nom,
events.L1PreFiringWeight.Up,
events.L1PreFiringWeight.Dn)
logger.debug("Weight statistics: %r" % weights.weightStatistics)
msd_matched = candidatejet.msdcorr * self._msdSF[self._year] * (
genflavor > 0) + candidatejet.msdcorr * (genflavor == 0)
regions = {
'signal': [
'noleptons', 'minjetkin', 'met', 'metfilter', 'jetid',
'antiak4btagMediumOppHem', 'n2ddt', 'trigger', 'lumimask'
],
'signal_noddt': [
'noleptons', 'minjetkin', 'met', 'jetid',
'antiak4btagMediumOppHem', 'trigger', 'lumimask', 'metfilter'
],
# 'muoncontrol': ['minjetkinmu', 'jetid', 'n2ddt', 'ak4btagMedium08', 'onemuon', 'muonkin', 'muonDphiAK8', 'muontrigger', 'lumimask', 'metfilter'],
'muoncontrol': [
'onemuon', 'muonkin', 'muonDphiAK8', 'metfilter',
'minjetkinmu', 'jetid', 'ak4btagMedium08', 'n2ddt',
'muontrigger', 'lumimask'
],
'muoncontrol_noddt': [
'onemuon', 'muonkin', 'muonDphiAK8', 'jetid', 'metfilter',
'minjetkinmu', 'jetid', 'ak4btagMedium08', 'muontrigger',
'lumimask'
],
'wtag': [
'onemuon', 'tightMuon', 'minjetkinw', 'jetid', 'met40p',
'metfilter', 'ptrecoW200', 'ak4btagOld', 'muontrigger',
'lumimask'
],
'wtag0': [
'onemuon', 'tightMuon', 'met40p', 'metfilter', 'ptrecoW200',
'ak4btagOld', 'muontrigger', 'lumimask'
],
'wtag2': [
'onemuon', 'tightMuon', 'minjetkinw', 'jetid',
'ak4btagMediumOppHem', 'met40p', 'metfilter', 'ptrecoW200',
'ak4btagOld', 'muontrigger', 'lumimask'
],
'noselection': [],
}
def normalize(val, cut):
if cut is None:
ar = ak.to_numpy(ak.fill_none(val, np.nan))
return ar
else:
ar = ak.to_numpy(ak.fill_none(val[cut], np.nan))
return ar
import time
tic = time.time()
if shift_name is None:
for region, cuts in regions.items():
allcuts = set([])
cut = selection.all(*allcuts)
output['cutflow_msd'].fill(region=region,
genflavor=normalize(
genflavor, None),
cut=0,
weight=weights.weight(),
msd=normalize(msd_matched, None))
output['cutflow_eta'].fill(region=region,
genflavor=normalize(genflavor, cut),
cut=0,
weight=weights.weight()[cut],
eta=normalize(
candidatejet.eta, cut))
output['cutflow_pt'].fill(region=region,
genflavor=normalize(genflavor, cut),
cut=0,
weight=weights.weight()[cut],
pt=normalize(candidatejet.pt, cut))
for i, cut in enumerate(cuts + ['ddcvbpass', 'ddcpass']):
allcuts.add(cut)
cut = selection.all(*allcuts)
output['cutflow_msd'].fill(region=region,
genflavor=normalize(
genflavor, cut),
cut=i + 1,
weight=weights.weight()[cut],
msd=normalize(msd_matched, cut))
output['cutflow_eta'].fill(
region=region,
genflavor=normalize(genflavor, cut),
cut=i + 1,
weight=weights.weight()[cut],
eta=normalize(candidatejet.eta, cut))
output['cutflow_pt'].fill(
region=region,
genflavor=normalize(genflavor, cut),
cut=i + 1,
weight=weights.weight()[cut],
pt=normalize(candidatejet.pt, cut))
if self._evtVizInfo and 'ddcpass' in allcuts and isRealData and region == 'signal':
if 'event' not in events.fields:
continue
_cut = selection.all(*allcuts, '_strict_mass',
'_high_score')
# _cut = selection.all('_strict_mass'')
output['to_check'][
'mass'] += processor.column_accumulator(
normalize(msd_matched, _cut))
nfatjet = ak.sum(
((fatjets.pt > 200) &
(abs(fatjets.eta) < 2.5) & fatjets.isTight),
axis=1)
output['to_check'][
'njet'] += processor.column_accumulator(
normalize(nfatjet, _cut))
output['to_check'][
'fname'] += processor.column_accumulator(
np.array([events.metadata['filename']] *
len(normalize(msd_matched, _cut))))
output['to_check'][
'event'] += processor.column_accumulator(
normalize(events.event, _cut))
output['to_check'][
'luminosityBlock'] += processor.column_accumulator(
normalize(events.luminosityBlock, _cut))
output['to_check'][
'run'] += processor.column_accumulator(
normalize(events.run, _cut))
if shift_name is None:
systematics = [None] + list(weights.variations)
else:
systematics = [shift_name]
def fill(region, systematic, wmod=None):
selections = regions[region]
cut = selection.all(*selections)
sname = 'nominal' if systematic is None else systematic
if wmod is None:
if systematic in weights.variations:
weight = weights.weight(modifier=systematic)[cut]
else:
weight = weights.weight()[cut]
else:
weight = weights.weight()[cut] * wmod[cut]
output['templates'].fill(
region=region,
systematic=sname,
runid=runmap(events.run)[cut],
genflavor=normalize(genflavor, cut),
pt=normalize(candidatejet.pt, cut),
msd=normalize(msd_matched, cut),
ddb=normalize(bvl, cut),
ddc=normalize(cvl, cut),
ddcvb=normalize(cvb, cut),
weight=weight,
)
if region in [
'wtag', 'wtag0', 'wtag2', 'wtag3', 'wtag4', 'wtag5',
'wtag6', 'wtag7', 'noselection'
]: # and sname in ['nominal', 'pileup_weightDown', 'pileup_weightUp', 'jet_triggerDown', 'jet_triggerUp']:
output['wtag'].fill(
region=region,
systematic=sname,
genflavor=normalize(genflavor, cut),
pt=normalize(candidatejet.pt, cut),
msd=normalize(msd_matched, cut),
n2ddt=normalize(candidatejet.n2ddt, cut),
ddc=normalize(cvl, cut),
ddcvb=normalize(cvb, cut),
weight=weight,
)
# if region in ['signal', 'noselection']:
# output['etaphi'].fill(
# region=region,
# systematic=sname,
# runid=runmap(events.run)[cut],
# genflavor=normalize(genflavor, cut),
# pt=normalize(candidatejet.pt, cut),
# eta=normalize(candidatejet.eta, cut),
# phi=normalize(candidatejet.phi, cut),
# ddc=normalize(cvl, cut),
# ddcvb=normalize(cvb, cut),
# ),
if not isRealData:
if wmod is not None:
_custom_weight = events.genWeight[cut] * wmod[cut]
else:
_custom_weight = np.ones_like(weight)
output['genresponse_noweight'].fill(
region=region,
systematic=sname,
pt=normalize(candidatejet.pt, cut),
genpt=normalize(genBosonPt, cut),
weight=_custom_weight,
)
output['genresponse'].fill(
region=region,
systematic=sname,
pt=normalize(candidatejet.pt, cut),
genpt=normalize(genBosonPt, cut),
weight=weight,
)
if systematic is None:
output['signal_opt'].fill(
region=region,
genflavor=normalize(genflavor, cut),
ddc=normalize(cvl, cut),
ddcvb=normalize(cvb, cut),
msd=normalize(msd_matched, cut),
weight=weight,
)
output['signal_optb'].fill(
region=region,
genflavor=normalize(genflavor, cut),
ddb=normalize(bvl, cut),
msd=normalize(msd_matched, cut),
weight=weight,
)
for region in regions:
cut = selection.all(*(set(regions[region]) - {'n2ddt'}))
if shift_name is None:
output['nminus1_n2ddt'].fill(
region=region,
n2ddt=normalize(candidatejet.n2ddt, cut),
weight=weights.weight()[cut],
)
for systematic in systematics:
if isRealData and systematic is not None:
continue
fill(region, systematic)
if shift_name is None and 'GluGluH' in dataset and 'LHEWeight' in events.fields:
for i in range(9):
fill(region, 'LHEScale_%d' % i, events.LHEScaleWeight[:,
i])
for c in events.LHEWeight.fields[1:]:
fill(region, 'LHEWeight_%s' % c, events.LHEWeight[c])
toc = time.time()
output["filltime"] = toc - tic
if shift_name is None:
output["weightStats"] = weights.weightStatistics
return {dataset: output}
def process(self, events):
def normalize(val, cut):
return ak.to_numpy(ak.fill_none(
val[cut],
np.nan)) #val[cut].pad(1, clip=True).fillna(0).flatten()
def fill(region, cuts, systematic=None, wmod=None):
print('filling %s' % region)
selections = cuts
cut = selection.all(*selections)
if 'signal' in region: weight = weights_signal.weight()[cut]
elif 'muonCR' in region: weight = weights_muonCR.weight()[cut]
elif 'VtaggingCR' in region:
weight = weights_VtaggingCR.weight()[cut]
output['templates'].fill(
dataset=dataset,
region=region,
pt=normalize(candidatejet.pt, cut),
msd=normalize(candidatejet.msdcorr, cut),
n2ddt=normalize(candidatejet.n2ddt, cut),
#gruddt=normalize(candidatejet.gruddt, cut),
in_v3_ddt=normalize(candidatejet.in_v3_ddt, cut),
hadW=normalize(candidatejet.nmatcheddau, cut),
weight=weight,
),
output['event'].fill(
dataset=dataset,
region=region,
MET=events.MET.pt[cut],
#nJet=fatjets.counts[cut],
nPFConstituents=normalize(candidatejet.nPFConstituents, cut),
weight=weight,
),
output['deepAK8'].fill(
dataset=dataset,
region=region,
deepTagMDWqq=normalize(candidatejet.deepTagMDWqq, cut),
deepTagMDZqq=normalize(candidatejet.deepTagMDZqq, cut),
msd=normalize(candidatejet.msdcorr, cut),
#genflavor=genflavor[cut],
weight=weight,
),
output['in_v3'].fill(
dataset=dataset,
region=region,
#genflavor=genflavor[cut],
in_v3=normalize(candidatejet.in_v3, cut),
n2=normalize(candidatejet.n2b1, cut),
gru=normalize(candidatejet.gru, cut),
weight=weight,
),
if 'muonCR' in dataset or 'VtaggingCR' in dataset:
output['muon'].fill(
dataset=dataset,
region=region,
mu_pt=normalize(candidatemuon.pt, cut),
mu_eta=normalize(candidatemuon.eta, cut),
mu_pfRelIso04_all=normalize(candidatemuon.pfRelIso04_all,
cut),
weight=weight,
),
#common jet kinematics
gru = events.GRU
IN = events.IN
fatjets = events.FatJet
fatjets['msdcorr'] = corrected_msoftdrop(fatjets)
fatjets['qcdrho'] = 2 * np.log(fatjets.msdcorr / fatjets.pt)
fatjets['gruddt'] = gru.v25 - shift(
fatjets, algo='gruddt', year='2017')
fatjets['gru'] = gru.v25
fatjets['in_v3'] = IN.v3
fatjets['in_v3_ddt'] = IN.v3 - shift(
fatjets, algo='inddt', year='2017')
fatjets['in_v3_ddt_90pctl'] = IN.v3 - shift(
fatjets, algo='inddt90pctl', year='2017')
fatjets['n2ddt'] = fatjets.n2b1 - n2ddt_shift(fatjets, year='2017')
fatjets['nmatcheddau'] = TTsemileptonicmatch(events)
dataset = events.metadata['dataset']
print('process dataset', dataset)
isRealData = not hasattr(events, 'genWeight')
output = self.accumulator.identity()
if (len(events) == 0): return output
selection = PackedSelection('uint64')
weights_signal = Weights(len(events))
weights_muonCR = Weights(len(events))
weights_VtaggingCR = Weights(len(events))
if not isRealData:
output['sumw'][dataset] += ak.sum(events.genWeight)
#######################
if 'signal' in self._region:
if isRealData:
trigger_fatjet = np.zeros(len(events), dtype='bool')
for t in self._triggers[self._year]:
try:
trigger_fatjet = trigger_fatjet | events.HLT[t]
except:
print('trigger %s not available' % t)
continue
else:
trigger_fatjet = np.ones(len(events), dtype='bool')
fatjets["genMatchFull"] = VQQgenmatch(events)
candidatejet = ak.firsts(fatjets)
candidatejet["genMatchFull"] = VQQgenmatch(events)
nelectrons = ak.sum(
(events.Electron.pt > 10.)
& (abs(events.Electron.eta) < 2.5)
& (events.Electron.cutBased >= events.Electron.VETO),
axis=1,
)
nmuons = ak.sum(
(events.Muon.pt > 10)
& (abs(events.Muon.eta) < 2.1)
& (events.Muon.pfRelIso04_all < 0.4)
& (events.Muon.looseId),
axis=1,
)
ntaus = ak.sum(
(events.Tau.pt > 20.)
& (events.Tau.idDecayMode)
& (events.Tau.rawIso < 5)
& (abs(events.Tau.eta) < 2.3),
axis=1,
)
cuts = {
"S_fatjet_trigger":
trigger_fatjet,
"S_pt":
candidatejet.pt > 525,
"S_eta": (abs(candidatejet.eta) < 2.5),
"S_msdcorr": (candidatejet.msdcorr > 40),
"S_rho":
((candidatejet.qcdrho > -5.5) & (candidatejet.qcdrho < -2.)),
"S_jetid": (candidatejet.isTight),
"S_VQQgenmatch": (candidatejet.genMatchFull),
"S_noelectron": (nelectrons == 0),
"S_nomuon": (nmuons == 0),
"S_notau": (ntaus == 0),
}
for name, cut in cuts.items():
print(name, cut)
selection.add(name, cut)
if isRealData:
genflavor = 0 #candidatejet.pt.zeros_like().pad(1, clip=True).fillna(-1).flatten()
if not isRealData:
weights_signal.add('genweight', events.genWeight)
#add_pileup_weight(weights_signal, events.Pileup.nPU, self._year, dataset)
add_jetTriggerWeight(weights_signal, candidatejet.msdcorr,
candidatejet.pt, self._year)
bosons = getBosons(events.GenPart)
genBosonPt = ak.fill_none(ak.firsts(bosons.pt), 0)
add_VJets_NLOkFactor(weights_signal, genBosonPt, self._year,
dataset)
#genflavor = matchedBosonFlavor(candidatejet, bosons).pad(1, clip=True).fillna(-1).flatten()
allcuts_signal = set()
output['cutflow_signal'][dataset]['none'] += float(
weights_signal.weight().sum())
for cut in cuts:
allcuts_signal.add(cut)
output['cutflow_signal'][dataset][cut] += float(
weights_signal.weight()[selection.all(
*allcuts_signal)].sum())
fill('signal', cuts.keys())
#######################
if 'muonCR' in self._region:
if isRealData:
trigger_muon = np.zeros(len(events), dtype='bool')
for t in self._muontriggers[self._year]:
trigger_muon = trigger_muon | events.HLT[t]
else:
trigger_muon = np.ones(len(events), dtype='bool')
candidatejet = ak.firsts(fatjets)
candidatemuon = events.Muon[:, :5]
jets = events.Jet[((events.Jet.pt > 50.)
& (abs(events.Jet.eta) < 2.5)
& (events.Jet.isTight))][:, :4]
dphi = abs(jets.delta_phi(candidatejet))
ak4_away = jets[(dphi > 0.8)]
nelectrons = ak.sum(
(events.Electron.pt > 10.)
& (abs(events.Electron.eta) < 2.5)
& (events.Electron.cutBased >= events.Electron.VETO),
axis=1,
)
nmuons = ak.sum(
(events.Muon.pt > 10)
& (abs(events.Muon.eta) < 2.4)
& (events.Muon.pfRelIso04_all < 0.25)
& (events.Muon.looseId),
axis=1,
)
ntaus = ak.sum(
(events.Tau.pt > 20.)
& (events.Tau.idDecayMode)
& (events.Tau.rawIso < 5)
& (abs(events.Tau.eta) < 2.3)
& (events.Tau.idMVAoldDM2017v1 >= 16),
axis=1,
)
cuts = {
"CR1_muon_trigger":
trigger_muon,
"CR1_jet_pt": (candidatejet.pt > 525),
"CR1_jet_eta": (abs(candidatejet.eta) < 2.5),
"CR1_jet_msd": (candidatejet.msdcorr > 40),
"CR1_jet_rho":
((candidatejet.qcdrho > -5.5) & (candidatejet.qcdrho < -2.)),
"CR1_mu_pt":
ak.any(candidatemuon.pt > 55, axis=1),
"CR1_mu_eta":
ak.any(abs(candidatemuon.eta) < 2.1, axis=1),
"CR1_mu_IDLoose":
ak.any(candidatemuon.looseId, axis=1),
"CR1_mu_isolationTight":
ak.any(candidatemuon.pfRelIso04_all < 0.15, axis=1),
"CR1_muonDphiAK8":
ak.any(
abs(candidatemuon.delta_phi(candidatejet)) > 2 * np.pi / 3,
axis=1),
"CR1_ak4btagMedium08":
(ak.max(ak4_away.btagCSVV2, axis=1, mask_identity=False) >
BTagEfficiency.btagWPs[self._year]['medium']
), #(ak4_away.btagCSVV2.max() > 0.8838),
"CR1_noelectron": (nelectrons == 0),
"CR1_onemuon": (nmuons == 1),
"CR1_notau": (ntaus == 0),
}
for name, cut in cuts.items():
selection.add(name, cut)
if isRealData:
genflavor = 0 #candidatejet.pt.zeros_like().pad(1, clip=True).fillna(-1).flatten()
if not isRealData:
weights_muonCR.add('genweight', events.genWeight)
#add_pileup_weight(weights_muonCR, events.Pileup.nPU, self._year, dataset)
#add_singleMuTriggerWeight(weights, candidatejet.msdcorr, candidatejet.pt, self._year)
bosons = getBosons(events.GenPart)
genBosonPt = ak.fill_none(ak.firsts(bosons.pt), 0)
#add_VJets_NLOkFactor(weights, genBosonPt, self._year, dataset)
#genflavor = matchedBosonFlavor(candidatejet, bosons).pad(1, clip=True).fillna(-1).flatten()
allcuts_ttbar_muoncontrol = set()
output['cutflow_muonCR'][dataset]['none'] += float(
weights_muonCR.weight().sum())
for cut in cuts:
allcuts_ttbar_muoncontrol.add(cut)
output['cutflow_muonCR'][dataset][cut] += float(
weights_muonCR.weight()[selection.all(
*allcuts_ttbar_muoncontrol)].sum())
fill('muonCR', cuts.keys())
#######################
if 'VtaggingCR' in self._region:
if isRealData:
trigger_muon = np.zeros(len(events), dtype='bool')
for t in self._muontriggers[self._year]:
trigger_muon = trigger_muon | events.HLT[t]
else:
trigger_muon = np.ones(len(events), dtype='bool')
candidatejet = ak.firsts(fatjets)
candidatemuon = ak.firsts(events.Muon)
jets = events.Jet[((events.Jet.pt > 30.)
& (abs(events.Jet.eta) < 2.4))][:, :4]
dr_ak4_ak8 = jets.delta_r(candidatejet)
dr_ak4_muon = jets.delta_r(candidatemuon)
ak4_away = jets[(dr_ak4_ak8 > 0.8)] # & (dr_ak4_muon > 0.4)]
mu_p4 = ak.zip(
{
"pt": ak.fill_none(candidatemuon.pt, 0),
"eta": ak.fill_none(candidatemuon.eta, 0),
"phi": ak.fill_none(candidatemuon.phi, 0),
"mass": ak.fill_none(candidatemuon.mass, 0),
},
with_name="PtEtaPhiMLorentzVector")
met_p4 = ak.zip(
{
"pt": ak.from_iter([[v] for v in events.MET.pt]),
"eta": ak.from_iter([[v] for v in np.zeros(len(events))]),
"phi": ak.from_iter([[v] for v in events.MET.phi]),
"mass": ak.from_iter([[v] for v in np.zeros(len(events))]),
},
with_name="PtEtaPhiMLorentzVector")
Wleptoniccandidate = mu_p4 + met_p4
nelectrons = ak.sum(
((events.Electron.pt > 10.)
& (abs(events.Electron.eta) < 2.5)
& (events.Electron.cutBased >= events.Electron.VETO)),
axis=1,
)
n_tight_muon = ak.sum(
((events.Muon.pt > 53)
& (abs(events.Muon.eta) < 2.1)
& (events.Muon.tightId)),
axis=1,
)
n_loose_muon = ak.sum(
((events.Muon.pt > 20)
& (events.Muon.looseId)
& (abs(events.Muon.eta) < 2.4)),
axis=1,
)
ntaus = ak.sum(
((events.Tau.pt > 20.)
& (events.Tau.idDecayMode)
& (events.Tau.rawIso < 5)
& (abs(events.Tau.eta) < 2.3)
& (events.Tau.idMVAoldDM2017v1 >= 16)),
axis=1,
)
cuts = {
"CR2_muon_trigger":
trigger_muon,
"CR2_jet_pt": (candidatejet.pt > 200),
"CR2_jet_eta": (abs(candidatejet.eta) < 2.5),
"CR2_jet_msd": (candidatejet.msdcorr > 40),
"CR2_mu_pt":
candidatemuon.pt > 53,
"CR2_mu_eta": (abs(candidatemuon.eta) < 2.1),
"CR2_mu_IDTight":
candidatemuon.tightId,
"CR2_mu_isolationTight": (candidatemuon.pfRelIso04_all < 0.15),
"CR2_muonDphiAK8":
abs(candidatemuon.delta_phi(candidatejet)) > 2 * np.pi / 3,
"CR2_ak4btagMedium08":
(ak.max(ak4_away.btagCSVV2, axis=1, mask_identity=False) >
BTagEfficiency.btagWPs[self._year]['medium']),
"CR2_leptonicW":
ak.flatten(Wleptoniccandidate.pt > 200),
"CR2_MET": (events.MET.pt > 40.),
"CR2_noelectron": (nelectrons == 0),
"CR2_one_tightMuon": (n_tight_muon == 1),
"CR2_one_looseMuon": (n_loose_muon == 1),
#"CR2_notau" : (ntaus==0),
}
for name, cut in cuts.items():
print(name, cut)
selection.add(name, cut)
#weights.add('metfilter', events.Flag.METFilters)
if isRealData:
genflavor = 0 #candidatejet.pt.zeros_like().pad(1, clip=True).fillna(-1).flatten()
if not isRealData:
weights_VtaggingCR.add('genweight', events.genWeight)
#add_pileup_weight(weights_VtaggingCR, events.Pileup.nPU, self._year, dataset)
#add_singleMuTriggerWeight(weights, abs(candidatemuon.eta), candidatemuon.pt, self._year)
bosons = getBosons(events.GenPart)
genBosonPt = ak.fill_none(ak.firsts(bosons.pt), 0)
#add_VJets_NLOkFactor(weights, genBosonPt, self._year, dataset)
#genflavor = matchedBosonFlavor(candidatejet, bosons).pad(1, clip=True).fillna(-1).flatten()
#b-tag weights
allcuts_vselection = set()
output['cutflow_VtaggingCR'][dataset]['none'] += float(
weights_VtaggingCR.weight().sum())
for cut in cuts:
allcuts_vselection.add(cut)
output['cutflow_VtaggingCR'][dataset][cut] += float(
weights_VtaggingCR.weight()[selection.all(
*allcuts_vselection)].sum())
fill('VtaggingCR', cuts.keys())
return output
