Update benchmarks/dis/analysis/truth_reconstruction.py

benchmarks/dis/analysis/truth_reconstruction.py

@@ -111,9 +111,17 @@ for i in range(len(MC_list)): #Repeat the following steps for each variable (mom
     Y_plot = list(np.zeros(len(X_list)))
 
     for j in range(len(X_list)): #Repeat the following steps for each particle (pions,protons,electrons,neutrons,photons)
-        X_s = np.array(ak.flatten(X_list[j])) 
-        Y_s = np.array(ak.flatten(Y_list[j]))
-        if i == 0:   #Momentum ratio
+        X = X_list[j]
+        Y = Y_list[j]
+        X_len = ak.count(X,axis=None)
+        Y_len = ak.count(Y,axis=None)
+        if X_len > Y_len: 
+            F_boolean = np.ones_like(Y) == 1
+        else: 
+            F_boolean = np.ones_like(X) == 1
+        X_s = np.array(ak.flatten(X[F_boolean])) #Filtered lists
+        Y_s = np.array(ak.flatten(Y[F_boolean]))
+        if i == 0:   #Momentum
             ratio = np.array((ak.Array(Y_s)/ak.Array(X_s)))
         else: #Angle difference
             ratio = np.array((ak.Array(Y_s)-(ak.Array(X_s))))
@@ -174,8 +182,16 @@ for i in range(len(MC_list)): #Repeat the following steps for each variable (mom
 
     if i > 0: #for each variable theta, phi, and eta
         for j in range(len(M_list)): #Repeat the following steps for each particle (pions,protons,electrons,neutrons,photons)
-            M_s = np.array(ak.flatten(M_list[j]))
+            X = X_list[j]
+            Y = Y_list[j]
+            M_mc = M_list[j]
+            boolean_M = np.ones_like(M_mc) == 1
+            X_s = np.array(ak.flatten(X[boolean_M])) 
+            Y_s = np.array(ak.flatten(Y[boolean_M])) 
+            M_s = np.array(ak.flatten(M_mc))
+            ratio = np.array((ak.Array(Y_s)-(ak.Array(X_s))))
             X_plot[j] = M_s
+            Y_plot[j] = ratio
 
         fig = plt.figure()
         gs = fig.add_gridspec(3, 2, wspace=0)
@@ -207,8 +223,15 @@ for i in range(len(MC_list)): #Repeat the following steps for each variable (mom
 ###################################################################################################
 
     #Repeat the following steps for each variable (momentum,theta,phi,eta)
-    X_s = np.array(ak.flatten(X1)) 
-    Y_s = np.array(ak.flatten(Y1)) 
+    X_len = ak.count(X1,axis=None)
+    Y_len = ak.count(Y1,axis=None)
+    if X_len > Y_len: 
+        F_boolean = np.ones_like(Y1) == 1
+    else: 
+        F_boolean = np.ones_like(X1) == 1
+    X_s = np.array(ak.flatten(X1[F_boolean])) 
+    Y_s = np.array(ak.flatten(Y1[F_boolean])) 
+ 
     #Histogram
     if i == 0 and particle in particle_dict.keys(): #Momentum in Single events
         h, xedges, yedges, image = plt.hist2d(x=X_s,y= Y_s,  bins = 11) 
@@ -250,11 +273,23 @@ for i in range(len(MC_list)): #Repeat the following steps for each variable (mom
 
 def particle_plots(boolean_particle):
     #filtered lists w.r.t the particle
-    theta_mc_fil = ak.Array(theta_mc[simID][booll])[boolean_particle] 
+    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]
-    ratio = np.array((ak.Array(theta_rc_fil)-(ak.Array(theta_mc_fil))))
+
+    theta_mc_fil_len = ak.count(theta_mc_fil,axis=None)
+    theta_rc_fil_len = ak.count(theta_rc_fil,axis=None)
+    if theta_mc_fil_len > theta_rc_fil_len:
+        F_boolean = np.ones_like(theta_rc_fil) == 1
+    else: 
+        F_boolean = np.ones_like(theta_mc_fil) == 1
+    #filtered lists w.r.t length
+    theta_mc_F = np.array(ak.flatten(theta_mc_fil[F_boolean]))
+    theta_rc_F = np.array(ak.flatten(theta_rc_fil[F_boolean]))
+    phi_mc_F = np.array(ak.flatten(phi_mc_fil[F_boolean]))
+    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.01)