zdc_lambda/analysis/lambda_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 plt.rcParams["figure.figsize"] = (7, 7)
0010
0011 outdir=sys.argv[1]+"/"
0012 config=outdir.split("/")[1]
0013 try:
0014     import os
0015     os.mkdir(outdir[:-1])
0016 except:
0017     pass
0018 import uproot as ur
0019 arrays_sim={}
0020 momenta=100, 125, 150, 175,200,225,250,275
0021 for p in momenta:
0022     arrays_sim[p] = ur.concatenate({
0023         f'sim_output/zdc_lambda/{config}_rec_lambda_dec_{p}GeV_{index}.edm4eic.root': 'events'
0024         for index in range(5)
0025     })
0026
0027 def gauss(x, A,mu, sigma):
0028     return A * np.exp(-(x-mu)**2/(2*sigma**2))
0029
0030 #keep track of the number of clusters per event
0031 nclusters={}
0032
0033 for p in momenta:
0034     plt.figure()
0035     nclusters[p]=[]
0036     for i in range(len(arrays_sim[p])):
0037         nclusters[p].append(len(arrays_sim[p]["HcalFarForwardZDCClusters.position.x"][i]))
0038     nclusters[p]=np.array(nclusters[p])
0039     plt.hist(nclusters[p],bins=20, range=(0,20))
0040     plt.xlabel("number of clusters")
0041     plt.yscale('log')
0042     plt.title(f"$p_\\Lambda={p}$ GeV")
0043     plt.ylim(1)
0044     plt.savefig(outdir+f"nclust_{p}GeV_recon.pdf")
0045     print("saved file ", outdir+f"nclust_{p}GeV_recon.pdf")
0046
0047
0048
0049 pt_truth={}
0050 theta_truth={}
0051
0052 for p in momenta:
0053     #get the truth value of theta* and pt*
0054     px=arrays_sim[p]["MCParticles.momentum.x"][:,2]
0055     py=arrays_sim[p]["MCParticles.momentum.y"][:,2]
0056     pz=arrays_sim[p]["MCParticles.momentum.z"][:,2]
0057     tilt=-0.025
0058     pt_truth[p]=np.hypot(px*np.cos(tilt)-pz*np.sin(tilt), py)
0059     theta_truth[p]=np.arctan2(pt_truth[p],pz*np.cos(tilt)+px*np.sin(tilt))
0060
0061 if "ReconstructedFarForwardZDCLambdas.momentum.x" not in arrays_sim[momenta[0]].fields:
0062     print("ReconstructedFarForwardZDCLambdas collection is not available (needs EICrecon 1.23)")
0063     import sys
0064     sys.exit(0)
0065
0066 theta_recon={}
0067 E_recon={}
0068 zvtx_recon={}
0069 mass_recon={}
0070 print(arrays_sim[p].fields)
0071 for p in momenta:
0072
0073     px,py,pz,m=(arrays_sim[p][f"ReconstructedFarForwardZDCLambdas.{a}"] for a in "momentum.x momentum.y momentum.z mass".split())
0074     theta_recon[p]=np.arctan2(np.hypot(px*np.cos(tilt)-pz*np.sin(tilt), py),pz*np.cos(tilt)+px*np.sin(tilt))
0075     E_recon[p]=np.sqrt(px**2+py**2+pz**2+m**2)
0076     zvtx_recon[p]=arrays_sim[p][f"ReconstructedFarForwardZDCLambdas.referencePoint.z"]*np.cos(tilt)+arrays_sim[p][f"ReconstructedFarForwardZDCLambdas.referencePoint.x"]*np.sin(tilt)
0077     mass_recon[p]=m
0078
0079 #theta plots
0080 fig,axs=plt.subplots(1,3, figsize=(24, 8))
0081 plt.sca(axs[0])
0082 plt.title(f"$E_{{\\Lambda}}=100-275$ GeV")
0083 x=[]
0084 y=[]
0085 import awkward as ak
0086 for p in momenta:
0087     x+=list(ak.flatten(theta_truth[p]+0*theta_recon[p])*1000)
0088     y+=list(ak.flatten(theta_recon[p]*1000))
0089 plt.scatter(x,y)
0090 plt.xlabel("$\\theta^{*\\rm truth}_{\\Lambda}$ [mrad]")
0091 plt.ylabel("$\\theta^{*\\rm recon}_{\\Lambda}$ [mrad]")
0092 plt.xlim(0,3.2)
0093 plt.ylim(0,3.2)
0094
0095 plt.sca(axs[1])
0096 plt.title(f"$E_{{\\Lambda}}=100-275$ GeV")
0097 y,x,_=plt.hist(y-np.array(x), bins=50, range=(-1,1))
0098 bc=(x[1:]+x[:-1])/2
0099
0100 from scipy.optimize import curve_fit
0101 slc=abs(bc)<0.3
0102 fnc=gauss
0103 p0=[100, 0, 0.05]
0104 coeff, var_matrix = curve_fit(fnc, bc[slc], y[slc], p0=p0,
0105                                  sigma=np.sqrt(y[slc])+(y[slc]==0), maxfev=10000)
0106 x=np.linspace(-1, 1)
0107 plt.plot(x, gauss(x, *coeff), color='tab:orange')
0108 plt.xlabel("$\\theta^{*\\rm recon}_{\\Lambda}-\\theta^{*\\rm truth}_{\\Lambda}$ [mrad]")
0109 plt.ylabel("events")
0110
0111 plt.sca(axs[2])
0112 sigmas=[]
0113 dsigmas=[]
0114 xvals=[]
0115 for p in momenta:
0116     y,x=np.histogram(ak.flatten(theta_recon[p]-theta_truth[p])*1000, bins=100, range=(-1,1))
0117     bc=(x[1:]+x[:-1])/2
0118
0119     from scipy.optimize import curve_fit
0120     slc=abs(bc)<0.3
0121     fnc=gauss
0122     p0=(100, 0, 0.06)
0123     sigma=np.sqrt(y[slc])+(y[slc]==0)
0124     try:
0125         coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0, sigma=list(sigma), maxfev=10000)
0126         sigmas.append(coeff[2])
0127         dsigmas.append(np.sqrt(var_matrix[2][2]))
0128         xvals.append(p)
0129     except:
0130         print("fit failed")
0131 plt.ylim(0, 0.3)
0132
0133 plt.errorbar(xvals, sigmas, dsigmas, ls='', marker='o', color='k')
0134 x=np.linspace(100, 275, 100)
0135 plt.plot(x, 3/np.sqrt(x), color='tab:orange')
0136 plt.text(170, .23, "YR requirement:\n   3 mrad/$\\sqrt{E}$")
0137 plt.xlabel("$E_{\\Lambda}$ [GeV]")
0138 plt.ylabel("$\\sigma[\\theta^*_{\\Lambda}]$ [mrad]")
0139 plt.tight_layout()
0140 plt.savefig(outdir+"thetastar_recon.pdf")
0141 #plt.show()
0142
0143 #vtx z
0144 fig,axs=plt.subplots(1,3, figsize=(24, 8))
0145 plt.sca(axs[0])
0146 plt.title(f"$E_{{\\Lambda}}=100-275$ GeV")
0147 x=[]
0148 y=[]
0149 for p in momenta:
0150     x+=list(ak.flatten(arrays_sim[p]['MCParticles.vertex.z'][:,3]+0*zvtx_recon[p])/1000)
0151     y+=list(ak.flatten(zvtx_recon[p])/1000)
0152 plt.scatter(x,y)
0153 #print(x,y)
0154 plt.xlabel("$z^{\\rm truth}_{\\rm vtx}$ [m]")
0155 plt.ylabel("$z^{\\rm recon}_{\\rm vtx}$  [m]")
0156 plt.xlim(0,40)
0157 plt.ylim(0,40)
0158
0159 plt.sca(axs[1])
0160 plt.title(f"$E_{{\\Lambda}}=100-275$ GeV")
0161 y,x,_=plt.hist(y-np.array(x), bins=50, range=(-10,10))
0162 bc=(x[1:]+x[:-1])/2
0163
0164 from scipy.optimize import curve_fit
0165 slc=abs(bc)<5
0166 fnc=gauss
0167 p0=[100, 0, 1]
0168 coeff, var_matrix = curve_fit(fnc, bc[slc], y[slc], p0=p0,
0169                                  sigma=np.sqrt(y[slc])+(y[slc]==0), maxfev=10000)
0170 x=np.linspace(-5, 5)
0171 plt.plot(x, gauss(x, *coeff), color='tab:orange')
0172 print(coeff[2], np.sqrt(var_matrix[2][2]))
0173 plt.xlabel("$z^{*\\rm recon}_{\\rm vtx}-z^{*\\rm truth}_{\\rm vtx}$ [m]")
0174 plt.ylabel("events")
0175
0176 plt.sca(axs[2])
0177 sigmas=[]
0178 dsigmas=[]
0179 xvals=[]
0180 for p in momenta:
0181
0182
0183     a=ak.flatten((zvtx_recon[p]-arrays_sim[p]['MCParticles.vertex.z'][:,3])/1000)
0184     y,x=np.histogram(a, bins=100, range=(-10,10))
0185     bc=(x[1:]+x[:-1])/2
0186
0187     from scipy.optimize import curve_fit
0188     slc=abs(bc)<5
0189     fnc=gauss
0190     p0=(100, 0, 1)
0191     #print(bc[slc],y[slc])
0192     sigma=np.sqrt(y[slc])+(y[slc]==0)
0193     try:
0194         coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0, sigma=list(sigma), maxfev=10000)
0195         sigmas.append(coeff[2])
0196         dsigmas.append(np.sqrt(var_matrix[2][2]))
0197         xvals.append(p)
0198     except:
0199         print("fit failed")
0200 plt.ylim(0, 2)
0201
0202 plt.errorbar(xvals, sigmas, dsigmas, ls='', marker='o', color='k')
0203 x=np.linspace(100, 275, 100)
0204
0205 avg=np.sum(sigmas/np.array(dsigmas)**2)/np.sum(1/np.array(dsigmas)**2)
0206 plt.axhline(avg, color='tab:orange')
0207 plt.text(150, 1.25,f"$\\sigma\\approx${avg:.1f} m")
0208
0209 plt.xlabel("$E_{\\Lambda}$ [GeV]")
0210 plt.ylabel("$\\sigma[z_{\\rm vtx}]$ [m]")
0211 plt.tight_layout()
0212 plt.savefig(outdir+"zvtx_recon.pdf")
0213 #plt.show()
0214
0215 p=100
0216 fig,axs=plt.subplots(1,2, figsize=(16, 8))
0217 plt.sca(axs[0])
0218 lambda_mass=1.115683
0219 vals=[]
0220 for p in momenta:
0221     vals+=list(ak.flatten(mass_recon[p]))
0222
0223 y,x,_= plt.hist(vals, bins=100, range=(1.0, 1.25))
0224 bc=(x[1:]+x[:-1])/2
0225 plt.axvline(lambda_mass, ls='--', color='tab:green', lw=3)
0226 plt.text(lambda_mass+.01, np.max(y)*1.05, "PDG mass", color='tab:green')
0227 plt.xlabel("$m_{\\Lambda}^{\\rm recon}$ [GeV]")
0228 plt.ylim(0, np.max(y)*1.2)
0229 plt.xlim(1.0, 1.25)
0230
0231 from scipy.optimize import curve_fit
0232 slc=abs(bc-lambda_mass)<0.05
0233 fnc=gauss
0234 p0=[100, lambda_mass, 0.04]
0235 coeff, var_matrix = curve_fit(fnc, bc[slc], y[slc], p0=p0,
0236                                  sigma=np.sqrt(y[slc])+(y[slc]==0), maxfev=10000)
0237 x=np.linspace(0.8, 1.3, 200)
0238 plt.plot(x, gauss(x, *coeff), color='tab:orange')
0239 print(coeff[2], np.sqrt(var_matrix[2][2]))
0240 plt.xlabel("$m^{\\rm recon}_{\\Lambda}$ [GeV]")
0241 plt.ylabel("events")
0242 plt.title(f"$E_{{\\Lambda}}=100-275$ GeV")
0243
0244 plt.sca(axs[1])
0245 xvals=[]
0246 sigmas=[]
0247 dsigmas=[]
0248 for p in momenta:
0249     y,x= np.histogram(ak.flatten(mass_recon[p]), bins=100, range=(0.6,1.4))
0250     bc=(x[1:]+x[:-1])/2
0251
0252     from scipy.optimize import curve_fit
0253     slc=abs(bc-lambda_mass)<0.05
0254     fnc=gauss
0255     p0=[100, lambda_mass, 0.05]
0256     try:
0257         coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,
0258                                        sigma=list(np.sqrt(y[slc])+(y[slc]==0)), maxfev=10000)
0259         x=np.linspace(0.8, 1.3, 200)
0260         sigmas.append(coeff[2])
0261         dsigmas.append(np.sqrt(var_matrix[2][2]))
0262         xvals.append(p)
0263     except:
0264         print("fit failed")
0265
0266 plt.errorbar(xvals, sigmas, dsigmas, ls='', marker='o', color='k')
0267 avg=np.sum(sigmas/np.array(dsigmas)**2)/np.sum(1/np.array(dsigmas)**2)
0268 plt.axhline(avg, color='tab:orange')
0269 plt.text(150, 0.01,f"$\\sigma\\approx${avg*1000:.0f} MeV")
0270 plt.xlabel("$E_{\\Lambda}$ [GeV]")
0271 plt.ylabel("$\\sigma[m_{\\Lambda}]$ [GeV]")
0272 plt.ylim(0, 0.02)
0273 plt.tight_layout()
0274 plt.savefig(outdir+"lambda_mass_rec.pdf")
0275
0276
0277 #now for the CM stuff:
0278 phi_residuals=[]
0279 theta_residuals=[]
0280 for p in momenta:
0281     isNeutron=arrays_sim[p]['ReconstructedFarForwardZDCLambdaDecayProductsCM.PDG']==2112
0282     pxcm=arrays_sim[p]['ReconstructedFarForwardZDCLambdaDecayProductsCM.momentum.x']
0283     pycm=arrays_sim[p]['ReconstructedFarForwardZDCLambdaDecayProductsCM.momentum.y']
0284     pzcm=arrays_sim[p]['ReconstructedFarForwardZDCLambdaDecayProductsCM.momentum.z']
0285
0286
0287     import ROOT
0288     px,py,pz,E=arrays_sim[p]['MCParticles.momentum.x'], arrays_sim[p]['MCParticles.momentum.y'], arrays_sim[p]['MCParticles.momentum.z'],np.sqrt(arrays_sim[p]['MCParticles.momentum.x']**2+arrays_sim[p]['MCParticles.momentum.y']**2+arrays_sim[p]['MCParticles.momentum.z']**2\
0289                 +arrays_sim[p]['MCParticles.mass']**2)
0290     phicmtruth=list(np.repeat(-9999, len(arrays_sim[p])))
0291     thetacmtruth=list(np.repeat(-9999, len(arrays_sim[p])))
0292     for i in range(len(arrays_sim[p])):
0293         l=ROOT.TLorentzVector(px[i,2], py[i,2],  pz[i,2], E[i,2])
0294         n=ROOT.TLorentzVector(px[i,3], py[i,3],  pz[i,3], E[i,3])
0295         ncm=n.Clone();
0296         ncm.Boost(-l.BoostVector())
0297         phicmtruth[i]=ncm.Phi()
0298         thetacmtruth[i]=ncm.Theta()
0299
0300     arrays_sim[p]["phicmtruth"]=phicmtruth
0301     arrays_sim[p]["thetacmtruth"]=thetacmtruth
0302
0303     phicmtruth=arrays_sim[p]["phicmtruth"]
0304     thetacmtruth=arrays_sim[p]["thetacmtruth"]
0305     phi_residuals=np.concatenate((phi_residuals, ak.flatten((np.arctan2(pycm,pxcm)[isNeutron]-phicmtruth)*np.sin(thetacmtruth))))
0306     theta_residuals=np.concatenate((theta_residuals, ak.flatten(np.arctan2(np.hypot(pycm,pxcm),pzcm)[isNeutron]-thetacmtruth)))
0307 plt.figure()
0308 plt.hist(phi_residuals*1000, bins=100, range=(-300, 300))
0309 plt.xlabel("$(\\phi^n_{cm,rec}-\\phi^n_{cm,truth})\\times\\sin\\theta^n_{cm,truth} [mrad]$")
0310 plt.savefig(outdir+"neutron_phi_cm_res.pdf")
0311
0312 plt.figure()
0313 plt.hist(1000*theta_residuals, bins=100, range=(-1000, 1000))
0314 plt.xlabel("$\\theta^n_{cm,rec}-\\theta^n_{cm,truth}$ [mrad]")
0315 plt.savefig(outdir+"neutron_theta_cm_res.pdf")