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}")