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 try:
0105     coeff, var_matrix = curve_fit(fnc, bc[slc], y[slc], p0=p0,
0106                                      sigma=np.sqrt(y[slc])+(y[slc]==0), maxfev=10000)
0107     x=np.linspace(-1, 1)
0108     plt.plot(x, gauss(x, *coeff), color='tab:orange')
0109     plt.xlabel("$\\theta^{*\\rm recon}_{\\Lambda}-\\theta^{*\\rm truth}_{\\Lambda}$ [mrad]")
0110     plt.ylabel("events")
0111 except RuntimeError:
0112     print("fit failed")
0113
0114 plt.sca(axs[2])
0115 sigmas=[]
0116 dsigmas=[]
0117 xvals=[]
0118 for p in momenta:
0119     y,x=np.histogram(ak.flatten(theta_recon[p]-theta_truth[p])*1000, bins=100, range=(-1,1))
0120     bc=(x[1:]+x[:-1])/2
0121
0122     from scipy.optimize import curve_fit
0123     slc=abs(bc)<0.3
0124     fnc=gauss
0125     p0=(100, 0, 0.06)
0126     sigma=np.sqrt(y[slc])+(y[slc]==0)
0127     try:
0128         coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0, sigma=list(sigma), maxfev=10000)
0129         sigmas.append(coeff[2])
0130         dsigmas.append(np.sqrt(var_matrix[2][2]))
0131         xvals.append(p)
0132     except:
0133         print("fit failed")
0134 plt.ylim(0, 0.3)
0135
0136 plt.errorbar(xvals, sigmas, dsigmas, ls='', marker='o', color='k')
0137 x=np.linspace(100, 275, 100)
0138 plt.plot(x, 3/np.sqrt(x), color='tab:orange')
0139 plt.text(170, .23, "YR requirement:\n   3 mrad/$\\sqrt{E}$")
0140 plt.xlabel("$E_{\\Lambda}$ [GeV]")
0141 plt.ylabel("$\\sigma[\\theta^*_{\\Lambda}]$ [mrad]")
0142 plt.tight_layout()
0143 plt.savefig(outdir+"thetastar_recon.pdf")
0144 #plt.show()
0145
0146 #vtx z
0147 fig,axs=plt.subplots(1,3, figsize=(24, 8))
0148 plt.sca(axs[0])
0149 plt.title(f"$E_{{\\Lambda}}=100-275$ GeV")
0150 x=[]
0151 y=[]
0152 for p in momenta:
0153     x+=list(ak.flatten(arrays_sim[p]['MCParticles.vertex.z'][:,3]+0*zvtx_recon[p])/1000)
0154     y+=list(ak.flatten(zvtx_recon[p])/1000)
0155 plt.scatter(x,y)
0156 #print(x,y)
0157 plt.xlabel("$z^{\\rm truth}_{\\rm vtx}$ [m]")
0158 plt.ylabel("$z^{\\rm recon}_{\\rm vtx}$  [m]")
0159 plt.xlim(0,40)
0160 plt.ylim(0,40)
0161
0162 plt.sca(axs[1])
0163 plt.title(f"$E_{{\\Lambda}}=100-275$ GeV")
0164 y,x,_=plt.hist(y-np.array(x), bins=50, range=(-10,10))
0165 bc=(x[1:]+x[:-1])/2
0166
0167 from scipy.optimize import curve_fit
0168 slc=abs(bc)<5
0169 fnc=gauss
0170 p0=[100, 0, 1]
0171 try:
0172     coeff, var_matrix = curve_fit(fnc, bc[slc], y[slc], p0=p0,
0173                                      sigma=np.sqrt(y[slc])+(y[slc]==0), maxfev=10000)
0174     x=np.linspace(-5, 5)
0175     plt.plot(x, gauss(x, *coeff), color='tab:orange')
0176     print(coeff[2], np.sqrt(var_matrix[2][2]))
0177     plt.xlabel("$z^{*\\rm recon}_{\\rm vtx}-z^{*\\rm truth}_{\\rm vtx}$ [m]")
0178     plt.ylabel("events")
0179 except RuntimeError:
0180     print("fit failed")
0181
0182 plt.sca(axs[2])
0183 sigmas=[]
0184 dsigmas=[]
0185 xvals=[]
0186 for p in momenta:
0187
0188
0189     a=ak.flatten((zvtx_recon[p]-arrays_sim[p]['MCParticles.vertex.z'][:,3])/1000)
0190     y,x=np.histogram(a, bins=100, range=(-10,10))
0191     bc=(x[1:]+x[:-1])/2
0192
0193     from scipy.optimize import curve_fit
0194     slc=abs(bc)<5
0195     fnc=gauss
0196     p0=(100, 0, 1)
0197     #print(bc[slc],y[slc])
0198     sigma=np.sqrt(y[slc])+(y[slc]==0)
0199     try:
0200         coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0, sigma=list(sigma), maxfev=10000)
0201         sigmas.append(coeff[2])
0202         dsigmas.append(np.sqrt(var_matrix[2][2]))
0203         xvals.append(p)
0204     except:
0205         print("fit failed")
0206 plt.ylim(0, 2)
0207
0208 plt.errorbar(xvals, sigmas, dsigmas, ls='', marker='o', color='k')
0209 x=np.linspace(100, 275, 100)
0210
0211 avg=np.sum(sigmas/np.array(dsigmas)**2)/np.sum(1/np.array(dsigmas)**2)
0212 plt.axhline(avg, color='tab:orange')
0213 plt.text(150, 1.25,f"$\\sigma\\approx${avg:.1f} m")
0214
0215 plt.xlabel("$E_{\\Lambda}$ [GeV]")
0216 plt.ylabel("$\\sigma[z_{\\rm vtx}]$ [m]")
0217 plt.tight_layout()
0218 plt.savefig(outdir+"zvtx_recon.pdf")
0219 #plt.show()
0220
0221 p=100
0222 fig,axs=plt.subplots(1,2, figsize=(16, 8))
0223 plt.sca(axs[0])
0224 lambda_mass=1.115683
0225 vals=[]
0226 for p in momenta:
0227     vals+=list(ak.flatten(mass_recon[p]))
0228
0229 y,x,_= plt.hist(vals, bins=100, range=(1.0, 1.25))
0230 bc=(x[1:]+x[:-1])/2
0231 plt.axvline(lambda_mass, ls='--', color='tab:green', lw=3)
0232 plt.text(lambda_mass+.01, np.max(y)*1.05, "PDG mass", color='tab:green')
0233 plt.xlabel("$m_{\\Lambda}^{\\rm recon}$ [GeV]")
0234 plt.ylim(0, np.max(y)*1.2)
0235 plt.xlim(1.0, 1.25)
0236
0237 from scipy.optimize import curve_fit
0238 slc=abs(bc-lambda_mass)<0.05
0239 fnc=gauss
0240 p0=[100, lambda_mass, 0.04]
0241 try:
0242     coeff, var_matrix = curve_fit(fnc, bc[slc], y[slc], p0=p0,
0243                                      sigma=np.sqrt(y[slc])+(y[slc]==0), maxfev=10000)
0244     x=np.linspace(0.8, 1.3, 200)
0245     plt.plot(x, gauss(x, *coeff), color='tab:orange')
0246     print(coeff[2], np.sqrt(var_matrix[2][2]))
0247     plt.xlabel("$m^{\\rm recon}_{\\Lambda}$ [GeV]")
0248     plt.ylabel("events")
0249     plt.title(f"$E_{{\\Lambda}}=100-275$ GeV")
0250 except RuntimeError:
0251     print("fit failed")
0252
0253 plt.sca(axs[1])
0254 xvals=[]
0255 sigmas=[]
0256 dsigmas=[]
0257 for p in momenta:
0258     y,x= np.histogram(ak.flatten(mass_recon[p]), bins=100, range=(0.6,1.4))
0259     bc=(x[1:]+x[:-1])/2
0260
0261     from scipy.optimize import curve_fit
0262     slc=abs(bc-lambda_mass)<0.05
0263     fnc=gauss
0264     p0=[100, lambda_mass, 0.05]
0265     try:
0266         coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,
0267                                        sigma=list(np.sqrt(y[slc])+(y[slc]==0)), maxfev=10000)
0268         x=np.linspace(0.8, 1.3, 200)
0269         sigmas.append(coeff[2])
0270         dsigmas.append(np.sqrt(var_matrix[2][2]))
0271         xvals.append(p)
0272     except:
0273         print("fit failed")
0274
0275 plt.errorbar(xvals, sigmas, dsigmas, ls='', marker='o', color='k')
0276 avg=np.sum(sigmas/np.array(dsigmas)**2)/np.sum(1/np.array(dsigmas)**2)
0277 plt.axhline(avg, color='tab:orange')
0278 plt.text(150, 0.01,f"$\\sigma\\approx${avg*1000:.0f} MeV")
0279 plt.xlabel("$E_{\\Lambda}$ [GeV]")
0280 plt.ylabel("$\\sigma[m_{\\Lambda}]$ [GeV]")
0281 plt.ylim(0, 0.02)
0282 plt.tight_layout()
0283 plt.savefig(outdir+"lambda_mass_rec.pdf")
0284
0285
0286 #now for the CM stuff:
0287 phi_residuals=[]
0288 theta_residuals=[]
0289 for p in momenta:
0290     isNeutron=arrays_sim[p]['ReconstructedFarForwardZDCLambdaDecayProductsCM.PDG']==2112
0291     pxcm=arrays_sim[p]['ReconstructedFarForwardZDCLambdaDecayProductsCM.momentum.x']
0292     pycm=arrays_sim[p]['ReconstructedFarForwardZDCLambdaDecayProductsCM.momentum.y']
0293     pzcm=arrays_sim[p]['ReconstructedFarForwardZDCLambdaDecayProductsCM.momentum.z']
0294
0295
0296     import ROOT
0297     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\
0298                 +arrays_sim[p]['MCParticles.mass']**2)
0299     phicmtruth=list(np.repeat(-9999, len(arrays_sim[p])))
0300     thetacmtruth=list(np.repeat(-9999, len(arrays_sim[p])))
0301     for i in range(len(arrays_sim[p])):
0302         l=ROOT.TLorentzVector(px[i,2], py[i,2],  pz[i,2], E[i,2])
0303         n=ROOT.TLorentzVector(px[i,3], py[i,3],  pz[i,3], E[i,3])
0304         ncm=n.Clone();
0305         ncm.Boost(-l.BoostVector())
0306         phicmtruth[i]=ncm.Phi()
0307         thetacmtruth[i]=ncm.Theta()
0308
0309     arrays_sim[p]["phicmtruth"]=phicmtruth
0310     arrays_sim[p]["thetacmtruth"]=thetacmtruth
0311
0312     phicmtruth=arrays_sim[p]["phicmtruth"]
0313     thetacmtruth=arrays_sim[p]["thetacmtruth"]
0314     phi_residuals=np.concatenate((phi_residuals, ak.flatten((np.arctan2(pycm,pxcm)[isNeutron]-phicmtruth)*np.sin(thetacmtruth))))
0315     theta_residuals=np.concatenate((theta_residuals, ak.flatten(np.arctan2(np.hypot(pycm,pxcm),pzcm)[isNeutron]-thetacmtruth)))
0316 plt.figure()
0317 plt.hist(phi_residuals*1000, bins=100, range=(-300, 300))
0318 plt.xlabel("$(\\phi^n_{cm,rec}-\\phi^n_{cm,truth})\\times\\sin\\theta^n_{cm,truth} [mrad]$")
0319 plt.savefig(outdir+"neutron_phi_cm_res.pdf")
0320
0321 plt.figure()
0322 plt.hist(1000*theta_residuals, bins=100, range=(-1000, 1000))
0323 plt.xlabel("$\\theta^n_{cm,rec}-\\theta^n_{cm,truth}$ [mrad]")
0324 plt.savefig(outdir+"neutron_theta_cm_res.pdf")