dis/analysis/kinematics_correlations.py

0001 #!/usr/bin/env python
0002 # coding: utf-8
0003
0004 import os
0005 import numpy as np
0006 import uproot as ur
0007 import awkward as ak
0008 import matplotlib.pyplot as plt
0009 import matplotlib as mpl
0010 import mplhep
0011 import argparse
0012
0013
0014
0015 parser = argparse.ArgumentParser()
0016 parser.add_argument('--rec_file', type=str, help='Reconstructed track file.')
0017 parser.add_argument('--config', type=str, help='Momentum configuration.')
0018 parser.add_argument('--nevents', type=float, help='Number of events to process.')
0019 parser.add_argument('--results_path', type=str, help='Output directory.')
0020 args = parser.parse_args()
0021 kwargs = vars(args)
0022
0023 rec_file = args.rec_file
0024 Nevents = args.nevents
0025 r_path = args.results_path  #Path for output figures and file.
0026 Dconfig = 'epic' + args.config.split('_epic')[1].strip() #Detector config
0027 config = args.config.split('_epic')[0].strip()
0028 minq2 = int(config.split('=')[1].strip())
0029 k = int(config.split('x')[0].split('_')[1].strip())    #ebeam Electron beam energy
0030 p = int(config.split('x')[1].split('_')[0].strip())     #pbeam Proton (or ion) beam energy
0031
0032
0033 #logarithmically spaced bins in Q2 and 𝑥
0034 Q2_bins = [0.40938507,0.64786184,1.025353587,1.626191868,2.581181894,4.091148983,6.478618696,10.25353599,16.26191868,25.81181894,40.91148983,64.78618696,102.5353599,162.6191868,258.1181894,409.1148983,647.8618696,1025.353599,1626.191868,2581.181894,4091.148983,6482.897648]
0035 x_bins = [4.09385E-05,6.47862E-05,0.000102535,0.000162619,0.000258118,0.000409115,0.000647862,0.001025354,0.001626192,0.002581182,0.004091149,0.006478619,0.010253536,0.016261919,0.025811819,0.04091149,0.064786187,0.10253536,0.162619187,0.25811819,0.409114898,0.647861868,1.025257131]
0036
0037
0038 ##########################################################################################
0039 #function to construct Q2 correlation plots
0040 ##########################################################################################
0041
0042 def Q2correlation(minq2,method): #minq2 can be 1,10,100, or 1000; method can be 'e','DA', or 'JB'
0043
0044     Q2values_Y = method_Q2values_dict['{}'.format(method)]  #Q2 values of the given method, that are mapped onto the y axis
0045
0046     Q2_List_T = Q2values_T #Truth Q2 values, mapped along x axis
0047     Q2_List_Y = Q2values_Y #method (E/DA/JB) Q2 values, mapped along y axis
0048
0049     T_len = ak.count(Q2_List_T,axis=0) #total number of events in Truth
0050     Y_len = ak.count(Q2_List_Y,axis=0) #total number of events in method
0051
0052     if T_len > Y_len: #if total number of events for Truth is greater
0053         Y_boolean = ak.count(Q2_List_Y,axis=-1) >= 1 #boolean to filter ak.Arrays wrt single events in method
0054         Q2_List_T_F = Q2_List_T[Y_boolean] #filtered Truth Q2 values
0055         Q2_List_Y_F = Q2_List_Y[Y_boolean] #filtered method Q2 values
0056     else: #if total number of events for method is greater
0057         T_boolean = ak.count(Q2_List_T,axis=-1) >= 1 #boolean to filter ak.Arrays wrt single events in Truth
0058         Q2_List_T_F = Q2_List_T[T_boolean] #filtered Truth Q2 values
0059         Q2_List_Y_F = Q2_List_Y[T_boolean] #filtered method Q2 values
0060
0061     T_Q2s = np.array(ak.flatten(Q2_List_T_F)) #Truth Q2 values, mapped along x axis
0062     Y_Q2s = np.array(ak.flatten(Q2_List_Y_F)) #method Q2 values, mapped along y axis
0063
0064     #2-dimensional histogram, h
0065     h, xedges, yedges = np.histogram2d(x=T_Q2s,y=Y_Q2s, bins=[Q2_bins,Q2_bins])
0066
0067     minq2_dict = {'1':2,'10':7,'100':12,'1000':17} #Q2 bin index at which minq2 starts
0068     h[0:minq2_dict['{}'.format(minq2)]]=0 #ignore values before minq2
0069
0070     #normalization of h:
0071     col_sum = ak.sum(h,axis=-1) #number of events in each (verticle) column
0072     norm_h = [] #norm_h is the normalized matrix
0073     norm_h_text = [] #display labels matrix
0074     for i in range(len(col_sum)):
0075         if col_sum[i] != 0:
0076             norm_c = h[i]/col_sum[i] #normalized column = column values divide by sum of the column
0077         else:
0078             norm_c = h[i]
0079         norm_h.append(norm_c)
0080         norm_c_text = [ '%.3f' % elem for elem in norm_c ] #display value to 3 dp
0081         norm_h_text.append(norm_c_text)
0082
0083     fig = plt.figure()
0084     mplhep.hist2dplot(H=norm_h,norm=mpl.colors.LogNorm(vmin= 1e-4, vmax= 1),labels=norm_h_text,xbins=Q2_bins,ybins=Q2_bins)
0085     plt.yscale('log')
0086     plt.xscale('log')
0087     fig.set_figwidth(11)
0088     fig.set_figheight(11)
0089     plt.xlabel('$Q^2$ [$GeV^2$] Truth')
0090     plt.ylabel('$Q^2$ [$GeV^2$] {}'.format(method_dict['{}'.format(method)]))
0091     plt.title('{}   $Q^2$ correlation   {}x{}   $minQ^2=${}$GeV^2$\n  {} events  DETECTOR_CONFIG: {} '.format(method_dict['{}'.format(method)],k,p,minq2,Nevents,Dconfig))
0092     plt.savefig(os.path.join(r_path, 'Q2_correlation_%s_%s.png' %(method,config)))
0093
0094
0095 ##########################################################################################
0096 #function to construct Bjorken-x correlation plots
0097 ##########################################################################################
0098
0099 def Xcorrelation(minq2,method): #minq2 can be 1,10,100, or 1000; method can be 'e','DA', or 'JB'
0100
0101     Xvalues_Y = method_Xvalues_dict['{}'.format(method)] #x values of the given method, that are mapped onto the y axis
0102
0103     X_List_T = Xvalues_T #Truth x values, mapped along x axis
0104     X_List_Y = Xvalues_Y #method (E/DA/JB) x values, mapped along y axis
0105
0106     T_len = ak.count(X_List_T,axis=0) #total number of events in Truth
0107     Y_len = ak.count(X_List_Y,axis=0) #total number of events in method
0108
0109     if T_len > Y_len: #if total number of events for Truth is greater
0110         Y_boolean = ak.count(X_List_Y,axis=-1) >= 1 #boolean to filter ak.Arrays wrt single events in method
0111         X_List_T_F = X_List_T[Y_boolean] #filtered Truth x values
0112         X_List_Y_F = X_List_Y[Y_boolean] #filtered method x values
0113     else: #if total number of events for method is greater
0114         T_boolean = ak.count(X_List_T,axis=-1) >= 1 #boolean to filter ak.Arrays wrt single events in Truth
0115         X_List_T_F = X_List_T[T_boolean] #filtered Truth x values
0116         X_List_Y_F = X_List_Y[T_boolean] #filtered method x values
0117
0118     T_Xs = np.array(ak.flatten(X_List_T_F)) #Truth Bjorken-x values, mapped along x axis
0119     Y_Xs = np.array(ak.flatten(X_List_Y_F)) #method Bjorken-x values, mapped along y axis
0120
0121     T_x_bool = T_Xs>=minq2/(4*k*p) #boolean to filter x values that satisfy bjorken-x equation for minq2, ebeam and pbeam
0122     T_Xs = T_Xs[T_x_bool]
0123     Y_Xs = Y_Xs[T_x_bool]
0124
0125     #2-dimensional histogram, h
0126     h, xedges, yedges = np.histogram2d(x=T_Xs,y=Y_Xs, bins=[x_bins,x_bins])
0127
0128     #normalization of h:
0129     col_sum = ak.sum(h,axis=-1) #number of events in each (verticle) column
0130     norm_h = [] #norm_h is the normalized matrix
0131     norm_h_text = [] #display labels matrix
0132     for i in range(len(col_sum)):
0133         if col_sum[i] != 0:
0134             norm_c = h[i]/col_sum[i] #normalized column = column values divide by sum of the column
0135         else:
0136             norm_c = h[i]
0137         norm_h.append(norm_c)
0138         norm_c_text = [ '%.2f' % elem for elem in norm_c ] #display value to 2 dp
0139         norm_h_text.append(norm_c_text)
0140
0141     fig = plt.figure()
0142     mplhep.hist2dplot(H=norm_h,norm=mpl.colors.LogNorm(vmin= 1e-4, vmax= 1),labels=norm_h_text,xbins=x_bins,ybins=x_bins)
0143     plt.yscale('log')
0144     plt.xscale('log')
0145     fig.set_figwidth(11)
0146     fig.set_figheight(11)
0147     plt.xlabel('x Truth')
0148     plt.ylabel('$x$   {}'.format(method_dict['{}'.format(method)]))
0149     plt.title('{}   $x$ correlation   {}x{}   $minQ^2=${}$GeV^2$\n  {} events  DETECTOR_CONFIG: {} '.format(method_dict['{}'.format(method)],k,p,minq2,Nevents,Dconfig))
0150     plt.savefig(os.path.join(r_path, 'x_correlation_%s_%s.png' %(method,config)))
0151
0152
0153
0154 keys = ur.concatenate(rec_file + ':events/' + 'InclusiveKinematicsTruth')
0155 Truth =   [keys['InclusiveKinematicsTruth.Q2'],keys['InclusiveKinematicsTruth.x']]
0156 keys = ur.concatenate(rec_file + ':events/' + 'InclusiveKinematicsElectron')
0157 Electron =   [keys['InclusiveKinematicsElectron.Q2'], keys['InclusiveKinematicsElectron.x']]
0158 keys = ur.concatenate(rec_file + ':events/' + 'InclusiveKinematicsDA')
0159 DoubleAngle =  [keys['InclusiveKinematicsDA.Q2'], keys['InclusiveKinematicsDA.x']]
0160 keys = ur.concatenate(rec_file + ':events/' + 'InclusiveKinematicsJB')
0161 JacquetBlondel =  [keys['InclusiveKinematicsJB.Q2'], keys['InclusiveKinematicsJB.x']]
0162 keys = ur.concatenate(rec_file + ':events/' + 'InclusiveKinematicsSigma')
0163 Sigma =  [keys['InclusiveKinematicsSigma.Q2'], keys['InclusiveKinematicsSigma.x']]
0164 try:
0165     keys = ur.concatenate(rec_file + ':events/' + 'InclusiveKinematicsESigma')
0166     ESigma =  [keys['InclusiveKinematicsESigma.Q2'], keys['InclusiveKinematicsESigma.x']]
0167 except ur.KeyInFileError:
0168     # Legacy compatibility
0169     keys = ur.concatenate(rec_file + ':events/' + 'InclusiveKinematicseSigma')
0170     ESigma =  [keys['InclusiveKinematicseSigma.Q2'], keys['InclusiveKinematicseSigma.x']]
0171
0172 Q2values_T = Truth[0]
0173 Q2values_E = Electron[0]
0174 Q2values_DA = DoubleAngle[0]
0175 Q2values_JB = JacquetBlondel[0]
0176 Q2values_S = Sigma[0]
0177 Q2values_eS = ESigma[0]
0178 Xvalues_T = Truth[1]
0179 Xvalues_E = Electron[1]
0180 Xvalues_DA = DoubleAngle[1]
0181 Xvalues_JB = JacquetBlondel[1]
0182 Xvalues_S = Sigma[1]
0183 Xvalues_eS = ESigma[1]
0184
0185 method_dict = {'e':'Electron','DA':'Double-Angle','JB':'Jacquet-Blondel','S':'Sigma','eS':'eSigma'}
0186 method_Q2values_dict = {'e':Q2values_E,'DA':Q2values_DA,'JB':Q2values_JB,'S':Q2values_S,'eS':Q2values_eS}
0187 method_Xvalues_dict = {'e':Xvalues_E,'DA':Xvalues_DA,'JB':Xvalues_JB,'S':Xvalues_S,'eS':Xvalues_eS}
0188
0189 Q2correlation(minq2,'e')
0190 Xcorrelation(minq2,'e')
0191 Q2correlation(minq2,'DA')
0192 Xcorrelation(minq2,'DA')
0193 Q2correlation(minq2,'JB')
0194 Xcorrelation(minq2,'JB')
0195 Q2correlation(minq2,'S')
0196 Xcorrelation(minq2,'S')
0197 Q2correlation(minq2,'eS')
0198 Xcorrelation(minq2,'eS')
0199
0200
0201 ##########################################################################################
0202 #Kinematic coverage plot
0203 ##########################################################################################
0204
0205 fig = plt.figure(figsize=(20,10))
0206 gs = fig.add_gridspec(2, 3, wspace=0.09, hspace = 0.2)
0207 (ax1, ax2, ax3), (ax4, ax5, ax6) = gs.subplots()
0208
0209 ax1.hist2d(ak.to_numpy(ak.flatten(Xvalues_T)), ak.to_numpy(ak.flatten(Q2values_T)), bins = [x_bins, Q2_bins])
0210 ax2.hist2d(ak.to_numpy(ak.flatten(Xvalues_S)), ak.to_numpy(ak.flatten(Q2values_S)), bins = [x_bins, Q2_bins])
0211 ax3.hist2d(ak.to_numpy(ak.flatten(Xvalues_DA)), ak.to_numpy(ak.flatten(Q2values_DA)), bins = [x_bins, Q2_bins])
0212 ax4.hist2d(ak.to_numpy(ak.flatten(Xvalues_E)), ak.to_numpy(ak.flatten(Q2values_E)), bins = [x_bins, Q2_bins])
0213 ax5.hist2d(ak.to_numpy(ak.flatten(Xvalues_eS)), ak.to_numpy(ak.flatten(Q2values_eS)), bins = [x_bins, Q2_bins])
0214 ax6.hist2d(ak.to_numpy(ak.flatten(Xvalues_JB)), ak.to_numpy(ak.flatten(Q2values_JB)), bins = [x_bins, Q2_bins])
0215
0216 for axes in [ax1,ax2,ax3,ax4,ax5,ax6]:
0217     axes.set_xscale('log')
0218     axes.set_yscale('log')
0219     axes.set_ylabel('$Q^2$ [$GeV^2$]')
0220     axes.set_xlabel('$x$')
0221 ax1.set_title('Truth')
0222 ax2.set_title('Sigma')
0223 ax3.set_title('Double Angle')
0224 ax4.set_title('Electron')
0225 ax5.set_title('eSigma')
0226 ax6.set_title('Jacquet Blondel')
0227
0228 fig.suptitle('Kinematic Coverage  {}x{}   $minQ^2=${}$GeV^2$\n  {} events  DETECTOR_CONFIG: {} '.format(k,p,minq2,Nevents,Dconfig))
0229 plt.savefig(os.path.join(r_path, 'kinematic_coverage_%s.png' %(config)))