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)
