def test_read_nanomc(suffix):
path = os.path.abspath(f'tests/samples/nano_dy.{suffix}')
factory = getattr(NanoEventsFactory, f'from_{suffix}')(path)
events = factory.events()
# test after views first
genroundtrips(events.GenPart.mask[events.GenPart.eta > 0])
genroundtrips(events.mask[ak.any(events.Electron.pt > 50, axis=1)].GenPart)
genroundtrips(events.GenPart)
genroundtrips(events.GenPart[events.GenPart.eta > 0])
genroundtrips(events[ak.any(events.Electron.pt > 50, axis=1)].GenPart)
# sane gen matching (note for electrons gen match may be photon(22))
assert ak.all((abs(events.Electron.matched_gen.pdgId) == 11) | (events.Electron.matched_gen.pdgId == 22))
assert ak.all(abs(events.Muon.matched_gen.pdgId) == 13)
genroundtrips(events.Electron.matched_gen)
crossref(events[ak.num(events.Jet) > 2])
crossref(events)
if suffix == 'root':
assert ak.any(events.Photon.isTight, axis=1).tolist()[:9] == [False, True, True, True, False, False, False, False, False]
if suffix == 'parquet':
assert ak.any(events.Photon.isTight, axis=1).tolist()[:9] == [False, True, False, True, False, False, False, False, True]
def match(first, second, **kwargs):
if 'dR_cutoff' not in kwargs:
raise Exception("Please specify dR cutoff for matching!")
dR_cutoff = kwargs.pop('dR_cutoff', 0.3)
return_match_properties = kwargs.pop('return_match_properties', False)
etas = ak.cartesian({
'first': first.eta,
'second': second.eta
},
axis=1,
nested=True)
phis = ak.cartesian({
'first': first.phi,
'second': second.phi
},
axis=1,
nested=True)
dR, deta, dphi = delta_r(etas['first'], etas['second'], phis['first'],
phis['second'])
min_idx = ak.argmin(dR, axis=2)
match_properties = {
'pt': second.pt[min_idx],
'eta': second.eta[min_idx],
'phi': second.phi[min_idx],
}
if return_match_properties:
return ak.any(dR < dR_cutoff, axis=2), match_properties
else:
return ak.any(dR < dR_cutoff, axis=2)
def test_any_all():
array = ak.Array(
[
[
[
np.datetime64("2022"),
np.datetime64("2023"),
np.datetime64("2025")
],
[],
[np.datetime64("2027"),
np.datetime64("2011")],
[np.datetime64("2013")],
],
[],
[[np.datetime64("2017"),
np.datetime64("2019")], [np.datetime64("2023")]],
],
check_valid=True,
)
assert ak.to_list(ak.any(array, axis=-1)) == [
[True, False, True, True],
[],
[True, True],
]
assert ak.to_list(ak.any(array, axis=-2)) == [[True, True, True], [],
[True, True]]
def test_read_nanomc(suffix):
path = os.path.abspath(f"tests/samples/nano_dy.{suffix}")
# parquet files were converted from even older nanoaod
nanoversion = NanoAODSchema.v6 if suffix == "root" else NanoAODSchema.v5
factory = getattr(NanoEventsFactory, f"from_{suffix}")(
path, schemaclass=nanoversion
)
events = factory.events()
# test after views first
genroundtrips(events.GenPart.mask[events.GenPart.eta > 0])
genroundtrips(events.mask[ak.any(events.Electron.pt > 50, axis=1)].GenPart)
genroundtrips(events.GenPart)
genroundtrips(events.GenPart[events.GenPart.eta > 0])
genroundtrips(events[ak.any(events.Electron.pt > 50, axis=1)].GenPart)
# sane gen matching (note for electrons gen match may be photon(22))
assert ak.all(
(abs(events.Electron.matched_gen.pdgId) == 11)
| (events.Electron.matched_gen.pdgId == 22)
)
assert ak.all(abs(events.Muon.matched_gen.pdgId) == 13)
genroundtrips(events.Electron.matched_gen)
crossref(events[ak.num(events.Jet) > 2])
crossref(events)
# test issue 409
assert ak.to_list(events[[]].Photon.mass) == []
if suffix == "root":
assert ak.any(events.Photon.isTight, axis=1).tolist()[:9] == [
False,
True,
True,
True,
False,
False,
False,
False,
False,
]
if suffix == "parquet":
assert ak.any(events.Photon.isTight, axis=1).tolist()[:9] == [
False,
True,
False,
True,
False,
False,
False,
False,
True,
]
def genroundtrips(genpart):
# check genpart roundtrip
assert ak.all(genpart.children.parent.pdgId == genpart.pdgId)
assert ak.all(ak.any(genpart.parent.children.pdgId == genpart.pdgId, axis=-1, mask_identity=True))
# distinctParent should be distinct and it should have a relevant child
assert ak.all(genpart.distinctParent.pdgId != genpart.pdgId)
assert ak.all(ak.any(genpart.distinctParent.children.pdgId == genpart.pdgId, axis=-1, mask_identity=True))
# exercise hasFlags
genpart.hasFlags(['isHardProcess'])
genpart.hasFlags(['isHardProcess', 'isDecayedLeptonHadron'])
def get_lepton_sf(year: str, lepton: str, corrections, pt, eta):
sf_dict = corrections[lepton]
eta_ranges = sf_dict["eta_ranges"]
reco_cen = np.ones(len(pt))
reco_err = np.ones(len(pt))
trig_cen = np.ones(len(pt))
trig_err = np.ones(len(pt))
for idx, eta_range in enumerate(eta_ranges):
mask = (eta >= eta_range[0]) & (eta < eta_range[1]) # find inds that are within given eta range
if not ak.any(mask): continue # no values fall within eta range
recoSF_cen = sf_dict["Reco_ID"]["Central"]["eta_bin%i" % idx]
recoSF_err = sf_dict["Reco_ID"]["Error"]["eta_bin%i" % idx]
trigSF_cen = sf_dict["Trig"]["Central"]["eta_bin%i" % idx]
trigSF_err = sf_dict["Trig"]["Error"]["eta_bin%i" % idx]
reco_cen[mask] = recoSF_cen(pt[mask])
reco_err[mask] = recoSF_err(pt[mask])
trig_cen[mask] = trigSF_cen(pt[mask])
trig_err[mask] = trigSF_err(pt[mask])
output_SFs = {
"RECO_CEN" : reco_cen,
"RECO_ERR" : reco_err,
"TRIG_CEN" : trig_cen,
"TRIG_ERR" : trig_err,
}
return output_SFs
def get_triggers(HLT, leptype, year, noIso=False, accumulator=None):
## event triggers to be used found here: https://twiki.cern.ch/twiki/bin/view/CMS/TopTriggerYear2016 or 2017, 2018...
allowed_leptypes = ['Muon', 'Electron']
if leptype not in allowed_leptypes:
raise ValueError(f"Select 'leptype' from {allowed_leptypes}")
iso_cat = 'noIso' if noIso else 'Iso'
pass_triggers = ak.any((HLT[i] for i in single_mu_trigger_paths[year][iso_cat] if i in HLT.fields), axis=0) if leptype == 'Muon'\
else ak.any((HLT[i] for i in single_el_trigger_paths[year][iso_cat] if i in HLT.fields), axis=0)
if accumulator:
accumulator['cutflow']['nEvts pass %s pass_triggers' % leptype] += ak.sum(pass_triggers)
return pass_triggers, accumulator
else:
return pass_triggers
def test_reducers():
# axis=None reducers are implemented in NumPy.
assert ak.sum(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]])) == 6 + 6j
assert ak.prod(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]])) == -12 + 12j
# axis != None reducers are implemented in libawkward; this should be ReducerSum.
assert ak.sum(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]),
axis=1).tolist() == [
3 + 3j,
0 + 0j,
3 + 3j,
]
# And this is in ReducerProd.
assert ak.prod(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]),
axis=1).tolist() == [
0 + 4j,
1 + 0j,
3 + 3j,
]
# ReducerCount, ReducerCountNonzero, ReducerAny, and ReducerAll work.
assert ak.count(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]),
axis=1).tolist() == [2, 0, 1]
assert ak.count_nonzero(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]),
axis=1).tolist() == [2, 0, 1]
assert ak.any(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]),
axis=1).tolist() == [
True,
False,
True,
]
assert ak.all(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]),
axis=1).tolist() == [
True,
True,
True,
]
assert ak.any(ak.from_iter([[1 + 1j, 2 + 2j, 0 + 0j], [], [3 + 3j]]),
axis=1).tolist() == [True, False, True]
assert ak.all(ak.from_iter([[1 + 1j, 2 + 2j, 0 + 0j], [], [3 + 3j]]),
axis=1).tolist() == [False, True, True]
def SS_selection(lep1, lep2):
selection = PackedSelection()
is_dilep = ((ak.num(lep1) + ak.num(lep2))==2)
pos_charge = ((ak.sum(lep1.pdgId, axis=1) + ak.sum(lep2.pdgId, axis=1))<0)
neg_charge = ((ak.sum(lep1.pdgId, axis=1) + ak.sum(lep2.pdgId, axis=1))>0)
dilep2 = choose(lep2, 2)
dilep1 = choose(lep1, 2)
dilep = cross(lep2, lep1)
is_SS = ( ak.any((dilep2['0'].charge * dilep2['1'].charge)>0, axis=1) | \
ak.any((dilep1['0'].charge * dilep1['1'].charge)>0, axis=1) | \
ak.any((dilep['0'].charge * dilep['1'].charge)>0, axis=1) )
selection.add('SS', is_SS)
ss_reqs = ['SS']
ss_reqs_d = {sel: True for sel in ss_reqs}
ss_selection = selection.require(**ss_reqs_d)
return ss_selection
def in_sensitive_volume(events, sensitive_volumes):
"""
Function which identifies which events are inside sensitive volumes.
Further, tests if sensitive volumes overlap. Assigns volume id, and
xenon density to interactions.
Args:
events (ak.records): Awkward record of the interactions.
sensitive_volumes (dict): Dictionary of the different volumes
defined via the SensitiveVolume class.
Returns:
ak.array: Awkward array containing the event ids.
"""
if len(events) == 0:
res_det_dtype = [('xe_density', 'float64'),
('vol_id', 'int64'),
('create_S2', 'bool'),
]
return ak.from_numpy(np.empty(0, dtype=res_det_dtype))
for ind, vol in enumerate(sensitive_volumes):
res = ak.ArrayBuilder()
res = _inside_sens_vol(events['x'],
events['y'],
events['z'],
vol.roi,
vol.volume_id,
vol.xe_density,
vol.create_S2,
res)
if ind:
# Now we add the other results, but first test if
# volumes overlap. Only possible if we explicitly loop
# over everything. This reduces performance but adds layer of
# safety.
m = (result['vol_id'] > 0) & (res['vol_id'] == vol.volume_id)
if ak.any(m):
overlapping_id = result[m][0]
# Get volume name:
name = [vol.name for vol in sensitive_volumes if vol.volume_id == overlapping_id][0]
raise ValueError(f'The volume {vol.name} is overlapping with'
f' volume {name}!')
new_results = res.snapshot()
for field in result.fields:
# Workaround since we cannot sum up records-arrays anymore
result[field] = result[field] + new_results[field]
else:
# First result initiates the array
result = res.snapshot()
return result
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, events):
output = pd.DataFrame()
events = events[ak.any(events.Muon.pt > 20, axis=-1)]
output['run'] = ak.to_numpy(events.run)
output['event'] = ak.to_numpy(events.event)
output['dataset'] = events.metadata['dataset']
output['mu1_pt'] = -999.0
output['mu2_pt'] = -999.0
muons = events.Muon[events.Muon.pt > 20]
counts = ak.num(muons)
one_muon = ak.to_numpy(counts > 0)
two_muons = ak.to_numpy(counts > 1)
output.loc[one_muon, 'mu1_pt'] = ak.to_numpy(muons[one_muon][:, 0].pt)
output.loc[two_muons, 'mu2_pt'] = ak.to_numpy(muons[two_muons][:,
1].pt)
return output
def studyAcc(ifile, title, acceptances_dict, mass):
dq_events = getData(ifile, "Truth")
# choosing the nevts for denominator (because all files have 10000 evts)
den = len(dq_events)
# choosing 500 evts for denominator
den500 = False
if den500:
dq_events = dq_events[:500]
den = 500
dq_showers = dq_events["Showers"]
dq_electrons = dq_events["Electrons"]
dq_sh_z = ak.fill_none(ak.pad_none(dq_showers.sz, 2, axis=1), 0)
dq_sh_e = ak.fill_none(ak.pad_none(dq_showers.sedep_ecal, 2, axis=1), -1)
dq_e_ge = ak.fill_none(ak.pad_none(dq_electrons.ge, 2, axis=1), -1)
shower_zero_mask = ak.any(dq_sh_e == 0, axis=1)
shower_zero_index = np.where(shower_zero_mask)
eshower_mask = ak.all(dq_sh_e > 0.1 * dq_e_ge, axis=1)
eshower_index = np.where(eshower_mask)
masks = {'shower_zero': shower_zero_index, 'eshower': eshower_index}
dq_evts = dq_events[ak.all(dq_e_ge > -1, axis=1)]
# print(acceptances_dict)
# we can print the mass, the number of events with gen e-, the number of events in the sample, and the denominator
print(mass, len(dq_evts), len(dq_events), den)
for maskstr, m in masks.items():
if maskstr in acceptances_dict:
acceptances_dict[maskstr].append(len(m[0]) / den)
else:
acceptances_dict[maskstr] = [len(m[0]) / den]
#print(maskstr+':', len(m[0])/den)
return acceptances_dict
def match2(first, second, deltaRCut=0.4):
drCut2 = deltaRCut**2
combs = ak.cartesian([first, second], nested=True)
return ak.any((combs['0'].delta_r2(combs['1']) < drCut2), axis=2)
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 bosonFlavor(bosons):
childid = abs(bosons.children.pdgId)
genflavor = ak.any(childid == 5, axis=-1) * 3 + ak.any(
childid == 4, axis=-1) * 2 + ak.all(childid < 4, axis=-1) * 1
return ak.fill_none(genflavor, 0)
def process(self, events):
output = self.accumulator.identity()
dataset = events.metadata['dataset']
output['sumw'][dataset] += ak.sum(np.sign(events.Generator.weight))
output['nevents'][dataset] += len(events)
output = self.accumulator.identity()
dataset = events.metadata['dataset']
if dataset not in ['singleelectron','singlemuon','egamma']:
output['sumw'][dataset] += ak.sum(np.sign(events.Generator.weight))
output['nevents'][dataset] += len(events)
if dataset in ['singleelectron','singlemuon','egamma']:
events = events[lumimask(events.run,events.luminosityBlock)]
events = events[(events.PuppiMET.pt > 30) | (events.PuppiMET.ptJERUp > 30) | (events.PuppiMET.ptJESUp > 30)]
if year == "2016":
if dataset == 'singlemuon':
events = events[events.HLT.IsoTkMu24 | events.HLT.IsoMu24]
elif dataset == 'singleelectron':
events = events[vents.HLT.IsoTkMu24 | events.HLT.IsoMu24 | events.HLT.Ele27_WPTight_Gsf]
else:
events = events[events.HLT.IsoTkMu24 | events.HLT.IsoMu24 | events.HLT.Ele27_WPTight_Gsf]
elif year == "2017":
if dataset == 'singlemuon':
events = events[events.HLT.IsoMu27]
elif dataset == 'singleelectron':
events = events[events.HLT.Ele32_WPTight_Gsf_L1DoubleEG]
else:
events = events[events.HLT.IsoMu27 | events.HLT.Ele32_WPTight_Gsf_L1DoubleEG]
elif year == "2018":
if dataset == 'singlemuon':
events = events[events.HLT.IsoMu24]
elif dataset == 'egamma':
events = events[events.HLT.Ele32_WPTight_Gsf]
else:
events = events[events.HLT.IsoMu24 |events.HLT.Ele32_WPTight_Gsf]
events = events[(ak.num(events.Jet) > 3) | ((ak.num(events.Jet) > 1) & (ak.num(events.FatJet) > 0))]
events = events[(ak.num(events.Electron) > 0) | (ak.num(events.Muon) > 0)]
tight_muons = events.Muon[events.Muon.tightId & (events.Muon.pfRelIso04_all < 0.15) & (events.Muon.pt > 26) & (abs(events.Muon.eta) < 2.4)]
loose_not_tight_muons = events.Muon[events.Muon.tightId & (events.Muon.pfRelIso04_all < 0.4) & (events.Muon.pfRelIso04_all > 0.15) & (events.Muon.pt > 20) & (abs(events.Muon.eta) < 2.4)]
tight_electrons = events.Electron[(events.Electron.pt > 30) & (events.Electron.cutBased >= 3) & (events.Electron.eta + events.Electron.deltaEtaSC < 2.5) & ((abs(events.Electron.dz) < 0.1) & (abs(events.Electron.dxy) < 0.05) & (events.Electron.eta + events.Electron.deltaEtaSC < 1.479)) | ((abs(events.Electron.dz) < 0.2) & (abs(events.Electron.dxy) < 0.1) & (events.Electron.eta + events.Electron.deltaEtaSC > 1.479))]
name_map = jec_stack.blank_name_map
name_map['JetPt'] = 'pt'
name_map['JetMass'] = 'mass'
name_map['JetEta'] = 'eta'
name_map['JetA'] = 'area'
jets = events.Jet
jets['pt_raw'] = (1 - jets['rawFactor']) * jets['pt']
jets['mass_raw'] = (1 - jets['rawFactor']) * jets['mass']
jets['pt_gen'] = ak.values_astype(ak.fill_none(jets.matched_gen.pt, 0), np.float32)
jets['rho'] = ak.broadcast_arrays(events.fixedGridRhoFastjetAll, jets.pt)[0]
name_map['ptGenJet'] = 'pt_gen'
name_map['ptRaw'] = 'pt_raw'
name_map['massRaw'] = 'mass_raw'
name_map['Rho'] = 'rho'
events_cache = events.caches[0]
jet_factory = CorrectedJetsFactory(name_map, jec_stack)
corrected_jets = jet_factory.build(jets, lazy_cache=events_cache)
jet_pt = corrected_jets.pt
jet_pt_jesup = corrected_jets.JES_jes.up.pt
jet_pt_jerup = corrected_jets.JER.up.pt
corrected_jets = ak.zip({
"pt": corrected_jets.pt,
"eta": corrected_jets.eta,
"phi": corrected_jets.phi,
"mass": corrected_jets.mass,
"charge": np.ones(len(corrected_jets.pt)),
"btagDeepB": corrected_jets.btagDeepB
}, with_name="PtEtaPhiMCandidate")
fatjets = events.FatJet[(events.FatJet.pt > 250) & (abs(events.FatJet.eta) < 2.5) & (events.FatJet.msoftdrop > 50) & (events.FatJet.msoftdrop < 150)]
b_jets = corrected_jets[(events.Jet.cleanmask == 1) & (jet_pt > 30) & (abs(events.Jet.eta) < 2.5) & (events.Jet.btagDeepB > 0.8953)]
vbf_jets = corrected_jets[(events.Jet.cleanmask == 1) & (jet_pt > 30) & (abs(events.Jet.eta) < 4.7) & (events.Jet.btagDeepB < 0.2217)]
nextrajets = ak.num(events.Jet[(events.Jet.cleanmask == 1) & (jet_pt > 30) & (abs(events.Jet.eta) < 4.7)]) - 4
nextrabjets = ak.num(events.Jet[(events.Jet.cleanmask == 1) & (jet_pt > 30) & (abs(events.Jet.eta) < 4.7) & (events.Jet.btagDeepB > 0.2217)]) - 2
basecut_merged = (ak.num(fatjets) > 0) & (ak.num(vbf_jets) > 1) & (ak.num(tight_muons) + ak.num(tight_electrons) == 1) & (ak.num(loose_not_tight_muons) == 0) & (events.PuppiMET.pt > 30)
events_merged = events[basecut_merged]
fatjets_merged = fatjets[basecut_merged]
vbf_jets_merged = vbf_jets[basecut_merged]
tight_muons_merged = tight_muons[basecut_merged]
tight_electrons_merged = tight_electrons[basecut_merged]
nextrajets_merged = nextrajets[basecut_merged]
nextrabjets_merged = nextrabjets[basecut_merged]
basecut = (ak.num(b_jets) > 1) & (ak.num(vbf_jets) > 1) & (ak.num(tight_muons) + ak.num(tight_electrons) == 1) & (ak.num(loose_not_tight_muons) == 0) & (events.PuppiMET.pt > 30)
events = events[basecut]
b_jets = b_jets[basecut]
vbf_jets = vbf_jets[basecut]
tight_muons = tight_muons[basecut]
tight_electrons = tight_electrons[basecut]
nextrajets = nextrajets[basecut]
nextrabjets = nextrabjets[basecut]
if dataset in ['singleelectron','singlemuon','egamma']:
dataset = 'data'
if ak.any(basecut_merged):
cut7 = (fatjets_merged[:,0].mass > 50) & (fatjets_merged[:,0].mass < 150) & ((vbf_jets_merged[:,0]+vbf_jets_merged[:,1]).mass > 500) & (abs(vbf_jets_merged[:,0].eta - vbf_jets_merged[:,1].eta) > 2.5) & (ak.num(tight_muons_merged) > 0)
cut8 = (fatjets_merged[:,0].mass > 50) & (fatjets_merged[:,0].mass < 150) & ((vbf_jets_merged[:,0]+vbf_jets_merged[:,1]).mass > 500) & (abs(vbf_jets_merged[:,0].eta - vbf_jets_merged[:,1].eta) > 2.5) & (ak.num(tight_electrons_merged) > 0)
# cut9 = cut7 | cut8
cut1 = ((b_jets[:,0] + b_jets[:,1]).mass > 50) & ((b_jets[:,0] + b_jets[:,1]).mass < 150) & ((vbf_jets[:,0] + vbf_jets[:,1]).mass > 500) & (abs(vbf_jets[:,0].eta - vbf_jets[:,1].eta) > 2.5) & (ak.num(tight_muons) > 0)
cut2 = ((b_jets[:,0] + b_jets[:,1]).mass > 50) & ((b_jets[:,0] + b_jets[:,1]).mass < 150) & ((vbf_jets[:,0] + vbf_jets[:,1]).mass > 500) & (abs(vbf_jets[:,0].eta - vbf_jets[:,1].eta) > 2.5) & (ak.num(tight_electrons) > 0)
# cut3 = cut1 | cut2
if ak.any(basecut_merged) and ak.any(cut7):
sel7_events = events_merged[cut7]
sel7_fatjets = fatjets_merged[cut7]
sel7_vbf_jets = vbf_jets_merged[cut7]
sel7_muons = tight_muons_merged[cut7][:,0]
sel7_nextrajets = nextrajets_merged[cut7]
sel7_nextrabjets = nextrabjets_merged[cut7]
output["weights_merged"][dataset] += processor.column_accumulator(np.sign(ak.to_numpy(sel7_events.Generator.weight).data))
output['variables_merged'][dataset] += processor.column_accumulator(np.transpose(np.vstack((
ak.to_numpy(sel7_fatjets[:,0].pt),
ak.to_numpy(sel7_fatjets[:,0].eta),
ak.to_numpy(sel7_fatjets[:,0].phi),
ak.to_numpy(sel7_fatjets[:,0].btagDeepB),
ak.to_numpy(sel7_fatjets[:,0].btagHbb),
ak.to_numpy(sel7_fatjets[:,0].msoftdrop),
ak.to_numpy(sel7_nextrajets),
ak.to_numpy(sel7_nextrabjets),
np.zeros(len(sel7_events)),
np.sign(ak.to_numpy(sel7_muons.charge)+1),
ak.to_numpy(sel7_muons.pt),
ak.to_numpy(sel7_muons.eta),
ak.to_numpy(sel7_muons.phi),
ak.to_numpy(sel7_events.PuppiMET.pt),
ak.to_numpy(sel7_events.PuppiMET.phi),
ak.to_numpy(sel7_vbf_jets[:,0].pt),
ak.to_numpy(sel7_vbf_jets[:,1].pt),
ak.to_numpy(sel7_vbf_jets[:,0].eta),
ak.to_numpy(sel7_vbf_jets[:,1].eta),
ak.to_numpy(sel7_vbf_jets[:,0].phi),
ak.to_numpy(sel7_vbf_jets[:,1].phi),
ak.to_numpy(sel7_vbf_jets[:,0].btagDeepB),
ak.to_numpy(sel7_vbf_jets[:,1].btagDeepB),
ak.to_numpy((sel7_vbf_jets[:,0]+sel7_vbf_jets[:,1]).mass),
ak.to_numpy(sel7_vbf_jets[:,0].eta - sel7_vbf_jets[:,1].eta),
ak.to_numpy(np.sqrt(2*(sel7_muons+sel7_vbf_jets[:,0]).pt*sel7_events.PuppiMET.pt*(1 - np.cos(sel7_events.PuppiMET.phi - (sel7_muons+sel7_vbf_jets[:,0]).phi)))),
ak.to_numpy(np.sqrt(2*(sel7_muons+sel7_vbf_jets[:,1]).pt*sel7_events.PuppiMET.pt*(1 - np.cos(sel7_events.PuppiMET.phi - (sel7_muons+sel7_vbf_jets[:,1]).phi))))))))
sel7_muonidsf = evaluator['muonidsf'](abs(sel7_muons.eta), sel7_muons.pt)
sel7_muonisosf = evaluator['muonisosf'](abs(sel7_muons.eta), sel7_muons.pt)
sel7_muonhltsf = evaluator['muonhltsf'](abs(sel7_muons.eta), sel7_muons.pt)
sel7_weight = np.sign(sel7_events.Generator.weight)*sel7_events.L1PreFiringWeight.Nom*sel7_muonidsf*sel7_muonisosf*sel7_muonhltsf
if ak.any(basecut_merged) and ak.any(cut8):
sel8_events = events_merged[cut8]
sel8_fatjets = fatjets_merged[cut8]
sel8_vbf_jets = vbf_jets_merged[cut8]
sel8_electrons = tight_electrons_merged[cut8][:,0]
sel8_nextrajets = nextrajets_merged[cut8]
sel8_nextrabjets = nextrabjets_merged[cut8]
output["weights_merged"][dataset] += processor.column_accumulator(np.sign(ak.to_numpy(sel8_events.Generator.weight).data))
output['variables_merged'][dataset] += processor.column_accumulator(np.transpose(np.vstack((
ak.to_numpy(sel8_fatjets[:,0].pt),
ak.to_numpy(sel8_fatjets[:,0].eta),
ak.to_numpy(sel8_fatjets[:,0].phi),
ak.to_numpy(sel8_fatjets[:,0].btagDeepB),
ak.to_numpy(sel8_fatjets[:,0].btagHbb),
ak.to_numpy(sel8_fatjets[:,0].msoftdrop),
ak.to_numpy(sel8_nextrajets),
ak.to_numpy(sel8_nextrabjets),
np.ones(len(sel8_events)),
np.sign(ak.to_numpy(sel8_electrons.charge)+1),
ak.to_numpy(sel8_electrons.pt),
ak.to_numpy(sel8_electrons.eta),
ak.to_numpy(sel8_electrons.phi),
ak.to_numpy(sel8_events.PuppiMET.pt),
ak.to_numpy(sel8_events.PuppiMET.phi),
ak.to_numpy(sel8_vbf_jets[:,0].pt),
ak.to_numpy(sel8_vbf_jets[:,1].pt),
ak.to_numpy(sel8_vbf_jets[:,0].eta),
ak.to_numpy(sel8_vbf_jets[:,1].eta),
ak.to_numpy(sel8_vbf_jets[:,0].phi),
ak.to_numpy(sel8_vbf_jets[:,1].phi),
ak.to_numpy(sel8_vbf_jets[:,0].btagDeepB),
ak.to_numpy(sel8_vbf_jets[:,1].btagDeepB),
ak.to_numpy((sel8_vbf_jets[:,0]+sel8_vbf_jets[:,1]).mass),
ak.to_numpy(sel8_vbf_jets[:,0].eta - sel8_vbf_jets[:,1].eta),
ak.to_numpy(np.sqrt(2*(sel8_electrons+sel8_vbf_jets[:,0]).pt*sel8_events.PuppiMET.pt*(1 - np.cos(sel8_events.PuppiMET.phi - (sel8_electrons+sel8_vbf_jets[:,0]).phi)))),
ak.to_numpy(np.sqrt(2*(sel8_electrons+sel8_vbf_jets[:,1]).pt*sel8_events.PuppiMET.pt*(1 - np.cos(sel8_events.PuppiMET.phi - (sel8_electrons+sel8_vbf_jets[:,1]).phi))))))))
sel8_electronidsf = evaluator['electronidsf'](sel8_electrons.eta, sel8_electrons.pt)
sel8_electronrecosf = evaluator['electronrecosf'](sel8_electrons.eta, sel8_electrons.pt)
sel8_weight = np.sign(sel8_events.Generator.weight)*sel8_events.L1PreFiringWeight.Nom*sel8_electronidsf*sel8_electronrecosf
if ak.any(basecut) and ak.any(cut1):
sel1_events = events[cut1]
sel1_b_jets = b_jets[cut1]
sel1_vbf_jets = vbf_jets[cut1]
sel1_muons = tight_muons[cut1][:,0]
sel1_nextrajets = nextrajets[cut1]
sel1_nextrabjets = nextrabjets[cut1]
output["weights"][dataset] += processor.column_accumulator(np.sign(ak.to_numpy(sel1_events.Generator.weight).data))
output['variables'][dataset] += processor.column_accumulator(np.transpose(np.vstack((
ak.to_numpy(sel1_nextrajets),
ak.to_numpy(sel1_nextrabjets),
np.zeros(len(sel1_events)),
np.sign(ak.to_numpy(sel1_muons.charge)+1),
ak.to_numpy(sel1_muons.pt),
ak.to_numpy(sel1_muons.eta),
ak.to_numpy(sel1_muons.phi),
ak.to_numpy(sel1_events.PuppiMET.pt),
ak.to_numpy(sel1_events.PuppiMET.phi),
ak.to_numpy(sel1_b_jets[:,0].pt),
ak.to_numpy(sel1_b_jets[:,1].pt),
ak.to_numpy(sel1_vbf_jets[:,0].pt),
ak.to_numpy(sel1_vbf_jets[:,1].pt),
ak.to_numpy(sel1_b_jets[:,0].eta),
ak.to_numpy(sel1_b_jets[:,1].eta),
ak.to_numpy(sel1_vbf_jets[:,0].eta),
ak.to_numpy(sel1_vbf_jets[:,1].eta),
ak.to_numpy(sel1_b_jets[:,0].phi),
ak.to_numpy(sel1_b_jets[:,1].phi),
ak.to_numpy(sel1_vbf_jets[:,0].phi),
ak.to_numpy(sel1_vbf_jets[:,1].phi),
ak.to_numpy(sel1_b_jets[:,0].btagDeepB),
ak.to_numpy(sel1_b_jets[:,1].btagDeepB),
ak.to_numpy(sel1_vbf_jets[:,0].btagDeepB),
ak.to_numpy(sel1_vbf_jets[:,1].btagDeepB),
ak.to_numpy((sel1_b_jets[:,0]+sel1_b_jets[:,1]).mass),
ak.to_numpy((sel1_vbf_jets[:,0]+sel1_vbf_jets[:,1]).mass),
ak.to_numpy(sel1_vbf_jets[:,0].eta - sel1_vbf_jets[:,1].eta),
ak.to_numpy(np.sqrt(2*(sel1_muons+sel1_b_jets[:,0]).pt*sel1_events.PuppiMET.pt*(1 - np.cos(sel1_events.PuppiMET.phi - (sel1_muons+sel1_b_jets[:,0]).phi)))),ak.to_numpy(np.sqrt(2*(sel1_muons+sel1_b_jets[:,1]).pt*sel1_events.PuppiMET.pt*(1 - np.cos(sel1_events.PuppiMET.phi - (sel1_muons+sel1_b_jets[:,1]).phi))))))))
sel1_pu_weight = evaluator['pileup'](sel1_events.Pileup.nTrueInt)
sel1_muonidsf = evaluator['muonidsf'](abs(sel1_muons.eta), sel1_muons.pt)
sel1_muonisosf = evaluator['muonisosf'](abs(sel1_muons.eta), sel1_muons.pt)
sel1_muonhltsf = evaluator['muonhltsf'](abs(sel1_muons.eta), sel1_muons.pt)
sel1_weight = np.sign(sel1_events.Generator.weight)*sel1_pu_weight*sel1_events.L1PreFiringWeight.Nom*sel1_muonidsf*sel1_muonisosf*sel1_muonhltsf
if ak.any(basecut) and ak.any(cut2):
sel2_events = events[cut2]
sel2_b_jets = b_jets[cut2]
sel2_vbf_jets = vbf_jets[cut2]
sel2_electrons = tight_electrons[cut2][:,0]
sel2_nextrajets = nextrajets[cut2]
sel2_nextrabjets = nextrabjets[cut2]
output["weights"][dataset] += processor.column_accumulator(np.sign(ak.to_numpy(sel2_events.Generator.weight).data))
output['variables'][dataset] += processor.column_accumulator(np.transpose(np.vstack((
ak.to_numpy(sel2_nextrajets),
ak.to_numpy(sel2_nextrabjets),
np.ones(len(sel2_events)),
np.sign(ak.to_numpy(sel2_electrons.charge)+1),
ak.to_numpy(sel2_electrons.pt),
ak.to_numpy(sel2_electrons.eta),
ak.to_numpy(sel2_electrons.phi),
ak.to_numpy(sel2_events.PuppiMET.pt),
ak.to_numpy(sel2_events.PuppiMET.phi),
ak.to_numpy(sel2_b_jets[:,0].pt),
ak.to_numpy(sel2_b_jets[:,1].pt),
ak.to_numpy(sel2_vbf_jets[:,0].pt),
ak.to_numpy(sel2_vbf_jets[:,1].pt),
ak.to_numpy(sel2_b_jets[:,0].eta),
ak.to_numpy(sel2_b_jets[:,1].eta),
ak.to_numpy(sel2_vbf_jets[:,0].eta),
ak.to_numpy(sel2_vbf_jets[:,1].eta),
ak.to_numpy(sel2_b_jets[:,0].phi),
ak.to_numpy(sel2_b_jets[:,1].phi),
ak.to_numpy(sel2_vbf_jets[:,0].phi),
ak.to_numpy(sel2_vbf_jets[:,1].phi),
ak.to_numpy(sel2_b_jets[:,0].btagDeepB),
ak.to_numpy(sel2_b_jets[:,1].btagDeepB),
ak.to_numpy(sel2_vbf_jets[:,0].btagDeepB),
ak.to_numpy(sel2_vbf_jets[:,1].btagDeepB),
ak.to_numpy((sel2_b_jets[:,0]+sel2_b_jets[:,1]).mass),
ak.to_numpy((sel2_vbf_jets[:,0]+sel2_vbf_jets[:,1]).mass),
ak.to_numpy(sel2_vbf_jets[:,0].eta - sel2_vbf_jets[:,1].eta),
ak.to_numpy(np.sqrt(2*(sel2_electrons+sel2_b_jets[:,0]).pt*sel2_events.PuppiMET.pt*(1 - np.cos(sel2_events.PuppiMET.phi - (sel2_electrons+sel2_b_jets[:,0]).phi)))),ak.to_numpy(np.sqrt(2*(sel2_electrons+sel2_b_jets[:,1]).pt*sel2_events.PuppiMET.pt*(1 - np.cos(sel2_events.PuppiMET.phi - (sel2_electrons+sel2_b_jets[:,1]).phi))))))))
sel2_pu_weight = evaluator['pileup'](sel2_events.Pileup.nTrueInt)
sel2_electronidsf = evaluator['electronidsf'](sel2_electrons.eta, sel2_electrons.pt)
sel2_electronrecosf = evaluator['electronrecosf'](sel2_electrons.eta, sel2_electrons.pt)
sel2_weight = np.sign(sel2_events.Generator.weight)*sel2_pu_weight*sel2_events.L1PreFiringWeight.Nom*sel2_electronidsf*sel2_electronrecosf
return output
def process(self, events):
output = self.accumulator.identity()
self.sample_name = events.metadata["dataset"]
## make event weights
# data or MC distinction made internally
evt_weights = MCWeights.get_event_weights(events,
year=args.year,
corrections=self.corrections)
all_gen_tops = events["GenPart"][events["GenPart"].pdgId == 6]
all_gen_tbars = events["GenPart"][events["GenPart"].pdgId == -6]
# find all stages of decay/hadronization process of tops (and tbars) for which there are only at most 1 top passing per event (but not all events are empty)
tops_dict = {
flag: all_gen_tops[all_gen_tops.hasFlags([flag])]
for flag in self.FLAGS if
(ak.all(ak.num(all_gen_tops[all_gen_tops.hasFlags([flag])]) < 2)
& ak.any(ak.num(all_gen_tops[all_gen_tops.hasFlags([flag])]) > 0))
}
tbars_dict = {
flag: all_gen_tbars[all_gen_tbars.hasFlags([flag])]
for flag in self.FLAGS
if (ak.all(
ak.num(all_gen_tbars[all_gen_tbars.hasFlags([flag])]) < 2)
& ak.any(
ak.num(all_gen_tbars[all_gen_tbars.hasFlags([flag])]) > 0))
}
for flag in tops_dict.keys():
if not ak.all(ak.num(tops_dict[flag]) == ak.num(tbars_dict[flag])):
print(
f"Gen tops and tbars don't have same number of entries for {flag}"
)
continue
evt_mask = ak.num(tops_dict[flag]) > 0
wts = evt_weights.weight()[evt_mask]
gen_tops = tops_dict[flag][evt_mask]
gen_tbars = tbars_dict[flag][evt_mask]
gen_ttbars = gen_tops + gen_tbars
#set_trace()
top_ctstar, tbar_ctstar = make_vars.ctstar(gen_tops, gen_tbars)
if "Int" in self.sample_name:
pos_evts = np.where(wts > 0)
# tops
output = self.fill_genp_hists(accumulator=output,
dname="%s_pos" %
self.sample_name,
genp_type="Top",
flag=flag,
obj=gen_tops[pos_evts],
evt_weights=wts[pos_evts])
# tbars
output = self.fill_genp_hists(accumulator=output,
dname="%s_pos" %
self.sample_name,
genp_type="Tbar",
flag=flag,
obj=gen_tbars[pos_evts],
evt_weights=wts[pos_evts])
# ttbars
output = self.fill_genp_hists(accumulator=output,
dname="%s_pos" %
self.sample_name,
genp_type="TTbar",
flag=flag,
obj=gen_ttbars[pos_evts],
evt_weights=wts[pos_evts])
neg_evts = np.where(wts < 0)
# tops
output = self.fill_genp_hists(accumulator=output,
dname="%s_neg" %
self.sample_name,
genp_type="Top",
flag=flag,
obj=gen_tops[neg_evts],
evt_weights=wts[neg_evts])
# tbars
output = self.fill_genp_hists(accumulator=output,
dname="%s_neg" %
self.sample_name,
genp_type="Tbar",
flag=flag,
obj=gen_tbars[neg_evts],
evt_weights=wts[neg_evts])
# ttbars
output = self.fill_genp_hists(accumulator=output,
dname="%s_neg" %
self.sample_name,
genp_type="TTbar",
flag=flag,
obj=gen_ttbars[neg_evts],
evt_weights=wts[neg_evts])
# fill ctstar hists
output["ctstar"].fill(dataset="%s_pos" % self.sample_name,
objtype="Top",
flag=flag,
ctstar=ak.flatten(top_ctstar[pos_evts],
axis=None),
weight=wts[pos_evts])
output["ctstar"].fill(dataset="%s_pos" % self.sample_name,
objtype="Tbar",
flag=flag,
ctstar=ak.flatten(tbar_ctstar[pos_evts],
axis=None),
weight=wts[pos_evts])
output["ctstar"].fill(dataset="%s_neg" % self.sample_name,
objtype="Top",
flag=flag,
ctstar=ak.flatten(top_ctstar[neg_evts],
axis=None),
weight=wts[neg_evts])
output["ctstar"].fill(dataset="%s_neg" % self.sample_name,
objtype="Tbar",
flag=flag,
ctstar=ak.flatten(tbar_ctstar[neg_evts],
axis=None),
weight=wts[neg_evts])
output["ctstar_abs"].fill(dataset="%s_pos" % self.sample_name,
objtype="Top",
flag=flag,
ctstar_abs=np.abs(
ak.flatten(top_ctstar[pos_evts],
axis=None)),
weight=wts[pos_evts])
output["ctstar_abs"].fill(dataset="%s_pos" % self.sample_name,
objtype="Tbar",
flag=flag,
ctstar_abs=np.abs(
ak.flatten(tbar_ctstar[pos_evts],
axis=None)),
weight=wts[pos_evts])
output["ctstar_abs"].fill(dataset="%s_neg" % self.sample_name,
objtype="Top",
flag=flag,
ctstar_abs=np.abs(
ak.flatten(top_ctstar[neg_evts],
axis=None)),
weight=wts[neg_evts])
output["ctstar_abs"].fill(dataset="%s_neg" % self.sample_name,
objtype="Tbar",
flag=flag,
ctstar_abs=np.abs(
ak.flatten(tbar_ctstar[neg_evts],
axis=None)),
weight=wts[neg_evts])
else:
# tops
output = self.fill_genp_hists(accumulator=output,
dname=self.sample_name,
genp_type="Top",
flag=flag,
obj=gen_tops,
evt_weights=wts)
# tbars
output = self.fill_genp_hists(accumulator=output,
dname=self.sample_name,
genp_type="Tbar",
flag=flag,
obj=gen_tbars,
evt_weights=wts)
# ttbars
output = self.fill_genp_hists(accumulator=output,
dname=self.sample_name,
genp_type="TTbar",
flag=flag,
obj=gen_ttbars,
evt_weights=wts)
# fill ctstar hists
output["ctstar"].fill(dataset=self.sample_name,
objtype="Top",
flag=flag,
ctstar=ak.flatten(top_ctstar, axis=None),
weight=wts)
output["ctstar"].fill(dataset=self.sample_name,
objtype="Tbar",
flag=flag,
ctstar=ak.flatten(tbar_ctstar,
axis=None),
weight=wts)
output["ctstar_abs"].fill(dataset=self.sample_name,
objtype="Top",
flag=flag,
ctstar_abs=np.abs(
ak.flatten(top_ctstar,
axis=None)),
weight=wts)
output["ctstar_abs"].fill(dataset=self.sample_name,
objtype="Tbar",
flag=flag,
ctstar_abs=np.abs(
ak.flatten(tbar_ctstar,
axis=None)),
weight=wts)
return output
def select_normal(genparts, w_decay_momid):
columns = ["pt", "eta", "phi", "mass", "pdgId", "charge", "decaytype"]
# get tops, defined as last copy
gen_tops = genparts[(genparts.hasFlags(['isLastCopy']))
& (genparts.pdgId == 6)]
gen_tops['charge'] = ak.ones_like(gen_tops.pt) * (2. / 3.)
gen_tbars = genparts[(genparts.hasFlags(['isLastCopy']))
& (genparts.pdgId == -6)]
gen_tbars['charge'] = ak.ones_like(gen_tbars.pt) * (-2. / 3.)
# get direct top decay products (children are "isHardProcess" and "isFirstCopy")
gen_bs = ak.flatten(gen_tops.children[(gen_tops.children.pdgId == 5)],
axis=2)
gen_bs['charge'] = ak.ones_like(gen_bs.pt) * (-1. / 3.)
gen_bbars = ak.flatten(
gen_tbars.children[(gen_tbars.children.pdgId == -5)], axis=2)
gen_bbars['charge'] = ak.ones_like(gen_bbars.pt) * (1. / 3.)
gen_wplus = ak.flatten(
gen_tops.children[(gen_tops.children.pdgId == w_decay_momid)], axis=2)
gen_wplus['charge'] = ak.ones_like(gen_wplus.pt)
gen_wminus = ak.flatten(
gen_tbars.children[(gen_tbars.children.pdgId == -1 * w_decay_momid)],
axis=2)
gen_wminus['charge'] = ak.ones_like(gen_wminus.pt) * (-1.)
# get w decay products
# get last copy of W bosons
last_gen_wplus = ak.flatten(gen_tops.distinctChildren[
(gen_tops.distinctChildren.pdgId == w_decay_momid)
& (gen_tops.distinctChildren.hasFlags(['isLastCopy']))],
axis=2)
last_gen_wminus = ak.flatten(gen_tbars.distinctChildren[
(gen_tbars.distinctChildren.pdgId == -1 * w_decay_momid)
& (gen_tbars.distinctChildren.hasFlags(['isLastCopy']))],
axis=2)
gen_wplus_decaytype = np.zeros(len(gen_wplus))
gen_wminus_decaytype = np.zeros(len(gen_wminus))
# up/down partons from last W+
gen_wpartons_up_fromWplus = ak.flatten(
last_gen_wplus.children[(np.mod(last_gen_wplus.children.pdgId, 2) == 0)
& (np.abs(last_gen_wplus.children.pdgId) < 6)],
axis=2)
gen_wpartons_up_fromWplus['charge'] = np.sign(
gen_wpartons_up_fromWplus.pdgId) * (2. / 3.)
gen_wplus_decaytype[ak.num(gen_wpartons_up_fromWplus) > 0] = np.ones(
ak.sum(ak.num(gen_wpartons_up_fromWplus) > 0)) * 2
gen_wpartons_dw_fromWplus = ak.flatten(
last_gen_wplus.children[(np.mod(last_gen_wplus.children.pdgId, 2) == 1)
& (np.abs(last_gen_wplus.children.pdgId) < 6)],
axis=2)
gen_wpartons_dw_fromWplus['charge'] = np.sign(
gen_wpartons_up_fromWplus.pdgId) * (-1. / 3.)
# up/down partons from last W-
gen_wpartons_up_fromWminus = ak.flatten(last_gen_wminus.children[
(np.mod(last_gen_wminus.children.pdgId, 2) == 0)
& (np.abs(last_gen_wminus.children.pdgId) < 6)],
axis=2)
gen_wpartons_up_fromWminus['charge'] = np.sign(
gen_wpartons_up_fromWminus.pdgId) * (2. / 3.)
gen_wminus_decaytype[ak.num(gen_wpartons_up_fromWminus) > 0] = np.ones(
ak.sum(ak.num(gen_wpartons_up_fromWminus) > 0)) * 2
gen_wpartons_dw_fromWminus = ak.flatten(last_gen_wminus.children[
(np.mod(last_gen_wminus.children.pdgId, 2) == 1)
& (np.abs(last_gen_wminus.children.pdgId) < 6)],
axis=2)
gen_wpartons_dw_fromWminus['charge'] = np.sign(
gen_wpartons_up_fromWminus.pdgId) * (-1. / 3.)
# charged leps from last W+
gen_charged_leps_fromWplus = ak.flatten(
last_gen_wplus.children[(np.abs(last_gen_wplus.children.pdgId) == 11) |
(np.abs(last_gen_wplus.children.pdgId) == 13) |
(np.abs(last_gen_wplus.children.pdgId) == 15)],
axis=2)
gen_charged_leps_fromWplus['charge'] = ak.ones_like(
gen_charged_leps_fromWplus.pdgId)
gen_wplus_decaytype[ak.num(gen_charged_leps_fromWplus) > 0] = np.ones(
ak.sum(ak.num(gen_charged_leps_fromWplus) > 0))
# add decaytype (0 is INVALID, 1 is LEPTONIC, 2 is HADRONIC)
gen_wplus['decaytype'] = ak.unflatten(gen_wplus_decaytype,
ak.num(gen_wplus))
gen_tops['decaytype'] = gen_wplus['decaytype'] # set decaytype for tops
gen_bs['decaytype'] = gen_wplus['decaytype'] # set decaytype for bs
# neutral leps from last W+
gen_neutral_leps_fromWplus = ak.flatten(
last_gen_wplus.children[(np.abs(last_gen_wplus.children.pdgId) == 12) |
(np.abs(last_gen_wplus.children.pdgId) == 14) |
(np.abs(last_gen_wplus.children.pdgId) == 16)],
axis=2)
gen_neutral_leps_fromWplus['charge'] = ak.zeros_like(
gen_neutral_leps_fromWplus.pt)
# charged leps from last W-
gen_charged_leps_fromWminus = ak.flatten(last_gen_wminus.children[
(np.abs(last_gen_wminus.children.pdgId) == 11) |
(np.abs(last_gen_wminus.children.pdgId) == 13) |
(np.abs(last_gen_wminus.children.pdgId) == 15)],
axis=2)
gen_charged_leps_fromWminus['charge'] = ak.ones_like(
gen_charged_leps_fromWminus.pdgId) * (-1.)
gen_wminus_decaytype[ak.num(gen_charged_leps_fromWminus) > 0] = np.ones(
ak.sum(ak.num(gen_charged_leps_fromWminus) > 0))
# add decaytype (0 is INVALID, 1 is LEPTONIC, 2 is HADRONIC)
gen_wminus['decaytype'] = ak.unflatten(gen_wminus_decaytype,
ak.num(gen_wminus))
gen_tbars['decaytype'] = gen_wminus['decaytype'] # set decaytype for tbars
gen_bbars['decaytype'] = gen_wminus['decaytype'] # set decaytype for bbars
# neutral leps from last W-
gen_neutral_leps_fromWminus = ak.flatten(last_gen_wminus.children[
(np.abs(last_gen_wminus.children.pdgId) == 12) |
(np.abs(last_gen_wminus.children.pdgId) == 14) |
(np.abs(last_gen_wminus.children.pdgId) == 16)],
axis=2)
gen_neutral_leps_fromWminus['charge'] = ak.zeros_like(
gen_neutral_leps_fromWminus.pt)
gen_wpartons_up = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
gen_wpartons_dw = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
gen_charged_leps = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
gen_neutral_leps = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
for column in columns:
if column == 'decaytype': continue
gen_wpartons_up[column] = ak.flatten(
ak.concatenate([
gen_wpartons_up_fromWplus[column],
gen_wpartons_up_fromWminus[column]
],
axis=1)) # (up-type partons from W+, W-)
gen_wpartons_dw[column] = ak.flatten(
ak.concatenate([
gen_wpartons_dw_fromWplus[column],
gen_wpartons_dw_fromWminus[column]
],
axis=1)) # (dw-type partons from W+, W-)
gen_charged_leps[column] = ak.flatten(
ak.concatenate([
gen_charged_leps_fromWplus[column],
gen_charged_leps_fromWminus[column]
],
axis=1)) # (charged leps from W+, W-)
gen_neutral_leps[column] = ak.flatten(
ak.concatenate([
gen_neutral_leps_fromWplus[column],
gen_neutral_leps_fromWminus[column]
],
axis=1)) # (neutral leps from W+, W-)
gen_wpartons_up = ak.unflatten(
gen_wpartons_up,
ak.num(gen_wpartons_up_fromWplus) + ak.num(gen_wpartons_up_fromWminus))
gen_wpartons_dw = ak.unflatten(
gen_wpartons_dw,
ak.num(gen_wpartons_dw_fromWplus) + ak.num(gen_wpartons_dw_fromWminus))
gen_charged_leps = ak.unflatten(
gen_charged_leps,
ak.num(gen_charged_leps_fromWplus) +
ak.num(gen_charged_leps_fromWminus))
gen_neutral_leps = ak.unflatten(
gen_neutral_leps,
ak.num(gen_neutral_leps_fromWplus) +
ak.num(gen_neutral_leps_fromWminus))
gen_taus = gen_charged_leps[np.abs(gen_charged_leps.pdgId) == 15]
gen_nu_taus = gen_neutral_leps[np.abs(gen_neutral_leps.pdgId) == 16]
# fully hadronic evts
had_evts = (ak.num(gen_charged_leps) == 0) & (
ak.num(gen_neutral_leps)
== 0) & (ak.num(gen_wpartons_up) == 2) & (ak.num(gen_wpartons_dw) == 2)
#set_trace()
# get direct tau decay products from hard processes (subset of gen_taus events above)
tau_decay_prods = genparts[genparts.hasFlags(
['isDirectHardProcessTauDecayProduct'])]
# only need decays to leptons (e/mu, tau nu) for event classification
tau_TO_tau_nu = tau_decay_prods[np.abs(tau_decay_prods.pdgId) == 16]
tau_TO_charged_lep = tau_decay_prods[(np.abs(tau_decay_prods.pdgId) == 11)
|
(np.abs(tau_decay_prods.pdgId) == 13)]
tau_TO_neutral_lep = tau_decay_prods[(np.abs(tau_decay_prods.pdgId) == 12)
|
(np.abs(tau_decay_prods.pdgId) == 14)]
# set decaytype for gen taus
charged_lep_decaytype_array = ak.to_numpy(
ak.flatten(ak.zeros_like(gen_charged_leps['pt'])))
# semilep evts
semilep_evts = (ak.num(gen_charged_leps) == 1) & (
ak.num(gen_neutral_leps)
== 1) & (ak.num(gen_wpartons_up) == 1) & (ak.num(gen_wpartons_dw) == 1)
tau_jets_evts = semilep_evts & (ak.num(gen_taus) == 1)
sl_evts_mask = np.repeat(ak.to_numpy(semilep_evts),
ak.to_numpy(ak.num(gen_charged_leps)))
semilep_decaytype_array = np.zeros(ak.to_numpy(semilep_evts).sum(),
dtype=int)
# tau -> l
semilep_tau_leptonic_decay = (tau_jets_evts) & (
ak.num(tau_TO_charged_lep) == 1) & (
ak.num(tau_TO_neutral_lep) == 1) & (ak.num(tau_TO_tau_nu) == 1)
semilep_tau_hadronic_decay = (tau_jets_evts) & (
~semilep_tau_leptonic_decay)
semilep_decaytype_array[
semilep_tau_leptonic_decay[semilep_evts]] = np.ones(
ak.to_numpy(semilep_tau_leptonic_decay).sum(), dtype=int)
semilep_decaytype_array[
semilep_tau_hadronic_decay[semilep_evts]] = np.ones(
ak.to_numpy(semilep_tau_hadronic_decay).sum(), dtype=int) * 2
# set charged_lep_decatype_array for semileptonic events
charged_lep_decaytype_array[sl_evts_mask] = semilep_decaytype_array
# dilep evts
dilep_evts = (ak.num(gen_charged_leps) == 2) & (
ak.num(gen_neutral_leps)
== 2) & (ak.num(gen_wpartons_up) == 0) & (ak.num(gen_wpartons_dw) == 0)
lep_lep_evts = dilep_evts & (ak.num(gen_taus) == 0)
lep_tau_evts = dilep_evts & (ak.num(gen_taus) == 1)
tau_tau_evts = dilep_evts & (ak.num(gen_taus) == 2)
dl_evts_mask = np.repeat(ak.to_numpy(dilep_evts),
ak.to_numpy(ak.num(gen_charged_leps)))
dilep_decaytype_array = np.zeros((ak.to_numpy(dilep_evts).sum(), 2),
dtype=int)
# tau + tau
# tau + tau -> ll
dilep_TauTau_ll_decay = (tau_tau_evts) & (
ak.num(tau_TO_charged_lep) == 2) & (
ak.num(tau_TO_neutral_lep) == 2) & (ak.num(tau_TO_tau_nu) == 2)
dilep_decaytype_array[dilep_TauTau_ll_decay[dilep_evts]] = np.ones(
(ak.to_numpy(dilep_TauTau_ll_decay).sum(), 2), dtype=int)
# tau + tau -> hh
dilep_TauTau_hh_decay = (tau_tau_evts) & (
(ak.num(tau_TO_charged_lep) + ak.num(tau_TO_neutral_lep))
== 0) & (ak.num(tau_TO_tau_nu) == 2)
dilep_decaytype_array[dilep_TauTau_hh_decay[dilep_evts]] = np.ones(
(ak.to_numpy(dilep_TauTau_hh_decay).sum(), 2), dtype=int) * 2
# tau + tau -> lh
dilep_TauTau_lh_decay = (
tau_tau_evts) & ~(dilep_TauTau_ll_decay | dilep_TauTau_hh_decay)
# set index corresponding to leptonically decaying tau to 1, default array is set to 2
dl_TauTau_to_lh_decaytype_array = np.ones(
ak.to_numpy(dilep_TauTau_lh_decay).sum() * 2, dtype=int) * 2
lep_tau_mask = (np.repeat(
ak.to_numpy(
ak.flatten(
tau_TO_charged_lep[dilep_TauTau_lh_decay].parent.pdgId)),
2) == ak.flatten(gen_charged_leps[dilep_TauTau_lh_decay].pdgId))
dl_TauTau_to_lh_decaytype_array[lep_tau_mask] = np.ones(lep_tau_mask.sum(),
dtype=int)
dilep_decaytype_array[dilep_TauTau_lh_decay[
dilep_evts]] = dl_TauTau_to_lh_decaytype_array.reshape(
ak.to_numpy(dilep_TauTau_lh_decay).sum(), 2)
# lep + tau
# tau -> l
dilep_LepTau_l_decay = (lep_tau_evts) & (
ak.num(tau_TO_charged_lep) == 1) & (
ak.num(tau_TO_neutral_lep) == 1) & (ak.num(tau_TO_tau_nu) == 1)
# set index corresponding to tau to 1
dl_LepTau_to_Lep_decaytype_array = np.zeros(
ak.to_numpy(dilep_LepTau_l_decay).sum() * 2, dtype=int)
dl_LepTau_to_Lep_decaytype_array[ak.flatten(
np.abs(gen_charged_leps[dilep_LepTau_l_decay].pdgId) == 15)] = np.ones(
ak.sum(dilep_LepTau_l_decay), dtype=int)
dilep_decaytype_array[dilep_LepTau_l_decay[
dilep_evts]] = dl_LepTau_to_Lep_decaytype_array.reshape(
ak.sum(dilep_LepTau_l_decay), 2)
# tau -> h
dilep_LepTau_h_decay = (lep_tau_evts) & ~(dilep_LepTau_l_decay)
# set index corresponding to tau to 2
dl_LepTau_to_Had_decaytype_array = np.zeros(ak.sum(dilep_LepTau_h_decay) *
2,
dtype=int)
dl_LepTau_to_Had_decaytype_array[ak.flatten(
np.abs(gen_charged_leps[dilep_LepTau_h_decay].pdgId) == 15)] = np.ones(
ak.sum(dilep_LepTau_h_decay), dtype=int) * 2
dilep_decaytype_array[dilep_LepTau_h_decay[
dilep_evts]] = dl_LepTau_to_Had_decaytype_array.reshape(
ak.sum(dilep_LepTau_h_decay), 2)
# set charged_lep_decatype_array for dileptonic events
charged_lep_decaytype_array[dl_evts_mask] = dilep_decaytype_array.flatten()
# set charged lepton decaytype (defined for taus only, e/mu are 0) (1 is LEPTONIC, 2 is HADRONIC)
gen_charged_leps['decaytype'] = ak.unflatten(charged_lep_decaytype_array,
ak.num(gen_charged_leps))
# make awkward arrays of (top decay prods, tbar decay prods)
Gen_Top_Pairs = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
Gen_B_Pairs = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
Gen_W_Pairs = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
Gen_Wparton_Pairs = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
for column in columns:
Gen_Top_Pairs[column] = ak.flatten(
ak.concatenate([gen_tops[column], gen_tbars[column]],
axis=1)) # (top, tbar)
Gen_B_Pairs[column] = ak.flatten(
ak.concatenate([gen_bs[column], gen_bbars[column]],
axis=1)) # (b, bbar)
Gen_W_Pairs[column] = ak.flatten(
ak.concatenate([gen_wplus[column], gen_wminus[column]],
axis=1)) # (W+, W-)
if column is not "decaytype":
Gen_Wparton_Pairs[column] = ak.flatten(
ak.concatenate(
[
ak.pad_none(gen_wpartons_up[column], 1, axis=1),
ak.pad_none(gen_wpartons_dw[column], 1, axis=1)
],
axis=1)) # (up-type wpartons, down-type wpartons)
Gen_Top_Pairs = ak.unflatten(Gen_Top_Pairs,
ak.num(gen_tops) + ak.num(gen_tbars))
Gen_B_Pairs = ak.unflatten(Gen_B_Pairs, ak.num(gen_bs) + ak.num(gen_bbars))
Gen_W_Pairs = ak.unflatten(Gen_W_Pairs,
ak.num(gen_wplus) + ak.num(gen_wminus))
Gen_Wparton_Pairs = ak.unflatten(
Gen_Wparton_Pairs,
ak.num(ak.pad_none(gen_wpartons_up, 1, axis=1)) +
ak.num(ak.pad_none(gen_wpartons_dw, 1, axis=1)))
Gen_Wparton_Pairs = Gen_Wparton_Pairs[ak.argsort(
Gen_Wparton_Pairs["pt"], ascending=False)] # sort by pt
Gen_TTbar = ak.Array(
{
"pt": (gen_tops + gen_tbars).pt,
"eta": (gen_tops + gen_tbars).eta,
"phi": (gen_tops + gen_tbars).phi,
"mass": (gen_tops + gen_tbars).mass,
"decaytype": gen_tops["decaytype"] + gen_tbars[
"decaytype"], # 0 is for INVALID, 2 for DILEP, 3 for SEMILEP, 4 for HADRONIC
},
with_name="PtEtaPhiMLorentzVector")
## make "table" of gen objects for certain decays
## DILEP
DILEP_evts = ak.zip({
"TTbar":
Gen_TTbar[dilep_evts] if ak.any(dilep_evts) else ak.unflatten(
Gen_TTbar[dilep_evts], ak.values_astype(dilep_evts, int)),
"Top":
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == 1][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == 1][dilep_evts],
ak.values_astype(dilep_evts, int)),
"Tbar":
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == -1][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == -1][dilep_evts],
ak.values_astype(dilep_evts, int)),
"B":
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == -1][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == -1][dilep_evts],
ak.values_astype(dilep_evts, int)),
"Bbar":
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == 1][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == 1][dilep_evts],
ak.values_astype(dilep_evts, int)),
"Wplus":
Gen_W_Pairs[Gen_W_Pairs.charge == 1][dilep_evts] if ak.any(dilep_evts)
else ak.unflatten(Gen_W_Pairs[Gen_W_Pairs.charge == 1][dilep_evts],
ak.values_astype(dilep_evts, int)),
"Wminus":
Gen_W_Pairs[Gen_W_Pairs.charge == -1][dilep_evts] if ak.any(dilep_evts)
else ak.unflatten(Gen_W_Pairs[Gen_W_Pairs.charge == -1][dilep_evts],
ak.values_astype(dilep_evts, int)),
"First_plus":
gen_charged_leps[gen_charged_leps.charge > 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_charged_leps[gen_charged_leps.charge > 0][dilep_evts],
ak.values_astype(dilep_evts,
int)), # charged lepton always made leading
"Second_plus":
gen_neutral_leps[gen_charged_leps.charge > 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_neutral_leps[gen_charged_leps.charge > 0][dilep_evts],
ak.values_astype(dilep_evts,
int)), # neutral lepton always made subleading
"First_minus":
gen_charged_leps[gen_charged_leps.charge < 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_charged_leps[gen_charged_leps.charge < 0][dilep_evts],
ak.values_astype(dilep_evts,
int)), # charged lepton always made leading
"Second_minus":
gen_neutral_leps[gen_charged_leps.charge < 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_neutral_leps[gen_charged_leps.charge < 0][dilep_evts],
ak.values_astype(dilep_evts,
int)), # neutral lepton always made subleading
"Up_plus":
gen_neutral_leps[gen_charged_leps.charge > 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_neutral_leps[gen_charged_leps.charge > 0][dilep_evts],
ak.values_astype(dilep_evts, int)), # same as second plus
"Down_plus":
gen_charged_leps[gen_charged_leps.charge > 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_charged_leps[gen_charged_leps.charge > 0][dilep_evts],
ak.values_astype(dilep_evts, int)), # same as first plus
"Up_minus":
gen_neutral_leps[gen_charged_leps.charge < 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_neutral_leps[gen_charged_leps.charge < 0][dilep_evts],
ak.values_astype(dilep_evts, int)), # same as second minus
"Down_minus":
gen_charged_leps[gen_charged_leps.charge < 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_charged_leps[gen_charged_leps.charge < 0][dilep_evts],
ak.values_astype(dilep_evts, int)), # same as first minus
})
## HAD
HAD_evts = ak.zip({
"TTbar":
Gen_TTbar[had_evts] if ak.any(had_evts) else ak.unflatten(
Gen_TTbar[had_evts], ak.values_astype(had_evts, int)),
"Top":
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == 1][had_evts]
if ak.any(had_evts) else ak.unflatten(
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == 1][had_evts],
ak.values_astype(had_evts, int)),
"Tbar":
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == -1][had_evts]
if ak.any(had_evts) else ak.unflatten(
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == -1][had_evts],
ak.values_astype(had_evts, int)),
"B":
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == -1][had_evts]
if ak.any(had_evts) else ak.unflatten(
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == -1][had_evts],
ak.values_astype(had_evts, int)),
"Bbar":
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == 1][had_evts]
if ak.any(had_evts) else ak.unflatten(
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == 1][had_evts],
ak.values_astype(had_evts, int)),
"Wplus":
Gen_W_Pairs[Gen_W_Pairs.charge == 1][had_evts] if ak.any(had_evts) else
ak.unflatten(Gen_W_Pairs[Gen_W_Pairs.charge == 1][had_evts],
ak.values_astype(had_evts, int)),
"Wminus":
Gen_W_Pairs[Gen_W_Pairs.charge == -1][had_evts] if ak.any(had_evts)
else ak.unflatten(Gen_W_Pairs[Gen_W_Pairs.charge == -1][had_evts],
ak.values_astype(had_evts, int)),
"First_plus":
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge > 0][:,
0]
if ak.any(had_evts) else ak.unflatten(
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge > 0]
[:, 0], ak.values_astype(
had_evts, int)), # leading positively-charged parton
"Second_plus":
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge > 0][:,
1]
if ak.any(had_evts) else ak.unflatten(
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge > 0]
[:, 1], ak.values_astype(
had_evts, int)), # subleading positively-charged parton
"First_minus":
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge < 0][:,
0]
if ak.any(had_evts) else ak.unflatten(
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge < 0]
[:, 0], ak.values_astype(
had_evts, int)), # leading negatively-charged parton
"Second_minus":
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge < 0][:,
1]
if ak.any(had_evts) else ak.unflatten(
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge < 0]
[:, 1], ak.values_astype(
had_evts, int)), # subleading negatively-charged parton
"Up_plus":
gen_wpartons_up[gen_wpartons_up.charge > 0][had_evts]
if ak.any(had_evts) else ak.unflatten(
gen_wpartons_up[gen_wpartons_up.charge > 0][had_evts],
ak.values_astype(had_evts,
int)), # positively-charged up-type parton
"Down_plus":
gen_wpartons_dw[gen_wpartons_dw.charge > 0][had_evts]
if ak.any(had_evts) else ak.unflatten(
gen_wpartons_dw[gen_wpartons_dw.charge > 0][had_evts],
ak.values_astype(had_evts,
int)), # positively-charged down-type parton
"Up_minus":
gen_wpartons_up[gen_wpartons_up.charge < 0][had_evts]
if ak.any(had_evts) else ak.unflatten(
gen_wpartons_up[gen_wpartons_up.charge < 0][had_evts],
ak.values_astype(had_evts,
int)), # negatively-charged up-type parton
"Down_minus":
gen_wpartons_dw[gen_wpartons_dw.charge < 0][had_evts]
if ak.any(had_evts) else ak.unflatten(
gen_wpartons_dw[gen_wpartons_dw.charge < 0][had_evts],
ak.values_astype(had_evts,
int)), # negatively-charged down-type parton
})
## SEMILEP
SEMILEP_evts = ak.zip({
"TTbar":
Gen_TTbar[semilep_evts] if ak.any(semilep_evts) else ak.unflatten(
Gen_TTbar[semilep_evts], ak.values_astype(semilep_evts, int)),
"THad":
Gen_Top_Pairs[Gen_Top_Pairs.decaytype == 2][semilep_evts]
if ak.any(semilep_evts) else ak.unflatten(
Gen_Top_Pairs[Gen_Top_Pairs.decaytype == 2][semilep_evts],
ak.values_astype(semilep_evts, int)),
"TLep":
Gen_Top_Pairs[Gen_Top_Pairs.decaytype == 1][semilep_evts]
if ak.any(semilep_evts) else ak.unflatten(
Gen_Top_Pairs[Gen_Top_Pairs.decaytype == 1][semilep_evts],
ak.values_astype(semilep_evts, int)),
"BHad":
Gen_B_Pairs[Gen_B_Pairs.decaytype == 2][semilep_evts]
if ak.any(semilep_evts) else ak.unflatten(
Gen_B_Pairs[Gen_B_Pairs.decaytype == 2][semilep_evts],
ak.values_astype(semilep_evts, int)),
"BLep":
Gen_B_Pairs[Gen_B_Pairs.decaytype == 1][semilep_evts]
if ak.any(semilep_evts) else ak.unflatten(
Gen_B_Pairs[Gen_B_Pairs.decaytype == 1][semilep_evts],
ak.values_astype(semilep_evts, int)),
"WHad":
Gen_W_Pairs[Gen_W_Pairs.decaytype == 2][semilep_evts]
if ak.any(semilep_evts) else ak.unflatten(
Gen_W_Pairs[Gen_W_Pairs.decaytype == 2][semilep_evts],
ak.values_astype(semilep_evts, int)),
"WLep":
Gen_W_Pairs[Gen_W_Pairs.decaytype == 1][semilep_evts]
if ak.any(semilep_evts) else ak.unflatten(
Gen_W_Pairs[Gen_W_Pairs.decaytype == 1][semilep_evts],
ak.values_astype(semilep_evts, int)),
"Lepton":
gen_charged_leps[semilep_evts] if ak.any(semilep_evts) else
ak.unflatten(gen_charged_leps[semilep_evts],
ak.values_astype(semilep_evts, int)),
"Nu":
gen_neutral_leps[semilep_evts] if ak.any(semilep_evts) else
ak.unflatten(gen_neutral_leps[semilep_evts],
ak.values_astype(semilep_evts, int)),
"WJa":
Gen_Wparton_Pairs[:, 0][semilep_evts] if ak.any(semilep_evts) else
ak.unflatten(Gen_Wparton_Pairs[:, 0][semilep_evts],
ak.values_astype(semilep_evts, int)),
"WJb":
Gen_Wparton_Pairs[:, 1][semilep_evts] if ak.any(semilep_evts) else
ak.unflatten(Gen_Wparton_Pairs[:, 1][semilep_evts],
ak.values_astype(semilep_evts, int)),
"Up_Had":
gen_wpartons_up[semilep_evts] if ak.any(semilep_evts) else
ak.unflatten(gen_wpartons_up[semilep_evts],
ak.values_astype(semilep_evts, int)),
"Down_Had":
gen_wpartons_dw[semilep_evts] if ak.any(semilep_evts) else
ak.unflatten(gen_wpartons_dw[semilep_evts],
ak.values_astype(semilep_evts, int)),
})
# make dictionary to return
GenObjects = dict({
"SL": SEMILEP_evts,
"DL": DILEP_evts,
"Had": HAD_evts,
})
return GenObjects
def test_highlevel():
array = ak.Array(
[[[2, 3, 5], [], [7, 11], [13]], [], [[17, 19], [23]]], check_valid=True
)
assert ak.count(array) == 9
assert ak.to_list(ak.count(array, axis=-1)) == [[3, 0, 2, 1], [], [2, 1]]
assert ak.to_list(ak.count(array, axis=2)) == [[3, 0, 2, 1], [], [2, 1]]
assert ak.to_list(ak.count(array, axis=-1, keepdims=True)) == [
[[3], [0], [2], [1]],
[],
[[2], [1]],
]
assert ak.to_list(ak.count(array, axis=-2)) == [[3, 2, 1], [], [2, 1]]
assert ak.to_list(ak.count(array, axis=1)) == [[3, 2, 1], [], [2, 1]]
assert ak.to_list(ak.count(array, axis=-2, keepdims=True)) == [
[[3, 2, 1]],
[[]],
[[2, 1]],
]
assert ak.count_nonzero(array) == 9
assert ak.to_list(ak.count_nonzero(array, axis=-1)) == [[3, 0, 2, 1], [], [2, 1]]
assert ak.to_list(ak.count_nonzero(array, axis=-2)) == [[3, 2, 1], [], [2, 1]]
assert ak.sum(array) == 2 + 3 + 5 + 7 + 11 + 13 + 17 + 19 + 23
assert ak.to_list(ak.sum(array, axis=-1)) == [
[2 + 3 + 5, 0, 7 + 11, 13],
[],
[17 + 19, 23],
]
assert ak.to_list(ak.sum(array, axis=-2)) == [
[2 + 7 + 13, 3 + 11, 5],
[],
[17 + 23, 19],
]
assert ak.prod(array) == 2 * 3 * 5 * 7 * 11 * 13 * 17 * 19 * 23
assert ak.to_list(ak.prod(array, axis=-1)) == [
[2 * 3 * 5, 1, 7 * 11, 13],
[],
[17 * 19, 23],
]
assert ak.to_list(ak.prod(array, axis=-2)) == [
[2 * 7 * 13, 3 * 11, 5],
[],
[17 * 23, 19],
]
assert ak.min(array) == 2
assert ak.to_list(ak.min(array, axis=-1)) == [[2, None, 7, 13], [], [17, 23]]
assert ak.to_list(ak.min(array, axis=-2)) == [[2, 3, 5], [], [17, 19]]
assert ak.max(array) == 23
assert ak.to_list(ak.max(array, axis=-1)) == [[5, None, 11, 13], [], [19, 23]]
assert ak.to_list(ak.max(array, axis=-2)) == [[13, 11, 5], [], [23, 19]]
array = ak.Array(
[
[[True, False, True], [], [False, False], [True]],
[],
[[False, True], [True]],
],
check_valid=True,
)
assert ak.any(array) == True
assert ak.to_list(ak.any(array, axis=-1)) == [
[True, False, False, True],
[],
[True, True],
]
assert ak.to_list(ak.any(array, axis=-2)) == [[True, False, True], [], [True, True]]
assert ak.all(array) == False
assert ak.to_list(ak.all(array, axis=-1)) == [
[False, True, False, True],
[],
[False, True],
]
assert ak.to_list(ak.all(array, axis=-2)) == [
[False, False, True],
[],
[False, True],
]
def trilep_baseline(self, omit=[], cutflow=None, tight=False):
'''
give it a cutflow object if you want it to be filed.
cuts in the omit list will not be applied
'''
self.selection = PackedSelection()
is_trilep = ((ak.num(self.ele) + ak.num(self.mu))==3)
los_trilep = ((ak.num(self.ele) + ak.num(self.mu))>=2)
pos_charge = ((ak.sum(self.ele.pdgId, axis=1) + ak.sum(self.mu.pdgId, axis=1))<0)
neg_charge = ((ak.sum(self.ele.pdgId, axis=1) + ak.sum(self.mu.pdgId, axis=1))>0)
lep0pt = ((ak.num(self.ele[(self.ele.pt>25)]) + ak.num(self.mu[(self.mu.pt>25)]))>0)
lep1pt = ((ak.num(self.ele[(self.ele.pt>20)]) + ak.num(self.mu[(self.mu.pt>20)]))>1)
lepveto = ((ak.num(self.ele_veto) + ak.num(self.mu_veto))==3)
dimu = choose(self.mu, 2)
diele = choose(self.ele, 2)
dimu_veto = choose(self.mu_veto,2)
diele_veto = choose(self.ele_veto,2)
#dilep = cross(self.mu, self.ele)
OS_dimu = dimu[(dimu['0'].charge*dimu['1'].charge < 0)]
OS_diele = diele[(diele['0'].charge*diele['1'].charge < 0)]
OS_dimu_veto = dimu_veto[(dimu_veto['0'].charge*dimu_veto['1'].charge < 0)]
OS_diele_veto = diele_veto[(diele_veto['0'].charge*diele_veto['1'].charge < 0)]
SFOS = ak.concatenate([OS_diele_veto, OS_dimu_veto], axis=1)
offZ = (ak.all(abs(OS_dimu.mass-91.2)>10, axis=1) & ak.all(abs(OS_diele.mass-91.2)>10, axis=1))
offZ_veto = (ak.all(abs(OS_dimu_veto.mass-91.2)>10, axis=1) & ak.all(abs(OS_diele_veto.mass-91.2)>10, axis=1))
lepton = ak.concatenate([self.ele, self.mu], axis=1)
lepton_pdgId_pt_ordered = ak.fill_none(ak.pad_none(lepton[ak.argsort(lepton.pt, ascending=False)].pdgId, 2, clip=True), 0)
dilep = choose(lepton,2)
SS_dilep = (dilep['0'].charge*dilep['1'].charge > 0)
los_trilep_SS = (ak.any(SS_dilep, axis=1))
vetolepton = ak.concatenate([self.ele_veto, self.mu_veto], axis=1)
vetotrilep = choose3(vetolepton, 3)
pos_trilep = ak.any((vetotrilep['0'].charge+vetotrilep['1'].charge+vetotrilep['2'].charge > 0),axis=1)
neg_trilep = ak.any((vetotrilep['0'].charge+vetotrilep['1'].charge+vetotrilep['2'].charge < 0),axis=1)
triggers = getTriggers(self.events,
ak.flatten(lepton_pdgId_pt_ordered[:,0:1]),
ak.flatten(lepton_pdgId_pt_ordered[:,1:2]), year=self.year, dataset=self.dataset)
ht = ak.sum(self.jet_all.pt, axis=1)
st = self.met.pt + ht + ak.sum(self.mu.pt, axis=1) + ak.sum(self.ele.pt, axis=1)
st_veto = self.met.pt + ht + ak.sum(self.mu_veto.pt, axis=1) + ak.sum(self.ele_veto.pt, axis=1)
lep0pt_veto = ((ak.num(self.ele_veto[(self.ele_veto.pt>25)]) + ak.num(self.mu_veto[(self.mu_veto.pt>25)]))>0)
lep1pt_veto = ((ak.num(self.ele_veto[(self.ele_veto.pt>20)]) + ak.num(self.mu_veto[(self.mu_veto.pt>20)]))>1)
self.selection.add('lepveto', lepveto)
self.selection.add('trilep', los_trilep_SS)
self.selection.add('filter', self.filters)
self.selection.add('trigger', triggers)
self.selection.add('p_T(lep0)>25', lep0pt_veto)
self.selection.add('p_T(lep1)>20', lep1pt_veto)
self.selection.add('N_jet>2', (ak.num(self.jet_all)>2) )
self.selection.add('N_jet>3', (ak.num(self.jet_all)>3) )
self.selection.add('N_central>1', (ak.num(self.jet_central)>1) )
self.selection.add('N_central>2', (ak.num(self.jet_central)>2) )
self.selection.add('N_btag>0', (ak.num(self.jet_btag)>0 ))
self.selection.add('N_fwd>0', (ak.num(self.jet_fwd)>0) )
self.selection.add('MET>50', (self.met.pt>50) )
self.selection.add('ST>600', (st_veto>600) )
self.selection.add('offZ', offZ_veto )
#self.selection.add('SFOS>=1', ak.num(SFOS)==0)
#self.selection.add('charge_sum', neg_trilep)
reqs = [
'filter',
'lepveto',
'trilep',
'p_T(lep0)>25',
'p_T(lep1)>20',
'trigger',
'offZ',
'MET>50',
'N_jet>2',
'N_central>1',
'N_btag>0',
'N_fwd>0',
#'SFOS>=1',
#'charge_sum'
]
if tight:
reqs += [
'N_jet>3',
'N_central>2',
'ST>600',
#'MET>50',
#'delta_eta',
]
reqs_d = { sel: True for sel in reqs if not sel in omit }
selection = self.selection.require(**reqs_d)
self.reqs = [ sel for sel in reqs if not sel in omit ]
if cutflow:
#
cutflow_reqs_d = {}
for req in reqs:
cutflow_reqs_d.update({req: True})
cutflow.addRow( req, self.selection.require(**cutflow_reqs_d) )
return selection
def process(self, events):
output = self.accumulator.identity()
output['total']['all'] += len(events)
# use a very loose preselection to filter the events
presel = ak.num(events.Jet) > 2
ev = events[presel]
dataset = ev.metadata['dataset']
# load the config - probably not needed anymore
cfg = loadConfig()
## Muons
muon = Collections(ev, "Muon", "vetoTTH").get()
tightmuon = Collections(ev, "Muon", "tightTTH").get()
dimuon = choose(muon, 2)
SSmuon = ak.any((dimuon['0'].charge * dimuon['1'].charge) > 0, axis=1)
leading_muon_idx = ak.singletons(ak.argmax(muon.pt, axis=1))
leading_muon = muon[leading_muon_idx]
## Electrons
electron = Collections(ev, "Electron", "vetoTTH").get()
tightelectron = Collections(ev, "Electron", "tightTTH").get()
dielectron = choose(electron, 2)
SSelectron = ak.any(
(dielectron['0'].charge * dielectron['1'].charge) > 0, axis=1)
leading_electron_idx = ak.singletons(ak.argmax(electron.pt, axis=1))
leading_electron = electron[leading_electron_idx]
## Merge electrons and muons - this should work better now in ak1
dilepton = cross(muon, electron)
SSlepton = ak.any((dilepton['0'].charge * dilepton['1'].charge) > 0,
axis=1)
lepton = ak.concatenate([muon, electron], axis=1)
leading_lepton_idx = ak.singletons(ak.argmax(lepton.pt, axis=1))
leading_lepton = lepton[leading_lepton_idx]
trailing_lepton_idx = ak.singletons(ak.argmin(lepton.pt, axis=1))
trailing_lepton = lepton[trailing_lepton_idx]
dilepton_mass = (leading_lepton + trailing_lepton).mass
dilepton_pt = (leading_lepton + trailing_lepton).pt
dilepton_dR = delta_r(leading_lepton, trailing_lepton)
mt_lep_met = mt(lepton.pt, lepton.phi, ev.MET.pt, ev.MET.phi)
min_mt_lep_met = ak.min(mt_lep_met, axis=1)
## Jets
jet = getJets(ev, minPt=25, maxEta=4.7, pt_var='pt_nom')
jet = jet[ak.argsort(
jet.pt_nom, ascending=False
)] # need to sort wrt smeared and recorrected jet pt
jet = jet[~match(jet, muon,
deltaRCut=0.4)] # remove jets that overlap with muons
jet = jet[~match(
jet, electron,
deltaRCut=0.4)] # remove jets that overlap with electrons
central = jet[(abs(jet.eta) < 2.4)]
btag = getBTagsDeepFlavB(
jet, year=self.year) # should study working point for DeepJet
light = getBTagsDeepFlavB(jet, year=self.year, invert=True)
fwd = getFwdJet(light)
fwd_noPU = getFwdJet(light, puId=False)
tau = getTaus(ev)
track = getIsoTracks(ev)
## forward jets
j_fwd = fwd[ak.singletons(ak.argmax(
fwd.p, axis=1))] # highest momentum spectator
high_score_btag = central[ak.argsort(central.btagDeepFlavB)][:, :2]
bl = cross(lepton, high_score_btag)
bl_dR = delta_r(bl['0'], bl['1'])
min_bl_dR = ak.min(bl_dR, axis=1)
jf = cross(j_fwd, jet)
mjf = (jf['0'] + jf['1']).mass
j_fwd2 = jf[ak.singletons(
ak.argmax(mjf, axis=1)
)]['1'] # this is the jet that forms the largest invariant mass with j_fwd
delta_eta = ak.fill_none(
ak.pad_none(abs(j_fwd2.eta - j_fwd.eta), 1, clip=True), 0)
## MET -> can switch to puppi MET
met_pt = ev.MET.pt
met_phi = ev.MET.phi
## other variables
ht = ak.sum(jet.pt, axis=1)
st = met_pt + ht + ak.sum(muon.pt, axis=1) + ak.sum(electron.pt,
axis=1)
## event selectors
filters = getFilters(ev, year=self.year, dataset=dataset)
dilep = ((ak.num(tightelectron) + ak.num(tightmuon)) == 2)
lep0pt = ((ak.num(electron[(electron.pt > 25)]) +
ak.num(muon[(muon.pt > 25)])) > 0)
lep1pt = ((ak.num(electron[(electron.pt > 20)]) +
ak.num(muon[(muon.pt > 20)])) > 1)
lepveto = ((ak.num(electron) + ak.num(muon)) == 2)
selection = PackedSelection()
selection.add('lepveto', lepveto)
selection.add('dilep', dilep)
selection.add('filter', (filters))
selection.add('p_T(lep0)>25', lep0pt)
selection.add('p_T(lep1)>20', lep1pt)
selection.add('SS', (SSlepton | SSelectron | SSmuon))
selection.add('N_jet>3', (ak.num(jet) >= 4))
selection.add('N_central>2', (ak.num(central) >= 3))
selection.add('N_btag>0', (ak.num(btag) >= 1))
selection.add('N_fwd>0', (ak.num(fwd) >= 1))
#ss_reqs = ['lepveto', 'dilep', 'filter', 'p_T(lep0)>25', 'p_T(lep1)>20', 'SS']
ss_reqs = [
'lepveto', 'dilep', 'filter', 'p_T(lep0)>25', 'p_T(lep1)>20', 'SS'
]
#bl_reqs = ss_reqs + ['N_jet>3', 'N_central>2', 'N_btag>0', 'N_fwd>0']
bl_reqs = ss_reqs + ['N_jet>3', 'N_central>2', 'N_btag>0']
ss_reqs_d = {sel: True for sel in ss_reqs}
ss_selection = selection.require(**ss_reqs_d)
bl_reqs_d = {sel: True for sel in bl_reqs}
BL = selection.require(**bl_reqs_d)
weight = Weights(len(ev))
if not dataset == 'MuonEG':
# lumi weight
weight.add("weight", ev.weight)
# PU weight - not in the babies...
weight.add("PU",
ev.puWeight,
weightUp=ev.puWeightUp,
weightDown=ev.puWeightDown,
shift=False)
# b-tag SFs
weight.add("btag", self.btagSF.Method1a(btag, light))
# lepton SFs
weight.add("lepton", self.leptonSF.get(electron, muon))
#cutflow = Cutflow(output, ev, weight=weight)
#cutflow_reqs_d = {}
#for req in bl_reqs:
# cutflow_reqs_d.update({req: True})
# cutflow.addRow( req, selection.require(**cutflow_reqs_d) )
labels = {
'topW_v3': 0,
'TTW': 1,
'TTZ': 2,
'TTH': 3,
'ttbar': 4,
'ttbar1l_MG': 4
}
if dataset in labels:
label_mult = labels[dataset]
else:
label_mult = 5
label = np.ones(len(ev[BL])) * label_mult
output["n_lep"] += processor.column_accumulator(
ak.to_numpy((ak.num(electron) + ak.num(muon))[BL]))
output["n_lep_tight"] += processor.column_accumulator(
ak.to_numpy((ak.num(tightelectron) + ak.num(tightmuon))[BL]))
output["lead_lep_pt"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(leading_lepton[BL].pt, axis=1)))
output["lead_lep_eta"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(leading_lepton[BL].eta, axis=1)))
output["lead_lep_phi"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(leading_lepton[BL].phi, axis=1)))
output["lead_lep_charge"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(leading_lepton[BL].charge, axis=1)))
output["sublead_lep_pt"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(trailing_lepton[BL].pt, axis=1)))
output["sublead_lep_eta"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(trailing_lepton[BL].eta, axis=1)))
output["sublead_lep_phi"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(trailing_lepton[BL].phi, axis=1)))
output["sublead_lep_charge"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(trailing_lepton[BL].charge, axis=1)))
output["lead_jet_pt"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(jet[:, 0:1][BL].pt, axis=1)))
output["lead_jet_eta"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(jet[:, 0:1][BL].eta, axis=1)))
output["lead_jet_phi"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(jet[:, 0:1][BL].phi, axis=1)))
output["sublead_jet_pt"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(jet[:, 1:2][BL].pt, axis=1)))
output["sublead_jet_eta"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(jet[:, 1:2][BL].eta, axis=1)))
output["sublead_jet_phi"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(jet[:, 1:2][BL].phi, axis=1)))
output["lead_btag_pt"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(high_score_btag[:, 0:1][BL].pt, axis=1)))
output["lead_btag_eta"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(high_score_btag[:, 0:1][BL].eta, axis=1)))
output["lead_btag_phi"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(high_score_btag[:, 0:1][BL].phi, axis=1)))
output["sublead_btag_pt"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(high_score_btag[:, 1:2][BL].pt, axis=1)))
output["sublead_btag_eta"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(high_score_btag[:, 1:2][BL].eta, axis=1)))
output["sublead_btag_phi"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(high_score_btag[:, 1:2][BL].phi, axis=1)))
output["fwd_jet_p"] += processor.column_accumulator(
ak.to_numpy(
ak.flatten(ak.fill_none(ak.pad_none(j_fwd[BL].p, 1, clip=True),
0),
axis=1)))
output["fwd_jet_pt"] += processor.column_accumulator(
ak.to_numpy(
ak.flatten(ak.fill_none(
ak.pad_none(j_fwd[BL].pt, 1, clip=True), 0),
axis=1)))
output["fwd_jet_eta"] += processor.column_accumulator(
ak.to_numpy(
ak.flatten(ak.fill_none(
ak.pad_none(j_fwd[BL].eta, 1, clip=True), 0),
axis=1)))
output["fwd_jet_phi"] += processor.column_accumulator(
ak.to_numpy(
ak.flatten(ak.fill_none(
ak.pad_none(j_fwd[BL].phi, 1, clip=True), 0),
axis=1)))
output["mjj_max"] += processor.column_accumulator(
ak.to_numpy(ak.fill_none(ak.max(mjf[BL], axis=1), 0)))
output["delta_eta_jj"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(delta_eta[BL], axis=1)))
output["met"] += processor.column_accumulator(ak.to_numpy(met_pt[BL]))
output["ht"] += processor.column_accumulator(ak.to_numpy(ht[BL]))
output["st"] += processor.column_accumulator(ak.to_numpy(st[BL]))
output["n_jet"] += processor.column_accumulator(
ak.to_numpy(ak.num(jet[BL])))
output["n_btag"] += processor.column_accumulator(
ak.to_numpy(ak.num(btag[BL])))
output["n_fwd"] += processor.column_accumulator(
ak.to_numpy(ak.num(fwd[BL])))
output["n_central"] += processor.column_accumulator(
ak.to_numpy(ak.num(central[BL])))
output["n_tau"] += processor.column_accumulator(
ak.to_numpy(ak.num(tau[BL])))
output["n_track"] += processor.column_accumulator(
ak.to_numpy(ak.num(track[BL])))
output["dilepton_pt"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(dilepton_pt[BL], axis=1)))
output["dilepton_mass"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(dilepton_mass[BL], axis=1)))
output["min_bl_dR"] += processor.column_accumulator(
ak.to_numpy(min_bl_dR[BL]))
output["min_mt_lep_met"] += processor.column_accumulator(
ak.to_numpy(min_mt_lep_met[BL]))
output["label"] += processor.column_accumulator(label)
output["weight"] += processor.column_accumulator(weight.weight()[BL])
output["presel"]["all"] += len(ev[ss_selection])
output["sel"]["all"] += len(ev[BL])
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
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 process(self, events):
def normalize(val, cut):
return ak.to_numpy(ak.fill_none(
val[cut],
np.nan)) #val[cut].pad(1, clip=True).fillna(0).flatten()
def fill(region, cuts, systematic=None, wmod=None):
print('filling %s' % region)
selections = cuts
cut = selection.all(*selections)
if 'signal' in region: weight = weights_signal.weight()[cut]
elif 'muonCR' in region: weight = weights_muonCR.weight()[cut]
elif 'VtaggingCR' in region:
weight = weights_VtaggingCR.weight()[cut]
output['templates'].fill(
dataset=dataset,
region=region,
pt=normalize(candidatejet.pt, cut),
msd=normalize(candidatejet.msdcorr, cut),
n2ddt=normalize(candidatejet.n2ddt, cut),
#gruddt=normalize(candidatejet.gruddt, cut),
in_v3_ddt=normalize(candidatejet.in_v3_ddt, cut),
hadW=normalize(candidatejet.nmatcheddau, cut),
weight=weight,
),
output['event'].fill(
dataset=dataset,
region=region,
MET=events.MET.pt[cut],
#nJet=fatjets.counts[cut],
nPFConstituents=normalize(candidatejet.nPFConstituents, cut),
weight=weight,
),
output['deepAK8'].fill(
dataset=dataset,
region=region,
deepTagMDWqq=normalize(candidatejet.deepTagMDWqq, cut),
deepTagMDZqq=normalize(candidatejet.deepTagMDZqq, cut),
msd=normalize(candidatejet.msdcorr, cut),
#genflavor=genflavor[cut],
weight=weight,
),
output['in_v3'].fill(
dataset=dataset,
region=region,
#genflavor=genflavor[cut],
in_v3=normalize(candidatejet.in_v3, cut),
n2=normalize(candidatejet.n2b1, cut),
gru=normalize(candidatejet.gru, cut),
weight=weight,
),
if 'muonCR' in dataset or 'VtaggingCR' in dataset:
output['muon'].fill(
dataset=dataset,
region=region,
mu_pt=normalize(candidatemuon.pt, cut),
mu_eta=normalize(candidatemuon.eta, cut),
mu_pfRelIso04_all=normalize(candidatemuon.pfRelIso04_all,
cut),
weight=weight,
),
#common jet kinematics
gru = events.GRU
IN = events.IN
fatjets = events.FatJet
fatjets['msdcorr'] = corrected_msoftdrop(fatjets)
fatjets['qcdrho'] = 2 * np.log(fatjets.msdcorr / fatjets.pt)
fatjets['gruddt'] = gru.v25 - shift(
fatjets, algo='gruddt', year='2017')
fatjets['gru'] = gru.v25
fatjets['in_v3'] = IN.v3
fatjets['in_v3_ddt'] = IN.v3 - shift(
fatjets, algo='inddt', year='2017')
fatjets['in_v3_ddt_90pctl'] = IN.v3 - shift(
fatjets, algo='inddt90pctl', year='2017')
fatjets['n2ddt'] = fatjets.n2b1 - n2ddt_shift(fatjets, year='2017')
fatjets['nmatcheddau'] = TTsemileptonicmatch(events)
dataset = events.metadata['dataset']
print('process dataset', dataset)
isRealData = not hasattr(events, 'genWeight')
output = self.accumulator.identity()
if (len(events) == 0): return output
selection = PackedSelection('uint64')
weights_signal = Weights(len(events))
weights_muonCR = Weights(len(events))
weights_VtaggingCR = Weights(len(events))
if not isRealData:
output['sumw'][dataset] += ak.sum(events.genWeight)
#######################
if 'signal' in self._region:
if isRealData:
trigger_fatjet = np.zeros(len(events), dtype='bool')
for t in self._triggers[self._year]:
try:
trigger_fatjet = trigger_fatjet | events.HLT[t]
except:
print('trigger %s not available' % t)
continue
else:
trigger_fatjet = np.ones(len(events), dtype='bool')
fatjets["genMatchFull"] = VQQgenmatch(events)
candidatejet = ak.firsts(fatjets)
candidatejet["genMatchFull"] = VQQgenmatch(events)
nelectrons = ak.sum(
(events.Electron.pt > 10.)
& (abs(events.Electron.eta) < 2.5)
& (events.Electron.cutBased >= events.Electron.VETO),
axis=1,
)
nmuons = ak.sum(
(events.Muon.pt > 10)
& (abs(events.Muon.eta) < 2.1)
& (events.Muon.pfRelIso04_all < 0.4)
& (events.Muon.looseId),
axis=1,
)
ntaus = ak.sum(
(events.Tau.pt > 20.)
& (events.Tau.idDecayMode)
& (events.Tau.rawIso < 5)
& (abs(events.Tau.eta) < 2.3),
axis=1,
)
cuts = {
"S_fatjet_trigger":
trigger_fatjet,
"S_pt":
candidatejet.pt > 525,
"S_eta": (abs(candidatejet.eta) < 2.5),
"S_msdcorr": (candidatejet.msdcorr > 40),
"S_rho":
((candidatejet.qcdrho > -5.5) & (candidatejet.qcdrho < -2.)),
"S_jetid": (candidatejet.isTight),
"S_VQQgenmatch": (candidatejet.genMatchFull),
"S_noelectron": (nelectrons == 0),
"S_nomuon": (nmuons == 0),
"S_notau": (ntaus == 0),
}
for name, cut in cuts.items():
print(name, cut)
selection.add(name, cut)
if isRealData:
genflavor = 0 #candidatejet.pt.zeros_like().pad(1, clip=True).fillna(-1).flatten()
if not isRealData:
weights_signal.add('genweight', events.genWeight)
#add_pileup_weight(weights_signal, events.Pileup.nPU, self._year, dataset)
add_jetTriggerWeight(weights_signal, candidatejet.msdcorr,
candidatejet.pt, self._year)
bosons = getBosons(events.GenPart)
genBosonPt = ak.fill_none(ak.firsts(bosons.pt), 0)
add_VJets_NLOkFactor(weights_signal, genBosonPt, self._year,
dataset)
#genflavor = matchedBosonFlavor(candidatejet, bosons).pad(1, clip=True).fillna(-1).flatten()
allcuts_signal = set()
output['cutflow_signal'][dataset]['none'] += float(
weights_signal.weight().sum())
for cut in cuts:
allcuts_signal.add(cut)
output['cutflow_signal'][dataset][cut] += float(
weights_signal.weight()[selection.all(
*allcuts_signal)].sum())
fill('signal', cuts.keys())
#######################
if 'muonCR' in self._region:
if isRealData:
trigger_muon = np.zeros(len(events), dtype='bool')
for t in self._muontriggers[self._year]:
trigger_muon = trigger_muon | events.HLT[t]
else:
trigger_muon = np.ones(len(events), dtype='bool')
candidatejet = ak.firsts(fatjets)
candidatemuon = events.Muon[:, :5]
jets = events.Jet[((events.Jet.pt > 50.)
& (abs(events.Jet.eta) < 2.5)
& (events.Jet.isTight))][:, :4]
dphi = abs(jets.delta_phi(candidatejet))
ak4_away = jets[(dphi > 0.8)]
nelectrons = ak.sum(
(events.Electron.pt > 10.)
& (abs(events.Electron.eta) < 2.5)
& (events.Electron.cutBased >= events.Electron.VETO),
axis=1,
)
nmuons = ak.sum(
(events.Muon.pt > 10)
& (abs(events.Muon.eta) < 2.4)
& (events.Muon.pfRelIso04_all < 0.25)
& (events.Muon.looseId),
axis=1,
)
ntaus = ak.sum(
(events.Tau.pt > 20.)
& (events.Tau.idDecayMode)
& (events.Tau.rawIso < 5)
& (abs(events.Tau.eta) < 2.3)
& (events.Tau.idMVAoldDM2017v1 >= 16),
axis=1,
)
cuts = {
"CR1_muon_trigger":
trigger_muon,
"CR1_jet_pt": (candidatejet.pt > 525),
"CR1_jet_eta": (abs(candidatejet.eta) < 2.5),
"CR1_jet_msd": (candidatejet.msdcorr > 40),
"CR1_jet_rho":
((candidatejet.qcdrho > -5.5) & (candidatejet.qcdrho < -2.)),
"CR1_mu_pt":
ak.any(candidatemuon.pt > 55, axis=1),
"CR1_mu_eta":
ak.any(abs(candidatemuon.eta) < 2.1, axis=1),
"CR1_mu_IDLoose":
ak.any(candidatemuon.looseId, axis=1),
"CR1_mu_isolationTight":
ak.any(candidatemuon.pfRelIso04_all < 0.15, axis=1),
"CR1_muonDphiAK8":
ak.any(
abs(candidatemuon.delta_phi(candidatejet)) > 2 * np.pi / 3,
axis=1),
"CR1_ak4btagMedium08":
(ak.max(ak4_away.btagCSVV2, axis=1, mask_identity=False) >
BTagEfficiency.btagWPs[self._year]['medium']
), #(ak4_away.btagCSVV2.max() > 0.8838),
"CR1_noelectron": (nelectrons == 0),
"CR1_onemuon": (nmuons == 1),
"CR1_notau": (ntaus == 0),
}
for name, cut in cuts.items():
selection.add(name, cut)
if isRealData:
genflavor = 0 #candidatejet.pt.zeros_like().pad(1, clip=True).fillna(-1).flatten()
if not isRealData:
weights_muonCR.add('genweight', events.genWeight)
#add_pileup_weight(weights_muonCR, events.Pileup.nPU, self._year, dataset)
#add_singleMuTriggerWeight(weights, candidatejet.msdcorr, candidatejet.pt, self._year)
bosons = getBosons(events.GenPart)
genBosonPt = ak.fill_none(ak.firsts(bosons.pt), 0)
#add_VJets_NLOkFactor(weights, genBosonPt, self._year, dataset)
#genflavor = matchedBosonFlavor(candidatejet, bosons).pad(1, clip=True).fillna(-1).flatten()
allcuts_ttbar_muoncontrol = set()
output['cutflow_muonCR'][dataset]['none'] += float(
weights_muonCR.weight().sum())
for cut in cuts:
allcuts_ttbar_muoncontrol.add(cut)
output['cutflow_muonCR'][dataset][cut] += float(
weights_muonCR.weight()[selection.all(
*allcuts_ttbar_muoncontrol)].sum())
fill('muonCR', cuts.keys())
#######################
if 'VtaggingCR' in self._region:
if isRealData:
trigger_muon = np.zeros(len(events), dtype='bool')
for t in self._muontriggers[self._year]:
trigger_muon = trigger_muon | events.HLT[t]
else:
trigger_muon = np.ones(len(events), dtype='bool')
candidatejet = ak.firsts(fatjets)
candidatemuon = ak.firsts(events.Muon)
jets = events.Jet[((events.Jet.pt > 30.)
& (abs(events.Jet.eta) < 2.4))][:, :4]
dr_ak4_ak8 = jets.delta_r(candidatejet)
dr_ak4_muon = jets.delta_r(candidatemuon)
ak4_away = jets[(dr_ak4_ak8 > 0.8)] # & (dr_ak4_muon > 0.4)]
mu_p4 = ak.zip(
{
"pt": ak.fill_none(candidatemuon.pt, 0),
"eta": ak.fill_none(candidatemuon.eta, 0),
"phi": ak.fill_none(candidatemuon.phi, 0),
"mass": ak.fill_none(candidatemuon.mass, 0),
},
with_name="PtEtaPhiMLorentzVector")
met_p4 = ak.zip(
{
"pt": ak.from_iter([[v] for v in events.MET.pt]),
"eta": ak.from_iter([[v] for v in np.zeros(len(events))]),
"phi": ak.from_iter([[v] for v in events.MET.phi]),
"mass": ak.from_iter([[v] for v in np.zeros(len(events))]),
},
with_name="PtEtaPhiMLorentzVector")
Wleptoniccandidate = mu_p4 + met_p4
nelectrons = ak.sum(
((events.Electron.pt > 10.)
& (abs(events.Electron.eta) < 2.5)
& (events.Electron.cutBased >= events.Electron.VETO)),
axis=1,
)
n_tight_muon = ak.sum(
((events.Muon.pt > 53)
& (abs(events.Muon.eta) < 2.1)
& (events.Muon.tightId)),
axis=1,
)
n_loose_muon = ak.sum(
((events.Muon.pt > 20)
& (events.Muon.looseId)
& (abs(events.Muon.eta) < 2.4)),
axis=1,
)
ntaus = ak.sum(
((events.Tau.pt > 20.)
& (events.Tau.idDecayMode)
& (events.Tau.rawIso < 5)
& (abs(events.Tau.eta) < 2.3)
& (events.Tau.idMVAoldDM2017v1 >= 16)),
axis=1,
)
cuts = {
"CR2_muon_trigger":
trigger_muon,
"CR2_jet_pt": (candidatejet.pt > 200),
"CR2_jet_eta": (abs(candidatejet.eta) < 2.5),
"CR2_jet_msd": (candidatejet.msdcorr > 40),
"CR2_mu_pt":
candidatemuon.pt > 53,
"CR2_mu_eta": (abs(candidatemuon.eta) < 2.1),
"CR2_mu_IDTight":
candidatemuon.tightId,
"CR2_mu_isolationTight": (candidatemuon.pfRelIso04_all < 0.15),
"CR2_muonDphiAK8":
abs(candidatemuon.delta_phi(candidatejet)) > 2 * np.pi / 3,
"CR2_ak4btagMedium08":
(ak.max(ak4_away.btagCSVV2, axis=1, mask_identity=False) >
BTagEfficiency.btagWPs[self._year]['medium']),
"CR2_leptonicW":
ak.flatten(Wleptoniccandidate.pt > 200),
"CR2_MET": (events.MET.pt > 40.),
"CR2_noelectron": (nelectrons == 0),
"CR2_one_tightMuon": (n_tight_muon == 1),
"CR2_one_looseMuon": (n_loose_muon == 1),
#"CR2_notau" : (ntaus==0),
}
for name, cut in cuts.items():
print(name, cut)
selection.add(name, cut)
#weights.add('metfilter', events.Flag.METFilters)
if isRealData:
genflavor = 0 #candidatejet.pt.zeros_like().pad(1, clip=True).fillna(-1).flatten()
if not isRealData:
weights_VtaggingCR.add('genweight', events.genWeight)
#add_pileup_weight(weights_VtaggingCR, events.Pileup.nPU, self._year, dataset)
#add_singleMuTriggerWeight(weights, abs(candidatemuon.eta), candidatemuon.pt, self._year)
bosons = getBosons(events.GenPart)
genBosonPt = ak.fill_none(ak.firsts(bosons.pt), 0)
#add_VJets_NLOkFactor(weights, genBosonPt, self._year, dataset)
#genflavor = matchedBosonFlavor(candidatejet, bosons).pad(1, clip=True).fillna(-1).flatten()
#b-tag weights
allcuts_vselection = set()
output['cutflow_VtaggingCR'][dataset]['none'] += float(
weights_VtaggingCR.weight().sum())
for cut in cuts:
allcuts_vselection.add(cut)
output['cutflow_VtaggingCR'][dataset][cut] += float(
weights_VtaggingCR.weight()[selection.all(
*allcuts_vselection)].sum())
fill('VtaggingCR', cuts.keys())
return output
