diff --git a/benchmarks/dis/analysis/truth_reconstruction.py b/benchmarks/dis/analysis/truth_reconstruction.py index 2bd7ea23cbe83e3c2297e8bbe6a8047145709124..9faaf9dccff2ad8eb83ddcf9f51c348948175a7b 100644 --- a/benchmarks/dis/analysis/truth_reconstruction.py +++ b/benchmarks/dis/analysis/truth_reconstruction.py @@ -8,7 +8,6 @@ import mplhep import argparse - parser = argparse.ArgumentParser() parser.add_argument('--rec_file', type=str, help='Reconstructed track file.') parser.add_argument('--ebeam', type=float, help='Electron beam energy.') @@ -25,7 +24,6 @@ k = int(args.ebeam) p = int(args.pbeam) Nevents = int(args.nevents) - for array in ur.iterate(rec_file + ':events',['MCParticles/MCParticles.generatorStatus', 'MCParticles/MCParticles.mass', 'MCParticles/MCParticles.PDG', @@ -61,9 +59,6 @@ phi_mc = np.arctan2(py_mc, px_mc) theta_rc = np.arctan2(np.sqrt(px_rc**2+py_rc**2), pz_rc) phi_rc = np.arctan2(py_rc, px_rc) - - - booll = (PDG_mc[simID])==(PDG_rc[recID]) MC_list = [ak.Array(momentum_mc[simID][booll]), ak.Array(theta_mc[simID][booll]), @@ -74,7 +69,16 @@ RC_list = [ak.Array(momentum_rc[recID][booll]), ak.Array(phi_rc[recID][booll]), -np.log(np.tan((ak.Array(theta_rc[recID][booll]))/2))] title_list = ['Momentum','Theta','Phi','Eta'] +boolean_pion = np.logical_or(ak.Array(PDG_mc[simID][booll])==-211, ak.Array(PDG_mc[simID][booll])==+211) +boolean_proton = np.logical_or(ak.Array(PDG_mc[simID][booll])==-2212, ak.Array(PDG_mc[simID][booll])==+2212) +boolean_electron = ak.Array(PDG_mc[simID][booll])==11 +boolean_neutron = ak.Array(PDG_mc[simID][booll])==2112 +boolean_photon = ak.Array(PDG_mc[simID][booll])==22 +if Nevents == 100: + ssize = 1 +else: + ssize = 0.01 ################### @@ -86,13 +90,11 @@ for i in range(len(MC_list)): Y_len = ak.count(Y,axis=None) if X_len > Y_len: - Y_boolean = np.ones_like(Y) == 1 - X_F = X[Y_boolean] - Y_F = Y[Y_boolean] + F_boolean = np.ones_like(Y) == 1 else: - X_boolean = np.ones_like(X) == 1 - X_F = X[X_boolean] - Y_F = Y[X_boolean] + F_boolean = np.ones_like(X) == 1 + X_F = X[F_boolean] + Y_F = Y[F_boolean] X_s = np.array(ak.flatten(X_F)) Y_s = np.array(ak.flatten(Y_F)) @@ -135,11 +137,7 @@ for i in range(len(MC_list)): ################# - boolean_pion = np.logical_or(ak.Array(PDG_mc[simID][booll])==-211, ak.Array(PDG_mc[simID][booll])==+211) -boolean_proton = np.logical_or(ak.Array(PDG_mc[simID][booll])==-2212, ak.Array(PDG_mc[simID][booll])==+2212) -boolean_electron = ak.Array(PDG_mc[simID][booll])==11 -boolean_neutron = ak.Array(PDG_mc[simID][booll])==2112 -boolean_photon = ak.Array(PDG_mc[simID][booll])==22 + for i in range(len(MC_list)): X1 = MC_list[i] @@ -167,13 +165,11 @@ for i in range(len(MC_list)): Y_len = ak.count(Y,axis=None) if X_len > Y_len: - Y_boolean = np.ones_like(Y) == 1 - X_F = X[Y_boolean] - Y_F = Y[Y_boolean] + F_boolean = np.ones_like(Y) == 1 else: - X_boolean = np.ones_like(X) == 1 - X_F = X[X_boolean] - Y_F = Y[X_boolean] + F_boolean = np.ones_like(X) == 1 + X_F = X[F_boolean] + Y_F = Y[F_boolean] X_s = np.array(ak.flatten(X_F)) Y_s = np.array(ak.flatten(Y_F)) @@ -190,24 +186,21 @@ for i in range(len(MC_list)): (ax1, ax2), (ax3, ax4),(ax5, ax6) = gs.subplots(sharex=True, sharey=True) # fig.suptitle('') if i == 1: - ax1.scatter(-X_plot[0], Y_plot[0], s = 0.01) - ax2.scatter(-X_plot[1], Y_plot[1], s = 0.01) - ax3.scatter(-X_plot[2], Y_plot[2], s = 0.01) - ax4.scatter(-X_plot[3], Y_plot[3], s = 0.01) - ax5.scatter(-X_plot[4], Y_plot[4], s = 0.01) - ax6.scatter(-X_plot[5], Y_plot[5], s = 0.01) - else: - ax1.scatter(X_plot[0], Y_plot[0], s = 0.01) - ax2.scatter(X_plot[1], Y_plot[1], s = 0.01) - ax3.scatter(X_plot[2], Y_plot[2], s = 0.01) - ax4.scatter(X_plot[3], Y_plot[3], s = 0.01) - ax5.scatter(X_plot[4], Y_plot[4], s = 0.01) - ax6.scatter(X_plot[5], Y_plot[5], s = 0.01) + X_plot[0],X_plot[1],X_plot[2],X_plot[3],X_plot[4],X_plot[5] = -X_plot[0],-X_plot[1],-X_plot[2],-X_plot[3],-X_plot[4],-X_plot[5] + + ax1.scatter(X_plot[0], Y_plot[0], s = ssize) + ax2.scatter(X_plot[1], Y_plot[1], s = ssize) + ax3.scatter(X_plot[2], Y_plot[2], s = ssize) + ax4.scatter(X_plot[3], Y_plot[3], s = ssize) + ax5.scatter(X_plot[4], Y_plot[4], s = ssize) + ax6.scatter(X_plot[5], Y_plot[5], s = ssize) + ax_list = [ax1,ax2,ax3,ax4,ax5] if i == 0: ax1.set_ylabel('rc/mc') ax3.set_ylabel('rc/mc') ax5.set_ylabel('rc/mc') + tratio ='ratio' for ax in ax_list: ax.set_yscale('log') ax.set_xscale('log') @@ -215,7 +208,7 @@ for i in range(len(MC_list)): ax1.set_ylabel('rc-mc') ax3.set_ylabel('rc-mc') ax5.set_ylabel('rc-mc') - + tratio ='difference' ax2.set_title('Pions') ax3.set_title('Protons') @@ -226,11 +219,9 @@ for i in range(len(MC_list)): ax6.set_xlabel('%s'%(title_list[i])) fig.set_figwidth(20) fig.set_figheight(10) - if i == 0: - ax1.set_title('%s Ratio (%g x %g)GeV %gGeV minQ2 %s events'%(title_list[i],k,p,minq2,Nevents)) - else: - ax1.set_title('%s Difference (%g x %g)GeV %gGeV minQ2 %s events'%(title_list[i],k,p,minq2,Nevents)) - plt.savefig(os.path.join(args.outdir, '%gon%g/minQ2=%g/truth_reconstruction/%s_ratio_%gx%g_minQ2=%g.png' % (k,p,minq2,title_list[i],k,p,minq2))) + + ax1.set_title('%s %s (%g x %g)GeV %gGeV minQ2 %s events'%(title_list[i],tratio,k,p,minq2,Nevents)) + plt.savefig(os.path.join(args.outdir, '%gon%g/minQ2=%g/truth_reconstruction/%s_%s_%gx%g_minQ2=%g.png' % (k,p,minq2,title_list[i],tratio,k,p,minq2))) ############### @@ -282,12 +273,13 @@ for i in range(1,len(MC_list),1): (ax1, ax2), (ax3, ax4),(ax5, ax6) = gs.subplots(sharex=True, sharey=True) # fig.suptitle('') - ax1.scatter(X_plot[0], Y_plot[0], s = 0.01) - ax2.scatter(X_plot[1], Y_plot[1], s = 0.01) - ax3.scatter(X_plot[2], Y_plot[2], s = 0.01) - ax4.scatter(X_plot[3], Y_plot[3], s = 0.01) - ax5.scatter(X_plot[4], Y_plot[4], s = 0.01) - ax6.scatter(X_plot[5], Y_plot[5], s = 0.01) + ax1.scatter(X_plot[0], Y_plot[0], s = ssize) + ax2.scatter(X_plot[1], Y_plot[1], s = ssize) + ax3.scatter(X_plot[2], Y_plot[2], s = ssize) + ax4.scatter(X_plot[3], Y_plot[3], s = ssize) + ax5.scatter(X_plot[4], Y_plot[4], s = ssize) + ax6.scatter(X_plot[5], Y_plot[5], s = ssize) + ax_list = [ax1,ax2,ax3,ax4,ax5] for ax in ax_list: ax.set_xscale('log') @@ -332,14 +324,14 @@ phi_rc_F = np.array(ak.flatten(phi_rc_fil[F_boolean])) ratio = np.array((ak.Array(theta_rc_F)-(ak.Array(theta_mc_F)))) fig = plt.figure() -gs = fig.add_gridspec(2, 2, wspace=0.03) +gs = fig.add_gridspec(2, 2, wspace=0.01) (ax1, ax2), (ax3, ax4) = gs.subplots(sharex=True, sharey=True) fig.suptitle('Photons in (%g x %g)GeV %gGeV minQ2 %s events'%(k,p,minq2,Nevents)) -ax1.scatter(-theta_mc_F, ratio, s = 0.01) -ax2.scatter(-theta_rc_F, ratio, s = 0.01) -ax3.scatter(-theta_mc_F, phi_mc_F, s = 0.01) -ax4.scatter(-theta_rc_F, phi_rc_F, s = 0.01) +ax1.scatter(-theta_mc_F, ratio, s = ssize) +ax2.scatter(-theta_rc_F, ratio, s = ssize) +ax3.scatter(-theta_mc_F, phi_mc_F, s = ssize) +ax4.scatter(-theta_rc_F, phi_rc_F, s = ssize) ax1.set_ylabel('rc-mc') ax2.set_ylabel('rc-mc') @@ -348,7 +340,7 @@ ax4.set_ylabel('Phi rc') ax3.set_xlabel('Theta mc') ax4.set_xlabel('Theta rc') fig.set_figwidth(20) -fig.set_figheight(20) +fig.set_figheight(10) plt.savefig(os.path.join(args.outdir, '%gon%g/minQ2=%g/truth_reconstruction/photons_%gx%g_minQ2=%g.png' % (k,p,minq2,k,p,minq2)))