Update benchmarks/dis/analysis/truth_reconstruction.py

benchmarks/dis/analysis/truth_reconstruction.py

@@ -168,49 +168,12 @@ for i in range(len(MC_list)): #Repeat the following steps for each variable (mom
     plt.savefig(os.path.join(r_path, '%s_%s_%s.png' %  (title_list[i],title,config)))
     plt.close()
 
-###################################################################################################     #Correlation
 ###################################################################################################
-#####For momentum variable
-    X_s = np.array(ak.flatten(X1)) 
-    Y_s = np.array(ak.flatten(Y1)) 
-
-    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:    
-        h, xedges, yedges, image = plt.hist2d(x=X_s,y= Y_s,  bins = 11, range = [x_range,x_range]) 
-    plt.close()
-
-    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 = []
-    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
-        else:
-            norm_c = h[j]
-        norm_h.append(norm_c)
-        norm_c_text = [ '%.3f' % elem for elem in norm_c ] #display value to 3 dp
-        norm_h_text.append(norm_c_text)
-    
-    fig, axs = plt.subplots(1, 2, figsize=(20, 10))
-    if i == 0 and particle in particle_dict.keys(): #Momentum in Single events
-        axs[0].hist2d(x=X_s,y=Y_s, bins = 11)
-    else: 
-        axs[0].hist2d(x=X_s,y=Y_s, bins = 11,range = [x_range,x_range]) 
-    mplhep.hist2dplot(H=norm_h,norm=mpl.colors.LogNorm(vmin= 1e-4, vmax= 1),labels=norm_h_text, xbins = xedges, ybins = yedges, ax=axs[1])
-    
-    axs[0].set_title('%s Histogram'%(title_list[i]))
-    axs[0].set_xlabel('%s_mc'%(title_list[i]))
-    axs[0].set_ylabel('%s_rc'%(title_list[i]))
-    axs[1].set_xlabel('%s_mc'%(title_list[i]))
-    axs[1].set_ylabel('%s_rc'%(title_list[i]))
-    axs[1].set_title('%s Correlation'%(title_list[i]))
-    fig.suptitle('%s  %s events'%(config,Nevents))
-    plt.savefig(os.path.join(r_path, '%s_correlation_%s.png' %  (title_list[i],config)))
+     #Ratio vs momentum
+###################################################################################################
 
-#####For theta, phi, eta variable
-    if i > 0:
-        for j in range(len(X_list)):
+    if i > 0: #for each variable theta, phi, and eta
+        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]
             M_mc = M_list[j]
@@ -247,34 +210,67 @@ for i in range(len(MC_list)): #Repeat the following steps for each variable (mom
         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
+###################################################################################################
+     #Correlation
+###################################################################################################
+
+    #Repeat the following steps for each variable (momentum,theta,phi,eta)
+    X_s = np.array(ak.flatten(X1)) 
+    Y_s = np.array(ak.flatten(Y1)) 
+    #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:    
+        h, xedges, yedges, image = plt.hist2d(x=X_s,y= Y_s,  bins = 11, range = [x_range,x_range]) 
+    plt.close()
+
+    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 = []
+    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
+        else:
+            norm_c = h[j]
+        norm_h.append(norm_c)
+        norm_c_text = [ '%.3f' % elem for elem in norm_c ] #display value to 3 dp
+        norm_h_text.append(norm_c_text)
+    
+    fig, axs = plt.subplots(1, 2, figsize=(20, 10))
+    if i == 0 and particle in particle_dict.keys(): #Momentum in Single events
+        axs[0].hist2d(x=X_s,y=Y_s, bins = 11)
+    else: 
+        axs[0].hist2d(x=X_s,y=Y_s, bins = 11,range = [x_range,x_range]) 
+    mplhep.hist2dplot(H=norm_h,norm=mpl.colors.LogNorm(vmin= 1e-4, vmax= 1),labels=norm_h_text, xbins = xedges, ybins = yedges, ax=axs[1])
+    
+    axs[0].set_title('%s Histogram'%(title_list[i]))
+    axs[0].set_xlabel('%s_mc'%(title_list[i]))
+    axs[0].set_ylabel('%s_rc'%(title_list[i]))
+    axs[1].set_xlabel('%s_mc'%(title_list[i]))
+    axs[1].set_ylabel('%s_rc'%(title_list[i]))
+    axs[1].set_title('%s Correlation'%(title_list[i]))
+    fig.suptitle('%s  %s events'%(config,Nevents))
+    plt.savefig(os.path.join(r_path, '%s_correlation_%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]
     phi_rc_fil = ak.Array(phi_rc[recID][booll])[boolean_particle]
-
-    theta_mc_fil_len = ak.count(theta_mc_fil,axis=None)
-    theta_rc_fil_len = ak.count(theta_rc_fil,axis=None)
-    if theta_mc_fil_len > theta_rc_fil_len:
-        F_boolean = np.ones_like(theta_rc_fil) == 1
-    else: 
-        F_boolean = np.ones_like(theta_mc_fil) == 1
-    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]))
-    phi_rc_F = np.array(ak.flatten(phi_rc_fil[F_boolean]))
-    ratio = np.array((ak.Array(theta_rc_F)-(ak.Array(theta_mc_F))))
+    ratio = np.array((ak.Array(theta_rc_fil)-(ak.Array(theta_mc_fil))))
 
     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')