zdc_pi0/analysis/zdc_pi0_plots.py

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     arrays_sim[p] = ur.concatenate({
0025         f'sim_output/zdc_pi0/{config}_rec_zdc_pi0_{p}GeV_{index}.edm4eic.root': 'events'
0026         for index in range(5)
0027     })
0028
0029 #energy res plot
0030 fig,axs=plt.subplots(1,3, figsize=(24, 8))
0031 pvals=[]
0032 resvals=[]
0033 dresvals=[]
0034 scalevals=[]
0035 dscalevals=[]
0036 for p in momenta:
0037     selection=[len(arrays_sim[p]["HcalFarForwardZDCClusters.energy"][i])==2 for i in range(len(arrays_sim[p]))]
0038     E=arrays_sim[p][selection]["HcalFarForwardZDCClusters.energy"]
0039
0040     Etot=np.sum(E, axis=-1)
0041     if len(Etot)<25:
0042         continue
0043     #print(p, res, mrecon)
0044     if p==100:
0045         plt.sca(axs[0])
0046         y, x, _=plt.hist(Etot, bins=100, range=(p*.5, p*1.5), histtype='step')
0047         plt.ylabel("events")
0048         plt.title(f"$p_{{\pi^0}}$={p} GeV")
0049         plt.xlabel("$E^{\\pi^{0}}_{recon}$ [GeV]")
0050     else:
0051         y, x = np.histogram(Etot, bins=100, range=(p*.5, p*1.5))
0052
0053     bc=(x[1:]+x[:-1])/2
0054     from scipy.optimize import curve_fit
0055     slc=abs(bc-p)<10
0056     fnc=gauss
0057     p0=[100, p, 10]
0058     #print(list(y), list(x))
0059     coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,
0060                                  sigma=list(np.sqrt(y[slc])+(y[slc]==0)), maxfev=10000)
0061     if p==100:
0062         xx=np.linspace(p*0.5,p*1.5, 100)
0063         plt.plot(xx, fnc(xx,*coeff))
0064     pvals.append(p)
0065     resvals.append(np.abs(coeff[2])/coeff[1])
0066     dresvals.append(np.sqrt(var_matrix[2][2])/coeff[1])
0067     scalevals.append(np.abs(coeff[1])/p)
0068     dscalevals.append(np.sqrt(var_matrix[2][2])/p)
0069
0070 plt.sca(axs[1])
0071 plt.errorbar(pvals, resvals, dresvals, ls='', marker='o')
0072
0073 plt.ylabel("$\\sigma[E_{\\pi^0}]/\\mu[E_{\\pi^0}]$")
0074 plt.xlabel("$p_{\\pi^0}$ [GeV]")
0075
0076 fnc=lambda E,a: a/np.sqrt(E)
0077 #pvals, resvals, dresvals
0078 coeff, var_matrix = curve_fit(fnc, pvals, resvals, p0=(1,),
0079                                  sigma=dresvals, maxfev=10000)
0080 xx=np.linspace(55, 200, 100)
0081 plt.plot(xx, fnc(xx, *coeff), label=f'fit:  $\\frac{{{coeff[0]:.2f}\\%}}{{\\sqrt{{E}}}}$')
0082 plt.legend()
0083 plt.ylim(0)
0084 plt.sca(axs[2])
0085 plt.errorbar(pvals, scalevals, dscalevals, ls='', marker='o')
0086 plt.ylim(0.8, 1.2)
0087 plt.ylabel("$\\mu[E_{\\pi^0}]/E_{\\pi^0}$")
0088 plt.xlabel("$p_{\\pi^0}$ [GeV]")
0089 plt.axhline(1, ls='--', alpha=0.7, color='0.5')
0090 plt.tight_layout()
0091 plt.savefig(outdir+"/pi0_energy_res.pdf")
0092
0093
0094 fig,axs=plt.subplots(1,2, figsize=(16, 8))
0095 pvals=[]
0096 resvals=[]
0097 dresvals=[]
0098 for p in momenta:
0099     selection=[len(arrays_sim[p]["HcalFarForwardZDCClusters.energy"][i])==2 for i in range(len(arrays_sim[p]))]
0100     x=arrays_sim[p][selection]["HcalFarForwardZDCClusters.position.x"]
0101     y=arrays_sim[p][selection]["HcalFarForwardZDCClusters.position.y"]
0102     z=arrays_sim[p][selection]["HcalFarForwardZDCClusters.position.z"]
0103     E=arrays_sim[p][selection]["HcalFarForwardZDCClusters.energy"]
0104     r=np.sqrt(x**2+y**2+z**2)
0105     px=np.sum(E*x/r, axis=-1)
0106     py=np.sum(E*y/r, axis=-1)
0107     pz=np.sum(E*z/r, axis=-1)
0108
0109     theta_recon=np.arctan2(np.hypot(px*np.cos(-.025)-pz*np.sin(-.025), py), pz*np.cos(-.025)+px*np.sin(-.025))
0110     if len(theta_recon)<25:
0111         continue
0112     px=arrays_sim[p][selection]["MCParticles.momentum.x"][::,2]
0113     py=arrays_sim[p][selection]["MCParticles.momentum.y"][::,2]
0114     pz=arrays_sim[p][selection]["MCParticles.momentum.z"][::,2]
0115
0116     theta_truth=np.arctan2(np.hypot(px*np.cos(-.025)-pz*np.sin(-.025), py), pz*np.cos(-.025)+px*np.sin(-.025))
0117
0118     Etot=np.sum(E, axis=-1)
0119     #print(p, res, mrecon)
0120     if p==100:
0121         plt.sca(axs[0])
0122         y, x, _=plt.hist(1000*(theta_recon-theta_truth), bins=100, range=(-0.5, 0.5), histtype='step')
0123         plt.ylabel("events")
0124         plt.title(f"$p_{{\\pi^0}}$={p} GeV")
0125         plt.xlabel("$\\theta^{\\pi^0}_{recon}$ [mrad]")
0126     else:
0127         y, x = np.histogram(1000*(theta_recon-theta_truth), bins=100, range=(-0.5, 0.5))
0128
0129     bc=(x[1:]+x[:-1])/2
0130     from scipy.optimize import curve_fit
0131     slc=abs(bc)<0.2#1.5*np.std(1000*(theta_recon-theta_truth))
0132     fnc=gauss
0133     p0=[100, 0, 0.1]
0134     #print(list(y), list(x))
0135     coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,
0136                                  sigma=list(np.sqrt(y[slc])+(y[slc]==0)), maxfev=10000)
0137     if p==100:
0138         xx=np.linspace(-0.5,0.5, 100)
0139         plt.plot(xx, fnc(xx,*coeff))
0140     pvals.append(p)
0141     resvals.append(np.abs(coeff[2]))
0142     dresvals.append(np.sqrt(var_matrix[2][2]))
0143
0144 plt.sca(axs[1])
0145 plt.errorbar(pvals, resvals, dresvals, ls='', marker='o')
0146 #print(dresvals)
0147
0148 fnc=lambda E,a: a/np.sqrt(E)
0149 #pvals, resvals, dresvals
0150 coeff, var_matrix = curve_fit(fnc, pvals, resvals, p0=(1,),
0151                                  sigma=dresvals, maxfev=10000)
0152
0153 xx=np.linspace(55, 200, 100)
0154
0155 plt.plot(xx, fnc(xx, *coeff), label=f'fit:  $\\frac{{{coeff[0]:.2f}}}{{\\sqrt{{E}}}}$ mrad')
0156
0157 plt.ylabel("$\\sigma[\\theta_{\\pi^0}]$ [mrad]")
0158 plt.xlabel("$p_{\\pi^0}$ [GeV]")
0159
0160 plt.ylim(0, 0.1)
0161 plt.legend()
0162 plt.tight_layout()
0163 plt.savefig(outdir+"/pi0_theta_res.pdf")
0164
0165 fig,axs=plt.subplots(1,2, figsize=(16, 8))
0166 pvals=[]
0167 resvals=[]
0168 dresvals=[]
0169 for p in momenta:
0170     selection=[len(arrays_sim[p]["HcalFarForwardZDCClusters.energy"][i])==2 for i in range(len(arrays_sim[p]))]
0171     E=arrays_sim[p][selection]["HcalFarForwardZDCClusters.energy"]
0172     cx=arrays_sim[p][selection]["HcalFarForwardZDCClusters.position.x"]
0173     cy=arrays_sim[p][selection]["HcalFarForwardZDCClusters.position.y"]
0174     cz=arrays_sim[p][selection]["HcalFarForwardZDCClusters.position.z"]
0175     r=np.sqrt(cx**2+cy**2+cz**2)
0176     px=E*cx/r
0177     py=E*cy/r
0178     pz=E*cz/r
0179
0180     cos_opening_angle=(cx/r)[::,0]*(cx/r)[::,1]+(cy/r)[::,0]*(cy/r)[::,1]+(cz/r)[::,0]*(cz/r)[::,1]
0181     mrecon=np.sqrt(2*E[::,0]*E[::,1]*(1-cos_opening_angle))
0182
0183     if len(mrecon)<25:
0184         continue
0185
0186     #print(p, res, mrecon)
0187     if p==100:
0188         plt.sca(axs[0])
0189         y, x, _=plt.hist(mrecon, bins=100, range=(0, 0.2), histtype='step')
0190         plt.ylabel("events")
0191         plt.title(f"$p_{{\pi^0}}$={p} GeV")
0192         plt.xlabel("$m^{\\pi^{0}}_{recon}$ [GeV]")
0193     else:
0194         #y, x, _=plt.hist(mrecon, bins=100, range=(0, 0.2), histtype='step')#y, x =np.histogram(mrecon, bins=100, range=(0, 0.2))
0195         y, x = np.histogram(mrecon, bins=100, range=(0, 0.2))
0196
0197     bc=(x[1:]+x[:-1])/2
0198     from scipy.optimize import curve_fit
0199     slc=abs(bc-.135)<.1
0200     fnc=gauss
0201     p0=[100, .135, 0.2]
0202     #print(list(y), list(x))
0203     coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,
0204                                  sigma=list(np.sqrt(y[slc])+(y[slc]==0)), maxfev=10000)
0205     if p==100:
0206         xx=np.linspace(0,0.2)
0207         plt.plot(xx, fnc(xx,*coeff))
0208     pvals.append(p)
0209     resvals.append(np.abs(coeff[2]))
0210     dresvals.append(np.sqrt(var_matrix[2][2]))
0211
0212 plt.sca(axs[1])
0213 plt.errorbar(pvals, resvals, dresvals, ls='', marker='o')
0214 plt.ylim(0)
0215 plt.ylabel("$\\sigma[m_{\\pi^0}]$ [GeV]")
0216 plt.xlabel("$p_{\\pi^0}$ [GeV]")
0217
0218 fnc=lambda E,a,b: a+b*E
0219 #pvals, resvals, dresvals
0220 coeff, var_matrix = curve_fit(fnc, pvals, resvals, p0=(1,1),
0221                                  sigma=dresvals, maxfev=10000)
0222 xx=np.linspace(55, 200, 100)
0223 #plt.plot(xx, fnc(xx, *coeff), label=f'fit:  ${coeff[0]*1000:.1f}+{coeff[1]*1000:.4f}\\times E$ MeV')
0224 plt.plot(xx, fnc(xx, *coeff), label=f'fit:  $({coeff[0]*1000:.1f}+{coeff[1]*1000:.4f}\\times [E\,in\,GeV])$ MeV')
0225 plt.legend()
0226
0227
0228 plt.tight_layout()
0229 plt.savefig(outdir+"/pi0_mass_res.pdf")