Skip to content

Updated event display code #32

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jacoblyons98 wants to merge 3 commits into
base: develop
Choose a base branch
from
Open

Updated event display code #32

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
612e924
Update EventDisplay.py
jacoblyons98 Aug 29, 2022
ee41a85
Update EventDisplay.py
jacoblyons98 Aug 29, 2022
960d8fb
Update EventDisplay.py
jacoblyons98 Aug 29, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
165 changes: 47 additions & 118 deletions root_gnn/EventDisplay.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,10 +33,14 @@ def plot_jet_rings(jet_eta,jet_phi,tau_tagged):
plt.plot(a_out, b_out, 'y')
plt.plot(jet_eta, jet_phi, 'y')

def evt_display(filename, evt_type, evt_id, jet_idxs = 'None'):
"""
Displays all tracks and towers in an event along with rings for each jet. Includes functionallity to only graph particular jets and their associated tracks/towers
"""
def evt_display(filename, evt_id, title, output, jet_idxs = 'None'):
'''
filename: root ntuple with jet data
evt_id: which event you want to display
title: title of graph
output: file destination to save display
jet_idxs: optional argument for which jets in event to display
'''
tree_name = "output"
chain = ROOT.TChain(tree_name, tree_name)
chain.Add(filename)
Expand All @@ -46,8 +50,12 @@ def evt_display(filename, evt_type, evt_id, jet_idxs = 'None'):
fig = plt.gcf()
fig.set_size_inches(12, 12)
if jet_idxs == 'None':
plt.plot(chain.TowerEta,chain.TowerPhi, 'mo', fillstyle ='none', markersize = 5, label = 'towers')
plt.plot( chain.TrackEta,chain.TrackPhi, 'cx', markersize = 5, label = 'tracks')
pt = list(chain.TowerEt) + list(chain.TrackPt)
mini, maxi = min(pt), max(pt)
norm = plt.Normalize(mini, maxi)
cm = plt.cm.get_cmap('Blues_r')
plt.scatter(chain.TowerEta,chain.TowerPhi, linewidths = 0.1, s=36, norm=norm, c=chain.TowerEt, cmap=cm, label = 'towers')
plt.scatter( chain.TrackEta,chain.TrackPhi, marker = 'x', s=36,norm=norm, c=chain.TrackPt, cmap=cm, label = 'tracks')

tower_phi = []
tower_eta = []
Expand All @@ -57,150 +65,71 @@ def evt_display(filename, evt_type, evt_id, jet_idxs = 'None'):
last_track_idx = 0 #keeps track of last final index for tracks and towers
#this for loop splits up all of the tower/track phi and eta data into different lists for each jet, compiles together as list of lists
for jet_id in range(chain.nJets):
jet_tower_n = chain.JetTowerN[jet_id]
tower_phi.append(chain.JetTowerPhi[last_tower_idx:last_tower_idx+jet_tower_n])
tower_eta.append(chain.JetTowerEta[last_tower_idx:last_tower_idx+jet_tower_n])
last_tower_idx += jet_tower_n

jet_track_n = chain.JetGhostTrackN[jet_id]
jet_track_idxs = chain.JetGhostTrackIdx[last_track_idx:last_track_idx+jet_track_n]

jet_track_eta = []
jet_track_phi = []
for jet_track_idx in jet_track_idxs:
jet_track_eta.append(chain.TrackEta[jet_track_idx])
jet_track_phi.append(chain.TrackPhi[jet_track_idx])
track_eta.append(jet_track_eta)
track_phi.append(jet_track_phi)
last_track_idx += jet_track_n

plot_jet_rings(chain.JetEta[jet_id],chain.JetPhi[jet_id], chain.TruthJetIsTautagged[jet_id])
plt.plot(tower_eta[jet_id],tower_phi[jet_id], 'bo', fillstyle ='none', markersize = 7, label = 'towers')
plt.plot(track_eta[jet_id],track_phi[jet_id], 'gx', markersize = 7, label = 'tracks')
plt.title('{} Event {}'.format(evt_type, evt_id))
plt.xlabel('Pseudorapiditiy ($\eta$)',fontsize = 18)
plt.ylabel('Azimuthal angle ($\phi$)', fontsize = 18)
# plt.legend()
plt.savefig('{}_Event_Display_{}'.format(evt_type, evt_id))
plt.title(title)
plt.xlabel('Pseudorapiditiy ($\eta$)',fontsize = 16)
plt.ylabel('Azimuthal angle ($\phi$)', fontsize = 16)
plt.colorbar(ax = ax, shrink = 0.5)
plt.tight_layout()
plt.savefig(output)
plt.show()

else:
tower_phi = []
tower_eta = []
tower_Et = []
track_phi = []
track_eta = []
track_pt = []
last_tower_idx = 0
last_track_idx = 0 #keeps track of last final index for tracks and towers
#this for loop splits up all of the tower/track phi and eta data into different lists for each jet, compiles together as list of lists
for jet_id in range(chain.nJets):
jet_tower_n = chain.JetTowerN[jet_id]
tower_phi.append(chain.JetTowerPhi[last_tower_idx:last_tower_idx+jet_tower_n])
tower_eta.append(chain.JetTowerEta[last_tower_idx:last_tower_idx+jet_tower_n])
tower_Et.append(chain.JetTowerEt[last_tower_idx:last_tower_idx+jet_tower_n])
last_tower_idx += jet_tower_n

jet_track_n = chain.JetGhostTrackN[jet_id]
jet_track_idxs = chain.JetGhostTrackIdx[last_track_idx:last_track_idx+jet_track_n]

jet_track_eta = []
jet_track_phi = []
jet_track_pt = []
for jet_track_idx in jet_track_idxs:
jet_track_eta.append(chain.TrackEta[jet_track_idx])
jet_track_phi.append(chain.TrackPhi[jet_track_idx])
jet_track_pt.append(chain.TrackPt[jet_track_idx])
track_eta.append(jet_track_eta)
track_phi.append(jet_track_phi)
track_pt.append(jet_track_pt)
last_track_idx += jet_track_n


#now, looks at each jet_idx given to function on plots that jet and it's associated towers and tracks
for jet_idx in jet_idxs:
plot_jet_rings(chain.JetEta[jet_idx],chain.JetPhi[jet_idx], chain.TruthJetIsTautagged[jet_idx])
plt.plot(tower_eta[jet_idx],tower_phi[jet_idx], 'bo', fillstyle ='none', markersize = 6, label = 'towers')
plt.plot(track_eta[jet_idx],track_phi[jet_idx], 'gx', markersize = 6, label = 'tracks')

pt = list(tower_Et[jet_idx]) + list(track_pt[jet_idx])
mini, maxi = min(pt), max(pt)
norm = plt.Normalize(mini, maxi)
cm = plt.cm.get_cmap('Blues_r')
plt.scatter(tower_eta[jet_idx],tower_phi[jet_idx], marker ='o', linewidths = 0.5, s=36, norm=norm, c=tower_Et[jet_idx], cmap=cm, label = 'towers')
plt.scatter(track_eta[jet_idx],track_phi[jet_idx], marker ='x', norm=norm, s=36, c=track_pt[jet_idx], cmap=cm, label = 'tracks')

#finishes plotting the graph
plt.title('{} Event {} Jets {}'.format(evt_type, evt_id,jet_idxs), fontsize = 18)
plt.xlabel('Pseudorapiditiy ($\eta$)',fontsize = 18)
plt.ylabel('Azimuthal angle ($\phi$)', fontsize = 18)
plt.title(title, fontsize = 16)
plt.xlabel('Pseudorapiditiy ($\eta$)',fontsize = 16)
plt.ylabel('Azimuthal angle ($\phi$)', fontsize = 16)
plt.colorbar(ax = ax, shrink = 0.5)
plt.tight_layout()
plt.legend()
plt.savefig('{}_Event_Display_{}_Jets_{}'.format(evt_type, evt_id,jet_idxs))
plt.savefig(output)
plt.show()

def jet_display_edges(filename, evt_type, evt_id, jet_idxs):
"""
Displays a single jet with color coded edges conntecting tracks-tracks and towers-towers, will add functionality of track-tower edges
"""
tree_name = "output"
chain = ROOT.TChain(tree_name, tree_name)
chain.Add(filename)
chain.GetEntry(evt_id)
ax = plt.gca() #you first need to get the axis handle
ax.set_aspect(1)
fig = plt.gcf()
fig.set_size_inches(12, 12)
tower_phi = []
tower_eta = []
track_phi = []
track_eta = []

edge_tower_phi = []
edge_tower_eta = []
edge_track_phi = []
edge_track_eta = []

last_tower_idx = 0
last_track_idx = 0 #keeps track of last final index for tracks and towers
#this for loop splits up all of the tower/track phi and eta data into different lists for each jet, compiles together as list of lists
for jet_id in range(chain.nJets):
jet_tower_n = chain.JetTowerN[jet_id]
jet_tower_phi = chain.JetTowerPhi[last_tower_idx:last_tower_idx+jet_tower_n]
jet_tower_eta = chain.JetTowerEta[last_tower_idx:last_tower_idx+jet_tower_n]
tower_phi.append(jet_tower_phi)
tower_eta.append(jet_tower_eta)
last_tower_idx += jet_tower_n

#tower-tower edges
jet_tower_phi_edges = list(itertools.permutations(jet_tower_phi, 2))
jet_tower_eta_edges = list(itertools.permutations(jet_tower_eta, 2))
edge_tower_phi.append(jet_tower_phi_edges)
edge_tower_eta.append(jet_tower_eta_edges)

jet_track_n = chain.JetGhostTrackN[jet_id]
jet_track_idxs = chain.JetGhostTrackIdx[last_track_idx:last_track_idx+jet_track_n]

jet_track_eta = []
jet_track_phi = []
for jet_track_idx in jet_track_idxs:
jet_track_eta.append(chain.TrackEta[jet_track_idx])
jet_track_phi.append(chain.TrackPhi[jet_track_idx])
track_eta.append(jet_track_eta)
track_phi.append(jet_track_phi)
last_track_idx += jet_track_n

#track-track edges
jet_track_phi_edges = list(itertools.permutations(jet_track_phi, 2))
jet_track_eta_edges = list(itertools.permutations(jet_track_eta, 2))
edge_track_phi.append(jet_track_phi_edges)
edge_track_eta.append(jet_track_eta_edges)

#now, looks at each jet_idx given to function on plots that jet and it's associated towers and tracks
for jet_idx in jet_idxs:
plot_jet_rings(chain.JetEta[jet_idx],chain.JetPhi[jet_idx], chain.TruthJetIsTautagged[jet_idx])
plt.plot(tower_eta[jet_idx],tower_phi[jet_idx], 'bo', fillstyle ='none', markersize = 6, label = 'towers')
plt.plot(track_eta[jet_idx],track_phi[jet_idx], 'gx', markersize = 6, label = 'tracks')
for edge_idx in range(len(edge_tower_phi[jet_idx])):
plt.plot(edge_tower_eta[jet_idx][edge_idx],edge_tower_phi[jet_idx][edge_idx], 'b')
for edge_idx in range(len(edge_track_phi[jet_idx])):
plt.plot(edge_track_eta[jet_idx][edge_idx],edge_track_phi[jet_idx][edge_idx], 'g')
#finishes plotting the graph
plt.title('{} Event {} Jets {}'.format(evt_type, evt_id,jet_idxs), fontsize = 18)
plt.xlabel('Pseudorapiditiy ($\eta$)',fontsize = 18)
plt.ylabel('Azimuthal angle ($\phi$)', fontsize = 18)
plt.legend()
plt.savefig('{}_Event_Display_{}_Jets_{}'.format(evt_type, evt_id,jet_idxs))
plt.show()

if __name__ == "__main__":
filename_tau = '/global/homes/j/jacobl/v0/Ntuple_ditau_processed.root'
filename_qcd = '/global/homes/j/jacobl/v0/Ntuple_qcd_processed.root'
filename_tau_pileup = '/global/cfs/cdirs/m3443/data/TauStudies/v2/Ntuple_ditau_PU.root'
evt_display(filename_tau,'tau',4)
filename_tau_pileup = '/global/cfs/cdirs/m3443/data/TauStudies/v4/Ntuple_ditau_PU_processed.root'
filename_qcd_pileup = '/global/cfs/cdirs/m3443/data/TauStudies/v4/Ntuple_qcd_processed_PU.root'
evt_display(filename_tau_pileup, 4, 'Jet from Hadronically Decayed Tau', 'tau_jet.pdf',[0])
evt_display(filename_qcd_pileup, 4, 'QCD Jet', 'qcd_jet.pdf',[0])