insert_tau/analysis/tau_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 config=sys.argv[1].split("/")[1]  #results/{config}/{benchmark_name}

0012 outdir=sys.argv[1]+"/"
0013 try:
0014     import os
0015     os.mkdir(outdir[:-1])
0016 except:
0017     pass
0018
0019 import uproot as ur
0020 arrays_sim={}
0021 momenta=20, 30, 40, 50, 60,80,100
0022 for p in momenta:
0023     arrays_sim[p] = ur.concatenate({
0024         f'sim_output/insert_tau/{config}_rec_tau-_{p}GeV_{index}.edm4eic.root': 'events'
0025         for index in range(5)
0026     })
0027
0028
0029 for a in arrays_sim.values():
0030     #recon

0031     Etot=0
0032     px=0
0033     py=0
0034     pz=0
0035     for det in "HcalEndcapPInsert", "EcalEndcapPInsert", "EcalEndcapP", "LFHCAL":
0036
0037         E=a[f'{det}Clusters.energy']
0038
0039         #todo apply corrections depending on whether this is an electromagnetic or hadronic shower.

0040
0041         x=a[f'{det}Clusters.position.x']
0042         y=a[f'{det}Clusters.position.y']
0043         z=a[f'{det}Clusters.position.z']
0044         r=np.sqrt(x**2+y**2+z**2)
0045         Etot=Etot+np.sum(E, axis=-1)
0046         px=px+np.sum(E*x/r,axis=-1)
0047         py=py+np.sum(E*y/r,axis=-1)
0048         pz=pz+np.sum(E*z/r,axis=-1)
0049
0050     a['jet_p_recon']=np.sqrt(px**2+py**2+pz**2)
0051     a['jet_E_recon']=Etot
0052
0053     a['jet_theta_recon']=np.arctan2(np.hypot(px*np.cos(-.025)-pz*np.sin(-.025),py),
0054                                     pz*np.cos(-.025)+px*np.sin(-.025))
0055
0056     #truth

0057     m=a['MCParticles.mass'][::,2:]
0058     px=a['MCParticles.momentum.x'][::,2:]
0059     py=a['MCParticles.momentum.y'][::,2:]
0060     pz=a['MCParticles.momentum.z'][::,2:]
0061     E=np.sqrt(m**2+px**2+py**2+pz**2)
0062     status=a['MCParticles.simulatorStatus'][::,2:]
0063     PDG=a['MCParticles.PDG'][::,2:]
0064
0065     #find the hadronic part: initial-state tau - all leptons

0066     selection=1*(PDG==15)-1*(np.abs(PDG)==16)
0067     is_hadronic=1*(np.sum((PDG==-14)+(PDG==-12), axis=-1)==0)
0068
0069     px_hfs, py_hfs, pz_hfs= np.sum(px*selection,axis=-1)*is_hadronic,np.sum(py*selection,axis=-1)*is_hadronic, np.sum(pz*selection,axis=-1)*is_hadronic
0070
0071     a['hfs_p_truth']=np.sqrt(px_hfs**2+py_hfs**2+pz_hfs**2)
0072     a['hfs_E_truth']=np.sum(E*selection,axis=-1)*is_hadronic
0073
0074
0075     a['hfs_theta_truth']=np.arctan2(np.hypot(px_hfs*np.cos(-.025)-pz_hfs*np.sin(-.025),py_hfs),
0076                                     pz_hfs*np.cos(-.025)+px_hfs*np.sin(-.025))
0077     a['hfs_eta_truth']=-np.log(np.tan(a['hfs_theta_truth']/2))
0078     a['n_mu']=np.sum(np.abs(PDG)==13, axis=-1)
0079     a['n_e']=np.sum(np.abs(PDG)==13, axis=-1)
0080     a['hfs_mass_truth']=np.sqrt(a['hfs_E_truth']**2-a['hfs_p_truth']**2)
0081
0082 for a in arrays_sim.values():
0083     selection=(a['hfs_eta_truth']>3.1) & (a['hfs_eta_truth']<3.8)\
0084             &(a['n_mu']==0)&(a['n_e']==0)&(a['hfs_mass_truth']>.140)&(a['jet_E_recon']>0)
0085
0086 Etruth=[]
0087 Erecon=[]
0088
0089 theta_truth=[]
0090 theta_recon=[]
0091
0092 eta_max=3.7
0093 eta_min=3.3
0094 for a in arrays_sim.values():
0095     selection=(a['hfs_eta_truth']>eta_min) & (a['hfs_eta_truth']<eta_max)\
0096             &(a['n_mu']==0)&(a['n_e']==0)&(a['hfs_mass_truth']>.140)&(a['jet_E_recon']>1)
0097     theta_truth=np.concatenate((theta_truth,a['hfs_theta_truth'][selection]))
0098     theta_recon=np.concatenate((theta_recon,a['jet_theta_recon'][selection]))
0099     Etruth=np.concatenate((Etruth,a['hfs_E_truth'][selection]))
0100     Erecon=np.concatenate((Erecon,a['jet_E_recon'][selection]))
0101
0102 plt.figure()
0103 plt.scatter(theta_truth, theta_recon, 1)
0104 plt.xlabel("$\\theta^{hfs}_{\\rm truth}$ [rad]")
0105 plt.ylabel("$\\theta^{hfs}_{\\rm rec}$ [rad]")
0106 plt.title(f"$E_{{\\tau}}$=20-100 GeV, ${eta_min}<\\eta_{{hfs}}<{eta_max}$")
0107 plt.plot([0.04,0.1], [0.04, 0.1], color='tab:orange')
0108 plt.ylim(0, 0.15)
0109 plt.savefig(outdir+"/theta_scatter.pdf")
0110
0111 plt.figure()
0112 plt.scatter(Etruth, Erecon, 1)
0113 plt.xlabel("$E^{hfs}_{\\rm truth}$ [GeV]")
0114 plt.ylabel("$E^{hfs}_{\\rm rec}$ [GeV]")
0115 plt.title(f"$E_{{\\tau}}$=20-100 GeV, ${eta_min}<\\eta_{{hfs}}<{eta_max}$")
0116 plt.plot((0,100), (0, 100), color='tab:orange')
0117 plt.savefig(outdir+"/energy_scatter.pdf")
0118
0119 def gauss(x, A,mu, sigma):
0120     return A * np.exp(-(x-mu)**2/(2*sigma**2))
0121 from scipy.optimize import curve_fit
0122
0123 res=[]
0124 dres=[]
0125 emid=[]
0126 ew=[]
0127 partitions=(20,30, 40, 60,80,100)
0128 for emin, emax in zip(partitions[:-1], partitions[1:]):
0129
0130     y,x = np.histogram((theta_recon-theta_truth)[(Etruth>emin)&(Etruth<emax)], bins=100, range=(-0.03,0.03))
0131     bc=(x[1:]+x[:-1])/2
0132     slc=abs(bc)<0.5
0133     # try:

0134     p0=(100, 0, 0.15)
0135     fnc=gauss
0136     sigma=np.sqrt(y[slc])+(y[slc]==0)
0137
0138     coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0, sigma=list(sigma), maxfev=10000)
0139     res.append(abs(coeff[2]))
0140     dres.append(np.sqrt(var_matrix[2][2]))
0141     emid.append((emin+emax)/2)
0142     ew.append((emax-emin)/2)
0143 plt.errorbar(emid, 1000*np.array(res),1000*np.array(dres), ew, ls='', label=f'{eta_min}<$\\eta_{{hfs}}$<{eta_max}')
0144 plt.xlabel('$E_{hfs}$ [GeV]')
0145 plt.ylabel('$\\theta$ resolution [mrad]')
0146 plt.ylim(0)
0147
0148 fnc=lambda E,B:B/E
0149 p0=[1,]
0150 coeff, var_matrix = curve_fit(fnc, emid, res, p0=p0, sigma=list(dres), maxfev=10000)
0151 xx=np.linspace(10, 100, 100)
0152 plt.plot(xx, 1000*fnc(xx, *coeff), label=f"fit: ${coeff[0]:.2f}/E$ mrad")
0153 plt.legend()
0154 plt.savefig(outdir+"/theta_res.pdf")