def test_generate_ghost_grid():
spacing = 10 # Only one point will fit
output = GhostParticles.generate_ghost_grid(spacing)
for var in output:
assert len(var) == 1
massless = output[0]**2 - output[1]**2 - output[2]**2 - output[3]**2
tst.assert_allclose(massless, 0, atol=1e-5)
mask = np.abs(output[2]) > 0.001
tst.assert_allclose(output[1][mask] / output[2][mask],
np.cos(output[5][mask]))
# many points
spacing = 1
output = GhostParticles.generate_ghost_grid(spacing)
assert len(output[0]) > 1
for var in output:
assert len(var) == len(output[0])
tst.assert_allclose(output[5][0], -np.pi)
tst.assert_allclose(sorted(set(output[5]))[1], -np.pi + spacing)
tst.assert_allclose(output[4][0], -2.5)
tst.assert_allclose(sorted(set(output[4]))[1], -2.5 + spacing)
massless = output[0]**2 - output[1]**2 - output[2]**2 - output[3]**2
tst.assert_allclose(massless, 0, atol=1e-5)
mask = np.abs(output[2]) > 0.001
tst.assert_allclose(output[1][mask] / output[2][mask],
np.cos(output[5][mask]))
def test_clustering_with_ghosts():
with TempTestDir("pseudojet") as dir_name:
# need a eventWise to call clustering
name = "empty.parquet"
ew = Components.EventWise(os.path.join(dir_name, name))
ew.selected_event = 0
ghost_indices, cluster_class = GhostParticles.clustering_with_ghosts(
FormJets.GeneralisedKT, 2)
# start by finding some data to cluster
floats = SimpleClusterSamples.two_degenerate['floats']
# must put this into the eventwise as calling with
# ints_floats causes the cluster to skip the ghosts
set_JetInputs(ew, floats)
ew.selected_event = 0
initial_indices = set(range(len(ew.JetInputs_Energy)))
# now call the clustering algorithm
pseudojets = cluster_class(ew)
jet_ghost_tests(pseudojets, initial_indices, ghost_indices)
pseudojets.run()
jet_ghost_tests(pseudojets, initial_indices, ghost_indices)
# try with non gridded
ghost_indices, cluster_class = GhostParticles.clustering_with_ghosts(
FormJets.GeneralisedKT, 10, False)
pseudojets = cluster_class(ew, run=True)
jet_ghost_tests(pseudojets, initial_indices, ghost_indices)
def collenar_splits(energies, pxs, pys, pzs, labels):
split_fractions = np.random.rand(len(pxs))
four_momentum = np.vstack((energies, pxs, pys, pzs))
four_momentum1 = four_momentum * split_fractions
four_momentum2 = four_momentum * (1 - split_fractions)
_, ints1, floats1 = GhostParticles.create_jet_internals(*four_momentum1,
labels=labels)
labels2 = -(np.array(labels) + 2)
_, ints2, floats2 = GhostParticles.create_jet_internals(*four_momentum2,
labels=labels2)
return ints1, floats1, ints2, floats2
def single_check(eventWise,
jet_class,
jet_name,
jet_params={},
num_particles=10,
compare_tree=None):
if compare_tree is None:
# make a guess baes on the jet name
if "SKM" in jet_name or "KMeans" in jet_name:
compare_tree = False
elif "Indicator" or "SI" in jet_name:
compare_tree = False
else:
compare_tree = True
print("Clustering without ghosts")
jet_name_no = "NoGhost" + jet_name
FormJets.cluster_multiapply(eventWise,
jet_class,
dict_jet_params=jet_params,
jet_name=jet_name_no,
batch_length=np.inf,
silent=False)
print("Clustering with ghosts")
GhostParticles.cluster_multiapply_with_ghosts(eventWise,
jet_class,
dict_jet_params=jet_params,
jet_name=jet_name,
batch_length=np.inf,
silent=False,
variable=num_particles,
on_grid=False)
print("Comparing")
if compare_tree:
changed = check_tree_changes(eventWise, jet_name)
else:
changed = check_leaf_changes(eventWise, jet_name)
if np.any(changed):
print("Not IR safe")
else:
print("IR safe")
eventWise.append(**{jet_name + "_GhostChanges": changed})
def test_create_jet_internals():
# a set of empty lists should result in empty lists
energy, px, py, pz, rapidity, phi =\
np.array([]), np.array([]), \
np.array([]), np.array([]), np.array([]), np.array([])
inputidxs, ints, floats = GhostParticles.create_jet_internals(
energy, px, py, pz, rapidity, phi)
assert len(inputidxs) == 0
assert len(ints) == 0
assert len(floats) == 0
# anything else should produce some regular grids
rows = 5
energy, px, py, pz, rapidity, phi =\
np.array([1]*rows), np.array([2]*rows), \
np.array([3]*rows), np.array([4]*rows), \
np.array([5]*rows), np.array([6]*rows)
inputidxs, ints, floats = GhostParticles.create_jet_internals(
energy, px, py, pz, rapidity, phi)
assert len(inputidxs) == rows
assert len(ints) == rows
assert len(floats) == rows
for i in range(rows):
tst.assert_allclose(floats[i][1:], [5, 6, 1, 2, 3, 4, 0, 1])
tst.assert_allclose(ints[i], [inputidxs[i]] + [-1] * 4)
def test_generate_ghost_random():
# shouldn't choke on 0 particles
output = GhostParticles.generate_ghost_random(0)
for var in output:
assert len(var) == 0
# just one particle
output = GhostParticles.generate_ghost_random(1)
for var in output:
assert len(var) == 1
massless = output[0]**2 - output[1]**2 - output[2]**2 - output[3]**2
tst.assert_allclose(massless, 0, atol=1e-5)
mask = np.abs(output[2]) > 0.001
tst.assert_allclose(output[1][mask] / output[2][mask],
np.cos(output[5][mask]))
# many points
n_points = 10
output = GhostParticles.generate_ghost_random(n_points)
for var in output:
assert len(var) == n_points
massless = output[0]**2 - output[1]**2 - output[2]**2 - output[3]**2
tst.assert_allclose(massless, 0, atol=1e-5)
mask = np.abs(output[2]) > 0.001
tst.assert_allclose(output[1][mask] / output[2][mask],
np.cos(output[5][mask]))
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 test_extract_particles_from_tree():
# empty events should return empty
label, old_parents, old_child1, old_child2, remove = [[]] * 5
keep_mask, parents, child1, child2 = \
GhostParticles.extract_particles_from_tree(
label, old_parents, old_child1, old_child2, remove)
assert len(keep_mask) + len(parents) + len(child1) + len(child2) == 0
# even if things are specified to remove
remove = [1]
keep_mask, parents, child1, child2 = \
GhostParticles.extract_particles_from_tree(
label, old_parents, old_child1, old_child2, remove)
assert len(keep_mask) + len(parents) + len(child1) + len(child2) == 0
# event with nothing from the remove list should return as it went in
label = [1]
old_parents = [-1]
old_child1 = [-1]
old_child2 = [-1]
remove = [5]
keep_mask, parents, child1, child2 = \
GhostParticles.extract_particles_from_tree(
label, old_parents, old_child1, old_child2, remove)
tst.assert_allclose(keep_mask, [True])
tst.assert_allclose(parents, [-1])
tst.assert_allclose(child1, [-1])
tst.assert_allclose(child2, [-1])
# events asking the only element to be removes should disappear
label = [1]
old_parents = [-1]
old_child1 = [-1]
old_child2 = [-1]
remove = [1]
keep_mask, parents, child1, child2 = \
GhostParticles.extract_particles_from_tree(
label, old_parents, old_child1, old_child2, remove)
tst.assert_allclose(keep_mask, [False])
tst.assert_allclose(parents, [])
tst.assert_allclose(child1, [])
tst.assert_allclose(child2, [])
# simple two child graph
label = [1, 2, 3]
old_parents = [2, -1, 2]
old_child1 = [-1, 1, -1]
old_child2 = [-1, 3, -1]
remove = []
keep_mask, parents, child1, child2 = \
GhostParticles.extract_particles_from_tree(
label, old_parents, old_child1, old_child2, remove)
tst.assert_allclose(keep_mask, [True, True, True])
tst.assert_allclose(parents, old_parents)
tst.assert_allclose(child1, old_child1)
tst.assert_allclose(child2, old_child2)
remove = [1]
keep_mask, parents, child1, child2 = \
GhostParticles.extract_particles_from_tree(
label, old_parents, old_child1, old_child2, remove)
tst.assert_allclose(keep_mask, [False, False, True])
tst.assert_allclose(parents, [-1])
tst.assert_allclose(child1, [-1])
tst.assert_allclose(child2, [-1])
remove = [1, 3]
keep_mask, parents, child1, child2 = \
GhostParticles.extract_particles_from_tree(
label, old_parents, old_child1, old_child2, remove)
tst.assert_allclose(keep_mask, [False, False, False])
tst.assert_allclose(parents, [])
tst.assert_allclose(child1, [])
tst.assert_allclose(child2, [])
# more complex graph
label = [1, 2, 3, 4, 5]
old_parents = [2, 4, 2, -1, 4]
old_child1 = [-1, 1, -1, 2, -1]
old_child2 = [-1, 3, -1, 5, -1]
remove = []
keep_mask, parents, child1, child2 = \
GhostParticles.extract_particles_from_tree(
label, old_parents, old_child1, old_child2, remove)
tst.assert_allclose(keep_mask, [True] * 5)
tst.assert_allclose(parents, old_parents)
tst.assert_allclose(child1, old_child1)
tst.assert_allclose(child2, old_child2)
remove = [1]
keep_mask, parents, child1, child2 = \
GhostParticles.extract_particles_from_tree(
label, old_parents, old_child1, old_child2, remove)
tst.assert_allclose(keep_mask, [False, False, True, True, True])
tst.assert_allclose(parents, [4, -1, 4])
tst.assert_allclose(child1, [-1, 3, -1])
tst.assert_allclose(child2, [-1, 5, -1])
remove = [5]
keep_mask, parents, child1, child2 = \
GhostParticles.extract_particles_from_tree(
label, old_parents, old_child1, old_child2, remove)
tst.assert_allclose(keep_mask, [True, True, True, False, False])
tst.assert_allclose(parents, [2, -1, 2])
tst.assert_allclose(child1, [-1, 1, -1])
tst.assert_allclose(child2, [-1, 3, -1])
remove = [3, 5]
keep_mask, parents, child1, child2 = \
GhostParticles.extract_particles_from_tree(
label, old_parents, old_child1, old_child2, remove)
tst.assert_allclose(keep_mask, [True, False, False, False, False])
tst.assert_allclose(parents, [-1])
tst.assert_allclose(child1, [-1])
tst.assert_allclose(child2, [-1])
remove = [1, 3, 5]
keep_mask, parents, child1, child2 = \
GhostParticles.extract_particles_from_tree(
label, old_parents, old_child1, old_child2, remove)
tst.assert_allclose(keep_mask, [False] * 5)
tst.assert_allclose(parents, [])
tst.assert_allclose(child1, [])
tst.assert_allclose(child2, [])