Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2026-09-29 08:26:40

0001 import numpy as np
0002 import matplotlib.pyplot as plt
0003 import mplhep as hep
0004 import uproot
0005 import pandas as pd
0006 from scipy.optimize import curve_fit
0007 from matplotlib.backends.backend_pdf import PdfPages
0008 import os
0009 import awkward as ak
0010 
0011 plt.figure()
0012 hep.set_style(hep.style.CMS)
0013 hep.set_style("CMS")
0014 
0015 def gaussian(x, amp, mean, sigma):
0016     return amp * np.exp( -(x - mean)**2 / (2*sigma**2) ) 
0017 
0018 def rotateY(xdata, zdata, angle):
0019     s = np.sin(angle)
0020     c = np.cos(angle)
0021     rotatedz = c*zdata - s*xdata
0022     rotatedx = s*zdata + c*xdata
0023     return rotatedx, rotatedz
0024     
0025 Energy = [0.005, 0.01, 0.05, 0.1, 0.5, 1.0]
0026 
0027 DETECTOR_CONFIG = os.environ["DETECTOR_CONFIG"]
0028 
0029 df = pd.DataFrame({})
0030 for eng in Energy:
0031     tree = uproot.open(f'sim_output/zdc_lyso/{DETECTOR_CONFIG}_gamma_{eng}GeV_theta_0deg_thru_0.3deg.eicrecon.edm4eic.root')['events']
0032     ecal_reco_energy = ak.sum(tree['EcalFarForwardZDCClusters/EcalFarForwardZDCClusters.energy'].array(), axis=-1)
0033     hcal_reco_energy = ak.sum(tree['HcalFarForwardZDCClusters/HcalFarForwardZDCClusters.energy'].array(), axis=-1)
0034     ecal_rec_energy = ak.sum(tree['EcalFarForwardZDCRecHits/EcalFarForwardZDCRecHits.energy'].array(), axis=-1)
0035     hcal_rec_energy = ak.sum(tree['HcalFarForwardZDCRecHits/HcalFarForwardZDCRecHits.energy'].array(), axis=-1)
0036     ecal_reco_clusters = [len(row) if len(row)>=1 else 0 for row in tree['EcalFarForwardZDCClusters/EcalFarForwardZDCClusters.nhits'].array()]
0037     ecal_reco_nhits = [row[0] if len(row)>=1 else 0 for row in tree['EcalFarForwardZDCClusters/EcalFarForwardZDCClusters.nhits'].array()]
0038     
0039     tree = uproot.open(f'sim_output/zdc_lyso/{DETECTOR_CONFIG}_gamma_{eng}GeV_theta_0deg_thru_0.3deg.edm4hep.root')['events']
0040     ecal_sim_energy = ak.sum(tree['EcalFarForwardZDCHits/EcalFarForwardZDCHits.energy'].array(), axis=-1)
0041     hcal_sim_energy = ak.sum(tree['HcalFarForwardZDCHits/HcalFarForwardZDCHits.energy'].array(), axis=-1)
0042 
0043     par_x = tree['MCParticles/MCParticles.momentum.x'].array()[:,2]
0044     par_y = tree['MCParticles/MCParticles.momentum.y'].array()[:,2]
0045     par_z = tree['MCParticles/MCParticles.momentum.z'].array()[:,2]
0046     
0047     eng = int(eng*1000)
0048 
0049     ecal_reco_energy = pd.DataFrame({f'ecal_reco_energy_{eng}': np.array(ecal_reco_energy, dtype=object)})
0050     hcal_reco_energy = pd.DataFrame({f'hcal_reco_energy_{eng}': np.array(hcal_reco_energy, dtype=object)})
0051     ecal_rec_energy = pd.DataFrame({f'ecal_rec_energy_{eng}': np.array(ecal_rec_energy, dtype=object)})
0052     hcal_rec_energy = pd.DataFrame({f'hcal_rec_energy_{eng}': np.array(hcal_rec_energy, dtype=object)})
0053     ecal_sim_energy = pd.DataFrame({f'ecal_sim_energy_{eng}': np.array(ecal_sim_energy, dtype=object)})
0054     hcal_sim_energy = pd.DataFrame({f'hcal_sim_energy_{eng}': np.array(hcal_sim_energy, dtype=object)})
0055     ecal_reco_nhits = pd.DataFrame({f'ecal_reco_nhits_{eng}': np.array(ecal_reco_nhits, dtype=object)})
0056     ecal_reco_clusters = pd.DataFrame({f'ecal_reco_clusters_{eng}': np.array(ecal_reco_clusters, dtype=object)})
0057     par_x = pd.DataFrame({f'par_x_{eng}': np.array(par_x.tolist(), dtype=object)})
0058     par_y = pd.DataFrame({f'par_y_{eng}': np.array(par_y.tolist(), dtype=object)})
0059     par_z = pd.DataFrame({f'par_z_{eng}': np.array(par_z.tolist(), dtype=object)})
0060 
0061 
0062     df = pd.concat([df,ecal_reco_energy,ecal_rec_energy,ecal_sim_energy,hcal_reco_energy,hcal_rec_energy,hcal_sim_energy,ecal_reco_clusters,ecal_reco_nhits,par_x,par_y,par_z],axis=1)
0063 
0064 
0065 mu = []
0066 sigma = []
0067 fig1, ax = plt.subplots(3,2,figsize=(20,10))
0068 fig1.suptitle('ZDC ECal Cluster Energy Reconstruction')
0069 
0070 plt.tight_layout()
0071 for i in range(6):
0072     x = df[f'par_x_{eng}'].astype(float).to_numpy()
0073     y = df[f'par_y_{eng}'].astype(float).to_numpy()
0074     z = df[f'par_z_{eng}'].astype(float).to_numpy()
0075     x, z = rotateY(x,z, 0.025)
0076     theta = np.arccos(z/np.sqrt((x**2+y**2+z**2)))*1000
0077     condition = theta <= 3.5
0078     
0079     plt.sca(ax[i%3,i//3])
0080     eng = int(Energy[i]*1000)
0081     plt.title(f'Gamma Energy: {eng} MeV')
0082     temp = np.array(df[f'ecal_reco_energy_{eng}'].astype(float).to_numpy()[condition])*1000
0083     hist, x = np.histogram(temp,bins=np.linspace(min(temp),max(temp)+np.std(abs(temp)),2*int(np.sqrt(len(temp)))))
0084     x = x[1:]/2 + x[:-1]/2
0085     plt.errorbar(x,hist,yerr=np.sqrt(hist),fmt='-o',label='Cluster')
0086     try:
0087         coeff, covar = curve_fit(gaussian,x[1:],hist[1:],p0=(max(hist[x>=np.std(abs(temp))]),np.mean(temp[temp!=0]),np.std(temp[temp!=0])),maxfev=10000)
0088         #plt.plot(np.linspace(coeff[1]-3*coeff[2],coeff[1]+3*coeff[2],50),gaussian(np.linspace(coeff[1]-3*coeff[2],coeff[1]+3*coeff[2],50),*coeff))
0089         mu.append(coeff[1])
0090         sigma.append(coeff[2])
0091     except RuntimeError:
0092         print("fit failed")
0093         mu.append(np.nan)
0094         sigma.append(np.nan)
0095     
0096     temp = np.array(df[f'ecal_rec_energy_{eng}'].astype(float).to_numpy()[condition])*1000
0097     hist, x = np.histogram(temp,bins=np.linspace(min(temp),max(temp)+np.std(abs(temp)),2*int(np.sqrt(len(temp)))))
0098     x = x[1:]/2 + x[:-1]/2
0099     plt.errorbar(x,hist,yerr=np.sqrt(hist),fmt='-o',label='Digitization')
0100     try:
0101         coeff, covar = curve_fit(gaussian,x[1:],hist[1:],p0=(max(hist[x>=np.std(abs(temp))]),np.mean(temp[temp!=0]),np.std(temp[temp!=0])),maxfev=10000)
0102         #plt.plot(np.linspace(coeff[1]-3*coeff[2],coeff[1]+3*coeff[2],50),gaussian(np.linspace(coeff[1]-3*coeff[2],coeff[1]+3*coeff[2],50),*coeff))
0103         mu.append(coeff[1])
0104         sigma.append(coeff[2])
0105     except RuntimeError:
0106         print("fit failed")
0107         mu.append(np.nan)
0108         sigma.append(np.nan)
0109     
0110     temp = np.array(df[f'ecal_sim_energy_{eng}'].astype(float).to_numpy()[condition])*1000
0111     hist, x = np.histogram(temp,bins=np.linspace(min(temp),max(temp)+np.std(abs(temp)),2*int(np.sqrt(len(temp)))))
0112     x = x[1:]/2 + x[:-1]/2
0113     plt.errorbar(x,hist,yerr=np.sqrt(hist),fmt='-o',label='Simulation')
0114     try:
0115         coeff, covar = curve_fit(gaussian,x[1:],hist[1:],p0=(max(hist[x>=np.std(abs(temp))]),np.mean(temp[temp!=0]),np.std(temp[temp!=0])),maxfev=10000)
0116         #plt.plot(np.linspace(coeff[1]-3*coeff[2],coeff[1]+3*coeff[2],50),gaussian(np.linspace(coeff[1]-3*coeff[2],coeff[1]+3*coeff[2],50),*coeff))
0117         mu.append(coeff[1])
0118         sigma.append(coeff[2])
0119     except RuntimeError:
0120         print("fit failed")
0121         mu.append(np.nan)
0122         sigma.append(np.nan)
0123     
0124     plt.xlabel('Energy (MeV)')
0125     plt.legend()
0126     
0127 #plt.savefig('results/Energy_reconstruction_cluster.pdf')
0128 
0129 mu = np.array(mu)
0130 sigma = np.array(sigma)
0131 
0132 plt.show()
0133 
0134 fig2, (ax1,ax2) = plt.subplots(2,1,figsize=(15,10),sharex=True)
0135 
0136 plt.tight_layout()
0137 # Plot data on primary axis
0138 ax1.scatter(np.array(Energy)*1000, mu[::3], label='cluster')
0139 ax1.scatter(np.array(Energy)*1000, mu[1::3], label='digitization')
0140 ax1.scatter(np.array(Energy)*1000, mu[2::3], label='simulation')
0141 
0142 ax1.plot([4.5,1000],[4.5,1000],c='black',label='x=y')
0143 ax1.set_ylabel('Reconstructed Energy (MeV)')
0144 ax1.set_yscale('log')
0145 ax1.legend()
0146 ax1.set_title('ECal Craterlake Cluster Energy Reconstruction')
0147 
0148 ax2.errorbar(np.array(Energy)*1000, abs(sigma[::3]/mu[::3])*100, fmt='-o', label='cluster')
0149 ax2.errorbar(np.array(Energy)*1000, abs(sigma[1::3]/mu[1::3])*100, fmt='-o', label='digitization')
0150 ax2.errorbar(np.array(Energy)*1000, abs(sigma[2::3]/mu[2::3])*100, fmt='-o', label='simulation')
0151 
0152 ax2.set_ylabel('Resolution (%)')
0153 ax2.set_xlabel('Gamma Energy (MeV)')
0154 ax2.set_xscale('log')
0155 ax2.legend()
0156 
0157 #plt.savefig('results/Energy_resolution.pdf')
0158 
0159 plt.show()
0160 
0161 
0162 htower = []
0163 herr = []
0164 hmean = []
0165 hhits = []
0166 hhits_cut = []
0167 emean = []
0168 ehits = []
0169 etower = []
0170 eerr = []
0171 ehits_cut = []
0172 
0173 fig3, ax = plt.subplots(2,3,figsize=(20,10))
0174 fig3.suptitle('ZDC Simulation Energy Reconstruction')
0175 for i in range(6):
0176     plt.sca(ax[i//3,i%3])
0177     eng = int(Energy[i]*1000)
0178 
0179     x = df[f'par_x_{eng}'].astype(float).to_numpy()
0180     y = df[f'par_y_{eng}'].astype(float).to_numpy()
0181     z = df[f'par_z_{eng}'].astype(float).to_numpy()
0182     x, z = rotateY(x,z, 0.025)
0183     theta = np.arccos(z/np.sqrt((x**2+y**2+z**2)))*1000
0184     condition = theta <= 3.5
0185 
0186     plt.title(f'Gamma Energy: {eng} MeV')
0187     energy1 = df[f'hcal_sim_energy_{eng}'].astype(float).to_numpy()#df.eval(f'hcal_sim_energy_{eng}').apply(lambda row: sum(row))
0188     hist, x = np.histogram(energy1*1000,bins=np.logspace(0,3,200))
0189     x = x[1:]/2 + x[:-1]/2
0190     plt.plot(x,hist,marker='o',label="HCal")
0191     hhits.append(len(energy1[energy1!=0]))
0192     condition1 = energy1!=0
0193     hhits_cut.append(len(energy1[condition & condition1])/len(condition[condition==True]))
0194     energy = df[f'ecal_sim_energy_{eng}'].astype(float).to_numpy()#df.eval(f'ecal_sim_energy_{eng}').apply(lambda row: sum(row))
0195     hist, x = np.histogram(energy*1000,bins=np.logspace(0,3,200))
0196     x = x[1:]/2 + x[:-1]/2
0197     plt.plot(x,hist,marker='o',label="ECal")
0198     emean.append(sum(energy[energy!=0])*1000/len(energy[energy!=0]))
0199     hmean.append(sum(energy1[energy!=0])*1000/len(energy[energy!=0]))
0200     condition1 = energy!=0
0201     ehits_cut.append(len(energy[condition & condition1])/len(condition[condition==True]))
0202     ehits.append(len(energy[energy!=0]))
0203     plt.legend()
0204     plt.xscale('log')
0205     plt.xlim(1e0,1e3)
0206 
0207 
0208 
0209     
0210 
0211 plt.xlabel('Energy (MeV)')
0212 
0213 #plt.savefig('results/Energy_deposition.pdf')
0214 plt.show()
0215 
0216 fig4, ax = plt.subplots(2,1,sharex=True,gridspec_kw={'height_ratios': [2,1]})
0217 plt.sca(ax[0])
0218 plt.errorbar(np.array(Energy)*1000,np.array(hmean)*47.619+np.array(emean),label='HCal/sf+ECal',fmt='-o')
0219 plt.errorbar(np.array(Energy)*1000,emean,label='ECal',fmt='-o')
0220 plt.legend()
0221 plt.yscale('log')
0222 plt.xscale('log')
0223 plt.ylabel('Simulation Energy (MeV)')
0224 plt.sca(ax[1])
0225 plt.errorbar(np.array(Energy)*1000,(1 - np.array(emean)/(np.array(hmean)*47.619+np.array(emean)))*100,label='Total/ECal',fmt='-o')
0226 plt.legend()
0227 plt.ylabel('Fraction of energy\n deposited in Hcal (%)')
0228 plt.xlabel('Truth Energy (MeV)')
0229 #plt.savefig('results/Energy_ratio_and_Leakage.pdf')
0230 plt.tight_layout()
0231 plt.show()
0232 
0233 fig5 = plt.figure()
0234 plt.errorbar(np.array(Energy)*1000,np.array(hhits)/1000*100,label='HCal Hits',fmt='-o')
0235 plt.errorbar(np.array(Energy)*1000,np.array(ehits)/1000*100,label='ECal Hits',fmt='-o')
0236 #plt.errorbar(np.array(Energy)*1000,np.array(hhits)/np.array(ehits)*100,label='HCal / ECal',fmt='-o',c='b')
0237 
0238 plt.errorbar(np.array(Energy)*1000,np.array(hhits_cut)*100,label='HCal Hits with 3.5 mRad cut',fmt='-^')
0239 plt.errorbar(np.array(Energy)*1000,np.array(ehits_cut)*100,label='ECal Hits with 3.5 mRad cut',fmt='-^')
0240 #plt.errorbar(np.array(Energy)*1000,np.array(hhits_cut)/np.array(ehits_cut)*100,label='HCal / ECal with 3.5 mRad cut',fmt='-^',c='b')
0241 ### 3mrad cuts
0242 
0243 plt.legend()
0244 plt.xlabel('Simulation Truth Gamma Energy (MeV)')
0245 plt.ylabel('Fraction of Events with non-zero energy (%)')
0246 #plt.savefig('results/Hits.pdf')
0247 plt.xscale('log')
0248 plt.show()
0249 
0250 fig6, ax = plt.subplots(2,3,figsize=(20,10))
0251 fig6.suptitle('ZDC Clustering')
0252 fig6.tight_layout(pad=1.8)
0253 for i in range(6):
0254     plt.sca(ax[i//3,i%3])
0255     eng = int(Energy[i]*1000)
0256     
0257     x = df[f'par_x_{eng}'].astype(float).to_numpy()
0258     y = df[f'par_y_{eng}'].astype(float).to_numpy()
0259     z = df[f'par_z_{eng}'].astype(float).to_numpy()
0260     x, z = rotateY(x,z, 0.025)
0261     theta = np.arccos(z/np.sqrt((x**2+y**2+z**2)))*1000
0262     condition = theta <= 3.5
0263     
0264     plt.hist(df[f'ecal_reco_clusters_{eng}'][condition],bins=np.linspace(0,5,6))
0265     plt.xlabel('Number of Clusters')
0266     plt.title(f'Gamma Energy: {eng} MeV')
0267 plt.show()
0268 
0269 fig7, ax = plt.subplots(2,3,figsize=(20,10))
0270 fig7.suptitle('ZDC Towering in Clusters')
0271 fig7.tight_layout(pad=1.8)
0272 for i in range(6):
0273     plt.sca(ax[i//3,i%3])
0274     eng = int(Energy[i]*1000)
0275     
0276     x = df[f'par_x_{eng}'].astype(float).to_numpy()
0277     y = df[f'par_y_{eng}'].astype(float).to_numpy()
0278     z = df[f'par_z_{eng}'].astype(float).to_numpy()
0279     x, z = rotateY(x,z, 0.025)
0280     theta = np.arccos(z/np.sqrt((x**2+y**2+z**2)))*1000
0281     condition = theta <= 3.5
0282     
0283     plt.hist(df[f'ecal_reco_nhits_{eng}'][condition],bins=np.linspace(0,max(df[f'ecal_reco_nhits_{eng}'][condition]),max(df[f'ecal_reco_nhits_{eng}'][condition])+1))
0284     plt.xlabel('Number of tower in Clusters')
0285     plt.title(f'Gamma Energy: {eng} MeV')
0286 plt.show()
0287 
0288 
0289 #pdfs = ['results/Energy_reconstruction_cluster.pdf','results/Energy_resolution.pdf','results/Energy_deposition.pdf','results/Energy_ratio_and_Leakage.pdf','results/Hits.pdf']
0290 with PdfPages(f'results/{DETECTOR_CONFIG}/zdc_lyso/plots.pdf') as pdf:
0291     pdf.savefig(fig1)
0292     pdf.savefig(fig2)
0293     pdf.savefig(fig3)
0294     pdf.savefig(fig4)
0295     pdf.savefig(fig5)
0296     pdf.savefig(fig6)
0297     pdf.savefig(fig7)