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_packed_selection(): from coffea.processor import PackedSelection sel = PackedSelection() counts, test_eta, test_pt = dummy_jagged_eta_pt() all_true = np.full(shape=counts.shape, fill_value=True, dtype=np.bool) all_false = np.full(shape=counts.shape, fill_value=False, dtype=np.bool) ones = np.ones(shape=counts.shape, dtype=np.uint64) wrong_shape = ones = np.ones(shape=(counts.shape[0] - 5, ), dtype=np.bool) sel.add('all_true', all_true) sel.add('all_false', all_false) assert (np.all(sel.require(all_true=True, all_false=False) == all_true)) assert (np.all(sel.all('all_true', 'all_false') == all_false)) try: sel.require(all_true=1, all_false=0) except ValueError: pass try: sel.add('wrong_shape', wrong_shape) except ValueError: pass try: sel.add('ones', ones) except ValueError: pass try: overpack = PackedSelection() for i in range(65): overpack.add('sel_%d', all_true) except RuntimeError: pass
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)