update plot

1 files changed, 25 insertions, 19 deletions

diff --git a/utils/plot b/utils/plot
index 06d9951..59eb4c2 100755
--- a/utils/plot
+++ b/utils/plot
@@ -83,6 +83,8 @@ if __name__ == '__main__':
             if id not in event['ev'][0]['fit']:
                 continue
 
+            fit = event['ev'][0]['fit']
+
             mass = SNOMAN_MASS[id]
             # for some reason it's the *second* track which seems to contain the
             # initial energy
@@ -90,37 +92,39 @@ if __name__ == '__main__':
             # The MCTK bank has the particle's total energy (except for neutrons)
             # so we need to convert it into kinetic energy
             ke = true_energy - mass
-            energy = event['ev'][0]['fit'][id]['energy']
+            energy = fit[id]['energy']
             dT.append(energy-ke)
             true_posx = event['mcvx'][0]['posx']
-            posx = event['ev'][0]['fit'][id]['posx']
+            posx = fit[id]['posx']
             dx.append(posx-true_posx)
             true_posy = event['mcvx'][0]['posy']
-            posy = event['ev'][0]['fit'][id]['posy']
+            posy = fit[id]['posy']
             dy.append(posy-true_posy)
             true_posz = event['mcvx'][0]['posz']
-            posz = event['ev'][0]['fit'][id]['posz']
+            posz = fit[id]['posz']
             dz.append(posz-true_posz)
             dirx = event['mctk'][-1]['dirx']
             diry = event['mctk'][-1]['diry']
             dirz = event['mctk'][-1]['dirz']
             true_dir = [dirx,diry,dirz]
             true_dir = np.array(true_dir)/np.linalg.norm(true_dir)
-            theta = event['ev'][0]['fit'][id]['theta']
-            phi = event['ev'][0]['fit'][id]['phi']
+            theta = fit[id]['theta']
+            phi = fit[id]['phi']
             dir = [np.sin(theta)*np.cos(phi),np.sin(theta)*np.sin(phi),np.cos(theta)]
             dir = np.array(dir)/np.linalg.norm(dir)
             thetas.append(np.degrees(np.arccos(np.dot(true_dir,dir))))
-            if IDP_E_MINUS in event['ev'][0]['fit'] and IDP_MU_MINUS in event['ev'][0]['fit']:
-                fmin_electron = event['ev'][0]['fit'][IDP_E_MINUS]['fmin']
-                fmin_muon = event['ev'][0]['fit'][IDP_MU_MINUS]['fmin']
+            if IDP_E_MINUS in fit and IDP_MU_MINUS in fit:
+                fmin_electron = fit[IDP_E_MINUS]['fmin']
+                fmin_muon = fit[IDP_MU_MINUS]['fmin']
                 likelihood_ratio.append(fmin_muon-fmin_electron)
-            if IDP_E_MINUS in event['ev'][0]['fit']:
-                t_electron.append(event['ev'][0]['fit'][IDP_E_MINUS]['time'])
-            if IDP_MU_MINUS in event['ev'][0]['fit']:
-                t_muon.append(event['ev'][0]['fit'][IDP_MU_MINUS]['time'])
-            nhit = event['ev'][0]['nhit']
-            psi.append(event['ev'][0]['fit'][id]['psi']/nhit)
+            if IDP_E_MINUS in fit:
+                t_electron.append(fit[IDP_E_MINUS]['time'])
+            if IDP_MU_MINUS in fit:
+                t_muon.append(fit[IDP_MU_MINUS]['time'])
+
+            if 'nhit' in event['ev'][0]:
+                nhit = event['ev'][0]['nhit']
+                psi.append(fit[id]['psi']/nhit)
 
         if len(t_electron) < 2 or len(t_muon) < 2:
             continue
@@ -164,9 +168,10 @@ if __name__ == '__main__':
         plt.figure(8)
         plot_hist(np.array(t_muon)/1e3/60.0, label=filename)
         plt.xlabel(r"Muon Fit time (minutes)")
-        plt.figure(9)
-        plot_hist(psi, label=filename)
-        plt.xlabel(r"$\Psi$/Nhit")
+        if len(psi):
+            plt.figure(9)
+            plot_hist(psi, label=filename)
+            plt.xlabel(r"$\Psi$/Nhit")
 
     plot_legend(1)
     plot_legend(2)
@@ -176,5 +181,6 @@ if __name__ == '__main__':
     plot_legend(6)
     plot_legend(7)
     plot_legend(8)
-    plot_legend(9)
+    if len(psi):
+        plot_legend(9)
     plt.show()