diff --git a/cli/dumpTree.py b/cli/dumpTree.py index 4f43f54c..1c58057a 100644 --- a/cli/dumpTree.py +++ b/cli/dumpTree.py @@ -115,15 +115,30 @@ 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'), ('start_process', 'u4'), + ('start_subprocess', 'u4'), + ('end_process', 'u4'), + ('end_subprocess', 'u4')]) + + segments_list = [] + trajectories_list = [] for jentry in range(entries): print(jentry) @@ -139,51 +154,67 @@ def dump(input_file, output_file): # printPrimaryVertex("PP", primaryVertex) # Dump the trajectories - 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]['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) # 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) diff --git a/cli/simulate_pixels.py b/cli/simulate_pixels.py index 8450f770..b9784aa6 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: @@ -313,6 +314,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)