def tags_to_quarks(eventWise, tag_idxs, quark_pdgids=[-5, 5]):
"""
Find the quark that is most strongly assocated with each tag particle.
Parameters
----------
eventWise : EventWise
dataset containing locations of particles and jets
tag_idx : iterable of ints
the idx of the tag particles as found in the EventWise
quark_pdgids : iterable of ints
list of mcpids considred to be quarks
(Default value = [-5, 5])
Returns
-------
quark_parents : numpy array of ints
the idx of the quark parents as found in the EventWise
"""
assert eventWise.selected_event is not None
if np.all([pid in quark_pdgids for pid in eventWise.MCPID[tag_idxs]]):
# check if the tags are already quarks
return tag_idxs
# fetch the angular variables for speed
rapidity = eventWise.Rapidity
phi = eventWise.Phi
pids = eventWise.MCPID
parents = eventWise.Parents
children = eventWise.Children
# fetch any quarks in the tag's parents
quark_parents = []
quark_distances = []
for tag_idx in tag_idxs:
this_qparents = []
parent_stack = [tag_idx]
while parent_stack:
idx = parent_stack.pop()
if pids[idx] in quark_pdgids:
this_qparents.append(idx)
else:
parent_stack += parents[idx].tolist()
this_qparents = list(set(this_qparents))
# if there are multiple parents, abandon any that have b-quark decendants
if len(this_qparents) > 1:
last_b = [
5 not in np.abs(pids[children[idx]]) for idx in this_qparents
]
if sum(last_b): # there must be at least 1 remaining
this_qparents = [
this_qparents[i] for i, l in enumerate(last_b) if l
]
quark_parents.append(this_qparents)
# if there is more than one quark on a tag calculate the deltaR to that quark
# if these is just one parent, call the deltaR 0
if len(this_qparents) == 1:
quark_distances.append(np.zeros(1))
elif len(this_qparents) > 1:
distances = np.sqrt((rapidity[this_qparents] -
rapidity[tag_idx])**2 +
Components.angular_distance(
phi[this_qparents], phi[tag_idx])**2)
quark_distances.append(distances.tolist())
else:
raise RuntimeError(
f"Why does this tag have no quark parents? event_n = {eventWise.selected_event}, tag_idx = {tag_idx}."
)
# go through the quark distances, assigning each quark to its closest tag
remaining = list(range(len(quark_parents)))
while remaining:
tag_num = np.argmin([np.min(quarks) for quarks in quark_distances])
remaining.remove(tag_num) # we found this one
quark_distances[tag_num] = [np.inf]
quark_num = np.argmin(quark_distances[tag_num])
quark_idx = quark_parents[tag_num].pop(quark_num)
quark_parents[tag_num] = quark_idx
# chekc if this quark_idx is found elsewhere
for num in remaining:
try:
quark_num2 = quark_parents[num].index(quark_idx)
except ValueError:
# it wasn't in this list, that's fine
pass
else:
quark_parents[num].pop(quark_num2)
quark_distances[num].pop(quark_num2)
# check if there is now only one parent left at num
if len(quark_distances[num]) == 1:
quark_distances[num][0] = 0. # set this distance to 0
assert np.all([isinstance(q, int) for q in quark_parents])
return quark_parents
"""Processing events in version 1 of the paper"""
#hepmc_dir = "/scratch/hadh1g17/MG5_aMC_v2_6_7/ttsemileptonic/Events/run_02/tag_1_pythia8_events.hepmc"
#hepmc_dir = "../megaIgnore/ttsemileptonic_pt4.hepmc"
hepmc_dir = "/home/henry/Programs/MG5_aMC_v2_9_2/ttsemileptonic3/Events/run_03/tag_1_pythia8_events.hepmc"
from jet_tools import ReadHepmc, Components, FormJets, TrueTag, MassPeaks, FormJetInputs
import numpy as np
hepmc = ReadHepmc.Hepmc(hepmc_dir, 0, np.inf)
hepmc.file_name = "ttsemileptonic_pt7.awkd"
#hepmc.dir_name = "/home/hadh1g17/jets/megaIgnore"
hepmc.dir_name = "../megaIgnore"
hepmc.write()
print()
print("Adding components")
print()
Components.add_all(hepmc)
print()
print("Creating jet inputs")
print()
FormJetInputs.create_jetInputs(hepmc)
jet_names = []
jet_name = "AntiKTp8Jet"
jet_names.append(jet_name)
jet_class = FormJets.Traditional
jet_params = {
'DeltaR': 0.8,
'ExpofPTFormat': 'min',
'ExpofPTMultiplier': 0,
'PhyDistance': 'angular'
}
print()
def test_remove_particles_from_jets():
contents = {}
# check it does the right thing with a non-empty event
contents["Event_n"] = [0]
contents["DogJet_Parent"] = [[[], [-1], [3, -1, 3]]]
contents["DogJet_Child1"] = [[[], [-1], [-1, 2, -1]]]
contents["DogJet_Child2"] = [[[], [-1], [-1, 4, -1]]]
contents["DogJet_Rapidity"] = [[[], [0.5], [0.1, 0.3, 0.4]]]
contents["DogJet_Label"] = [[[], [1], [2, 3, 4]]]
contents["DogJet_Tag"] = [[[], [10, 20], [34, 50]]]
# check it manages an empty event
contents["Event_n"] += [1]
contents["DogJet_Parent"] += [[]]
contents["DogJet_Child1"] += [[]]
contents["DogJet_Child2"] += [[]]
contents["DogJet_Rapidity"] += [[]]
contents["DogJet_Label"] += [[]]
contents["DogJet_Tag"] += [[]]
# what we want to remove
remove = [1, 2]
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**contents)
GhostParticles.remove_particles_from_jets(eventWise, remove)
eventWise.selected_event = 1
assert "GhostDogJet_Rapidity" in eventWise.columns
# the first event shoudl still be empty
for name in eventWise.columns:
value = getattr(eventWise, name)
if name == "Event_n":
assert value == 1
else:
assert len(value) == 0
# the second event has a predictable structure
eventWise.selected_event = 0
assert eventWise.Event_n == 0
# all relations will have the same format
relations = [
eventWise.DogJet_Parent, eventWise.DogJet_Child1,
eventWise.DogJet_Child2
]
for values in relations:
assert len(values[0]) == 0
assert len(values[1]) == 0
tst.assert_allclose(values[2], [-1])
# only one value of rapidity/Label should be left
rapidity = eventWise.DogJet_Rapidity
assert len(rapidity[0]) == 0
assert len(rapidity[1]) == 0
tst.assert_allclose(rapidity[2], [0.4])
inputidx = eventWise.DogJet_Label
assert len(inputidx[0]) == 0
assert len(inputidx[1]) == 0
tst.assert_allclose(inputidx[2], [4])
# the tags should be unchanged
tags = eventWise.DogJet_Tag
assert len(tags[0]) == 0
tst.assert_allclose(tags[1], [10, 20])
tst.assert_allclose(tags[2], [34, 50])
# test the right ghosts have been added
ghost_inputs = eventWise.GhostDogJet_Label
assert len(ghost_inputs[0]) == 0
tst.assert_allclose(ghost_inputs[1], [1])
tst.assert_allclose(ghost_inputs[2], [2])
ghost_rapidity = eventWise.GhostDogJet_Rapidity
assert len(ghost_rapidity[0]) == 0
tst.assert_allclose(ghost_rapidity[1], [0.5])
tst.assert_allclose(ghost_rapidity[2], [0.1])
def tree_motion(start, root, steps_between):
"""
Parameters
----------
start :
param root:
steps_between :
root :
Returns
-------
"""
if root.is_leaf:
location = [[[start[0]], [start[1]]]] * (steps_between + 1)
size = [[root.size]] * (steps_between + 1)
colour = [[root.colour]] * (steps_between + 1)
return location, size, colour
else:
# find out what is beflow this point
next_left = [root.left.rap, root.left.phi]
below_left, below_left_size, below_left_colour = tree_motion(
next_left, root.left, steps_between)
next_right = [root.right.rap, root.right.phi]
below_right, below_right_size, below_right_colour = tree_motion(
next_right, root.right, steps_between)
# if either were leaves pad to the right depth
if len(below_left) < len(below_right):
pad_height = int(len(below_right) - len(below_left))
pad = pad_height * [below_left[-1]]
below_left += pad
pad_colour = pad_height * [below_left_colour[-1]]
below_left_colour += pad_colour
pad_size = pad_height * [below_left_size[-1]]
below_left_size += pad_size
elif len(below_right) < len(below_left):
pad_height = int(len(below_left) - len(below_right))
pad = pad_height * [below_right[-1]]
below_right += pad
pad_colour = pad_height * [below_right_colour[-1]]
below_right_colour += pad_colour
pad_size = pad_height * [below_right_size[-1]]
below_right_size += pad_size
assert len(below_left) == len(below_right)
levels = [[r_rap + l_rap, r_phi + l_phi]
for (r_rap, r_phi), (l_rap,
l_phi) in zip(below_left, below_right)]
sizes = [
l_size + r_size
for l_size, r_size in zip(below_left_size, below_right_size)
]
colours = [
l_col + r_col
for l_col, r_col in zip(below_left_colour, below_right_colour)
]
# now make the this level
rap_left = np.linspace(next_left[0], start[0], steps_between)
rap_right = np.linspace(next_right[0], start[0], steps_between)
# this coordinate is cyclic
# so this next bit is a shuffle to get it to link by the shortest route
distance = next_left[1] - start[1]
if distance > np.pi:
adjusted_next_left = -next_left[1]
elif distance < -np.pi:
adjusted_next_left = next_left[1] + 2 * np.pi
else:
adjusted_next_left = next_left[1]
phi_left = np.linspace(adjusted_next_left, start[1], steps_between)
phi_left = Components.confine_angle(phi_left)
distance = next_right[1] - start[1]
if distance > np.pi:
adjusted_next_right = -next_right[1]
elif distance < -np.pi:
adjusted_next_right = next_right[1] + 2 * np.pi
else:
adjusted_next_right = next_right[1]
phi_right = np.linspace(adjusted_next_right, start[1], steps_between)
phi_right = Components.confine_angle(phi_right)
this_level = [[[eleft, eright], [pleft, pright]]
for eleft, eright, pleft, pright in zip(
rap_left, rap_right, phi_left, phi_right)]
this_level += [[[root.rap], [root.phi]]]
# add all the motion together
levels += this_level
# pick size
l_size = root.left.size
r_size = root.right.size
this_size = [[l_size, r_size]] * steps_between + [[root.size]]
sizes += this_size
l_colour = root.left.colour
r_colour = root.right.colour
this_colour = [[l_colour, r_colour]] * steps_between + [[
(*root.colour[:3], 0.5)
]]
colours += this_colour
return levels, sizes, colours
def test_add_detectable_fourvector():
# will need the TagIndex, Px, Py, Pz, Energy, JetInputs_SourceIdx, Child, Parent
# empty event should do nothing
params = {}
params["Children"] = [ak.from_iter([])]
params["Parents"] = [ak.from_iter([])]
params["Energy"] = [ak.from_iter([])]
params["Px"] = [ak.from_iter([])]
params["Py"] = [ak.from_iter([])]
params["Pz"] = [ak.from_iter([])]
params["TagIndex"] = [ak.from_iter([])]
params["JetInputs_SourceIdx"] = [ak.from_iter([])]
# idx 0 1 2 3 4 5 6 7 8 9 10
# dtectable x v v v v
children = [[], [2, 3], [5], [6, 5], [], [], [7, 8, 9], [], [], [], []]
parents = [[], [], [1], [1], [], [2, 3], [3], [6], [6], [6], []]
energy = [0, 11, 9, 11, 10, 23, 1, 2, 0, 5, 0]
px = [4, 2, 30, 0, 1, 0, 1, 0, 0, -4, 0]
py = [0, 0, 2, 6, 0, 4, 1, 0, 0, 0, 0]
pz = [0, 1, 0, 0, 0, 2, 0, 1, 0, 1, 0]
bquark = [2, 3, 4, 10]
source = [0, 4, 5, 8, 9]
expected_roots = [{2, 3}, {4}]
expected_leaves = [{5, 8, 9}, {4}]
expected_px = [-4, 1]
expected_py = [4, 0]
expected_pz = [3, 0]
expected_energy = [28, 10]
params["Children"] += [ak.from_iter(children)]
params["Parents"] += [ak.from_iter(parents)]
params["Energy"] += [ak.from_iter(energy)]
params["Px"] += [ak.from_iter(px)]
params["Py"] += [ak.from_iter(py)]
params["Pz"] += [ak.from_iter(pz)]
params["TagIndex"] += [ak.from_iter(bquark)]
params["JetInputs_SourceIdx"] += [ak.from_iter(source)]
params = {k: ak.from_iter(v) for k, v in params.items()}
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**params)
columns_before = [x for x in eventWise.columns]
TrueTag.add_detectable_fourvector(eventWise)
columns_after = [x for x in eventWise.columns]
# the first event should contain nothing
eventWise.selected_event = 0
try:
roots = eventWise.DetectableTag_Roots
except Exception as e:
message = f"Error; {e}\ncolumns_before={columns_before}\ncolumns_after={columns_after}"
raise Exception(message)
assert len(roots) == 0
energy = eventWise.DetectableTag_Energy
assert len(energy) == 0
# the second event should contain the predicted values
eventWise.selected_event = 1
roots = [set(r) for r in eventWise.DetectableTag_Roots]
assert len(roots) == len(expected_roots)
# this will thro an error if anythign is missing
order = [roots.index(r) for r in expected_roots]
for e, f in enumerate(order):
assert expected_leaves[e] == set(eventWise.DetectableTag_Leaves[f])
tst.assert_allclose(expected_energy,
eventWise.DetectableTag_Energy[order])
tst.assert_allclose(expected_px, eventWise.DetectableTag_Px[order])
tst.assert_allclose(expected_py, eventWise.DetectableTag_Py[order])
tst.assert_allclose(expected_pz, eventWise.DetectableTag_Pz[order])
def test_make_n_working_fragments():
with TempTestDir("tst") as dir_name:
# calling this one a directory that dosn't contain any eventWise objects
# should raise a file not found erroj
with pytest.raises(FileNotFoundError):
ParallelFormJets.make_n_working_fragments(dir_name, 3, "squibble")
# trying to split something that isn't an eventwise shoudl raise a NotADirectoryError
wrong_path = os.path.join(dir_name, "wrong.ods")
open(wrong_path, 'w').close() # equivalent of touch
with pytest.raises(NotADirectoryError):
ParallelFormJets.make_n_working_fragments(wrong_path, 3, "flob")
os.remove(wrong_path)
# now make a real eventWise to play with
file_name = "test.parquet"
ew = Components.EventWise(os.path.join(dir_name, file_name))
ew_path = os.path.join(dir_name, file_name)
# try with 12 events
n_events = 12
params = {}
params['Event_n'] = ak.from_iter(np.arange(n_events))
params['JetInputs_Energy'] = ak.from_iter(np.arange(n_events))
unfinished_jet = 'DogJet'
n_unfinished = 6
params[unfinished_jet + '_Energy'] = ak.from_iter(
np.random.rand(n_events - n_unfinished))
params[unfinished_jet + '_Food'] = ak.from_iter([
np.random.rand(np.random.randint(5))
for _ in range(n_events - n_unfinished)
])
finished_jet = 'CatJet'
params[finished_jet + '_Energy'] = ak.from_iter([[
ak.from_iter(np.random.rand(np.random.randint(5)))
for _ in range(np.random.randint(5))
] for _ in range(n_events)])
ew.append(**params)
# making fragments of the finished jet should result in no change
paths = ParallelFormJets.make_n_working_fragments(
ew_path, 3, finished_jet)
assert isinstance(paths, bool)
assert paths
# check nothing else has been made in the dir
assert len(os.listdir(dir_name)) == 1
# now split the unfinished jet
paths = ParallelFormJets.make_n_working_fragments(
ew_path, 3, unfinished_jet)
assert len(paths) == 3
ew.selected_event = None
# there is no garentee that the split events will hold the same order
expected_indices = list(range(n_events - n_unfinished, n_events))
for path in paths:
ew_part = Components.EventWise.from_file(path)
indices_here = ew_part.JetInputs_Energy.tolist()
for i in indices_here:
expected_indices.remove(i)
flat_here = ak.flatten(ew.CatJet_Energy[indices_here])
flat_part = ak.flatten(ew_part.CatJet_Energy)
try:
flat_here = ak.flatten(flat_here)
flat_part = ak.flatten(flat_part)
except ValueError: # already flat
pass
tst.assert_allclose(flat_here.tolist(), flat_part.tolist())
assert not expected_indices, "Didn't find all the expected indices"
# if we ask for the same split again it would not do anything
paths2 = ParallelFormJets.make_n_working_fragments(
ew_path, 3, unfinished_jet)
assert set(paths2) == set(paths)
paths2 = ParallelFormJets.make_n_working_fragments(
os.path.split(paths[0])[0], 3, unfinished_jet)
assert set(paths2) == set(paths)
# if we ask for a diferent number of paths it should recluster and then split
paths3 = ParallelFormJets.make_n_working_fragments(
ew_path, 2, unfinished_jet)
assert len(paths3) == 2
# there is no garentee that the split events will hold the same order
expected_indices = list(range(n_events - n_unfinished, n_events))
for path in paths3:
ew_part = Components.EventWise.from_file(path)
indices_here = ew_part.JetInputs_Energy.tolist()
for i in indices_here:
expected_indices.remove(i)
flat_here = ak.flatten(ew.CatJet_Energy[indices_here])
flat_part = ak.flatten(ew_part.CatJet_Energy)
try:
flat_here = ak.flatten(flat_here)
flat_part = ak.flatten(flat_part)
except ValueError: # already flat
pass
tst.assert_allclose(flat_here.tolist(), flat_part.tolist())
assert not expected_indices, "Didn't find all the expected indices"
# remove any existing directories
[
shutil.rmtree(os.path.join(dir_name, name))
for name in os.listdir(dir_name) if '.' not in name
]
# try fragmenting and then joining
ew.fragment("JetInputs_Energy", n_fragments=3)
fragment_dir = next(
os.path.join(dir_name, name) for name in os.listdir(dir_name)
if "fragment" in name)
Components.EventWise.combine(fragment_dir,
ew.file_name.split('.', 1)[0])
ParallelFormJets.make_n_working_fragments(ew_path, 4, finished_jet)
# check it has reconstructed the original eventWise
found = os.listdir(fragment_dir)
assert len(found) == 1
new_ew = Components.EventWise.from_file(
os.path.join(fragment_dir, found[0]))
order = np.argsort(new_ew.JetInputs_Energy)
dog_order = order[order < n_events - n_unfinished]
tst.assert_allclose(ew.JetInputs_Energy,
new_ew.JetInputs_Energy[order])
tst.assert_allclose(ew.DogJet_Energy, new_ew.DogJet_Energy[dog_order])
for i, j in enumerate(dog_order):
tst.assert_allclose(ew.DogJet_Food[i], new_ew.DogJet_Food[j])
for i, j in enumerate(order):
for evt, new_evt in zip(ew.CatJet_Energy[i],
new_ew.CatJet_Energy[j]):
tst.assert_allclose(evt, new_evt)
else:
root_present = False
print(
"It is often useful to calculate some values from the momentum and energy."
)
print("In the Components module there are a number of methods for this; ")
print(" > Components.add_thetas(eventwise)")
print(" > Components.add_pt(eventwise)")
print(" > Components.add_phi(eventwise)")
print(" > Components.add_rapidity(eventwise)")
print(" > Components.add_pseudorapidity(eventwise)")
print(" > Components.add_mass(eventwise)")
print("These have been combined in Components.add_all(eventWise)")
print(" > Components.add_all(eventWise, inc_mass=True)")
Components.add_all(eventWise, inc_mass=True)
print(
"All these variables now exist as attributes of the eventWise, have a look; "
)
interactive()
if InputTools.yesNo_question("Do you want to cluster the jets with Anti-kt? "):
print("~~ Example of adding JetInputs ~~")
if root_present:
print(
"As the root file has been added we can form jets based on the particles"
)
print(
"based on the particles that actually have been picked up by the detector."
)
print(" > filter1 = FormJets.filter_obs")
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_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_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)
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 create_jet_internals(energy, px, py, pz,
rapidity=None, phi=None, labels=None):
""" For a requested set of particle 4 momentum, rapidity, and phi
create the list structure that would be needed to place those particles
into FormJets.Clustering object
Parameters
----------
energy : array of floats
List of energy values of the the desired particles
px : array of floats
List of px values of the the desired particles
py : array of floats
List of py values of the the desired particles
pz : array of floats
List of pz values of the the desired particles
rapidity : array of floats (optional)
List of rapidity values of the the desired particles
phi : array of floats (optional)
List of phi values of the the desired particles
labels : array of ints (optional)
List of labels to give the new particles,
if not given negative numbers below -2 will be assigned
Returns
-------
labels : numpy array of ints
the Label given to the new particles
new_ints : list of list of ints
the list of list of ints for the new particles,
in the same format as in the input to FormJets.Agglomerative
new_floats : list of list of floats
the list of list of floats for the new particles,
in the same format as in the input to FormJets.Agglomerative
"""
# jets.int_columns =
#["Pseudojet_Label",
# "Pseudojet_Parent", "Pseudojet_Child1", "Pseudojet_Child2",
# "Pseudojet_Rank"]
# jets.float_columns =
#["Pseudojet_PT",
# "Pseudojet_Rapidity",
# "Pseudojet_Phi",
# "Pseudojet_Energy",
# "Pseudojet_Px",
# "Pseudojet_Py",
# "Pseudojet_Pz",
# "Pseudojet_JoinDistance",
# "Pseudojet_Size"]
# check the jet interals have the shape you expected
num_ints = 5
# give the particles a set of labels that are negative
num_additions = px.shape[0]
if labels is None:
labels = np.arange(-2, -num_additions - 2, -1)
new_ints = [[i] + [-1]*(num_ints-1) for i in labels]
if phi is None:
phi, pt = Components.pxpy_to_phipt(px, py)
else:
_, pt = Components.pxpy_to_phipt(px, py)
if rapidity is None:
rapidity = Components.ptpze_to_rapidity(pt, pz, energy)
# work out the float values
new_floats = [pt, rapidity, phi, energy, px, py, pz,
np.zeros(num_additions), # Join distance
np.ones(num_additions)] # Size
new_floats = list(map(list, zip(*new_floats)))
return labels, new_ints, new_floats
def fix_phis(eventWise, jet_name):
eventWise.selected_event = None
px = getattr(eventWise, jet_name + "_Px")
py = getattr(eventWise, jet_name + "_Py")
phis, _ = Components.pxpy_to_phipt(px, py)
return {jet_name + "_Phi": phis}
def test_combine():
with TempTestDir("tst") as dir_name:
empty_path = os.path.join(dir_name, "empty.parquet")
# an empty eventWise
empty = Components.EventWise(empty_path)
# putting in an empty eventwise should have no effect
returns = AddPileup.combine(empty, empty, [])
assert len(returns.columns) == 0
assert len(returns.hyperparameter_columns) == 0
# if the base ew has hyper params they should be held
hypers = {"Dog": "woof"}
hypers_path = os.path.join(dir_name, "hypers.parquet")
# an hypers eventWise
hyper_ew = Components.EventWise(hypers_path)
hyper_ew.append_hyperparameters(**hypers)
returns = AddPileup.combine(hyper_ew, empty, [])
assert len(returns.columns) == 0
assert len(returns.hyperparameter_columns) == 1
assert returns.Dog == "woof"
# alternativly, hyperparameters in the pileup don't get kept
returns = AddPileup.combine(empty, hyper_ew, [])
assert len(returns.columns) == 0
assert len(returns.hyperparameter_columns) == 0
# make a simple base with just one event
content = {}
content["Parents"] = [[[], [0]]]
content["Children"] = [[[1], []]]
content["Particle_barcode"] = [[1, 2]]
content["Start_vertex_barcode"] = [[0, -1]]
content["End_vertex_barcode"] = [[-1, 0]]
content["Vertex_barcode"] = [[-1]]
content["Px"] = [[0., 3.]]
content["Z"] = [[0.]]
content["MCPID"] = [[1, 2]]
content["Event_n"] = [0]
ak_content = {name: ak.from_iter(val) for name, val in content.items()}
simple_path = os.path.join(dir_name, "simple.parquet")
simple = Components.EventWise(simple_path)
simple.append(**ak_content)
# putting it in with an empty pileup shoudl result in no change
returns = AddPileup.combine(simple, empty, [[]])
assert len(returns.Px) == 1
for col in returns.columns:
if col == "Is_pileup":
continue
found = getattr(returns, col)
assert np.all(found == ak.from_iter(content[col]))
assert np.all(~returns.Is_pileup[0])
# simplist addition
returns = AddPileup.combine(simple, simple, [[0]])
returns.selected_event = 0
parents = ak.from_iter([[], [0], [], [2]])
assert np.all(returns.Parents == parents)
children = ak.from_iter([[1], [], [3], []])
assert np.all(returns.Children == children)
particle_barcode = ak.from_iter([1, 2, 3, 4])
assert np.all(returns.Particle_barcode == particle_barcode)
start_vertex_barcode = ak.from_iter([0, -1, 0, -2])
assert np.all(returns.Start_vertex_barcode == start_vertex_barcode)
end_vertex_barcode = ak.from_iter([-1, 0, -2, 0])
assert np.all(returns.End_vertex_barcode == end_vertex_barcode)
vertex_barcode = ak.from_iter([-1, -2])
assert np.all(returns.Vertex_barcode == vertex_barcode)
px = ak.from_iter([0., 3., 0., 3.])
assert np.all(returns.Px == px)
MCPID = ak.from_iter([1, 2, 1, 2])
assert np.all(returns.MCPID == MCPID)
assert returns.Event_n == 0
assert np.all(
returns.Is_pileup == np.array([False, False, True, True]))
# now make a more complex, 2 event pileup data
content["Parents"] += [[[], [0], [1]]]
content["Children"] += [[[1], [2], []]]
content["Particle_barcode"] += [[2, 1, 3]]
content["Start_vertex_barcode"] += [[0, -1, -2]]
content["End_vertex_barcode"] += [[-1, -2, 0]]
content["Vertex_barcode"] += [[-1, -2]]
content["Px"] += [[2., 5., 6.]]
content["Z"] += [[2., 5.]]
content["MCPID"] += [[9, 8, 10]]
content["Event_n"] += [1]
ak_content = {name: ak.from_iter(val) for name, val in content.items()}
complex_path = os.path.join(dir_name, "complex.parquet")
comple = Components.EventWise(complex_path)
comple.append(**ak_content)
# check we can add it in a non trivial pattern
returns = AddPileup.combine(comple, comple, [[], [1, 0]])
# the first event should be unchanged
returns.selected_event = 0
for col in returns.columns:
found = getattr(returns, col)
if col == "Is_pileup":
continue
if col == "Event_n":
assert found == content[col][0]
else:
assert np.all(found == ak.from_iter(content[col][0]))
assert np.all(~returns.Is_pileup)
# the second event shoudl have additions
returns.selected_event = 1
parents = ak.from_iter([[], [0], [1], [], [3], [4], [], [6]])
assert np.all(returns.Parents == parents)
children = ak.from_iter([[1], [2], [], [4], [5], [], [7], []])
assert np.all(returns.Children == children)
particle_barcode = ak.from_iter([2, 1, 3, 5, 4, 6, 7, 8])
assert np.all(returns.Particle_barcode == particle_barcode)
start_vertex_barcode = ak.from_iter([0, -1, -2, 0, -3, -4, 0, -5])
assert np.all(returns.Start_vertex_barcode == start_vertex_barcode)
end_vertex_barcode = ak.from_iter([-1, -2, 0, -3, -4, 0, -5, 0])
assert np.all(returns.End_vertex_barcode == end_vertex_barcode)
vertex_barcode = ak.from_iter([-1, -2, -3, -4, -5])
assert np.all(returns.Vertex_barcode == vertex_barcode)
px = ak.from_iter([2., 5., 6., 2., 5., 6., 0., 3])
assert np.all(returns.Px == px)
MCPID = ak.from_iter([9, 8, 10, 9, 8, 10, 1, 2])
assert np.all(returns.MCPID == MCPID)
assert returns.Event_n == 1
assert np.all(returns.Is_pileup == np.array(
[False, False, False, True, True, True, True, True]))
def test_filter_jets():
# will need
# Jet_Parent, Jet_Child1, Jet_PT
min_pt = 0.1
min_ntracks = 2
params = {}
# an empty event should return nothing
params['Jet_Parent'] = [ak.from_iter([])]
params['Jet_Child1'] = [ak.from_iter([])]
params['Jet_PT'] = [ak.from_iter([])]
params['MCPID'] = [ak.from_iter([])]
params['Generated_mass'] = [ak.from_iter([])]
# an event with nothing that passes cuts
params['Jet_Parent'] += [ak.from_iter([[-1], [-1, 1, 1, 2, 2]])]
params['Jet_Child1'] += [ak.from_iter([[-1], [1, 2, -1, -1, -1]])]
params['Jet_PT'] += [ak.from_iter([[
50.,
], [0.2, 0.1, 0., 0., .1]])]
params['MCPID'] += [ak.from_iter([25])]
params['Generated_mass'] += [ak.from_iter([40.])]
# an event with somthing that passes cuts
params['Jet_Parent'] += [ak.from_iter([[-1], [-1, 1, 1]])]
params['Jet_Child1'] += [ak.from_iter([[-1], [1, -1, -1]])]
params['Jet_PT'] += [ak.from_iter([[
50.,
], [21., 0.1, 0.]])]
params['MCPID'] += [ak.from_iter([25])]
params['Generated_mass'] += [ak.from_iter([40.])]
# an event to make use of variable cuts
params['Jet_Parent'] += [ak.from_iter([[-1, 0], [-1, 1, 1]])]
params['Jet_Child1'] += [ak.from_iter([[-1, -1], [1, -1, -1]])]
params['Jet_PT'] += [ak.from_iter([[50., 20.], [16., 0.1, 0.]])]
params['MCPID'] += [ak.from_iter([25])]
params['Generated_mass'] += [ak.from_iter([40.])]
# an event to make use of variable cuts
params['Jet_Parent'] += [ak.from_iter([[-1, 0], [-1, 1, 1]])]
params['Jet_Child1'] += [ak.from_iter([[-1, -1], [1, -1, -1]])]
params['Jet_PT'] += [ak.from_iter([[50., 20.], [14., 0.1, 0.]])]
params['MCPID'] += [ak.from_iter([25])]
params['Generated_mass'] += [ak.from_iter([40.])]
# an event to make use of variable cuts
params['Jet_Parent'] += [ak.from_iter([[-1, 0], [-1, 1, 1]])]
params['Jet_Child1'] += [ak.from_iter([[-1, -1], [1, -1, -1]])]
params['Jet_PT'] += [ak.from_iter([[17., 20.], [16., 0.1, 0.]])]
params['MCPID'] += [ak.from_iter([25])]
params['Generated_mass'] += [ak.from_iter([40.])]
params = {key: ak.from_iter(val) for key, val in params.items()}
with TempTestDir("tst") as dir_name:
# this will raise a value error if given an empty eventWise
file_name = "test.parquet"
ew = Components.EventWise(os.path.join(dir_name, file_name))
ew.append(**params)
# using defaults
jet_idxs = FormJets.filter_jets(ew, "Jet")
assert len(jet_idxs[0]) == 0
assert len(jet_idxs[1]) == 0
assert len(jet_idxs[2]) == 1
assert len(jet_idxs[3]) == 2
assert len(jet_idxs[4]) == 1
assert len(jet_idxs[5]) == 0
assert 1 in jet_idxs[2]
assert 0 in jet_idxs[3]
assert 1 in jet_idxs[3]
assert 0 in jet_idxs[4]
# using selected values
jet_idxs = FormJets.filter_jets(ew, "Jet", min_pt, min_ntracks)
assert len(jet_idxs[0]) == 0
assert len(jet_idxs[1]) == 1
assert 1 in jet_idxs[1]
assert len(jet_idxs[2]) == 1
assert 1 in jet_idxs[2]
def test_generate_pool():
# mock _worker, this is tested above
with unittest.mock.patch('jet_tools.ParallelFormJets._worker',
new=fake_worker):
with TempTestDir("tst") as temp_dir:
ew = Components.EventWise(os.path.join(temp_dir, "file.parquet"))
ew.append(JetInputs_Energy=np.arange(100))
eventWise_path = os.path.join(temp_dir, ew.file_name)
# get rid of a continue file if there is one or we will get stuck
try:
os.remove('continue')
except FileNotFoundError:
pass
end_time = time.time() + 100
jet_class = FormJets.Spectral
jet_params = {"Bark": 3}
n_cores = psutil.cpu_count()
function = lambda x: x
found = ParallelFormJets.generate_pool(
eventWise_path,
function,
function_args=[jet_class, "Bali", jet_params],
function_kwargs=dict(batch_size=500),
leave_one_free=True,
end_time=end_time)
assert found, "One of the processes crashed"
# now we expect to find a subdirectory containing n_cores - 1 eventWise
subdir = next(
os.path.join(temp_dir, name) for name in os.listdir(temp_dir)
if ".parquet" not in name)
paths = [
os.path.join(subdir, name) for name in os.listdir(subdir)
if Components.EventWise.pottential_file(name)
]
assert len(paths) == n_cores - 1
# check the ake worker has been run on all of them
for path in paths:
ew = Components.EventWise.from_file(path)
tst.assert_allclose(ew.Catto, [2, 3])
assert ew.Run_condition == end_time
assert ew.Jet_class == jet_class.__name__
assert ew.Cluster_parameters["Bark"] == 3
assert ew.Batch_size == 500
# tidy up
shutil.rmtree(os.path.split(path)[0])
# it should work fine with a continue file too
open('continue', 'w').close()
found = ParallelFormJets.generate_pool(
eventWise_path,
function,
function_args=[jet_class, "Bali", jet_params],
function_kwargs=dict(batch_size=500),
leave_one_free=False,
end_time=None)
assert found, "One of the processes crashed"
# now we expect to find a subdirectory containing n_cores eventWise
subdir = next(
os.path.join(temp_dir, name) for name in os.listdir(temp_dir)
if ".parquet" not in name)
paths = [
os.path.join(subdir, name) for name in os.listdir(subdir)
if Components.EventWise.pottential_file(name)
]
assert len(paths) == n_cores, \
f"paths = {len(paths)}, cores = {n_cores}"
# check the ake worker has been run on all of them
for path in paths:
ew = Components.EventWise.from_file(path)
tst.assert_allclose(ew.Catto, [2, 3])
assert ew.Run_condition == 'continue'
assert ew.Jet_class == jet_class.__name__
assert ew.Cluster_parameters["Bark"] == 3
assert ew.Batch_size == 500
def test_create_JetInputs():
with TempTestDir("create_JetInputs") as dir_name:
name = "test.parquet"
columns = [
name for name in FormJets.Clustering.float_columns
if "Distance" not in name and "Size" not in name
]
columns_unchanged = [c for c in columns
] # because the ew will change the columns list
floats = SimpleClusterSamples.two_oposite['floats']
contents = {
name: ak.from_iter([x])
for name, x in zip(columns, floats.T)
}
def return_all(_, current_idx):
return current_idx
def return_second(_, current_idx):
return current_idx[1:2]
ew = Components.EventWise(os.path.join(dir_name, name),
columns=columns,
contents=contents)
FormJetInputs.create_jetInputs(ew, filter_functions=[return_all])
for name in columns_unchanged:
ji_name = "JetInputs_" + name
idx = columns.index(name)
assert hasattr(ew, ji_name)
tst.assert_allclose(getattr(ew, ji_name).tolist(), [floats.T[idx]],
err_msg=f"In two_oposite {name} not matching")
ew.remove_prefix("JetInputs")
FormJetInputs.create_jetInputs(ew, filter_functions=[return_second])
for name in columns_unchanged:
ji_name = "JetInputs_" + name
idx = columns.index(name)
assert hasattr(ew, ji_name)
tst.assert_allclose(
getattr(ew, ji_name).tolist(), [floats.T[idx][1:2]])
# test batching
columns = [
name for name in FormJets.Clustering.float_columns
if "Distance" not in name
]
contents = {
name: ak.from_iter([x, x, x])
for name, x in zip(columns, floats.T)
}
ew = Components.EventWise(os.path.join(dir_name, name),
columns=columns,
contents=contents)
FormJetInputs.create_jetInputs(ew,
filter_functions=[return_all],
batch_length=0)
for name in columns_unchanged:
ji_name = "JetInputs_" + name
assert len(getattr(ew, ji_name)) == 0
FormJetInputs.create_jetInputs(ew,
filter_functions=[return_all],
batch_length=1)
for name in columns_unchanged:
ji_name = "JetInputs_" + name
assert hasattr(ew, ji_name)
idx = columns.index(name)
tst.assert_allclose(getattr(ew, ji_name).tolist(), [floats.T[idx]])
FormJetInputs.create_jetInputs(ew,
filter_functions=[return_all],
batch_length=1)
for name in columns_unchanged:
ji_name = "JetInputs_" + name
assert hasattr(ew, ji_name)
idx = columns.index(name)
tst.assert_allclose(
getattr(ew, ji_name).tolist(), [floats.T[idx], floats.T[idx]])
hepmc_save_dir = os.path.join(os.path.split(pileup_name)[0], higgs_weight)
print(f"Will be saving in {hepmc_save_dir}")
print("Reading hepmc")
ew = ReadHepmc.Hepmc(hepmc_name)
ew.selected_event = None
num_events = len(ew.X)
print(f"Contains {num_events} events")
ew.dir_name = hepmc_save_dir
i = 0
while True:
ew.file_name = f"{i:02}_{file_num:02}_signal.parquet"
try:
# will throw an error if the fiel exists
open(ew.path_name, 'x').close()
# otherwise, works like touch
break
except FileExistsError:
i += 1
ew.write()
print("Adding components")
Components.add_all(ew)
#print(f"Taking pileup from {pileup_name}")
#pileup = Components.EventWise.from_file(pileup_name)
#average_count = 50
#print(f"Adding an average of {average_count} pileup events")
#noised = AddPileup.add_pileup(ew, pileup, average_count)
#noised.write()
#print("done")
def create_jetInputs(eventWise,
filter_functions=None,
batch_length=np.inf,
silent=False):
"""
Add to the eventWise a set of particles prefixed by JetInputs
which pass all criteria required to be used in jet clustering.
Parameters
----------
eventWise : EventWise
data structure
filter_functions : list of callabels
callabels with the same signature as filter_pt_eta
that should reduce down the list of list of particle indices
to those that are sutable as jet inputs
(Default value = [filter_obs, filter_neutrinos, filter_pt_eta, filter_pileup])
batch_length : int
how many particles to checks
(Default value = np.inf)
silent : bool (optional)
don't print progress?
(Default; False)
"""
if filter_functions is None or filter_functions == "remove pileup":
filter_functions = [
filter_ends, filter_neutrinos, filter_pt_eta,
RemovePileup.filter_pileup
]
elif filter_functions == "ignore pileup":
filter_functions = [filter_ends, filter_neutrinos, filter_pt_eta]
#filter_functions = [filter_ends, filter_pt_eta]
# decide on run range
eventWise.selected_event = None
n_events = len(eventWise.Energy)
start_point = len(getattr(eventWise, "JetInputs_Energy", []))
if start_point >= n_events:
print("Finished")
return True
end_point = min(n_events, start_point + batch_length)
if not silent:
print(f" Will stop at {end_point/n_events:.1%}")
# sort out olumn names
sources = ["PT", "Rapidity", "Phi", "Energy", "Px", "Py", "Pz"]
for s in sources:
if not hasattr(eventWise, s):
Components.add_all(eventWise)
break
columns = ["JetInputs_" + c for c in sources]
columns.append("JetInputs_SourceIdx")
# the source column gives indices in the origin
# construct the observable filter in advance
contents = {
"JetInputs_SourceIdx":
list(getattr(eventWise, "JetInputs_SourceIdx", []))
}
for name in columns:
contents[name] = list(getattr(eventWise, name, []))
mask = []
for event_n in range(start_point, end_point):
if not silent and event_n % 100 == 0:
print(f"{event_n/n_events:.1%}", end='\r', flush=True)
eventWise.selected_event = event_n
idx_selection = np.arange(len(eventWise.PT))
for filter_func in filter_functions:
idx_selection = filter_func(eventWise, idx_selection)
contents["JetInputs_SourceIdx"].append(idx_selection)
mask_here = np.full_like(ak.to_numpy(eventWise.PT), False, dtype=bool)
mask_here[idx_selection] = True
mask.append(mask_here)
mask = ak.from_iter(mask)
eventWise.selected_event = None
for name, source_name in zip(columns, sources):
contents[name] += list(
getattr(eventWise, source_name)[start_point:end_point][mask])
contents = {k: ak.from_iter(v) for k, v in contents.items()}
eventWise.append(**contents)
if not silent:
print("Completed batch")
return end_point == n_events
def test_sorted_masses():
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
# only intrested in 0 1
expected = []
# 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'] += [[[8, 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'] += [[[2, 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.append(13)
# 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, 10, 10],
[13, -1, 13]]]
# PT order is 2, 0, 1, 4
expected.append(15)
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 = JetQuality.sorted_masses(eventWise,
jet_name,
mass_function='highest pt pair',
jet_pt_cut=0)
tst.assert_allclose(masses, sorted(expected))
masses = JetQuality.sorted_masses(eventWise,
jet_name,
mass_function='highest pt pair',
jet_pt_cut=2.1)
tst.assert_allclose(masses, sorted(expected[1:]))
def python_shape(energies, pxs, pys, pzs, thrust_axis=None):
"""
Parameters
----------
energies :
pxs :
pys :
pzs :
Returns
-------
"""
shapes = {}
# https://home.fnal.gov/~mrenna/lutp0613man2/node234.html
# precalculate some quantities
momentums = np.vstack((pxs, pys, pzs)).T
birrs2 = np.sum(momentums**2, axis=1)
# This way round gets Fortran like results
#sum_birr = np.sqrt(np.sum(birrs2))
#sum_Tbirr = np.sqrt(np.sum(momentums[:, :2]**2))
sum_birr = np.sum(np.sqrt(birrs2))
sum_Tbirr = np.sum(np.sqrt(np.sum(momentums[:, :2]**2, axis=1)))
# Thrust
if thrust_axis is None:
def to_minimise(theta_phi):
"""
Parameters
----------
theta_phi :
Returns
-------
"""
sin_theta = np.sin(theta_phi[0])
thrust_axis = [sin_theta*np.cos(theta_phi[1]),
sin_theta*np.sin(theta_phi[1]),
np.cos(theta_phi[0])]
return -np.sum(np.abs(np.dot(momentums, thrust_axis)))
theta_phi_bounds = ((0, np.pi), (-np.pi, np.pi))
best_theta_phi = minimzer(to_minimise, bounds=theta_phi_bounds,
objective_name="Thrust")
sin_theta = np.sin(best_theta_phi[0])
thrust_axis = [sin_theta*np.cos(best_theta_phi[1]),
sin_theta*np.sin(best_theta_phi[1]),
np.cos(best_theta_phi[0])]
else:
phi, pt = Components.pxpy_to_phipt(thrust_axis[0], thrust_axis[1])
theta = Components.ptpz_to_theta(pt, thrust_axis[2])
best_theta_phi = (theta, phi)
shapes['thrustVector[1]'] = thrust_axis[0]
shapes['thrustVector[2]'] = thrust_axis[1]
shapes['thrustVector[3]'] = thrust_axis[2]
momentum_dot_thrust = np.dot(momentums, thrust_axis)
shapes['Thrust'] = np.sum(np.abs(momentum_dot_thrust))/sum_birr
# transverse Thrust
def to_minimise_transverse(phi):
"""
Parameters
----------
phi :
Returns
-------
"""
transverse_thrust_axis = [np.cos(phi), np.sin(phi)]
return -np.sum(np.abs(np.dot(momentums[:, :2], transverse_thrust_axis)))
phi_bounds = (-np.pi, np.pi)
best_phi = minimzer(to_minimise_transverse, bounds=phi_bounds,
objective_name="TransThrust")
transverse_thrust_axis = [np.cos(best_phi), np.sin(best_phi)]
momentum_dot_Tthrust = np.dot(momentums[:, :2], transverse_thrust_axis)
shapes['Transversethrust'] = np.sum(np.abs(momentum_dot_Tthrust))/sum_Tbirr
# the major thrust has an exist in the plane perpendicular to the thrust
if best_theta_phi[0] in (np.pi, 0):
# along the z axis
perp1 = np.array([1, 0, 0])
perp2 = np.array([0, 1, 0])
else:
perp1 = np.cross(np.array([0, 0, 1]), thrust_axis)
perp1 /= np.sqrt(np.sum(perp1**2))
perp2 = np.cross(perp1, thrust_axis)
perp2 /= np.sqrt(np.sum(perp2**2))
if not (np.isclose(np.dot(perp1, thrust_axis), 0) and \
np.isclose(np.dot(perp2, thrust_axis), 0) and \
np.isclose(np.dot(perp1, perp2), 0)):
print("?")
#st()
def to_minimise_major(alpha):
"""
Parameters
----------
alpha :
Returns
-------
"""
major_thrust_axis = np.cos(alpha)*perp1 + np.sin(alpha)*perp2
return -np.sum(np.abs(momentums*major_thrust_axis))
best_alpha = minimzer(to_minimise_major, bounds=(-np.pi, np.pi),
objective_name="Major")
major_thrust_axis = np.cos(best_alpha)*perp1 + np.sin(best_alpha)*perp2
shapes['Major[1]'] = major_thrust_axis[0]
shapes['Major[2]'] = major_thrust_axis[1]
shapes['Major[3]'] = major_thrust_axis[2]
momentum_dot_major = np.dot(momentums, major_thrust_axis)
shapes['Major'] = np.sum(np.abs(momentum_dot_major))/sum_birr
minor_direction = np.cross(major_thrust_axis, thrust_axis)
minor_direction /= np.sqrt(np.sum(minor_direction**2))
shapes['Minor[1]'] = minor_direction[0]
shapes['Minor[2]'] = minor_direction[1]
shapes['Minor[3]'] = minor_direction[2]
sum_momentum_dot_minor = max(np.sum(np.abs(np.dot(momentums, minor_direction))),
np.sum(np.abs(np.dot(momentums, -minor_direction))))
shapes['Minor'] = sum_momentum_dot_minor/sum_birr
shapes['Oblateness'] = shapes['Major'] - shapes['Minor']
# sphericity
dimension = 3
per_mom_tensor = np.ones((dimension, dimension, len(momentums)))
for i in range(dimension):
per_mom_tensor[:, i, :] *= momentums.T
per_mom_tensor[i, :, :] *= momentums.T
mom_tensor = np.sum(per_mom_tensor, axis=2)
eigenvalues, _ = scipy.linalg.eig(mom_tensor/sum_birr**2)
eigenvalues = np.real(np.sort(eigenvalues)/np.sum(eigenvalues))
shapes['Sphericity'] = 1.5*(eigenvalues[0] + eigenvalues[1])
shapes['Alpanarity'] = 1.5*eigenvalues[0]
shapes['Planarity'] = eigenvalues[1] - eigenvalues[0]
lin_mom_tensor = np.sum(per_mom_tensor/np.sqrt(birrs2), axis=2)
eigenvalues, _ = scipy.linalg.eig(lin_mom_tensor/sum_birr)
eigenvalues = np.real(np.sort(eigenvalues)/np.sum(eigenvalues))
shapes['Cparameter'] = 3*(eigenvalues[2]*eigenvalues[1] +
eigenvalues[2]*eigenvalues[0] +
eigenvalues[1]*eigenvalues[0])
shapes['Dparameter'] = 27*np.product(eigenvalues)
# spherocity
def to_minimise_spherocity(phi):
"""
Parameters
----------
phi :
Returns
-------
"""
spherocity_axis = [np.cos(phi), np.sin(phi)]
return -np.abs(np.sum(np.cross(momentums[:, :2], spherocity_axis)))
phi_bounds = (-np.pi, np.pi)
best_phi = minimzer(to_minimise_spherocity, bounds=phi_bounds, objective_name="Spherocity")
spherocity_axis = [np.cos(best_phi), np.sin(best_phi)]
momentum_cross_sphro = np.cross(momentums[:, :2], spherocity_axis)
shapes['Spherocity'] = 0.25*np.pi**2*(np.sum(momentum_cross_sphro)/sum_Tbirr)**2
# Acoplanarity
def to_minimise_aco(theta_phi):
"""
Parameters
----------
theta_phi :
Returns
-------
"""
sin_theta = np.sin(theta_phi[0])
acoplanarity_axis = [sin_theta*np.cos(theta_phi[1]),
sin_theta*np.sin(theta_phi[1]),
np.cos(theta_phi[0])]
return np.sum(np.abs(momentums*acoplanarity_axis))
theta_phi_bounds = ((0, np.pi), (-np.pi, np.pi))
best_theta_phi = minimzer(to_minimise_aco, bounds=theta_phi_bounds, objective_name="Acoplanarity")
sin_theta = np.sin(best_theta_phi[0])
acoplanarity_axis = [sin_theta*np.cos(best_theta_phi[1]),
sin_theta*np.sin(best_theta_phi[1]),
np.cos(best_theta_phi[0])]
momentum_dot_aco = np.dot(momentums, acoplanarity_axis)
shapes['Acoplanarity'] = 4*np.sum(np.abs(momentum_dot_aco))/sum_birr
# jet masses
upper = np.dot(momentums[:, :2], transverse_thrust_axis) > 0
upper_mass2 = np.sum(energies[upper]**2) - np.sum(momentums[upper]**2)
lower_mass2 = np.sum(energies[~upper]**2) - np.sum(momentums[~upper]**2)
shapes['Lightjetmass2'], shapes['Heavyjetmass2'] = sorted((upper_mass2, lower_mass2))
shapes['Differencejetmass2'] = abs(upper_mass2 - lower_mass2)
return shapes
def plot_rapidity_phi_offset(jet_name,
eventWise,
rapidity_in,
phi_in,
tag_rapidity_in,
tag_phi_in,
ax_arr=None):
if ax_arr is None:
fig, ax_arr = plt.subplots(1, 3)
fig.suptitle(jet_name)
colour = ak.flatten(eventWise.DetectableTag_Mass)
alpha = 0.2
x_range = (-3.5, 3.5)
y_range = (0, np.pi + 0.1)
# zeroth axis has the full jet kinematics
ax = ax_arr[0]
ax.set_title("Entire jets")
rapidity_distance = ak.flatten(
(ak.from_iter(rapidity_in) - eventWise.DetectableTag_Rapidity))
phi_distance = ak.flatten(
ak.from_iter(phi_in) - eventWise.DetectableTag_Phi)
phi_distance = Components.raw_to_angular_distance(phi_distance)
# print the mean distance
ax.scatter([0], [0], s=5, c='k', marker='o')
mean_distance = np.nanmean(np.sqrt(rapidity_distance**2 + phi_distance**2))
# put a dot at the origin
ax.text(0,
y_range[1] * 0.75,
f"mean distance={mean_distance:.2f}",
horizontalalignment='center')
points = ax.scatter(rapidity_distance, phi_distance, c=colour, alpha=alpha)
ax.set_xlabel("Jet rapidity offset")
ax.set_ylabel("Jet $|\\phi|$ offset")
ax.set_ylim(y_range)
ax.set_xlim(x_range)
# add a colourbar
cbar = plt.colorbar(points, ax=ax)
cbar.set_label("True mass of detectable signal")
# the first axis has only the signal content of the jet
ax = ax_arr[1]
ax.set_title("Signal content of jets")
tag_rapidity_distance = ak.flatten(
ak.from_iter(tag_rapidity_in) - eventWise.DetectableTag_Rapidity)
tag_phi_distance = ak.flatten(
(ak.from_iter(tag_phi_in) - eventWise.DetectableTag_Phi))
tag_phi_distance = Components.raw_to_angular_distance(tag_phi_distance)
# print the mean distance
tag_mean_distance = np.nanmean(
np.sqrt(tag_rapidity_distance**2 + tag_phi_distance**2))
ax.text(0,
y_range[1] * 0.75,
f"mean distance={tag_mean_distance:.2f}",
horizontalalignment='center')
# put a dot at the origin
ax.scatter([0], [0], s=5, c='k', marker='o')
points = ax.scatter(tag_rapidity_distance,
tag_phi_distance,
c=colour,
alpha=alpha)
ax.set_xlabel("Jet rapidity offset")
ax.set_ylabel("Jet $|\\phi|$ offset")
ax.set_ylim(y_range)
ax.set_xlim(x_range)
# add a colourbar
cbar = plt.colorbar(points, ax=ax)
cbar.set_label("True mass of detectable signal")
if len(ax_arr) > 2:
# then in the last axis discribe the jet
PlottingTools.discribe_jet(eventWise, jet_name, ax=ax_arr[2])
fig.subplots_adjust(top=0.88,
bottom=0.11,
left=0.05,
right=1.0,
hspace=0.2,
wspace=0.255)
fig.set_size_inches(12, 4.)
return mean_distance, tag_mean_distance
def allocate(jet_phis,
jet_rapidities,
tag_phis,
tag_rapidities,
max_angle2,
valid_jets=None):
"""
In a given event each tag is allocated to a jet.
Each tag may only be allocated up to once, jets may recive multiple tags.
If no valid jet is found inside the max_angle the tag will not be allocated.
Unallocated tags are returned as -1.
Parameters
----------
jet_phis : array of floats
the center of mass phi of each jet
jet_rapidities : array of floats
the center of mass rapidity of each jet
tag_phis : array of floats
the phi of each tag
tag_rapidities : array of floats
the rapidity of each tag
valid_jets : array like of ints
The idx of the jets that can be tagged, as found in the eventWise
If None then all jets are valid
(Default value = None)
jet_name : str
The prefix of the jet vairables in the eventWise
max_angle2 : float
the maximium angle that a tag may be alloated to squared
(a deltaR mesurment)
Returns
-------
closest : numpy array of ints
The indices of the closest jet to each tag,
-1 if no sutable jet found
"""
if valid_jets is not None:
# check that a list of indices not a bool mask has been passed
if len(valid_jets) == len(jet_rapidities):
# should they be the same length every index must appear
# just check for the last one
assert len(jet_rapidities) - 1 in valid_jets
# select only valid jets for comparison
jet_phis = jet_phis[valid_jets]
jet_rapidities = jet_rapidities[valid_jets]
tag_phis = ak.to_numpy(tag_phis).reshape((-1, 1))
phi_distance = tag_phis - jet_phis
phi_distance = Components.raw_to_angular_distance(phi_distance)
tag_rapidities = ak.to_numpy(tag_rapidities).reshape((-1, 1))
rap_distance = tag_rapidities - jet_rapidities
dist2 = np.square(phi_distance) + np.square(rap_distance)
try:
closest = np.argmin(dist2, axis=1)
except (np.AxisError, ValueError):
assert len(dist2) == 0
return dist2.reshape((0, 0))
if valid_jets is not None:
# revert to true indices
closest = ak.to_numpy(valid_jets)[closest]
# remove anything with too large an angle
dist2_closest = np.min(dist2, axis=1)
closest[dist2_closest > max_angle2] = -1
return closest
def test_add_tags_particles():
params = {}
jet_name = "Jet"
# event 0
params['Jet_Label'] = [ak.from_iter([])]
params['Jet_Parent'] = [ak.from_iter([])]
params['Jet_Phi'] = [ak.from_iter([])]
params['Jet_Rapidity'] = [ak.from_iter([])]
params['Jet_Energy'] = [ak.from_iter([])]
params['Jet_PT'] = [ak.from_iter([])]
params['Phi'] = [ak.from_iter([])]
params['Rapidity'] = [ak.from_iter([])]
params['Children'] = [ak.from_iter([])]
params['Parents'] = [ak.from_iter([])]
params['MCPID'] = [ak.from_iter([])]
params['DetectableTag_Roots'] = [ak.from_iter([])]
# event 1
params['Jet_Label'] += [ak.from_iter([[0, 1, 2], [3, 4, 5]])]
params['Jet_Parent'] += [ak.from_iter([[1, -1, 1], [1, -1, 1]])]
params['Jet_Phi'] += [ak.from_iter([[0., 1., 1.], [1., np.pi, 1.]])]
params['Jet_Rapidity'] += [ak.from_iter([[0., 1., 1.], [1., 2., 1.]])]
params['Jet_Energy'] += [ak.from_iter([[1., 1., 1.], [1., 1., 1.]])]
params['Jet_PT'] += [ak.from_iter([[1., 5., 1.], [1., 5., 1.]])]
params['Phi'] += [ak.from_iter([4., 0., -np.pi, -1., np.pi, 2.])]
params['Rapidity'] += [ak.from_iter([0., 0., -2., 10., 2., 0.])]
params['Children'] += [
ak.from_iter([[], [2, 3], [5], [6, 5], [], [], [7, 8, 9], [], [], [],
[]])
]
# 0 1 2 3 4 5 6 7
params['Parents'] += [
ak.from_iter([[], [], [1], [1], [], [2, 3], [3], [6], [6], [6], []])
]
# 0 1 2 3 4 5 6 7 8
params['MCPID'] += [ak.from_iter([4, -5, 5, 44, 2, 1, -5, -1, 7, 11, 12])]
params['DetectableTag_Roots'] += [ak.from_iter([[2]])]
params['X'] = [[], []]
expected = [[], [1]]
expected_creators = [np.array([]), np.array([[]])]
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**params)
# try with append falst
tag_angle = 8.
hyperparameter_content, content =\
TrueTag.add_tags(eventWise, jet_name, tag_angle,
overwrite=True, append=False)
assert hyperparameter_content[jet_name + "_TagAngle"] == tag_angle
# as we didn't append the eventWise content should have stayed the same
assert jet_name + "_Tags" not in eventWise.columns
# the add tags method will call the add_tag_particles
tag_angle = 10.
TrueTag.add_tags(eventWise, jet_name, tag_angle, append=True)
hyperparameter_content, content =\
TrueTag.add_tags(eventWise, jet_name, tag_angle, overwrite=False, append=False)
assert hyperparameter_content[jet_name + "_TagAngle"] == tag_angle
# check the tagged jets
eventWise.selected_event = 0
tst.assert_allclose(eventWise.Jet_Tags.tolist(), [])
# in event2 the only tag is particle 2
# particle 2 has phi=-pi rap=-2
# this is closest to the first jet at 0. 0.
eventWise.selected_event = 1
assert len(ak.flatten(eventWise.Jet_Tags)) == 1
assert eventWise.Jet_Tags[0][0] == 2
assert content[jet_name + "_Tags"][1][0][0] == 2
# add tag particles ~ colud be a sperate test really....
TrueTag.add_tag_particles(eventWise)
# check the tag particles
eventWise.selected_event = 0
tst.assert_allclose(eventWise.TagIndex, expected[0])
tst.assert_allclose(ak.flatten(eventWise.TagCreators),
expected_creators[0].flatten())
eventWise.selected_event = 1
tst.assert_allclose(eventWise.TagIndex, expected[1])
tst.assert_allclose(ak.flatten(eventWise.TagCreators),
ak.flatten(expected_creators[1]))
def add_tags(eventWise,
jet_name,
max_angle,
batch_length=100,
min_tracks=None,
silent=False,
append=True,
overwrite=False):
"""
Calculate and allocate the tags in the traditional way, using add_detectable_fourvector
as we only wish to tag with particles that are in principle detectable.
Parameters
----------
eventWise : EventWise
dataset containing locations of particles and jets
jet_name : str
The prefix of the jet vairables in the eventWise
max_angle : float
the maximium angle that a tag may be alloated to a jet
(a deltaR mesurment)
If None then max_angle is drawn from Constants.py
(Default value = None)
batch_length: int
max number of events to process
(Default value = 100)
min_tracks : int
the minimum number of track for a jet to eb considered for tagging.
If None then min_tracks is drawn from Constants.py
(Default value = None)
silent : bool
Should the progress be printed?
(Default value = False)
append : bool
Should the results be appended to the eventWise?
(Default value = True)
overwrite : bool
Should existing results be abandoned?
(Default value = True)
Returns
-------
(if not appending)
hyperparameter_content : dict
hyperparameter values for eventWise
content: dict of awkward arrays
content for eventWise
"""
if min_tracks is None:
min_tracks = Constants.min_ntracks
if max_angle is None:
max_angle = Constants.max_tagangle
eventWise.selected_event = None
name = jet_name + "_Tags"
namePID = jet_name + "_TagPIDs"
n_events = len(getattr(eventWise, jet_name + "_Energy", []))
if "DetectableTag_Roots" not in eventWise.columns:
add_detectable_fourvector(eventWise, silent=silent)
if overwrite:
jet_tags = []
jet_tagpids = []
else:
jet_tags = list(getattr(eventWise, name, []))
jet_tagpids = list(getattr(eventWise, namePID, []))
start_point = len(jet_tags)
name_tagangle = jet_name + "_TagAngle"
hyperparameter_content = {name_tagangle: max_angle}
if start_point >= n_events:
print("Finished")
content = {}
content[name] = ak.from_iter(jet_tags)
content[namePID] = ak.from_iter(jet_tagpids)
if append:
return
else:
return hyperparameter_content, content
end_point = min(n_events, start_point + batch_length)
if not silent:
print(f" Will stop at {end_point/n_events:.1%}")
# name the vaiables to be cut on
# will actually compare the square of the angle for speed
max_angle2 = max_angle**2
for event_n in range(start_point, end_point):
if event_n % 10 == 0 and not silent:
print(f"{event_n/n_events:.1%}", end='\r', flush=True)
if os.path.exists("stop"):
print(f"Completed event {event_n-1}")
break
eventWise.selected_event = event_n
# get the tags
tags = ak.flatten(eventWise.DetectableTag_Roots).tolist()
# Jacan need this to not be a BQuark!!!
tag_phis = eventWise.Phi[tags]
tag_raps = eventWise.Rapidity[tags]
# get the valiables to cut on
jet_roots = getattr(eventWise, jet_name + "_Parent") == -1
jet_pt = ak.flatten(getattr(eventWise, jet_name + "_PT")[jet_roots])
if len(jet_pt) == 0:
num_tracks = []
valid_jets = []
else:
# note this actually counts num pesudojets,
# but for more than 2 that is sufficient
num_tracks = \
Components.apply_array_func(len, getattr(eventWise,
jet_name+"_PT"))
valid_jets = np.where(num_tracks > min_tracks - 0.1)[0]
jets_tags = [[] for _ in jet_pt]
if len(tags) > 0 and len(valid_jets) > 0:
# there may not be any of the particles we wish to tag in the event
# or there may not be any jets
jet_phis = ak.flatten(
getattr(eventWise, jet_name + "_Phi")[jet_roots])
jet_raps = ak.flatten(
getattr(eventWise, jet_name + "_Rapidity")[jet_roots])
closest_matches = allocate(jet_phis, jet_raps, tag_phis, tag_raps,
max_angle2, valid_jets)
else:
closest_matches = []
# keep only the indices for space reasons
for match, particle in zip(closest_matches, tags):
if match != -1:
jets_tags[match].append(particle)
jet_tags.append(ak.from_iter(jets_tags))
tagpids = [eventWise.MCPID[jet] for jet in tags]
jet_tagpids.append(ak.from_iter(tagpids))
content = {}
content[name] = ak.from_iter(jet_tags)
content[namePID] = ak.from_iter(jet_tagpids)
hyperparameter_content = {name_tagangle: max_angle}
if append:
eventWise.append(**content)
eventWise.append_hyperparameters(**hyperparameter_content)
else:
return hyperparameter_content, content
def test_create_eigenvectors():
params = {}
# event 0 # no particles, should not have issues
params['X'] = [ak.from_iter([])]
params['JetInputs_SourceIdx'] = [ak.from_iter([])]
params['JetInputs_Energy'] = [ak.from_iter([])]
params['JetInputs_Px'] = [ak.from_iter([])]
params['JetInputs_Py'] = [ak.from_iter([])]
params['JetInputs_Pz'] = [ak.from_iter([])]
# event 1 # just one particle should still work
params['X'] += [ak.from_iter(np.arange(1))]
params['JetInputs_SourceIdx'] += [ak.from_iter(np.arange(1))]
params['JetInputs_Energy'] += [ak.from_iter([30.])]
params['JetInputs_Px'] += [ak.from_iter([3.])]
params['JetInputs_Py'] += [ak.from_iter([3.])]
params['JetInputs_Pz'] += [ak.from_iter([3.])]
# event 2 # two totally seperated particles shoudl produce totaly sperated results
params['X'] += [ak.from_iter(np.arange(1))]
params['JetInputs_SourceIdx'] += [ak.from_iter(np.arange(1))]
params['JetInputs_Energy'] += [ak.from_iter([30., 40.])]
params['JetInputs_Px'] += [ak.from_iter([3., 3.])]
params['JetInputs_Py'] += [ak.from_iter([3., -3.])]
params['JetInputs_Pz'] += [ak.from_iter([3., -3.])]
# event 3 # standard event
params['X'] += [ak.from_iter(np.arange(6))]
params['JetInputs_SourceIdx'] += [ak.from_iter(np.arange(6))]
params['JetInputs_Energy'] += [
ak.from_iter([30., 10., 20., 70., 20., 10.])
]
params['JetInputs_Px'] += [ak.from_iter([3., 1., 2., 1., 2., -1.])]
params['JetInputs_Py'] += [ak.from_iter([3., 1., 2., 2., 2., 1.])]
params['JetInputs_Pz'] += [ak.from_iter([3., 0., 2., 0., 2., 2.])]
# invarient_mass 873 sqrt99 388 4859, 388, 94
# shifted energy 30 10 20 70 sqrt(393) sqrt(103)
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**params)
Components.add_phi(eventWise, "JetInputs")
Components.add_PT(eventWise, "JetInputs")
Components.add_rapidity(eventWise, "JetInputs")
#TODO renitroduce this
#jet_params = dict(CutoffDistance=1.)
jet_params = {"EigenvalueLimit": np.inf}
values, vectors = ParameterInvestigation.create_eigenvectors(
eventWise, jet_params)
# first event was empty
assert len(values[0]) == 0
assert len(vectors[0]) == 0
# second event only had one object
assert len(values[1]) == 0
tst.assert_allclose(vectors[1].shape, [1, 0])
# third event should be totaly seperated
assert len(values[2]) == 1
tst.assert_allclose(vectors[2].shape, [2, 1])
# TODO need the CutoffDistance for this
#tst.assert_allclose(vectors[2][0, 1], 0)
#tst.assert_allclose(vectors[2][1, 0], 0)
# forth event is regular
assert len(values[3]) == 5
assert len(vectors[3]) == 6
assert len(vectors[3][0]) == 5
def marry(hepmc, root_particles):
"""
Combine the information in a hepmc file and a root file
in one eventWise, to have richer information about the particles
Parameters
----------
hepmc : str or EventWise
if a str should be the path to the hepmc file on disk
if an EventWise should be a dataset with the
hepmc event data
root_particles : str or EventWise
if a str should be the path to the root file on disk
if an EventWise should be a dataset with the
root event data
Returns
-------
new_eventWise : EventWise
dataset containing matched data from teh root file and
the hepmc file.
"""
if isinstance(root_particles, str):
root_particles = Components.RootReadout(root_particles,
['Particle', 'Track', 'Tower'])
if isinstance(hepmc, str):
if Components.EventWise.pottential_file(hepmc):
hepmc = ReadHepmc.Hepmc.from_file(hepmc)
else:
hepmc = ReadHepmc.Hepmc(hepmc)
# first we assert that they both contain the same number of events
n_events = len(root_particles.PID)
assert n_events == len(
hepmc.MCPID), "Files cotain doferent number of events"
# this being extablished we compare them eventwise
# (hepmc name, root name)
precicely_equivalent = [('MCPID', 'PID'), ('Status_code', 'Status')]
close_equivalent = [('Px', 'Px'), ('Py', 'Py'), ('Pz', 'Pz'),
('Energy', 'Energy'), ('Generated_mass', 'Mass')]
for event_n in range(n_events):
hepmc.selected_event = event_n
root_particles.selected_event = event_n
# the particles are expected to have the same order in both files
for hepmc_name, root_name in precicely_equivalent:
assert np.all(root_particles.__getattr__(root_name)
== hepmc.__getattr__(hepmc_name)), \
f"{root_name} in root file not equal to {hepmc_name}"+\
" in hepmc file"
for hepmc_name, root_name in close_equivalent:
np.testing.assert_allclose(
root_particles.__getattr__(root_name),
hepmc.__getattr__(hepmc_name),
err_msg=f"{root_name} in root file not close" +
"to {hepmc_name} in hepmc file")
# remove all selected indices
hepmc.selected_event = None
root_particles.selected_event = None
# we will keep all of the root columns, but only a selection of the hepmc columns
# ensure a copy is made
columns = [name for name in root_particles.columns]
# forcefully read in here
contents = {key: getattr(root_particles, key) for key in columns}
per_event_hepmc_cols = hepmc.event_information_cols + hepmc.weight_cols + \
hepmc.units_cols + hepmc.cross_section_cols
columns += sorted(per_event_hepmc_cols)
for name in per_event_hepmc_cols:
contents[name] = getattr(hepmc, name)
# some get renamed
per_vertex_hepmc_cols = {
'Vertex_barcode': 'Vertex_barcode',
'X': 'Vertex_X',
'Y': 'Vertex_Y',
'Z': 'Vertex_Z',
'Ctau': 'Vertex_Ctau'
}
columns += sorted(per_vertex_hepmc_cols.values())
for name, new_name in per_vertex_hepmc_cols.items():
contents[new_name] = getattr(hepmc, name)
per_particle_hepmc_cols = [
'End_vertex_barcode', 'Start_vertex_barcode', 'Parents', 'Children',
'Is_root', 'Is_leaf'
]
columns += sorted(per_particle_hepmc_cols)
for name in per_particle_hepmc_cols:
contents[name] = hepmc.__getattr__(name)
# firgure out which of the root columns are per particle and which are per event
per_particle_root_cols = []
per_event_root_cols = []
for name in root_particles.columns:
values = getattr(root_particles, name)
_, depth = Components.detect_depth(values[:])
if depth == 0:
per_event_root_cols.append(name)
else:
per_particle_root_cols.append(name)
# record what has what level of granularity
contents['per_event'] = ak.from_iter(per_event_root_cols +
per_event_hepmc_cols)
contents['per_vertex'] = ak.from_iter(per_vertex_hepmc_cols.values())
contents['per_particle'] = ak.from_iter(per_particle_root_cols +
per_particle_hepmc_cols)
# make the new object and save it
file_name = (root_particles.file_name.split('.', 1)[0] + '_particles' +
Components.EventWise.FILE_EXTENTION)
dir_name = root_particles.dir_name
path_name = os.path.join(dir_name, file_name)
new_eventWise = Components.EventWise(path_name, columns, contents)
new_eventWise.write()
return new_eventWise