rename plot.py -> plot

1 files changed, 167 insertions, 0 deletions

diff --git a/utils/plot b/utils/plot
new file mode 100755
index 0000000..c6d8b4f
--- /dev/null
+++ b/utils/plot
@@ -0,0 +1,167 @@
+#!/usr/bin/env python
+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
+# Qt5 will scale everything using the dpi in ~/.Xresources
+import matplotlib
+matplotlib.use("Qt5Agg")
+
+IDP_E_MINUS  =    20
+IDP_MU_MINUS =    22
+
+SNOMAN_MASS = {
+    20: 0.511,
+    21: 0.511,
+    22: 105.658,
+    23: 105.658
+}
+
+def plot_hist(x, label=None):
+    # determine the bin width using the Freedman Diaconis rule
+    # see https://en.wikipedia.org/wiki/Freedman%E2%80%93Diaconis_rule
+    h = 2*iqr(x)/len(x)**(1/3)
+    n = max(int((np.max(x)-np.min(x))/h),10)
+    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.
+
+    The formula for the standard error on the standard deviation comes from
+    https://stats.stackexchange.com/questions/156518.
+    """
+    mean = np.mean(x)
+    std = np.std(x)
+    n = len(x)
+    u4 = np.mean((x-mean)**4)
+    error = np.sqrt((u4-(n-3)*std**4/(n-1))/n)/(2*std)
+    return mean, std/np.sqrt(n), std, error
+
+if __name__ == '__main__':
+    import argparse
+    import matplotlib.pyplot as plt
+    import numpy as np
+
+    parser = argparse.ArgumentParser("plot fit results")
+    parser.add_argument("filenames", nargs='+', help="input files")
+    args = parser.parse_args()
+
+    for filename in args.filenames:
+        print(filename)
+        with open(filename) as f:
+            data = yaml.load(f.read())
+
+        dx = []
+        dy = []
+        dz = []
+        dT = []
+        thetas = []
+        likelihood_ratio = []
+        t_electron = []
+        t_muon = []
+        for event in data['data']:
+            # get the particle ID
+            id = event['mctk'][-1]['id']
+
+            if id not in event['ev'][0]['fit']:
+                continue
+
+            mass = SNOMAN_MASS[id]
+            # for some reason it's the *second* track which seems to contain the
+            # initial energy
+            true_energy = event['mctk'][-1]['energy']
+            # 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']
+            dT.append(energy-ke)
+            true_posx = event['mcvx'][0]['posx']
+            posx = event['ev'][0]['fit'][id]['posx']
+            dx.append(posx-true_posx)
+            true_posy = event['mcvx'][0]['posy']
+            posy = event['ev'][0]['fit'][id]['posy']
+            dy.append(posy-true_posy)
+            true_posz = event['mcvx'][0]['posz']
+            posz = event['ev'][0]['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']
+            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']
+                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))
+        print("std(dT) = %.2g +/- %.2g" % (std, std_error))
+        mean, mean_error, std, std_error = get_stats(dx)
+        print("dx      = %4.2g +/- %.2g" % (mean, mean_error))
+        print("std(dx) = %4.2g +/- %.2g" % (std, std_error))
+        mean, mean_error, std, std_error = get_stats(dy)
+        print("dy      = %4.2g +/- %.2g" % (mean, mean_error))
+        print("std(dy) = %4.2g +/- %.2g" % (std, std_error))
+        mean, mean_error, std, std_error = get_stats(dz)
+        print("dz      = %4.2g +/- %.2g" % (mean, mean_error))
+        print("std(dz) = %4.2g +/- %.2g" % (std, std_error))
+        mean, mean_error, std, std_error = get_stats(thetas)
+        print("std(theta) = %4.2g +/- %.2g" % (std, std_error))
+
+        plt.figure(1)
+        plot_hist(dT, label=filename)
+        plt.xlabel("Kinetic Energy difference (MeV)")
+        plt.figure(2)
+        plot_hist(dx, label=filename)
+        plt.xlabel("X Position difference (cm)")
+        plt.figure(3)
+        plot_hist(dy, label=filename)
+        plt.xlabel("Y Position difference (cm)")
+        plt.figure(4)
+        plot_hist(dz, label=filename)
+        plt.xlabel("Z Position difference (cm)")
+        plt.figure(5)
+        plot_hist(thetas, label=filename)
+        plt.xlabel(r"$\theta$ (deg)")
+        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)")
+
+    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()