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