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