update plot.py to plot the distribution of fit times

1 files changed, 28 insertions, 12 deletions

diff --git a/utils/plot.py b/utils/plot.py
index 7abb7e3..c6d8b4f 100755
--- a/utils/plot.py
+++ b/utils/plot.py
@@ -3,6 +3,7 @@ from __future__ import print_function, division
 import yaml
 import numpy as np
 from scipy.stats import iqr
+from matplotlib.lines import Line2D
 
 # on retina screens, the default plots are way too small
 # by using Qt5 and setting QT_AUTO_SCREEN_SCALE_FACTOR=1
@@ -28,6 +29,13 @@ def plot_hist(x, label=None):
     bins = np.linspace(np.min(x),np.max(x),n)
     plt.hist(x, bins=bins, histtype='step', label=label)
 
+def plot_legend(n):
+    plt.figure(n)
+    ax = plt.gca()
+    handles, labels = ax.get_legend_handles_labels()
+    new_handles = [Line2D([],[],c=h.get_edgecolor()) for h in handles]
+    plt.legend(handles=new_handles,labels=labels)
+
 def get_stats(x):
     """
     Returns a tuple (mean, error mean, std, error std) for the values in x.
@@ -62,6 +70,8 @@ if __name__ == '__main__':
         dT = []
         thetas = []
         likelihood_ratio = []
+        t_electron = []
+        t_muon = []
         for event in data['data']:
             # get the particle ID
             id = event['mctk'][-1]['id']
@@ -101,6 +111,10 @@ if __name__ == '__main__':
                 fmin_electron = event['ev'][0]['fit'][IDP_E_MINUS]['fmin']
                 fmin_muon = event['ev'][0]['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'])
 
         mean, mean_error, std, std_error = get_stats(dT)
         print("dT      = %.2g +/- %.2g" % (mean, mean_error))
@@ -135,17 +149,19 @@ if __name__ == '__main__':
         plt.figure(6)
         plot_hist(likelihood_ratio, label=filename)
         plt.xlabel(r"Log Likelihood Ratio ($e/\mu$)")
+        plt.figure(7)
+        plot_hist(np.array(t_electron)/1e3/60.0, label=filename)
+        plt.xlabel(r"Electron Fit time (minutes)")
+        plt.figure(8)
+        plot_hist(np.array(t_muon)/1e3/60.0, label=filename)
+        plt.xlabel(r"Muon Fit time (minutes)")
 
-    plt.figure(1)
-    plt.legend()
-    plt.figure(2)
-    plt.legend()
-    plt.figure(3)
-    plt.legend()
-    plt.figure(4)
-    plt.legend()
-    plt.figure(5)
-    plt.legend()
-    plt.figure(6)
-    plt.legend()
+    plot_legend(1)
+    plot_legend(2)
+    plot_legend(3)
+    plot_legend(4)
+    plot_legend(5)
+    plot_legend(6)
+    plot_legend(7)
+    plot_legend(8)
     plt.show()