Skip to content
Snippets Groups Projects

Truth reconstruction

Merged Truth reconstruction
truth_reconstruction into master
Merged Tooba Ali requested to merge truth_reconstruction into master
Compare
and
1 file
+ 99
− 118
Compare changes
  • Side-by-side
  • Inline
benchmarks/dis/analysis/truth_reconstruction.py
+ 99
− 118
@@ -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,24 +107,15 @@ 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)):
X = X_list[j]
Y = Y_list[j]
X_len = ak.count(X,axis=None)
Y_len = ak.count(Y,axis=None)
if X_len > Y_len:
F_boolean = np.ones_like(Y) == 1
else:
F_boolean = np.ones_like(X) == 1
X_s = np.array(ak.flatten(X[F_boolean]))
Y_s = np.array(ak.flatten(Y[F_boolean]))
if i == 0: #Momentum
for j in range(len(X_list)): #Repeat the following steps for each particle (pions,protons,electrons,neutrons,photons)
X_s = np.array(ak.flatten(X_list[j]))
Y_s = np.array(ak.flatten(Y_list[j]))
if i == 0: #Momentum ratio
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 +124,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 +133,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,18 +167,49 @@ 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
X_len = ak.count(X1,axis=None)
Y_len = ak.count(Y1,axis=None)
if X_len > Y_len:
F_boolean = np.ones_like(Y1) == 1
else:
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]))
###################################################################################################
#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)
M_s = np.array(ak.flatten(M_list[j]))
X_plot[j] = M_s
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_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:
@@ -188,7 +218,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,74 +244,25 @@ 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):
theta_mc_fil = ak.Array(theta_mc[simID][booll])[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')