analysis/TPC-ClusterAnimation/TPC_Cluster_Drift_Animator_beam.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     effective_time=iteration
0061     if(effective_time>=0):
0062         effective_z = 0
0063         if(dataPoint[2]>0):
0064             effective_z=dataPoint[2]+drift_speed_posz[2]*effective_time
0065         if(dataPoint[2]<0):
0066             effective_z=dataPoint[2]+drift_speed_negz[2]*effective_time
0067         if(abs(effective_z)<len_TPC+drift_speed_posz[2]):
0068             scatter._offsets3d = (dataPoint[0:1], dataPoint[1:2], [effective_z])
0069             color=['r','g','b','c','m','y']
0070
0071             if(abs(effective_z)<len_TPC):
0072                 #scatter.set_color(color[int(dataPoint[3]%6)])
0073                 scatter.set_color('r')
0074                 scatter.set_sizes([0.1])
0075             if(abs(effective_z)>=len_TPC):
0076                 scatter.set_color('white')
0077                 scatter.set(alpha=1.0)
0078                 scatter.set_sizes([0.1])
0079
0080         elif(abs(effective_z)<len_TPC+2*drift_speed_posz[2]):
0081                 scatter._offsets3d = ([100], [-100], [100]) #to plot all points outside TPC at one point
0082                 scatter.set_color('black')
0083                 scatter.set_sizes([0.1])
0084     return 0
0085
0086 #Parallel processing
0087 def animate_scatters(iteration, data,
0088     scatters,fig_text,time_scale,iteration_time,start_iteration):
0089     iteration=iteration+start_iteration
0090     if(iteration%1==0):
0091         print("iteration=")
0092         print(iteration)
0093     iter_array=[iteration]*len(data)
0094     time=(iteration)*iteration_time/time_scale*(10**3)
0095     theLoop_return = [theLoop(iteration,da,scat) for da,scat in zip(data,scatters)]
0096     fig_text.set_text(str(round(time,1))+"$\mu$s")
0097
0098     return scatters,fig_text
0099
0100 #Normal Processing
0101 #def animate_scatters(iteration, data, scatters,fig_text,time_scale,iteration_time,start_iteration):
0102 #    print("iteration=")
0103 #    print(iteration)
0104 #    for i in range(data.shape[0]):
0105 #        time=(iteration+start_iteration)*iteration_time/time_scale*(10**3)
0106 #        effective_time=iteration-data[i,4]
0107 #        if(effective_time>=0):
0108 #            if(data[i,2]>0):
0109 #                effective_z=data[i,2]+drift_speed_posz[2]*effective_time
0110 #            if(data[i,2]<0):
0111 #                effective_z=data[i,2]+drift_speed_negz[2]*effective_time
0112 #            if(abs(effective_z)<len_TPC+drift_speed_posz[2]):
0113 #                scatters[i]._offsets3d = (data[i,0:1], data[i,1:2], [effective_z])
0114 #                color=['r','g','b','c','m','y']
0115 #
0116 #                if(abs(effective_z)<len_TPC):
0117 #                    scatters[i].set_color(color[int(data[i,3]%6)])
0118 #
0119 #                if(abs(effective_z)>=len_TPC):
0120 #                    scatters[i].set_color('white')
0121 #                    scatters[i].set(alpha=1.0)
0122 #
0123 #            if(abs(effective_z)>=len_TPC+2*drift_speed_posz[2]):
0124 #                    scatters[i]._offsets3d = ([100], [-100], [100]) #to plot all points outside TPC at one point
0125 #                    scatters[i].set_color('black')
0126 #                    scatters[i].set_sizes([0.01])
0127 #        else:
0128 #            scatters[i]._offsets3d = ([100], [-100], [100]) #clusters from event not yet taken place
0129 #            scatters[i].set_color('black')
0130 #
0131 #    fig_text.set_text(str(round(time,2))+"$\mu$s")
0132 #
0133 #    return scatters,fig_text
0134
0135 def animate_clusters(data,animation_name="Animated_clusters_TPC.mp4",save=False,start_iteration=0,no_iterations=1):
0136
0137     # Attaching 3D axis to the figure
0138     fig = plt.figure(figsize=[7.5,7.5],layout='constrained')
0139     ax = fig.add_subplot(111, projection='3d',facecolor='black',alpha=0.0)
0140     ax.grid(False)
0141     ax.margins(0.0)
0142     ax.xaxis.set_pane_color((1,1,1,0))
0143     ax.xaxis.line.set_color('w')
0144     ax.spines['top'].set_color('w')
0145     ax.spines['bottom'].set_color('w')
0146     ax.spines['left'].set_color('w')
0147     ax.spines['right'].set_color('w')
0148
0149     ax.yaxis.set_pane_color((1,1,1,0))
0150     ax.yaxis.line.set_color('w')
0151     ax.zaxis.set_pane_color((1,1,1,0))
0152     ax.zaxis.line.set_color('w')
0153
0154     #Drawing TPC
0155     endcap1=Wedge((0, 0), 80, 0, 360, color='blue',alpha=0.7)
0156     endcap2=Wedge((0, 0), 80, 0, 360, color='blue',alpha=0.7)
0157     endcap1.set_width(60)
0158     endcap2.set_width(60)
0159
0160     ax.add_artist(endcap1)
0161     ax.add_artist(endcap2)
0162
0163     art3d.pathpatch_2d_to_3d(endcap1, z=len_TPC, zdir="z")
0164     art3d.pathpatch_2d_to_3d(endcap2, z=-len_TPC, zdir="z")
0165
0166     Xc_in,Yc_in,Xc_out,Yc_out,Zc = TPC_surface(20,80,len_TPC)
0167     ax.plot_surface(Xc_in, Yc_in, Zc, alpha=0.5)
0168     ax.plot_surface(Xc_out, Yc_out, Zc, alpha=0.5)
0169
0170     # Setting the axes properties
0171     ax.set_xlim3d([-100, 100])
0172     ax.set_xlabel('X (cm)',color='white',fontsize=7)
0173     ax.xaxis.set_tick_params(colors='white',labelsize=7)
0174
0175     ax.set_ylim3d([-100, 100])
0176     ax.set_ylabel('Y (cm)',color='white',fontsize=7)
0177     ax.yaxis.set_tick_params(colors='white',labelsize=7)
0178
0179     ax.set_zlim3d([-120, 120])
0180     ax.set_zlabel('Z (cm)',color='white',fontsize=7)
0181     ax.zaxis.set_tick_params(colors='white',labelsize=7)
0182
0183     ax.set_title('Clusters drifting in TPC') #(time scaled by $2*10^{5}$)')
0184     fig_text=ax.text(-100,80,100,  'time', size=10,color='w',alpha=0.9)
0185     fig_text_sPhenix=ax.text(-100,130,100,  'sPHENIX', size=10,fontweight='bold',style='italic',color='w',alpha=0.9)
0186     fig_text_TPC=ax.text(-60,134,70,  'Time Projection Chamber', size=10,color='w',alpha=0.9)
0187     fig_text_collisions=ax.text(-130,110,90,  'Diffused laser with TPC HV on and Laser 98% ALL ON Trigger modebit@2', size=10,color='w',alpha=0.9)
0188     fig_text_frame=ax.text(-130,95,90,  '2023-07-06, Run 11011', size=10,color='w',alpha=0.9)
0189
0190     #fig_text_sPhenix=ax.text(-130,100,90,  'Au+Au, $\sqrt{s_{NN}}$ = 200 GeV', size=10,color='w',alpha=0.9)
0191
0192
0193     # Providing the starting angle for the view
0194     ax.view_init(10,70,0,'y')
0195
0196     # Initializing scatters
0197     scatters = [ ax.scatter([100],[-100],[100]) for i in range(data.shape[0])]
0198     for i in range(data.shape[0]): scatters[i].set_color('black')
0199     print("Plot initialized")
0200     #start = timer()
0201     ani = animation.FuncAnimation(fig, animate_scatters, iterations, fargs=(data, scatters,fig_text,time_scale,iteration_time,start_iteration),
0202                                        interval=iteration_time, blit=False, repeat=False) #interval is in milliseconds and is the time between each frame
0203     if save:
0204         print("Saving animation as"+animation_name)
0205         ani.save(animation_name,writer='ffmpeg')
0206         print("Animation saved")
0207     #end = timer()
0208     #print("Time for process=")
0209     #print(end-start)
0210     plt.show()
0211
0212 def read_cluster_pos(inFile):
0213     if(inFile.lower().endswith('.json')):
0214         print("Reading data from json file")
0215         file=open(inFile)
0216         data_json=json.load(file)
0217         file.close()
0218         dict_json = data_json.items()
0219         numpy_array_json = np.array(list(dict_json))
0220         #print(numpy_array_json[2][1]['TRACKHITS'])
0221         mom_dict=numpy_array_json[2][1] #only works for this format the 3rd row is TRACKS and and 1 refers to INNERTRACKER
0222         #print(mom_dict)
0223         innertracker_arr=mom_dict['TRACKHITS'] #stores the tracks information as an array
0224         #print(innertracker_arr[0])
0225         #print(innertracker_arr)
0226         print("Reading clusters")
0227         #Getting the cluster positions of a track
0228         x_y_z_clusters=np.array([])
0229         #print(x_y_z_clusters)
0230         no_hits=len(innertracker_arr) #set no_tracks=10 for testing
0231         print(no_hits)
0232         #no_hits=1000
0233         for hit_no in range(no_hits):
0234             x_y_z_dict_track=innertracker_arr[hit_no].copy() #innertracker_arr[i] reads the ith track
0235             #print(list(x_y_z_dict_track.values()))
0236             x_y_z_clusters_track=np.array(list(x_y_z_dict_track.values())) #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
0237             x_y_z_clusters_track=x_y_z_clusters_track[:3]
0238             #print(x_y_z_clusters_track)
0239             #hit_no=0
0240             if(raddist_cluster(x_y_z_clusters_track)<30 or x_y_z_clusters_track[2]==0.0):
0241                continue
0242             else:
0243                if(hit_no==0):
0244                 x_y_z_clusters=[np.copy(x_y_z_clusters_track)]
0245
0246                else:
0247                 #if(hit_no==1):
0248                 #x_y_z_clusters=np.append(x_y_z_clusters,x_y_z_clusters_track,axis=0)
0249                 #x_y_z_clusters=np.append(x_y_z_clusters,x_y_z_clusters_track)
0250                   x_y_z_clusters=np.vstack([x_y_z_clusters,x_y_z_clusters_track])
0251                 #np.vstack([a,l])
0252         #print(x_y_z_clusters)
0253         return x_y_z_clusters
0254
0255     if(inFile.lower().endswith('.root')):
0256         print("Reading data from root file")
0257         file = uproot.open(inFile)
0258         ntp_cluster_tree=file['hitTree']
0259         branches=ntp_cluster_tree.arrays(["x","y","z"])
0260         #branches=branches[~np.isnan(branches.gvt)]
0261         branches=branches[((branches.x)**2+(branches.y)**2)>900]
0262         #branches=branches[branches.layer>6]
0263         #branches=branches[branches.layer<55]
0264         #branches=branches[branches.event>=75]
0265         #branches=branches[branches.event<55]
0266
0267         print("Reading clusters")
0268         x_y_z_clusters_run=np.array([])
0269         len_events=len(np.unique(branches.event))
0270         len_events=200
0271         event_times=[0]
0272         event_times=np.append(event_times,np.random.poisson(mean_eventtime*1000,len_events-1))
0273         event_times=np.cumsum(event_times,dtype=float)
0274         for cluster in range(len(branches)):
0275             #gvt_event=event_times[int(branches[cluster]['event'])]
0276             gvt_event=1
0277             gvt_event=gvt_event*(10**(-6))*time_scale #Time in milliseconds scaled for animation
0278             gvt_event=gvt_event/(iteration_time) #this is the time in terms of iterations
0279             x_y_z_clusters_track=np.array([[branches[cluster]['x'], branches[cluster]['y'], branches[cluster]['z'],branches[cluster]['event'],gvt_event]])
0280             if(cluster==0):
0281                 x_y_z_clusters_run=np.copy(x_y_z_clusters_track)
0282             else:
0283                 x_y_z_clusters_run=np.append(x_y_z_clusters_run,x_y_z_clusters_track,axis=0)
0284         return x_y_z_clusters_run
0285
0286 print("Generating data for animation")
0287
0288
0289 # Main Program starts from here
0290 np.random.seed(3)
0291 #User defined values
0292 time_scale=4.0*(10.0**5) #inverse of speed scale
0293 collision_rate=4.0 #MHz #Set fig_text_sPhenix in line 188
0294 mean_eventtime=1/collision_rate
0295 print("Mean event time (microseconds)=")
0296 print(mean_eventtime)
0297 TPC_drift_speed=8.0*(10.0**3) #Actual TPC drift speed =8cm/microsecond=8*10^3cm/millisecond
0298 iteration_time=100.0 #actual duration between frames in FuncAnimation
0299 len_TPC=105.0
0300
0301 drift_speed_posz=np.array([0.0,0.0,TPC_drift_speed/time_scale*iteration_time,0.0,0.0])
0302 drift_speed_negz=np.array([0.0,0.0,-TPC_drift_speed/time_scale*iteration_time,0.0,0.0])
0303 print("Drift_speed_posz(cm/iteration)=")
0304 print(drift_speed_posz)
0305 print("Drift_speed_negz(cm/iteration)=")
0306 print(drift_speed_negz)
0307
0308 #ANIMATION
0309 #data=read_cluster_pos("Data_files/G4sPHENIX_g4svtx_eval.root")
0310 data=read_cluster_pos("Data_files/TPCEventDisplay_11011.json")
0311 print(data)
0312 # Number of iterations
0313 #no_events=np.max(data[:,3])+1
0314 no_events=1
0315 print("no_events=")
0316 print(no_events)
0317
0318 iterations = int((no_events-1)*mean_eventtime*time_scale/(iteration_time*1000)+len_TPC/drift_speed_posz[2])+5
0319 #iterations=10
0320
0321 print("number of iterations=")
0322 print(iterations)
0323 print("Animation starting!")
0324
0325 #Saving takes a long time so use Save=True only when necessary
0326 #increase drift_speed_posz and drift_speed_negz if desired
0327 animate_clusters(data,"Animated_clusters_TPC_beam.mp4",save=True,start_iteration=50,no_iterations=iterations)
0328
0329
0330 #Merge mp4 files using ffmpeg -f concat -safe 0 -i filelist.txt -c copy mergedVideo.mp4