zdc_photon/analysis/zdc_photon_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
0010 outdir=sys.argv[1]+"/"
0011 config=outdir.split("/")[1]
0012 try:
0013     import os
0014     os.mkdir(outdir[:-1])
0015 except:
0016     pass
0017
0018 def gauss(x, A,mu, sigma):
0019     return A * np.exp(-(x-mu)**2/(2*sigma**2))
0020
0021 #load the files
0022 import uproot as ur
0023 arrays_sim={}
0024 momenta=20, 30, 50, 70, 100, 150, 200, 275
0025 for p in momenta:
0026     arrays_sim[p] = ur.concatenate({
0027         f'sim_output/zdc_photon/{config}_rec_zdc_photon_{p}GeV_{index}.edm4eic.root': 'events'
0028         for index in range(5)
0029     })
0030
0031 if "ReconstructedFarForwardZDCNeutrals.PDG" not in arrays_sim[p][momenta[0]].fields:
0032     print("ReconstructedFarForwardZDCNeutrals collection is not available (needs EICrecon 1.23)")
0033     import sys
0034     sys.exit(0)
0035
0036 import awkward as ak
0037 fig,axs=plt.subplots(1,3, figsize=(24, 8))
0038 pvals=[]
0039 resvals=[]
0040 dresvals=[]
0041 scalevals=[]
0042 dscalevals=[]
0043 for p in momenta:
0044     selection=arrays_sim[p]["ReconstructedFarForwardZDCNeutrals.PDG"]==22
0045     Etot=arrays_sim[p]["ReconstructedFarForwardZDCNeutrals.energy"][selection]
0046     if p==100:
0047         plt.sca(axs[0])
0048         y, x, _=plt.hist(ak.flatten(Etot), bins=100, range=(p*.75, p*1.25), histtype='step')
0049         plt.ylabel("events")
0050         plt.title(f"$p_{{\\gamma}}$={p} GeV")
0051         plt.xlabel("$E^{\\gamma}_{recon}$ [GeV]")
0052     else:
0053         y, x = np.histogram(ak.flatten(Etot), bins=100, range=(p*.75, p*1.25))
0054
0055     bc=(x[1:]+x[:-1])/2
0056     from scipy.optimize import curve_fit
0057     slc=abs(bc-p)<10
0058     fnc=gauss
0059     p0=[100, p, 10]
0060     try:
0061         coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,
0062                                       sigma=list(np.sqrt(y[slc])+(y[slc]==0)), maxfev=10000)
0063         if p==100:
0064             xx=np.linspace(p*0.75,p*1.25, 100)
0065             plt.plot(xx, fnc(xx,*coeff))
0066         pvals.append(p)
0067         resvals.append(np.abs(coeff[2])/coeff[1])
0068         dresvals.append(np.sqrt(var_matrix[2][2])/coeff[1])
0069         scalevals.append(np.abs(coeff[1])/p)
0070         dscalevals.append(np.sqrt(var_matrix[2][2])/p)
0071     except RuntimeError as e:
0072         print(f"fit failed for p={p}", e, list(bc[slc]), list(y[slc]))
0073
0074 plt.sca(axs[1])
0075 plt.errorbar(pvals, resvals, dresvals, ls='', marker='o')
0076 plt.ylim(0)
0077 plt.ylabel("$\\sigma[E_{\\gamma}]/\\mu[E_{\\gamma}]$")
0078 plt.xlabel("$p_{\\gamma}$ [GeV]")
0079
0080 xx=np.linspace(15, 275, 100)
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
0088     xx=np.linspace(15, 275, 100)
0089     plt.plot(xx, fnc(xx, *coeff), label=f'fit:  $\\frac{{{coeff[0]*100:.0f}\\%}}{{\\sqrt{{E}}}}$')
0090 except RuntimeError as e:
0091     print("fit failed", e)
0092 plt.legend()
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_{\\gamma}]/E_{\\gamma}$")
0097 plt.xlabel("$p_{\\gamma}$ [GeV]")
0098 plt.axhline(1, ls='--', alpha=0.7, color='0.5')
0099 plt.tight_layout()
0100 plt.savefig(outdir+"photon_energy_res.pdf")
0101
0102 #theta res
0103 fig,axs=plt.subplots(1,2, figsize=(16, 8))
0104 pvals=[]
0105 resvals=[]
0106 dresvals=[]
0107 for p in momenta:
0108     selection=arrays_sim[p]["ReconstructedFarForwardZDCNeutrals.PDG"]==22
0109     px_recon=arrays_sim[p]["ReconstructedFarForwardZDCNeutrals.momentum.x"][selection]
0110     py_recon=arrays_sim[p]["ReconstructedFarForwardZDCNeutrals.momentum.y"][selection]
0111     pz_recon=arrays_sim[p]["ReconstructedFarForwardZDCNeutrals.momentum.z"][selection]
0112
0113     theta_recon=np.arctan2(np.hypot(px_recon*np.cos(-.025)-pz_recon*np.sin(-.025), py_recon), pz_recon*np.cos(-.025)+px_recon*np.sin(-.025))
0114
0115     px=arrays_sim[p]["MCParticles.momentum.x"][::,2]
0116     py=arrays_sim[p]["MCParticles.momentum.y"][::,2]
0117     pz=arrays_sim[p]["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     if p==100:
0121         plt.sca(axs[0])
0122         y, x, _=plt.hist(1000*ak.flatten(theta_recon-theta_truth), bins=100, range=(-0.5, 0.5), histtype='step')
0123         plt.ylabel("events")
0124         plt.title(f"$p_{{\\gamma}}$={p} GeV")
0125         plt.xlabel("$\\theta^{\\gamma}_{recon}$ [mrad]")
0126     else:
0127         y, x = np.histogram(1000*ak.flatten(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     try:
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     except:
0144         pass
0145 plt.sca(axs[1])
0146 plt.errorbar(pvals, resvals, dresvals, ls='', marker='o')
0147
0148 fnc=lambda E,a, b: np.hypot(a/np.sqrt(E), b)
0149 #pvals, resvals, dresvals
0150 coeff, var_matrix = curve_fit(fnc, pvals, resvals, p0=(1,.1),
0151                                  sigma=dresvals, maxfev=10000)
0152
0153 xx=np.linspace(15, 275, 100)
0154
0155 plt.plot(xx, fnc(xx, *coeff), label=f'fit:  $\\frac{{{coeff[0]:.2f}}}{{\\sqrt{{E}}}}\\oplus {coeff[1]:.3f}$ mrad')
0156
0157 plt.ylabel("$\\sigma[\\theta_{\\gamma}]$ [mrad]")
0158 plt.xlabel("$p_{\\gamma}$ [GeV]")
0159
0160 plt.ylim(0, 0.1)
0161 plt.legend()
0162 plt.tight_layout()
0163 plt.savefig(outdir+"photon_theta_res.pdf")