def plot_MassPeaks(jet_name, eventWise, ax_arr=None):
if ax_arr is None:
fig, ax_arr = plt.subplots(3, 3)
fig.suptitle(jet_name)
MassPeaks.plot_correct_pairs(eventWise, [jet_name],
show=False,
plot_type='hist',
signal_background='both',
ax_array=ax_arr[0])
MassPeaks.plot_correct_pairs(eventWise, [jet_name],
show=False,
plot_type='hist',
signal_background='signal',
ax_array=ax_arr[1])
if len(ax_arr) > 2:
PlottingTools.discribe_jet(eventWise,
jet_name,
font_size=10,
ax=ax_arr[2, 1])
PlottingTools.hide_axis(ax_arr[2, 0])
PlottingTools.hide_axis(ax_arr[2, 2])
fig.subplots_adjust(top=0.93,
bottom=0.005,
left=0.060,
right=0.97,
hspace=0.35,
wspace=0.255)
fig.set_size_inches(12, 9)
def test_find_additional_decay_product():
# need TagCreators, Children, MCPID
contents = {}
# empty event
for name in ["TagCreators", "Children", "MCPID"]:
contents[name] = [[]]
wanted = [1]
expected = [[]]
# one particle, not tagged
contents["TagCreators"] += [[]]
contents["Children"] += [[[]]]
contents["MCPID"] += [[5]]
wanted += [0]
expected += [[]]
# one particle, tagged
contents["TagCreators"] += [[[]]]
contents["Children"] += [[[]]]
contents["MCPID"] += [[5]]
wanted += [7]
expected += [[[]]]
# two particles, one tagged
contents["TagCreators"] += [[[0]]]
contents["Children"] += [[[1], []]]
contents["MCPID"] += [[10, 5]]
wanted += [4]
expected += [[[]]]
# three particles, one hard creates two, one tagged one wanted
contents["TagCreators"] += [[[0]]]
contents["Children"] += [[[1, 2], [], []]]
contents["MCPID"] += [[10, -5, 7]]
wanted = [7]
expected += [[[2]]]
# four particles, one hard creates three, one tagged two wanted
contents["TagCreators"] += [[[0]]]
contents["Children"] += [[[1, 2, 3], [], []]]
contents["MCPID"] += [[10, -5, 7, 7]]
wanted = [7]
expected += [[[2, 3]]]
contents = {k: ak.from_iter(v) for k, v in contents.items()}
n_events = len(wanted)
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**contents)
for i in range(n_events):
eventWise.selected_event = i
additional = MassPeaks.find_additional_decay_product(
eventWise, wanted[i])
for add_idxs, exp_idxs in zip(additional, expected[i]):
assert set(add_idxs) == set(exp_idxs[i])
# doing it again should take the data from cache,
# so the columns are not needed
eventWise.remove("MCPID")
for i in range(n_events):
eventWise.selected_event = i
additional = MassPeaks.find_additional_decay_product(
eventWise, wanted[i])
for add_idxs, exp_idxs in zip(additional, expected[i]):
assert set(add_idxs) == set(exp_idxs[i])
def sorted_masses(eventWise,
jet_name,
mass_function='correct allocation',
jet_pt_cut=None,
max_tag_angle=None):
"""
Get mass predictions for this jet.
Parameters
----------
eventWise : EventWise
dataset containing the jets
jet_name : str
prefix of the jet's variables in the eventWise
mass_function : str or callable
function that predicts a particle masses from the events
(Default value = 'highest pt pair')
jet_pt_cut : float
required minimum jet PT for the jet to be selected
if None the value s taken from Constants.py
(Default = None)
max_tag_angle : float
The maximum deltaR betweeen a tag and its jet
(Default value = 0.8)
Returns
-------
masses : numpy array of floats
the sorted list of particle mass predictions
"""
if mass_function == 'highest pt pair':
all_masses, pairs, pair_masses = MassPeaks.all_PT_pairs(
eventWise, jet_name, jet_pt_cut, max_tag_angle=max_tag_angle)
idx = next(i for i, p in enumerate(pairs) if set(p) == {0, 1})
masses = pair_masses[idx]
elif mass_function == 'correct allocation':
# need to get the assigned tags
assigned_tags, descendants_mass = \
MassPeaks.descendants_masses(eventWise, [25])
tag_combinations_list = assigned_tags[25]
masses = MassPeaks.mass_tag_combinations(eventWise,
jet_name,
tag_combinations_list,
jet_pt_cut=jet_pt_cut)
masses = np.array(masses)
# zero masses should be dropped as unreconstructed
masses = masses[masses > 0.00001]
else:
masses = mass_function(eventWise, jet_name, jet_pt_cut)
masses = np.sort(masses)
return masses
def plot_results(eventWise, jet_name, pretag_jet_pt_cut, img_base):
"""
Parameters
----------
eventWise :
jet_name :
pretag_jet_pt_cut :
img_base :
Returns
-------
"""
tag_before_pt_cut = pretag_jet_pt_cut is None
jet_pt_cut = pretag_jet_pt_cut
if not tag_before_pt_cut:
if InputTools.yesNo_question("Apply pt cut to jets after tagging? "):
jet_pt_cut = InputTools.get_literal("Jet pt cut; ", float)
if InputTools.yesNo_question("Plot shape variables? "):
ShapeVariables.append_tagshapes(eventWise, batch_length=BATCH_LENGTH, jet_pt_cut=jet_pt_cut, tag_before_pt_cut=tag_before_pt_cut)
ShapeVariables.append_jetshapes(eventWise, jet_name, batch_length=BATCH_LENGTH, jet_pt_cut=jet_pt_cut, tag_before_pt_cut=tag_before_pt_cut)
if tag_before_pt_cut:
ShapeVariables.plot_shapevars(eventWise, jet_name)
else:
ShapeVariables.plot_shapevars(eventWise, jet_name, jet_pt_cut)
if img_base:
plt.tight_layout()
fig = plt.gcf()
fig.set_size_inches(10, 10)
plt.savefig(img_base + '_shape.png')
else:
plt.show()
# mass peaks
if not hasattr(eventWise, jet_name + "_Tags"):
raise Exception
#st()
MassPeaks.plot_PT_pairs(eventWise, jet_name, jet_pt_cut=jet_pt_cut, show=not img_base)
if img_base:
plt.tight_layout()
fig = plt.gcf()
fig.set_size_inches(10, 10)
plt.savefig(img_base + '_PTpairs.png')
else:
plt.show()
def test_smallest_angle_parings():
params = {}
jet_name = "Jet"
params['Jet_Label'] = []
params['Jet_Phi'] = []
params['Jet_PT'] = []
params['Jet_Rapidity'] = []
params['Jet_Tags'] = []
params['Jet_Energy'] = []
params['Jet_Child1'] = []
params['Jet_Parent'] = []
params = {key: [ak.from_iter(v)] for key, v in params.items()}
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**params)
eventWise.selected_event = 0
found = MassPeaks.smallest_angle_parings(eventWise, jet_name, [])
assert len(found) == 0
# try a regular event
params = {}
params['Jet_Label'] = [[0, 1, 2], [3, 4, 5], [7, 8, 9], [10, 11]]
params['Jet_Phi'] = [[0, 0, 0], [40, 0, 50], [40, 0, 0], [1, 3]]
params['Jet_PT'] = [[10, 10, 10], [40, 50, 50], [40, 0, 0], [1, 7]]
params['Jet_Rapidity'] = [[0, 0, 0], [100, 0, 50], [70, 0, 0], [0, 5]]
params['Jet_Child1'] = [[1, -1, -1], [-1, 5, -1], [-1, -1, 7], [-1, -1]]
params['Jet_Parent'] = [[-1, 0, 0], [1, -1, 1], [2, 2, -1], [-1, -1]]
params['Jet_Tags'] = [[0], [], [1], []]
params = {key: [ak.from_iter(v)] for key, v in params.items()}
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**params)
eventWise.selected_event = 0
found = MassPeaks.smallest_angle_parings(eventWise, jet_name,
[0, 1, 2])
assert len(found) == 1
assert set(found[0]) == {0, 2}
found = MassPeaks.smallest_angle_parings(eventWise, jet_name, [1, 2])
assert len(found) == 0
params['Jet_Tags'] = [ak.from_iter([[0], [5], [1], [7]])]
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**params)
eventWise.selected_event = 0
found = MassPeaks.smallest_angle_parings(eventWise, jet_name,
[0, 1, 2, 3])
assert len(found) == 2
found = [set(pair) for pair in found]
assert {0, 1} in found
assert {2, 3} in found
def test_all_doubleTagged_jets():
params = {}
jet_name = "Jet"
# event 0 is empty - no entries
params['Jet_Label'] = [[]]
params['Jet_Px'] = [[]]
params['Jet_Py'] = [[]]
params['Jet_Pz'] = [[]]
params['Jet_Energy'] = [[]]
params['Jet_Phi'] = [[]]
params['Jet_PT'] = [[]]
params['Jet_Rapidity'] = [[]]
params['Jet_Tags'] = [[]]
params['Jet_Child1'] = [[]]
params['Jet_Parent'] = [[]]
# event 1 has one double tagged jet blow the pt cut and one tagged jet - no entries
params['Jet_Label'] += [[[0, 1, 2], [3, 4, 5]]]
params['Jet_Px'] += [[[1, 1, 1], [2, 1, 1]]]
params['Jet_Py'] += [[[1, 1, 1], [2, 1, 1]]]
params['Jet_Pz'] += [[[1, 1, 1], [2, 1, 1]]]
params['Jet_Energy'] += [[[10, 10, 10], [11, 10, 10]]]
params['Jet_Phi'] += [[[0, 0, 0], [0, 0, 0]]]
params['Jet_PT'] += [[[1, 1, 1], [2, 1, 1]]]
params['Jet_Rapidity'] += [[[1, 1, 1], [1, 1, 1]]]
params['Jet_Tags'] += [[[1, 2], []]]
params['Jet_Child1'] += [[[1, -1, -1], [4, -1, -1]]]
params['Jet_Parent'] += [[[-1, 0, 0], [-1, 3, 3]]]
# event 2 has one single tagged and two double tagged jets - two entiries
params['Jet_Label'] += [[[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11]]]
params['Jet_Px'] += [[[1, 1, 1], [3, 1, 1], [-1, -1, -1], [-1, -1, -1]]]
params['Jet_Py'] += [[[1, 1, 1], [3, 1, 1], [-1, -1, -1], [-1, -1, -1]]]
params['Jet_Pz'] += [[[1, 1, 1], [3, 1, 1], [-1, -1, -1], [-1, -1, -1]]]
params['Jet_Energy'] += [[[10, 10, 10], [12, 12, 12], [5, 5, 5],
[10, 10, 10]]]
params['Jet_Phi'] += [[[0, 0, 0], [0, 0, 0], [3, 3, 3], [1, 3, 3]]]
params['Jet_PT'] += [[[4, 5, 4], [10, 10, 10], [11, 11, 11], [7, 7, 7]]]
params['Jet_Rapidity'] += [[[1, 1, 1], [1, 1, 1], [-1, -1, -1],
[10, 10, 10]]]
params['Jet_Tags'] += [[[1, 12], [1], [4, 5], [3]]]
params['Jet_Child1'] += [[[1, -1, -1], [4, -1, -1], [-1, 8, -1],
[11, -1, -1]]]
params['Jet_Parent'] += [[[-1, 0, 0], [-1, 3, 3], [7, -1, 7], [-1, 9, 9]]]
expected = sorted([np.sqrt(10**2 - 3), np.sqrt(5**2 - 3)])
prep_params = {
key: ak.from_iter([ak.from_iter(e) for e in v])
for key, v in params.items()
}
# check this results in the predicted masses
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**prep_params)
masses = MassPeaks.all_doubleTagged_jets(eventWise, jet_name, 2)
tst.assert_allclose(sorted(masses), expected)
def test_order_tagged_jets():
# need Tags, Label and a ranking_variable, by default PT
# try an empty event
params = {}
jet_name = "Jet"
params['Jet_Label'] = []
params['Jet_Parent'] = []
params['Jet_Tags'] = []
params['Jet_PT'] = []
params['Jet_Child1'] = []
params = {key: [ak.from_iter(v)] for key, v in params.items()}
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**params)
eventWise.selected_event = 0
filtered_event_idxs = []
found = MassPeaks.order_tagged_jets(eventWise, jet_name,
filtered_event_idxs)
assert len(found) == 0
# try a regular event
params = {}
params['Jet_Label'] = [[0, 1, 2], [3, 4, 5], [7, 8, 9]]
params['Jet_Tags'] = [[0], [1], []]
params['Jet_PT'] = [[60, 10, 40], [100, 40, 50], [40, 40, 40]]
params['Jet_Bobble'] = [[-5, 1, 1], [0, 0, -10], [12, 20, 40]]
params['Jet_Child1'] = [[1, -1, -1], [-1, 5, -1], [-1, -1, 7]]
params['Jet_Parent'] = [[-1, 0, 0], [1, -1, 1], [2, 2, -1]]
params = {key: [ak.from_iter(v)] for key, v in params.items()}
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**params)
eventWise.selected_event = 0
filtered_event_idxs = [0, 1]
found = MassPeaks.order_tagged_jets(eventWise, jet_name,
filtered_event_idxs)
tst.assert_allclose(found, [1, 0])
found = MassPeaks.order_tagged_jets(eventWise, jet_name,
filtered_event_idxs, "Bobble")
tst.assert_allclose(found, [0, 1])
def test_combined_jet_mass():
# need jet Px, Py, Pz, Energy, Label,
# try an empty event
params = {}
jet_name = "Jet"
params['Jet_Label'] = []
params['Jet_Parent'] = []
params['Jet_Px'] = []
params['Jet_Py'] = []
params['Jet_Pz'] = []
params['Jet_Energy'] = []
params = {key: [ak.from_iter(v)] for key, v in params.items()}
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**params)
eventWise.selected_event = 0
found = MassPeaks.combined_jet_mass(eventWise, jet_name, [])
assert found == 0
# try a regular event
params = {}
params['Jet_Label'] = [[0, 1, 2], [3, 4, 5], [7, 8, 9], [10]]
params['Jet_Parent'] = [[1, -1, 1], [1, -1, 1], [-1, 0, 0], [-1]]
params['Jet_Px'] = [[0, 0, 0], [40, 0, 50], [40, 0, 0], [1]]
params['Jet_Py'] = [[0, 0, 0], [100, 0, 50], [70, 0, 0], [0]]
params['Jet_Pz'] = [[0, 0, 0], [100, 0, 50], [1, 0, 0], [0]]
params['Jet_Energy'] = [[0, 0, 0], [10, 40, 10], [400, 40, 40], [10]]
params = {key: [ak.from_iter(v)] for key, v in params.items()}
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**params)
eventWise.selected_event = 0
for i, expected in enumerate(
[0, 40, np.sqrt(400**2 - 1 - 70**2 - 40**2),
np.sqrt(99)]):
found = MassPeaks.combined_jet_mass(eventWise, jet_name, [i])
tst.assert_allclose(found, expected)
found = MassPeaks.combined_jet_mass(eventWise, jet_name, [0, i])
tst.assert_allclose(found, expected)
def test_cluster_mass():
# need Px, Py, Pz, Energy
# try an empty event
params = {}
params['Px'] = []
params['Py'] = []
params['Pz'] = []
params['Energy'] = []
params = {key: [ak.from_iter(v)] for key, v in params.items()}
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**params)
eventWise.selected_event = 0
found = MassPeaks.cluster_mass(eventWise, [])
assert found == 0
# try a regular event
params = {}
params['Px'] = [0, 0, 40, 1]
params['Py'] = [0, 0, -70, 0]
params['Pz'] = [0, 0, 1, 0]
params['Energy'] = [0, 40, 400, 10]
params = {key: [ak.from_iter(v)] for key, v in params.items()}
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**params)
eventWise.selected_event = 0
for i, expected in enumerate(
[0, 40, np.sqrt(400**2 - 1 - 70**2 - 40**2),
np.sqrt(99)]):
found = MassPeaks.cluster_mass(eventWise, [i])
tst.assert_allclose(found, expected)
found = MassPeaks.cluster_mass(eventWise, [0, i])
tst.assert_allclose(found, expected)
found = MassPeaks.cluster_mass(eventWise, [0, 1, 2, 3])
expected = np.sqrt(450**2 - 41**2 - 70**2 - 1)
tst.assert_allclose(found, expected)
def get_normed_points(eventWise, jet_params, first_event, last_event):
eventWise.selected_event = None
allocated_tags, tag_masses = MassPeaks.descendants_masses(eventWise, [25])
allocated_tags, tag_masses = allocated_tags[25], tag_masses[25]
events = []
for event_n in range(first_event, last_event):
eventWise.selected_event = event_n
tag_idxs = list(eventWise.TagIndex)
tag_pairs = [[tag_idxs.index(i) for i in pair]
for pair in allocated_tags[event_n]]
tag_rapidity = eventWise.Rapidity[tag_idxs]
tag_phi = eventWise.Phi[tag_idxs]
jets = FormJets.Spectral(eventWise, dict_jet_params=jet_params)
jets.setup_internal()
distances, masses = get_jet_masses(jets, tag_rapidity, tag_phi,
tag_pairs)
masses -= list(tag_masses[event_n])
events.append((distances, masses))
return events
'Laplacien': 'symmetric',
'NumEigenvectors': np.inf,
'PhyDistance': 'angular',
'Sigma': 0.15,
'StoppingCondition': 'meandistance'
}
print()
print(f"Clustering {jet_name}")
print()
FormJets.cluster_multiapply(hepmc, jet_class, jet_params, jet_name, np.inf)
print()
print("Adding tags")
print()
TrueTag.semileptonic_top_tag_particles(hepmc)
for name in jet_names:
print(name)
TrueTag.add_tags(hepmc, name, 0.8, np.inf)
print()
print("Calculating mass peaks")
print()
hepmc.selected_event = None
MassPeaks.semileptonic_descendants_masses(hepmc)
for name in jet_names:
print(name)
hepmc.selected_event = None
MassPeaks.semileptonic_reconstructed_masses(hepmc, name)
print()
print("Done")
def test_all_jet_masses():
params = {}
jet_name = "Jet"
# event 0 is empty - should be 0
params['Jet_Label'] = [[]]
params['Jet_Px'] = [[]]
params['Jet_Py'] = [[]]
params['Jet_Pz'] = [[]]
params['Jet_Energy'] = [[]]
params['Jet_Phi'] = [[]]
params['Jet_PT'] = [[]]
params['Jet_Rapidity'] = [[]]
params['Jet_Tags'] = [[]]
params['Jet_Child1'] = [[]]
params['Jet_Parent'] = [[]]
expected_masses = [0.]
# event 1 has just 1 jet
params['Jet_Label'] += [[[0, 1, 2]]]
params['Jet_Px'] += [[[1, 1, 1]]]
params['Jet_Py'] += [[[1, 1, 1]]]
params['Jet_Pz'] += [[[1, 1, 1]]]
params['Jet_Energy'] += [[[10, 10, 10]]]
params['Jet_Phi'] += [[[0, 0, 0]]]
params['Jet_PT'] += [[[1, 1, 1]]]
params['Jet_Rapidity'] += [[[1, 1, 1]]]
params['Jet_Tags'] += [[[1]]]
params['Jet_Child1'] += [[[1, -1, -1]]]
params['Jet_Parent'] += [[[-1, 0, 0]]]
expected_masses += [np.sqrt(97)]
# event 2 has just 1 jet but it is untagged - should be 0
params['Jet_Label'] += [[[0, 1, 2]]]
params['Jet_Px'] += [[[1, 1, 1]]]
params['Jet_Py'] += [[[1, 1, 1]]]
params['Jet_Pz'] += [[[1, 1, 1]]]
params['Jet_Energy'] += [[[10, 10, 10]]]
params['Jet_Phi'] += [[[0, 0, 0]]]
params['Jet_PT'] += [[[1, 1, 1]]]
params['Jet_Rapidity'] += [[[1, 1, 1]]]
params['Jet_Tags'] += [[[]]]
params['Jet_Child1'] += [[[1, -1, -1]]]
params['Jet_Parent'] += [[[-1, 0, 0]]]
expected_masses += [0.]
# event 3 has one untagged and one tagged jet
params['Jet_Label'] += [[[0, 1, 2], [3, 4, 5]]]
params['Jet_Px'] += [[[1, 1, 1], [1, 1, 1]]]
params['Jet_Py'] += [[[1, 1, 1], [1, 1, 1]]]
params['Jet_Pz'] += [[[1, 1, 1], [1, 1, 1]]]
params['Jet_Energy'] += [[[10, 10, 10], [10, 10, 10]]]
params['Jet_Phi'] += [[[0, 0, 0], [0, 0, 0]]]
params['Jet_PT'] += [[[1, 1, 1], [1, 1, 1]]]
params['Jet_Rapidity'] += [[[1, 1, 1], [1, 1, 1]]]
params['Jet_Tags'] += [[[1], []]]
params['Jet_Child1'] += [[[1, -1, -1], [4, -1, -1]]]
params['Jet_Parent'] += [[[-1, 0, 0], [-1, 3, 3]]]
expected_masses += [np.sqrt(10**2 - 3)]
# event 4 has one untagged and two tagged jets
params['Jet_Label'] += [[[0, 1, 2], [3, 4, 5], [6, 7, 8]]]
params['Jet_Px'] += [[[1, 1, 1], [1, 1, 1], [-1, -1, -1]]]
params['Jet_Py'] += [[[1, 1, 1], [1, 1, 1], [-1, -1, -1]]]
params['Jet_Pz'] += [[[1, 1, 1], [1, 1, 1], [-1, -1, -1]]]
params['Jet_Energy'] += [[[10, 10, 10], [10, 10, 10], [5, 5, 5]]]
params['Jet_Phi'] += [[[0, 0, 0], [0, 0, 0], [3, 3, 3]]]
params['Jet_PT'] += [[[1, 1, 1], [1, 1, 1], [1, 1, 1]]]
params['Jet_Rapidity'] += [[[1, 1, 1], [1, 1, 1], [-1, -1, -1]]]
params['Jet_Tags'] += [[[1], [], [4]]]
params['Jet_Child1'] += [[[1, -1, -1], [4, -1, -1], [-1, 8, -1]]]
params['Jet_Parent'] += [[[-1, 0, 0], [-1, 3, 3], [7, -1, 7]]]
prep_params = {
key: ak.from_iter([ak.from_iter(e) for e in v])
for key, v in params.items()
}
expected_masses += [15.]
# check this results in the predicted masses
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**prep_params)
found = MassPeaks.all_jet_masses(eventWise, jet_name, 0)
tst.assert_allclose(found, expected_masses)
def test_all_smallest_angles():
params = {}
jet_name = "Jet"
# event 0 is empty - no contribution
params['Jet_Label'] = [[]]
params['Jet_Px'] = [[]]
params['Jet_Py'] = [[]]
params['Jet_Pz'] = [[]]
params['Jet_Energy'] = [[]]
params['Jet_Phi'] = [[]]
params['Jet_PT'] = [[]]
params['Jet_Rapidity'] = [[]]
params['Jet_Tags'] = [[]]
params['Jet_Child1'] = [[]]
params['Jet_Parent'] = [[]]
# event 1 has just 1 jet - no contribution
params['Jet_Label'] += [[[0, 1, 2]]]
params['Jet_Px'] += [[[1, 1, 1]]]
params['Jet_Py'] += [[[1, 1, 1]]]
params['Jet_Pz'] += [[[1, 1, 1]]]
params['Jet_Energy'] += [[[10, 10, 10]]]
params['Jet_Phi'] += [[[0, 0, 0]]]
params['Jet_PT'] += [[[1, 1, 1]]]
params['Jet_Rapidity'] += [[[1, 1, 1]]]
params['Jet_Tags'] += [[[1]]]
params['Jet_Child1'] += [[[1, -1, -1]]]
params['Jet_Parent'] += [[[-1, 0, 0]]]
# event 2 has one untagged and one tagged jet - no contribution
params['Jet_Label'] += [[[0, 1, 2], [3, 4, 5]]]
params['Jet_Px'] += [[[1, 1, 1], [1, 1, 1]]]
params['Jet_Py'] += [[[1, 1, 1], [1, 1, 1]]]
params['Jet_Pz'] += [[[1, 1, 1], [1, 1, 1]]]
params['Jet_Energy'] += [[[10, 10, 10], [10, 10, 10]]]
params['Jet_Phi'] += [[[0, 0, 0], [0, 0, 0]]]
params['Jet_PT'] += [[[1, 1, 1], [1, 1, 1]]]
params['Jet_Rapidity'] += [[[1, 1, 1], [1, 1, 1]]]
params['Jet_Tags'] += [[[1], []]]
params['Jet_Child1'] += [[[1, -1, -1], [4, -1, -1]]]
params['Jet_Parent'] += [[[-1, 0, 0], [-1, 3, 3]]]
prep_params = {
key: ak.from_iter([ak.from_iter(e) for e in v])
for key, v in params.items()
}
# check this results in no masses
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**prep_params)
found = MassPeaks.all_smallest_angles(eventWise, jet_name, 0)
assert len(found) == 0
# now start adding things that will contribute
# event 3 has one untagged and two tagged jets - two tagged jets should combine
params['Jet_Label'] += [[[0, 1, 2], [3, 4, 5], [6, 7, 8]]]
params['Jet_Px'] += [[[1, 1, 1], [1, 1, 1], [-1, -1, -1]]]
params['Jet_Py'] += [[[1, 1, 1], [1, 1, 1], [-1, -1, -1]]]
params['Jet_Pz'] += [[[1, 1, 1], [1, 1, 1], [-1, -1, -1]]]
params['Jet_Energy'] += [[[10, 10, 10], [10, 10, 10], [5, 5, 5]]]
params['Jet_Phi'] += [[[0, 0, 0], [0, 0, 0], [3, 3, 3]]]
params['Jet_PT'] += [[[1, 1, 1], [1, 1, 1], [1, 1, 1]]]
params['Jet_Rapidity'] += [[[1, 1, 1], [1, 1, 1], [-1, -1, -1]]]
params['Jet_Tags'] += [[[1], [], [4]]]
params['Jet_Child1'] += [[[1, -1, -1], [4, -1, -1], [-1, 8, -1]]]
params['Jet_Parent'] += [[[-1, 0, 0], [-1, 3, 3], [7, -1, 7]]]
expected_masses = [15.]
# event 4 has two untagged and three tagged jets - two closest tagged jets should combine
params['Jet_Label'] += [[[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11],
[12, 13, 14]]]
params['Jet_Px'] += [[[1, 1, 1], [1, 1, 1], [-1, -1, -1], [-1, -1, -1],
[-1, -1, -1]]]
params['Jet_Py'] += [[[1, 1, 1], [1, 1, 1], [-1, -1, -1], [-1, -1, -1],
[-1, -1, -1]]]
params['Jet_Pz'] += [[[1, 1, 1], [1, 1, 1], [-1, -1, -1], [-1, -1, -1],
[-1, -1, -1]]]
params['Jet_Energy'] += [[[10, 10, 10], [10, 10, 10], [5, 5, 5], [6, 6, 6],
[7, 7, 7]]]
params['Jet_Phi'] += [[[0, 0, 0], [0, 0, 0], [3, 3, 3], [3, 3, 3],
[3, 3, 3]]]
params['Jet_PT'] += [[[1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1],
[1, 1, 1]]]
params['Jet_Rapidity'] += [[[1, 1, 1], [1, 1, 1], [-1, -1, -1],
[-1, -1, -1], [-1, -1, -1]]]
params['Jet_Tags'] += [[[1], [], [4], [10, 11], []]]
params['Jet_Child1'] += [[[1, -1, -1], [4, -1, -1], [-1, 8, -1],
[10, -1, -1], [13, -1, -1]]]
params['Jet_Parent'] += [[[-1, 0, 0], [-1, 3, 3], [7, -1, 7], [-1, 9, 9],
[12, -1, 12]]]
expected_masses += [np.sqrt(11**2 - 3 * 4)]
# event 5 has one untagged and four tagged jets - all tagged jets will combine
params['Jet_Label'] += [[[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11],
[12, 13, 14]]]
params['Jet_Px'] += [[[1, 1, 1], [1, 1, 1], [-1, -1, -1], [-1, -1, -1],
[-1, -1, -1]]]
params['Jet_Py'] += [[[1, 1, 1], [1, 1, 1], [-1, -1, -1], [-1, -1, -1],
[-1, -1, -1]]]
params['Jet_Pz'] += [[[1, 1, 1], [1, 1, 1], [-1, -1, -1], [-1, -1, -1],
[-1, -1, -1]]]
params['Jet_Energy'] += [[[10, 10, 10], [10, 10, 10], [5, 5, 5], [6, 6, 6],
[7, 7, 7]]]
params['Jet_Phi'] += [[[0, 0, 0], [0, 0, 0], [3, 3, 3], [3, 3, 3],
[3, 3, 3]]]
params['Jet_PT'] += [[[1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1],
[1, 1, 1]]]
params['Jet_Rapidity'] += [[[1, 1, 1], [1, 1, 1], [-1, -1, -1],
[-1, -1, -1], [-1, -1, -1]]]
params['Jet_Tags'] += [[[1], [3], [4], [10, 11], []]]
params['Jet_Child1'] += [[[1, -1, -1], [4, -1, -1], [-1, 8, -1],
[10, -1, -1], [13, -1, -1]]]
params['Jet_Parent'] += [[[-1, 0, 0], [-1, 3, 3], [7, -1, 7], [-1, 9, 9],
[12, -1, 12]]]
expected_masses += [np.sqrt(20**2 - 3 * 4), np.sqrt(11**2 - 3 * 4)]
prep_params = {
key: ak.from_iter([ak.from_iter(e) for e in v])
for key, v in params.items()
}
# check this results in the predicted masses
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**prep_params)
found = MassPeaks.all_smallest_angles(eventWise, jet_name, 0)
tst.assert_allclose(sorted(found), sorted(expected_masses))
def test_all_PT_pairs():
params = {}
jet_name = "Jet"
# event 0 is empty - no pairs
params['Jet_Label'] = [[]]
params['Jet_Px'] = [[]]
params['Jet_Py'] = [[]]
params['Jet_Pz'] = [[]]
params['Jet_Energy'] = [[]]
params['Jet_Phi'] = [[]]
params['Jet_PT'] = [[]]
params['Jet_Rapidity'] = [[]]
params['Jet_Tags'] = [[]]
params['Jet_Child1'] = [[]]
params['Jet_Parent'] = [[]]
# 0 1, 0 2, 0 3, 1 2, 1 3, 2 3
expected_order = [(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)]
expected = [[] for _ in expected_order]
# event 1 has one untagged and one tagged jet - no pairs
params['Jet_Label'] += [[[0, 1, 2], [3, 4, 5]]]
params['Jet_Px'] += [[[1, 1, 1], [2, 1, 1]]]
params['Jet_Py'] += [[[1, 1, 1], [2, 1, 1]]]
params['Jet_Pz'] += [[[1, 1, 1], [2, 1, 1]]]
params['Jet_Energy'] += [[[10, 10, 10], [11, 10, 10]]]
params['Jet_Phi'] += [[[0, 0, 0], [0, 0, 0]]]
params['Jet_PT'] += [[[1, 1, 1], [2, 1, 1]]]
params['Jet_Rapidity'] += [[[1, 1, 1], [1, 1, 1]]]
params['Jet_Tags'] += [[[1], []]]
params['Jet_Child1'] += [[[1, -1, -1], [4, -1, -1]]]
params['Jet_Parent'] += [[[-1, 0, 0], [-1, 3, 3]]]
# event 2 has one untagged and two tagged jets - only the 0 1 pair
params['Jet_Label'] += [[[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11]]]
params['Jet_Px'] += [[[1, 1, 1], [3, 1, 1], [-1, -1, -1], [-1, -1, -1]]]
params['Jet_Py'] += [[[1, 1, 1], [3, 1, 1], [-1, -1, -1], [-1, -1, -1]]]
params['Jet_Pz'] += [[[1, 1, 1], [3, 1, 1], [-1, -1, -1], [-1, -1, -1]]]
params['Jet_Energy'] += [[[10, 10, 10], [12, 12, 12], [5, 5, 5],
[10, 10, 10]]]
params['Jet_Phi'] += [[[0, 0, 0], [0, 0, 0], [3, 3, 3], [1, 3, 3]]]
params['Jet_PT'] += [[[4, 5, 4], [10, 10, 10], [1, 1, 1], [7, 7, 7]]]
params['Jet_Rapidity'] += [[[1, 1, 1], [1, 1, 1], [-1, -1, -1],
[10, 10, 10]]]
params['Jet_Tags'] += [[[1], [], [4], []]]
params['Jet_Child1'] += [[[1, -1, -1], [4, -1, -1], [-1, 8, -1],
[11, -1, -1]]]
params['Jet_Parent'] += [[[-1, 0, 0], [-1, 3, 3], [7, -1, 7], [-1, 9, 9]]]
expected[expected_order.index((0, 1))].append(15)
# event 3 has five tagged jet - the 4 with highest PT will contribute to every combination
params['Jet_Label'] += [[[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11],
[12, 13, 14]]]
params['Jet_Px'] += [[[1, 1, 1], [1, 1, 1], [-1, -1, -1], [1, 1, 1],
[-1, -1, -1]]]
params['Jet_Py'] += [[[1, 1, 1], [2, 2, 2], [-1, -1, -1], [1, 1, 1],
[-1, -1, -1]]]
params['Jet_Pz'] += [[[1, 1, 1], [1, 1, 1], [-1, -1, -1], [1, 1, 1],
[3, 3, 3]]]
params['Jet_Energy'] += [[[10, 10, 10], [10, 10, 10], [5, 5, 5],
[10, 10, 10], [6, 6, 6]]]
params['Jet_Phi'] += [[[0, 0, 0], [0, 0, 0], [3, 3, 3], [0, 0, 0],
[3, 3, 3]]]
params['Jet_PT'] += [[[2.5, 2, 2], [2, 2, 2], [3, 3, 3], [1, 1, 4],
[2, 1.5, 1]]]
params['Jet_Rapidity'] += [[[1, 1, 1], [1, 1, 1], [-1, -1, -1], [1, 1, 1],
[-1, -1, -1]]]
params['Jet_Tags'] += [[[1], [2], [4], [5, 6], [10]]]
params['Jet_Child1'] += [[[1, -1, -1], [4, -1, -1], [-1, 8, -1],
[-1, 11, -1], [-1, 14, -1]]]
params['Jet_Parent'] += [[[-1, 0, 0], [-1, 3, 3], [7, -1, 7], [-1, 9, 9],
[12, -1, 12]]]
# PT order is 2, 0, 1, 4
expected[expected_order.index((0, 1))].append(15)
mass = np.sqrt(15**2 - 1)
expected[expected_order.index((0, 2))].append(mass)
mass = np.sqrt(11**2 - 4 * 3)
expected[expected_order.index((0, 3))].append(mass)
mass = np.sqrt(20**2 - 4 * 2 - 3**2)
expected[expected_order.index((1, 2))].append(mass)
mass = np.sqrt(16**2 - 4**2)
expected[expected_order.index((1, 3))].append(mass)
mass = np.sqrt(16**2 - 4**2 - 1)
expected[expected_order.index((2, 3))].append(mass)
prep_params = {
key: ak.from_iter([ak.from_iter(e) for e in v])
for key, v in params.items()
}
# check this results in the predicted masses
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**prep_params)
_, pairs, pair_masses = MassPeaks.all_PT_pairs(eventWise,
jet_name,
jet_pt_cut=0)
err_start = f"All expected = {list(zip(expected_order, expected))}\n"
err_start += f"All found = {list(zip(pairs, pair_masses))}\n"
for found_pair, found_masses in zip(pairs, pair_masses):
expected_mass = expected[expected_order.index(
tuple(sorted(found_pair)))]
err_msg = err_start + f"found pair = {found_pair}, found_masses = {found_masses}\n"
err_msg = err_start + f"expected masses = {expected_mass}\n"
tst.assert_allclose(found_masses, expected_mass, err_msg=err_msg)
from jet_tools import MassPeaks
import awkward as ak
import numpy as np
from matplotlib import pyplot as plt
from jet_tools import Components, MassPeaks, TrueTag
if __name__ == '__main__':
print("Getting mass gap tags")
ew = Components.EventWise.from_file(
"../megaIgnore/run_01/clustered.parquet")
jet_pt_cut = 0
all_masses4, chosen_pairs4, chosen_tags4 = MassPeaks.all_best_mass(
ew, "AntiKTp4Jet", jet_pt_cut, 0.4, require_all=False, dijet_mass=60)
#all_masses8, chosen_pairs8, chosen_tags8 = MassPeaks.all_best_mass(ew, "AntiKTp8Jet", jet_pt_cut, 0.8, require_all=False)
chosen_tags4 = [[set(t) for t in event] for event in chosen_tags4]
#chosen_tags8 = [[set(t) for t in event] for event in chosen_tags8]
#assert len(chosen_tags4) == len(chosen_tags8)
print("Getting correct found tags")
assigned_tags, decendant_masses = MassPeaks.descendants_masses(
ew, [25, 35])
n_events = len(chosen_tags4)
found_correct_tags4 = [[] for _ in range(n_events)]
#found_correct_tags8 = [[] for _ in range(n_events)]
for event_n, event_correct in enumerate(assigned_tags[25]):
ew.selected_event = event_n
tags = ew.AntiKTp4Jet_Tags
pt = ew.AntiKTp4Jet_PT
parent = ew.AntiKTp4Jet_Parent
def test_link_tags_to_hard():
# need TagCreators, TagIndex, Parents, Children, MCPID
contents = {}
# empty event
for name in ["TagCreators", "TagIndex", "Parents", "Children", "MCPID"]:
contents[name] = [[]]
hard = [[]]
expected_hard = [[]]
expected_created = [[]]
# one particle, not tagged
contents["TagIndex"] += [[]]
contents["TagCreators"] += [[]]
contents["Parents"] += [[[]]]
contents["Children"] += [[[]]]
contents["MCPID"] += [[5]]
hard += [[]]
expected_hard += [[]]
expected_created += [[]]
# one particle, tagged
contents["TagIndex"] += [[0]]
contents["TagCreators"] += [[[]]]
contents["Parents"] += [[[]]]
contents["Children"] += [[[]]]
contents["MCPID"] += [[5]]
hard += [[5]]
expected_hard += [[[0]]]
expected_created += [[[[]]]]
# two particles, one tagged
contents["TagIndex"] += [[1]]
contents["TagCreators"] += [[[0]]]
contents["Parents"] += [[[], [0]]]
contents["Children"] += [[[1], []]]
contents["MCPID"] += [[10, 5]]
hard += [[10, 1]]
expected_hard += [[[0], []]]
expected_created += [[[[1]], []]]
# two particles, related but not hard
contents["TagIndex"] += [[1]]
contents["TagCreators"] += [[[0]]]
contents["Parents"] += [[[], [0]]]
contents["Children"] += [[[1], []]]
contents["MCPID"] += [[1, 6]]
hard += [[10, 5]]
expected_hard += [[[], []]]
expected_created += [[[], []]]
# three particles, one hard creates two
contents["TagIndex"] += [[1, 2]]
contents["TagCreators"] += [[[0], [0]]]
contents["Parents"] += [[[], [0], [0]]]
contents["Children"] += [[[1, 2], [], []]]
contents["MCPID"] += [[10, -5, -5]]
hard += [[10, 6]]
expected_hard += [[[0], []]]
expected_created += [[[[1, 2]], []]]
# three particles, two hard creates one
contents["TagIndex"] += [[2]]
contents["TagCreators"] += [[[0, 1]]]
contents["Parents"] += [[[], [], [0, 1]]]
contents["Children"] += [[[2], [2], []]]
contents["MCPID"] += [[10, 5, -4]]
hard += [[10, 5]]
expected_hard += [[[0], [1]]]
expected_created += [[[[2]], [[2]]]]
# three particles, hard particle is up a chain
contents["TagIndex"] += [[2]]
contents["TagCreators"] += [[[1]]]
contents["Parents"] += [[[], [0], [1], [2]]]
contents["Children"] += [[[1], [2], [3], []]]
contents["MCPID"] += [[10, 10, 5, -5]]
hard += [[10]]
expected_hard += [[[0]]]
expected_created += [[[[2]]]]
contents = {k: ak.from_iter(v) for k, v in contents.items()}
n_events = len(hard)
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**contents)
for i in range(n_events):
eventWise.selected_event = i
found_hard, found_created = MassPeaks.link_tags_to_hard(
eventWise, hard[i])
for p, pid in enumerate(hard[i]):
for h, c in zip(found_hard[p], found_created[p]):
assert h in expected_hard[i][p],\
f"{i}; Expected {expected_hard[i][p]}, found {h}"
position = expected_hard[i][p].index(h)
assert set(c) == set(expected_created[i][p][position]),\
f"{i}; Expected {expected_created[i][p][position]}, found {c}"
def test_descendants_masses():
# need TagCreators, TagIndex, Parents, Children, MCPID,
# Px, Py, Px, Energy
# JetInput_SourceIdx
# also, Additional3
contents = {}
# empty event
hard = [10, 6]
# with jet inputs
expected_tags_ji = {pid: [] for pid in hard}
expected_mass_ji = {pid: [] for pid in hard}
# with additional particles
expected_tags_ad = {pid: [] for pid in hard}
expected_mass_ad = {pid: [] for pid in hard}
# with jet inputs and additional particles
expected_tags_jiad = {pid: [] for pid in hard}
expected_mass_jiad = {pid: [] for pid in hard}
for name in [
"TagCreators", "TagIndex", "Parents", "Children", "MCPID", "Px",
"Py", "Pz", "Energy", "Additional3", "JetInputs_SourceIdx"
]:
contents[name] = [[]]
for pid in hard:
expected_tags_ji[pid] += [[]]
expected_mass_ji[pid] += [[]]
expected_tags_ad[pid] += [[]]
expected_mass_ad[pid] += [[]]
expected_tags_jiad[pid] += [[]]
expected_mass_jiad[pid] += [[]]
# four particles, one hard creates two tags, one additional
contents["TagIndex"] += [[1, 2]]
contents["TagCreators"] += [[[0], [0]]]
contents["Parents"] += [[[], [0], [0], [0]]]
contents["Children"] += [[[1, 2, 3], [], [], []]]
contents["MCPID"] += [[10, -5, -5, 3]]
contents["Additional3"] += [[[3], [3]]]
# One tag is not in the jet inputs
contents["JetInputs_SourceIdx"] += [[1]]
# kinematics
contents["Px"] += [[0., 1., 1., 1.]]
contents["Py"] += [[0., 1., 1., 1.]]
contents["Pz"] += [[0., 1., 1., 1.]]
contents["Energy"] += [[0., 10., 10., 10.]]
# one hard undetected
expected_tags_ji[6] += [[]]
expected_mass_ji[6] += [[]]
expected_tags_ad[6] += [[]]
expected_mass_ad[6] += [[]]
expected_tags_jiad[6] += [[]]
expected_mass_jiad[6] += [[]]
# the tag will still be assocated with that hard pid, even if it is unseen
expected_tags_ji[10] += [[{1, 2}]]
expected_mass_ji[10] += [[np.sqrt(100 - 3)]]
expected_tags_ad[10] += [[{1, 2}]]
expected_mass_ad[10] += [[np.sqrt(30**2 - 3 * 3**2)]]
expected_tags_jiad[10] += [[{1, 2}]]
expected_mass_jiad[10] += [[np.sqrt(20**2 - 3 * 2**2)]]
# same again, but nothing is in the jet inputs
contents["TagIndex"] += [[1, 2]]
contents["TagCreators"] += [[[0], [0]]]
contents["Parents"] += [[[], [0], [0], [0]]]
contents["Children"] += [[[1, 2, 3], [], [], []]]
contents["MCPID"] += [[10, -5, -5, 3]]
contents["Additional3"] += [[[3], [3]]]
# nothing in the jet inputs
contents["JetInputs_SourceIdx"] += [[]]
# kinematics
contents["Px"] += [[0., 1., 1., 1.]]
contents["Py"] += [[0., 1., 1., 1.]]
contents["Pz"] += [[0., 1., 1., 1.]]
contents["Energy"] += [[0., 10., 10., 10.]]
# one hard undetected
expected_tags_ji[6] += [[]]
expected_mass_ji[6] += [[]]
expected_tags_ad[6] += [[]]
expected_mass_ad[6] += [[]]
expected_tags_jiad[6] += [[]]
expected_mass_jiad[6] += [[]]
# the tag will still be assocated with that hard pid, even if it is unseen
expected_tags_ji[10] += [[{1, 2}]]
expected_mass_ji[10] += [[0.]]
expected_tags_ad[10] += [[{1, 2}]]
# still get mass from add
expected_mass_ad[10] += [[np.sqrt(100 - 3)]]
expected_tags_jiad[10] += [[{1, 2}]]
expected_mass_jiad[10] += [[np.sqrt(100 - 3)]]
contents = {k: ak.from_iter(v) for k, v in contents.items()}
n_events = len(hard)
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**contents)
found_tags_ji, found_mass_ji = MassPeaks.descendants_masses(
eventWise, hard, additional_decay_pid=None, use_jetInputs=True)
for i in range(n_events):
for pid in hard:
for tags, mass in zip(found_tags_ji[pid][i],
found_mass_ji[pid][i]):
tags = set(tags)
assert tags in expected_tags_ji[pid][i]
pos = expected_tags_ji[pid][i].index(tags)
err_message = f"Jet inputs, no additionals\n" + \
f"event = {i}, pid = {pid}, tags = {tags}\n" + \
f"expected_mass = {expected_mass_ji[pid][i][pos]}\n" + \
f"found mass = {mass}\n"
tst.assert_allclose(expected_mass_ji[pid][i][pos],
mass,
err_msg=err_message)
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**contents)
found_tags_ad, found_mass_ad = MassPeaks.descendants_masses(
eventWise, hard, additional_decay_pid=3, use_jetInputs=False)
for i in range(n_events):
for pid in hard:
for tags, mass in zip(found_tags_ad[pid][i],
found_mass_ad[pid][i]):
tags = set(tags)
assert tags in expected_tags_ad[pid][i]
pos = expected_tags_ad[pid][i].index(tags)
err_message = f"no Jet inputs, with additionals\n" + \
f"event = {i}, pid = {pid}, tags = {tags}\n" + \
f"expected_mass = {expected_mass_ad[pid][i][pos]}\n" + \
f"found mass = {mass}\n"
tst.assert_allclose(expected_mass_ad[pid][i][pos],
mass,
err_msg=err_message)
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**contents)
found_tags_jiad, found_mass_jiad = MassPeaks.descendants_masses(
eventWise, hard, additional_decay_pid=3, use_jetInputs=True)
for i in range(n_events):
for pid in hard:
for tags, mass in zip(found_tags_jiad[pid][i],
found_mass_jiad[pid][i]):
tags = set(tags)
assert tags in expected_tags_jiad[pid][i]
pos = expected_tags_jiad[pid][i].index(tags)
err_message = f"Jet inputs, with additionals\n" + \
f"event = {i}, pid = {pid}, tags = {tags}\n" + \
f"expected_mass = {expected_mass_jiad[pid][i][pos]}\n" + \
f"found mass = {mass}\n"
tst.assert_allclose(expected_mass_jiad[pid][i][pos],
mass,
err_msg=err_message)
def test_inverse_condensed_indices():
assert MassPeaks.inverse_condensed_indices(0, 0) is None
assert set(MassPeaks.inverse_condensed_indices(0, 2)) == {0, 1}
assert set(MassPeaks.inverse_condensed_indices(3, 4)) == {1, 2}
