def gen_reco_TLV():
from coffea.nanoevents.methods import vector
ak.behavior.update(vector.behavior)
gen_pt = ak.Array([[10.0, 20.0, 30.0], [], [40.0, 50.0]])
reco_pt = ak.Array([[20.2, 10.1, 30.3, 50.5], [50.5], [60]])
gen_eta = ak.Array([[-3.0, -2.0, 2.0], [], [-1.0, 1.0]])
reco_eta = ak.Array([[-2.2, -3.3, 2.2, 0.0], [0.0], [1.1]])
gen_phi = ak.Array([[-1.5, 0.0, 1.5], [], [0.78, -0.78]])
reco_phi = ak.Array([[0.1, -1.4, 1.4, 0.78], [0.78], [-0.77]])
gen = ak.zip(
{
"pt": gen_pt,
"eta": gen_eta,
"phi": gen_phi,
"mass": ak.full_like(gen_pt, 0.2),
},
with_name="PtEtaPhiMLorentzVector",
)
reco = ak.zip(
{
"pt": reco_pt,
"eta": reco_eta,
"phi": reco_phi,
"mass": ak.full_like(reco_pt, 0.2),
},
with_name="PtEtaPhiMLorentzVector",
)
return (gen, reco)
def jet_features(jets, mask_bool=False, mask=None):
vecs = ak.zip(
{
"pt": jets[:, :, 2:3],
"eta": jets[:, :, 0:1],
"phi": jets[:, :, 1:2],
"mass": ak.full_like(jets[:, :, 2:3], 0),
},
with_name="PtEtaPhiMLorentzVector")
sum_vecs = vecs.sum(axis=1)
jf = np.nan_to_num(
np.array(np.concatenate((sum_vecs.mass, sum_vecs.pt), axis=1)))
return jf
def test():
array = ak.Array(
[
[{"x": 0.0, "y": []}, {"x": 1.1, "y": [1]}, {"x": 2.2, "y": [1, 2]}],
[],
[
{"x": 3.3, "y": [1, 2, None, 3]},
False,
False,
True,
{"x": 4.4, "y": [1, 2, None, 3, 4]},
],
]
)
def assert_array_type(new_array, intended_type):
"""helper function to check each part of the array took the intended type"""
assert isinstance(new_array[0][0]["x"], intended_type)
assert isinstance(new_array[0][1]["x"], intended_type)
assert isinstance(new_array[0][1]["y"][0], intended_type)
assert isinstance(new_array[0][2]["x"], intended_type)
assert isinstance(new_array[0][2]["y"][0], intended_type)
assert isinstance(new_array[0][2]["y"][1], intended_type)
assert isinstance(new_array[2][0]["x"], intended_type)
assert isinstance(new_array[2][0]["y"][0], intended_type)
assert isinstance(new_array[2][0]["y"][1], intended_type)
assert isinstance(new_array[2][0]["y"][3], intended_type)
assert isinstance(new_array[2][1], intended_type)
assert isinstance(new_array[2][2], intended_type)
assert isinstance(new_array[2][3], intended_type)
assert isinstance(new_array[2][4]["x"], intended_type)
assert isinstance(new_array[2][4]["y"][0], intended_type)
assert isinstance(new_array[2][4]["y"][1], intended_type)
assert isinstance(new_array[2][4]["y"][3], intended_type)
assert isinstance(new_array[2][4]["y"][4], intended_type)
int_type = ak.full_like(array, 12, dtype=int)
float_type = ak.full_like(array, 12, dtype=float)
assert int_type.tolist() == float_type.tolist()
assert_array_type(int_type, int)
assert_array_type(float_type, float)
bool_type = ak.full_like(array, 12, dtype=bool)
assert_array_type(bool_type, bool)
int_type = ak.full_like(array, -1.2, dtype=int)
float_type = ak.full_like(array, -1.2, dtype=float)
bool_type = ak.full_like(array, -1.2, dtype=bool)
assert_array_type(int_type, int)
assert_array_type(float_type, float)
assert_array_type(bool_type, bool)
int_type = ak.zeros_like(array, dtype=int)
float_type = ak.zeros_like(array, dtype=float)
bool_type = ak.zeros_like(array, dtype=bool)
assert int_type.tolist() == float_type.tolist()
assert int_type.tolist() == bool_type.tolist()
assert_array_type(int_type, int)
assert_array_type(float_type, float)
assert_array_type(bool_type, bool)
int_type = ak.ones_like(array, dtype=int)
float_type = ak.ones_like(array, dtype=float)
bool_type = ak.ones_like(array, dtype=bool)
assert int_type.tolist() == float_type.tolist()
assert int_type.tolist() == bool_type.tolist()
assert_array_type(int_type, int)
assert_array_type(float_type, float)
assert_array_type(bool_type, bool)
array = ak.Array([["one", "two", "three"], [], ["four", "five"]])
def assert_array_type(new_array, intended_type):
"""helper function to check each part of the array took the intended type"""
assert isinstance(new_array[0][0], intended_type)
assert isinstance(new_array[0][1], intended_type)
assert isinstance(new_array[0][2], intended_type)
assert isinstance(new_array[2][0], intended_type)
assert isinstance(new_array[2][1], intended_type)
int_type = ak.full_like(array, 12, dtype=int)
float_type = ak.full_like(array, 12, dtype=float)
assert int_type.tolist() == float_type.tolist()
assert_array_type(int_type, int)
assert_array_type(float_type, float)
bool_type = ak.full_like(array, 12, dtype=bool)
assert_array_type(bool_type, bool)
int_type = ak.full_like(array, -1.2, dtype=int)
float_type = ak.full_like(array, -1.2, dtype=float)
bool_type = ak.full_like(array, -1.2, dtype=bool)
assert_array_type(int_type, int)
assert_array_type(float_type, float)
assert_array_type(bool_type, bool)
int_type = ak.zeros_like(array, dtype=int)
float_type = ak.zeros_like(array, dtype=float)
bool_type = ak.zeros_like(array, dtype=bool)
assert int_type.tolist() == float_type.tolist()
assert int_type.tolist() == bool_type.tolist()
assert_array_type(int_type, int)
assert_array_type(float_type, float)
assert_array_type(bool_type, bool)
int_type = ak.ones_like(array, dtype=int)
float_type = ak.ones_like(array, dtype=float)
bool_type = ak.ones_like(array, dtype=bool)
assert int_type.tolist() == float_type.tolist()
assert int_type.tolist() == bool_type.tolist()
assert_array_type(int_type, int)
assert_array_type(float_type, float)
assert_array_type(bool_type, bool)
def test():
array = ak.Array([
[{
"x": 0.0,
"y": []
}, {
"x": 1.1,
"y": [1]
}, {
"x": 2.2,
"y": [1, 2]
}],
[],
[
{
"x": 3.3,
"y": [1, 2, None, 3]
},
False,
False,
True,
{
"x": 4.4,
"y": [1, 2, None, 3, 4]
},
],
])
assert ak.full_like(array, 12.3).tolist() == [
[{
"x": 12.3,
"y": []
}, {
"x": 12.3,
"y": [12]
}, {
"x": 12.3,
"y": [12, 12]
}],
[],
[
{
"x": 12.3,
"y": [12, 12, None, 12]
},
True,
True,
True,
{
"x": 12.3,
"y": [12, 12, None, 12, 12]
},
],
]
assert ak.zeros_like(array).tolist() == [
[{
"x": 0.0,
"y": []
}, {
"x": 0.0,
"y": [0]
}, {
"x": 0.0,
"y": [0, 0]
}],
[],
[
{
"x": 0.0,
"y": [0, 0, None, 0]
},
False,
False,
False,
{
"x": 0.0,
"y": [0, 0, None, 0, 0]
},
],
]
assert ak.ones_like(array).tolist() == [
[{
"x": 1.0,
"y": []
}, {
"x": 1.0,
"y": [1]
}, {
"x": 1.0,
"y": [1, 1]
}],
[],
[
{
"x": 1.0,
"y": [1, 1, None, 1]
},
True,
True,
True,
{
"x": 1.0,
"y": [1, 1, None, 1, 1]
},
],
]
array = ak.Array([["one", "two", "three"], [], ["four", "five"]])
assert ak.full_like(array, "hello").tolist() == [
["hello", "hello", "hello"],
[],
["hello", "hello"],
]
assert ak.full_like(array, 1).tolist() == [["1", "1", "1"], [], ["1", "1"]]
assert ak.full_like(array, 0).tolist() == [["0", "0", "0"], [], ["0", "0"]]
assert ak.ones_like(array).tolist() == [["1", "1", "1"], [], ["1", "1"]]
assert ak.zeros_like(array).tolist() == [["", "", ""], [], ["", ""]]
array = ak.Array([[b"one", b"two", b"three"], [], [b"four", b"five"]])
assert ak.full_like(array, b"hello").tolist() == [
[b"hello", b"hello", b"hello"],
[],
[b"hello", b"hello"],
]
assert ak.full_like(array, 1).tolist() == [[b"1", b"1", b"1"], [],
[b"1", b"1"]]
assert ak.full_like(array, 0).tolist() == [[b"0", b"0", b"0"], [],
[b"0", b"0"]]
assert ak.ones_like(array).tolist() == [[b"1", b"1", b"1"], [],
[b"1", b"1"]]
assert ak.zeros_like(array).tolist() == [[b"", b"", b""], [], [b"", b""]]
def smear_factor(jetPt, pt_gen, jersf):
return (ak.full_like(jetPt, 1.0) +
(jersf[:, 0] - ak.full_like(jetPt, 1.0)) *
(jetPt - pt_gen) / jetPt)
def test_factorized_jet_corrector():
from coffea.jetmet_tools import FactorizedJetCorrector
counts, test_eta, test_pt = dummy_jagged_eta_pt()
test_Rho = np.full_like(test_eta, 100.0)
test_A = np.full_like(test_eta, 5.0)
# Check that the FactorizedJetCorrector is functional
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})
print(corrector)
pt_copy = np.copy(test_pt)
# Check that the corrector can be evaluated for flattened arrays
corrs = corrector.getCorrection(JetEta=test_eta,
Rho=test_Rho,
JetPt=test_pt,
JetA=test_A)
assert (np.abs(pt_copy - test_pt) < 1e-6).all()
test_pt_jag = ak.unflatten(test_pt, counts)
test_eta_jag = ak.unflatten(test_eta, counts)
test_Rho_jag = ak.unflatten(test_Rho, counts)
test_A_jag = ak.unflatten(test_A, counts)
# Check that the corrector can be evaluated for jagges arrays
corrs_jag = corrector.getCorrection(JetEta=test_eta_jag,
Rho=test_Rho_jag,
JetPt=test_pt_jag,
JetA=test_A_jag)
assert ak.all(np.abs(pt_copy - ak.flatten(test_pt_jag)) < 1e-6)
assert ak.all(np.abs(corrs - ak.flatten(corrs_jag)) < 1e-6)
# Check that the corrector returns the correct answers for each level of correction
# Use a subset of the values so that we can check the corrections by hand
test_pt_jag = test_pt_jag[0:3]
test_eta_jag = test_eta_jag[0:3]
test_Rho_jag = test_Rho_jag[0:3]
test_A_jag = test_A_jag[0:3]
counts = counts[0:3]
print("Raw jet values:")
print("pT:", test_pt_jag)
print("eta:", test_eta_jag)
print("rho:", test_Rho_jag)
print("area:", test_A_jag, "\n")
# Start by checking the L1 corrections
corrs_L1_jag_ref = ak.full_like(test_pt_jag, 1.0)
corrector = FactorizedJetCorrector(
**{name: evaluator[name]
for name in jec_names[0:1]})
corrs_L1_jag = corrector.getCorrection(JetEta=test_eta_jag,
Rho=test_Rho_jag,
JetPt=test_pt_jag,
JetA=test_A_jag)
print("Reference L1 corrections:", corrs_L1_jag_ref)
print("Calculated L1 corrections:", corrs_L1_jag)
assert ak.all(
np.abs(ak.flatten(corrs_L1_jag_ref) - ak.flatten(corrs_L1_jag)) < 1e-6)
# Apply the L1 corrections and save the result
test_ptL1_jag = test_pt_jag * corrs_L1_jag
print("L1 corrected pT values:", test_ptL1_jag, "\n")
assert ak.all(
np.abs(ak.flatten(test_pt_jag) - ak.flatten(test_ptL1_jag)) < 1e-6)
# Check the L2 corrections on a subset of jets
# Look up the parameters for the L2 corrections by hand and calculate the corrections
# [(1.37906,35.8534,-0.00829227,7.96644e-05,5.18988e-06),
# (1.38034,17.9841,-0.00729638,-0.000127141,5.70889e-05),
# (1.74466,18.6372,-0.0367036,0.00310864,-0.000277062),
# (1.4759,24.8882,-0.0155333,0.0020836,-0.000198039),
# (1.14606,36.4215,-0.00174801,-1.76393e-05,1.91863e-06),
# (0.999657,4.02981,1.06597,-0.619679,-0.0494)],
# [(1.54524,23.9023,-0.0162807,0.000665243,-4.66608e-06),
# (1.48431,8.68725,0.00642424,0.0252104,-0.0335696)]])
corrs_L2_jag_ref = ak.unflatten(
np.array([
1.37038741364,
1.37710384514,
1.65148641108,
1.46840446827,
1.1328319784,
1.0,
1.50762056349,
1.48719866989,
]),
counts,
)
corrector = FactorizedJetCorrector(
**{name: evaluator[name]
for name in jec_names[1:2]})
corrs_L2_jag = corrector.getCorrection(JetEta=test_eta_jag,
JetPt=test_pt_jag)
print("Reference L2 corrections:", corrs_L2_jag_ref.tolist())
print("Calculated L2 corrections:", corrs_L2_jag.tolist())
assert ak.all(
np.abs(ak.flatten(corrs_L2_jag_ref) - ak.flatten(corrs_L2_jag)) < 1e-6)
# Apply the L2 corrections and save the result
test_ptL1L2_jag = test_ptL1_jag * corrs_L2_jag
print("L1L2 corrected pT values:", test_ptL1L2_jag, "\n")
# Apply the L3 corrections and save the result
corrs_L3_jag = ak.full_like(test_pt_jag, 1.0)
test_ptL1L2L3_jag = test_ptL1L2_jag * corrs_L3_jag
print("L1L2L3 corrected pT values:", test_ptL1L2L3_jag, "\n")
# Check that the corrections can be chained together
corrs_L1L2L3_jag_ref = ak.unflatten(
np.array([
1.37038741364,
1.37710384514,
1.65148641108,
1.46840446827,
1.1328319784,
1.0,
1.50762056349,
1.48719866989,
]),
counts,
)
corrector = FactorizedJetCorrector(
**{name: evaluator[name]
for name in (jec_names[0:2] + jec_names[3:])})
corrs_L1L2L3_jag = corrector.getCorrection(JetEta=test_eta_jag,
Rho=test_Rho_jag,
JetPt=test_pt_jag,
JetA=test_A_jag)
print("Reference L1L2L3 corrections:", corrs_L1L2L3_jag_ref)
print("Calculated L1L2L3 corrections:", corrs_L1L2L3_jag)
assert ak.all(
np.abs(
ak.flatten(corrs_L1L2L3_jag_ref) -
ak.flatten(corrs_L1L2L3_jag)) < 1e-6)
# Apply the L1L2L3 corrections and save the result
test_ptL1L2L3chain_jag = test_pt_jag * corrs_L1L2L3_jag
print("Chained L1L2L3 corrected pT values:", test_ptL1L2L3chain_jag, "\n")
assert ak.all(
np.abs(
ak.flatten(test_ptL1L2L3_jag) -
ak.flatten(test_ptL1L2L3chain_jag)) < 1e-6)