insert_neutron/analysis/neutron_plots.py

0001 import numpy as np, pandas as pd, matplotlib.pyplot as plt, matplotlib as mpl, awkward as ak, sys, uproot as ur
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 plt.rcParams["figure.figsize"] = (7, 7)
0010 config=sys.argv[1].split("/")[1]  #results/{config}/neutron
0011 outdir=sys.argv[1]+"/"
0012 try:
0013     import os
0014     os.mkdir(outdir[:-1])
0015 except:
0016     pass
0017
0018 #read files
0019 arrays_sim={}
0020 for p in 20,30,40,50,60,70,80:
0021     arrays_sim[p] = ur.open(f'sim_output/insert_neutron/{config}_rec_neutron_{p}GeV.edm4eic.root:events')\
0022                     .arrays()
0023
0024 def gauss(x, A,mu, sigma):
0025     return A * np.exp(-(x-mu)**2/(2*sigma**2))
0026
0027 #get the truth pseudorapidity and energy
0028 for array in arrays_sim.values():
0029     tilt=-0.025
0030     px=array['MCParticles.momentum.x'][:,2]
0031     py=array['MCParticles.momentum.y'][:,2]
0032     pz=array['MCParticles.momentum.z'][:,2]
0033     p=np.sqrt(px**2+py**2+pz**2)
0034
0035     pxp=px*np.cos(tilt)-pz*np.sin(tilt)
0036     pyp=py
0037     pzp=pz*np.cos(tilt)+px*np.sin(tilt)
0038
0039     array['E_truth']=np.hypot(p, 0.9406)
0040     array['eta_truth']=1/2*np.log((p+pzp)/(p-pzp))
0041     array['theta_truth']=np.arccos(pzp/p)
0042
0043 #
0044 # get reconstructed theta as avg of theta of cluster centers, weighted by energy
0045 for array in arrays_sim.values():
0046     tilt=-0.025
0047     x=array['HcalEndcapPInsertClusters.position.x']
0048     y=array['HcalEndcapPInsertClusters.position.y']
0049     z=array['HcalEndcapPInsertClusters.position.z']
0050     E=array['HcalEndcapPInsertClusters.energy']
0051     r=np.sqrt(x**2+y**2+z**2)
0052
0053     xp=x*np.cos(tilt)-z*np.sin(tilt)
0054     yp=y
0055     zp=z*np.cos(tilt)+x*np.sin(tilt)
0056
0057     w=E
0058
0059     array['theta_recon']=np.sum(np.arccos(zp/r)*w, axis=-1)/np.sum(w, axis=-1)
0060     array['eta_recon']=-np.log(np.tan(array['theta_recon']/2))
0061
0062
0063 #plot theta residuals:
0064 print("making theta recon plot")
0065 from scipy.optimize import curve_fit
0066
0067 fig, axs=plt.subplots(1,2, figsize=(16,8))
0068 plt.sca(axs[0])
0069 p=40
0070 eta_min=3.4; eta_max=3.6
0071 y,x,_=plt.hist(1000*(arrays_sim[p]['theta_recon']-arrays_sim[p]['theta_truth'])\
0072                [(arrays_sim[p]['eta_truth']>eta_min)&(arrays_sim[p]['eta_truth']<eta_max)], bins=50,
0073                     range=(-10,10), histtype='step')
0074 bc=(x[1:]+x[:-1])/2
0075 slc=abs(bc)<3
0076 # try:
0077 fnc=gauss
0078 sigma=np.sqrt(y[slc])+(y[slc]==0)
0079 p0=(100, 0, 5)
0080 coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,sigma=list(sigma))
0081 xx=np.linspace(-5,5,100)
0082 plt.plot(xx,fnc(xx,*coeff))
0083 # except:
0084 #     pass
0085 plt.xlabel("$\\theta_{rec}-\\theta_{truth}$ [mrad]")
0086 plt.ylabel("events")
0087 plt.title(f"$p={p}$ GeV, ${eta_min}<\\eta<{eta_max}$")
0088
0089 r=[3.2,3.4,3.6,3.8,4.0]
0090 for eta_min, eta_max in zip(r[:-1],r[1:]):
0091     xvals=[]
0092     sigmas=[]
0093     dsigmas=[]
0094     for p in 20,30,40, 50, 60, 70, 80:
0095         y,x=np.histogram(1000*(arrays_sim[p]['theta_recon']-arrays_sim[p]['theta_truth'])\
0096                          [(arrays_sim[p]['eta_truth']>eta_min)&(arrays_sim[p]['eta_truth']<eta_max)],
0097                          bins=50, range=(-10,10))
0098         bc=(x[1:]+x[:-1])/2
0099         slc=abs(bc)<3
0100         fnc=gauss
0101         p0=(100, 0, 5)
0102         #print(bc[slc],y[slc])
0103         sigma=np.sqrt(y[slc])+(y[slc]==0)
0104         try:
0105             coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,sigma=list(sigma))
0106             sigmas.append(np.abs(coeff[2]))
0107             dsigmas.append(np.sqrt(var_matrix[2][2]))
0108             xvals.append(p)
0109         except:
0110             pass
0111     plt.sca(axs[1])
0112     plt.errorbar(xvals, sigmas, dsigmas, ls='', marker='o', label=f"${eta_min}<\\eta<{eta_max}$")
0113 plt.xlabel("$p_{n}$ [GeV]")
0114 plt.ylabel("$\\sigma[\\theta]$ [mrad]")
0115 plt.ylim(0)
0116 plt.legend()
0117 plt.tight_layout()
0118 plt.savefig(outdir+"neutron_theta_recon.pdf")
0119
0120 #now determine the energy recon parameters
0121 pvals=[]
0122 resvals=[]
0123 reserrs=[]
0124 wvals=[]
0125 svals=[]
0126 Euncorr=[]
0127
0128 print("determining the energy recon correction parameters")
0129 fig,axs=plt.subplots(1,2, figsize=(20,10))
0130 eta_min=3.4;eta_max=3.6
0131 for p in 20, 30,40,50,60,70, 80:
0132     best_res=1000
0133     res_err=1000
0134     best_s=1000
0135     wrange=np.linspace(30, 70, 41)*0.0257
0136     coeff_best=None
0137
0138     wbest=0
0139     a=arrays_sim[p]
0140     h=np.sum(a[f'HcalEndcapPInsertClusters.energy'], axis=-1)
0141     e=np.sum(a[f'EcalEndcapPInsertClusters.energy'], axis=-1)
0142     for w in wrange:
0143
0144         r=(e/w+h)[(h>0)&(a['eta_truth']>eta_min)&(a['eta_truth']<eta_max)]
0145         y,x=np.histogram(r,bins=50)
0146         bcs=(x[1:]+x[:-1])/2
0147         fnc=gauss
0148         slc=abs(bcs-np.mean(r)*1.25)<2*np.std(r)
0149         sigma=np.sqrt(y[slc])+0.5*(y[slc]==0)
0150         p0=(100, np.mean(r), np.std(r))
0151         try:
0152             coeff, var_matrix = curve_fit(fnc, list(bcs[slc]), list(y[slc]), p0=p0,sigma=list(sigma))
0153             res=np.abs(coeff[2]/coeff[1])
0154
0155             if res<best_res:
0156                 best_res=res
0157                 coeff_best=coeff
0158                 res_err=res*np.hypot(np.sqrt(var_matrix[2][2])/coeff[2], np.sqrt(var_matrix[1][1])/coeff[1])
0159                 wbest=w
0160                 best_s=np.hypot(p,0.9406)/coeff[1]
0161                 Euncorr_best=coeff[1]
0162         except :
0163             print("fit failed")
0164
0165     if p==50:
0166         r=(e/wbest+h)[(h>0)&(a['eta_truth']>3.4)&(a['eta_truth']<3.6)]
0167         plt.sca(axs[0])
0168         y, x, _= plt.hist(r, histtype='step', bins=50)
0169         xx=np.linspace(20, 55, 100)
0170         plt.plot(xx,fnc(xx, *coeff_best), ls='-')
0171         plt.xlabel("$E_{uncorr}=E_{Hcal}+E_{Ecal}/w$ [GeV]")
0172         plt.title(f"p=50 GeV, ${eta_min}<\\eta<{eta_max}$, w={wbest:.2f}")
0173         plt.axvline(np.sqrt(50**2+.9406**2), color='g', ls=':')
0174         plt.text(40, max(y)*0.9, "generated\nenergy", color='g', fontsize=20)
0175
0176     Euncorr.append(Euncorr_best)
0177     resvals.append(best_res)
0178     reserrs.append(res_err)
0179     pvals.append(p)
0180     svals.append(best_s)
0181     wvals.append(wbest)
0182
0183 pvals=np.array(pvals)
0184 svals=np.array(svals)
0185 Euncorr=np.array(Euncorr)
0186 plt.sca(axs[1])
0187 plt.plot(Euncorr,wvals, label="w")
0188 w_avg=np.mean(wvals)
0189 plt.axhline(w_avg, label=f'w avg: {w_avg:.2f}', ls=':')
0190 plt.plot(Euncorr,svals, label="s")
0191 m=(np.sum(svals*Euncorr)*len(Euncorr)-np.sum(Euncorr)*np.sum(svals))/(np.sum(Euncorr**2)*len(Euncorr)-np.sum(Euncorr)**2)
0192 b=np.mean(svals)-np.mean(Euncorr)*m
0193 plt.plot(Euncorr,Euncorr*m+b, label=f"s fit: ${m:.4f}E_{{uncorr}}+{b:.2f}$", ls=':')
0194
0195 plt.xlabel("$E_{uncorr}=E_{Hcal}+E_{Ecal}/w$ [GeV]")
0196 plt.title("$E_{n,recon}=s\\times(E_{Hcal}+E_{Ecal}/w)$")
0197 plt.ylabel('parameter values')
0198 plt.legend()
0199 plt.ylim(0)
0200 plt.savefig(outdir+"neutron_energy_params.pdf")
0201
0202 #now make the energy plot
0203 print("making energy recon plot")
0204 fig, axs=plt.subplots(1,3, figsize=(24,8))
0205 partitions=[3.2,3.4, 3.6, 3.8, 4.0]
0206 for eta_min, eta_max in zip(partitions[:-1],partitions[1:]):
0207     pvals=[]
0208     resvals=[]
0209     reserrs=[]
0210     scalevals=[]
0211     scaleerrs=[]
0212     for p in 20, 30,40,50,60,70, 80:
0213         best_res=1000
0214         res_err=1000
0215
0216
0217         wrange=np.linspace(30, 70, 30)*0.0257
0218
0219         w=w_avg
0220         a=arrays_sim[p]
0221         h=np.sum(a[f'HcalEndcapPInsertClusters.energy'], axis=-1)
0222         e=np.sum(a[f'EcalEndcapPInsertClusters.energy'], axis=-1)
0223         #phi=a['phi_truth']
0224         uncorr=(e/w+h)
0225         s=-0.0064*uncorr+1.80
0226         r=uncorr*s #reconstructed energy with correction
0227         r=r[(h>0)&(a['eta_truth']>eta_min)&(a['eta_truth']<eta_max)]#&(abs(phi)>np.pi/2)]
0228         y,x=np.histogram(r,bins=50)
0229         bcs=(x[1:]+x[:-1])/2
0230         fnc=gauss
0231         slc=abs(bcs-np.mean(r)*1.25)<2*np.std(r)
0232         sigma=np.sqrt(y[slc])+0.5*(y[slc]==0)
0233         p0=(100, np.mean(r), np.std(r))
0234         try:
0235             coeff, var_matrix = curve_fit(fnc, list(bcs[slc]), list(y[slc]), p0=p0,sigma=list(sigma))
0236             res=np.abs(coeff[2]/coeff[1])
0237
0238             if res<best_res:
0239                 best_res=res
0240                 res_err=res*np.hypot(np.sqrt(var_matrix[2][2])/coeff[2], np.sqrt(var_matrix[1][1])/coeff[1])
0241                 wbest=w
0242                 scale=coeff[1]/np.sqrt(p**2+0.9406**2)
0243                 dscale=np.sqrt(var_matrix[1][1]/np.sqrt(p**2+0.9406**2))
0244         except :
0245             print("fit failed")
0246         if p==50 and eta_min==3.4:
0247             plt.sca(axs[0])
0248             plt.errorbar(bcs, y, np.sqrt(y)+(y==0),marker='o', ls='', )
0249             plt.title(f'p={p} GeV, ${eta_min}<\\eta<{eta_max}$')
0250             plt.xlabel("$E_{recon}$ [GeV]")
0251             plt.ylabel("events")
0252             #plt.ylim(0)
0253             xx=np.linspace(40, 70, 50)
0254             plt.plot(xx, fnc(xx, *coeff))
0255         resvals.append(best_res)
0256         reserrs.append(res_err)
0257         scalevals.append(scale)
0258         scaleerrs.append(dscale)
0259         pvals.append(p)
0260     plt.sca(axs[1])
0261     plt.errorbar(pvals, resvals, reserrs, marker='o', ls='', label=f"${eta_min}<\\eta<{eta_max}$")
0262     #plt.ylim(0)
0263     plt.ylabel("$\\sigma[E]/\\mu[E]$")
0264     plt.xlabel("$p_{n}$ [GeV]")
0265
0266     plt.sca(axs[2])
0267     plt.errorbar(pvals, scalevals, scaleerrs, marker='o', ls='', label=f"${eta_min}<\\eta<{eta_max}$")
0268
0269
0270     plt.ylabel("$\\mu[E]/E$")
0271
0272
0273 axs[2].set_xlabel("$p_n$ [GeV]")
0274 axs[2].axhline(1, ls='--', color='0.5', alpha=0.7)
0275 axs[0].set_ylim(0)
0276 axs[1].set_ylim(0, 0.35)
0277 axs[2].set_ylim(0)
0278 axs[1].legend()
0279 axs[2].legend()
0280 plt.tight_layout()
0281 plt.savefig(outdir+"neutron_energy_recon.pdf")