remove same_length_lists() function in benchmarks/dis/analysis/truth_reconstruction.py

benchmarks/dis/analysis/truth_reconstruction.py

@@ -53,31 +53,31 @@ 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 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
+boolean_pion = np.logical_or(ak.Array(PDG_mc[simID])==-211, ak.Array(PDG_mc[simID])==+211) #boolean that allows events involving pions
+boolean_proton = np.logical_or(ak.Array(PDG_mc[simID])==-2212, ak.Array(PDG_mc[simID])==+2212) #boolean that allows events involving protons
+boolean_electron = ak.Array(PDG_mc[simID])==11 #boolean that allows events involving electrons
+boolean_neutron = ak.Array(PDG_mc[simID])==2112 #boolean that allows events involving neutrons
+boolean_photon = ak.Array(PDG_mc[simID])==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
+MC_list = [ak.Array(momentum_mc[simID]),                     #Momentum
+           ak.Array(theta_mc[simID]),                        #Theta
+           ak.Array(phi_mc[simID]),                          #Phi
+           -np.log(np.tan((ak.Array(theta_mc[simID]))/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
+RC_list = [ak.Array(momentum_rc[recID]),                     #Momentum
+           ak.Array(theta_rc[recID]),                        #Theta
+           ak.Array(phi_rc[recID]),                          #Phi
+           -np.log(np.tan((ak.Array(theta_rc[recID]))/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])]
+M_list = [np.array(ak.flatten(ak.Array(momentum_mc[simID]))),
+          np.array(ak.flatten(ak.Array(momentum_mc[simID][boolean_pion]))),
+          np.array(ak.flatten(ak.Array(momentum_mc[simID][boolean_proton]))),
+          np.array(ak.flatten(ak.Array(momentum_mc[simID][boolean_electron]))),
+          np.array(ak.flatten(ak.Array(momentum_mc[simID][boolean_neutron]))),
+          np.array(ak.flatten(ak.Array(momentum_mc[simID][boolean_photon])))]
+momentum_range = (min(M_list[0]),max(M_list[0]))
 
 #Marker Size in plots
 if Nevents == 100:
@@ -116,52 +116,41 @@ def error_bars(plot_x, plot_y, x_axis_range):
     bin_HalfWidth = (xedges[1:] - xedges[:-1])/2
     return bin_Midpoint, mean, bin_HalfWidth, std, min_std
     
-def same_length_lists(plot_x, plot_y):
-    X_length = ak.count(plot_x,axis=None)
-    Y_length = ak.count(plot_y,axis=None)
-    if X_length > Y_length: 
-        F_boolean = np.ones_like(plot_y) == 1
-    else: 
-        F_boolean = np.ones_like(plot_x) == 1
-    return F_boolean
 
+boolean_tilt = list(np.zeros(len(M_list)))
 for i in range(len(MC_list)): #Repeat the following steps for each variable (momentum,theta,phi,eta)
     MCparts = MC_list[i] #MCParticles events to be plotted on x-axis
     RCparts = RC_list[i] #ReconstructedParticles events
-    X_list = [ak.Array(MCparts),
-          ak.Array(MCparts[boolean_pion]),
-          ak.Array(MCparts[boolean_proton]),
-          ak.Array(MCparts[boolean_electron]),
-          ak.Array(MCparts[boolean_neutron]),
-          ak.Array(MCparts[boolean_photon])]
-    Y_list = [ak.Array(RCparts),
-          ak.Array(RCparts[boolean_pion]),
-          ak.Array(RCparts[boolean_proton]),
-          ak.Array(RCparts[boolean_electron]),
-          ak.Array(RCparts[boolean_neutron]),
-          ak.Array(RCparts[boolean_photon])]
+    X_list = [np.array(ak.flatten(MCparts)),
+          np.array(ak.flatten(MCparts[boolean_pion])),
+          np.array(ak.flatten(MCparts[boolean_proton])),
+          np.array(ak.flatten(MCparts[boolean_electron])),
+          np.array(ak.flatten(MCparts[boolean_neutron])),
+          np.array(ak.flatten(MCparts[boolean_photon]))]
+    Y_list = [np.array(ak.flatten(RCparts)),
+          np.array(ak.flatten(RCparts[boolean_pion])),
+          np.array(ak.flatten(RCparts[boolean_proton])),
+          np.array(ak.flatten(RCparts[boolean_electron])),
+          np.array(ak.flatten(RCparts[boolean_neutron])),
+          np.array(ak.flatten(RCparts[boolean_photon]))]
     X_plot = list(np.zeros(len(X_list)))
     Y_plot = list(np.zeros(len(X_list)))
     Y_error = list(np.zeros(len(X_list)))
+    for j in range(len(X_list)): #Repeat the following steps for each particle (pions,protons,electrons,neutrons,photons)
+        if i == 0:   #Momentum
+            ratio = Y_list[j]/X_list[j]
+        else: #Angle difference
+            ratio = Y_list[j]-X_list[j]
+        Y_plot[j] = ratio
 
 
 ####################################################################################################
     #Ratio 
 ####################################################################################################
 
-    for j in range(len(X_list)): #Repeat the following steps for each particle (pions,protons,electrons,neutrons,photons)
-        F_boolean = same_length_lists(X_list[j],Y_list[j])
-        X_s = np.array(ak.flatten(X_list[j][F_boolean])) #Filtered lists
-        Y_s = np.array(ak.flatten(Y_list[j][F_boolean]))
-        if i == 0:   #Momentum
-            ratio = np.array((ak.Array(Y_s)/ak.Array(X_s)))
-        else: #Angle difference
-            ratio = np.array((ak.Array(Y_s)-(ak.Array(X_s))))
-        X_plot[j] = X_s
-        Y_plot[j] = ratio
-    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]
-    
+    X_plot = X_list
+    if i == 1:  #Theta
+        X_plot = -1*np.array(X_plot)
     particle_nlist = ['All','Pions','Protons','Electrons','Neutrons','Photons']
     for iterate in [0,1]:
         fig = plt.figure()
@@ -203,8 +192,6 @@ for i in range(len(MC_list)): #Repeat the following steps for each variable (mom
             ax5.set_xlabel('%s mc'%(title_list[i]))
             ax6.set_xlabel('%s mc'%(title_list[i]))
             x_range = list(ax1.get_xlim())
-            if i == 0:
-                momentum_range = x_range
         fig.set_figwidth(20)
         fig.set_figheight(10)
         
@@ -235,25 +222,14 @@ for i in range(len(MC_list)): #Repeat the following steps for each variable (mom
         else:
             ax1.set_title('%s %s  %s  %s events\n DETECTOR_CONFIG: %s'%(title_list[i],title,config,Nevents,Dconfig))
             plt.savefig(os.path.join(r_path, '%s_%s_%s.png' %  (title_list[i],title,config)))
-        plt.close()
-    
     
+
 ###################################################################################################
      #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
+        X_plot = M_list
 
         title ='difference vs momentum'    
         particle_nlist = ['All','Pions','Protons','Electrons','Neutrons','Photons']
@@ -298,7 +274,6 @@ for i in range(len(MC_list)): #Repeat the following steps for each variable (mom
             else:
                 ax1.set_title('%s Difference Vs Momentum  %s  %s events\n DETECTOR_CONFIG: %s'%(title_list[i],config,Nevents,Dconfig))
                 plt.savefig(os.path.join(r_path, '%s_difference_vs_momentum_%s.png' %  (title_list[i],config)))
-            plt.close()
         
 
 ###################################################################################################
@@ -306,15 +281,17 @@ for i in range(len(MC_list)): #Repeat the following steps for each variable (mom
 ###################################################################################################
 
     #Repeat the following steps for each variable (momentum,theta,phi,eta)
-    F_boolean = same_length_lists(MCparts,RCparts)
-    X_s = np.array(ak.flatten(MCparts[F_boolean])) #Filtered lists
-    Y_s = np.array(ak.flatten(RCparts[F_boolean])) 
+    X_s = X_list[0]
+    Y_s = Y_list[0]
 
     #Histogram
+    fig, axs = plt.subplots(1, 2, figsize=(20, 10))
     if i == 0 and particle in particle_dict.keys(): #Momentum in Single events
-        h, xedges, yedges = np.histogram2d(x=X_s,y= Y_s,  bins = 11) 
+        h, xedges, yedges = np.histogram2d(x=X_s,y= Y_s,  bins = 11)
+        axs[0].hist2d(x=X_s,y=Y_s, bins = 11)
     else:    
-        h, xedges, yedges = np.histogram2d(x=X_s,y= Y_s,  bins = 11, range = [x_range,x_range]) 
+        h, xedges, yedges = np.histogram2d(x=X_s,y= Y_s,  bins = 11, range = [x_range,x_range])
+        axs[0].hist2d(x=X_s,y=Y_s, bins = 11,range = [x_range,x_range])
 
     col_sum = ak.sum(h,axis=-1) #number of events in each (verticle) column 
     norm_h = [] #norm_h is the normalized matrix
@@ -328,11 +305,6 @@ for i in range(len(MC_list)): #Repeat the following steps for each variable (mom
         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_ylabel('%s_rc'%(title_list[i]))
@@ -342,8 +314,7 @@ for i in range(len(MC_list)): #Repeat the following steps for each variable (mom
     axs[1].set_title('%s Correlation'%(title_list[i]))
     fig.suptitle('%s  %s events\n DETECTOR_CONFIG: %s'%(config,Nevents,Dconfig))
     plt.savefig(os.path.join(r_path, '%s_correlation_%s.png' %  (title_list[i],config)))
-    plt.close()
-    
+
 
 ###################################################################################################
      #Phi vs Theta plots
@@ -351,29 +322,23 @@ for i in range(len(MC_list)): #Repeat the following steps for each variable (mom
 
 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_filtered = np.array(ak.flatten(ak.Array(theta_mc[simID])[boolean_particle]))
+    theta_rc_filtered = np.array(ak.flatten(ak.Array(theta_rc[recID])[boolean_particle]))
+    phi_mc_filtered = np.array(ak.flatten(ak.Array(phi_mc[simID])[boolean_particle]))
+    phi_rc_filtered = np.array(ak.flatten(ak.Array(phi_rc[recID])[boolean_particle]))
     
-    F_boolean = same_length_lists(theta_mc_fil, theta_rc_fil)
-    #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]))
-    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 = theta_rc_filtered-theta_mc_filtered
     x_range = [0,np.pi]
-    Y_error = [error_bars(theta_mc_F, ratio, x_range),error_bars(theta_rc_F, ratio, x_range)]
+    Y_error = [error_bars(theta_mc_filtered, ratio, x_range),error_bars(theta_rc_filtered, ratio, x_range)]
     fig = plt.figure()
     gs = fig.add_gridspec(3, 2, wspace=0, hspace = 0.3)
     (ax1, ax2), (ax3, ax4), (ax5, ax6) = gs.subplots(sharex=True, sharey='row')
-    ax1.scatter(-theta_mc_F, ratio, s = ssize)
-    ax2.scatter(-theta_rc_F, ratio, s = ssize)
-    ax3.scatter(-theta_mc_F, ratio, s = ssize)
-    ax4.scatter(-theta_rc_F, ratio, s = ssize)
-    ax5.scatter(-theta_mc_F, phi_mc_F, s = ssize)
-    ax6.scatter(-theta_rc_F, phi_rc_F, s = ssize)
+    ax1.scatter(-theta_mc_filtered, ratio, s = ssize)
+    ax2.scatter(-theta_rc_filtered, ratio, s = ssize)
+    ax3.scatter(-theta_mc_filtered, ratio, s = ssize)
+    ax4.scatter(-theta_rc_filtered, ratio, s = ssize)
+    ax5.scatter(-theta_mc_filtered, phi_mc_filtered, s = ssize)
+    ax6.scatter(-theta_rc_filtered, phi_rc_filtered, s = ssize)
     ax1.set_ylabel('Theta rc-mc')
     ax2.set_ylabel('Theta rc-mc')
     ax3.set_ylabel('Theta rc-mc')