From ae8ca2e57fb99f6da4a5c3676b4a2a903323f087 Mon Sep 17 00:00:00 2001
From: Tooba Ali <alit1@myumanitoba.ca>
Date: Fri, 5 May 2023 18:02:45 +0000
Subject: [PATCH] remove same_length_lists() function in
benchmarks/dis/analysis/truth_reconstruction.py
---
.../dis/analysis/truth_reconstruction.py | 163 +++++++-----------
1 file changed, 64 insertions(+), 99 deletions(-)
diff --git a/benchmarks/dis/analysis/truth_reconstruction.py b/benchmarks/dis/analysis/truth_reconstruction.py
index 38b0d066..33f426a5 100644
--- a/benchmarks/dis/analysis/truth_reconstruction.py
+++ b/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')
--
GitLab