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     filename=f'sim_output/zdc_photon/{config}_rec_zdc_photon_{p}GeV.edm4hep.root'
0027     print("opening file", filename)
0028     events = ur.open(filename+':events')
0029     arrays_sim[p] = events.arrays()#[:-1] #remove last event, which for some reason is blank
0030     import gc
0031     gc.collect()
0032     print("read", filename)
0033
0034 fig,axs=plt.subplots(1,3, figsize=(24, 8))
0035 pvals=[]
0036 resvals=[]
0037 dresvals=[]
0038 scalevals=[]
0039 dscalevals=[]
0040 for p in momenta:
0041     selection=[len(arrays_sim[p]["HcalFarForwardZDCClusters.energy"][i])==1 for i in range(len(arrays_sim[p]))]
0042     E=arrays_sim[p][selection]["HcalFarForwardZDCClusters.energy"]
0043
0044     Etot=np.sum(E, axis=-1)
0045     #print(len(Etot))
0046     #print(p, res, mrecon)
0047     if p==100:
0048         plt.sca(axs[0])
0049         y, x, _=plt.hist(Etot, bins=100, range=(p*.75, p*1.25), histtype='step')
0050         plt.ylabel("events")
0051         plt.title(f"$p_{{\gamma}}$={p} GeV")
0052         plt.xlabel("$E^{\\gamma}_{recon}$ [GeV]")
0053     else:
0054         y, x = np.histogram(Etot, bins=100, range=(p*.75, p*1.25))
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     #print(list(y), list(x))
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.75,p*1.25, 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 plt.ylim(0)
0076 plt.ylabel("$\\sigma[E_{\\gamma}]/\\mu[E_{\\gamma}]$")
0077 plt.xlabel("$p_{\\gamma}$ [GeV]")
0078
0079 xx=np.linspace(15, 275, 100)
0080
0081 fnc=lambda E,a: a/np.sqrt(E)
0082 #pvals, resvals, dresvals
0083 coeff, var_matrix = curve_fit(fnc, pvals, resvals, p0=(1,),
0084                                  sigma=dresvals)
0085
0086 xx=np.linspace(15, 275, 100)
0087 plt.plot(xx, fnc(xx, *coeff), label=f'fit:  $\\frac{{{coeff[0]*100:.0f}\\%}}{{\\sqrt{{E}}}}$')
0088 plt.legend()
0089 plt.sca(axs[2])
0090 plt.errorbar(pvals, scalevals, dscalevals, ls='', marker='o')
0091 plt.ylim(0.8, 1.2)
0092 plt.ylabel("$\\mu[E_{\\gamma}]/E_{\\gamma}$")
0093 plt.xlabel("$p_{\\gamma}$ [GeV]")
0094 plt.axhline(1, ls='--', alpha=0.7, color='0.5')
0095 plt.tight_layout()
0096 plt.savefig(outdir+"photon_energy_res.pdf")
0097
0098 #theta res
0099 fig,axs=plt.subplots(1,2, figsize=(16, 8))
0100 pvals=[]
0101 resvals=[]
0102 dresvals=[]
0103 for p in momenta:
0104     selection=[len(arrays_sim[p]["HcalFarForwardZDCClusters.energy"][i])==1 for i in range(len(arrays_sim[p]))]
0105     x=arrays_sim[p][selection]["HcalFarForwardZDCClusters.position.x"][::,0]
0106     y=arrays_sim[p][selection]["HcalFarForwardZDCClusters.position.y"][::,0]
0107     z=arrays_sim[p][selection]["HcalFarForwardZDCClusters.position.z"][::,0]
0108
0109     theta_recon=np.arctan2(np.hypot(x*np.cos(-.025)-z*np.sin(-.025), y), z*np.cos(-.025)+x*np.sin(-.025))
0110
0111     px=arrays_sim[p][selection]["MCParticles.momentum.x"][::,2]
0112     py=arrays_sim[p][selection]["MCParticles.momentum.y"][::,2]
0113     pz=arrays_sim[p][selection]["MCParticles.momentum.z"][::,2]
0114
0115     theta_truth=np.arctan2(np.hypot(px*np.cos(-.025)-pz*np.sin(-.025), py), pz*np.cos(-.025)+px*np.sin(-.025))
0116
0117     Etot=np.sum(E, axis=-1)
0118     #print(p, res, mrecon)
0119     if p==100:
0120         plt.sca(axs[0])
0121         y, x, _=plt.hist(1000*(theta_recon-theta_truth), bins=100, range=(-0.5, 0.5), histtype='step')
0122         plt.ylabel("events")
0123         plt.title(f"$p_{{\gamma}}$={p} GeV")
0124         plt.xlabel("$\\theta^{\\gamma}_{recon}$ [mrad]")
0125     else:
0126         y, x = np.histogram(1000*(theta_recon-theta_truth), bins=100, range=(-0.5, 0.5))
0127
0128     bc=(x[1:]+x[:-1])/2
0129     from scipy.optimize import curve_fit
0130     slc=abs(bc)<0.2#1.5*np.std(1000*(theta_recon-theta_truth))
0131     fnc=gauss
0132     p0=[100, 0, 0.1]
0133     #print(list(y), list(x))
0134     coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,
0135                                  sigma=list(np.sqrt(y[slc])+(y[slc]==0)))
0136     if p==100:
0137         xx=np.linspace(-0.5,0.5, 100)
0138         plt.plot(xx, fnc(xx,*coeff))
0139     pvals.append(p)
0140     resvals.append(np.abs(coeff[2]))
0141     dresvals.append(np.sqrt(var_matrix[2][2]))
0142 plt.sca(axs[1])
0143 plt.errorbar(pvals, resvals, dresvals, ls='', marker='o')
0144 #print(dresvals)
0145
0146 fnc=lambda E,a, b: np.hypot(a/np.sqrt(E), b)
0147 #pvals, resvals, dresvals
0148 coeff, var_matrix = curve_fit(fnc, pvals, resvals, p0=(1,.1),
0149                                  sigma=dresvals)
0150
0151 xx=np.linspace(15, 275, 100)
0152
0153 plt.plot(xx, fnc(xx, *coeff), label=f'fit:  $\\frac{{{coeff[0]:.2f}}}{{\\sqrt{{E}}}}\\oplus {coeff[1]:.3f}$ mrad')
0154
0155 plt.ylabel("$\\sigma[\\theta_{\\gamma}]$ [mrad]")
0156 plt.xlabel("$p_{\\gamma}$ [GeV]")
0157
0158 plt.ylim(0, 0.1)
0159 plt.legend()
0160 plt.tight_layout()
0161 plt.savefig(outdir+"photon_theta_res.pdf")