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 # Format-agnostic data loading: uproot 5.x automatically handles both TTree and RNTuple formats
0151 # Works with:
0152 #   - .edm4eic.root files (TTree format)
0153 #   - .edm4eic.rnt.root files (RNTuple format)
0154 # Both formats created by podio use the same 'events' tree/RNTuple name and branch structure
0155 file = ur.open(rec_file)
0156 events = file["events"]
0157
0158 # Use direct array access instead of concatenate() to avoid RNTuple compatibility issues
0159 Truth =   [events['InclusiveKinematicsTruth.Q2'].array(library="ak"),
0160            events['InclusiveKinematicsTruth.x'].array(library="ak")]
0161 Electron = [events['InclusiveKinematicsElectron.Q2'].array(library="ak"),
0162             events['InclusiveKinematicsElectron.x'].array(library="ak")]
0163 DoubleAngle = [events['InclusiveKinematicsDA.Q2'].array(library="ak"),
0164                events['InclusiveKinematicsDA.x'].array(library="ak")]
0165 JacquetBlondel = [events['InclusiveKinematicsJB.Q2'].array(library="ak"),
0166                   events['InclusiveKinematicsJB.x'].array(library="ak")]
0167 Sigma = [events['InclusiveKinematicsSigma.Q2'].array(library="ak"),
0168          events['InclusiveKinematicsSigma.x'].array(library="ak")]
0169 try:
0170     ESigma = [events['InclusiveKinematicsESigma.Q2'].array(library="ak"),
0171               events['InclusiveKinematicsESigma.x'].array(library="ak")]
0172 except ur.KeyInFileError:
0173     # Legacy compatibility
0174     ESigma = [events['InclusiveKinematicseSigma.Q2'].array(library="ak"),
0175               events['InclusiveKinematicseSigma.x'].array(library="ak")]
0176
0177 Q2values_T = Truth[0]
0178 Q2values_E = Electron[0]
0179 Q2values_DA = DoubleAngle[0]
0180 Q2values_JB = JacquetBlondel[0]
0181 Q2values_S = Sigma[0]
0182 Q2values_eS = ESigma[0]
0183 Xvalues_T = Truth[1]
0184 Xvalues_E = Electron[1]
0185 Xvalues_DA = DoubleAngle[1]
0186 Xvalues_JB = JacquetBlondel[1]
0187 Xvalues_S = Sigma[1]
0188 Xvalues_eS = ESigma[1]
0189
0190 method_dict = {'e':'Electron','DA':'Double-Angle','JB':'Jacquet-Blondel','S':'Sigma','eS':'eSigma'}
0191 method_Q2values_dict = {'e':Q2values_E,'DA':Q2values_DA,'JB':Q2values_JB,'S':Q2values_S,'eS':Q2values_eS}
0192 method_Xvalues_dict = {'e':Xvalues_E,'DA':Xvalues_DA,'JB':Xvalues_JB,'S':Xvalues_S,'eS':Xvalues_eS}
0193
0194 Q2correlation(minq2,'e')
0195 Xcorrelation(minq2,'e')
0196 Q2correlation(minq2,'DA')
0197 Xcorrelation(minq2,'DA')
0198 Q2correlation(minq2,'JB')
0199 Xcorrelation(minq2,'JB')
0200 Q2correlation(minq2,'S')
0201 Xcorrelation(minq2,'S')
0202 Q2correlation(minq2,'eS')
0203 Xcorrelation(minq2,'eS')
0204
0205
0206 ##########################################################################################
0207 #Kinematic coverage plot
0208 ##########################################################################################
0209
0210 fig = plt.figure(figsize=(20,10))
0211 gs = fig.add_gridspec(2, 3, wspace=0.09, hspace = 0.2)
0212 (ax1, ax2, ax3), (ax4, ax5, ax6) = gs.subplots()
0213
0214 ax1.hist2d(ak.to_numpy(ak.flatten(Xvalues_T)), ak.to_numpy(ak.flatten(Q2values_T)), bins = [x_bins, Q2_bins])
0215 ax2.hist2d(ak.to_numpy(ak.flatten(Xvalues_S)), ak.to_numpy(ak.flatten(Q2values_S)), bins = [x_bins, Q2_bins])
0216 ax3.hist2d(ak.to_numpy(ak.flatten(Xvalues_DA)), ak.to_numpy(ak.flatten(Q2values_DA)), bins = [x_bins, Q2_bins])
0217 ax4.hist2d(ak.to_numpy(ak.flatten(Xvalues_E)), ak.to_numpy(ak.flatten(Q2values_E)), bins = [x_bins, Q2_bins])
0218 ax5.hist2d(ak.to_numpy(ak.flatten(Xvalues_eS)), ak.to_numpy(ak.flatten(Q2values_eS)), bins = [x_bins, Q2_bins])
0219 ax6.hist2d(ak.to_numpy(ak.flatten(Xvalues_JB)), ak.to_numpy(ak.flatten(Q2values_JB)), bins = [x_bins, Q2_bins])
0220
0221 for axes in [ax1,ax2,ax3,ax4,ax5,ax6]:
0222     axes.set_xscale('log')
0223     axes.set_yscale('log')
0224     axes.set_ylabel('$Q^2$ [$GeV^2$]')
0225     axes.set_xlabel('$x$')
0226 ax1.set_title('Truth')
0227 ax2.set_title('Sigma')
0228 ax3.set_title('Double Angle')
0229 ax4.set_title('Electron')
0230 ax5.set_title('eSigma')
0231 ax6.set_title('Jacquet Blondel')
0232
0233 fig.suptitle('Kinematic Coverage  {}x{}   $minQ^2=${}$GeV^2$\n  {} events  DETECTOR_CONFIG: {} '.format(k,p,minq2,Nevents,Dconfig))
0234 plt.savefig(os.path.join(r_path, 'kinematic_coverage_%s.png' %(config)))