analysis/TPC-ClusterAnimation/TPC_Cluster_Drift_Animator.py

0001 import numpy as np
0002 import json
0003 import matplotlib.pyplot as plt
0004 import pandas as pd
0005 import uproot
0006 import awkward as ak
0007 import mpl_toolkits.mplot3d.art3d as art3d
0008 from matplotlib.patches import Wedge
0009 import matplotlib.animation as animation
0010 import array
0011 #import time
0012 #from timeit import default_timer as timer
0013
0014 #/*************************************************************/
0015 #/*              TPC Cluster Drift Animator                  */
0016 #/*         Aditya Prasad Dash, Thomas Marshall              */
0017 #/*      aditya55@physics.ucla.edu, rosstom@ucla.edu         */
0018 #/*************************************************************/
0019 #Code in python
0020 #Input:
0021 # root or json file containing TPC cluster positions
0022 #Output:
0023 # Animation of drifting of TPC clusters with user defined speed and option to save in .mp4 format
0024
0025 def TPC_surface(inner_radius,outer_radius, length_z):
0026     ngridpoints=30
0027     z = np.linspace(-length_z, length_z, ngridpoints)
0028     phi = np.linspace(0, 2*np.pi, ngridpoints)
0029     phi_grid, z_grid=np.meshgrid(phi, z)
0030     x_grid_inner = inner_radius*np.cos(phi_grid)
0031     y_grid_inner = inner_radius*np.sin(phi_grid)
0032     x_grid_outer = outer_radius*np.cos(phi_grid)
0033     y_grid_outer = outer_radius*np.sin(phi_grid)
0034
0035     return x_grid_inner,y_grid_inner,x_grid_outer,y_grid_outer,z_grid
0036
0037 def TPC_endcap(inner_radius,outer_radius, length_z):
0038     ngridpoints=30
0039     radius = np.linspace(inner_radius, outer_radius, 5)
0040     phi = np.linspace(0, 2*np.pi, ngridpoints)
0041     phi_grid, r_grid=np.meshgrid(phi, radius)
0042     x_grid=[]
0043     y_grid=[]
0044     for r in radius:
0045         x_grid_radius = r*np.cos(phi_grid)
0046         y_grid_radius = r*np.sin(phi_grid)
0047         x_grid=np.append(x_grid,x_grid_radius)
0048         y_grid=np.append(y_grid,y_grid_radius)
0049
0050     z_grid=x_grid
0051     return x_grid,y_grid,z_grid
0052
0053
0054 def raddist_cluster(cluster_pos):
0055     radius=np.sqrt(cluster_pos[:,0]*cluster_pos[:,0]+cluster_pos[:,1]*cluster_pos[:,1])
0056     return radius
0057
0058 def theLoop(iteration,dataPoint,scatter):
0059     effective_time=iteration-dataPoint[4]
0060     if(effective_time>=0):
0061         effective_z = 0
0062         if(dataPoint[2]>0):
0063             effective_z=dataPoint[2]+drift_speed_posz[2]*effective_time
0064         if(dataPoint[2]<0):
0065             effective_z=dataPoint[2]+drift_speed_negz[2]*effective_time
0066         if(abs(effective_z)<len_TPC+drift_speed_posz[2]):
0067             scatter._offsets3d = (dataPoint[0:1], dataPoint[1:2], [effective_z])
0068             color=['r','g','b','c','m','y']
0069
0070             if(abs(effective_z)<len_TPC):
0071                 scatter.set_color(color[int(dataPoint[3]%6)])
0072
0073             if(abs(effective_z)>=len_TPC):
0074                 scatter.set_color('white')
0075                 scatter.set(alpha=1.0)
0076
0077         elif(abs(effective_z)<len_TPC+2*drift_speed_posz[2]):
0078                 scatter._offsets3d = ([100], [-100], [100]) #to plot all points outside TPC at one point
0079                 scatter.set_color('black')
0080                 scatter.set_sizes([10])
0081     return 0
0082
0083 #Parallel processing
0084 def animate_scatters(iteration, data,
0085     scatters,fig_text,time_scale,iteration_time,start_iteration):
0086     iteration=iteration+start_iteration
0087     if(iteration%1==0):
0088         print("iteration=")
0089         print(iteration)
0090     iter_array=[iteration]*len(data)
0091     time=(iteration)*iteration_time/time_scale*(10**3)
0092     theLoop_return = [theLoop(iteration,da,scat) for da,scat in zip(data,scatters)]
0093     fig_text.set_text(str(round(time,1))+"$\mu$s")
0094
0095     return scatters,fig_text
0096
0097 #Normal Processing
0098 #def animate_scatters(iteration, data, scatters,fig_text,time_scale,iteration_time,start_iteration):
0099 #    print("iteration=")
0100 #    print(iteration)
0101 #    for i in range(data.shape[0]):
0102 #        time=(iteration+start_iteration)*iteration_time/time_scale*(10**3)
0103 #        effective_time=iteration-data[i,4]
0104 #        if(effective_time>=0):
0105 #            if(data[i,2]>0):
0106 #                effective_z=data[i,2]+drift_speed_posz[2]*effective_time
0107 #            if(data[i,2]<0):
0108 #                effective_z=data[i,2]+drift_speed_negz[2]*effective_time
0109 #            if(abs(effective_z)<len_TPC+drift_speed_posz[2]):
0110 #                scatters[i]._offsets3d = (data[i,0:1], data[i,1:2], [effective_z])
0111 #                color=['r','g','b','c','m','y']
0112 #
0113 #                if(abs(effective_z)<len_TPC):
0114 #                    scatters[i].set_color(color[int(data[i,3]%6)])
0115 #
0116 #                if(abs(effective_z)>=len_TPC):
0117 #                    scatters[i].set_color('white')
0118 #                    scatters[i].set(alpha=1.0)
0119 #
0120 #            if(abs(effective_z)>=len_TPC+2*drift_speed_posz[2]):
0121 #                    scatters[i]._offsets3d = ([100], [-100], [100]) #to plot all points outside TPC at one point
0122 #                    scatters[i].set_color('black')
0123 #                    scatters[i].set_sizes([10])
0124 #        else:
0125 #            scatters[i]._offsets3d = ([100], [-100], [100]) #clusters from event not yet taken place
0126 #            scatters[i].set_color('black')
0127 #
0128 #    fig_text.set_text(str(round(time,2))+"$\mu$s")
0129 #
0130 #    return scatters,fig_text
0131
0132 def animate_clusters(data,animation_name="Animated_clusters_TPC.mp4",save=False,start_iteration=0,no_iterations=1):
0133
0134     # Attaching 3D axis to the figure
0135     fig = plt.figure(figsize=[7.5,7.5],layout='constrained')
0136     ax = fig.add_subplot(111, projection='3d',facecolor='black',alpha=0.0)
0137     ax.grid(False)
0138     ax.margins(0.0)
0139     ax.xaxis.set_pane_color((1,1,1,0))
0140     ax.xaxis.line.set_color('w')
0141     ax.spines['top'].set_color('w')
0142     ax.spines['bottom'].set_color('w')
0143     ax.spines['left'].set_color('w')
0144     ax.spines['right'].set_color('w')
0145
0146     ax.yaxis.set_pane_color((1,1,1,0))
0147     ax.yaxis.line.set_color('w')
0148     ax.zaxis.set_pane_color((1,1,1,0))
0149     ax.zaxis.line.set_color('w')
0150
0151     #Drawing TPC
0152     endcap1=Wedge((0, 0), 80, 0, 360, color='blue',alpha=0.7)
0153     endcap2=Wedge((0, 0), 80, 0, 360, color='blue',alpha=0.7)
0154     endcap1.set_width(60)
0155     endcap2.set_width(60)
0156
0157     ax.add_artist(endcap1)
0158     ax.add_artist(endcap2)
0159
0160     art3d.pathpatch_2d_to_3d(endcap1, z=len_TPC, zdir="z")
0161     art3d.pathpatch_2d_to_3d(endcap2, z=-len_TPC, zdir="z")
0162
0163     Xc_in,Yc_in,Xc_out,Yc_out,Zc = TPC_surface(20,80,len_TPC)
0164     ax.plot_surface(Xc_in, Yc_in, Zc, alpha=0.5)
0165     ax.plot_surface(Xc_out, Yc_out, Zc, alpha=0.5)
0166
0167     # Setting the axes properties
0168     ax.set_xlim3d([-100, 100])
0169     ax.set_xlabel('X (cm)',color='white',fontsize=7)
0170     ax.xaxis.set_tick_params(colors='white',labelsize=7)
0171
0172     ax.set_ylim3d([-100, 100])
0173     ax.set_ylabel('Y (cm)',color='white',fontsize=7)
0174     ax.yaxis.set_tick_params(colors='white',labelsize=7)
0175
0176     ax.set_zlim3d([-120, 120])
0177     ax.set_zlabel('Z (cm)',color='white',fontsize=7)
0178     ax.zaxis.set_tick_params(colors='white',labelsize=7)
0179
0180     ax.set_title('Clusters drifting in TPC') #(time scaled by $2*10^{5}$)')
0181     fig_text=ax.text(-100,90,100,  'time', size=10,color='w',alpha=0.9)
0182     fig_text_sPhenix=ax.text(-100,130,100,  'sPHENIX', size=10,fontweight='bold',style='italic',color='w',alpha=0.9)
0183     fig_text_TPC=ax.text(-60,135,70,  'TPC simulation', size=10,color='w',alpha=0.9)
0184
0185     fig_text_sPhenix=ax.text(-100,110,100,  'p+p, $\sqrt{s_{NN}}$ = 200 GeV 4MHz', size=10,color='w',alpha=0.9)
0186
0187
0188     # Providing the starting angle for the view
0189     ax.view_init(10,70,0,'y')
0190
0191     # Initializing scatters
0192     scatters = [ ax.scatter([100],[-100],[100]) for i in range(data.shape[0])]
0193     for i in range(data.shape[0]): scatters[i].set_color('black')
0194     print("Plot initialized")
0195     #start = timer()
0196     ani = animation.FuncAnimation(fig, animate_scatters, iterations, fargs=(data, scatters,fig_text,time_scale,iteration_time,start_iteration),
0197                                        interval=iteration_time, blit=False, repeat=False) #interval is in milliseconds and is the time between each frame
0198     if save:
0199         print("Saving animation as"+animation_name)
0200         ani.save(animation_name,writer='ffmpeg')
0201         print("Animation saved")
0202     #end = timer()
0203     #print("Time for process=")
0204     #print(end-start)
0205     #plt.show()
0206
0207 def read_cluster_pos(inFile):
0208     if(inFile.lower().endswith('.json')):
0209         print("Reading data from json file")
0210         file=open(inFile)
0211         data_json=json.load(file)
0212         file.close()
0213         dict_json = data_json.items()
0214         numpy_array_json = np.array(list(dict_json))
0215         mom_dict=numpy_array_json[3][1] #only works for this format the 3rd row is TRACKS and and 1 refers to INNERTRACKER
0216         innertracker_arr=mom_dict['INNERTRACKER'] #stores the tracks information as an array
0217
0218         print("Reading clusters")
0219         #Getting the cluster positions of a track
0220         x_y_z_clusters=np.array([])
0221         no_tracks=len(innertracker_arr) #set no_tracks=10 for testing
0222
0223         for track_no in range(no_tracks):
0224             x_y_z_dict_track=innertracker_arr[track_no].copy() #innertracker_arr[i] reads the ith track
0225             x_y_z_clusters_track=np.array(x_y_z_dict_track['pts']) #2D array the x,y,z positions corresponding to each cluster e.g. x_y_z[0][0] is the x coordinate of the 1st cluster
0226             event_and_times=np.zeros((len(x_y_z_clusters_track),1))
0227             x_y_z_clusters_track=np.append(x_y_z_clusters_track,event_and_times,axis=1)
0228             x_y_z_clusters_track=np.append(x_y_z_clusters_track,event_and_times,axis=1)
0229             x_y_z_clusters_track=x_y_z_clusters_track[raddist_cluster(x_y_z_clusters_track[:,:2])>30]
0230
0231             if(track_no==0):
0232                 x_y_z_clusters=np.copy(x_y_z_clusters_track)
0233
0234             else:
0235                 x_y_z_clusters=np.append(x_y_z_clusters,x_y_z_clusters_track,axis=0)
0236         return x_y_z_clusters
0237
0238     if(inFile.lower().endswith('.root')):
0239         print("Reading data from root file")
0240         file = uproot.open(inFile)
0241         ntp_cluster_tree=file['ntp_cluster']
0242         branches=ntp_cluster_tree.arrays(["x","y","z","event","gvt","layer"])
0243         branches=branches[~np.isnan(branches.gvt)]
0244         branches=branches[((branches.x)**2+(branches.y)**2)>900]
0245         branches=branches[branches.layer>6]        #layers corresponding to TPC
0246         branches=branches[branches.layer<55]       #layers corresponding to TPC
0247         #branches=branches[branches.event>=75]
0248         #branches=branches[branches.event<55]
0249
0250         print("Reading clusters")
0251         x_y_z_clusters_run=np.array([])
0252         len_events=len(np.unique(branches.event))
0253         len_events=200
0254         event_times=[0]
0255         event_times=np.append(event_times,np.random.poisson(mean_eventtime*1000,len_events-1))
0256         event_times=np.cumsum(event_times,dtype=float)
0257         for cluster in range(len(branches)):
0258             gvt_event=event_times[int(branches[cluster]['event'])]
0259             gvt_event=gvt_event*(10**(-6))*time_scale #Time in milliseconds scaled for animation
0260             gvt_event=gvt_event/(iteration_time) #this is the time in terms of iterations
0261             x_y_z_clusters_track=np.array([[branches[cluster]['x'], branches[cluster]['y'], branches[cluster]['z'],branches[cluster]['event'],gvt_event]])
0262             if(cluster==0):
0263                 x_y_z_clusters_run=np.copy(x_y_z_clusters_track)
0264             else:
0265                 x_y_z_clusters_run=np.append(x_y_z_clusters_run,x_y_z_clusters_track,axis=0)
0266         return x_y_z_clusters_run
0267
0268 print("Generating data for animation")
0269
0270
0271 # Main Program starts from here
0272 np.random.seed(3)
0273 #User defined values
0274 time_scale=4.0*(10.0**5) #inverse of speed scale
0275 collision_rate=4.0 #MHz #Set fig_text_sPhenix in line 188
0276 mean_eventtime=1/collision_rate
0277 print("Mean event time (microseconds)=")
0278 print(mean_eventtime)
0279 TPC_drift_speed=8.0*(10.0**3) #Actual TPC drift speed =8cm/microsecond=8*10^3cm/millisecond
0280 iteration_time=100.0 #actual duration between frames in FuncAnimation
0281 len_TPC=105.0
0282
0283 drift_speed_posz=np.array([0.0,0.0,TPC_drift_speed/time_scale*iteration_time,0.0,0.0])
0284 drift_speed_negz=np.array([0.0,0.0,-TPC_drift_speed/time_scale*iteration_time,0.0,0.0])
0285 print("Drift_speed_posz(cm/iteration)=")
0286 print(drift_speed_posz)
0287 print("Drift_speed_negz(cm/iteration)=")
0288 print(drift_speed_negz)
0289
0290 #ANIMATION
0291 data=read_cluster_pos("G4sPHENIX_g4svtx_eval.root")
0292
0293 # Number of iterations
0294 no_events=np.max(data[:,3])+1
0295 print("no_events=")
0296 print(no_events)
0297
0298 iterations = int((no_events-1)*mean_eventtime*time_scale/(iteration_time*1000)+len_TPC/drift_speed_posz[2])+5
0299 #iterations=20
0300
0301 print("number of iterations=")
0302 print(iterations)
0303 print("Animation starting!")
0304
0305 #Saving takes a long time so use Save=True only when necessary
0306 #increase drift_speed_posz and drift_speed_negz if desired
0307 animate_clusters(data,"Animated_clusters_TPC.mp4",save=True,start_iteration=0,no_iterations=iterations)
0308
0309
0310 #Merge mp4 files using ffmpeg -f concat -safe 0 -i filelist.txt -c copy mergedVideo.mp4