def test_jet_transformer(): from coffea.analysis_objects import JaggedCandidateArray as CandArray from coffea.jetmet_tools import (FactorizedJetCorrector, JetResolution, JetResolutionScaleFactor, JetCorrectionUncertainty, JetTransformer) counts, test_px, test_py, test_pz, test_e = dummy_four_momenta() test_Rho = np.full(shape=(np.sum(counts),), fill_value=100.) test_A = np.full(shape=(np.sum(counts),), fill_value=5.) jets = CandArray.candidatesfromcounts(counts, px=test_px, py=test_py, pz=test_pz, energy=test_e) jets.add_attributes(ptRaw=jets.pt, massRaw=jets.mass, rho=test_Rho, area=test_A) jec_names = ['Summer16_23Sep2016V3_MC_L1FastJet_AK4PFPuppi', 'Summer16_23Sep2016V3_MC_L2Relative_AK4PFPuppi', 'Summer16_23Sep2016V3_MC_L2L3Residual_AK4PFPuppi', 'Summer16_23Sep2016V3_MC_L3Absolute_AK4PFPuppi'] corrector = FactorizedJetCorrector(**{name: evaluator[name] for name in jec_names}) junc_names = [] for name in dir(evaluator): if 'Summer16_23Sep2016V3_MC_UncertaintySources_AK4PFPuppi' in name: junc_names.append(name) junc = JetCorrectionUncertainty(**{name: evaluator[name] for name in junc_names}) jer_names = ['Spring16_25nsV10_MC_PtResolution_AK4PFPuppi'] reso = JetResolution(**{name: evaluator[name] for name in jer_names}) jersf_names = ['Spring16_25nsV10_MC_SF_AK4PFPuppi'] resosf = JetResolutionScaleFactor(**{name: evaluator[name] for name in jersf_names}) xform = JetTransformer(jec=corrector, junc=junc, jer=reso, jersf=resosf) print(xform.uncertainties) xform.transform(jets) print(jets.columns) assert('pt_jer_up' in jets.columns) assert('pt_jer_down' in jets.columns) assert('mass_jer_up' in jets.columns) assert('mass_jer_down' in jets.columns) for unc in xform.uncertainties: assert('pt_'+unc+'_up' in jets.columns) assert('pt_'+unc+'_down' in jets.columns) assert('mass_'+unc+'_up' in jets.columns) assert('mass_'+unc+'_down' in jets.columns)
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 test_jet_transformer(): import numpy as np import awkward as ak import math from coffea.analysis_objects import JaggedCandidateArray as CandArray from coffea.jetmet_tools import (FactorizedJetCorrector, JetResolution, JetResolutionScaleFactor, JetCorrectionUncertainty, JetTransformer) counts, test_px, test_py, test_pz, test_e = dummy_four_momenta() test_Rho = np.full(shape=(np.sum(counts), ), fill_value=100.) test_A = np.full(shape=(np.sum(counts), ), fill_value=5.) jets = CandArray.candidatesfromcounts(counts, px=test_px, py=test_py, pz=test_pz, energy=test_e) jets.add_attributes(ptRaw=jets.pt, massRaw=jets.mass, rho=test_Rho, area=test_A) fakemet = np.random.exponential(scale=1.0, size=counts.size) metphi = np.random.uniform(low=-math.pi, high=math.pi, size=counts.size) syst_up = 0.001 * fakemet syst_down = -0.001 * fakemet met = CandArray.candidatesfromcounts( np.ones_like(counts), pt=fakemet, eta=np.zeros_like(counts), phi=metphi, mass=np.zeros_like(counts), MetUnclustEnUpDeltaX=syst_up * np.cos(metphi), MetUnclustEnUpDeltaY=syst_down * np.sin(metphi)) jec_names = [ 'Summer16_23Sep2016V3_MC_L1FastJet_AK4PFPuppi', 'Summer16_23Sep2016V3_MC_L2Relative_AK4PFPuppi', 'Summer16_23Sep2016V3_MC_L2L3Residual_AK4PFPuppi', 'Summer16_23Sep2016V3_MC_L3Absolute_AK4PFPuppi' ] corrector = FactorizedJetCorrector( **{name: evaluator[name] for name in jec_names}) junc_names = [] for name in dir(evaluator): if 'Summer16_23Sep2016V3_MC_UncertaintySources_AK4PFPuppi' in name: junc_names.append(name) junc = JetCorrectionUncertainty( **{name: evaluator[name] for name in junc_names}) jer_names = ['Spring16_25nsV10_MC_PtResolution_AK4PFPuppi'] reso = JetResolution(**{name: evaluator[name] for name in jer_names}) jersf_names = ['Spring16_25nsV10_MC_SF_AK4PFPuppi'] resosf = JetResolutionScaleFactor( **{name: evaluator[name] for name in jersf_names}) xform = JetTransformer(jec=corrector, junc=junc, jer=reso, jersf=resosf) print(xform.uncertainties) xform.transform(jets, met=met) print('jets', jets.columns) print('met', met.columns) assert ('pt_jer_up' in jets.columns) assert ('pt_jer_down' in jets.columns) assert ('mass_jer_up' in jets.columns) assert ('mass_jer_down' in jets.columns) assert ('pt_UnclustEn_up' in met.columns) assert ('pt_UnclustEn_down' in met.columns) assert ('phi_UnclustEn_up' in met.columns) assert ('phi_UnclustEn_down' in met.columns) for unc in xform.uncertainties: assert ('pt_' + unc + '_up' in jets.columns) assert ('pt_' + unc + '_down' in jets.columns) assert ('mass_' + unc + '_up' in jets.columns) assert ('mass_' + unc + '_down' in jets.columns) assert ('pt_' + unc + '_up' in met.columns) assert ('phi_' + unc + '_up' in met.columns)
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()