diff --git a/Tracking/python/eventDisplay.py b/Tracking/python/eventDisplay.py new file mode 100644 index 000000000..b8841407f --- /dev/null +++ b/Tracking/python/eventDisplay.py @@ -0,0 +1,154 @@ +from LDMX.Framework import EventTree +import numpy as np +import matplotlib.pyplot as plt +import uproot +import math +import time +import copy +import os +import sys +import ROOT as r +import statistics as stat +from mpl_toolkits.mplot3d import Axes3D, art3d +import matplotlib.animation as animation +from sklearn.linear_model import LinearRegression + +def pause(): + programPause = input("Press the key to continue...") + +def trackPlotter(tree, event_number, tag, save): + recoilSimHits = addBranch(tree, 'SimTrackerHit', 'RecoilSimHits_{}'.format(tag)) + digiRecoil = addBranch(tree, 'Measurement', 'DigiRecoilSimHits_{}'.format(tag)) + ecalRecHit = addBranch(tree, 'EcalHit', 'EcalRecHits_{}'.format(tag)) + #recoil = addBranch(tree, 'ReducedTrack', 'ReducedTracks_{}'.format(tag)) + recoil = addBranch(tree, 'ReducedTrack', 'ReducedSeedTracks_{}'.format(tag)) + + tree.GetEntry(event_number) + + recoilSim_x = [] + recoilSim_y = [] + recoilSim_z = [] + + xpos_digi_tot = [] + ypos_digi_tot = [] + zpos_digi_tot = [] + + ecal_end_x = [] + ecal_end_y = [] + ecal_end_z = [] + + trackParams = [] + + for particle in recoilSimHits: + recoilSim_z.append(particle.getPosition()[2]) + recoilSim_x.append(particle.getPosition()[0]) + recoilSim_y.append(particle.getPosition()[1]) + + for x_digi in digiRecoil: + zpos_digi_tot.append(x_digi.getGlobalPosition()[0]) + xpos_digi_tot.append(x_digi.getGlobalPosition()[1]) + ypos_digi_tot.append(x_digi.getGlobalPosition()[2]) + + for x_ecal in ecalRecHit: + if (x_ecal.getZPos() < 250): + ecal_end_x.append(x_ecal.getXPos()) + ecal_end_y.append(x_ecal.getYPos()) + ecal_end_z.append(x_ecal.getZPos()) + + for x in recoil: + trackParams.append((x.getAX(), x.getBX(), x.getAY(), x.getBY(), x.getDistanceToRecHit())) + + digiPoints = np.column_stack((zpos_digi_tot, xpos_digi_tot, ypos_digi_tot)) + ecalRecHits = np.column_stack((ecal_end_z, ecal_end_x, ecal_end_y)) + recoilSim = np.column_stack((recoilSim_z, recoilSim_x, recoilSim_y)) + + fig = plt.figure(figsize=(8, 8)) + ax = fig.add_subplot(111, projection='3d') + + x_unc = 0.006 #mm + y_unc = 0.12 # mm + + for zi, xi, yi in zip(digiPoints[:,0], digiPoints[:,1], digiPoints[:,2]): + ax.plot([zi, zi], [xi - x_unc, xi + x_unc], [yi, yi], color='black',linewidth=2) # x error + ax.plot([zi, zi], [xi, xi], [yi - y_unc, yi + y_unc], color='black',linewidth=2) # y error + + ax.scatter(digiPoints[:,0], digiPoints[:,1], digiPoints[:,2], c='b', marker='o', label='DigiRecoil SimHit', s=25) + ax.scatter(ecalRecHits[:, 0], ecalRecHits[:, 1], ecalRecHits[:, 2], c='purple', label='ECalRecHit', s=50, alpha=0.5) + ax.scatter(recoilSim[:, 0], recoilSim[:, 1], recoilSim[:, 2], marker='o', c='gray', label='RecoilSimHits', s=50, alpha=0.5) + + i=0 + for track in trackParams: + ax.plot([0.0, ecalRecHits[0][0]], [track[1], track[0]*ecalRecHits[0][0]+track[1]], [track[3], track[2]*ecalRecHits[0][0]+track[3]], 'r--', label=f'Reconstructed Track {i+1}, d_RecHit = {track[4]:.2f} mm') + i += 1 + + x = ecalRecHits[0][0] + y_range = np.linspace(min(ecalRecHits[:,1]) - 5, max(ecalRecHits[:,1]) + 5, 50) + z_range = np.linspace(min(ecalRecHits[:,2]) - 5, max(ecalRecHits[:,2]) + 5, 50) + Y, Z = np.meshgrid(y_range, z_range) + X = np.full(Y.shape, ecalRecHits[0][0]) + + plane = ax.plot_surface(X, Y, Z, color='red', alpha=0.1, edgecolor='none') + + for recHit in ecalRecHits: + center_z, center_x, center_y = recHit[0], recHit[1], recHit[2] + radius = 3.87 + theta = np.linspace(0, 2 * np.pi, 100) + circle_x = center_x + radius * np.cos(theta) + circle_y = center_y + radius * np.sin(theta) + circle_z = np.full_like(circle_x, center_z) + ax.plot(circle_z, circle_x, circle_y, c='purple') + + # Labels and title + ax.set_xlabel('z [mm]') + ax.set_ylabel('x [mm]') + ax.set_zlabel('y [mm]') + + fig.text(0.15, 0.9, 'LDMX', ha='center', fontsize=36, fontweight='bold') + fig.text(0.43, 0.9, 'Work in Progress', ha='center', fontsize=24, fontstyle='italic') + fig.text(0.85, 0.9, '(4 GeV)', ha='center', fontsize=12, fontstyle='italic') + fig.subplots_adjust(right=0.99) # Adjust the value as needed + left, bottom, width, height = [0.06, 0.05, 0.85, 0.85] # Adjust these values as needed + ax.set_position([left, bottom, width, height]) + +# ax.legend(fontsize=10) + + plt.savefig(f'{save}/{tag}_eventDisplay_eventID_{event_number}.png') + + ax.set_xlim(0, 50) + ax.set_ylim(-10, 10) + ax.set_zlim(-10, 10) + plane.set_visible(False) + + plt.savefig(f'{save}/{tag}_eventDisplay_eventID_{event_number}_ZOOM.png') + +def addBranch(tree, ldmx_class, branch_name): + if tree == None: + sys.exit('Set tree') + + if ldmx_class == 'EventHeader': branch = r.ldmx.EventHeader() + elif ldmx_class == 'EcalVetoResult': branch = r.ldmx.EcalVetoResult() + elif ldmx_class == 'HcalVetoResult': branch = r.ldmx.HcalVetoResult() + elif ldmx_class == 'TriggerResult': branch = r.ldmx.TriggerResult() + elif ldmx_class == 'SimParticle': branch = r.std.map(int, 'ldmx::'+ldmx_class)() + else: branch = r.std.vector('ldmx::'+ldmx_class)() + + tree.SetBranchAddress(branch_name,r.AddressOf(branch)) + + return branch + +def main(): + tree = r.TChain("LDMX_Events") + tree.Add('events_1_reducedAlgoV1Test_seederONLY.root') + + nentries = tree.GetEntries() + print("nentries = ", nentries) + + tag = 'reduced_algo_v1_test' + save_loc = '/Users/fdelzanno/Desktop/Incandela/Coding/ldmx-sw' + + trackPlotter(tree, 0, tag, save_loc) + + +if __name__ == "__main__": + main() +