def get_objects(ievent):
visible_momentum = MadMinerParticle()
for p in (
electrons_all_events[ievent]
+ jets_all_events[ievent]
+ muons_all_events[ievent]
+ leptons_all_events[ievent]
):
visible_momentum += p
all_momentum = visible_momentum + met_all_events[ievent][0]
objects = math_commands()
objects.update(
{
"e": electrons_all_events[ievent],
"j": jets_all_events[ievent],
"a": photons_all_events[ievent],
"mu": muons_all_events[ievent],
"l": leptons_all_events[ievent],
"met": met_all_events[ievent][0],
"visible": visible_momentum,
"all": all_momentum,
"boost_to_com": lambda momentum: momentum.boost(all_momentum.boost_vector()),
}
)
return objects
def _get_objects(particles, met_resolution=None):
# Find visible particles
electrons = []
muons = []
taus = []
photons = []
jets = []
leptons = []
neutrinos = []
unstables = []
invisibles = []
for particle in particles:
pdgid = abs(particle.pdgid)
if pdgid in [1, 2, 3, 4, 5, 6, 9, 21]:
jets.append(particle)
elif pdgid == 11:
electrons.append(particle)
leptons.append(particle)
elif pdgid == 13:
muons.append(particle)
leptons.append(particle)
elif pdgid == 15:
taus.append(particle)
elif pdgid == 22:
photons.append(particle)
elif pdgid in [12, 14, 16]:
neutrinos.append(particle)
invisibles.append(particle)
elif pdgid in [23, 24, 25]:
unstables.append(particle)
else:
logger.warning(
"Unknown particle with PDG id %s, treating as invisible!")
invisibles.append(particle)
# Sort by pT
electrons = sorted(electrons, reverse=True, key=lambda x: x.pt)
muons = sorted(muons, reverse=True, key=lambda x: x.pt)
taus = sorted(taus, reverse=True, key=lambda x: x.pt)
photons = sorted(photons, reverse=True, key=lambda x: x.pt)
leptons = sorted(leptons, reverse=True, key=lambda x: x.pt)
neutrinos = sorted(neutrinos, reverse=True, key=lambda x: x.pt)
jets = sorted(jets, reverse=True, key=lambda x: x.pt)
# Sum over all particles
ht = 0.0
visible_sum = MadMinerParticle()
visible_sum.setpxpypze(0.0, 0.0, 0.0, 0.0)
for particle in particles:
ht += particle.pt
pdgid = abs(particle.pdgid)
if pdgid in [1, 2, 3, 4, 5, 6, 9, 11, 13, 15, 21, 22, 23, 24, 25]:
visible_sum += particle
# Soft noise
if met_resolution is not None:
noise_std = met_resolution[0] + met_resolution[1] * ht
noise_x = np.random.normal(0.0, noise_std, size=None)
noise_y = np.random.normal(0.0, noise_std, size=None)
else:
noise_x = 0.0
noise_y = 0.0
# MET
met_x = -visible_sum.px + noise_x
met_y = -visible_sum.px + noise_y
met = MadMinerParticle()
met.setpxpypze(met_x, met_y, 0.0, (met_x**2 + met_y**2)**0.5)
# Build objects
objects = math_commands()
objects.update({
"p": particles,
"e": electrons,
"j": jets,
"a": photons,
"mu": muons,
"tau": taus,
"l": leptons,
"met": met,
"v": neutrinos,
})
return objects
def _get_objects(particles):
# Find visible particles
electrons = []
muons = []
taus = []
photons = []
jets = []
leptons = []
neutrinos = []
unstables = []
invisibles = []
for particle in particles:
pdgid = abs(particle.pdgid)
if pdgid in [1, 2, 3, 4, 5, 6, 9, 21]:
jets.append(particle)
elif pdgid == 11:
electrons.append(particle)
leptons.append(particle)
elif pdgid == 13:
muons.append(particle)
leptons.append(particle)
elif pdgid == 15:
taus.append(particle)
elif pdgid == 22:
photons.append(particle)
elif pdgid in [12, 14, 16]:
neutrinos.append(particle)
invisibles.append(particle)
elif pdgid in [23, 24, 25]:
unstables.append(particle)
else:
logger.warning("Unknown particle with PDG id %s, treating as invisible!")
invisibles.append(particle)
# Sort by pT
electrons = sorted(electrons, reverse=True, key=lambda x: x.pt)
muons = sorted(muons, reverse=True, key=lambda x: x.pt)
taus = sorted(taus, reverse=True, key=lambda x: x.pt)
photons = sorted(photons, reverse=True, key=lambda x: x.pt)
leptons = sorted(leptons, reverse=True, key=lambda x: x.pt)
neutrinos = sorted(neutrinos, reverse=True, key=lambda x: x.pt)
jets = sorted(jets, reverse=True, key=lambda x: x.pt)
# MET
visible_sum = MadMinerParticle()
visible_sum.setpxpypze(0.0, 0.0, 0.0, 0.0)
for particle in particles:
pdgid = abs(particle.pdgid)
if pdgid in [1, 2, 3, 4, 5, 6, 9, 11, 13, 15, 21, 22, 23, 24, 25]:
visible_sum += particle
met = MadMinerParticle()
met.setpxpypze(-visible_sum.px, -visible_sum.px, 0.0, visible_sum.pt)
# Build objects
objects = math_commands()
objects.update(
{
"p": particles,
"e": electrons,
"j": jets,
"a": photons,
"mu": muons,
"tau": taus,
"l": leptons,
"met": met,
"v": neutrinos,
}
)
return objects
def get_objects(ievent):
visible_momentum = MadMinerParticle()
for p in (
electrons_all_events[ievent]
+ jets_all_events[ievent]
+ muons_all_events[ievent]
+ photons_all_events[ievent]
):
visible_momentum += p
all_momentum = visible_momentum + met_all_events[ievent][0]
objects = math_commands()
objects.update(
{
"e": electrons_all_events[ievent],
"j": jets_all_events[ievent],
"a": photons_all_events[ievent],
"mu": muons_all_events[ievent],
"l": leptons_all_events[ievent],
"met": met_all_events[ievent][0],
"visible": visible_momentum,
"all": all_momentum,
"boost_to_com": lambda momentum: momentum.boost(all_momentum.boost_vector()),
}
)
##################### PREPARE FUNCTIONS TO FIND BEST ZZ CANDIDATE #####################
#find same flavour opposite sign lepton pairs,
#given the list of electrons or muons in an event
#returns all possible SFOS pairs ([pos_idx, neg_idx],...)
#which satisfy invariant mass cuts
def find_Z_cand(cand_list, part_list):
#positively and negatively charged leptons, and their
idx_pos = []
cand_pos = []
idx_neg = []
cand_neg = []
for idx, (is_cand, part) in enumerate(zip(cand_list, part_list)):
#checking if the lepton is a candidate
if part.charge > 0:
idx_pos.append(idx)
cand_pos.append(is_cand)
if part.charge < 0:
idx_neg.append(idx)
cand_neg.append(is_cand)
z_cand_list = []
for candpair, pair in zip(itertools.product(cand_pos, cand_neg), itertools.product(idx_pos, idx_neg)):
iscand = candpair[0]*candpair[1] #1 only if both are candidates
#cut invariant mass
invm = (part_list[pair[0]] + part_list[pair[1]]).m
if invm < 120 and invm > 12:
z_cand_list.append([pair[0],pair[1]])
return z_cand_list #[[pos_idx,neg_idx], [pos_idx,neg_idx], ...]
#find pairs of SFOS pairs with no leptons in common
#keep track of their different flavours
#order of Z1 and Z2 is not important (each only counted once)
def find_ZZ_cand(Z_cand_e_list, Z_cand_mu_list):
ZZ_cand_ee = []; ZZ_cand_mm = []; ZZ_cand_em = []
#same flavour e
for Zcand1,Zcand2 in itertools.product(Z_cand_e_list,Z_cand_e_list):
if not any(x in Zcand1 for x in Zcand2):
ZZ_cand_ee.append([Zcand1[0], Zcand1[1], Zcand2[0], Zcand2[1]])
#same flavour mu
for Zcand1,Zcand2 in itertools.product(Z_cand_mu_list,Z_cand_mu_list):
if not any(x in Zcand1 for x in Zcand2):
ZZ_cand_mm.append([Zcand1[0], Zcand1[1], Zcand2[0], Zcand2[1]])
#different flavour
for Zcand1,Zcand2 in itertools.product(Z_cand_e_list,Z_cand_mu_list):
ZZ_cand_em.append([Zcand1[0], Zcand1[1], Zcand2[0], Zcand2[1]])
return ZZ_cand_ee, ZZ_cand_mm, ZZ_cand_em
#return the list of particles in the correct order:
#first list with flavour of first Z, second flavour of second Z
def set_flavours_lists(part_list_e, part_list_mu, flavours):
if flavours == "ee":
part_list1 = part_list_e
part_list2 = part_list_e
elif flavours == "mm":
part_list1 = part_list_mu
part_list2 = part_list_mu
elif flavours == "em":
part_list1 = part_list_e
part_list2 = part_list_mu
elif flavours == "me":
part_list1 = part_list_mu
part_list2 = part_list_e
else:
print("set_flavours_lists: flavour not valid: choose between ee, mm, em")
return
return part_list1, part_list2
#return the leptons corresponding to the index in ZZcand
#taking into account the flavours
def ZZidx_to_leps(ZZcand, part_list_e, part_list_mu, flavours):
part_list1, part_list2 = set_flavours_lists(part_list_e, part_list_mu, flavours)
leps = [part_list1[ZZcand[0]], part_list1[ZZcand[1]], part_list2[ZZcand[2]], part_list2[ZZcand[3]]]
return leps
#reorder the two Z in each pair: first the one with closest mass to Z
#keep track of the possible swapping in case of different flavours
def Z1Z2_ordering(ZZcand, part_list_e, part_list_mu, flavours):
leps = ZZidx_to_leps(ZZcand, part_list_e, part_list_mu, flavours)
#compute invariant masses of each Z
mz_p1 = (leps[0] + leps[1]).m
mz_p2 = (leps[2] + leps[3]).m
#order accordingly to closest mass to Z mass
if min([mz_p1,mz_p2], key=lambda x:abs(x-91.2), default=-1) == mz_p1:
z1z2 = ZZcand
swapped = False
else:
z1z2 = [ZZcand[2],ZZcand[3],ZZcand[0],ZZcand[1]]
swapped = True
return(z1z2, swapped)
#apply cuts to ZZ candidate
#consider differently the case of same flavour ZZ
def ZZ_cuts(ZZcand, part_list_e, part_list_mu, flavours, isSF):
leps = ZZidx_to_leps(ZZcand, part_list_e, part_list_mu, flavours)
#Z1 mass
mz1 = (leps[0] + leps[1]).m
if not mz1 > 40:
#print("fail mz1 mass cut")
return False
#leptons pt
leps_pt_sort = sorted(leps, key=lambda x:x.pt)
if not (leps_pt_sort[3].pt > 20 and leps_pt_sort[2].pt > 10):
#print("fail leptons pt cut")
#print("leptons pt: ", [a.pt for a in leps_pt_sort])
return False
#OS invariant mass (pos, neg, pos, neg)
mza = (leps[0]+leps[3]).m
mzb = (leps[1]+leps[2]).m
mz2 = (leps[2]+leps[3]).m
if not (mza > 4 and mzb > 4 and mz2 > 4):
#print("fail os cut")
return False
#4l invariant mass
m4l = (leps[0]+leps[1]+leps[2]+leps[3]).m
#print("m4l: ", m4l)
if not m4l>70:
#print("fail m4l cut")
return False
#same flavour cut
if isSF == True:
mz_pdg = 91.2
mzab_ord = [mza, mzb] if min([mza,mzb], key=lambda x:abs(x-mz_pdg)) == mza else [mzb,mza]
if (abs(mzab_ord[0]-mz_pdg) < abs(mz1-mz_pdg) and mzb < mzab_ord[0]):
#print("fail SF cut")
return False
#if candidate passes all cuts
return True
#choose ZZ candidate for which Z1 has mass closest to Z
#keep track of the flavour of Z1 and Z2
def choose_final_ZZ(ZZlist, part_list_e, part_list_mu, flavourslist):
if len(ZZlist) == 1:
return ZZlist[0], flavourslist[0]
else:
mz1list = []
for ZZcand, flavours in zip(ZZlist,flavourslist):
part_list1, part_list2 = set_flavours_lists(part_list_e, part_list_mu, flavours)
mz1 = (part_list1[ZZcand[0]] + part_list1[ZZcand[1]]).m
mz1list.append(mz1)
mz1min = min(mz1list, key=lambda x:abs(x-91.2))
idx = mz1list.index(mz1min)
return ZZlist[idx], flavourslist[idx]
##################### MAIN CODE TO FIND BEST ZZ CANDIDATE #####################
#default values
isZZcand = 0
lep1 = MadMinerParticle(); lep2 = MadMinerParticle(); lep3 = MadMinerParticle(); lep4 = MadMinerParticle()
#list of candidates electrons/muons: 1 if candidate, 0 if not
candidate_es = np.ones(len(objects["e"]))
candidate_mus = np.ones(len(objects["mu"]))
#print(len(objects["e"]), len(objects["mu"]))
#find candidate leptons: eta and pt cuts (to be checked)
for idx, el in enumerate(objects["e"]):
if not (abs(el.eta) < 2.5 and el.pt > 7):
candidate_es[idx] = 0
for idx, mu in enumerate(objects["mu"]):
if not (abs(mu.eta) < 2.4 and mu.pt > 5):
candidate_mus[idx] = 0
#find Z candidates
e_Z_cand = find_Z_cand(candidate_es, objects["e"])
mu_Z_cand = find_Z_cand(candidate_mus, objects["mu"])
#find ZZ candidates
ZZ_cand_ee_list, ZZ_cand_mm_list, ZZ_cand_em_list = find_ZZ_cand(e_Z_cand,mu_Z_cand)
#initialise lists for ZZ candidates which pass cuts, and keep track of flavours
ZZ_cands_final = []
ZZ_cands_final_flavours = []
#same flavours, electrons
for ZZ_cand_ee in ZZ_cand_ee_list:
ZZ_cand_ee, swapped = Z1Z2_ordering(ZZ_cand_ee, objects["e"], objects["mu"], flavours="ee")
#apply cuts
if ZZ_cuts(ZZ_cand_ee, objects["e"], objects["mu"], flavours="ee", isSF=True) == True:
ZZ_cands_final.append(ZZ_cand_ee)
ZZ_cands_final_flavours.append("ee")
#same flavours, muons
for ZZ_cand_mm in ZZ_cand_mm_list:
ZZ_cand_mm, swapped = Z1Z2_ordering(ZZ_cand_mm, objects["e"], objects["mu"], flavours="mm")
#apply cuts
if ZZ_cuts(ZZ_cand_mm, objects["e"], objects["mu"], flavours="mm", isSF=True) == True:
ZZ_cands_final.append(ZZ_cand_mm)
ZZ_cands_final_flavours.append("mm")
#different flavours
for ZZ_cand_em in ZZ_cand_em_list:
ZZ_cand_em, swapped = Z1Z2_ordering(ZZ_cand_em, objects["e"], objects["mu"], flavours="em")
#apply cuts, careful when swapping em/me
flavours_OF = "em" if swapped==False else "me"
if ZZ_cuts(ZZ_cand_em, objects["e"], objects["mu"], flavours=flavours_OF, isSF=False) == True:
ZZ_cands_final.append(ZZ_cand_em)
ZZ_cands_final_flavours.append(flavours_OF)
#find final best ZZ candidate
#if more than one ZZ candidate is left: choose the one with Z1 mass closest to MZ
def printlep(lep, name):
#print(name, " = [", lep.px, ",", lep.py, ",", lep.pz, "]")
print("[", lep.e,",", lep.px, ",", lep.py, ",", lep.pz, "]")
return
if len(ZZ_cands_final) > 0:
isZZcand = 1
ZZfinal, flavours = choose_final_ZZ(ZZ_cands_final, objects["e"], objects["mu"], ZZ_cands_final_flavours)
lep1, lep2, lep3, lep4 = ZZidx_to_leps(ZZfinal, objects["e"], objects["mu"], flavours)
#print("chosen leptons: (flavours", flavours)
#for lep, nlep in zip([lep1,lep2,lep3,lep4],["a","b","c","d"]):
#printlep(lep, nlep)
#print((lep1+lep2+lep3+lep4).m)
#print((lep1+lep2).m, (lep3+lep4).m)
#print("original leptons: e")
#for idx, lep in enumerate(objects["e"]):
#nlep = "ne" + str(idx)
#printlep(lep, nlep)
#print("original leptons: mu")
#for idx, lep in enumerate(objects["mu"]):
#nlep = "nmu" + str(idx)
#printlep(lep, nlep)
#logger.debug("Value of isZZcand: %d", isZZcand)
#debugging
#print(len(objects["e"]), len(objects["mu"]))
#if(len(ZZ_cand_ee_list) + len(ZZ_cand_mm_list) + len(ZZ_cand_em_list)) >= 1:
#print(len(objects["e"]), len(objects["mu"]))
#if isZZcand != 1:
#print("no ZZ cand found after cuts")
#else:
#print("ok candidate and found")
#update objects dictionary
objects.update(
{
"isZZcand": isZZcand,
"lep1ZZ": lep1,
"lep2ZZ": lep2,
"lep3ZZ": lep3,
"lep4ZZ": lep4,
}
)
return objects
def _get_objects(particles, particles_truth, met_resolution=None, global_event_data=None):
# Find visible particles
electrons = []
muons = []
taus = []
photons = []
jets = []
leptons = []
neutrinos = []
unstables = []
invisibles = []
for particle in particles:
pdgid = abs(particle.pdgid)
if pdgid in [1, 2, 3, 4, 5, 6, 9, 21]:
if pdgid == 5:
particle.set_tags(False, True, False)
if pdgid == 6:
particle.set_tags(False, False, True)
jets.append(particle)
elif pdgid == 11:
electrons.append(particle)
leptons.append(particle)
elif pdgid == 13:
muons.append(particle)
leptons.append(particle)
elif pdgid == 15:
taus.append(particle)
elif pdgid == 22:
photons.append(particle)
elif pdgid in [12, 14, 16]:
neutrinos.append(particle)
invisibles.append(particle)
elif pdgid in [23, 24, 25]:
unstables.append(particle)
else:
logger.warning("Unknown particle with PDG id %s, treating as invisible!")
invisibles.append(particle)
# Sort by pT
electrons = sorted(electrons, reverse=True, key=lambda x: x.pt)
muons = sorted(muons, reverse=True, key=lambda x: x.pt)
taus = sorted(taus, reverse=True, key=lambda x: x.pt)
photons = sorted(photons, reverse=True, key=lambda x: x.pt)
leptons = sorted(leptons, reverse=True, key=lambda x: x.pt)
neutrinos = sorted(neutrinos, reverse=True, key=lambda x: x.pt)
jets = sorted(jets, reverse=True, key=lambda x: x.pt)
# Sum over all particles
ht = 0.0
visible_sum = MadMinerParticle.from_xyzt(0.0, 0.0, 0.0, 0.0)
standard_ids = set(get_elementary_pdg_ids())
neutrino_ids = set(
int(p.pdgid)
for p in Particle.findall(lambda p: p.pdgid.is_sm_lepton and p.charge == 0)
)
for particle in particles:
if particle.pdgid in standard_ids and particle.pdgid not in neutrino_ids:
visible_sum += particle
ht += particle.pt
# Soft noise
if met_resolution is not None:
noise_std = met_resolution[0] + met_resolution[1] * ht
noise_x = np.random.normal(0.0, noise_std, size=None)
noise_y = np.random.normal(0.0, noise_std, size=None)
else:
noise_x = 0.0
noise_y = 0.0
# MET
met_x = -visible_sum.x + noise_x
met_y = -visible_sum.y + noise_y
met = MadMinerParticle.from_xyzt(
x=met_x,
y=met_y,
z=0.0,
t=(met_x ** 2 + met_y ** 2) ** 0.5,
)
# Build objects
objects = math_commands()
objects.update(
{
"p": particles,
"p_truth": particles_truth,
"e": electrons,
"j": jets,
"a": photons,
"mu": muons,
"tau": taus,
"l": leptons,
"met": met,
"v": neutrinos,
}
)
# Global event_data
if global_event_data is not None:
objects.update(global_event_data)
return objects
