def test_match():
# first try getting an exact match
desired = 'b'
pid_list = [3, 5, -5, 5122, 111]
tst.assert_allclose(PDGNames.match(pid_list, desired, False),
[False, True, False, False, False])
tst.assert_allclose(PDGNames.match(pid_list, desired, True),
[False, True, True, True, False])
tst.assert_allclose(PDGNames.match([], desired, True), [])
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 charged_indices(pids):
""" Indicate charged tracks which can have their vertices reconstructed.
Parameters
----------
pids : listlike of int
monte carlo particle ids of the particles;
https://pdg.lbl.gov/2007/reviews/montecarlorpp.pdf
"""
identifier = PDGNames.Identities()
indices = [i for i, pid in enumerate(pids) if identifier.charges[pid]]
return indices
def test_Identities():
ids = PDGNames.Identities()
str(ids) # check the string method runs
known_particles = [{'id': 5, 'name': 'b', 'charge': -1/3, 'spin':1/2},
{'id': -5, 'name': 'bbar', 'charge': 1/3, 'spin':1/2},
{'id': 25, 'name': 'h0', 'charge': 0, 'spin': 0}]
for particle in known_particles:
found = ids[particle['id']]
for key in particle:
assert particle[key] == found[key], f"For test particle {particle} found {found}"
assert 'cbar' == ids.antiNames[4]
assert 'cbar' == ids.antiNames[-4]
assert 2/3 == ids.charges[4]
assert -2/3 == ids.charges[-4]
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 test_ReadHepmc():
check_lines()
hepmc_file = os.path.join(data_dir, "mini.hepmc")
n_events = 4
# try starting from 0
hepmc1 = ReadHepmc.Hepmc(hepmc_file, 0, n_events)
# try startng from 1
start2 = 1
hepmc2 = ReadHepmc.Hepmc(hepmc_file, start=start2, stop=n_events)
idents = PDGNames.Identities()
all_ids = set(idents.particle_data[:, idents.columns["id"]])
for event_n in range(n_events): # loop over out selection of events
hepmc1.selected_event = event_n
# sanity checks on the components
assert hepmc1.Event_n == event_n
# again for the other one
check_event_consistancy(hepmc1, all_ids)
try:
hepmc2.selected_event = event_n
assert hepmc2.Event_n == event_n + 1
except (IndexError, ValueError):
assert event_n + start2 == n_events
continue # reach the end of 1
check_event_consistancy(hepmc2, all_ids)
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
def void_test_ReadHepmc2():
hepmc_file = os.path.join(data_dir, "billy_tag_2_pythia8_events.hepmc")
n_events = 4
hepmc = ReadHepmc.Hepmc(hepmc_file, 0, n_events)
idents = PDGNames.Identities()
mom_diff = []
all_ids = set(idents.particle_data[:, idents.columns["id"]])
for event_n in range(n_events): # loop over out selection of events
hepmc.selected_event = event_n
# sanity checks on the components
assert hepmc.Event_n == event_n
# check the beam particles
beam_idx1 = np.where(hepmc.Particle_barcode == hepmc.Barcode_beam_particle1)[0][0]
beam_idx2 = np.where(hepmc.Particle_barcode == hepmc.Barcode_beam_particle2)[0][0]
assert hepmc.Is_root[beam_idx1]
assert hepmc.Is_root[beam_idx2]
assert len(hepmc.Parents[beam_idx1]) == 0
assert len(hepmc.Parents[beam_idx2]) == 0
# valid PID
assert set(np.abs(hepmc.MCPID)).issubset(all_ids)
# conservation of momentum at the vertices
parent_p4 = np.zeros(4)
child_p4 = np.zeros(4)
issues = 0
relations = 0
conserved_multiplicity = []
unconserved_multiplicity = []
for v_barcode in hepmc.Vertex_barcode:
parent_p4[:] = 0
child_p4[:] = 0
parents_idx = np.where(hepmc.End_vertex_barcode == v_barcode)[0]
children_idx = np.where(hepmc.Start_vertex_barcode == v_barcode)[0]
if len(parents_idx) == 0:
for cidx in children_idx:
assert hepmc.Is_root[cidx]
elif len(children_idx) == 0:
for pidx in parents_idx:
assert hepmc.Is_leaf[pidx]
else:
for pidx in parents_idx:
parent_p4 += (hepmc.Energy[pidx], hepmc.Px[pidx], hepmc.Py[pidx], hepmc.Pz[pidx])
for cidx in children_idx:
child_p4 += (hepmc.Energy[cidx], hepmc.Px[cidx], hepmc.Py[cidx], hepmc.Pz[cidx])
relations += 1
mom_diff.append(parent_p4 - child_p4)
try:
tst.assert_allclose(parent_p4[1:], child_p4[1:], rtol=0.05, atol=0.001)
conserved_multiplicity.append(len(children_idx) + len(parents_idx))
except AssertionError as e:
if issues < 1 and event_n==0:
parent_barcodes = hepmc.Particle_barcode[parents_idx]
children_barcodes = hepmc.Particle_barcode[children_idx]
print(f"Vertex {v_barcode} has incoming particles with barcodes {parent_barcodes} and outgoing particles with barcodes {children_barcodes}")
print(f"The incoming total momentum is {parent_p4}")
print(f"The outgoing total momentum is {child_p4}")
issues += 1
unconserved_multiplicity.append(len(children_idx) + len(parents_idx))
print(f"Event {event_n}")
print(f"non-conservation {issues}; {100*issues/relations:.1f}%")
print(f"multipicity, conserved {np.mean(conserved_multiplicity):.2f}, unconserved {np.mean(unconserved_multiplicity):.2f}")
return np.array(mom_diff)
def void_test_ReadHepmc():
hepmc_file = os.path.join(data_dir, "mini.hepmc")
n_events = 1
hepmc = ReadHepmc.Hepmc(hepmc_file, 0, n_events)
idents = PDGNames.Identities()
all_ids = set(idents.particle_data[:, idents.columns["id"]])
for event_n in range(n_events): # loop over out selection of events
hepmc.selected_event = event_n
# sanity checks on the components
assert hepmc.Event_n == event_n
# check the beam particles
beam_idx1 = np.where(hepmc.Particle_barcode == hepmc.Barcode_beam_particle1)[0][0]
beam_idx2 = np.where(hepmc.Particle_barcode == hepmc.Barcode_beam_particle2)[0][0]
assert hepmc.Is_root[beam_idx1]
assert hepmc.Is_root[beam_idx2]
assert len(hepmc.Parents[beam_idx1]) == 0
assert len(hepmc.Parents[beam_idx2]) == 0
# valid PID
assert set(np.abs(hepmc.MCPID)).issubset(all_ids)
# conservation of momentum and parent child reflection
num_particles = len(hepmc.MCPID)
parent_p4 = np.zeros(4)
child_p4 = np.zeros(4)
self_p4 = np.zeros(4)
children_issues = 0
children_relations = 0
parent_issues = 0
parent_relations = 0
conserved_parent_multiplicity = []
conserved_children_multiplicity = []
unconserved_parent_multiplicity = []
unconserved_children_multiplicity = []
for idx in range(num_particles): # loop over particles in event
parent_p4[:] = 0
child_p4[:] = 0
self_p4[:] = (hepmc.Energy[idx], hepmc.Px[idx], hepmc.Py[idx], hepmc.Pz[idx])
if not hepmc.Is_root[idx]:
parent_relations += 1
# we have already checked the parents of roots
parent_idx = hepmc.Parents[idx]
assert len(parent_idx) > 0
for pidx in parent_idx:
assert idx in hepmc.Children[pidx]
parent_p4 += (hepmc.Energy[pidx], hepmc.Px[pidx], hepmc.Py[pidx], hepmc.Pz[pidx])
# if these parents have children besides this child
# subtract their momentum out
for cidx in hepmc.Children[pidx]:
if cidx == idx:
continue
else:
parent_p4 -= (hepmc.Energy[cidx], hepmc.Px[cidx], hepmc.Py[cidx], hepmc.Pz[cidx])
try:
tst.assert_allclose(parent_p4[1:], self_p4[1:], rtol=0.05, atol=0.001)
conserved_parent_multiplicity.append(len(parent_idx))
except AssertionError as e:
#print(f"Event {event_n}, particle {idx}, parents; {e}")
parent_issues += 1
unconserved_parent_multiplicity.append(len(parent_idx))
if hepmc.Is_leaf[idx]:
assert len(hepmc.Children[idx]) == 0
else:
children_relations += 1
children_idx = hepmc.Children[idx]
assert len(children_idx) > 0
for cidx in children_idx:
assert idx in hepmc.Parents[cidx]
child_p4 += (hepmc.Energy[cidx], hepmc.Px[cidx], hepmc.Py[cidx], hepmc.Pz[cidx])
# if these children have parents besides this one
# subtract their momentum out
for pidx in hepmc.Parents[cidx]:
if pidx == idx:
continue
else:
child_p4 -= (hepmc.Energy[pidx], hepmc.Px[pidx], hepmc.Py[pidx], hepmc.Pz[pidx])
try:
tst.assert_allclose(child_p4[1:], self_p4[1:], rtol=0.05, atol = 0.001)
conserved_children_multiplicity.append(len(children_idx))
except AssertionError as e:
#print(f"Event {event_n}, particle {idx}, children; {e}")
children_issues += 1
unconserved_children_multiplicity.append(len(children_idx))
print(f"Event {event_n}")
print(f"Children non-conservation {children_issues}; {100*children_issues/children_relations:.1f}%")
print(f"Parent non-conservation {parent_issues}; {100*parent_issues/parent_relations:.1f}%")
print(f"Any non-conservation {100*(parent_issues+children_issues)/(children_relations+parent_relations):.1f}%")
print(f"Children multipicity, conserved {np.mean(conserved_children_multiplicity):.2f}, unconserved {np.mean(unconserved_children_multiplicity):.2f}")
print(f"Parent multipicity, conserved {np.mean(conserved_parent_multiplicity):.2f}, unconserved {np.mean(unconserved_parent_multiplicity):.2f}")