Skip to content

Commit

Permalink
Merge pull request #31 from peter-madigan/upgrade/add-process-truth
Browse files Browse the repository at this point in the history 
Upgrade/add process truth
  • Loading branch information
@soleti
soleti authored Sep 20, 2021
2 parents de637f0 + 78370e3 commit e676302
Show file tree
Hide file tree
Showing 2 changed files with 77 additions and 44 deletions.
119 changes: 75 additions & 44 deletions cli/dumpTree.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)
2 changes: 2 additions & 0 deletions cli/simulate_pixels.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit e676302

Please sign in to comment.