def test_add_inheritance():
params = {}
jet_name = "Jet"
# event 0
params['Jet_Label'] = [ak.from_iter([])]
params['Jet_Parent'] = [ak.from_iter([])]
params['Jet_Energy'] = [ak.from_iter([])]
params['Jet_Px'] = [ak.from_iter([])]
params['Jet_Py'] = [ak.from_iter([])]
params['Jet_Pz'] = [ak.from_iter([])]
params['Children'] = [ak.from_iter([])]
params['Parents'] = [ak.from_iter([])]
params['MCPID'] = [ak.from_iter([])]
params['PT'] = [ak.from_iter([])]
params['JetInputs_SourceIdx'] = [ak.from_iter([])]
params['TagIndex'] = [ak.from_iter([])]
# event 1
params['JetInputs_SourceIdx'] += [ak.from_iter(np.arange(11))]
params['Jet_Label'] += [ak.from_iter([[0, 101, 2], [102, 4, 5]])]
params['Jet_Parent'] += [ak.from_iter([[101, -1, 101], [-1, 102, 102]])]
params['Jet_Energy'] += [ak.from_iter([[30., 10., 20.], [70., 20., 10.]])]
params['Jet_Px'] += [ak.from_iter([[3., 0., 2.], [1., 2., -1.]])]
params['Jet_Py'] += [ak.from_iter([[3., 0., 2.], [2., 2., 1.]])]
params['Jet_Pz'] += [ak.from_iter([[3., 0., 2.], [0., 2., 2.]])]
# invarient_mass 873 100 388 4859, 388, 94
# shifted energy 30 10 20 70 sqrt(393) sqrt(103)
params['Children'] += [
ak.from_iter([[], [3], [], [5], [], [], [2, 7, 8, 9], [], [], [], []])
]
params['Parents'] += [
ak.from_iter([[], [], [6], [1], [], [3], [], [6], [6], [6], []])
]
params['PT'] += [ak.from_iter([3, 1, 2, 1, 2, 1, 3, 1, 2, 1, 2])]
params['MCPID'] += [ak.from_iter([4, -5, 5, 3, 2, 1, -5, -1, 7, 11, 12])]
params['TagIndex'] += [ak.from_iter([1, 6])]
# the positivity will be 0 1 0 1 0 1 0 0 0 0 0
# 0 0 1 0 0 0 0 0 0 0 0
# 0 0 0 0 0 0 0 0 0 0 0
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**params)
FormShower.append_b_idxs(eventWise)
TrueTag.add_inheritance(eventWise, jet_name)
# the first event is empty
eventWise.selected_event = 0
assert len(eventWise.Jet_Inheritance) == 0
assert len(eventWise.Jet_ITags) == 0
# the second event has two jets
eventWise.selected_event = 1
expected0 = [[0, 0, 0],
[np.sqrt(103) / (np.sqrt(393) + np.sqrt(103)), 0, 1]]
expected1 = [[0, 2 / 5, 1], [0, 0, 0]]
tst.assert_allclose(eventWise.Jet_Inheritance.tolist()[0], expected0)
tst.assert_allclose(eventWise.Jet_Inheritance.tolist()[1], expected1)
expected_tags = [[1], [6]]
def test_add_mass_share():
params = {}
jet_name = "Jet"
# event 0
params['Jet_Label'] = [ak.from_iter([])]
params['Energy'] = [ak.from_iter([])]
params['Px'] = [ak.from_iter([])]
params['Py'] = [ak.from_iter([])]
params['Pz'] = [ak.from_iter([])]
params['Children'] = [ak.from_iter([])]
params['Parents'] = [ak.from_iter([])]
params['MCPID'] = [ak.from_iter([])]
params['JetInputs_SourceIdx'] = [ak.from_iter([])]
params['TagIndex'] = [ak.from_iter([])]
# event 1
params['JetInputs_SourceIdx'] += [ak.from_iter(np.arange(6))]
params['Jet_Label'] += [ak.from_iter([[0, 101, 2], [102, 4, 5]])]
params['Energy'] += [
ak.from_iter([30., 10., 20., 70., 20., 10., 45., 56., 40., 25.])
]
params['Px'] += [ak.from_iter([3., 0., 2., 1., 2., -1., 0., 3., -1., 0.])]
params['Py'] += [ak.from_iter([3., 0., 2., 2., 2., 1., -1., -3., 0., -1.])]
params['Pz'] += [ak.from_iter([3., 0., 2., 0., 2., 2., -5., -2., 1., 0.])]
# invarient_mass 873 100 388 4859, 388, 94
# shifted energy 30 10 20 70 sqrt(393) sqrt(103)
# 0 1 2 3 4 5 6 7 8 9 10
params['Children'] += [
ak.from_iter([[], [3], [], [5], [], [], [2, 7, 8, 9], [], [], [], []])
]
params['Parents'] += [
ak.from_iter([[], [], [6], [1], [], [3], [], [6], [6], [6], []])
]
params['MCPID'] += [ak.from_iter([4, -5, 5, 3, 2, 1, -5, -1, 7, 11, 12])]
params['TagIndex'] += [ak.from_iter([1, 6])]
# the positivity will be 0 1 0 1 0 1 0 0 0 0 0
# 0 0 1 0 0 0 0 0 0 0 0
# in jet 1 are decendents 2 from tag 1
# in jet 2 are decendents 5 from tag 0
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**params)
FormShower.append_b_idxs(eventWise)
TrueTag.add_mass_share(eventWise, jet_name)
# the first event is empty
eventWise.selected_event = 0
assert len(eventWise.Jet_TagMass) == 0
assert len(eventWise.Jet_MTags) == 0
# the second event has two jets
eventWise.selected_event = 1
expected = [[0., np.sqrt(388.)], [np.sqrt(94.), 0.]]
expected_tags = [[6], [1]]
tst.assert_allclose(eventWise.Jet_TagMass.tolist(), expected)
tst.assert_allclose(eventWise.Jet_MTags.tolist(), expected_tags)
def test_append_b_idxs():
# 0 1 2 3 4 5 6 7 8 9 10
children = [[], [2, 3], [5], [6, 5], [], [], [7, 8, 9], [], [], [], []]
mcpid = [4, 5, 5, 3, 2, 1, -5, 1, 7, 11, 12]
expected = [2, 6]
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(Children=[ak.from_iter(children)],
MCPID=[ak.from_iter(mcpid)])
found = FormShower.append_b_idxs(eventWise, append=False)
tst.assert_allclose(sorted(found['BQuarkIdx'][0]), expected)
found = FormShower.append_b_idxs(eventWise, silent=False, append=True)
eventWise.selected_event = 0
tst.assert_allclose(sorted(eventWise.BQuarkIdx), expected)
found = FormShower.append_b_idxs(eventWise, silent=False, append=True)
assert found == True