def test_IDConverter():
ids = PDGNames.IDConverter()
assert 'c' == ids[4]
assert 'cbar' == ids[-4]
assert 2/3 == ids.charges[4]
assert -2/3 == ids.charges[-4]
# try somethign that won't be found
large_num = 10**100
assert large_num == ids[large_num]
def test_upper_layers():
# will need an eventwise with Parents, Children, MCPID
# layer -1 0 1 1 -1 2 2 3 3 3 -1
# idx 0 1 2 3 4 5 6 7 8 9 10
children = [[], [2, 3], [5], [6, 5], [], [], [7, 8, 9], [], [], [], []]
parents = [[], [], [1], [1], [], [2, 3], [3], [6], [6], [6], []]
mcpid = [4, 5, 5, 3, 2, 1, -5, -1, 7, 11, 12]
expected = [2, 6]
labeler = PDGNames.IDConverter()
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(Children=[ak.from_iter(children)],
Parents=[ak.from_iter(parents)],
MCPID=[ak.from_iter(mcpid)])
eventWise.selected_event = 0
expected_particle_idx = [0, 1, 2, 3, 4, 10]
expected_children = ak.from_iter(
[c for i in expected_particle_idx for c in children[i]])
expected_parents = ak.from_iter(
[p for i in expected_particle_idx for p in parents[i]])
expected_labels = [labeler[mcpid[i]] for i in expected_particle_idx]
shower = FormShower.upper_layers(eventWise, n_layers=2)
order = np.argsort(shower.particle_idxs)
tst.assert_allclose(shower.particle_idxs[order], expected_particle_idx)
tst.assert_allclose(ak.flatten(ak.from_iter(shower.children[order])),
expected_children)
tst.assert_allclose(ak.flatten(ak.from_iter(shower.parents[order])),
expected_parents)
for a, b in zip(shower.labels[order], expected_labels):
assert a == b
# try with capture pids
expected_particle_idx = [0, 1, 2, 3, 4, 5, 6, 10]
expected_children = ak.from_iter(
[c for i in expected_particle_idx for c in children[i]])
expected_parents = ak.from_iter(
[p for i in expected_particle_idx for p in parents[i]])
expected_labels = [labeler[mcpid[i]] for i in expected_particle_idx]
shower = FormShower.upper_layers(eventWise,
n_layers=2,
capture_pids=[1])
order = np.argsort(shower.particle_idxs)
tst.assert_allclose(shower.particle_idxs[order], expected_particle_idx)
tst.assert_allclose(ak.flatten(ak.from_iter(shower.children[order])),
expected_children)
tst.assert_allclose(ak.flatten(ak.from_iter(shower.parents[order])),
expected_parents)
for a, b in zip(shower.labels[order], expected_labels):
assert a == b
def upper_layers(eventWise, n_layers=5, capture_pids=[]):
"""
Make a shower of just the topmost layers of the event
Parameters
----------
eventWise :
n_layers :
(Default value = 5)
capture_pids :
(Default value = [])
Returns
-------
"""
assert eventWise.selected_event is not None
n_particles = len(eventWise.Parents)
# start from the roots
particle_idxs = set(get_roots(list(range(n_particles)), eventWise.Parents))
current_layer = [*particle_idxs] # copy it into a new layer
locations_to_capture = {i for i, pid in enumerate(eventWise.MCPID) if pid in capture_pids}
layer_reached = 1
while locations_to_capture or layer_reached < n_layers:
children = set(ak.flatten(eventWise.Children[current_layer]))
locations_to_capture.difference_update(children)
particle_idxs.update(children)
current_layer = list(children)
layer_reached += 1
particle_idxs = list(particle_idxs)
labeler = PDGNames.IDConverter()
labels = [labeler[pid] for pid in eventWise.MCPID[particle_idxs]]
shower = Shower(particle_idxs, eventWise.Parents[particle_idxs],
eventWise.Children[particle_idxs], labels, amalgam=True)
return shower
def get_showers(eventWise, exclude_pids=True):
"""
From each root split the decendants into showers
Each particle can only belong to a single shower.
Parameters
----------
eventWise :
param exclude_pids: (Default value = [2212, 25, 35])
exclude_pids :
(Default value = True)
Returns
-------
"""
if exclude_pids is True:
exclude_pids = [25, 35]
# problem, protons can actually appear later on in the shower, if they want
mask = []
for i, p in enumerate(eventWise.MCPID):
if p == 2122 and len(eventWise.Parents[i]) == 0:
mask.append(False)
else:
mask.append(p in exclude_pids)
mask = [p not in exclude_pids for p in eventWise.MCPID]
else:
if exclude_pids is None:
exclude_pids = []
mask = [p not in exclude_pids for p in eventWise.MCPID]
particle_idxs = np.where(mask)[0] # this messes up the indexing
parent_ids = eventWise.Parents[mask] # these refer to particle_idxs
child_ids = eventWise.Children[mask] # not neat indices
pids = eventWise.MCPID[mask]
# now where possible convert labels to names
pdg = PDGNames.IDConverter()
labels = np.array([pdg[x] for x in pids])
# now we have values for the whole event,
# but we want to split the event into showers
# at start all particles are allocated to a diferent shower
showers = []
root_gids = get_roots(particle_idxs, parent_ids)
list_particle_idxs = list(particle_idxs)
for root_gid in root_gids:
root_idx = list_particle_idxs.index(root_gid)
shower_indices = []
stack = [root_idx]
while stack:
next_idx = stack.pop()
if next_idx in shower_indices:
continue # prevents looping
shower_indices.append(next_idx)
stack += [list_particle_idxs.index(child) for child in child_ids[next_idx]
if child in particle_idxs]
assert len(set(shower_indices)) == len(shower_indices)
new_shower = Shower(particle_idxs[shower_indices],
parent_ids[shower_indices].tolist(),
child_ids[shower_indices].tolist(),
labels[shower_indices])
showers.append(new_shower)
return showers