diff --git a/benchmarks/neutron/analysis/neutron_plots.py b/benchmarks/neutron/analysis/neutron_plots.py
index 41583b9c9e592e96b2aee498cbda9e8a14355b2d..fc1d5b30e6249a1ee5cc9ee51a1b1cce4225f238 100644
--- a/benchmarks/neutron/analysis/neutron_plots.py
+++ b/benchmarks/neutron/analysis/neutron_plots.py
@@ -21,7 +21,6 @@ for p in 20,30,40,50,60,70,80:
arrays_sim[p] = ur.open(f'sim_output/neutron/{config}_rec_neutron_{p}GeV.edm4hep.root:events')\
.arrays()
-#get reconstructed theta, eta, and E
def gauss(x, A,mu, sigma):
return A * np.exp(-(x-mu)**2/(2*sigma**2))
@@ -41,7 +40,8 @@ for array in arrays_sim.values():
array['eta_truth']=1/2*np.log((p+pzp)/(p-pzp))
array['theta_truth']=np.arccos(pzp/p)
-#weighted avg of theta of cluster centers, by energy
+#
+# get reconstructed theta as avg of theta of cluster centers, weighted by energy
for array in arrays_sim.values():
tilt=-0.025
x=array['HcalEndcapPInsertClusters.position.x']
@@ -59,14 +59,9 @@ for array in arrays_sim.values():
array['theta_recon']=np.sum(np.arccos(zp/r)*w, axis=-1)/np.sum(w, axis=-1)
array['eta_recon']=-np.log(np.tan(array['theta_recon']/2))
-
- array['E_Hcal']=np.sum(array['HcalEndcapPInsertClusters.energy'], axis=-1)#*20/12.5
- array['E_Ecal']=np.sum(array['EcalEndcapPInsertClusters.energy'], axis=-1)#*27/20
-
- array['E_recon']=array['E_Hcal']/(0.70-.30/np.sqrt(array['E_Hcal']))\
- +array['E_Ecal']/(0.82-0.35/np.sqrt(array['E_Ecal']))
#plot theta residuals:
+print("making theta recon plot")
from scipy.optimize import curve_fit
fig, axs=plt.subplots(1,2, figsize=(16,8))
@@ -91,7 +86,7 @@ plt.xlabel("$\\theta_{rec}-\\theta_{truth}$ [mrad]")
plt.ylabel("events")
plt.title(f"$p={p}$ GeV, ${eta_min}<\\eta<{eta_max}$")
-r=[3.0,3.2,3.4,3.6,3.8]
+r=[3.2,3.4,3.6,3.8,4.0]
for eta_min, eta_max in zip(r[:-1],r[1:]):
xvals=[]
sigmas=[]
@@ -122,54 +117,165 @@ plt.legend()
plt.tight_layout()
plt.savefig(outdir+"neutron_theta_recon.pdf")
-fig, axs=plt.subplots(1,2, figsize=(16,8))
-plt.sca(axs[0])
-p=50
-eta_min=3.4; eta_max=3.6
-y,x,_=plt.hist(((arrays_sim[p]['E_recon']-arrays_sim[p]['E_truth'])/arrays_sim[p]['E_truth'])\
- [(arrays_sim[p]['eta_truth']>eta_min)&(arrays_sim[p]['eta_truth']<eta_max)&(arrays_sim[p]['E_Hcal']>0)], bins=50,
- range=(-.5,.5), histtype='step')
-bc=(x[1:]+x[:-1])/2
-slc=abs(bc)<0.4
-# try:
-p0=(100, 0, 0.3)
-fnc=gauss
-sigma=np.sqrt(y[slc])+(y[slc]==0)
-coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,sigma=list(sigma))
-xx=np.linspace(-.5,.5,100)
-plt.plot(xx,fnc(xx,*coeff))
-# except:
-# pass
-plt.xlabel("$(E_{rec}-E_{truth})/E_{truth}$")
-plt.ylabel("events")
-plt.title(f"$p={p}$ GeV, ${eta_min}<\\eta<{eta_max}$")
+#now determine the energy recon parameters
+pvals=[]
+resvals=[]
+reserrs=[]
+wvals=[]
+svals=[]
+Euncorr=[]
-r=[3.0,3.2,3.4,3.6,3.8]
-for eta_min, eta_max in zip(r[:-1],r[1:]):
- xvals=[]
- sigmas=[]
- dsigmas=[]
- for p in 20,30,40, 50, 60, 70, 80:
- y,x=np.histogram(((arrays_sim[p]['E_recon']-arrays_sim[p]['E_truth'])/arrays_sim[p]['E_truth'])\
- [(arrays_sim[p]['eta_truth']>eta_min)&(arrays_sim[p]['eta_truth']<eta_max)], bins=50, range=(-.5,.5))
- bc=(x[1:]+x[:-1])/2
- slc=abs(bc)<0.5
+print("determining the energy recon correction parameters")
+fig,axs=plt.subplots(1,2, figsize=(20,10))
+eta_min=3.4;eta_max=3.6
+for p in 20, 30,40,50,60,70, 80:
+ best_res=1000
+ res_err=1000
+ best_s=1000
+ wrange=np.linspace(30, 70, 41)*0.0257
+ coeff_best=None
+
+ wbest=0
+ a=arrays_sim[p]
+ h=np.sum(a[f'HcalEndcapPInsertClusters.energy'], axis=-1)
+ e=np.sum(a[f'EcalEndcapPInsertClusters.energy'], axis=-1)
+ for w in wrange:
+
+ r=(e/w+h)[(h>0)&(a['eta_truth']>eta_min)&(a['eta_truth']<eta_max)]
+ y,x=np.histogram(r,bins=50)
+ bcs=(x[1:]+x[:-1])/2
fnc=gauss
- p0=(100, 0, 0.3)
- #print(bc[slc],y[slc])
- sigma=np.sqrt(y[slc])+(y[slc]==0)
+ slc=abs(bcs-np.mean(r)*1.25)<2*np.std(r)
+ sigma=np.sqrt(y[slc])+0.5*(y[slc]==0)
+ p0=(100, np.mean(r), np.std(r))
try:
- coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,sigma=list(sigma))
- sigmas.append(np.abs(coeff[2]))
- dsigmas.append(np.sqrt(var_matrix[2][2]))
- xvals.append(p)
- except:
- pass
- plt.sca(axs[1])
- plt.errorbar(xvals, sigmas, dsigmas, ls='', marker='o', label=f"${eta_min}<\\eta<{eta_max}$")
-plt.xlabel("$p_{n}$ [GeV]")
-plt.ylabel("$\\sigma[E]/E$")
-plt.ylim(0)
+ coeff, var_matrix = curve_fit(fnc, list(bcs[slc]), list(y[slc]), p0=p0,sigma=list(sigma))
+ res=np.abs(coeff[2]/coeff[1])
+
+ if res<best_res:
+ best_res=res
+ coeff_best=coeff
+ res_err=res*np.hypot(np.sqrt(var_matrix[2][2])/coeff[2], np.sqrt(var_matrix[1][1])/coeff[1])
+ wbest=w
+ best_s=np.hypot(p,0.9406)/coeff[1]
+ Euncorr_best=coeff[1]
+ except :
+ print("fit failed")
+
+ if p==50:
+ r=(e/wbest+h)[(h>0)&(a['eta_truth']>3.4)&(a['eta_truth']<3.6)]
+ plt.sca(axs[0])
+ y, x, _= plt.hist(r, histtype='step', bins=50)
+ xx=np.linspace(20, 55, 100)
+ plt.plot(xx,fnc(xx, *coeff_best), ls='-')
+ plt.xlabel("$E_{uncorr}=E_{Hcal}+E_{Ecal}/w$ [GeV]")
+ plt.title(f"p=50 GeV, ${eta_min}<\\eta<{eta_max}$, w={wbest:.2f}")
+ plt.axvline(np.sqrt(50**2+.9406**2), color='g', ls=':')
+ plt.text(40, max(y)*0.9, "generated\nenergy", color='g', fontsize=20)
+
+ Euncorr.append(Euncorr_best)
+ resvals.append(best_res)
+ reserrs.append(res_err)
+ pvals.append(p)
+ svals.append(best_s)
+ wvals.append(wbest)
+
+pvals=np.array(pvals)
+svals=np.array(svals)
+Euncorr=np.array(Euncorr)
+plt.sca(axs[1])
+plt.plot(Euncorr,wvals, label="w")
+w_avg=np.mean(wvals)
+plt.axhline(w_avg, label=f'w avg: {w_avg:.2f}', ls=':')
+plt.plot(Euncorr,svals, label="s")
+m=(np.sum(svals*Euncorr)*len(Euncorr)-np.sum(Euncorr)*np.sum(svals))/(np.sum(Euncorr**2)*len(Euncorr)-np.sum(Euncorr)**2)
+b=np.mean(svals)-np.mean(Euncorr)*m
+plt.plot(Euncorr,Euncorr*m+b, label=f"s fit: ${m:.4f}E_{{uncorr}}+{b:.2f}$", ls=':')
+
+plt.xlabel("$E_{uncorr}=E_{Hcal}+E_{Ecal}/w$ [GeV]")
+plt.title("$E_{n,recon}=s\\times(E_{Hcal}+E_{Ecal}/w)$")
+plt.ylabel('parameter values')
plt.legend()
+plt.ylim(0)
+plt.savefig(outdir+"neutron_energy_params.pdf")
+
+#now make the energy plot
+print("making energy recon plot")
+fig, axs=plt.subplots(1,3, figsize=(24,8))
+partitions=[3.2,3.4, 3.6, 3.8, 4.0]
+for eta_min, eta_max in zip(partitions[:-1],partitions[1:]):
+ pvals=[]
+ resvals=[]
+ reserrs=[]
+ scalevals=[]
+ scaleerrs=[]
+ for p in 20, 30,40,50,60,70, 80:
+ best_res=1000
+ res_err=1000
+
+
+ wrange=np.linspace(30, 70, 30)*0.0257
+
+ w=w_avg
+ a=arrays_sim[p]
+ h=np.sum(a[f'HcalEndcapPInsertClusters.energy'], axis=-1)
+ e=np.sum(a[f'EcalEndcapPInsertClusters.energy'], axis=-1)
+ #phi=a['phi_truth']
+ uncorr=(e/w+h)
+ s=-0.0064*uncorr+1.80
+ r=uncorr*s #reconstructed energy with correction
+ r=r[(h>0)&(a['eta_truth']>eta_min)&(a['eta_truth']<eta_max)]#&(abs(phi)>np.pi/2)]
+ y,x=np.histogram(r,bins=50)
+ bcs=(x[1:]+x[:-1])/2
+ fnc=gauss
+ slc=abs(bcs-np.mean(r)*1.25)<2*np.std(r)
+ sigma=np.sqrt(y[slc])+0.5*(y[slc]==0)
+ p0=(100, np.mean(r), np.std(r))
+ try:
+ coeff, var_matrix = curve_fit(fnc, list(bcs[slc]), list(y[slc]), p0=p0,sigma=list(sigma))
+ res=np.abs(coeff[2]/coeff[1])
+
+ if res<best_res:
+ best_res=res
+ res_err=res*np.hypot(np.sqrt(var_matrix[2][2])/coeff[2], np.sqrt(var_matrix[1][1])/coeff[1])
+ wbest=w
+ scale=coeff[1]/np.sqrt(p**2+0.9406**2)
+ dscale=np.sqrt(var_matrix[1][1]/np.sqrt(p**2+0.9406**2))
+ except :
+ print("fit failed")
+ if p==50 and eta_min==3.4:
+ plt.sca(axs[0])
+ plt.errorbar(bcs, y, np.sqrt(y)+(y==0),marker='o', ls='', )
+ plt.title(f'p={p} GeV, ${eta_min}<\\eta<{eta_max}$')
+ plt.xlabel("$E_{recon}$ [GeV]")
+ plt.ylabel("events")
+ #plt.ylim(0)
+ xx=np.linspace(40, 70, 50)
+ plt.plot(xx, fnc(xx, *coeff))
+ resvals.append(best_res)
+ reserrs.append(res_err)
+ scalevals.append(scale)
+ scaleerrs.append(dscale)
+ pvals.append(p)
+ plt.sca(axs[1])
+ plt.errorbar(pvals, resvals, reserrs, marker='o', ls='', label=f"${eta_min}<\\eta<{eta_max}$")
+ #plt.ylim(0)
+ plt.ylabel("$\\sigma[E]/\\mu[E]$")
+ plt.xlabel("$p_{n}$ [GeV]")
+
+ plt.sca(axs[2])
+ plt.errorbar(pvals, scalevals, scaleerrs, marker='o', ls='', label=f"${eta_min}<\\eta<{eta_max}$")
+
+
+ plt.ylabel("$\\mu[E]/E$")
+
+
+axs[2].set_xlabel("$p_n$ [GeV]")
+axs[2].axhline(1, ls='--', color='0.5', alpha=0.7)
+axs[0].set_ylim(0)
+axs[1].set_ylim(0, 0.35)
+axs[2].set_ylim(0)
+axs[1].legend()
+axs[2].legend()
plt.tight_layout()
plt.savefig(outdir+"neutron_energy_recon.pdf")