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     # Filter based on which dataset has fewer events
0053     filter_boolean = ak.where(T_len > Y_len,
0054                               ak.count(Q2_List_Y,axis=-1) >= 1,  # if Truth has more, filter by method presence
0055                               ak.count(Q2_List_T,axis=-1) >= 1)  # else filter by Truth presence
0056     Q2_List_T_F = Q2_List_T[filter_boolean] #filtered Truth Q2 values
0057     Q2_List_Y_F = Q2_List_Y[filter_boolean] #filtered method Q2 values
0058
0059     T_Q2s = np.array(ak.flatten(Q2_List_T_F)) #Truth Q2 values, mapped along x axis
0060     Y_Q2s = np.array(ak.flatten(Q2_List_Y_F)) #method Q2 values, mapped along y axis
0061
0062     #2-dimensional histogram, h
0063     h, xedges, yedges = np.histogram2d(x=T_Q2s,y=Y_Q2s, bins=[Q2_bins,Q2_bins])
0064
0065     minq2_dict = {'1':2,'10':7,'100':12,'1000':17} #Q2 bin index at which minq2 starts
0066     h[0:minq2_dict['{}'.format(minq2)]]=0 #ignore values before minq2
0067
0068     #normalization of h:
0069     col_sum = ak.sum(h,axis=-1) #number of events in each (verticle) column
0070     norm_h = [] #norm_h is the normalized matrix
0071     norm_h_text = [] #display labels matrix
0072     for i in range(len(col_sum)):
0073         if col_sum[i] != 0:
0074             norm_c = h[i]/col_sum[i] #normalized column = column values divide by sum of the column
0075         else:
0076             norm_c = h[i]
0077         norm_h.append(norm_c)
0078         norm_c_text = [ '%.3f' % elem for elem in norm_c ] #display value to 3 dp
0079         norm_h_text.append(norm_c_text)
0080
0081     fig = plt.figure()
0082     mplhep.hist2dplot(H=norm_h,norm=mpl.colors.LogNorm(vmin= 1e-4, vmax= 1),labels=norm_h_text,xbins=Q2_bins,ybins=Q2_bins)
0083     plt.yscale('log')
0084     plt.xscale('log')
0085     fig.set_figwidth(11)
0086     fig.set_figheight(11)
0087     plt.xlabel('$Q^2$ [$GeV^2$] Truth')
0088     plt.ylabel('$Q^2$ [$GeV^2$] {}'.format(method_dict['{}'.format(method)]))
0089     plt.title('{}   $Q^2$ correlation   {}x{}   $minQ^2=${}$GeV^2$\n  {} events  DETECTOR_CONFIG: {} '.format(method_dict['{}'.format(method)],k,p,minq2,Nevents,Dconfig))
0090     plt.savefig(os.path.join(r_path, 'Q2_correlation_%s_%s.png' %(method,config)))
0091
0092
0093 ##########################################################################################
0094 #function to construct Bjorken-x correlation plots
0095 ##########################################################################################
0096
0097 def Xcorrelation(minq2,method): #minq2 can be 1,10,100, or 1000; method can be 'e','DA', or 'JB'
0098
0099     Xvalues_Y = method_Xvalues_dict['{}'.format(method)] #x values of the given method, that are mapped onto the y axis
0100
0101     X_List_T = Xvalues_T #Truth x values, mapped along x axis
0102     X_List_Y = Xvalues_Y #method (E/DA/JB) x values, mapped along y axis
0103
0104     T_len = ak.count(X_List_T,axis=0) #total number of events in Truth
0105     Y_len = ak.count(X_List_Y,axis=0) #total number of events in method
0106
0107     # Filter based on which dataset has fewer events
0108     filter_boolean = ak.where(T_len > Y_len,
0109                               ak.count(X_List_Y,axis=-1) >= 1,  # if Truth has more, filter by method presence
0110                               ak.count(X_List_T,axis=-1) >= 1)  # else filter by Truth presence
0111     X_List_T_F = X_List_T[filter_boolean] #filtered Truth x values
0112     X_List_Y_F = X_List_Y[filter_boolean] #filtered method x values
0113
0114     T_Xs = np.array(ak.flatten(X_List_T_F)) #Truth Bjorken-x values, mapped along x axis
0115     Y_Xs = np.array(ak.flatten(X_List_Y_F)) #method Bjorken-x values, mapped along y axis
0116
0117     T_x_bool = T_Xs>=minq2/(4*k*p) #boolean to filter x values that satisfy bjorken-x equation for minq2, ebeam and pbeam
0118     T_Xs = T_Xs[T_x_bool]
0119     Y_Xs = Y_Xs[T_x_bool]
0120
0121     #2-dimensional histogram, h
0122     h, xedges, yedges = np.histogram2d(x=T_Xs,y=Y_Xs, bins=[x_bins,x_bins])
0123
0124     #normalization of h:
0125     col_sum = ak.sum(h,axis=-1) #number of events in each (verticle) column
0126     norm_h = [] #norm_h is the normalized matrix
0127     norm_h_text = [] #display labels matrix
0128     for i in range(len(col_sum)):
0129         if col_sum[i] != 0:
0130             norm_c = h[i]/col_sum[i] #normalized column = column values divide by sum of the column
0131         else:
0132             norm_c = h[i]
0133         norm_h.append(norm_c)
0134         norm_c_text = [ '%.2f' % elem for elem in norm_c ] #display value to 2 dp
0135         norm_h_text.append(norm_c_text)
0136
0137     fig = plt.figure()
0138     mplhep.hist2dplot(H=norm_h,norm=mpl.colors.LogNorm(vmin= 1e-4, vmax= 1),labels=norm_h_text,xbins=x_bins,ybins=x_bins)
0139     plt.yscale('log')
0140     plt.xscale('log')
0141     fig.set_figwidth(11)
0142     fig.set_figheight(11)
0143     plt.xlabel('x Truth')
0144     plt.ylabel('$x$   {}'.format(method_dict['{}'.format(method)]))
0145     plt.title('{}   $x$ correlation   {}x{}   $minQ^2=${}$GeV^2$\n  {} events  DETECTOR_CONFIG: {} '.format(method_dict['{}'.format(method)],k,p,minq2,Nevents,Dconfig))
0146     plt.savefig(os.path.join(r_path, 'x_correlation_%s_%s.png' %(method,config)))
0147
0148
0149
0150 keys = ur.concatenate(rec_file + ':events/' + 'InclusiveKinematicsTruth')
0151 Truth =   [keys['InclusiveKinematicsTruth.Q2'],keys['InclusiveKinematicsTruth.x']]
0152 keys = ur.concatenate(rec_file + ':events/' + 'InclusiveKinematicsElectron')
0153 Electron =   [keys['InclusiveKinematicsElectron.Q2'], keys['InclusiveKinematicsElectron.x']]
0154 keys = ur.concatenate(rec_file + ':events/' + 'InclusiveKinematicsDA')
0155 DoubleAngle =  [keys['InclusiveKinematicsDA.Q2'], keys['InclusiveKinematicsDA.x']]
0156 keys = ur.concatenate(rec_file + ':events/' + 'InclusiveKinematicsJB')
0157 JacquetBlondel =  [keys['InclusiveKinematicsJB.Q2'], keys['InclusiveKinematicsJB.x']]
0158 keys = ur.concatenate(rec_file + ':events/' + 'InclusiveKinematicsSigma')
0159 Sigma =  [keys['InclusiveKinematicsSigma.Q2'], keys['InclusiveKinematicsSigma.x']]
0160 try:
0161     keys = ur.concatenate(rec_file + ':events/' + 'InclusiveKinematicsESigma')
0162     ESigma =  [keys['InclusiveKinematicsESigma.Q2'], keys['InclusiveKinematicsESigma.x']]
0163 except ur.KeyInFileError:
0164     # Legacy compatibility
0165     keys = ur.concatenate(rec_file + ':events/' + 'InclusiveKinematicseSigma')
0166     ESigma =  [keys['InclusiveKinematicseSigma.Q2'], keys['InclusiveKinematicseSigma.x']]
0167
0168 Q2values_T = Truth[0]
0169 Q2values_E = Electron[0]
0170 Q2values_DA = DoubleAngle[0]
0171 Q2values_JB = JacquetBlondel[0]
0172 Q2values_S = Sigma[0]
0173 Q2values_eS = ESigma[0]
0174 Xvalues_T = Truth[1]
0175 Xvalues_E = Electron[1]
0176 Xvalues_DA = DoubleAngle[1]
0177 Xvalues_JB = JacquetBlondel[1]
0178 Xvalues_S = Sigma[1]
0179 Xvalues_eS = ESigma[1]
0180
0181 method_dict = {'e':'Electron','DA':'Double-Angle','JB':'Jacquet-Blondel','S':'Sigma','eS':'eSigma'}
0182 method_Q2values_dict = {'e':Q2values_E,'DA':Q2values_DA,'JB':Q2values_JB,'S':Q2values_S,'eS':Q2values_eS}
0183 method_Xvalues_dict = {'e':Xvalues_E,'DA':Xvalues_DA,'JB':Xvalues_JB,'S':Xvalues_S,'eS':Xvalues_eS}
0184
0185 Q2correlation(minq2,'e')
0186 Xcorrelation(minq2,'e')
0187 Q2correlation(minq2,'DA')
0188 Xcorrelation(minq2,'DA')
0189 Q2correlation(minq2,'JB')
0190 Xcorrelation(minq2,'JB')
0191 Q2correlation(minq2,'S')
0192 Xcorrelation(minq2,'S')
0193 Q2correlation(minq2,'eS')
0194 Xcorrelation(minq2,'eS')
0195
0196
0197 ##########################################################################################
0198 #Kinematic coverage plot
0199 ##########################################################################################
0200
0201 fig = plt.figure(figsize=(20,10))
0202 gs = fig.add_gridspec(2, 3, wspace=0.09, hspace = 0.2)
0203 (ax1, ax2, ax3), (ax4, ax5, ax6) = gs.subplots()
0204
0205 ax1.hist2d(ak.to_numpy(ak.flatten(Xvalues_T)), ak.to_numpy(ak.flatten(Q2values_T)), bins = [x_bins, Q2_bins])
0206 ax2.hist2d(ak.to_numpy(ak.flatten(Xvalues_S)), ak.to_numpy(ak.flatten(Q2values_S)), bins = [x_bins, Q2_bins])
0207 ax3.hist2d(ak.to_numpy(ak.flatten(Xvalues_DA)), ak.to_numpy(ak.flatten(Q2values_DA)), bins = [x_bins, Q2_bins])
0208 ax4.hist2d(ak.to_numpy(ak.flatten(Xvalues_E)), ak.to_numpy(ak.flatten(Q2values_E)), bins = [x_bins, Q2_bins])
0209 ax5.hist2d(ak.to_numpy(ak.flatten(Xvalues_eS)), ak.to_numpy(ak.flatten(Q2values_eS)), bins = [x_bins, Q2_bins])
0210 ax6.hist2d(ak.to_numpy(ak.flatten(Xvalues_JB)), ak.to_numpy(ak.flatten(Q2values_JB)), bins = [x_bins, Q2_bins])
0211
0212 for axes in [ax1,ax2,ax3,ax4,ax5,ax6]:
0213     axes.set_xscale('log')
0214     axes.set_yscale('log')
0215     axes.set_ylabel('$Q^2$ [$GeV^2$]')
0216     axes.set_xlabel('$x$')
0217 ax1.set_title('Truth')
0218 ax2.set_title('Sigma')
0219 ax3.set_title('Double Angle')
0220 ax4.set_title('Electron')
0221 ax5.set_title('eSigma')
0222 ax6.set_title('Jacquet Blondel')
0223
0224 fig.suptitle('Kinematic Coverage  {}x{}   $minQ^2=${}$GeV^2$\n  {} events  DETECTOR_CONFIG: {} '.format(k,p,minq2,Nevents,Dconfig))
0225 plt.savefig(os.path.join(r_path, 'kinematic_coverage_%s.png' %(config)))