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()