def get_taus(self, apply_selection=True):
events = self.get_events()
taus_dict = {"e": events.tau_e, "pt": events.tau_pt, "eta": events.tau_eta, "phi": events.tau_phi,
"looseIsoAbs": events.tau_looseIsoAbs, "looseIsoRel": events.tau_looseIsoRel,
"mediumIsoAbs": events.tau_mediumIsoAbs, "mediumIsoRel": events.tau_mediumIsoRel,
"tightIsoAbs": events.tau_tightIsoAbs, "tightIsoRel": events.tau_tightIsoRel,
"gen_e": events.gen_tau_e, "gen_pt": events.gen_tau_pt, "gen_eta": events.gen_tau_eta,
"gen_phi": events.gen_tau_phi,
"lepton_gen_match": events.lepton_gen_match, "deepTau_VSjet": events.deepTau_VSjet}
if self.is_old:
taus = ak.zip(taus_dict)
index = ak.argsort(taus.pt, ascending=False)
taus = taus[index]
tau_1, tau_2 = ak.unzip(ak.combinations(taus, 2, axis=1))
else:
taus_dict["vz"] = events.tau_vz
taus = ak.zip(taus_dict)
index = ak.argsort(taus.pt, ascending=False)
taus = taus[index]
tau_1, tau_2 = ak.unzip(ak.combinations(taus, 2, axis=1))
if apply_selection:
L1taus = ak.zip({"e": events.L1tau_e, "pt": events.L1tau_pt, "eta": events.L1tau_eta,
"phi": events.L1tau_phi})
# apply L1seed correction in case Pt28 and Pt30 seeds are considered
L1taus, taus = L1seed_correction(L1taus, taus)
# match taus with L1 taus
taus = L1THLTTauMatching(L1taus, taus)
tau_1, tau_2 = HLTJetPairDzMatchFilter(taus)
# Return all possible pairs of tau which pass preselection
return tau_1, tau_2
def test_axis0():
array = ak.Array([0.0, 1.1, 2.2, 3.3])
assert ak.to_list(ak.combinations(array, 2, replacement=False,
axis=0)) == [
(0.0, 1.1),
(0.0, 2.2),
(0.0, 3.3),
(1.1, 2.2),
(1.1, 3.3),
(2.2, 3.3),
]
assert ak.to_list(
ak.combinations(array, 2, replacement=False, axis=0,
fields=["x", "y"])) == [
{
"x": 0.0,
"y": 1.1
},
{
"x": 0.0,
"y": 2.2
},
{
"x": 0.0,
"y": 3.3
},
{
"x": 1.1,
"y": 2.2
},
{
"x": 1.1,
"y": 3.3
},
{
"x": 2.2,
"y": 3.3
},
]
assert (ak.combinations(array,
2,
replacement=False,
axis=0,
parameters={
"some": "param"
}).layout.parameters["some"] == "param")
assert ak.to_list(ak.combinations(array, 3, replacement=False,
axis=0)) == [
(0.0, 1.1, 2.2),
(0.0, 1.1, 3.3),
(0.0, 2.2, 3.3),
(1.1, 2.2, 3.3),
]
def process(self, events):
jets = ak.zip(
{k: getattr(events.Jet, k)
for k in ["x", "y", "z", "t", "btag"]},
with_name="LorentzVector",
behavior=events.Jet.behavior,
)
trijet = ak.combinations(jets, 3, fields=["j1", "j2", "j3"])
trijet["p4"] = trijet.j1 + trijet.j2 + trijet.j3
trijet = ak.flatten(trijet[ak.singletons(
ak.argmin(abs(trijet.p4.mass - 172.5), axis=1))])
maxBtag = np.maximum(
trijet.j1.btag,
np.maximum(
trijet.j2.btag,
trijet.j3.btag,
),
)
return {
"trijetpt":
hist.Hist.new.Reg(100,
0,
200,
name="pt3j",
label="Trijet $p_{T}$ [GeV]").Double().fill(
trijet.p4.pt),
"maxbtag":
hist.Hist.new.Reg(
100, 0, 1, name="btag",
label="Max jet b-tag score").Double().fill(maxBtag),
}
def test_IndexedArray():
array = ak.Array([
[0.0, 1.1, 2.2, 3.3],
[],
[4.4, 5.5, 6.6],
None,
[7.7],
None,
[8.8, 9.9, 10.0, 11.1, 12.2],
])
assert ak.to_list(ak.combinations(array, 2, replacement=False)) == [
[(0.0, 1.1), (0.0, 2.2), (0.0, 3.3), (1.1, 2.2), (1.1, 3.3),
(2.2, 3.3)],
[],
[(4.4, 5.5), (4.4, 6.6), (5.5, 6.6)],
None,
[],
None,
[
(8.8, 9.9),
(8.8, 10.0),
(8.8, 11.1),
(8.8, 12.2),
(9.9, 10.0),
(9.9, 11.1),
(9.9, 12.2),
(10.0, 11.1),
(10.0, 12.2),
(11.1, 12.2),
],
]
def process(self, df):
ak.behavior.update(vector.behavior)
output = self.accumulator.identity()
dataset = df.metadata["dataset"]
print(df.metadata)
if "checkusermeta" in df.metadata:
metaname, metavalue = self.expected_usermeta[dataset]
assert metavalue == df.metadata[metaname]
muon = ak.zip(
{
"pt": df.Muon_pt,
"eta": df.Muon_eta,
"phi": df.Muon_phi,
"mass": df.Muon_mass,
},
with_name="PtEtaPhiMLorentzVector",
)
dimuon = ak.combinations(muon, 2)
dimuon = dimuon["0"] + dimuon["1"]
output["pt"].fill(dataset=dataset, pt=ak.flatten(muon.pt))
output["mass"].fill(dataset=dataset, mass=ak.flatten(dimuon.mass))
output["cutflow"]["%s_pt" % dataset] += np.sum(ak.num(muon))
output["cutflow"]["%s_mass" % dataset] += np.sum(ak.num(dimuon))
return output
def choose3(first, n=3):
tmp = ak.combinations(first, n)
combs = (tmp['0'] + tmp['1'] + tmp['2'])
combs['0'] = tmp['0']
combs['1'] = tmp['1']
combs['2'] = tmp['2']
return combs
def process(self, events):
output = self.accumulator.identity()
dataset = events.metadata["dataset"]
print(events.metadata)
if "checkusermeta" in events.metadata:
metaname, metavalue = self.expected_usermeta[dataset]
assert metavalue == events.metadata[metaname]
mapping = events.behavior["__events_factory__"]._mapping
muon_pt = events.Muon.pt
if isinstance(mapping, nanoevents.mapping.CachedMapping):
keys_in_cache = list(mapping.cache.cache.keys())
has_canaries = [
canary in keys_in_cache for canary in self._canaries
]
if has_canaries:
try:
from distributed import get_worker
worker = get_worker()
output["worker"].add(worker.name)
except ValueError:
pass
dimuon = ak.combinations(events.Muon, 2)
dimuon = dimuon["0"] + dimuon["1"]
output["pt"].fill(dataset=dataset, pt=ak.flatten(muon_pt))
output["mass"].fill(dataset=dataset, mass=ak.flatten(dimuon.mass))
output["cutflow"]["%s_pt" % dataset] += sum(ak.num(events.Muon))
output["cutflow"]["%s_mass" % dataset] += sum(ak.num(dimuon))
return output
def test_axis2():
array = ak.Array([
[[0.0, 1.1, 2.2, 3.3], [], [4.4, 5.5, 6.6]],
[],
[[7.7], [8.8, 9.9, 10.0, 11.1, 12.2]],
])
assert ak.to_list(ak.combinations(array, 2, axis=1,
replacement=False)) == [
[
([0.0, 1.1, 2.2, 3.3], []),
([0.0, 1.1, 2.2,
3.3], [4.4, 5.5, 6.6]),
([], [4.4, 5.5, 6.6]),
],
[],
[([7.7],
[8.8, 9.9, 10.0, 11.1, 12.2])],
]
assert ak.to_list(ak.combinations(array, 2, axis=2,
replacement=False)) == [
[
[(0.0, 1.1), (0.0, 2.2),
(0.0, 3.3), (1.1, 2.2),
(1.1, 3.3), (2.2, 3.3)],
[],
[(4.4, 5.5), (4.4, 6.6),
(5.5, 6.6)],
],
[],
[
[],
[
(8.8, 9.9),
(8.8, 10.0),
(8.8, 11.1),
(8.8, 12.2),
(9.9, 10.0),
(9.9, 11.1),
(9.9, 12.2),
(10.0, 11.1),
(10.0, 12.2),
(11.1, 12.2),
],
],
]
def get_gen_taus(self):
events = self.get_gen_events()
gen_taus = ak.zip({"gen_e": events.gen_tau_e, "gen_pt": events.gen_tau_pt, "gen_eta": events.gen_tau_eta,
"gen_phi": events.gen_tau_phi,
"lepton_gen_match": events.lepton_gen_match})
index = ak.argsort(gen_taus.gen_pt, ascending=False)
gen_taus = gen_taus[index]
gen_tau_1, gen_tau_2 = ak.unzip(ak.combinations(gen_taus, 2, axis=1))
return gen_tau_1, gen_tau_2
def test():
one = ak.Array([1, 2, 3, 4])
two = ak.Array(["aa", "bb", "cc", "dd"])
with pytest.raises(ValueError):
ak.cartesian([one, two], axis=1)
with pytest.raises(ValueError):
ak.combinations(two, 2, axis=1)
two = ak.Array([["aa", "bb"], ["cc"], [], ["dd"]])
assert ak.to_list(ak.cartesian([one, two], axis=1)) == [
[(1, "aa"), (1, "bb")],
[(2, "cc")],
[],
[(4, "dd")],
]
assert ak.to_list(ak.combinations(two, 2, axis=1)) == [[("aa", "bb")], [], [], []]
with pytest.raises(ValueError):
ak.cartesian([one, two], axis=2)
with pytest.raises(ValueError):
ak.combinations(two, 2, axis=2)
def get_dat_photons(self):
arrs = self.tree.arrays(
['pt', 'theta', 'phi', 'mass'],
cut='(nt>=2)&(nks>0)&(phen>0)',
aliases={
'pt': 'phen*sin(phth)',
'theta': 'phth',
'phi': 'phphi',
'mass': '0*phen'
})
vecs = vector.Array(arrs)
df = ak.to_pandas(ak.combinations(vecs.px, 2))
df = df.rename({'0': 'px0', '1': 'px1'}, axis=1)
df_len = len(df)
df = df.join(ak.to_pandas(ak.combinations(vecs.py, 2)))
assert df_len == len(df)
df_len = len(df)
df = df.rename({'0': 'py0', '1': 'py1'}, axis=1)
df = df.join(ak.to_pandas(ak.combinations(vecs.pz, 2)))
assert df_len == len(df)
df_len = len(df)
df = df.rename({'0': 'pz0', '1': 'pz1'}, axis=1)
df = df.join(ak.to_pandas(ak.combinations(vecs.E, 2)))
assert df_len == len(df)
df_len = len(df)
df = df.rename({'0': 'E0', '1': 'E1'}, axis=1)
for coord in ('x', 'y', 'z'):
df[f'P{coord}'] = df[f'p{coord}0'] + df[f'p{coord}1']
df['P'] = np.sqrt(df['Px']**2 + df['Py']**2 + df['Pz']**2)
df['E'] = df['E0'] + df['E1']
M2 = df['E']**2 - df['P']**2
df['M'] = np.where(M2 > 0, np.sqrt(np.abs(M2)), -np.sqrt(np.abs(M2)))
return df
def calculate_selection(self, syst_tag, events):
"""
"""
electron_cut = lepton_selections.select_electrons(
electrons=events.Electron,
options=self.options["electrons"],
clean={},
name="ele",
tagger=self)
electrons = awkward_utils.add_field(events=events,
name="ele",
data=events.Electron[electron_cut])
electrons = awkward.Array(electrons, with_name="Momentum4D")
ee_pairs = awkward.combinations(electrons,
2,
fields=["LeadEle", "SubleadEle"])
ee_pairs["ZCand"] = ee_pairs.LeadEle + ee_pairs.SubleadEle
ee_pairs[("ZCand", "mass")] = ee_pairs.ZCand.mass
ee_pairs[("ZCand", "pt")] = ee_pairs.ZCand.pt
ee_pairs[("ZCand", "phi")] = ee_pairs.ZCand.phi
ee_pairs[("ZCand", "eta")] = ee_pairs.ZCand.eta
events["ZCand"] = ee_pairs.ZCand
os_cut = ee_pairs.LeadEle.charge * ee_pairs.SubleadEle.charge == -1
mass_cut = (ee_pairs.ZCand.mass > self.options["z_window"][0]) & (
ee_pairs.ZCand.mass < self.options["z_window"][1])
id_cut = (ee_pairs.LeadEle.mvaFall17V2Iso_WP80
== True) | (ee_pairs.SubleadEle.mvaFall17V2Iso_WP80 == True)
pair_cut = os_cut & mass_cut & id_cut
evt_os_cut = awkward.num(ee_pairs[os_cut]) == 1
evt_mass_cut = awkward.num(ee_pairs[mass_cut]) == 1
evt_id_cut = awkward.num(ee_pairs[id_cut]) == 1
met_cut = events.MET_pt <= self.options["met"]
presel_cut = (awkward.num(ee_pairs[pair_cut]) == 1) & met_cut
self.register_cuts(
names=["met cut", "os cut", "mass cut", "id cut", "all"],
results=[
met_cut, evt_os_cut, evt_mass_cut, evt_id_cut, presel_cut
])
return presel_cut, events
def HLTJetPairDzMatchFilter(L2taus):
jetMinPt = 20.0
jetMaxEta = 2.1
jetMinDR = 0.5
jetMaxDZ = 0.2
L2taus = L2taus[reco_tau_selection(L2taus,
minPt=jetMinPt,
maxEta=jetMaxEta)]
# Take all possible pairs of L2 taus
L2tau_1, L2tau_2 = ak.unzip(ak.combinations(L2taus, 2, axis=1))
dr2 = delta_r2(L2tau_1, L2tau_2)
dz = delta_z(L2tau_1, L2tau_2)
pair_mask = (dr2 >= jetMinDR * jetMinDR) & (abs(dz) <= jetMaxDZ)
# ev_mask = ak.sum(pair_mask, axis=1) > 0
return L2tau_1[pair_mask], L2tau_2[pair_mask]
def process(self, events):
mupair = ak.combinations(events.Muon, 2)
with np.errstate(invalid="ignore"):
pairmass = (mupair.slot0 + mupair.slot1).mass
goodevent = ak.any(
(pairmass > 60)
& (pairmass < 120)
& (mupair.slot0.charge == -mupair.slot1.charge),
axis=1,
)
return (hist.Hist.new.Reg(100,
0,
200,
name="met",
label="$E_{T}^{miss}$ [GeV]").Double().fill(
events[goodevent].MET.pt))
def process(self, df):
ak.behavior.update(vector.behavior)
output = self.accumulator.identity()
dataset = df.metadata["dataset"]
muon = ak.zip(
{
'pt': df.Muon_pt,
'eta': df.Muon_eta,
'phi': df.Muon_phi,
'mass': df.Muon_mass
},
with_name="PtEtaPhiMLorentzVector")
dimuon = ak.combinations(muon, 2)
dimuon = dimuon["0"] + dimuon["1"]
output["pt"].fill(dataset=dataset, pt=ak.flatten(muon.pt))
output["mass"].fill(dataset=dataset, mass=ak.flatten(dimuon.mass))
output["cutflow"]["%s_pt" % dataset] += np.sum(ak.num(muon))
output["cutflow"]["%s_mass" % dataset] += np.sum(ak.num(dimuon))
return output
def test_ByteMaskedArray_combinations():
content = ak.from_iter(
[[[0, 1, 2], [], [3, 4]], [], [[5]], [[6, 7, 8, 9]], [[], [10, 11, 12]]
],
highlevel=False,
)
mask = ak.layout.Index8(np.array([0, 0, 1, 1, 0], dtype=np.int8))
array = ak.Array(ak.layout.ByteMaskedArray(mask, content,
valid_when=False))
assert ak.to_list(array) == [
[[0, 1, 2], [], [3, 4]],
[],
None,
None,
[[], [10, 11, 12]],
]
assert ak.to_list(ak.combinations(array, 2, axis=0)) == [
([[0, 1, 2], [], [3, 4]], []),
([[0, 1, 2], [], [3, 4]], None),
([[0, 1, 2], [], [3, 4]], None),
([[0, 1, 2], [], [3, 4]], [[], [10, 11, 12]]),
([], None),
([], None),
([], [[], [10, 11, 12]]),
(None, None),
(None, [[], [10, 11, 12]]),
(None, [[], [10, 11, 12]]),
]
assert ak.to_list(ak.combinations(array, 2, axis=-3)) == [
([[0, 1, 2], [], [3, 4]], []),
([[0, 1, 2], [], [3, 4]], None),
([[0, 1, 2], [], [3, 4]], None),
([[0, 1, 2], [], [3, 4]], [[], [10, 11, 12]]),
([], None),
([], None),
([], [[], [10, 11, 12]]),
(None, None),
(None, [[], [10, 11, 12]]),
(None, [[], [10, 11, 12]]),
]
assert ak.to_list(ak.combinations(array, 2, axis=1)) == [
[([0, 1, 2], []), ([0, 1, 2], [3, 4]), ([], [3, 4])],
[],
None,
None,
[([], [10, 11, 12])],
]
assert ak.to_list(ak.combinations(array, 2, axis=-2)) == [
[([0, 1, 2], []), ([0, 1, 2], [3, 4]), ([], [3, 4])],
[],
None,
None,
[([], [10, 11, 12])],
]
assert ak.to_list(ak.combinations(array, 2, axis=2)) == [
[[(0, 1), (0, 2), (1, 2)], [], [(3, 4)]],
[],
None,
None,
[[], [(10, 11), (10, 12), (11, 12)]],
]
assert ak.to_list(ak.combinations(array, 2, axis=-1)) == [
[[(0, 1), (0, 2), (1, 2)], [], [(3, 4)]],
[],
None,
None,
[[], [(10, 11), (10, 12), (11, 12)]],
]
def test_combinations():
assert ak.combinations(empty, 2, axis=0).tolist() == []
assert ak.combinations(empty, 2, axis=1).tolist() == []
assert ak.combinations(empty, 2, axis=2).tolist() == []
def process(self, events):
output = self._accumulator.identity()
jets=events.Jet
jetSel = (jets.pt>30) & (abs(jets.eta)<2.4)
tightJet = jets[jetSel]
bJet = tightJet[tightJet.btagDeepFlavB > 0.642]
muons = events.Muon
muonSel = (muons.pt>30) & (abs(muons.eta)<2.4)
tightMuon = muons[muonSel]
ele = events.Electron
eleSel = (ele.pt>35)&(abs(ele.eta)<2.4)
tightEle = ele[eleSel]
eventSel = (((ak.num(tightMuon)==1) | (ak.num(tightEle)==1)) &
(ak.num(tightJet)>= 3) & (ak.num(bJet)>=1)
)
final = events[eventSel]
#####GENPART MATCHING ######
genPart = final.GenPart
tops = genPart[abs(genPart.pdgId)==6]
#The isLastCopy Flag filters out copy Genparticles:
tops = tops[tops.hasFlags('isLastCopy')]
tDecay = tops.distinctChildren
tDecay = tDecay[tDecay.hasFlags('isLastCopy')]
t_Events=tDecay[abs(tDecay.pdgId)==5]
W = tDecay[abs(tDecay.pdgId)==24]
W = W[W.hasFlags('isLastCopy')]
WDecay = W.distinctChildren
WDecay = WDecay[WDecay.hasFlags('isLastCopy')]
#t_events is the lone bottom, W_events is the -> two jets
#select the hadronically decaying W
W_Events=ak.flatten(WDecay[ak.all(abs(WDecay.pdgId)<=8,axis=-1)],axis=3)
#print(qqb)
#HadW is mask for Quark deacying W boson
hadW = ak.num(W_Events,axis=2)==2
#filters out t_events that have a hadronically decayign W Boson
hadB = t_Events[hadW]
hadB = ak.flatten(hadB,axis=2)
W_quarks = W_Events[hadW]
W_quarks = ak.flatten(W_quarks,axis=2)
#concatentating these two arrays make an array of events with the correctly decaying GenParticles.
qqb = ak.concatenate([hadB,W_quarks],axis=1)
#####GEN JET MATCHING ######
final=final[(ak.count(qqb.pdgId,axis=1)==3)]
finaljets=final.Jet
qqb=qqb[(ak.count(qqb.pdgId,axis=1)==3)]
#Implementing Tight Jet Cuts on Training Data
finaljetSel=(abs(finaljets.eta)<2.4)&(finaljets.pt>30)
finalJets=finaljets[finaljetSel]
#Match Gen part to gen jet
matchedGenJets=qqb.nearest(final.GenJet)
#match gen to reco
matchedJets=matchedGenJets.nearest(finalJets)
### VALIDATION ###
test=matchedJets.genJetIdx
combs=ak.combinations(finalJets,3,replacement=False)
t1=(combs['0'].genJetIdx==test[:,0])|(combs['0'].genJetIdx==test[:,1])|(combs['0'].genJetIdx==test[:,2])
t2=(combs['1'].genJetIdx==test[:,0])|(combs['1'].genJetIdx==test[:,1])|(combs['1'].genJetIdx==test[:,2])
t3=(combs['2'].genJetIdx==test[:,0])|(combs['2'].genJetIdx==test[:,1])|(combs['2'].genJetIdx==test[:,2])
t=t1&t2&t3
trutharray=ak.flatten(t)
jetcombos=ak.flatten(combs)
j1,j2,j3=ak.unzip(jetcombos)
output["dR12"]+=processor.column_accumulator(ak.to_numpy(j1.delta_r(j2)))
output["dR13"]+=processor.column_accumulator(ak.to_numpy(j1.delta_r(j3)))
output["dR23"]+=processor.column_accumulator(ak.to_numpy(j2.delta_r(j3)))
output["j1btag"]+=processor.column_accumulator(ak.to_numpy(j1.btagCSVV2))
output["j2btag"]+=processor.column_accumulator(ak.to_numpy(j1.btagCSVV2))
output["j3btag"]+=processor.column_accumulator(ak.to_numpy(j1.btagCSVV2))
output["j1area"]+=processor.column_accumulator(ak.to_numpy(j1.area))
output["j2area"]+=processor.column_accumulator(ak.to_numpy(j2.area))
output["j3area"]+=processor.column_accumulator(ak.to_numpy(j3.area))
output["j12deta"]+=processor.column_accumulator(ak.to_numpy(j1.eta-j2.eta))
output["j23deta"]+=processor.column_accumulator(ak.to_numpy(j2.eta-j3.eta))
output["j13deta"]+=processor.column_accumulator(ak.to_numpy(j1.eta-j3.eta))
output["j12dphi"]+=processor.column_accumulator(ak.to_numpy(j1.phi-j2.phi))
output["j23dphi"]+=processor.column_accumulator(ak.to_numpy(j2.phi-j3.phi))
output["j13dphi"]+=processor.column_accumulator(ak.to_numpy(j1.phi-j3.phi))
output["j1j2mass"]+=processor.column_accumulator(ak.to_numpy(j1.mass+j2.mass))
output["j2j3mass"]+=processor.column_accumulator(ak.to_numpy(j2.mass+j3.mass))
output["j1j3mass"]+=processor.column_accumulator(ak.to_numpy(j1.mass+j3.mass))
output["j1pt"]+=processor.column_accumulator(ak.to_numpy(j1.pt))
output["j1phi"]+=processor.column_accumulator(ak.to_numpy(j1.phi))
output["j1eta"]+=processor.column_accumulator(ak.to_numpy(abs(j1.eta)))
output["j1mass"]+=processor.column_accumulator(ak.to_numpy(j1.mass))
output["j2pt"]+=processor.column_accumulator(ak.to_numpy(j2.pt))
output["j2phi"]+=processor.column_accumulator(ak.to_numpy(j2.phi))
output["j2eta"]+=processor.column_accumulator(ak.to_numpy(abs(j2.eta)))
output["j2mass"]+=processor.column_accumulator(ak.to_numpy(j2.mass))
output["j3pt"]+=processor.column_accumulator(ak.to_numpy(j3.pt))
output["j3phi"]+=processor.column_accumulator(ak.to_numpy(j3.phi))
output["j3eta"]+=processor.column_accumulator(ak.to_numpy(abs(j3.eta)))
output["j3mass"]+=processor.column_accumulator(ak.to_numpy(j3.mass))
output["event"]+=processor.column_accumulator(ak.to_numpy(ak.flatten(ak.broadcast_arrays(final.event,combs['0'].pt)[0])))
output["truth"]+=processor.column_accumulator(ak.to_numpy(trutharray).astype(int))
return output
def process(self, events):
# Dataset parameters
dataset = events.metadata['dataset']
histAxisName = self._samples[dataset]['histAxisName']
year = self._samples[dataset]['year']
xsec = self._samples[dataset]['xsec']
sow = self._samples[dataset]['nSumOfWeights']
isData = self._samples[dataset]['isData']
datasets = [
'SingleMuon', 'SingleElectron', 'EGamma', 'MuonEG', 'DoubleMuon',
'DoubleElectron'
]
for d in datasets:
if d in dataset: dataset = dataset.split('_')[0]
# Initialize objects
met = events.MET
e = events.Electron
mu = events.Muon
tau = events.Tau
j = events.Jet
# Muon selection
mu['isPres'] = isPresMuon(mu.dxy, mu.dz, mu.sip3d, mu.looseId)
mu['isTight'] = isTightMuon(mu.pt,
mu.eta,
mu.dxy,
mu.dz,
mu.pfRelIso03_all,
mu.sip3d,
mu.mvaTTH,
mu.mediumPromptId,
mu.tightCharge,
mu.looseId,
minpt=10)
mu['isGood'] = mu['isPres'] & mu['isTight']
leading_mu = mu[ak.argmax(mu.pt, axis=-1, keepdims=True)]
leading_mu = leading_mu[leading_mu.isGood]
mu = mu[mu.isGood]
mu_pres = mu[mu.isPres]
# Electron selection
e['isPres'] = isPresElec(e.pt,
e.eta,
e.dxy,
e.dz,
e.miniPFRelIso_all,
e.sip3d,
e.lostHits,
minpt=15)
e['isTight'] = isTightElec(e.pt,
e.eta,
e.dxy,
e.dz,
e.miniPFRelIso_all,
e.sip3d,
e.mvaTTH,
e.mvaFall17V2Iso,
e.lostHits,
e.convVeto,
e.tightCharge,
e.sieie,
e.hoe,
e.eInvMinusPInv,
minpt=15)
e['isClean'] = isClean(e, mu, drmin=0.05)
e['isGood'] = e['isPres'] & e['isTight'] & e['isClean']
leading_e = e[ak.argmax(e.pt, axis=-1, keepdims=True)]
leading_e = leading_e[leading_e.isGood]
e = e[e.isGood]
e_pres = e[e.isPres & e.isClean]
# Tau selection
tau['isPres'] = isPresTau(tau.pt,
tau.eta,
tau.dxy,
tau.dz,
tau.leadTkPtOverTauPt,
tau.idAntiMu,
tau.idAntiEle,
tau.rawIso,
tau.idDecayModeNewDMs,
minpt=20)
tau['isClean'] = isClean(tau, e_pres, drmin=0.4) & isClean(
tau, mu_pres, drmin=0.4)
tau['isGood'] = tau['isPres'] # & tau['isClean'], for the moment
tau = tau[tau.isGood]
nElec = ak.num(e)
nMuon = ak.num(mu)
nTau = ak.num(tau)
twoLeps = (nElec + nMuon) == 2
threeLeps = (nElec + nMuon) == 3
twoElec = (nElec == 2)
twoMuon = (nMuon == 2)
e0 = e[ak.argmax(e.pt, axis=-1, keepdims=True)]
m0 = mu[ak.argmax(mu.pt, axis=-1, keepdims=True)]
# Attach the lepton SFs to the electron and muons collections
AttachElectronSF(e, year=year)
AttachMuonSF(mu, year=year)
# Create a lepton (muon+electron) collection and calculate a per event lepton SF
leps = ak.concatenate([e, mu], axis=-1)
events['lepSF_nom'] = ak.prod(leps.sf_nom, axis=-1)
events['lepSF_hi'] = ak.prod(leps.sf_hi, axis=-1)
events['lepSF_lo'] = ak.prod(leps.sf_lo, axis=-1)
# Jet selection
jetptname = 'pt_nom' if hasattr(j, 'pt_nom') else 'pt'
### Jet energy corrections
if not isData:
j["pt_raw"] = (1 - j.rawFactor) * j.pt
j["mass_raw"] = (1 - j.rawFactor) * j.mass
j["pt_gen"] = ak.values_astype(ak.fill_none(j.matched_gen.pt, 0),
np.float32)
j["rho"] = ak.broadcast_arrays(events.fixedGridRhoFastjetAll,
j.pt)[0]
events_cache = events.caches[0]
corrected_jets = jet_factory.build(j, lazy_cache=events_cache)
#print('jet pt: ',j.pt)
#print('cor pt: ',corrected_jets.pt)
#print('jes up: ',corrected_jets.JES_jes.up.pt)
#print('jes down: ',corrected_jets.JES_jes.down.pt)
#print(ak.fields(corrected_jets))
'''
# SYSTEMATICS
jets = corrected_jets
if(self.jetSyst == 'JERUp'):
jets = corrected_jets.JER.up
elif(self.jetSyst == 'JERDown'):
jets = corrected_jets.JER.down
elif(self.jetSyst == 'JESUp'):
jets = corrected_jets.JES_jes.up
elif(self.jetSyst == 'JESDown'):
jets = corrected_jets.JES_jes.down
'''
j['isGood'] = isTightJet(getattr(j,
jetptname), j.eta, j.jetId, j.neHEF,
j.neEmEF, j.chHEF, j.chEmEF, j.nConstituents)
#j['isgood'] = isGoodJet(j.pt, j.eta, j.jetId)
#j['isclean'] = isClean(j, e, mu)
j['isClean'] = isClean(j, e, drmin=0.4) & isClean(
j, mu, drmin=0.4) # & isClean(j, tau, drmin=0.4)
goodJets = j[(j.isClean) & (j.isGood)]
njets = ak.num(goodJets)
ht = ak.sum(goodJets.pt, axis=-1)
j0 = goodJets[ak.argmax(goodJets.pt, axis=-1, keepdims=True)]
#nbtags = ak.num(goodJets[goodJets.btagDeepFlavB > 0.2770])
# Loose DeepJet WP
if year == 2017: btagwpl = 0.0532 #WP loose
else: btagwpl = 0.0490 #WP loose
isBtagJetsLoose = (goodJets.btagDeepB > btagwpl)
isNotBtagJetsLoose = np.invert(isBtagJetsLoose)
nbtagsl = ak.num(goodJets[isBtagJetsLoose])
# Medium DeepJet WP
if year == 2017: btagwpm = 0.3040 #WP medium
else: btagwpm = 0.2783 #WP medium
isBtagJetsMedium = (goodJets.btagDeepB > btagwpm)
isNotBtagJetsMedium = np.invert(isBtagJetsMedium)
nbtagsm = ak.num(goodJets[isBtagJetsMedium])
# Btag SF following 1a) in https://twiki.cern.ch/twiki/bin/viewauth/CMS/BTagSFMethods
btagSF = np.ones_like(ht)
btagSFUp = np.ones_like(ht)
btagSFDo = np.ones_like(ht)
if not isData:
pt = goodJets.pt
abseta = np.abs(goodJets.eta)
flav = goodJets.hadronFlavour
bJetSF = GetBTagSF(abseta, pt, flav)
bJetSFUp = GetBTagSF(abseta, pt, flav, sys=1)
bJetSFDo = GetBTagSF(abseta, pt, flav, sys=-1)
bJetEff = GetBtagEff(abseta, pt, flav, year)
bJetEff_data = bJetEff * bJetSF
bJetEff_dataUp = bJetEff * bJetSFUp
bJetEff_dataDo = bJetEff * bJetSFDo
pMC = ak.prod(bJetEff[isBtagJetsMedium], axis=-1) * ak.prod(
(1 - bJetEff[isNotBtagJetsMedium]), axis=-1)
pData = ak.prod(bJetEff_data[isBtagJetsMedium], axis=-1) * ak.prod(
(1 - bJetEff_data[isNotBtagJetsMedium]), axis=-1)
pDataUp = ak.prod(
bJetEff_dataUp[isBtagJetsMedium], axis=-1) * ak.prod(
(1 - bJetEff_dataUp[isNotBtagJetsMedium]), axis=-1)
pDataDo = ak.prod(
bJetEff_dataDo[isBtagJetsMedium], axis=-1) * ak.prod(
(1 - bJetEff_dataDo[isNotBtagJetsMedium]), axis=-1)
pMC = ak.where(pMC == 0, 1,
pMC) # removeing zeroes from denominator...
btagSF = pData / pMC
btagSFUp = pDataUp / pMC
btagSFDo = pDataUp / pMC
##################################################################
### 2 same-sign leptons
##################################################################
# emu
singe = e[(nElec == 1) & (nMuon == 1) & (e.pt > -1)]
singm = mu[(nElec == 1) & (nMuon == 1) & (mu.pt > -1)]
em = ak.cartesian({"e": singe, "m": singm})
emSSmask = (em.e.charge * em.m.charge > 0)
emSS = em[emSSmask]
nemSS = len(ak.flatten(emSS))
# ee and mumu
# pt>-1 to preserve jagged dimensions
ee = e[(nElec == 2) & (nMuon == 0) & (e.pt > -1)]
mm = mu[(nElec == 0) & (nMuon == 2) & (mu.pt > -1)]
sumcharge = ak.sum(e.charge, axis=-1) + ak.sum(mu.charge, axis=-1)
eepairs = ak.combinations(ee, 2, fields=["e0", "e1"])
eeSSmask = (eepairs.e0.charge * eepairs.e1.charge > 0)
eeonZmask = (np.abs((eepairs.e0 + eepairs.e1).mass - 91.2) < 10)
eeoffZmask = (eeonZmask == 0)
mmpairs = ak.combinations(mm, 2, fields=["m0", "m1"])
mmSSmask = (mmpairs.m0.charge * mmpairs.m1.charge > 0)
mmonZmask = (np.abs((mmpairs.m0 + mmpairs.m1).mass - 91.2) < 10)
mmoffZmask = (mmonZmask == 0)
eeSSonZ = eepairs[eeSSmask & eeonZmask]
eeSSoffZ = eepairs[eeSSmask & eeoffZmask]
mmSSonZ = mmpairs[mmSSmask & mmonZmask]
mmSSoffZ = mmpairs[mmSSmask & mmoffZmask]
neeSS = len(ak.flatten(eeSSonZ)) + len(ak.flatten(eeSSoffZ))
nmmSS = len(ak.flatten(mmSSonZ)) + len(ak.flatten(mmSSoffZ))
print('Same-sign events [ee, emu, mumu] = [%i, %i, %i]' %
(neeSS, nemSS, nmmSS))
# Cuts
eeSSmask = (ak.num(eeSSmask[eeSSmask]) > 0)
mmSSmask = (ak.num(mmSSmask[mmSSmask]) > 0)
eeonZmask = (ak.num(eeonZmask[eeonZmask]) > 0)
eeoffZmask = (ak.num(eeoffZmask[eeoffZmask]) > 0)
mmonZmask = (ak.num(mmonZmask[mmonZmask]) > 0)
mmoffZmask = (ak.num(mmoffZmask[mmoffZmask]) > 0)
emSSmask = (ak.num(emSSmask[emSSmask]) > 0)
##################################################################
### 3 leptons
##################################################################
# eem
muon_eem = mu[(nElec == 2) & (nMuon == 1) & (mu.pt > -1)]
elec_eem = e[(nElec == 2) & (nMuon == 1) & (e.pt > -1)]
ee_eem = ak.combinations(elec_eem, 2, fields=["e0", "e1"])
ee_eemZmask = (ee_eem.e0.charge * ee_eem.e1.charge < 1) & (np.abs(
(ee_eem.e0 + ee_eem.e1).mass - 91.2) < 10)
ee_eemOffZmask = (ee_eem.e0.charge * ee_eem.e1.charge < 1) & (np.abs(
(ee_eem.e0 + ee_eem.e1).mass - 91.2) > 10)
ee_eemZmask = (ak.num(ee_eemZmask[ee_eemZmask]) > 0)
ee_eemOffZmask = (ak.num(ee_eemOffZmask[ee_eemOffZmask]) > 0)
eepair_eem = (ee_eem.e0 + ee_eem.e1)
trilep_eem = eepair_eem + muon_eem #ak.cartesian({"e0":ee_eem.e0,"e1":ee_eem.e1, "m":muon_eem})
# mme
muon_mme = mu[(nElec == 1) & (nMuon == 2) & (mu.pt > -1)]
elec_mme = e[(nElec == 1) & (nMuon == 2) & (e.pt > -1)]
mm_mme = ak.combinations(muon_mme, 2, fields=["m0", "m1"])
mm_mmeZmask = (mm_mme.m0.charge * mm_mme.m1.charge < 1) & (np.abs(
(mm_mme.m0 + mm_mme.m1).mass - 91.2) < 10)
mm_mmeOffZmask = (mm_mme.m0.charge * mm_mme.m1.charge < 1) & (np.abs(
(mm_mme.m0 + mm_mme.m1).mass - 91.2) > 10)
mm_mmeZmask = (ak.num(mm_mmeZmask[mm_mmeZmask]) > 0)
mm_mmeOffZmask = (ak.num(mm_mmeOffZmask[mm_mmeOffZmask]) > 0)
mmpair_mme = (mm_mme.m0 + mm_mme.m1)
trilep_mme = mmpair_mme + elec_mme
mZ_mme = mmpair_mme.mass
mZ_eem = eepair_eem.mass
m3l_eem = trilep_eem.mass
m3l_mme = trilep_mme.mass
# eee and mmm
eee = e[(nElec == 3) & (nMuon == 0) & (e.pt > -1)]
mmm = mu[(nElec == 0) & (nMuon == 3) & (mu.pt > -1)]
eee_leps = ak.combinations(eee, 3, fields=["e0", "e1", "e2"])
mmm_leps = ak.combinations(mmm, 3, fields=["m0", "m1", "m2"])
ee_pairs = ak.combinations(eee, 2, fields=["e0", "e1"])
mm_pairs = ak.combinations(mmm, 2, fields=["m0", "m1"])
ee_pairs_index = ak.argcombinations(eee, 2, fields=["e0", "e1"])
mm_pairs_index = ak.argcombinations(mmm, 2, fields=["m0", "m1"])
mmSFOS_pairs = mm_pairs[
(np.abs(mm_pairs.m0.pdgId) == np.abs(mm_pairs.m1.pdgId))
& (mm_pairs.m0.charge != mm_pairs.m1.charge)]
offZmask_mm = ak.all(
np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2) > 10.,
axis=1,
keepdims=True) & (ak.num(mmSFOS_pairs) > 0)
onZmask_mm = ak.any(
np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2) < 10.,
axis=1,
keepdims=True)
eeSFOS_pairs = ee_pairs[
(np.abs(ee_pairs.e0.pdgId) == np.abs(ee_pairs.e1.pdgId))
& (ee_pairs.e0.charge != ee_pairs.e1.charge)]
offZmask_ee = ak.all(
np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2) > 10,
axis=1,
keepdims=True) & (ak.num(eeSFOS_pairs) > 0)
onZmask_ee = ak.any(
np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2) < 10,
axis=1,
keepdims=True)
# Create masks **for event selection**
eeeOnZmask = (ak.num(onZmask_ee[onZmask_ee]) > 0)
eeeOffZmask = (ak.num(offZmask_ee[offZmask_ee]) > 0)
mmmOnZmask = (ak.num(onZmask_mm[onZmask_mm]) > 0)
mmmOffZmask = (ak.num(offZmask_mm[offZmask_mm]) > 0)
# Now we need to create masks for the leptons in order to select leptons from the Z boson candidate (in onZ categories)
ZeeMask = ak.argmin(np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass -
91.2),
axis=1,
keepdims=True)
ZmmMask = ak.argmin(np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass -
91.2),
axis=1,
keepdims=True)
Zee = eeSFOS_pairs[ZeeMask]
Zmm = mmSFOS_pairs[ZmmMask]
eZ0 = Zee.e0[ak.num(eeSFOS_pairs) > 0]
eZ1 = Zee.e1[ak.num(eeSFOS_pairs) > 0]
eZ = eZ0 + eZ1
mZ0 = Zmm.m0[ak.num(mmSFOS_pairs) > 0]
mZ1 = Zmm.m1[ak.num(mmSFOS_pairs) > 0]
mZ = mZ0 + mZ1
mZ_eee = eZ.mass
mZ_mmm = mZ.mass
# And for the W boson
ZmmIndices = mm_pairs_index[ZmmMask]
ZeeIndices = ee_pairs_index[ZeeMask]
eW = eee[~ZeeIndices.e0 | ~ZeeIndices.e1]
mW = mmm[~ZmmIndices.m0 | ~ZmmIndices.m1]
triElec = eee_leps.e0 + eee_leps.e1 + eee_leps.e2
triMuon = mmm_leps.m0 + mmm_leps.m1 + mmm_leps.m2
m3l_eee = triElec.mass
m3l_mmm = triMuon.mass
##################################################################
### >=4 leptons
##################################################################
# 4lep cat
is4lmask = ((nElec + nMuon) >= 4)
muon_4l = mu[(is4lmask) & (mu.pt > -1)]
elec_4l = e[(is4lmask) & (e.pt > -1)]
# selecting 4 leading leptons
leptons = ak.concatenate([e, mu], axis=-1)
leptons_sorted = leptons[ak.argsort(leptons.pt,
axis=-1,
ascending=False)]
lep4l = leptons_sorted[:, 0:4]
e4l = lep4l[abs(lep4l.pdgId) == 11]
mu4l = lep4l[abs(lep4l.pdgId) == 13]
nElec4l = ak.num(e4l)
nMuon4l = ak.num(mu4l)
# Triggers
trig_eeSS = passTrigger(events, 'ee', isData, dataset)
trig_mmSS = passTrigger(events, 'mm', isData, dataset)
trig_emSS = passTrigger(events, 'em', isData, dataset)
trig_eee = passTrigger(events, 'eee', isData, dataset)
trig_mmm = passTrigger(events, 'mmm', isData, dataset)
trig_eem = passTrigger(events, 'eem', isData, dataset)
trig_mme = passTrigger(events, 'mme', isData, dataset)
trig_4l = triggerFor4l(events, nMuon, nElec, isData, dataset)
# MET filters
# Weights
genw = np.ones_like(events['event']) if (
isData or len(self._wc_names_lst) > 0) else events['genWeight']
### We need weights for: normalization, lepSF, triggerSF, pileup, btagSF...
weights = {}
for r in [
'all', 'ee', 'mm', 'em', 'eee', 'mmm', 'eem', 'mme', 'eeee',
'eeem', 'eemm', 'mmme', 'mmmm'
]:
# weights[r] = coffea.analysis_tools.Weights(len(events))
weights[r] = coffea.analysis_tools.Weights(len(events),
storeIndividual=True)
if len(self._wc_names_lst) > 0:
sow = np.ones_like(
sow
) # Not valid in nanoAOD for EFT samples, MUST use SumOfEFTweights at analysis level
weights[r].add('norm', genw if isData else (xsec / sow) * genw)
weights[r].add('btagSF', btagSF, btagSFUp, btagSFDo)
weights[r].add('lepSF', events.lepSF_nom, events.lepSF_hi,
events.lepSF_lo)
# Extract the EFT quadratic coefficients and optionally use them to calculate the coefficients on the w**2 quartic function
# eft_coeffs is never Jagged so convert immediately to numpy for ease of use.
eft_coeffs = ak.to_numpy(events['EFTfitCoefficients']) if hasattr(
events, "EFTfitCoefficients") else None
if eft_coeffs is not None:
# Check to see if the ordering of WCs for this sample matches what want
if self._samples[dataset]['WCnames'] != self._wc_names_lst:
eft_coeffs = efth.remap_coeffs(
self._samples[dataset]['WCnames'], self._wc_names_lst,
eft_coeffs)
eft_w2_coeffs = efth.calc_w2_coeffs(eft_coeffs, self._dtype) if (
self._do_errors and eft_coeffs is not None) else None
# Selections and cuts
selections = PackedSelection() #(dtype='uint64')
channels2LSS = ['eeSSonZ', 'eeSSoffZ', 'mmSSonZ', 'mmSSoffZ', 'emSS']
selections.add('eeSSonZ', (eeonZmask) & (eeSSmask) & (trig_eeSS))
selections.add('eeSSoffZ', (eeoffZmask) & (eeSSmask) & (trig_eeSS))
selections.add('mmSSonZ', (mmonZmask) & (mmSSmask) & (trig_mmSS))
selections.add('mmSSoffZ', (mmoffZmask) & (mmSSmask) & (trig_mmSS))
selections.add('emSS', (emSSmask) & (trig_emSS))
channels3L = ['eemSSonZ', 'eemSSoffZ', 'mmeSSonZ', 'mmeSSoffZ']
selections.add('eemSSonZ', (ee_eemZmask) & (trig_eem))
selections.add('eemSSoffZ', (ee_eemOffZmask) & (trig_eem))
selections.add('mmeSSonZ', (mm_mmeZmask) & (trig_mme))
selections.add('mmeSSoffZ', (mm_mmeOffZmask) & (trig_mme))
channels3L += ['eeeSSonZ', 'eeeSSoffZ', 'mmmSSonZ', 'mmmSSoffZ']
selections.add('eeeSSonZ', (eeeOnZmask) & (trig_eee))
selections.add('eeeSSoffZ', (eeeOffZmask) & (trig_eee))
selections.add('mmmSSonZ', (mmmOnZmask) & (trig_mmm))
selections.add('mmmSSoffZ', (mmmOffZmask) & (trig_mmm))
channels4L = ['eeee', 'eeem', 'eemm', 'mmme', 'mmmm']
selections.add('eeee', ((nElec4l == 4) & (nMuon4l == 0)) & (trig_4l))
selections.add('eeem', ((nElec4l == 3) & (nMuon4l == 1)) & (trig_4l))
selections.add('eemm', ((nElec4l == 2) & (nMuon4l == 2)) & (trig_4l))
selections.add('mmme', ((nElec4l == 1) & (nMuon4l == 3)) & (trig_4l))
selections.add('mmmm', ((nElec4l == 0) & (nMuon4l == 4)) & (trig_4l))
selections.add('ch+', (sumcharge > 0))
selections.add('ch-', (sumcharge < 0))
selections.add('ch0', (sumcharge == 0))
levels = ['base', '1+bm2+bl', '1bm', '2+bm']
selections.add('base', (nElec + nMuon >= 2))
selections.add('1+bm2+bl', (nElec + nMuon >= 2) & ((nbtagsm >= 1) &
(nbtagsl >= 2)))
selections.add('1bm', (nElec + nMuon >= 2) & (nbtagsm == 1))
selections.add('2+bm', (nElec + nMuon >= 2) & (nbtagsm >= 2))
# Variables
invMass_eeSSonZ = (eeSSonZ.e0 + eeSSonZ.e1).mass
invMass_eeSSoffZ = (eeSSoffZ.e0 + eeSSoffZ.e1).mass
invMass_mmSSonZ = (mmSSonZ.m0 + mmSSonZ.m1).mass
invMass_mmSSoffZ = (mmSSoffZ.m0 + mmSSoffZ.m1).mass
invMass_emSS = (emSS.e + emSS.m).mass
varnames = {}
varnames['met'] = met.pt
varnames['ht'] = ht
varnames['njets'] = njets
varnames['invmass'] = {
'eeSSonZ': invMass_eeSSonZ,
'eeSSoffZ': invMass_eeSSoffZ,
'mmSSonZ': invMass_mmSSonZ,
'mmSSoffZ': invMass_mmSSoffZ,
'emSS': invMass_emSS,
'eemSSonZ': mZ_eem,
'eemSSoffZ': mZ_eem,
'mmeSSonZ': mZ_mme,
'mmeSSoffZ': mZ_mme,
'eeeSSonZ': mZ_eee,
'eeeSSoffZ': mZ_eee,
'mmmSSonZ': mZ_mmm,
'mmmSSoffZ': mZ_mmm,
}
varnames['m3l'] = {
'eemSSonZ': m3l_eem,
'eemSSoffZ': m3l_eem,
'mmeSSonZ': m3l_mme,
'mmeSSoffZ': m3l_mme,
'eeeSSonZ': m3l_eee,
'eeeSSoffZ': m3l_eee,
'mmmSSonZ': m3l_mmm,
'mmmSSoffZ': m3l_mmm,
}
varnames['e0pt'] = e0.pt
varnames['e0eta'] = e0.eta
varnames['m0pt'] = m0.pt
varnames['m0eta'] = m0.eta
varnames['j0pt'] = j0.pt
varnames['j0eta'] = j0.eta
varnames['counts'] = np.ones_like(events['event'])
# systematics
systList = []
if isData == False:
systList = ['nominal']
if self._do_systematics:
systList = systList + [
'lepSFUp', 'lepSFDown', 'btagSFUp', 'btagSFDown'
]
else:
systList = ['noweight']
# fill Histos
hout = self.accumulator.identity()
normweights = weights['all'].weight().flatten(
) # Why does it not complain about .flatten() here?
sowweights = np.ones_like(normweights) if len(
self._wc_names_lst) > 0 else normweights
hout['SumOfEFTweights'].fill(sample=histAxisName,
SumOfEFTweights=varnames['counts'],
weight=sowweights,
eft_coeff=eft_coeffs,
eft_err_coeff=eft_w2_coeffs)
for syst in systList:
for var, v in varnames.items():
for ch in channels2LSS + channels3L + channels4L:
for sumcharge in ['ch+', 'ch-', 'ch0']:
for lev in levels:
#find the event weight to be used when filling the histograms
weightSyst = syst
#in the case of 'nominal', or the jet energy systematics, no weight systematic variation is used (weightSyst=None)
if syst in [
'nominal', 'JERUp', 'JERDown', 'JESUp',
'JESDown'
]:
weightSyst = None # no weight systematic for these variations
if syst == 'noweight':
weight = np.ones(len(events)) # for data
else:
# call weights.weight() with the name of the systematic to be varied
if ch in channels3L: ch_w = ch[:3]
elif ch in channels2LSS: ch_w = ch[:2]
else: ch_w = ch
weight = weights['all'].weight(
weightSyst
) if isData else weights[ch_w].weight(
weightSyst)
cuts = [ch] + [lev] + [sumcharge]
cut = selections.all(*cuts)
weights_flat = weight[cut].flatten(
) # Why does it not complain about .flatten() here?
weights_ones = np.ones_like(weights_flat,
dtype=np.int)
eft_coeffs_cut = eft_coeffs[
cut] if eft_coeffs is not None else None
eft_w2_coeffs_cut = eft_w2_coeffs[
cut] if eft_w2_coeffs is not None else None
# filling histos
if var == 'invmass':
if ((ch in [
'eeeSSoffZ', 'mmmSSoffZ', 'eeeSSonZ',
'mmmSSonZ'
]) or (ch in channels4L)):
continue
else:
values = ak.flatten(v[ch][cut])
hout['invmass'].fill(
eft_coeff=eft_coeffs_cut,
eft_err_coeff=eft_w2_coeffs_cut,
sample=histAxisName,
channel=ch,
cut=lev,
sumcharge=sumcharge,
invmass=values,
weight=weights_flat,
systematic=syst)
elif var == 'm3l':
if ((ch in channels2LSS) or (ch in [
'eeeSSoffZ', 'mmmSSoffZ', 'eeeSSonZ',
'mmmSSonZ'
]) or (ch in channels4L)):
continue
values = ak.flatten(v[ch][cut])
hout['m3l'].fill(
eft_coeff=eft_coeffs_cut,
eft_err_coeff=eft_w2_coeffs_cut,
sample=histAxisName,
channel=ch,
cut=lev,
sumcharge=sumcharge,
m3l=values,
weight=weights_flat,
systematic=syst)
else:
values = v[cut]
# These all look identical, do we need if/else here?
if var == 'ht':
hout[var].fill(
eft_coeff=eft_coeffs_cut,
eft_err_coeff=eft_w2_coeffs_cut,
ht=values,
sample=histAxisName,
channel=ch,
cut=lev,
sumcharge=sumcharge,
weight=weights_flat,
systematic=syst)
elif var == 'met':
hout[var].fill(
eft_coeff=eft_coeffs_cut,
eft_err_coeff=eft_w2_coeffs_cut,
met=values,
sample=histAxisName,
channel=ch,
cut=lev,
sumcharge=sumcharge,
weight=weights_flat,
systematic=syst)
elif var == 'njets':
hout[var].fill(
eft_coeff=eft_coeffs_cut,
eft_err_coeff=eft_w2_coeffs_cut,
njets=values,
sample=histAxisName,
channel=ch,
cut=lev,
sumcharge=sumcharge,
weight=weights_flat,
systematic=syst)
elif var == 'nbtags':
hout[var].fill(
eft_coeff=eft_coeffs_cut,
eft_err_coeff=eft_w2_coeffs_cut,
nbtags=values,
sample=histAxisName,
channel=ch,
cut=lev,
sumcharge=sumcharge,
weight=weights_flat,
systematic=syst)
elif var == 'counts':
hout[var].fill(counts=values,
sample=histAxisName,
channel=ch,
cut=lev,
sumcharge=sumcharge,
weight=weights_ones,
systematic=syst)
elif var == 'j0eta':
if lev == 'base': continue
values = ak.flatten(values)
#values=np.asarray(values)
hout[var].fill(
eft_coeff=eft_coeffs_cut,
eft_err_coeff=eft_w2_coeffs_cut,
j0eta=values,
sample=histAxisName,
channel=ch,
cut=lev,
sumcharge=sumcharge,
weight=weights_flat,
systematic=syst)
elif var == 'e0pt':
if ch in [
'mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ',
'mmmSSonZ', 'mmmm'
]:
continue
values = ak.flatten(values)
#values=np.asarray(values)
hout[var].fill(
eft_coeff=eft_coeffs_cut,
eft_err_coeff=eft_w2_coeffs_cut,
e0pt=values,
sample=histAxisName,
channel=ch,
cut=lev,
sumcharge=sumcharge,
weight=weights_flat,
systematic=syst
) # Crashing here, not sure why. Related to values?
elif var == 'm0pt':
if ch in [
'eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ',
'eeeSSonZ', 'eeee'
]:
continue
values = ak.flatten(values)
#values=np.asarray(values)
hout[var].fill(
eft_coeff=eft_coeffs_cut,
eft_err_coeff=eft_w2_coeffs_cut,
m0pt=values,
sample=histAxisName,
channel=ch,
cut=lev,
sumcharge=sumcharge,
weight=weights_flat,
systematic=syst)
elif var == 'e0eta':
if ch in [
'mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ',
'mmmSSonZ', 'mmmm'
]:
continue
values = ak.flatten(values)
#values=np.asarray(values)
hout[var].fill(
eft_coeff=eft_coeffs_cut,
eft_err_coeff=eft_w2_coeffs_cut,
e0eta=values,
sample=histAxisName,
channel=ch,
cut=lev,
sumcharge=sumcharge,
weight=weights_flat,
systematic=syst)
elif var == 'm0eta':
if ch in [
'eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ',
'eeeSSonZ', 'eeee'
]:
continue
values = ak.flatten(values)
#values=np.asarray(values)
hout[var].fill(
eft_coeff=eft_coeffs_cut,
eft_err_coeff=eft_w2_coeffs_cut,
m0eta=values,
sample=histAxisName,
channel=ch,
cut=lev,
sumcharge=sumcharge,
weight=weights_flat,
systematic=syst)
elif var == 'j0pt':
if lev == 'base': continue
values = ak.flatten(values)
#values=np.asarray(values)
hout[var].fill(
eft_coeff=eft_coeffs_cut,
eft_err_coeff=eft_w2_coeffs_cut,
j0pt=values,
sample=histAxisName,
channel=ch,
cut=lev,
sumcharge=sumcharge,
weight=weights_flat,
systematic=syst)
return hout
#Truth-matching
Kst_parent_cut_pi = abs(Kst["pi"]["pdg"]) == abs(lp.Kst_892_0.pdgid)
Kst_parent_cut_k = abs(Kst["k"]["pdg"]) == abs(lp.Kst_892_0.pdgid)
Kst_parent_cut_all = np.logical_and(Kst_parent_cut_pi, Kst_parent_cut_k)
Kst = Kst[Kst_parent_cut_all]
#Ensure kaon and pion come from the same K*
Kst_parent_cut_kpi = Kst["pi","MC_parentindex"] == Kst["k","MC_parentindex"]
Kst = Kst[Kst_parent_cut_kpi]
for comp in ["x","y","z"]:
Kst[f"p{comp}"] = Kst["k",f"p{comp}"] + Kst["pi",f"p{comp}"]
Kst["p"] = np.sqrt(Kst["px"]**2 + Kst["py"]**2 + Kst["pz"]**2)
Kst["e"] = np.sqrt(Kst["p"]**2 + Kst["mass"]**2)
tau = ak.combinations(pions,3)
tau["pi_1"], tau["pi_2"], tau["pi_3"] = ak.unzip(tau)
#Charge cut on the pions
tau_charge_cut_1 = np.sign(tau["pi_1","charge"]) != np.sign(tau["pi_3","charge"])
tau_charge_cut_2 = np.sign(tau["pi_2","charge"]) != np.sign(tau["pi_3","charge"])
tau_charge_cut = np.logical_and(tau_charge_cut_1, tau_charge_cut_2)
tau = tau[tau_charge_cut]
PDG_pi_m = lp.pi_plus.mass/1000.
tau["mass"] = kinematics_flat.mass([tau["pi_1"], tau["pi_2"], tau["pi_3"]], [PDG_pi_m, PDG_pi_m, PDG_pi_m])
tau_m_low_cut_val = 3*PDG_pi_m
tau_m_high_cut_val = lp.tau_plus.mass/1000.
tau_m_low_cut = tau["mass"] >= tau_m_low_cut_val
tau_m_high_cut = tau["mass"] <= lp.tau_plus.mass/1000.
def test_IndexedOptionArray_combinations():
content = ak.from_iter(
[[[0, 1, 2], [], [3, 4]], [], [[5]], [[6, 7, 8, 9]], [[], [10, 11, 12]]
],
highlevel=False,
)
index = ak.layout.Index64(np.array([0, 1, -1, -1, 4], dtype=np.int64))
array = ak.Array(ak.layout.IndexedOptionArray64(index, content))
assert ak.to_list(array) == [
[[0, 1, 2], [], [3, 4]],
[],
None,
None,
[[], [10, 11, 12]],
]
assert ak.to_list(ak.combinations(array, 2, axis=0)) == [
([[0, 1, 2], [], [3, 4]], []),
([[0, 1, 2], [], [3, 4]], None),
([[0, 1, 2], [], [3, 4]], None),
([[0, 1, 2], [], [3, 4]], [[], [10, 11, 12]]),
([], None),
([], None),
([], [[], [10, 11, 12]]),
(None, None),
(None, [[], [10, 11, 12]]),
(None, [[], [10, 11, 12]]),
]
assert ak.to_list(ak.combinations(array, 2, axis=-3)) == [
([[0, 1, 2], [], [3, 4]], []),
([[0, 1, 2], [], [3, 4]], None),
([[0, 1, 2], [], [3, 4]], None),
([[0, 1, 2], [], [3, 4]], [[], [10, 11, 12]]),
([], None),
([], None),
([], [[], [10, 11, 12]]),
(None, None),
(None, [[], [10, 11, 12]]),
(None, [[], [10, 11, 12]]),
]
assert ak.to_list(ak.combinations(array, 2, axis=1)) == [
[([0, 1, 2], []), ([0, 1, 2], [3, 4]), ([], [3, 4])],
[],
None,
None,
[([], [10, 11, 12])],
]
assert ak.to_list(ak.combinations(array, 2, axis=-2)) == [
[([0, 1, 2], []), ([0, 1, 2], [3, 4]), ([], [3, 4])],
[],
None,
None,
[([], [10, 11, 12])],
]
assert ak.to_list(ak.combinations(array, 2, axis=2)) == [
[[(0, 1), (0, 2), (1, 2)], [], [(3, 4)]],
[],
None,
None,
[[], [(10, 11), (10, 12), (11, 12)]],
]
assert ak.to_list(ak.combinations(array, 2, axis=-1)) == [
[[(0, 1), (0, 2), (1, 2)], [], [(3, 4)]],
[],
None,
None,
[[], [(10, 11), (10, 12), (11, 12)]],
]
def test_ListOffsetArray():
array = ak.Array([[0.0, 1.1, 2.2, 3.3], [], [4.4, 5.5, 6.6], [7.7],
[8.8, 9.9, 10.0, 11.1, 12.2]])
assert ak.to_list(ak.combinations(array, 2, replacement=False)) == [
[(0.0, 1.1), (0.0, 2.2), (0.0, 3.3), (1.1, 2.2), (1.1, 3.3),
(2.2, 3.3)],
[],
[(4.4, 5.5), (4.4, 6.6), (5.5, 6.6)],
[],
[
(8.8, 9.9),
(8.8, 10.0),
(8.8, 11.1),
(8.8, 12.2),
(9.9, 10.0),
(9.9, 11.1),
(9.9, 12.2),
(10.0, 11.1),
(10.0, 12.2),
(11.1, 12.2),
],
]
assert ak.to_list(
ak.combinations(array, 2, replacement=False, fields=["x", "y"])) == [
[
{
"x": 0.0,
"y": 1.1
},
{
"x": 0.0,
"y": 2.2
},
{
"x": 0.0,
"y": 3.3
},
{
"x": 1.1,
"y": 2.2
},
{
"x": 1.1,
"y": 3.3
},
{
"x": 2.2,
"y": 3.3
},
],
[],
[{
"x": 4.4,
"y": 5.5
}, {
"x": 4.4,
"y": 6.6
}, {
"x": 5.5,
"y": 6.6
}],
[],
[
{
"x": 8.8,
"y": 9.9
},
{
"x": 8.8,
"y": 10.0
},
{
"x": 8.8,
"y": 11.1
},
{
"x": 8.8,
"y": 12.2
},
{
"x": 9.9,
"y": 10.0
},
{
"x": 9.9,
"y": 11.1
},
{
"x": 9.9,
"y": 12.2
},
{
"x": 10.0,
"y": 11.1
},
{
"x": 10.0,
"y": 12.2
},
{
"x": 11.1,
"y": 12.2
},
],
]
tmp = ak.combinations(array,
2,
replacement=False,
parameters={
"some": "param"
}).layout
if isinstance(tmp, ak.partition.PartitionedArray):
assert ak.partition.first(tmp).content.parameters["some"] == "param"
else:
assert tmp.content.parameters["some"] == "param"
assert ak.to_list(ak.combinations(array, 2, replacement=True)) == [
[
(0.0, 0.0),
(0.0, 1.1),
(0.0, 2.2),
(0.0, 3.3),
(1.1, 1.1),
(1.1, 2.2),
(1.1, 3.3),
(2.2, 2.2),
(2.2, 3.3),
(3.3, 3.3),
],
[],
[(4.4, 4.4), (4.4, 5.5), (4.4, 6.6), (5.5, 5.5), (5.5, 6.6),
(6.6, 6.6)],
[(7.7, 7.7)],
[
(8.8, 8.8),
(8.8, 9.9),
(8.8, 10.0),
(8.8, 11.1),
(8.8, 12.2),
(9.9, 9.9),
(9.9, 10.0),
(9.9, 11.1),
(9.9, 12.2),
(10.0, 10.0),
(10.0, 11.1),
(10.0, 12.2),
(11.1, 11.1),
(11.1, 12.2),
(12.2, 12.2),
],
]
assert ak.to_list(ak.combinations(array, 3, replacement=False)) == [
[(0.0, 1.1, 2.2), (0.0, 1.1, 3.3), (0.0, 2.2, 3.3), (1.1, 2.2, 3.3)],
[],
[(4.4, 5.5, 6.6)],
[],
[
(8.8, 9.9, 10.0),
(8.8, 9.9, 11.1),
(8.8, 9.9, 12.2),
(8.8, 10.0, 11.1),
(8.8, 10.0, 12.2),
(8.8, 11.1, 12.2),
(9.9, 10.0, 11.1),
(9.9, 10.0, 12.2),
(9.9, 11.1, 12.2),
(10.0, 11.1, 12.2),
],
]
assert ak.to_list(ak.combinations(array, 3, replacement=True)) == [
[
(0.0, 0.0, 0.0),
(0.0, 0.0, 1.1),
(0.0, 0.0, 2.2),
(0.0, 0.0, 3.3),
(0.0, 1.1, 1.1),
(0.0, 1.1, 2.2),
(0.0, 1.1, 3.3),
(0.0, 2.2, 2.2),
(0.0, 2.2, 3.3),
(0.0, 3.3, 3.3),
(1.1, 1.1, 1.1),
(1.1, 1.1, 2.2),
(1.1, 1.1, 3.3),
(1.1, 2.2, 2.2),
(1.1, 2.2, 3.3),
(1.1, 3.3, 3.3),
(2.2, 2.2, 2.2),
(2.2, 2.2, 3.3),
(2.2, 3.3, 3.3),
(3.3, 3.3, 3.3),
],
[],
[
(4.4, 4.4, 4.4),
(4.4, 4.4, 5.5),
(4.4, 4.4, 6.6),
(4.4, 5.5, 5.5),
(4.4, 5.5, 6.6),
(4.4, 6.6, 6.6),
(5.5, 5.5, 5.5),
(5.5, 5.5, 6.6),
(5.5, 6.6, 6.6),
(6.6, 6.6, 6.6),
],
[(7.7, 7.7, 7.7)],
[
(8.8, 8.8, 8.8),
(8.8, 8.8, 9.9),
(8.8, 8.8, 10.0),
(8.8, 8.8, 11.1),
(8.8, 8.8, 12.2),
(8.8, 9.9, 9.9),
(8.8, 9.9, 10.0),
(8.8, 9.9, 11.1),
(8.8, 9.9, 12.2),
(8.8, 10.0, 10.0),
(8.8, 10.0, 11.1),
(8.8, 10.0, 12.2),
(8.8, 11.1, 11.1),
(8.8, 11.1, 12.2),
(8.8, 12.2, 12.2),
(9.9, 9.9, 9.9),
(9.9, 9.9, 10.0),
(9.9, 9.9, 11.1),
(9.9, 9.9, 12.2),
(9.9, 10.0, 10.0),
(9.9, 10.0, 11.1),
(9.9, 10.0, 12.2),
(9.9, 11.1, 11.1),
(9.9, 11.1, 12.2),
(9.9, 12.2, 12.2),
(10.0, 10.0, 10.0),
(10.0, 10.0, 11.1),
(10.0, 10.0, 12.2),
(10.0, 11.1, 11.1),
(10.0, 11.1, 12.2),
(10.0, 12.2, 12.2),
(11.1, 11.1, 11.1),
(11.1, 11.1, 12.2),
(11.1, 12.2, 12.2),
(12.2, 12.2, 12.2),
],
]
def run(B_type):
#Load the MC
file_name = f"p8_ee_Zbb_ecm91_EvtGen_{B_type}2TauNuTAUHADNU"
file = uproot.open(f"{loc.IN}/{mode}/{file_name}.root")
tree = file['events']
#Awkward array of the MC - filter event variables if running on Bu sample
if (B_type == "Bc"):
events = tree.arrays(library="ak", how="zip")
#Keep a manageable subset of the Bu events (1M too many to handle in memory)
else:
events = tree.arrays(library="ak", how="zip", entry_stop=20000)
#Reco particles
rp = events["RP"]
#Index cut to remove negative indices (particles produced before hadronisation)
parentindex_cut = rp["MC_parentindex"] >= 0
rp = rp[parentindex_cut]
grandparentindex_cut = rp["MC_grandparentindex"] >= 0
rp = rp[grandparentindex_cut]
#Match to MC truth partners
rp["parent_pdg"] = events["MC", "pdg"][rp["MC_parentindex"]]
rp["grandparent_pdg"] = events["MC", "pdg"][rp["MC_grandparentindex"]]
rp["pdg"] = events["MC", "pdg"][rp["MC_index"]]
#Get the production vertex of the track (truth-level)
for v in ["x", "y", "z"]:
rp[f"vertex_{v}"] = events["MC", f"vertex_{v}"][rp["MC_index"]]
#Get the production vertex of the parent of the tracks (for the pions, this will give the tau production vertex i.e. the Bc decay vertex)
for v in ["x", "y", "z"]:
rp[f"parent_vertex_{v}"] = events["MC",
f"vertex_{v}"][rp["MC_parentindex"]]
#Separation of vertex from PV (combined tau and Bc flight)
rp["PVsep"] = np.sqrt(rp[f"vertex_x"]**2 + rp[f"vertex_y"]**2 +
rp[f"vertex_z"]**2)
#Separation of the Bc vertex from the PV
rp["PVsep_parent"] = np.sqrt(rp[f"parent_vertex_x"]**2 +
rp[f"parent_vertex_y"]**2 +
rp[f"parent_vertex_z"]**2)
#Charged tracks to make tau from, will truth match them to pi+/- ID below
p_cut = rp["p"] > 0.
pions = rp[p_cut]
#Keep charged tracks
charge_cut = abs(rp["charge"]) == 1
pions = pions[charge_cut]
#Build the tau -> 3pi
tau = ak.combinations(pions, 3)
tau["pi_1"], tau["pi_2"], tau["pi_3"] = ak.unzip(tau)
PDG_pi_m = lp.pi_plus.mass / 1000.
tau["mass"] = kinematics_flat.mass([tau["pi_1"], tau["pi_2"], tau["pi_3"]],
[PDG_pi_m, PDG_pi_m, PDG_pi_m])
tau_m_low_cut_val = 3 * PDG_pi_m
tau_m_high_cut_val = lp.tau_plus.mass / 1000.
tau_m_low_cut = tau["mass"] >= tau_m_low_cut_val
tau_m_high_cut = tau["mass"] <= lp.tau_plus.mass / 1000.
tau_m_cut = np.logical_and(tau_m_low_cut, tau_m_high_cut)
tau = tau[tau_m_cut]
#Truth-matching for signal
#Pions are all pions
pi1_cut = abs(tau["pi_1"]["pdg"]) == abs(lp.pi_plus.pdgid)
pi2_cut = abs(tau["pi_2"]["pdg"]) == abs(lp.pi_plus.pdgid)
pi3_cut = abs(tau["pi_3"]["pdg"]) == abs(lp.pi_plus.pdgid)
pi12_cut = np.logical_and(pi1_cut, pi2_cut)
pi_cut_all = np.logical_and(pi12_cut, pi3_cut)
tau = tau[pi_cut_all]
#Pion parents are all tau
tau_parent_cut_pi1 = abs(tau["pi_1"]["parent_pdg"]) == abs(
lp.tau_plus.pdgid)
tau_parent_cut_pi2 = abs(tau["pi_2"]["parent_pdg"]) == abs(
lp.tau_plus.pdgid)
tau_parent_cut_pi3 = abs(tau["pi_3"]["parent_pdg"]) == abs(
lp.tau_plus.pdgid)
tau_parent_cut_12 = np.logical_and(tau_parent_cut_pi1, tau_parent_cut_pi2)
tau_parent_cut_all = np.logical_and(tau_parent_cut_12, tau_parent_cut_pi3)
tau = tau[tau_parent_cut_all]
#Pion grandparent are all Bu / Bc
parent_id = {"Bu": lp.B_plus.pdgid, "Bc": lp.B_c_plus.pdgid}
tau_grandparent_cut_pi1 = abs(tau["pi_1"]["grandparent_pdg"]) == abs(
parent_id[B_type])
tau_grandparent_cut_pi2 = abs(tau["pi_2"]["grandparent_pdg"]) == abs(
parent_id[B_type])
tau_grandparent_cut_pi3 = abs(tau["pi_3"]["grandparent_pdg"]) == abs(
parent_id[B_type])
tau_grandparent_cut_12 = np.logical_and(tau_grandparent_cut_pi1,
tau_grandparent_cut_pi2)
tau_grandparent_cut_all = np.logical_and(tau_grandparent_cut_12,
tau_grandparent_cut_pi3)
tau = tau[tau_grandparent_cut_all]
#Ensure pions come from the same tau
tau_parent_cut_12 = tau["pi_1", "MC_parentindex"] == tau["pi_2",
"MC_parentindex"]
tau_parent_cut_13 = tau["pi_1", "MC_parentindex"] == tau["pi_3",
"MC_parentindex"]
tau_parent_cut = np.logical_and(tau_parent_cut_12, tau_parent_cut_13)
tau = tau[tau_parent_cut]
#Net charge of the tau
tau["charge"] = tau["pi_1", "charge"] + tau["pi_2",
"charge"] + tau["pi_3",
"charge"]
#Tau energy (visible energy of the signal)
for comp in ["x", "y", "z"]:
tau[f"p{comp}"] = tau["pi_1", f"p{comp}"] + tau[
"pi_2", f"p{comp}"] + tau["pi_1", f"p{comp}"]
tau["p"] = np.sqrt(tau["px"]**2 + tau["py"]**2 + tau["pz"]**2)
tau["e"] = np.sqrt(tau["mass"]**2 + tau["p"]**2)
#Tau flight and Bc flight from PV (using true vertices)
tau["PVsep"] = tau["pi_1", "PVsep"]
#x, y, and z coordinates of the tau vertex
for v in ["x", "y", "z"]:
tau[f"PVsep_{v}"] = tau["pi_1", f"vertex_{v}"]
tau["PVsep_parent"] = tau["pi_1", "PVsep_parent"]
#Thrust axis co-ordinates to persist
for v in ["x", "y", "z"]:
tau[f"EVT_thrust_{v}"] = events[f"EVT_thrust_{v}"]
#Hemisphere energies
for hem in ["0", "1"]:
events[f"EVT_thrutshemis{hem}_e"] = events[
f"EVT_thrutshemis{hem}_echarged"] + events[
f"EVT_thrutshemis{hem}_eneutral"]
#Total energy (sum of both hemispheres)
events[
"EVT_e"] = events["EVT_thrutshemis0_e"] + events["EVT_thrutshemis1_e"]
#Min and max hemisphere energies per-event
events["EVT_thrutshemis_e_min"] = np.minimum(events["EVT_thrutshemis0_e"],
events["EVT_thrutshemis1_e"])
events["EVT_thrutshemis_e_max"] = np.maximum(events["EVT_thrutshemis0_e"],
events["EVT_thrutshemis1_e"])
#Difference in energy between hemispheres (max - min so always positive)
events["EVT_thrutshemis_e_diff"] = events[
"EVT_thrutshemis_e_max"] - events["EVT_thrutshemis_e_min"]
fig, ax = plt.subplots(figsize=(9, 9))
plt.hist(events["EVT_thrutshemis_e_min"],
bins=50,
range=(0, 60),
color="crimson",
histtype='step',
linewidth=2,
label="Lower energy hemisphere per-event")
plt.hist(events["EVT_thrutshemis_e_max"],
bins=50,
range=(0, 60),
color="dodgerblue",
histtype='step',
linewidth=2,
label="Higher energy hemisphere per-event")
ax.tick_params(axis='both', which='major', labelsize=25)
plt.xlim(0, 60)
plt.axvline(lp.Z_0.mass / 2000.,
color='k',
linestyle='--',
label="$m(Z^0)/2$")
plt.xlabel("Hemisphere energy [GeV]", fontsize=30)
plt.legend(fontsize=18, loc="upper left")
ut.create_dir(loc.PLOTS)
fig.savefig(f"{loc.PLOTS}/{B_type}2TauNu_min_max_hemisphere_E.pdf")
fig, ax = plt.subplots(figsize=(9, 9))
plt.hist(events["EVT_thrutshemis_e_diff"],
bins=30,
range=(0, 50),
color="k",
histtype='step',
linewidth=2)
ax.tick_params(axis='both', which='major', labelsize=25)
plt.xlim(0, 50)
plt.xlabel("Hemisphere energy difference [GeV]", fontsize=30)
fig.savefig(f"{loc.PLOTS}/{B_type}2TauNu_diff_hemisphere_E.pdf")
#Store hemisphere energy info in tau container for writing out later
tau["EVT_e"] = events["EVT_e"]
tau["EVT_thrutshemis_e_min"] = events["EVT_thrutshemis_e_min"]
tau["EVT_thrutshemis_e_max"] = events["EVT_thrutshemis_e_max"]
tau["EVT_thrutshemis_e_diff"] = events["EVT_thrutshemis_e_diff"]
tau["EVT_thrutshemis0_e"] = events["EVT_thrutshemis0_e"]
tau["EVT_thrutshemis1_e"] = events["EVT_thrutshemis1_e"]
#Nominal B energy = m(Z) - sum of all energy apart from signal
tau["nominal_B_e"] = lp.Z_0.mass / 1000. - (tau["EVT_e"] - tau["e"])
fig, ax = plt.subplots(figsize=(9, 9))
plt.hist(ak.flatten(tau["nominal_B_e"]),
range=(0, 80),
bins=30,
color="k",
histtype='step',
linewidth=2)
ax.tick_params(axis='both', which='major', labelsize=25)
plt.title(
"$m(Z) - \\Sigma$(all visible E apart from true $\\tau \\to 3\\pi$)",
fontsize=25)
plt.xlim(0, 80)
plt.xlabel("Nominal $B$ energy [GeV]", fontsize=30)
fig.savefig(f"{loc.PLOTS}/{B_type}2TauNu_nominal_B_E.pdf")
#Plot the tau mass for truth-matched signal
fig, ax = plt.subplots(figsize=(9, 9))
plt.hist(ak.flatten(tau["mass"]),
range=(tau_m_low_cut_val, tau_m_high_cut_val),
bins=30,
color="k",
histtype='step',
linewidth=2)
ax.tick_params(axis='both', which='major', labelsize=25)
plt.xlim(tau_m_low_cut_val, tau_m_high_cut_val)
plt.xlabel("$m(3\\pi)$ [GeV]", fontsize=30)
fig.savefig(f"{loc.PLOTS}/{B_type}2TauNu_3pi_M.pdf")
#Look at charged and neutral multiplicities in the two hemispheres
#Events where hemisphere 0 is the minimum energy
events_hem0_min_cut = events["EVT_thrutshemis0_e"] == events[
"EVT_thrutshemis_e_min"]
events_hem0_min = events[events_hem0_min_cut]
#Events where hemisphere 1 is the minimum energy
events_hem1_min_cut = events["EVT_thrutshemis1_e"] == events[
"EVT_thrutshemis_e_min"]
events_hem1_min = events[events_hem1_min_cut]
#Get the charged and neutral energy and multiplicity
for ptype in ["charged", "neutral"]:
for var in ["n", "e"]:
events_hem0_min[f"{var}{ptype}_min"] = events_hem0_min[
f"EVT_thrutshemis0_{var}{ptype}"]
events_hem0_min[f"{var}{ptype}_max"] = events_hem0_min[
f"EVT_thrutshemis1_{var}{ptype}"]
events_hem1_min[f"{var}{ptype}_min"] = events_hem1_min[
f"EVT_thrutshemis1_{var}{ptype}"]
events_hem1_min[f"{var}{ptype}_max"] = events_hem1_min[
f"EVT_thrutshemis0_{var}{ptype}"]
#Recombine
events = ak.concatenate([events_hem0_min, events_hem1_min], axis=0)
#Plot the charged and neutral energies and multiplicities, where the low and high energy hemispheres are shown separately
#Axis ranges, titles, and bins
plot_config = {
"echarged": [0, 50, "Charged energy [GeV]", 50],
"eneutral": [0, 40, "Neutral energy [GeV]", 50],
"ncharged": [0, 25, "Charged multiplicity", 25],
"nneutral": [0, 20, "Neutral multiplicity", 20]
}
for p in plot_config:
#Store the values in the tau container, so we can write it out
tau[f"{p}_min"] = events[f"{p}_min"]
tau[f"{p}_max"] = events[f"{p}_max"]
fig, ax = plt.subplots(figsize=(9, 9))
plt.hist(events[f"{p}_min"],
bins=plot_config[p][3],
range=(plot_config[p][0], plot_config[p][1]),
color="crimson",
histtype='step',
linewidth=2,
label="Lower energy hemisphere per-event")
plt.hist(events[f"{p}_max"],
bins=plot_config[p][3],
range=(plot_config[p][0], plot_config[p][1]),
color="dodgerblue",
histtype='step',
linewidth=2,
label="Higher energy hemisphere per-event")
ax.tick_params(axis='both', which='major', labelsize=25)
plt.xlim(plot_config[p][0], plot_config[p][1])
plt.xlabel(plot_config[p][2], fontsize=30)
ymin, ymax = plt.ylim()
plt.ylim(0, 1.15 * ymax)
plt.legend(fontsize=18, loc="upper left")
fig.savefig(f"{loc.PLOTS}/{B_type}2TauNu_min_max_hemisphere_{p}.pdf")
#Plot the separation of the tau pion vertex from the PV (a measure of the B + tau flight)
fig, ax = plt.subplots(figsize=(9, 9))
plt.hist(ak.flatten(tau["pi_1", "PVsep"]),
range=(0, 10),
bins=30,
color="k",
histtype='step',
linewidth=2)
ax.tick_params(axis='both', which='major', labelsize=25)
plt.xlim(0, 10)
plt.xlabel("$3\pi$ separation from PV [mm]", fontsize=30)
fig.savefig(f"{loc.PLOTS}/{B_type}2TauNu_3pi_vertex_PV_sep.pdf")
#Plot the separation of the true Bc decay vertex from the PV - this quantity cannot be reconstructed experimentally
fig, ax = plt.subplots(figsize=(9, 9))
plt.hist(ak.flatten(tau["pi_1", "PVsep_parent"]),
range=(0, 4),
bins=30,
color="k",
histtype='step',
linewidth=2)
ax.tick_params(axis='both', which='major', labelsize=25)
plt.xlim(0, 4)
B_name = {"Bu": "$B^\\pm$", "Bc": "$B_c^\\pm$"}
plt.xlabel(f"{B_name[B_type]} separation from PV [mm]", fontsize=30)
fig.savefig(f"{loc.PLOTS}/{B_type}2TauNu_Bc_vertex_PV_sep.pdf")
#Persist information into a flat dataframe for use elsewhere e.g. MVA training vs. inclusive Z -> qq/cc/bb background
df = pd.DataFrame()
persist_vars = [
"EVT_e", #Total event visible energy
"EVT_thrutshemis_e_min", #Lowest energy hemisphere in event
"EVT_thrutshemis_e_max", #Highest energy hemisphere in event
"EVT_thrutshemis0_e", #costheta < 0 hemisphere energy
"EVT_thrutshemis1_e", #costheta > 0 hemisphere energy
"nominal_B_e", #Nominal signal B energy calculated using reco energy of true tau -> 3pi
"px",
"py",
"pz",
"e",
"mass",
"PVsep", #Distance of true tau decay vertex from the PV
"PVsep_x",
"PVsep_y",
"PVsep_z",
"PVsep_parent", #Distance of true Bc vertex from the PV (Bc flight which is not reconstructible)
"EVT_thrust_x",
"EVT_thrust_y",
"EVT_thrust_z"
]
for var in ["e", "n"]:
for ptype in ["charged", "neutral"]:
for m in ["min", "max"]:
persist_vars.append(f"{var}{ptype}_{m}")
for var in persist_vars:
df[var] = ak.flatten(tau[var]).tolist()
#Save to CSV
ut.create_dir(loc.CSV)
df.to_csv(f"{loc.CSV}/{B_type}2TauNu.csv")
'''
def process(self, events):
# Dataset parameters
dataset = events.metadata['dataset']
year = self._samples[dataset]['year']
xsec = self._samples[dataset]['xsec']
sow = self._samples[dataset]['nSumOfWeights' ]
isData = self._samples[dataset]['isData']
datasets = ['SingleMuon', 'SingleElectron', 'EGamma', 'MuonEG', 'DoubleMuon', 'DoubleElectron']
for d in datasets:
if d in dataset: dataset = dataset.split('_')[0]
# Inittialize objects
met = events.GenMET
e = events.GenPart[abs(events.GenPart.pdgId)==11]
mu = events.GenPart[abs(events.GenPart.pdgId)==13]
tau = events.GenPart[abs(events.GenPart.pdgId)==15]
j = events.GenJet
leading_mu = mu[ak.argmax(mu.pt,axis=-1,keepdims=True)]
leading_e = e[ak.argmax(e.pt,axis=-1,keepdims=True)]
nElec = ak.num(e)
nMuon = ak.num(mu)
nTau = ak.num(tau)
twoLeps = (nElec+nMuon) == 2
threeLeps = (nElec+nMuon) == 3
twoElec = (nElec == 2)
twoMuon = (nMuon == 2)
e0 = e[ak.argmax(e.pt,axis=-1,keepdims=True)]
m0 = mu[ak.argmax(mu.pt,axis=-1,keepdims=True)]
elecs = e[ak.argsort(e.pt, ascending=False)]
muons = mu[ak.argsort(mu.pt, ascending=False)]
e1 = elecs
e2 = elecs
m1 = muons
m2 = muons
# Jet selection
jetptname = 'pt_nom' if hasattr(j, 'pt_nom') else 'pt'
njets = ak.num(j)
ht = ak.sum(j.pt,axis=-1)
jets = j[ak.argsort(j.pt, ascending=False)]
j0 = j[ak.argmax(j.pt,axis=-1,keepdims=True)]
j1 = jets
j2 = jets
j3 = jets
nbtags = ak.num(j[abs(j.hadronFlavour)==5])
##################################################################
### 2 same-sign leptons
##################################################################
# emu
singe = e [(nElec==1)&(nMuon==1)&(e .pt>-1)]
singm = mu[(nElec==1)&(nMuon==1)&(mu.pt>-1)]
em = ak.cartesian({"e":singe,"m":singm})
emSSmask = (em.e.pdgId*em.m.pdgId>0)
emSS = em[emSSmask]
nemSS = len(ak.flatten(emSS))
year = 2018
lepSF_emSS = GetLeptonSF(mu.pt, mu.eta, 'm', e.pt, e.eta, 'e', year=year)
# ee and mumu
# pt>-1 to preserve jagged dimensions
ee = e [(nElec==2)&(nMuon==0)&(e.pt>-1)]
mm = mu[(nElec==0)&(nMuon==2)&(mu.pt>-1)]
eepairs = ak.combinations(ee, 2, fields=["e0","e1"])
eeSSmask = (eepairs.e0.pdgId*eepairs.e1.pdgId>0)
eeonZmask = (np.abs((eepairs.e0+eepairs.e1).mass-91.2)<10)
eeoffZmask = (eeonZmask==0)
mmpairs = ak.combinations(mm, 2, fields=["m0","m1"])
mmSSmask = (mmpairs.m0.pdgId*mmpairs.m1.pdgId>0)
mmonZmask = (np.abs((mmpairs.m0+mmpairs.m1).mass-91.2)<10)
mmoffZmask = (mmonZmask==0)
eeSSonZ = eepairs[eeSSmask & eeonZmask]
eeSSoffZ = eepairs[eeSSmask & eeoffZmask]
mmSSonZ = mmpairs[mmSSmask & mmonZmask]
mmSSoffZ = mmpairs[mmSSmask & mmoffZmask]
neeSS = len(ak.flatten(eeSSonZ)) + len(ak.flatten(eeSSoffZ))
nmmSS = len(ak.flatten(mmSSonZ)) + len(ak.flatten(mmSSoffZ))
lepSF_eeSS = GetLeptonSF(eepairs.e0.pt, eepairs.e0.eta, 'e', eepairs.e1.pt, eepairs.e1.eta, 'e', year=year)
lepSF_mumuSS = GetLeptonSF(mmpairs.m0.pt, mmpairs.m0.eta, 'm', mmpairs.m1.pt, mmpairs.m1.eta, 'm', year=year)
print('Same-sign events [ee, emu, mumu] = [%i, %i, %i]'%(neeSS, nemSS, nmmSS))
# Cuts
eeSSmask = (ak.num(eeSSmask[eeSSmask])>0)
mmSSmask = (ak.num(mmSSmask[mmSSmask])>0)
eeonZmask = (ak.num(eeonZmask[eeonZmask])>0)
eeoffZmask = (ak.num(eeoffZmask[eeoffZmask])>0)
mmonZmask = (ak.num(mmonZmask[mmonZmask])>0)
mmoffZmask = (ak.num(mmoffZmask[mmoffZmask])>0)
emSSmask = (ak.num(emSSmask[emSSmask])>0)
##################################################################
### 3 leptons
##################################################################
# eem
muon_eem = mu[(nElec==2)&(nMuon==1)&(mu.pt>-1)]
elec_eem = e[(nElec==2)&(nMuon==1)&( e.pt>-1)]
ee_eem = ak.combinations(elec_eem, 2, fields=["e0", "e1"])
ee_eemZmask = (ee_eem.e0.pdgId*ee_eem.e1.pdgId<1)&(np.abs((ee_eem.e0+ee_eem.e1).mass-91.2)<10)
ee_eemOffZmask = (ee_eem.e0.pdgId*ee_eem.e1.pdgId<1)&(np.abs((ee_eem.e0+ee_eem.e1).mass-91.2)>10)
ee_eemZmask = (ak.num(ee_eemZmask[ee_eemZmask])>0)
ee_eemOffZmask = (ak.num(ee_eemOffZmask[ee_eemOffZmask])>0)
eepair_eem = (ee_eem.e0+ee_eem.e1)
trilep_eem = eepair_eem+muon_eem #ak.cartesian({"e0":ee_eem.e0,"e1":ee_eem.e1, "m":muon_eem})
lepSF_eem = GetLeptonSF(ee_eem.e0.pt, ee_eem.e0.eta, 'e', ee_eem.e1.pt, ee_eem.e1.eta, 'e', mu.pt, mu.eta, 'm', year)
# mme
muon_mme = mu[(nElec==1)&(nMuon==2)&(mu.pt>-1)]
elec_mme = e[(nElec==1)&(nMuon==2)&( e.pt>-1)]
mm_mme = ak.combinations(muon_mme, 2, fields=["m0", "m1"])
mm_mmeZmask = (mm_mme.m0.pdgId*mm_mme.m1.pdgId<1)&(np.abs((mm_mme.m0+mm_mme.m1).mass-91.2)<10)
mm_mmeOffZmask = (mm_mme.m0.pdgId*mm_mme.m1.pdgId<1)&(np.abs((mm_mme.m0+mm_mme.m1).mass-91.2)>10)
mm_mmeZmask = (ak.num(mm_mmeZmask[mm_mmeZmask])>0)
mm_mmeOffZmask = (ak.num(mm_mmeOffZmask[mm_mmeOffZmask])>0)
mmpair_mme = (mm_mme.m0+mm_mme.m1)
trilep_mme = mmpair_mme+elec_mme
mZ_mme = mmpair_mme.mass
mZ_eem = eepair_eem.mass
m3l_eem = trilep_eem.mass
m3l_mme = trilep_mme.mass
lepSF_mme = GetLeptonSF(mm_mme.m0.pt, mm_mme.m0.eta, 'm', mm_mme.m1.pt, mm_mme.m1.eta, 'm', e.pt, e.eta, 'e', year)
# eee and mmm
eee = e[(nElec==3)&(nMuon==0)&( e.pt>-1)]
mmm = mu[(nElec==0)&(nMuon==3)&(mu.pt>-1)]
eee_leps = ak.combinations(eee, 3, fields=["e0", "e1", "e2"])
mmm_leps = ak.combinations(mmm, 3, fields=["m0", "m1", "m2"])
ee_pairs = ak.combinations(eee, 2, fields=["e0", "e1"])
mm_pairs = ak.combinations(mmm, 2, fields=["m0", "m1"])
ee_pairs_index = ak.argcombinations(eee, 2, fields=["e0", "e1"])
mm_pairs_index = ak.argcombinations(mmm, 2, fields=["m0", "m1"])
lepSF_eee = GetLeptonSF(eee_leps.e0.pt, eee_leps.e0.eta, 'e', eee_leps.e1.pt, eee_leps.e1.eta, 'e', eee_leps.e2.pt, eee_leps.e2.eta, 'e', year)
lepSF_mmm = GetLeptonSF(mmm_leps.m0.pt, mmm_leps.m0.eta, 'm', mmm_leps.m1.pt, mmm_leps.m1.eta, 'm', mmm_leps.m2.pt, mmm_leps.m2.eta, 'm', year)
mmSFOS_pairs = mm_pairs[(np.abs(mm_pairs.m0.pdgId) == np.abs(mm_pairs.m1.pdgId)) & (mm_pairs.m0.pdgId != mm_pairs.m1.pdgId)]
offZmask_mm = ak.all(np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2)>10., axis=1, keepdims=True) & (ak.num(mmSFOS_pairs)>0)
onZmask_mm = ak.any(np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2)<10., axis=1, keepdims=True)
eeSFOS_pairs = ee_pairs[(np.abs(ee_pairs.e0.pdgId) == np.abs(ee_pairs.e1.pdgId)) & (ee_pairs.e0.pdgId != ee_pairs.e1.pdgId)]
offZmask_ee = ak.all(np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2)>10, axis=1, keepdims=True) & (ak.num(eeSFOS_pairs)>0)
onZmask_ee = ak.any(np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2)<10, axis=1, keepdims=True)
# Create masks **for event selection**
eeeOnZmask = (ak.num(onZmask_ee[onZmask_ee])>0)
eeeOffZmask = (ak.num(offZmask_ee[offZmask_ee])>0)
mmmOnZmask = (ak.num(onZmask_mm[onZmask_mm])>0)
mmmOffZmask = (ak.num(offZmask_mm[offZmask_mm])>0)
# Now we need to create masks for the leptons in order to select leptons from the Z boson candidate (in onZ categories)
ZeeMask = ak.argmin(np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2),axis=1,keepdims=True)
ZmmMask = ak.argmin(np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2),axis=1,keepdims=True)
Zee = eeSFOS_pairs[ZeeMask]
Zmm = mmSFOS_pairs[ZmmMask]
eZ0= Zee.e0[ak.num(eeSFOS_pairs)>0]
eZ1= Zee.e1[ak.num(eeSFOS_pairs)>0]
eZ = eZ0+eZ1
mZ0= Zmm.m0[ak.num(mmSFOS_pairs)>0]
mZ1= Zmm.m1[ak.num(mmSFOS_pairs)>0]
mZ = mZ0+mZ1
mZ_eee = eZ.mass
mZ_mmm = mZ.mass
# And for the W boson
ZmmIndices = mm_pairs_index[ZmmMask]
ZeeIndices = ee_pairs_index[ZeeMask]
eW = eee[~ZeeIndices.e0 | ~ZeeIndices.e1]
mW = mmm[~ZmmIndices.m0 | ~ZmmIndices.m1]
triElec = eee_leps.e0+eee_leps.e1+eee_leps.e2
triMuon = mmm_leps.m0+mmm_leps.m1+mmm_leps.m2
m3l_eee = triElec.mass
m3l_mmm = triMuon.mass
# Triggers
trig_eeSS = passTrigger(events,'ee',isData,dataset)
trig_mmSS = passTrigger(events,'mm',isData,dataset)
trig_emSS = passTrigger(events,'em',isData,dataset)
trig_eee = passTrigger(events,'eee',isData,dataset)
trig_mmm = passTrigger(events,'mmm',isData,dataset)
trig_eem = passTrigger(events,'eem',isData,dataset)
trig_mme = passTrigger(events,'mme',isData,dataset)
# MET filters
# Weights
genw = np.ones_like(events['MET_pt']) if isData else events['genWeight']
### We need weights for: normalization, lepSF, triggerSF, pileup, btagSF...
weights = {}
for r in ['all', 'ee', 'mm', 'em', 'eee', 'mmm', 'eem', 'mme']:
weights[r] = coffea.analysis_tools.Weights(len(events))
weights[r].add('norm',genw if isData else (xsec/sow)*genw)
weights['ee'].add('lepSF_eeSS', lepSF_eeSS)
weights['em'].add('lepSF_emSS', lepSF_emSS)
weights['mm'].add('lepSF_mmSS', lepSF_mumuSS)
weights['eee'].add('lepSF_eee', lepSF_eee)
weights['mmm'].add('lepSF_mmm', lepSF_mmm)
weights['mme'].add('lepSF_mme', lepSF_mme)
weights['eem'].add('lepSF_eem', lepSF_eem)
# Extract the EFT quadratic coefficients and optionally use them to calculate the coefficients on the w**2 quartic function
# eft_coeffs is never Jagged so convert immediately to numpy for ease of use.
eft_coeffs = ak.to_numpy(events['EFTfitCoefficients']) if hasattr(events, "EFTfitCoefficients") else None
eft_w2_coeffs = efth.calc_w2_coeffs(eft_coeffs,self._dtype) if (self._do_errors and eft_coeffs is not None) else None
# Selections and cuts
selections = PackedSelection()
channels2LSS = ['eeSSonZ', 'eeSSoffZ', 'mmSSonZ', 'mmSSoffZ', 'emSS']
selections.add('eeSSonZ', (eeonZmask)&(eeSSmask)&(trig_eeSS))
selections.add('eeSSoffZ', (eeoffZmask)&(eeSSmask)&(trig_eeSS))
selections.add('mmSSonZ', (mmonZmask)&(mmSSmask)&(trig_mmSS))
selections.add('mmSSoffZ', (mmoffZmask)&(mmSSmask)&(trig_mmSS))
selections.add('emSS', (emSSmask)&(trig_emSS))
channels3L = ['eemSSonZ', 'eemSSoffZ', 'mmeSSonZ', 'mmeSSoffZ']
selections.add('eemSSonZ', (ee_eemZmask)&(trig_eem))
selections.add('eemSSoffZ', (ee_eemOffZmask)&(trig_eem))
selections.add('mmeSSonZ', (mm_mmeZmask)&(trig_mme))
selections.add('mmeSSoffZ', (mm_mmeOffZmask)&(trig_mme))
channels3L += ['eeeSSonZ', 'eeeSSoffZ', 'mmmSSonZ', 'mmmSSoffZ']
selections.add('eeeSSonZ', (eeeOnZmask)&(trig_eee))
selections.add('eeeSSoffZ', (eeeOffZmask)&(trig_eee))
selections.add('mmmSSonZ', (mmmOnZmask)&(trig_mmm))
selections.add('mmmSSoffZ', (mmmOffZmask)&(trig_mmm))
levels = ['base', '2jets', '4jets', '4j1b', '4j2b']
selections.add('base', (nElec+nMuon>=2))
selections.add('2jets',(njets>=2))
selections.add('4jets',(njets>=4))
selections.add('4j1b',(njets>=4)&(nbtags>=1))
selections.add('4j2b',(njets>=4)&(nbtags>=2))
# Variables
invMass_eeSSonZ = ( eeSSonZ.e0+ eeSSonZ.e1).mass
invMass_eeSSoffZ = (eeSSoffZ.e0+eeSSoffZ.e1).mass
invMass_mmSSonZ = ( mmSSonZ.m0+ mmSSonZ.m1).mass
invMass_mmSSoffZ = (mmSSoffZ.m0+mmSSoffZ.m1).mass
invMass_emSS = (emSS.e+emSS.m).mass
varnames = {}
varnames['met'] = met.pt
varnames['ht'] = ht
varnames['njets'] = njets
varnames['nbtags'] = nbtags
varnames['invmass'] = {
'eeSSonZ' : invMass_eeSSonZ,
'eeSSoffZ' : invMass_eeSSoffZ,
'mmSSonZ' : invMass_mmSSonZ,
'mmSSoffZ' : invMass_mmSSoffZ,
'emSS' : invMass_emSS,
'eemSSonZ' : mZ_eem,
'eemSSoffZ' : mZ_eem,
'mmeSSonZ' : mZ_mme,
'mmeSSoffZ' : mZ_mme,
'eeeSSonZ' : mZ_eee,
'eeeSSoffZ' : mZ_eee,
'mmmSSonZ' : mZ_mmm,
'mmmSSoffZ' : mZ_mmm,
}
varnames['m3l'] = {
'eemSSonZ' : m3l_eem,
'eemSSoffZ' : m3l_eem,
'mmeSSonZ' : m3l_mme,
'mmeSSoffZ' : m3l_mme,
'eeeSSonZ' : m3l_eee,
'eeeSSoffZ' : m3l_eee,
'mmmSSonZ' : m3l_mmm,
'mmmSSoffZ' : m3l_mmm,
}
varnames['e0pt' ] = e0.pt
varnames['e0eta'] = e0.eta
varnames['m0pt' ] = m0.pt
varnames['m0eta'] = m0.eta
varnames['e1pt' ] = e1
varnames['e1eta'] = e1
varnames['e2pt' ] = e2
varnames['e2eta'] = e2
varnames['m1pt' ] = m1
varnames['m1eta'] = m1
varnames['m2pt' ] = m2
varnames['m2eta'] = m2
varnames['j0pt' ] = j0.pt
varnames['j0eta'] = j0.eta
varnames['j1pt'] = j1
varnames['j1eta'] = j1
varnames['j2pt'] = j2
varnames['j2eta'] = j2
varnames['j3pt'] = j3
varnames['j3eta'] = j3
varnames['counts'] = np.ones_like(events.GenMET.pt)
# fill Histos
hout = self.accumulator.identity()
normweights = weights['all'].weight().flatten() # Why does it not complain about .flatten() here?
hout['SumOfEFTweights'].fill(sample=dataset, SumOfEFTweights=varnames['counts'], weight=normweights, eft_coeff=eft_coeffs, eft_err_coeff=eft_w2_coeffs)
for var, v in varnames.items():
for ch in channels2LSS+channels3L:
for lev in levels:
weight = weights[ ch[:3] if (ch.startswith('eee') or ch.startswith('mmm') or ch.startswith('eem') or ch.startswith('mme')) else ch[:2]].weight()
cuts = [ch] + [lev]
cut = selections.all(*cuts)
weights_flat = weight[cut].flatten() # Why does it not complain about .flatten() here?
weights_ones = np.ones_like(weights_flat, dtype=np.int)
eft_coeffs_cut = eft_coeffs[cut] if eft_coeffs is not None else None
eft_w2_coeffs_cut = eft_w2_coeffs[cut] if eft_w2_coeffs is not None else None
if var == 'invmass':
if ch in ['eeeSSoffZ', 'mmmSSoffZ']: continue
elif ch in ['eeeSSonZ' , 'mmmSSonZ' ]: continue #values = v[ch]
else : values = ak.flatten(v[ch][cut])
hout['invmass'].fill(eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, sample=dataset, channel=ch, cut=lev, invmass=values, weight=weights_flat)
elif var == 'm3l':
if ch in ['eeSSonZ','eeSSoffZ', 'mmSSonZ', 'mmSSoffZ','emSS', 'eeeSSoffZ', 'mmmSSoffZ', 'eeeSSonZ' , 'mmmSSonZ']: continue
values = ak.flatten(v[ch][cut])
hout['m3l'].fill(sample=dataset, channel=ch, cut=lev, m3l=values, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut)
else:
values = v[cut]
if var == 'ht' : hout[var].fill(ht=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut)
elif var == 'met' : hout[var].fill(met=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut)
elif var == 'njets' : hout[var].fill(njets=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut)
elif var == 'nbtags': hout[var].fill(nbtags=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut)
elif var == 'counts': hout[var].fill(counts=values, sample=dataset, channel=ch, cut=lev, weight=weights_ones)
elif var == 'j0eta' :
if lev == 'base': continue
values = ak.flatten(values)
#values=np.asarray(values)
hout[var].fill(j0eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut)
elif var == 'e0pt' :
if ch in ['mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ', 'mmmSSonZ']: continue
values = ak.flatten(values)
#values=np.asarray(values)
hout[var].fill(e0pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut)
elif var == 'm0pt' :
if ch in ['eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ', 'eeeSSonZ']: continue
values = ak.flatten(values)
#values=np.asarray(values)
hout[var].fill(m0pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut)
elif var == 'e0eta' :
if ch in ['mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ', 'mmmSSonZ']: continue
values = ak.flatten(values)
#values=np.asarray(values)
hout[var].fill(e0eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut)
elif var == 'm0eta':
if ch in ['eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ', 'eeeSSonZ']: continue
values = ak.flatten(values)
#values=np.asarray(values)
hout[var].fill(m0eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut)
elif var == 'j0pt' :
if lev == 'base': continue
values = ak.flatten(values)
#values=np.asarray(values)
hout[var].fill(j0pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut)
elif var == 'j1pt':
if lev == "base": continue
values = values.pt[:,1]
#values = ak.flatten(values)
hout[var].fill(j1pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut)
elif var =='j1eta':
if lev == 'base': continue
values = values.eta[:,1]
hout[var].fill(j1eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut)
elif var == 'j2pt':
if lev in ['base', "2jets"]: continue
values = values.pt[:,2]
hout[var].fill(j2pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut)
elif var == 'j2eta':
if lev in ['base', "2jets"]: continue
values = values.eta[:,2]
hout[var].fill(j2eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut)
elif var == 'j3pt':
if lev in ['base', "2jets"]: continue
values = values.pt[:,3]
hout[var].fill(j3pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut)
elif var == 'j3eta':
if lev in ['base', "2jets"]: continue
values = values.eta[:,3]
hout[var].fill(j3eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut)
elif var == 'e1pt':
if ch in ['mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ', 'mmmSSonZ', 'mmeSSonZ', 'mmeSSoffZ', 'emSS']: continue
values = values.pt[:,1]
hout[var].fill(e1pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut)
elif var == 'e1eta':
if ch in ['mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ', 'mmmSSonZ', 'mmeSSonZ', 'mmeSSoffZ', 'emSS']: continue
values = values.eta[:,1]
hout[var].fill(e1eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut)
elif var == 'e2pt':
if ch in ['eeeSSonZ', 'eeeSSoffZ']:
values = values.pt[:,2]
hout[var].fill(e2pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut)
elif var == 'e2eta':
if ch in ['eeeSSonZ', 'eeeSSoffZ']:
values = values.eta[:,2]
hout[var].fill(e2eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut)
elif var == 'm1pt':
if ch in ['eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ', 'eeeSSonZ', 'eemSSonZ', 'eemSSoffZ', 'emSS']: continue
values = values.pt[:,1]
hout[var].fill(m1pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut)
elif var == 'm1eta':
if ch in ['eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ', 'eeeSSonZ', 'eemSSonZ', 'eemSSoffZ', 'emSS']: continue
values = values.eta[:,1]
hout[var].fill(m1eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut)
elif var == 'm2pt':
if ch in ['mmmSSonZ', 'mmmSSoffZ']:
values = values.pt[:,2]
hout[var].fill(m2pt=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut)
elif var == 'm2eta':
if ch in ['mmmSSonZ', 'mmmSSoffZ']:
values = values.eta[:,2]
hout[var].fill(m2eta=values, sample=dataset, channel=ch, cut=lev, weight=weights_flat, eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut)
return hout
def test_RegularArray():
array = ak.Array(np.array([[0.0, 1.1, 2.2, 3.3], [4.4, 5.5, 6.6, 7.7]]))
assert ak.to_list(ak.combinations(array, 2,
replacement=False)) == [
[(0.0, 1.1), (0.0, 2.2), (0.0, 3.3),
(1.1, 2.2), (1.1, 3.3), (2.2, 3.3)],
[(4.4, 5.5), (4.4, 6.6), (4.4, 7.7),
(5.5, 6.6), (5.5, 7.7), (6.6, 7.7)],
]
assert ak.to_list(
ak.combinations(array, 2, replacement=False, fields=["x", "y"])) == [
[
{
"x": 0.0,
"y": 1.1
},
{
"x": 0.0,
"y": 2.2
},
{
"x": 0.0,
"y": 3.3
},
{
"x": 1.1,
"y": 2.2
},
{
"x": 1.1,
"y": 3.3
},
{
"x": 2.2,
"y": 3.3
},
],
[
{
"x": 4.4,
"y": 5.5
},
{
"x": 4.4,
"y": 6.6
},
{
"x": 4.4,
"y": 7.7
},
{
"x": 5.5,
"y": 6.6
},
{
"x": 5.5,
"y": 7.7
},
{
"x": 6.6,
"y": 7.7
},
],
]
tmp = ak.combinations(array,
2,
replacement=False,
parameters={
"some": "param"
}).layout
if isinstance(tmp, ak.partition.PartitionedArray):
assert ak.partition.first(tmp).content.parameters["some"] == "param"
else:
assert tmp.content.parameters["some"] == "param"
assert ak.to_list(ak.combinations(array, 2, replacement=True)) == [
[
(0.0, 0.0),
(0.0, 1.1),
(0.0, 2.2),
(0.0, 3.3),
(1.1, 1.1),
(1.1, 2.2),
(1.1, 3.3),
(2.2, 2.2),
(2.2, 3.3),
(3.3, 3.3),
],
[
(4.4, 4.4),
(4.4, 5.5),
(4.4, 6.6),
(4.4, 7.7),
(5.5, 5.5),
(5.5, 6.6),
(5.5, 7.7),
(6.6, 6.6),
(6.6, 7.7),
(7.7, 7.7),
],
]
assert ak.to_list(ak.combinations(array, 3, replacement=False)) == [
[(0.0, 1.1, 2.2), (0.0, 1.1, 3.3), (0.0, 2.2, 3.3), (1.1, 2.2, 3.3)],
[(4.4, 5.5, 6.6), (4.4, 5.5, 7.7), (4.4, 6.6, 7.7), (5.5, 6.6, 7.7)],
]
assert ak.to_list(ak.combinations(array, 3, replacement=True)) == [
[
(0.0, 0.0, 0.0),
(0.0, 0.0, 1.1),
(0.0, 0.0, 2.2),
(0.0, 0.0, 3.3),
(0.0, 1.1, 1.1),
(0.0, 1.1, 2.2),
(0.0, 1.1, 3.3),
(0.0, 2.2, 2.2),
(0.0, 2.2, 3.3),
(0.0, 3.3, 3.3),
(1.1, 1.1, 1.1),
(1.1, 1.1, 2.2),
(1.1, 1.1, 3.3),
(1.1, 2.2, 2.2),
(1.1, 2.2, 3.3),
(1.1, 3.3, 3.3),
(2.2, 2.2, 2.2),
(2.2, 2.2, 3.3),
(2.2, 3.3, 3.3),
(3.3, 3.3, 3.3),
],
[
(4.4, 4.4, 4.4),
(4.4, 4.4, 5.5),
(4.4, 4.4, 6.6),
(4.4, 4.4, 7.7),
(4.4, 5.5, 5.5),
(4.4, 5.5, 6.6),
(4.4, 5.5, 7.7),
(4.4, 6.6, 6.6),
(4.4, 6.6, 7.7),
(4.4, 7.7, 7.7),
(5.5, 5.5, 5.5),
(5.5, 5.5, 6.6),
(5.5, 5.5, 7.7),
(5.5, 6.6, 6.6),
(5.5, 6.6, 7.7),
(5.5, 7.7, 7.7),
(6.6, 6.6, 6.6),
(6.6, 6.6, 7.7),
(6.6, 7.7, 7.7),
(7.7, 7.7, 7.7),
],
]
def process(self, events):
# Dataset parameters
dataset = events.metadata['dataset']
year = self._samples[dataset]['year']
xsec = self._samples[dataset]['xsec']
sow = self._samples[dataset]['nSumOfWeights']
isData = self._samples[dataset]['isData']
datasets = [
'SingleMuon', 'SingleElectron', 'EGamma', 'MuonEG', 'DoubleMuon',
'DoubleElectron'
]
for d in datasets:
if d in dataset: dataset = dataset.split('_')[0]
# Initialize objects
met = events.MET
e = events.Electron
mu = events.Muon
tau = events.Tau
j = events.Jet
# Muon selection
#mu['isGood'] = isMuonMVA(mu.pt, mu.eta, mu.dxy, mu.dz, mu.miniPFRelIso_all, mu.sip3d, mu.mvaTTH, mu.mediumPromptId, mu.tightCharge, minpt=10)
mu['isPres'] = isPresMuon(mu.dxy, mu.dz, mu.sip3d, mu.looseId)
mu['isTight'] = isTightMuon(mu.pt,
mu.eta,
mu.dxy,
mu.dz,
mu.pfRelIso03_all,
mu.sip3d,
mu.mvaTTH,
mu.mediumPromptId,
mu.tightCharge,
mu.looseId,
minpt=10)
mu['isGood'] = mu['isPres'] & mu['isTight']
leading_mu = mu[ak.argmax(mu.pt, axis=-1, keepdims=True)]
leading_mu = leading_mu[leading_mu.isGood]
mu = mu[mu.isGood]
mu_pres = mu[mu.isPres]
# Electron selection
#e['isGood'] = isElecMVA(e.pt, e.eta, e.dxy, e.dz, e.miniPFRelIso_all, e.sip3d, e.mvaTTH, e.mvaFall17V2Iso, e.lostHits, e.convVeto, e.tightCharge, minpt=10)
e['isPres'] = isPresElec(e.pt,
e.eta,
e.dxy,
e.dz,
e.miniPFRelIso_all,
e.sip3d,
e.lostHits,
minpt=15)
e['isTight'] = isTightElec(e.pt,
e.eta,
e.dxy,
e.dz,
e.miniPFRelIso_all,
e.sip3d,
e.mvaTTH,
e.mvaFall17V2Iso,
e.lostHits,
e.convVeto,
e.tightCharge,
e.sieie,
e.hoe,
e.eInvMinusPInv,
minpt=15)
e['isClean'] = isClean(e, mu, drmin=0.05)
e['isGood'] = e['isPres'] & e['isTight'] & e['isClean']
leading_e = e[ak.argmax(e.pt, axis=-1, keepdims=True)]
leading_e = leading_e[leading_e.isGood]
e = e[e.isGood]
e_pres = e[e.isPres & e.isClean]
# Tau selection
tau['isPres'] = isPresTau(tau.pt,
tau.eta,
tau.dxy,
tau.dz,
tau.leadTkPtOverTauPt,
tau.idAntiMu,
tau.idAntiEle,
tau.rawIso,
tau.idDecayModeNewDMs,
minpt=20)
tau['isClean'] = isClean(tau, e_pres, drmin=0.4) & isClean(
tau, mu_pres, drmin=0.4)
tau['isGood'] = tau['isPres'] # & tau['isClean'], for the moment
tau = tau[tau.isGood]
nElec = ak.num(e)
nMuon = ak.num(mu)
nTau = ak.num(tau)
twoLeps = (nElec + nMuon) == 2
threeLeps = (nElec + nMuon) == 3
twoElec = (nElec == 2)
twoMuon = (nMuon == 2)
e0 = e[ak.argmax(e.pt, axis=-1, keepdims=True)]
m0 = mu[ak.argmax(mu.pt, axis=-1, keepdims=True)]
# Jet selection
jetptname = 'pt_nom' if hasattr(j, 'pt_nom') else 'pt'
j['isGood'] = isTightJet(getattr(j,
jetptname), j.eta, j.jetId, j.neHEF,
j.neEmEF, j.chHEF, j.chEmEF, j.nConstituents)
#j['isgood'] = isGoodJet(j.pt, j.eta, j.jetId)
#j['isclean'] = isClean(j, e, mu)
j['isClean'] = isClean(j, e, drmin=0.4) & isClean(
j, mu, drmin=0.4) # & isClean(j, tau, drmin=0.4)
goodJets = j[(j.isClean) & (j.isGood)]
njets = ak.num(goodJets)
ht = ak.sum(goodJets.pt, axis=-1)
j0 = goodJets[ak.argmax(goodJets.pt, axis=-1, keepdims=True)]
#nbtags = ak.num(goodJets[goodJets.btagDeepFlavB > 0.2770])
nbtags = ak.num(goodJets[goodJets.btagDeepB > 0.4941])
##################################################################
### 2 same-sign leptons
##################################################################
# emu
singe = e[(nElec == 1) & (nMuon == 1) & (e.pt > -1)]
singm = mu[(nElec == 1) & (nMuon == 1) & (mu.pt > -1)]
em = ak.cartesian({"e": singe, "m": singm})
emSSmask = (em.e.charge * em.m.charge > 0)
emSS = em[emSSmask]
nemSS = len(ak.flatten(emSS))
# ee and mumu
# pt>-1 to preserve jagged dimensions
ee = e[(nElec == 2) & (nMuon == 0) & (e.pt > -1)]
mm = mu[(nElec == 0) & (nMuon == 2) & (mu.pt > -1)]
eepairs = ak.combinations(ee, 2, fields=["e0", "e1"])
eeSSmask = (eepairs.e0.charge * eepairs.e1.charge > 0)
eeonZmask = (np.abs((eepairs.e0 + eepairs.e1).mass - 91.2) < 10)
eeoffZmask = (eeonZmask == 0)
mmpairs = ak.combinations(mm, 2, fields=["m0", "m1"])
mmSSmask = (mmpairs.m0.charge * mmpairs.m1.charge > 0)
mmonZmask = (np.abs((mmpairs.m0 + mmpairs.m1).mass - 91.2) < 10)
mmoffZmask = (mmonZmask == 0)
eeSSonZ = eepairs[eeSSmask & eeonZmask]
eeSSoffZ = eepairs[eeSSmask & eeoffZmask]
mmSSonZ = mmpairs[mmSSmask & mmonZmask]
mmSSoffZ = mmpairs[mmSSmask & mmoffZmask]
neeSS = len(ak.flatten(eeSSonZ)) + len(ak.flatten(eeSSoffZ))
nmmSS = len(ak.flatten(mmSSonZ)) + len(ak.flatten(mmSSoffZ))
print('Same-sign events [ee, emu, mumu] = [%i, %i, %i]' %
(neeSS, nemSS, nmmSS))
# Cuts
eeSSmask = (ak.num(eeSSmask[eeSSmask]) > 0)
mmSSmask = (ak.num(mmSSmask[mmSSmask]) > 0)
eeonZmask = (ak.num(eeonZmask[eeonZmask]) > 0)
eeoffZmask = (ak.num(eeoffZmask[eeoffZmask]) > 0)
mmonZmask = (ak.num(mmonZmask[mmonZmask]) > 0)
mmoffZmask = (ak.num(mmoffZmask[mmoffZmask]) > 0)
emSSmask = (ak.num(emSSmask[emSSmask]) > 0)
##################################################################
### 3 leptons
##################################################################
# eem
muon_eem = mu[(nElec == 2) & (nMuon == 1) & (mu.pt > -1)]
elec_eem = e[(nElec == 2) & (nMuon == 1) & (e.pt > -1)]
ee_eem = ak.combinations(elec_eem, 2, fields=["e0", "e1"])
ee_eemZmask = (ee_eem.e0.charge * ee_eem.e1.charge < 1) & (np.abs(
(ee_eem.e0 + ee_eem.e1).mass - 91.2) < 10)
ee_eemOffZmask = (ee_eem.e0.charge * ee_eem.e1.charge < 1) & (np.abs(
(ee_eem.e0 + ee_eem.e1).mass - 91.2) > 10)
ee_eemZmask = (ak.num(ee_eemZmask[ee_eemZmask]) > 0)
ee_eemOffZmask = (ak.num(ee_eemOffZmask[ee_eemOffZmask]) > 0)
eepair_eem = (ee_eem.e0 + ee_eem.e1)
trilep_eem = eepair_eem + muon_eem #ak.cartesian({"e0":ee_eem.e0,"e1":ee_eem.e1, "m":muon_eem})
# mme
muon_mme = mu[(nElec == 1) & (nMuon == 2) & (mu.pt > -1)]
elec_mme = e[(nElec == 1) & (nMuon == 2) & (e.pt > -1)]
mm_mme = ak.combinations(muon_mme, 2, fields=["m0", "m1"])
mm_mmeZmask = (mm_mme.m0.charge * mm_mme.m1.charge < 1) & (np.abs(
(mm_mme.m0 + mm_mme.m1).mass - 91.2) < 10)
mm_mmeOffZmask = (mm_mme.m0.charge * mm_mme.m1.charge < 1) & (np.abs(
(mm_mme.m0 + mm_mme.m1).mass - 91.2) > 10)
mm_mmeZmask = (ak.num(mm_mmeZmask[mm_mmeZmask]) > 0)
mm_mmeOffZmask = (ak.num(mm_mmeOffZmask[mm_mmeOffZmask]) > 0)
mmpair_mme = (mm_mme.m0 + mm_mme.m1)
trilep_mme = mmpair_mme + elec_mme
mZ_mme = mmpair_mme.mass
mZ_eem = eepair_eem.mass
m3l_eem = trilep_eem.mass
m3l_mme = trilep_mme.mass
# eee and mmm
eee = e[(nElec == 3) & (nMuon == 0) & (e.pt > -1)]
mmm = mu[(nElec == 0) & (nMuon == 3) & (mu.pt > -1)]
eee_leps = ak.combinations(eee, 3, fields=["e0", "e1", "e2"])
mmm_leps = ak.combinations(mmm, 3, fields=["m0", "m1", "m2"])
ee_pairs = ak.combinations(eee, 2, fields=["e0", "e1"])
mm_pairs = ak.combinations(mmm, 2, fields=["m0", "m1"])
ee_pairs_index = ak.argcombinations(eee, 2, fields=["e0", "e1"])
mm_pairs_index = ak.argcombinations(mmm, 2, fields=["m0", "m1"])
mmSFOS_pairs = mm_pairs[
(np.abs(mm_pairs.m0.pdgId) == np.abs(mm_pairs.m1.pdgId))
& (mm_pairs.m0.charge != mm_pairs.m1.charge)]
offZmask_mm = ak.all(
np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2) > 10.,
axis=1,
keepdims=True) & (ak.num(mmSFOS_pairs) > 0)
onZmask_mm = ak.any(
np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2) < 10.,
axis=1,
keepdims=True)
eeSFOS_pairs = ee_pairs[
(np.abs(ee_pairs.e0.pdgId) == np.abs(ee_pairs.e1.pdgId))
& (ee_pairs.e0.charge != ee_pairs.e1.charge)]
offZmask_ee = ak.all(
np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2) > 10,
axis=1,
keepdims=True) & (ak.num(eeSFOS_pairs) > 0)
onZmask_ee = ak.any(
np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2) < 10,
axis=1,
keepdims=True)
# Create masks **for event selection**
eeeOnZmask = (ak.num(onZmask_ee[onZmask_ee]) > 0)
eeeOffZmask = (ak.num(offZmask_ee[offZmask_ee]) > 0)
mmmOnZmask = (ak.num(onZmask_mm[onZmask_mm]) > 0)
mmmOffZmask = (ak.num(offZmask_mm[offZmask_mm]) > 0)
# Now we need to create masks for the leptons in order to select leptons from the Z boson candidate (in onZ categories)
ZeeMask = ak.argmin(np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass -
91.2),
axis=1,
keepdims=True)
ZmmMask = ak.argmin(np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass -
91.2),
axis=1,
keepdims=True)
Zee = eeSFOS_pairs[ZeeMask]
Zmm = mmSFOS_pairs[ZmmMask]
eZ0 = Zee.e0[ak.num(eeSFOS_pairs) > 0]
eZ1 = Zee.e1[ak.num(eeSFOS_pairs) > 0]
eZ = eZ0 + eZ1
mZ0 = Zmm.m0[ak.num(mmSFOS_pairs) > 0]
mZ1 = Zmm.m1[ak.num(mmSFOS_pairs) > 0]
mZ = mZ0 + mZ1
mZ_eee = eZ.mass
mZ_mmm = mZ.mass
# And for the W boson
ZmmIndices = mm_pairs_index[ZmmMask]
ZeeIndices = ee_pairs_index[ZeeMask]
eW = eee[~ZeeIndices.e0 | ~ZeeIndices.e1]
mW = mmm[~ZmmIndices.m0 | ~ZmmIndices.m1]
triElec = eee_leps.e0 + eee_leps.e1 + eee_leps.e2
triMuon = mmm_leps.m0 + mmm_leps.m1 + mmm_leps.m2
m3l_eee = triElec.mass
m3l_mmm = triMuon.mass
# Triggers
trig_eeSS = passTrigger(events, 'ee', isData, dataset)
trig_mmSS = passTrigger(events, 'mm', isData, dataset)
trig_emSS = passTrigger(events, 'em', isData, dataset)
trig_eee = passTrigger(events, 'eee', isData, dataset)
trig_mmm = passTrigger(events, 'mmm', isData, dataset)
trig_eem = passTrigger(events, 'eem', isData, dataset)
trig_mme = passTrigger(events, 'mme', isData, dataset)
# MET filters
# Weights
genw = np.ones_like(
events['MET_pt']) if isData else events['genWeight']
weights = coffea.analysis_tools.Weights(len(events))
weights.add('norm', genw if isData else (xsec / sow) * genw)
eftweights = events['EFTfitCoefficients'] if hasattr(
events, "EFTfitCoefficients") else []
# Selections and cuts
selections = PackedSelection()
channels2LSS = ['eeSSonZ', 'eeSSoffZ', 'mmSSonZ', 'mmSSoffZ', 'emSS']
selections.add('eeSSonZ', (eeonZmask) & (eeSSmask) & (trig_eeSS))
selections.add('eeSSoffZ', (eeoffZmask) & (eeSSmask) & (trig_eeSS))
selections.add('mmSSonZ', (mmonZmask) & (mmSSmask) & (trig_mmSS))
selections.add('mmSSoffZ', (mmoffZmask) & (mmSSmask) & (trig_mmSS))
selections.add('emSS', (emSSmask) & (trig_emSS))
channels3L = ['eemSSonZ', 'eemSSoffZ', 'mmeSSonZ', 'mmeSSoffZ']
selections.add('eemSSonZ', (ee_eemZmask) & (trig_eem))
selections.add('eemSSoffZ', (ee_eemOffZmask) & (trig_eem))
selections.add('mmeSSonZ', (mm_mmeZmask) & (trig_mme))
selections.add('mmeSSoffZ', (mm_mmeOffZmask) & (trig_mme))
channels3L += ['eeeSSonZ', 'eeeSSoffZ', 'mmmSSonZ', 'mmmSSoffZ']
selections.add('eeeSSonZ', (eeeOnZmask) & (trig_eee))
selections.add('eeeSSoffZ', (eeeOffZmask) & (trig_eee))
selections.add('mmmSSonZ', (mmmOnZmask) & (trig_mmm))
selections.add('mmmSSoffZ', (mmmOffZmask) & (trig_mmm))
levels = ['base', '2jets', '4jets', '4j1b', '4j2b']
selections.add('base', (nElec + nMuon >= 2))
selections.add('2jets', (njets >= 2))
selections.add('4jets', (njets >= 4))
selections.add('4j1b', (njets >= 4) & (nbtags >= 1))
selections.add('4j2b', (njets >= 4) & (nbtags >= 2))
# Variables
invMass_eeSSonZ = (eeSSonZ.e0 + eeSSonZ.e1).mass
invMass_eeSSoffZ = (eeSSoffZ.e0 + eeSSoffZ.e1).mass
invMass_mmSSonZ = (mmSSonZ.m0 + mmSSonZ.m1).mass
invMass_mmSSoffZ = (mmSSoffZ.m0 + mmSSoffZ.m1).mass
invMass_emSS = (emSS.e + emSS.m).mass
varnames = {}
varnames['met'] = met.pt
varnames['ht'] = ht
varnames['njets'] = njets
varnames['nbtags'] = nbtags
varnames['invmass'] = {
'eeSSonZ': invMass_eeSSonZ,
'eeSSoffZ': invMass_eeSSoffZ,
'mmSSonZ': invMass_mmSSonZ,
'mmSSoffZ': invMass_mmSSoffZ,
'emSS': invMass_emSS,
'eemSSonZ': mZ_eem,
'eemSSoffZ': mZ_eem,
'mmeSSonZ': mZ_mme,
'mmeSSoffZ': mZ_mme,
'eeeSSonZ': mZ_eee,
'eeeSSoffZ': mZ_eee,
'mmmSSonZ': mZ_mmm,
'mmmSSoffZ': mZ_mmm,
}
varnames['m3l'] = {
'eemSSonZ': m3l_eem,
'eemSSoffZ': m3l_eem,
'mmeSSonZ': m3l_mme,
'mmeSSoffZ': m3l_mme,
'eeeSSonZ': m3l_eee,
'eeeSSoffZ': m3l_eee,
'mmmSSonZ': m3l_mmm,
'mmmSSoffZ': m3l_mmm,
}
varnames['e0pt'] = e0.pt
varnames['e0eta'] = e0.eta
varnames['m0pt'] = m0.pt
varnames['m0eta'] = m0.eta
varnames['j0pt'] = j0.pt
varnames['j0eta'] = j0.eta
varnames['counts'] = np.ones_like(events.MET.pt)
# fill Histos
hout = self.accumulator.identity()
allweights = weights.weight().flatten(
) # Why does it not complain about .flatten() here?
hout['SumOfEFTweights'].fill(eftweights,
sample=dataset,
SumOfEFTweights=varnames['counts'],
weight=allweights)
for var, v in varnames.items():
for ch in channels2LSS + channels3L:
for lev in levels:
weight = weights.weight()
cuts = [ch] + [lev]
cut = selections.all(*cuts)
weights_flat = weight[cut].flatten(
) # Why does it not complain about .flatten() here?
weights_ones = np.ones_like(weights_flat, dtype=np.int)
eftweightsvalues = eftweights[cut] if len(
eftweights) > 0 else []
if var == 'invmass':
if ch in ['eeeSSoffZ', 'mmmSSoffZ']: continue
elif ch in ['eeeSSonZ', 'mmmSSonZ']:
continue #values = v[ch]
else:
values = ak.flatten(v[ch][cut])
hout['invmass'].fill(sample=dataset,
channel=ch,
cut=lev,
invmass=values,
weight=weights_flat)
elif var == 'm3l':
if ch in [
'eeSSonZ', 'eeSSoffZ', 'mmSSonZ', 'mmSSoffZ',
'emSS', 'eeeSSoffZ', 'mmmSSoffZ', 'eeeSSonZ',
'mmmSSonZ'
]:
continue
values = ak.flatten(v[ch][cut])
hout['m3l'].fill(eftweightsvalues,
sample=dataset,
channel=ch,
cut=lev,
m3l=values,
weight=weights_flat)
else:
values = v[cut]
if var == 'ht':
hout[var].fill(eftweightsvalues,
ht=values,
sample=dataset,
channel=ch,
cut=lev,
weight=weights_flat)
elif var == 'met':
hout[var].fill(eftweightsvalues,
met=values,
sample=dataset,
channel=ch,
cut=lev,
weight=weights_flat)
elif var == 'njets':
hout[var].fill(eftweightsvalues,
njets=values,
sample=dataset,
channel=ch,
cut=lev,
weight=weights_flat)
elif var == 'nbtags':
hout[var].fill(eftweightsvalues,
nbtags=values,
sample=dataset,
channel=ch,
cut=lev,
weight=weights_flat)
elif var == 'counts':
hout[var].fill(counts=values,
sample=dataset,
channel=ch,
cut=lev,
weight=weights_ones)
elif var == 'j0eta':
if lev == 'base': continue
values = ak.flatten(values)
#values=np.asarray(values)
hout[var].fill(eftweightsvalues,
j0eta=values,
sample=dataset,
channel=ch,
cut=lev,
weight=weights_flat)
elif var == 'e0pt':
if ch in [
'mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ',
'mmmSSonZ'
]:
continue
values = ak.flatten(values)
#values=np.asarray(values)
hout[var].fill(
eftweightsvalues,
e0pt=values,
sample=dataset,
channel=ch,
cut=lev,
weight=weights_flat
) # Crashing here, not sure why. Related to values?
elif var == 'm0pt':
if ch in [
'eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ',
'eeeSSonZ'
]:
continue
values = ak.flatten(values)
#values=np.asarray(values)
hout[var].fill(eftweightsvalues,
m0pt=values,
sample=dataset,
channel=ch,
cut=lev,
weight=weights_flat)
elif var == 'e0eta':
if ch in [
'mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ',
'mmmSSonZ'
]:
continue
values = ak.flatten(values)
#values=np.asarray(values)
hout[var].fill(eftweightsvalues,
e0eta=values,
sample=dataset,
channel=ch,
cut=lev,
weight=weights_flat)
elif var == 'm0eta':
if ch in [
'eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ',
'eeeSSonZ'
]:
continue
values = ak.flatten(values)
#values=np.asarray(values)
hout[var].fill(eftweightsvalues,
m0eta=values,
sample=dataset,
channel=ch,
cut=lev,
weight=weights_flat)
elif var == 'j0pt':
if lev == 'base': continue
values = ak.flatten(values)
#values=np.asarray(values)
hout[var].fill(eftweightsvalues,
j0pt=values,
sample=dataset,
channel=ch,
cut=lev,
weight=weights_flat)
return hout
def get_evt_p4(p4, num=6):
combos = ak.combinations(p4, num)
part0, part1, part2, part3, part4, part5 = ak.unzip(combos)
evt_p4 = part0 + part1 + part2 + part3 + part4 + part5
return evt_p4
