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