def almir(
tree
): # Return a DataFrame that contains the information about WW invariant mass, lepton pair p_T and DeltaPli between jet_MET and leptonic-W_hadronic-W
Mw = 80.379 # Boson W mass
k = ((Mw**2) / 2 + get_branch(tree, 'muon_px') * get_branch(tree, 'METPx')
) + (get_branch(tree, 'muon_py') * get_branch(tree, 'METPy'))
raiz = (
(((k * get_branch(tree, 'muon_pz'))**2) /
(get_branch(tree, 'muon_pt')**4) -
((get_branch(tree, 'muon_E') * get_branch(tree, 'METPt'))**2 - k) /
get_branch(tree, 'muon_pt')**2)**0.5
).fillna(
0
) # .fillna(0) - replaces DataFrame NaN's that presents imaginary roots by 0
Pz_nu = (
(k * get_branch(tree, 'muon_pz') / (get_branch(tree, 'muon_pt')**2)) +
raiz) # Reconstructed neutrino's momentum z-component
W_lep_energy = get_branch(tree, 'muon_E') + (
get_branch(tree, 'METPx')**2 + get_branch(tree, 'METPy')**2 +
Pz_nu**2)**0.5 # Lepton pair energy
TLV_lep = uproot_methods.TLorentzVectorArray(
get_branch(tree, 'muon_px') + get_branch(tree, 'METPx'),
get_branch(tree, 'muon_py') + get_branch(tree, 'METPy'),
get_branch(tree, 'muon_pz') + Pz_nu,
W_lep_energy) # Lepton pair 4-vector
TLV_jet = uproot_methods.TLorentzVectorArray(get_branch(tree, 'jetAK8_px'),
get_branch(tree, 'jetAK8_py'),
get_branch(tree, 'jetAK8_pz'),
get_branch(tree, 'jetAK8_E'))
W_mass = (TLV_lep + TLV_jet).mass # WW invariant mass
W_lep_pt = (TLV_lep).pt # Lepton pair p_T
dphi_jet_lep = TLV_lep.phi - TLV_jet.phi
dphi_jet_lep = np.where(dphi_jet_lep >= scipy.constants.pi,
dphi_jet_lep - 2 * scipy.constants.pi,
dphi_jet_lep)
dphi_jet_lep = np.where(
dphi_jet_lep < -scipy.constants.pi,
dphi_jet_lep + 2 * scipy.constants.pi,
dphi_jet_lep) # delta phi between the jet and the lepton pair
dphi_jet_MET = get_branch(tree, 'METphi') - TLV_jet.phi
dphi_jet_MET = np.where(dphi_jet_MET >= scipy.constants.pi,
dphi_jet_MET - 2 * scipy.constants.pi,
dphi_jet_MET)
dphi_jet_MET = np.where(
dphi_jet_MET < -scipy.constants.pi,
dphi_jet_MET + 2 * scipy.constants.pi,
dphi_jet_MET) # delta phi between the jet e the MET
#acopl = (1. - np.abs(dphi)/scipy.constants.pi)
DataFrame = pd.DataFrame(
pd.concat([
pd.DataFrame(W_mass, columns=['Mass_WW']),
pd.DataFrame(W_lep_pt, columns=["W_lep_Pt"]),
pd.DataFrame(dphi_jet_lep, columns=['Dphi_jet_lep']),
pd.DataFrame(dphi_jet_MET, columns=['Dphi_jet_MET'])
],
axis=1))
return DataFrame
def boost_and_rotate(events):
"""
Given a list of events, boosts into the center of mass frame,
then rotates jets phi=0 for the vector sum of the first 3 jets in each event.
"""
# we assume events is rectangular (i.e. all events have same # of jets
njets = len(events.resolved_lv.pt[0])
# get vectors
vectors = events.resolved_lv
# get sum vectors
x_sum = np.repeat(np.sum(vectors.x, axis=1).reshape(-1, 1), njets, axis=1)
y_sum = np.repeat(np.sum(vectors.y, axis=1).reshape(-1, 1), njets, axis=1)
z_sum = np.repeat(np.sum(vectors.z, axis=1).reshape(-1, 1), njets, axis=1)
t_sum = np.repeat(np.sum(vectors.t, axis=1).reshape(-1, 1), njets, axis=1)
v_sum = urm.TLorentzVectorArray(x_sum, y_sum, z_sum, t_sum)
# b for boosted
vectors_b = vectors.boost(-v_sum.boostp3)
v_sum_b = v_sum.boost(-v_sum.boostp3)
# for filler data where eta = 0, we'll have NaN for eta, replace that
eta = np.nan_to_num(vectors_b.eta, nan=0.0)
vectors_b = urm.TLorentzVectorArray.from_ptetaphie(vectors_b.pt, eta,
vectors_b.phi,
vectors_b.E)
# now rotate the system based on the first 3 jets
# get sum of the first 3, similarly to before
x_sum3 = np.repeat(np.sum(vectors_b.x[:, :3], axis=1).reshape(-1, 1),
njets,
axis=1)
y_sum3 = np.repeat(np.sum(vectors_b.y[:, :3], axis=1).reshape(-1, 1),
njets,
axis=1)
z_sum3 = np.repeat(np.sum(vectors_b.z[:, :3], axis=1).reshape(-1, 1),
njets,
axis=1)
t_sum3 = np.repeat(np.sum(vectors_b.t[:, :3], axis=1).reshape(-1, 1),
njets,
axis=1)
v_sum3 = urm.TLorentzVectorArray(x_sum3, y_sum3, z_sum3, t_sum3)
# rotate about z so that phi=0 for v_sum3
vectors_r = vectors_b.rotatez(-v_sum3.phi)
v_sum3 = v_sum3.rotatez(-v_sum3.phi)
# and again replace filler etas with 0
print(
"Don't worry if you see a warning about dividing by zero, fixing that!"
)
eta = np.nan_to_num(vectors_r.eta, nan=0.0)
vectors_final = urm.TLorentzVectorArray.from_ptetaphie(
vectors_r.pt, eta, vectors_r.phi, vectors_r.E)
events.resolved_lv = vectors_final
return events
def candidatesfromoffsets(cls,offsets,p4,**kwargs):
items = {}
if isinstance(p4,uproot_methods.TLorentzVectorArray):
items['p4'] = p4
else:
items['p4'] = uproot_methods.TLorentzVectorArray(p4[:,0],p4[:,1],
p4[:,2],p4[:,3])
thep4 = items['p4']
items['__fast_pt'] = fast_pt(thep4)
items['__fast_eta'] = fast_eta(thep4)
items['__fast_phi'] = fast_phi(thep4)
items['__fast_mass'] = fast_mass(thep4)
items.update(kwargs)
return cls.fromoffsets(offsets,awkward.Table(items))
def test_analysis_objects():
counts, px, py, pz, energy = dummy_four_momenta()
thep4 = np.stack((px, py, pz, energy)).T
#test JaggedTLorentzVectorArray
tlva1 = uproot_methods.TLorentzVectorArray(px, py, pz, energy)
tlva2 = uproot_methods.TLorentzVectorArray(thep4[:, 0], thep4[:, 1],
thep4[:, 2], thep4[:, 3])
jtlva1 = JaggedTLorentzVectorArray.fromcounts(counts, tlva1)
jtlva2 = JaggedTLorentzVectorArray.fromcounts(counts, tlva2)
jtlva1_selection1 = jtlva1[jtlva1.counts > 0]
jtlva1_selection2 = jtlva1_selection1[jtlva1_selection1.pt > 5]
jtlva2_selection1 = jtlva2[jtlva2.counts > 0]
jtlva2_selection2 = jtlva1_selection1[jtlva2_selection1.pt > 5]
diffx = np.abs(jtlva1.x - jtlva2.x)
diffy = np.abs(jtlva1.y - jtlva2.y)
diffz = np.abs(jtlva1.z - jtlva2.z)
difft = np.abs(jtlva1.t - jtlva2.t)
assert (diffx < 1e-8).flatten().all()
assert (diffy < 1e-8).flatten().all()
assert (diffz < 1e-8).flatten().all()
assert (difft < 1e-8).flatten().all()
#test JaggedCandidateArray
jca1 = JaggedCandidateArray.candidatesfromcounts(counts, p4=thep4)
jca2 = JaggedCandidateArray.candidatesfromcounts(counts, p4=thep4)
assert ((jca1.offsets == jca2.offsets).all())
addon1 = jca1.zeros_like()
addon2 = jca2.ones_like()
jca1['addon'] = addon1
jca2['addon'] = addon2
jca1.add_attributes(addonFlat=addon1.flatten(), addonJagged=addon1)
diffm = np.abs(jca1.p4.mass - jca2.p4.mass)
assert ((jca1.offsets == jca2.offsets).all())
diffpt = np.abs(jca1.p4.pt - jca2.p4.pt)
assert ((jca1.offsets == jca2.offsets).all())
eta2 = jca2.p4.eta
eta1 = jca1.p4.eta
print(np.sum(eta1.counts), np.sum(eta2.counts))
diffeta_temp = np.abs(eta1 - eta2)
diffeta = np.abs(jca1.p4.eta - jca2.p4.eta)
assert ((jca1.offsets == jca2.offsets).all())
assert (diffm < 1e-8).flatten().all()
assert (diffpt < 1e-8).flatten().all()
assert (diffeta < 1e-8).flatten().all()
#test fast functions
fastfs = ['pt', 'eta', 'phi', 'mass']
for func in fastfs:
func1 = getattr(jca1, func)
func2 = getattr(jca1.p4, func)
dfunc = np.abs(func1 - func2)
assert (dfunc < 1e-8).flatten().all()
adistinct = jca1.distincts()
apair = jca1.pairs()
across = jca1.cross(jca2)
assert 'p4' in adistinct.columns
assert 'p4' in apair.columns
assert 'p4' in across.columns
admsum = (adistinct.i0.p4 + adistinct.i1.p4).mass
apmsum = (apair.i0.p4 + apair.i1.p4).mass
acmsum = (across.i0.p4 + across.i1.p4).mass
diffadm = np.abs(adistinct.p4.mass - admsum)
diffapm = np.abs(apair.p4.mass - apmsum)
diffacm = np.abs(across.p4.mass - acmsum)
assert (diffadm < 1e-8).flatten().all()
assert (diffapm < 1e-8).flatten().all()
assert (diffacm < 1e-8).flatten().all()
selection11 = jca1[jca1.counts > 0]
selection12 = selection11[selection11.p4.pt > 5]
selection21 = jca2[jca2.counts > 0]
selection22 = selection21[selection21.p4.pt > 5]
diffcnts = selection12.counts - jtlva1_selection2.counts
diffm = np.abs(selection12.p4.mass - jtlva1_selection2.mass)
diffaddon = selection12.addon - selection22.addon
assert (diffcnts == 0).flatten().all()
assert (diffm < 1e-8).flatten().all()
assert (diffaddon == -1).flatten().all()
#test gen-reco matching
gen, reco = gen_reco_TLV()
flat_gen = gen.flatten()
gen_px, gen_py, gen_pz, gen_e = flat_gen.x, flat_gen.y, flat_gen.z, flat_gen.t
flat_reco = reco.flatten()
reco_px, reco_py, reco_pz, reco_e = flat_reco.x, flat_reco.y, flat_reco.z, flat_reco.t
jca_gen = JaggedCandidateArray.candidatesfromcounts(gen.counts,
px=gen_px,
py=gen_py,
pz=gen_pz,
energy=gen_e)
jca_reco = JaggedCandidateArray.candidatesfromcounts(reco.counts,
px=reco_px,
py=reco_py,
pz=reco_pz,
energy=reco_e)
print('gen eta: ', jca_gen.p4.eta, '\n gen phi:', jca_gen.p4.phi)
print('reco eta: ', jca_reco.p4.eta, '\n reco phi:', jca_reco.p4.phi)
print('match mask: ', jca_reco.match(jca_gen, deltaRCut=0.3))
print('arg matches: ', jca_reco.argmatch(jca_gen, deltaRCut=0.3))
argmatch_nocut = jca_gen.argmatch(jca_reco).flatten()
argmatch_dr03 = jca_gen.argmatch(jca_reco, deltaRCut=0.3).flatten()
argmatch_dr03_dpt01 = jca_gen.argmatch(jca_reco,
deltaRCut=0.3,
deltaPtCut=0.1).flatten()
assert (argmatch_nocut.size == 5)
assert (argmatch_dr03[argmatch_dr03 != -1].size == 3)
assert (argmatch_dr03_dpt01[argmatch_dr03_dpt01 != -1].size == 2)
assert (jca_gen.match(jca_reco,
deltaRCut=0.3).flatten().flatten().sum() == 3)
assert (jca_gen.match(jca_reco, deltaRCut=0.3,
deltaPtCut=0.1).flatten().flatten().sum() == 2)
def almir(
file
): # Funcao que retorna um DataFrame que contem a massa invariante do WW, pt do par de le
trigger = (open_files_PF(file, 'HLT_pass') == 1)[:, 3]
Mw = 80.379 # massa do boson W
jetAK8_pt = open_files_muon(file, 'jetAK8_pt')[trigger]
jetAK8_prunedMass = open_files_muon(file, 'jetAK8_prunedMass')[trigger]
jetAK8_tau21 = open_files_muon(file, 'jetAK8_tau21')[trigger]
jetAK8_eta = open_files_muon(file, 'jetAK8_eta')[trigger]
jetAK8_px = open_files_muon(file, 'jetAK8_px')[trigger]
jetAK8_py = open_files_muon(file, 'jetAK8_py')[trigger]
jetAK8_pz = open_files_muon(file, 'jetAK8_pz')[trigger]
jetAK8_E = open_files_muon(file, 'jetAK8_E')[trigger]
print('jetAK8_pt -->', jetAK8_pt)
METPt = open_files_muon(file, 'METPt')[trigger]
METPx = open_files_muon(file, 'METPx')[trigger]
METPy = open_files_muon(file, 'METPy')[trigger]
METphi = open_files_muon(file, 'METphi')[trigger]
print('METphi -->', METphi)
muon_pt = open_files_muon(file, 'muon_pt')[trigger]
muon_eta = open_files_muon(file, 'muon_eta')[trigger]
muon_phi = open_files_muon(file, 'muon_phi')[trigger]
muon_px = open_files_muon(file, 'muon_px')[trigger]
muon_py = open_files_muon(file, 'muon_py')[trigger]
muon_pz = open_files_muon(file, 'muon_pz')[trigger]
muon_E = open_files_muon(file, 'muon_E')[trigger]
print('Muon_pt --> ', muon_pt)
k = ((Mw**2) / 2 + muon_px * METPx) + (muon_py * METPy)
raiz_ = ((((k * muon_pz)**2) / (muon_pt**4) -
((muon_E * METPt)**2 - k) / muon_pt**2)**0.5)
raiz = np.nan_to_num(
raiz_
) # Os valores de NaN, causados pela divisão por 0 ou pelo resultado de uma raiz imaginária, é substituida por NaN
Pz_nu = ((k * muon_pz / (muon_pt**2)) + raiz
) # coordenada z do momentum do neutrino reconstruido
W_lep_energy = (muon_E + (METPx**2 + METPy**2 + Pz_nu**2)**0.5
) # Energia do par de léptons
print('energia do W leptonico -->', W_lep_energy)
# Usamos o TLorenctzVector do Python
TLV_lep = uproot_methods.TLorentzVectorArray(
muon_px + METPx, muon_py + METPy, muon_pz + Pz_nu,
W_lep_energy) # 4-vector do par de lepton
TLV_jet = uproot_methods.TLorentzVectorArray(jetAK8_px, jetAK8_py,
jetAK8_pz, jetAK8_E)
W_mass = (TLV_lep + TLV_jet).mass # Massa invariante do WW
W_lep_pt = (TLV_lep).pt # Pt do par de lepton
print('Massa W -->', W_mass)
dphi_jet_lep = TLV_lep.phi - TLV_jet.phi
dphi_jet_lep = np.where(dphi_jet_lep >= scipy.constants.pi,
dphi_jet_lep - 2 * scipy.constants.pi,
dphi_jet_lep)
dphi_jet_lep = np.where(
dphi_jet_lep < -scipy.constants.pi,
dphi_jet_lep + 2 * scipy.constants.pi,
dphi_jet_lep) # delta phi entre o jato e o par de lepton
dphi_jet_MET = METphi - TLV_jet.phi
dphi_jet_MET = np.where(dphi_jet_MET >= scipy.constants.pi,
dphi_jet_MET - 2 * scipy.constants.pi,
dphi_jet_MET)
dphi_jet_MET = np.where(dphi_jet_MET < -scipy.constants.pi,
dphi_jet_MET + 2 * scipy.constants.pi,
dphi_jet_MET) # delta phi entre o jato e o MET
'''
** numeração das colunas do numpy array ** ( para facilitar na hora de fazer os cortes )
0 --> massa do WW
1 --> Pt do W leptônico
2 --> DeltaPhi entre W_hadrônico e W_leptônico
3 --> DeltaPhi entre Jatos e o MET
4 --> jetAK8_pt
5 --> jetAK8_eta
6 --> jetAK8_prunedMass
7 --> jetAK8_tau21
8 --> METPt
9 --> muon_pt
10 --> muon_eta
11 --> ExtraTracks
'''
pfeta = open_files_PF(file, 'pfeta')[trigger]
pfphi = open_files_PF(file, 'pfphi')[trigger]
pffromPV = open_files_PF(file, 'pffromPV')[trigger]
dR_muon = ((pfeta - muon_eta)**2 + (pfphi - muon_phi)**2)**0.5
dR_jet = ((pfeta - TLV_jet.eta)**2 + (pfphi - TLV_jet.phi)**2)**0.5
print('dR_muon', dR_muon)
dR_muon = dR_muon[pffromPV == 3]
dR_jet = dR_jet[pffromPV == 3]
dR_jet = dR_jet[dR_muon > 0.3]
dR_jet = dR_jet[dR_jet > 0.8]
ExtraTracks = [len(arr) for arr in dR_jet]
ExtraTracks = np.array(ExtraTracks).reshape(-1, 1)
print('Tracos Extras -->', ExtraTracks)
events_all = np.concatenate(
(W_mass.reshape(-1, 1), W_lep_pt.reshape(
-1, 1), dphi_jet_lep.reshape(-1, 1), dphi_jet_MET.reshape(
-1, 1), jetAK8_pt.reshape(-1, 1), TLV_jet.eta.reshape(-1, 1),
jetAK8_prunedMass.reshape(-1, 1), jetAK8_tau21.reshape(
-1, 1), METPt.reshape(-1, 1), muon_pt.reshape(
-1, 1), muon_eta.reshape(-1, 1), ExtraTracks),
axis=1) # concatenando todos as variáveis
events_all_cut = (events_all[:, 4] >= 200) & (events_all[:, 5] <= 2.4) & (
events_all[:, 8] >= 40) & (events_all[:, 9] >= 53) & (
events_all[:, 10] <= 2.4) # realizando os cortes nas variáveis
array_numpy = events_all[events_all_cut]
#columns = ['Mww','Pt_W_lep','dPhi_Whad_Wlep','dPhi_jatos_MET','jetAK8_pt','jetAK8_eta','jetAK8_prunedMass','jetAK8_tau21','METPt','muon_pt','muon_eta','ExtraTracks']
#DataFrame = pd.DataFrame( array_numpy , columns = columns )
print('array completo --> ', array_numpy)
#return DataFrame # ou retorna o DataFrame com as colunas
return array_numpy # ou retorna um matriz numpy aninhada para economizar memória
def test_analysis_objects():
counts, px, py, pz, energy = dummy_four_momenta()
thep4 = np.stack((px, py, pz, energy)).T
#test JaggedTLorentzVectorArray
tlva1 = uproot_methods.TLorentzVectorArray(px, py, pz, energy)
tlva2 = uproot_methods.TLorentzVectorArray(thep4[:, 0], thep4[:, 1],
thep4[:, 2], thep4[:, 3])
jtlva1 = JaggedTLorentzVectorArray.fromcounts(counts, tlva1)
jtlva2 = JaggedTLorentzVectorArray.fromcounts(counts, tlva2)
jtlva1_selection1 = jtlva1[jtlva1.counts > 0]
jtlva1_selection2 = jtlva1_selection1[jtlva1_selection1.pt > 5]
jtlva2_selection1 = jtlva2[jtlva2.counts > 0]
jtlva2_selection2 = jtlva1_selection1[jtlva2_selection1.pt > 5]
diffx = np.abs(jtlva1.x - jtlva2.x)
diffy = np.abs(jtlva1.y - jtlva2.y)
diffz = np.abs(jtlva1.z - jtlva2.z)
difft = np.abs(jtlva1.t - jtlva2.t)
assert (diffx < 1e-8).flatten().all()
assert (diffy < 1e-8).flatten().all()
assert (diffz < 1e-8).flatten().all()
assert (difft < 1e-8).flatten().all()
#test JaggedCandidateArray
jca1 = JaggedCandidateArray.candidatesfromcounts(counts, p4=thep4)
jca2 = JaggedCandidateArray.candidatesfromcounts(counts, p4=thep4)
assert ((jca1.offsets == jca2.offsets).all())
addon1 = jca1.zeros_like()
addon2 = jca2.ones_like()
jca1['addon'] = addon1
jca2['addon'] = addon2
diffm = np.abs(jca1.p4.mass - jca2.p4.mass)
assert ((jca1.offsets == jca2.offsets).all())
diffpt = np.abs(jca1.p4.pt - jca2.p4.pt)
assert ((jca1.offsets == jca2.offsets).all())
eta2 = jca2.p4.eta
eta1 = jca1.p4.eta
print(np.sum(eta1.counts), np.sum(eta2.counts))
diffeta_temp = np.abs(eta1 - eta2)
diffeta = np.abs(jca1.p4.eta - jca2.p4.eta)
assert ((jca1.offsets == jca2.offsets).all())
assert (diffm < 1e-8).flatten().all()
assert (diffpt < 1e-8).flatten().all()
assert (diffeta < 1e-8).flatten().all()
#test fast functions
fastfs = ['pt', 'eta', 'phi', 'mass']
for func in fastfs:
func1 = getattr(jca1, func)
func2 = getattr(jca1.p4, func)
dfunc = np.abs(func1 - func2)
assert (dfunc < 1e-8).flatten().all()
adistinct = jca1.distincts()
apair = jca1.pairs()
across = jca1.cross(jca2)
assert 'p4' in adistinct.columns
assert 'p4' in apair.columns
assert 'p4' in across.columns
admsum = (adistinct.i0.p4 + adistinct.i1.p4).mass
apmsum = (apair.i0.p4 + apair.i1.p4).mass
acmsum = (across.i0.p4 + across.i1.p4).mass
diffadm = np.abs(adistinct.p4.mass - admsum)
diffapm = np.abs(apair.p4.mass - apmsum)
diffacm = np.abs(across.p4.mass - acmsum)
assert (diffadm < 1e-8).flatten().all()
assert (diffapm < 1e-8).flatten().all()
assert (diffacm < 1e-8).flatten().all()
selection11 = jca1[jca1.counts > 0]
selection12 = selection11[selection11.p4.pt > 5]
selection21 = jca2[jca2.counts > 0]
selection22 = selection21[selection21.p4.pt > 5]
diffcnts = selection12.counts - jtlva1_selection2.counts
diffm = np.abs(selection12.p4.mass - jtlva1_selection2.mass)
diffaddon = selection12.addon - selection22.addon
assert (diffcnts == 0).flatten().all()
assert (diffm < 1e-8).flatten().all()
assert (diffaddon == -1).flatten().all()