From 7eb25e2eae22d5becb43da9f8556f0582bb08906 Mon Sep 17 00:00:00 2001 From: peter-madigan Date: Wed, 8 Sep 2021 16:27:06 -0700 Subject: [PATCH 1/4] Add trajectory truth info to dumpTree.py output --- cli/dumpTree.py | 113 ++++++++++++++++++++++++++++++------------------ 1 file changed, 71 insertions(+), 42 deletions(-) diff --git a/cli/dumpTree.py b/cli/dumpTree.py index 4f43f54c..cee79f4b 100644 --- a/cli/dumpTree.py +++ b/cli/dumpTree.py @@ -115,15 +115,28 @@ def dump(input_file, output_file): # Read all of the events. entries = inputTree.GetEntriesFast() - dtype = np.dtype([('eventID','u4'),('z_end','f4'),('trackID','u4'), - ('tran_diff','f4'),('z_start','f4'),('x_end','f4'), - ('y_end','f4'),('n_electrons','u4'),('pdgId','i4'), - ('x_start','f4'),('y_start','f4'),('t_start','f4'), - ('dx','f4'),('long_diff','f4'),('pixel_plane','u4'), - ('t_end','f4'),('dEdx','f4'),('dE','f4'),('t','f4'), - ('y','f4'),('x','f4'),('z','f4')]) - - segments = [] + segments_dtype = np.dtype([('eventID', 'u4'), ('z_end', 'f4'), + ('trackID', 'u4'), ('tran_diff', 'f4'), + ('z_start', 'f4'), ('x_end', 'f4'), + ('y_end', 'f4'), ('n_electrons', 'u4'), + ('pdgId', 'i4'), ('x_start', 'f4'), + ('y_start', 'f4'), ('t_start', 'f4'), + ('dx', 'f4'), ('long_diff', 'f4'), + ('pixel_plane', 'u4'), ('t_end', 'f4'), + ('dEdx', 'f4'), ('dE', 'f4'), ('t', 'f4'), + ('y', 'f4'), ('x', 'f4'), ('z', 'f4')]) + + trajectories_dtype = np.dtype([('eventID', 'u4'), ('trackID', 'u4'), + ('parentID', 'i4'), + ('pxyz_start', 'f4', (3,)), + ('xyz_start', 'f4', (3,)), ('t_start', 'f4'), + ('pxyz_end', 'f4', (3,)), + ('xyz_end', 'f4', (3,)), ('t_end', 'f4'), + ('pdgId', 'i4'), ('process', 'u4'), + ('subprocess', 'u4')]) + + segments_list = [] + trajectories_list = [] for jentry in range(entries): print(jentry) @@ -142,48 +155,64 @@ def dump(input_file, output_file): trajectories = {'pdgId': [], 'trackId': []} print("Number of trajectories ", len(event.Trajectories)) - for trajectory in event.Trajectories: - trajectories['pdgId'].append(trajectory.GetPDGCode()) - trajectories['trackId'].append(trajectory.GetTrackId()) + trajectories = np.empty(len(event.Trajectories), dtype=trajectories_dtype) + for iTraj, trajectory in enumerate(event.Trajectories): + start_pt, end_pt = trajectory.Points[0], trajectory.Points[-1] + trajectories[iTraj]['eventID'] = jentry + trajectories[iTraj]['trackID'] = trajectory.GetTrackId() + trajectories[iTraj]['parentID'] = trajectory.GetParentId() + trajectories[iTraj]['pxyz_start'] = (start_pt.GetMomentum().X(), start_pt.GetMomentum().Y(), start_pt.GetMomentum().Z()) + trajectories[iTraj]['pxyz_end'] = (end_pt.GetMomentum().X(), end_pt.GetMomentum().Y(), end_pt.GetMomentum().Z()) + trajectories[iTraj]['xyz_start'] = (start_pt.GetPosition().X(), start_pt.GetPosition().Y(), start_pt.GetPosition().Z()) + trajectories[iTraj]['xyz_end'] = (end_pt.GetPosition().X(), end_pt.GetPosition().Y(), end_pt.GetPosition().Z()) + trajectories[iTraj]['t_start'] = start_pt.GetPosition().T() + trajectories[iTraj]['t_end'] = end_pt.GetPosition().T() + trajectories[iTraj]['process'] = start_pt.GetProcess() + trajectories[iTraj]['subprocess'] = start_pt.GetSubprocess() + trajectories[iTraj]['pdgId'] = trajectory.GetPDGCode() + trajectories_list.append(trajectories) # Dump the segment containers print("Number of segment containers:", event.SegmentDetectors.size()) for containerName, hitSegments in event.SegmentDetectors: - for hitSegment in hitSegments: - segment = np.empty(1, dtype=dtype) - segment['eventID'] = jentry - segment['trackID'] = trajectories['trackId'][hitSegment.Contrib[0]] - segment['x_start'] = hitSegment.GetStart().X()/10 - segment['y_start'] = hitSegment.GetStart().Y()/10 - segment['z_start'] = hitSegment.GetStart().Z()/10 - segment['x_end'] = hitSegment.GetStop().X()/10 - segment['y_end'] = hitSegment.GetStop().Y()/10 - segment['z_end'] = hitSegment.GetStop().Z()/10 - segment['dE'] = hitSegment.GetEnergyDeposit() - segment['t'] = 0 - segment['t_start'] = 0 - segment['t_end'] = 0 - xd = segment['x_end'] - segment['x_start'] - yd = segment['y_end'] - segment['y_start'] - zd = segment['z_end'] - segment['z_start'] + segment = np.empty(len(hitSegments), dtype=segments_dtype) + for iHit, hitSegment in enumerate(hitSegments): + segment[iHit]['eventID'] = jentry + segment[iHit]['trackID'] = trajectories[hitSegment.Contrib[0]]['trackId'] + segment[iHit]['x_start'] = hitSegment.GetStart().X() / 10 + segment[iHit]['y_start'] = hitSegment.GetStart().Y() / 10 + segment[iHit]['z_start'] = hitSegment.GetStart().Z() / 10 + segment[iHit]['x_end'] = hitSegment.GetStop().X() / 10 + segment[iHit]['y_end'] = hitSegment.GetStop().Y() / 10 + segment[iHit]['z_end'] = hitSegment.GetStop().Z() / 10 + segment[iHit]['dE'] = hitSegment.GetEnergyDeposit() + segment[iHit]['t'] = 0 + segment[iHit]['t_start'] = 0 + segment[iHit]['t_end'] = 0 + xd = segment[iHit]['x_end'] - segment[iHit]['x_start'] + yd = segment[iHit]['y_end'] - segment[iHit]['y_start'] + zd = segment[iHit]['z_end'] - segment[iHit]['z_start'] dx = sqrt(xd**2 + yd**2 + zd**2) - segment['dx'] = dx - segment['x'] = (segment['x_start'] + segment['x_end'])/2. - segment['y'] = (segment['y_start'] + segment['y_end'])/2. - segment['z'] = (segment['z_start'] + segment['z_end'])/2. - segment['dEdx'] = hitSegment.GetEnergyDeposit()/dx if dx > 0 else 0 - segment['pdgId'] = trajectories['pdgId'][hitSegment.Contrib[0]] - segment['n_electrons'] = 0 - segment['long_diff'] = 0 - segment['tran_diff'] = 0 - segment['pixel_plane'] = 0 - segments.append(segment) - segments = np.concatenate(segments, axis=0) + segment[iHit]['dx'] = dx + segment[iHit]['x'] = (segment[iHit]['x_start'] + segment[iHit]['x_end']) / 2. + segment[iHit]['y'] = (segment[iHit]['y_start'] + segment[iHit]['y_end']) / 2. + segment[iHit]['z'] = (segment[iHit]['z_start'] + segment[iHit]['z_end']) / 2. + segment[iHit]['dEdx'] = hitSegment.GetEnergyDeposit() / dx if dx > 0 else 0 + segment[iHit]['pdgId'] = trajectories[hitSegment.Contrib[0]]['pdgId'] + segment[iHit]['n_electrons'] = 0 + segment[iHit]['long_diff'] = 0 + segment[iHit]['tran_diff'] = 0 + segment[iHit]['pixel_plane'] = 0 + segments_list.append(segment) + trajectories_list = np.concatenate(trajectories_list, axis=0) + segments_list = np.concatenate(segments_list, axis=0) with h5py.File(output_file, 'w') as f: - f.create_dataset("segments", data=segments) + f.create_dataset("trajectories", data=trajectories_list) + f.create_dataset("segments", data=segments_list) + if __name__ == "__main__": fire.Fire(dump) From 43b760f365cb13c644d7333c5505fa0aa2a07d22 Mon Sep 17 00:00:00 2001 From: Peter Madigan Date: Wed, 8 Sep 2021 18:54:54 -0500 Subject: [PATCH 2/4] Fix typo: --- cli/dumpTree.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/cli/dumpTree.py b/cli/dumpTree.py index cee79f4b..ca0e72d6 100644 --- a/cli/dumpTree.py +++ b/cli/dumpTree.py @@ -152,9 +152,7 @@ def dump(input_file, output_file): # printPrimaryVertex("PP", primaryVertex) # Dump the trajectories - trajectories = {'pdgId': [], 'trackId': []} print("Number of trajectories ", len(event.Trajectories)) - trajectories = np.empty(len(event.Trajectories), dtype=trajectories_dtype) for iTraj, trajectory in enumerate(event.Trajectories): start_pt, end_pt = trajectory.Points[0], trajectory.Points[-1] @@ -180,7 +178,7 @@ def dump(input_file, output_file): segment = np.empty(len(hitSegments), dtype=segments_dtype) for iHit, hitSegment in enumerate(hitSegments): segment[iHit]['eventID'] = jentry - segment[iHit]['trackID'] = trajectories[hitSegment.Contrib[0]]['trackId'] + segment[iHit]['trackID'] = trajectories[hitSegment.Contrib[0]]['trackID'] segment[iHit]['x_start'] = hitSegment.GetStart().X() / 10 segment[iHit]['y_start'] = hitSegment.GetStart().Y() / 10 segment[iHit]['z_start'] = hitSegment.GetStart().Z() / 10 From 4218dd5775ddc479eff90cfd035b788b7eb801cd Mon Sep 17 00:00:00 2001 From: peter-madigan Date: Wed, 8 Sep 2021 16:58:21 -0700 Subject: [PATCH 3/4] Add trajectory data to simulate_pixels.py --- cli/simulate_pixels.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/cli/simulate_pixels.py b/cli/simulate_pixels.py index b95dcf32..1b5dfd6a 100755 --- a/cli/simulate_pixels.py +++ b/cli/simulate_pixels.py @@ -89,6 +89,7 @@ def run_simulation(input_filename, # For this sample we need to invert $z$ and $y$ axes with h5py.File(input_filename, 'r') as f: tracks = np.array(f['segments']) + trajectories = np.array(f['trajectories']) RangePop() if tracks.size == 0: @@ -308,6 +309,7 @@ def run_simulation(input_filename, with h5py.File(output_filename, 'a') as f: f.create_dataset("tracks", data=tracks) + f.create_dataset("trajectories", data=trajectories) f['configs'].attrs['pixel_layout'] = pixel_layout print("Output saved in:", output_filename) From 78370e3f32b30e9d9cd175603ac2321fa77bc81f Mon Sep 17 00:00:00 2001 From: peter-madigan Date: Wed, 8 Sep 2021 18:31:10 -0700 Subject: [PATCH 4/4] Keep both start and end process info --- cli/dumpTree.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/cli/dumpTree.py b/cli/dumpTree.py index ca0e72d6..1c58057a 100644 --- a/cli/dumpTree.py +++ b/cli/dumpTree.py @@ -132,8 +132,10 @@ def dump(input_file, output_file): ('xyz_start', 'f4', (3,)), ('t_start', 'f4'), ('pxyz_end', 'f4', (3,)), ('xyz_end', 'f4', (3,)), ('t_end', 'f4'), - ('pdgId', 'i4'), ('process', 'u4'), - ('subprocess', 'u4')]) + ('pdgId', 'i4'), ('start_process', 'u4'), + ('start_subprocess', 'u4'), + ('end_process', 'u4'), + ('end_subprocess', 'u4')]) segments_list = [] trajectories_list = [] @@ -165,8 +167,10 @@ def dump(input_file, output_file): trajectories[iTraj]['xyz_end'] = (end_pt.GetPosition().X(), end_pt.GetPosition().Y(), end_pt.GetPosition().Z()) trajectories[iTraj]['t_start'] = start_pt.GetPosition().T() trajectories[iTraj]['t_end'] = end_pt.GetPosition().T() - trajectories[iTraj]['process'] = start_pt.GetProcess() - trajectories[iTraj]['subprocess'] = start_pt.GetSubprocess() + trajectories[iTraj]['start_process'] = start_pt.GetProcess() + trajectories[iTraj]['start_subprocess'] = start_pt.GetSubprocess() + trajectories[iTraj]['end_process'] = end_pt.GetProcess() + trajectories[iTraj]['end_subprocess'] = end_pt.GetSubprocess() trajectories[iTraj]['pdgId'] = trajectory.GetPDGCode() trajectories_list.append(trajectories)