Truth reconstruction

Updates to benchmarks/dis/analysis/truth_reconstruction.py

benchmarks/dis/analysis/truth_reconstruction.py

@@ -18,7 +18,7 @@ kwargs = vars(args)
 rec_file = args.rec_file
 config = args.config
 Nevents = int(args.nevents)
-r_path = args.results_path + '/truth_reconstruction/'
+r_path = args.results_path + '/truth_reconstruction/' #Path for output figures and file.
 
 for array in ur.iterate(rec_file + ':events',['MCParticles/MCParticles.generatorStatus',
                                           'MCParticles/MCParticles.PDG',
@@ -31,7 +31,7 @@ for array in ur.iterate(rec_file + ':events',['MCParticles/MCParticles.generator
                                           'ReconstructedParticles/ReconstructedParticles.momentum.z',
                                           'ReconstructedParticlesAssoc/ReconstructedParticlesAssoc.simID',
                                           'ReconstructedParticlesAssoc/ReconstructedParticlesAssoc.recID',],step_size=Nevents):
-    PDG_mc = array['MCParticles/MCParticles.PDG']
+    PDG_mc = array['MCParticles/MCParticles.PDG']  #Monte Carlo (MC) particle numbering scheme.
     px_mc = array['MCParticles/MCParticles.momentum.x']
     py_mc = array['MCParticles/MCParticles.momentum.y']
     pz_mc = array['MCParticles/MCParticles.momentum.z']
@@ -40,51 +40,61 @@ for array in ur.iterate(rec_file + ':events',['MCParticles/MCParticles.generator
     py_rc = array['ReconstructedParticles/ReconstructedParticles.momentum.y']
     pz_rc = array['ReconstructedParticles/ReconstructedParticles.momentum.z']
     simID = array['ReconstructedParticlesAssoc/ReconstructedParticlesAssoc.simID']
-    recID = array['ReconstructedParticlesAssoc/ReconstructedParticlesAssoc.recID']
+    recID = array['ReconstructedParticlesAssoc/ReconstructedParticlesAssoc.recID'] 
+    #SimID and recID contain the indices of the MCParticles and ReconstructedParticles entry for that event.
 
+### MCParticles Variables
 momentum_mc = np.sqrt(((px_mc**2)+(py_mc**2)+(pz_mc**2)))
 theta_mc = np.arctan2(np.sqrt(px_mc**2+py_mc**2), pz_mc)
 phi_mc = np.arctan2(py_mc, px_mc)
-
+### ReconstructedParticles Variables
 momentum_rc = np.sqrt(((px_rc**2)+(py_rc**2)+(pz_rc**2)))
 theta_rc = np.arctan2(np.sqrt(px_rc**2+py_rc**2), pz_rc)
 phi_rc = np.arctan2(py_rc, px_rc)
 
-booll = (PDG_mc[simID])==(PDG_rc[recID])
-boolean_pion = np.logical_or(ak.Array(PDG_mc[simID][booll])==-211, ak.Array(PDG_mc[simID][booll])==+211)
-boolean_proton = np.logical_or(ak.Array(PDG_mc[simID][booll])==-2212, ak.Array(PDG_mc[simID][booll])==+2212)
-boolean_electron = ak.Array(PDG_mc[simID][booll])==11
-boolean_neutron = ak.Array(PDG_mc[simID][booll])==2112
-boolean_photon = ak.Array(PDG_mc[simID][booll])==22
-
-MC_list = [ak.Array(momentum_mc[simID][booll]),
-           ak.Array(theta_mc[simID][booll]),
-           ak.Array(phi_mc[simID][booll]),
-           -np.log(np.tan((ak.Array(theta_mc[simID][booll]))/2))]
-RC_list = [ak.Array(momentum_rc[recID][booll]),
-           ak.Array(theta_rc[recID][booll]),
-           ak.Array(phi_rc[recID][booll]),
-           -np.log(np.tan((ak.Array(theta_rc[recID][booll]))/2))]
+booll = (PDG_mc[simID])==(PDG_rc[recID]) #boolean that allows events where the same particle is reconstructed
+boolean_pion = np.logical_or(ak.Array(PDG_mc[simID][booll])==-211, ak.Array(PDG_mc[simID][booll])==+211) #boolean that allows events involving pions
+boolean_proton = np.logical_or(ak.Array(PDG_mc[simID][booll])==-2212, ak.Array(PDG_mc[simID][booll])==+2212) #boolean that allows events involving protons
+boolean_electron = ak.Array(PDG_mc[simID][booll])==11 #boolean that allows events involving electrons
+boolean_neutron = ak.Array(PDG_mc[simID][booll])==2112 #boolean that allows events involving neutrons
+boolean_photon = ak.Array(PDG_mc[simID][booll])==22 #boolean that allows events involving photons
+
+### MCParticles variables list
+MC_list = [ak.Array(momentum_mc[simID][booll]),                     #Momentum
+           ak.Array(theta_mc[simID][booll]),                        #Theta
+           ak.Array(phi_mc[simID][booll]),                          #Phi
+           -np.log(np.tan((ak.Array(theta_mc[simID][booll]))/2))]   #Eta
+### ReconstructedParticles variables list
+RC_list = [ak.Array(momentum_rc[recID][booll]),                     #Momentum
+           ak.Array(theta_rc[recID][booll]),                        #Theta
+           ak.Array(phi_rc[recID][booll]),                          #Phi
+           -np.log(np.tan((ak.Array(theta_rc[recID][booll]))/2))]   #Eta
+title_list = ['Momentum','Theta','Phi','Eta']
+### MC Momentum for different particles list
 M_list = [ak.Array(momentum_mc[simID][booll]),
           ak.Array(momentum_mc[simID][booll][boolean_pion]),
           ak.Array(momentum_mc[simID][booll][boolean_proton]),
           ak.Array(momentum_mc[simID][booll][boolean_electron]),
           ak.Array(momentum_mc[simID][booll][boolean_neutron]),
           ak.Array(momentum_mc[simID][booll][boolean_photon])]
-title_list = ['Momentum','Theta','Phi','Eta']
 
+#Marker Size in plots
 if Nevents == 100:
     ssize = 1
 else:
     ssize = 0.01
 
+#Particle type for Single events
 particle = config.split('-')[0].strip()
 particle_dict = {'e':[boolean_electron,'Electrons'],'pi':[boolean_pion,'Pions']}
-###################
 
-for i in range(len(MC_list)):
-    X1 = MC_list[i]
-    Y1 = RC_list[i]
+####################################################################################################
+    #Ratio 
+####################################################################################################
+
+for i in range(len(MC_list)): #Repeat the following steps for each variable (momentum,theta,phi,eta)
+    X1 = MC_list[i] #MCParticles events to be plotted on x-axis
+    Y1 = RC_list[i] #ReconstructedParticles events
     X_list = [ak.Array(X1),
           ak.Array(X1[boolean_pion]),
           ak.Array(X1[boolean_proton]),
@@ -97,11 +107,10 @@ for i in range(len(MC_list)):
           ak.Array(Y1[boolean_electron]),
           ak.Array(Y1[boolean_neutron]),
           ak.Array(Y1[boolean_photon])]
-
     X_plot = list(np.zeros(len(X_list)))
     Y_plot = list(np.zeros(len(X_list)))
 
-    for j in range(len(X_list)):
+    for j in range(len(X_list)): #Repeat the following steps for each particle (pions,protons,electrons,neutrons,photons)
         X = X_list[j]
         Y = Y_list[j]
         X_len = ak.count(X,axis=None)
@@ -110,11 +119,11 @@ for i in range(len(MC_list)):
             F_boolean = np.ones_like(Y) == 1
         else: 
             F_boolean = np.ones_like(X) == 1
-        X_s = np.array(ak.flatten(X[F_boolean])) 
+        X_s = np.array(ak.flatten(X[F_boolean])) #Filtered lists
         Y_s = np.array(ak.flatten(Y[F_boolean]))
         if i == 0:   #Momentum
             ratio = np.array((ak.Array(Y_s)/ak.Array(X_s)))
-        else:
+        else: #Angle difference
             ratio = np.array((ak.Array(Y_s)-(ak.Array(X_s))))
         X_plot[j] = X_s
         Y_plot[j] = ratio
@@ -123,7 +132,7 @@ for i in range(len(MC_list)):
     gs = fig.add_gridspec(3, 2, wspace=0)
     (ax1, ax2), (ax3, ax4),(ax5, ax6) = gs.subplots(sharex=True, sharey=True)
     # fig.suptitle('')
-    if i == 1:   #theta
+    if i == 1:   # for theta
         X_plot[0],X_plot[1],X_plot[2],X_plot[3],X_plot[4],X_plot[5] = -X_plot[0],-X_plot[1],-X_plot[2],-X_plot[3],-X_plot[4],-X_plot[5]
     ax1.scatter(X_plot[0], Y_plot[0], s = ssize)
     ax2.scatter(X_plot[1], Y_plot[1], s = ssize)
@@ -132,17 +141,15 @@ for i in range(len(MC_list)):
     ax5.scatter(X_plot[4], Y_plot[4], s = ssize)
     ax6.scatter(X_plot[5], Y_plot[5], s = ssize)
 
-    ax_list = [ax1,ax2,ax3,ax4,ax5]
-    if i == 0:
-        ax1.set_ylabel('rc/mc')
+    if i == 0:  # for momentum
+        ax1.set_ylabel('rc/mc') #ratio
         ax3.set_ylabel('rc/mc')
         ax5.set_ylabel('rc/mc') 
         title ='ratio'
-        for ax in ax_list:
-            ax.set_yscale('log')
-            ax.set_xscale('log')
-    else:
-        ax1.set_ylabel('rc-mc')
+        ax1.set_yscale('log')
+        ax1.set_xscale('log')
+    else:       # for angles
+        ax1.set_ylabel('rc-mc') #difference
         ax3.set_ylabel('rc-mc')
         ax5.set_ylabel('rc-mc') 
         title ='difference'
@@ -168,9 +175,54 @@ for i in range(len(MC_list)):
     ax1.set_title('%s %s  %s  %s events'%(title_list[i],title,config,Nevents))
     plt.savefig(os.path.join(r_path, '%s_%s_%s.png' %  (title_list[i],title,config)))
     plt.close()
-    ############
-    
-    #Correlation
+
+###################################################################################################
+     #Ratio vs momentum
+###################################################################################################
+
+    if i > 0: #for each variable theta, phi, and eta
+        for j in range(len(M_list)): #Repeat the following steps for each particle (pions,protons,electrons,neutrons,photons)
+            X = X_list[j]
+            Y = Y_list[j]
+            M_mc = M_list[j]
+            boolean_M = np.ones_like(M_mc) == 1
+            X_s = np.array(ak.flatten(X[boolean_M])) 
+            Y_s = np.array(ak.flatten(Y[boolean_M])) 
+            M_s = np.array(ak.flatten(M_mc))
+            ratio = np.array((ak.Array(Y_s)-(ak.Array(X_s))))
+            X_plot[j] = M_s
+            Y_plot[j] = ratio
+
+        fig = plt.figure()
+        gs = fig.add_gridspec(3, 2, wspace=0)
+        (ax1, ax2), (ax3, ax4),(ax5, ax6) = gs.subplots(sharex=True, sharey=True)
+        ax1.scatter(X_plot[0], Y_plot[0], s = ssize)
+        ax2.scatter(X_plot[1], Y_plot[1], s = ssize)
+        ax3.scatter(X_plot[2], Y_plot[2], s = ssize)
+        ax4.scatter(X_plot[3], Y_plot[3], s = ssize)
+        ax5.scatter(X_plot[4], Y_plot[4], s = ssize)
+        ax6.scatter(X_plot[5], Y_plot[5], s = ssize)
+        ax1.set_xscale('log')
+        ax1.set_ylabel('rc-mc')
+        ax3.set_ylabel('rc-mc')
+        ax5.set_ylabel('rc-mc')
+        ax5.set_xlabel('Momentum mc')
+        ax6.set_xlabel('Momentum mc')
+        ax2.set_title('Pions')
+        ax3.set_title('Protons')
+        ax4.set_title('Electrons')
+        ax5.set_title('Neutrons')
+        ax6.set_title('Photons')
+        fig.set_figwidth(20)
+        fig.set_figheight(10)
+        ax1.set_title('%s Difference Vs Momentum  %s  %s events'%(title_list[i],config,Nevents))
+        plt.savefig(os.path.join(r_path, '%s_difference_vs_momentum_%s.png' %  (title_list[i],config)))
+
+###################################################################################################
+     #Correlation
+###################################################################################################
+
+    #Repeat the following steps for each variable (momentum,theta,phi,eta)
     X_len = ak.count(X1,axis=None)
     Y_len = ak.count(Y1,axis=None)
     if X_len > Y_len: 
@@ -179,7 +231,8 @@ for i in range(len(MC_list)):
         F_boolean = np.ones_like(X1) == 1
     X_s = np.array(ak.flatten(X1[F_boolean])) 
     Y_s = np.array(ak.flatten(Y1[F_boolean])) 
-
+ 
+    #Histogram
     if i == 0 and particle in particle_dict.keys(): #Momentum in Single events
         h, xedges, yedges, image = plt.hist2d(x=X_s,y= Y_s,  bins = 11) 
     else:    
@@ -188,7 +241,7 @@ for i in range(len(MC_list)):
 
     col_sum = ak.sum(h,axis=-1) #number of events in each (verticle) column 
     norm_h = [] #norm_h is the normalized matrix
-    norm_h_text = [] #display labels matrix
+    norm_h_text = []
     for j in range(len(col_sum)):
         if col_sum[j] != 0:
             norm_c = h[j]/col_sum[j] #normalized column = column values divide by sum of the column
@@ -214,50 +267,12 @@ for i in range(len(MC_list)):
     fig.suptitle('%s  %s events'%(config,Nevents))
     plt.savefig(os.path.join(r_path, '%s_correlation_%s.png' %  (title_list[i],config)))
 
-###############
-
-    if i > 0:
-        for j in range(len(X_list)):
-            X = X_list[j]
-            Y = Y_list[j]
-            M_mc = M_list[j]
-            boolean_M = np.ones_like(M_mc) == 1
-            X_s = np.array(ak.flatten(X[boolean_M])) 
-            Y_s = np.array(ak.flatten(Y[boolean_M])) 
-            M_s = np.array(ak.flatten(M_mc))
-            ratio = np.array((ak.Array(Y_s)-(ak.Array(X_s))))
-            X_plot[j] = M_s
-            Y_plot[j] = ratio
-
-        fig = plt.figure()
-        gs = fig.add_gridspec(3, 2, wspace=0)
-        (ax1, ax2), (ax3, ax4),(ax5, ax6) = gs.subplots(sharex=True, sharey=True)
-        # fig.suptitle('')
-        ax1.scatter(X_plot[0], Y_plot[0], s = ssize)
-        ax2.scatter(X_plot[1], Y_plot[1], s = ssize)
-        ax3.scatter(X_plot[2], Y_plot[2], s = ssize)
-        ax4.scatter(X_plot[3], Y_plot[3], s = ssize)
-        ax5.scatter(X_plot[4], Y_plot[4], s = ssize)
-        ax6.scatter(X_plot[5], Y_plot[5], s = ssize)
-        ax1.set_xscale('log')
-        ax1.set_ylabel('rc-mc')
-        ax3.set_ylabel('rc-mc')
-        ax5.set_ylabel('rc-mc')
-        ax5.set_xlabel('Momentum mc')
-        ax6.set_xlabel('Momentum mc')
-        ax2.set_title('Pions')
-        ax3.set_title('Protons')
-        ax4.set_title('Electrons')
-        ax5.set_title('Neutrons')
-        ax6.set_title('Photons')
-        fig.set_figwidth(20)
-        fig.set_figheight(10)
-        ax1.set_title('%s Difference Vs Momentum  %s  %s events'%(title_list[i],config,Nevents))
-        plt.savefig(os.path.join(r_path, '%s_difference_vs_momentum_%s.png' %  (title_list[i],config)))
-
-################
+###################################################################################################
+     #Phi vs Theta plots
+###################################################################################################
 
 def particle_plots(boolean_particle):
+    #filtered lists w.r.t the particle
     theta_mc_fil = ak.Array(theta_mc[simID][booll])[boolean_particle]
     theta_rc_fil = ak.Array(theta_rc[recID][booll])[boolean_particle]
     phi_mc_fil = ak.Array(phi_mc[simID][booll])[boolean_particle]
@@ -269,6 +284,7 @@ def particle_plots(boolean_particle):
         F_boolean = np.ones_like(theta_rc_fil) == 1
     else: 
         F_boolean = np.ones_like(theta_mc_fil) == 1
+    #filtered lists w.r.t length
     theta_mc_F = np.array(ak.flatten(theta_mc_fil[F_boolean]))
     theta_rc_F = np.array(ak.flatten(theta_rc_fil[F_boolean]))
     phi_mc_F = np.array(ak.flatten(phi_mc_fil[F_boolean]))
@@ -278,10 +294,10 @@ def particle_plots(boolean_particle):
     fig = plt.figure()
     gs = fig.add_gridspec(2, 2, wspace=0.01)
     (ax1, ax2), (ax3, ax4) = gs.subplots(sharex=True, sharey=True)
-    ax1.scatter(-theta_mc_F, ratio, s = ssize)
-    ax2.scatter(-theta_rc_F, ratio, s = ssize)
-    ax3.scatter(-theta_mc_F, phi_mc_F, s = ssize)
-    ax4.scatter(-theta_rc_F, phi_rc_F, s = ssize)
+    ax1.scatter(-theta_mc_fil, ratio, s = ssize)
+    ax2.scatter(-theta_rc_fil, ratio, s = ssize)
+    ax3.scatter(-theta_mc_fil, phi_mc_fil, s = ssize)
+    ax4.scatter(-theta_rc_fil, phi_rc_fil, s = ssize)
     ax1.set_ylabel('rc-mc')
     ax2.set_ylabel('rc-mc')
     ax3.set_ylabel('Phi mc')