taking account of single matrix jobs in benchmarks/dis/analysis/truth_reconstruction.py

benchmarks/dis/analysis/truth_reconstruction.py

@@ -75,6 +75,9 @@ if Nevents == 100:
     ssize = 1
 else:
     ssize = 0.01
+
+particle = config.split('-')[0].strip()
+particle_dict = {'e':[boolean_electron,'Electrons'],'pi':[boolean_pion,'Pions']}
 ###################
 
 
@@ -214,8 +217,8 @@ for i in range(len(MC_list)):
     ax5.set_title('Neutrons')
     ax6.set_title('Photons')
     if i == 1:
-        ax5.set_xlabel('- %s'%(title_list[i]))
-        ax6.set_xlabel('- %s'%(title_list[i]))
+        ax5.set_xlabel('- %s mc'%(title_list[i]))
+        ax6.set_xlabel('- %s mc'%(title_list[i]))
     else:
         ax5.set_xlabel('%s'%(title_list[i]))
         ax6.set_xlabel('%s'%(title_list[i]))
@@ -295,8 +298,8 @@ for i in range(1,len(MC_list),1):
     ax4.set_title('Electrons')
     ax5.set_title('Neutrons')
     ax6.set_title('Photons')
-    ax5.set_xlabel('Momentum')
-    ax6.set_xlabel('Momentum')
+    ax5.set_xlabel('Momentum mc')
+    ax6.set_xlabel('Momentum mc')
     fig.set_figwidth(20)
     fig.set_figheight(10)
     
@@ -306,11 +309,14 @@ for i in range(1,len(MC_list),1):
 
 
 ################
-
-theta_mc_fil = ak.Array(theta_mc[simID][booll])[boolean_photon]
-theta_rc_fil = ak.Array(theta_rc[recID][booll])[boolean_photon]
-phi_mc_fil = ak.Array(phi_mc[simID][booll])[boolean_photon]
-phi_rc_fil = ak.Array(phi_rc[recID][booll])[boolean_photon]
+if particle in particle_dict.keys():
+    boolean_particle = particle_dict[particle][0]
+else:
+    boolean_particle = boolean_photon
+theta_mc_fil = ak.Array(theta_mc[simID][booll])[boolean_particle]
+theta_rc_fil = ak.Array(theta_rc[recID][booll])[boolean_particle]
+phi_mc_fil = ak.Array(phi_mc[simID][booll])[boolean_particle]
+phi_rc_fil = ak.Array(phi_rc[recID][booll])[boolean_particle]
 
 theta_mc_fil_len = ak.count(theta_mc_fil,axis=None)
 theta_rc_fil_len = ak.count(theta_rc_fil,axis=None)
@@ -343,9 +349,14 @@ ax4.set_xlabel('- Theta rc')
 fig.set_figwidth(20)
 fig.set_figheight(10)
 
-fig.suptitle('Photons in %s  %s events'%(config,Nevents))
 
-plt.savefig(os.path.join(r_path, 'photons_%s.png' %  (config)))
+if particle in particle_dict.keys():
+    particle_name = particle_dict[particle][1]
+else:
+    particle_name = Photons
+
+fig.suptitle('%s in %s  %s events'%(particle_name,config,Nevents))
+plt.savefig(os.path.join(r_path, '%s_%s.png' %  (particle_name,config)))