File indexing completed on 2024-09-27 07:02:40
0001 import numpy as np, pandas as pd, matplotlib.pyplot as plt, matplotlib as mpl, awkward as ak, sys
0002 import mplhep as hep
0003 hep.style.use("CMS")
0004
0005 plt.rcParams['figure.facecolor']='white'
0006 plt.rcParams['savefig.facecolor']='white'
0007 plt.rcParams['savefig.bbox']='tight'
0008
0009 outdir=sys.argv[1]+"/"
0010 config=outdir.split("/")[1]
0011 try:
0012 import os
0013 os.mkdir(outdir[:-1])
0014 except:
0015 pass
0016
0017 def gauss(x, A,mu, sigma):
0018 return A * np.exp(-(x-mu)**2/(2*sigma**2))
0019
0020 import uproot as ur
0021 arrays_sim={}
0022 momenta=60, 80, 100, 130, 160,
0023 for p in momenta:
0024 filename=f'sim_output/zdc_pi0/{config}_rec_zdc_pi0_{p}GeV.edm4hep.root'
0025 print("opening file", filename)
0026 events = ur.open(filename+':events')
0027 arrays_sim[p] = events.arrays()
0028 import gc
0029 gc.collect()
0030 print("read", filename)
0031
0032
0033 fig,axs=plt.subplots(1,3, figsize=(24, 8))
0034 pvals=[]
0035 resvals=[]
0036 dresvals=[]
0037 scalevals=[]
0038 dscalevals=[]
0039 for p in momenta:
0040 selection=[len(arrays_sim[p]["HcalFarForwardZDCClusters.energy"][i])==2 for i in range(len(arrays_sim[p]))]
0041 E=arrays_sim[p][selection]["HcalFarForwardZDCClusters.energy"]
0042
0043 Etot=np.sum(E, axis=-1)
0044 if len(Etot)<25:
0045 continue
0046
0047 if p==100:
0048 plt.sca(axs[0])
0049 y, x, _=plt.hist(Etot, bins=100, range=(p*.5, p*1.5), histtype='step')
0050 plt.ylabel("events")
0051 plt.title(f"$p_{{\pi^0}}$={p} GeV")
0052 plt.xlabel("$E^{\\pi^{0}}_{recon}$ [GeV]")
0053 else:
0054 y, x = np.histogram(Etot, bins=100, range=(p*.5, p*1.5))
0055
0056 bc=(x[1:]+x[:-1])/2
0057 from scipy.optimize import curve_fit
0058 slc=abs(bc-p)<10
0059 fnc=gauss
0060 p0=[100, p, 10]
0061
0062 coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,
0063 sigma=list(np.sqrt(y[slc])+(y[slc]==0)))
0064 if p==100:
0065 xx=np.linspace(p*0.5,p*1.5, 100)
0066 plt.plot(xx, fnc(xx,*coeff))
0067 pvals.append(p)
0068 resvals.append(np.abs(coeff[2])/coeff[1])
0069 dresvals.append(np.sqrt(var_matrix[2][2])/coeff[1])
0070 scalevals.append(np.abs(coeff[1])/p)
0071 dscalevals.append(np.sqrt(var_matrix[2][2])/p)
0072
0073 plt.sca(axs[1])
0074 plt.errorbar(pvals, resvals, dresvals, ls='', marker='o')
0075
0076 plt.ylabel("$\\sigma[E_{\\pi^0}]/\\mu[E_{\\pi^0}]$")
0077 plt.xlabel("$p_{\\pi^0}$ [GeV]")
0078
0079 fnc=lambda E,a: a/np.sqrt(E)
0080
0081 coeff, var_matrix = curve_fit(fnc, pvals, resvals, p0=(1,),
0082 sigma=dresvals)
0083 xx=np.linspace(55, 200, 100)
0084 plt.plot(xx, fnc(xx, *coeff), label=f'fit: $\\frac{{{coeff[0]:.2f}\\%}}{{\\sqrt{{E}}}}$')
0085 plt.legend()
0086 plt.ylim(0)
0087 plt.sca(axs[2])
0088 plt.errorbar(pvals, scalevals, dscalevals, ls='', marker='o')
0089 plt.ylim(0.8, 1.2)
0090 plt.ylabel("$\\mu[E_{\\pi^0}]/E_{\\pi^0}$")
0091 plt.xlabel("$p_{\\pi^0}$ [GeV]")
0092 plt.axhline(1, ls='--', alpha=0.7, color='0.5')
0093 plt.tight_layout()
0094 plt.savefig(outdir+"/pi0_energy_res.pdf")
0095
0096
0097 fig,axs=plt.subplots(1,2, figsize=(16, 8))
0098 pvals=[]
0099 resvals=[]
0100 dresvals=[]
0101 for p in momenta:
0102 selection=[len(arrays_sim[p]["HcalFarForwardZDCClusters.energy"][i])==2 for i in range(len(arrays_sim[p]))]
0103 x=arrays_sim[p][selection]["HcalFarForwardZDCClusters.position.x"]
0104 y=arrays_sim[p][selection]["HcalFarForwardZDCClusters.position.y"]
0105 z=arrays_sim[p][selection]["HcalFarForwardZDCClusters.position.z"]
0106 E=arrays_sim[p][selection]["HcalFarForwardZDCClusters.energy"]
0107 r=np.sqrt(x**2+y**2+z**2)
0108 px=np.sum(E*x/r, axis=-1)
0109 py=np.sum(E*y/r, axis=-1)
0110 pz=np.sum(E*z/r, axis=-1)
0111
0112 theta_recon=np.arctan2(np.hypot(px*np.cos(-.025)-pz*np.sin(-.025), py), pz*np.cos(-.025)+px*np.sin(-.025))
0113 if len(theta_recon)<25:
0114 continue
0115 px=arrays_sim[p][selection]["MCParticles.momentum.x"][::,2]
0116 py=arrays_sim[p][selection]["MCParticles.momentum.y"][::,2]
0117 pz=arrays_sim[p][selection]["MCParticles.momentum.z"][::,2]
0118
0119 theta_truth=np.arctan2(np.hypot(px*np.cos(-.025)-pz*np.sin(-.025), py), pz*np.cos(-.025)+px*np.sin(-.025))
0120
0121 Etot=np.sum(E, axis=-1)
0122
0123 if p==100:
0124 plt.sca(axs[0])
0125 y, x, _=plt.hist(1000*(theta_recon-theta_truth), bins=100, range=(-0.5, 0.5), histtype='step')
0126 plt.ylabel("events")
0127 plt.title(f"$p_{{\\pi^0}}$={p} GeV")
0128 plt.xlabel("$\\theta^{\\pi^0}_{recon}$ [mrad]")
0129 else:
0130 y, x = np.histogram(1000*(theta_recon-theta_truth), bins=100, range=(-0.5, 0.5))
0131
0132 bc=(x[1:]+x[:-1])/2
0133 from scipy.optimize import curve_fit
0134 slc=abs(bc)<0.2
0135 fnc=gauss
0136 p0=[100, 0, 0.1]
0137
0138 coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,
0139 sigma=list(np.sqrt(y[slc])+(y[slc]==0)))
0140 if p==100:
0141 xx=np.linspace(-0.5,0.5, 100)
0142 plt.plot(xx, fnc(xx,*coeff))
0143 pvals.append(p)
0144 resvals.append(np.abs(coeff[2]))
0145 dresvals.append(np.sqrt(var_matrix[2][2]))
0146
0147 plt.sca(axs[1])
0148 plt.errorbar(pvals, resvals, dresvals, ls='', marker='o')
0149
0150
0151 fnc=lambda E,a: a/np.sqrt(E)
0152
0153 coeff, var_matrix = curve_fit(fnc, pvals, resvals, p0=(1,),
0154 sigma=dresvals)
0155
0156 xx=np.linspace(55, 200, 100)
0157
0158 plt.plot(xx, fnc(xx, *coeff), label=f'fit: $\\frac{{{coeff[0]:.2f}}}{{\\sqrt{{E}}}}$ mrad')
0159
0160 plt.ylabel("$\\sigma[\\theta_{\\pi^0}]$ [mrad]")
0161 plt.xlabel("$p_{\\pi^0}$ [GeV]")
0162
0163 plt.ylim(0, 0.1)
0164 plt.legend()
0165 plt.tight_layout()
0166 plt.savefig(outdir+"/pi0_theta_res.pdf")
0167
0168 fig,axs=plt.subplots(1,2, figsize=(16, 8))
0169 pvals=[]
0170 resvals=[]
0171 dresvals=[]
0172 for p in momenta:
0173 selection=[len(arrays_sim[p]["HcalFarForwardZDCClusters.energy"][i])==2 for i in range(len(arrays_sim[p]))]
0174 E=arrays_sim[p][selection]["HcalFarForwardZDCClusters.energy"]
0175 cx=arrays_sim[p][selection]["HcalFarForwardZDCClusters.position.x"]
0176 cy=arrays_sim[p][selection]["HcalFarForwardZDCClusters.position.y"]
0177 cz=arrays_sim[p][selection]["HcalFarForwardZDCClusters.position.z"]
0178 r=np.sqrt(cx**2+cy**2+cz**2)
0179 px=E*cx/r
0180 py=E*cy/r
0181 pz=E*cz/r
0182
0183 cos_opening_angle=(cx/r)[::,0]*(cx/r)[::,1]+(cy/r)[::,0]*(cy/r)[::,1]+(cz/r)[::,0]*(cz/r)[::,1]
0184 mrecon=np.sqrt(2*E[::,0]*E[::,1]*(1-cos_opening_angle))
0185
0186 if len(mrecon)<25:
0187 continue
0188
0189
0190 if p==100:
0191 plt.sca(axs[0])
0192 y, x, _=plt.hist(mrecon, bins=100, range=(0, 0.2), histtype='step')
0193 plt.ylabel("events")
0194 plt.title(f"$p_{{\pi^0}}$={p} GeV")
0195 plt.xlabel("$m^{\\pi^{0}}_{recon}$ [GeV]")
0196 else:
0197
0198 y, x = np.histogram(mrecon, bins=100, range=(0, 0.2))
0199
0200 bc=(x[1:]+x[:-1])/2
0201 from scipy.optimize import curve_fit
0202 slc=abs(bc-.135)<.1
0203 fnc=gauss
0204 p0=[100, .135, 0.2]
0205
0206 coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,
0207 sigma=list(np.sqrt(y[slc])+(y[slc]==0)))
0208 if p==100:
0209 xx=np.linspace(0,0.2)
0210 plt.plot(xx, fnc(xx,*coeff))
0211 pvals.append(p)
0212 resvals.append(np.abs(coeff[2]))
0213 dresvals.append(np.sqrt(var_matrix[2][2]))
0214
0215 plt.sca(axs[1])
0216 plt.errorbar(pvals, resvals, dresvals, ls='', marker='o')
0217 plt.ylim(0)
0218 plt.ylabel("$\\sigma[m_{\\pi^0}]$ [GeV]")
0219 plt.xlabel("$p_{\\pi^0}$ [GeV]")
0220
0221 fnc=lambda E,a,b: a+b*E
0222
0223 coeff, var_matrix = curve_fit(fnc, pvals, resvals, p0=(1,1),
0224 sigma=dresvals)
0225 xx=np.linspace(55, 200, 100)
0226
0227 plt.plot(xx, fnc(xx, *coeff), label=f'fit: $({coeff[0]*1000:.1f}+{coeff[1]*1000:.4f}\\times [E\,in\,GeV])$ MeV')
0228 plt.legend()
0229
0230
0231 plt.tight_layout()
0232 plt.savefig(outdir+"/pi0_mass_res.pdf")