add a script to plot the fit results as a function of energy

1 files changed, 293 insertions, 0 deletions

diff --git a/utils/plot-fit-results b/utils/plot-fit-results
new file mode 100755
index 0000000..b568075
--- /dev/null
+++ b/utils/plot-fit-results
@@ -0,0 +1,293 @@
+#!/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")
+
+matplotlib.rc('font', size=22)
+
+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()
+
+    events = []
+
+    for filename in args.filenames:
+        print(filename)
+        with open(filename) as f:
+            data = yaml.load(f.read())
+
+        a = np.ma.empty(len(data['data']),
+                        dtype=[('id',np.int),       # particle id
+                               ('T', np.double),    # true energy
+                               ('dx',np.double),    # dx
+                               ('dy',np.double),    # dy
+                               ('dz',np.double),    # dz
+                               ('dT',np.double),    # dT
+                               ('theta',np.double), # theta
+                               ('ratio',np.double), # likelihood ratio
+                               ('te',np.double),    # time electron
+                               ('tm',np.double),    # time muon
+                               ('Te',np.double)]    # electron energy
+                    )
+
+        for i, event in enumerate(data['data']):
+            # get the particle ID
+            id = event['mctk'][-1]['id']
+
+            a[i]['id'] = id
+
+            if id not in event['ev'][0]['fit']:
+                a[i] = np.ma.masked
+                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
+
+            fit = event['ev'][0]['fit']
+
+            a[i]['T'] = ke
+            energy = fit[id]['energy']
+            a[i]['dT'] = energy-ke
+
+            # store the fit position residuals
+            true_posx = event['mcvx'][0]['posx']
+            posx = fit[id]['posx']
+            a[i]['dx'] = posx-true_posx
+            true_posy = event['mcvx'][0]['posy']
+            posy = fit[id]['posy']
+            a[i]['dy'] = posy-true_posy
+            true_posz = event['mcvx'][0]['posz']
+            posz = fit[id]['posz']
+            a[i]['dz'] = posz-true_posz
+
+            # compute the angle between the fit direction and the true
+            # direction
+            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 = 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)
+            a[i]['theta'] = np.degrees(np.arccos(np.dot(true_dir,dir)))
+
+            # compute the log likelihood ratio
+            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']
+                a[i]['ratio'] = fmin_muon-fmin_electron
+            else:
+                a[i]['ratio'] = np.ma.masked
+
+            # store the time taken for electron and muon fits
+            if IDP_E_MINUS in fit:
+                a[i]['te'] = fit[IDP_E_MINUS]['time']
+                a[i]['Te'] = fit[IDP_E_MINUS]['energy']
+            else:
+                a[i]['te'] = np.ma.masked
+                a[i]['Te'] = np.ma.masked
+            if IDP_MU_MINUS in fit:
+                a[i]['tm'] = fit[IDP_MU_MINUS]['time']
+            else:
+                a[i]['tm'] = np.ma.masked
+
+        events.append(a)
+
+    a = np.concatenate(events)
+
+    bins = np.arange(50,1000,100)
+
+    stats_array = np.ma.empty(len(bins)-1,
+                     dtype=[('T',             np.double),
+                            ('dT',            np.double),
+                            ('dT_err',        np.double),
+                            ('dT_std',        np.double),
+                            ('dT_std_err',    np.double),
+                            ('dx',            np.double),
+                            ('dx_err',        np.double),
+                            ('dx_std',        np.double),
+                            ('dx_std_err',    np.double),
+                            ('dy',            np.double),
+                            ('dy_err',        np.double),
+                            ('dy_std',        np.double),
+                            ('dy_std_err',    np.double),
+                            ('dz',            np.double),
+                            ('dz_err',        np.double),
+                            ('dz_std',        np.double),
+                            ('dz_std_err',    np.double),
+                            ('theta',         np.double),
+                            ('theta_err',     np.double),
+                            ('theta_std',     np.double),
+                            ('theta_std_err', np.double)])
+
+    stats = {IDP_E_MINUS: stats_array, IDP_MU_MINUS: stats_array.copy()}
+
+    for id in stats:
+        electron_events = a[a['id'] == id]
+
+        for i, (ablah, b) in enumerate(zip(bins[:-1], bins[1:])):
+            events = electron_events[(electron_events['T'] >= ablah) & (electron_events['T'] < b)]
+
+            if len(events) < 2:
+                stats[id][i] = np.ma.masked
+                continue
+
+            stats[id][i]['T'] = (ablah+b)/2
+            mean, mean_error, std, std_error = get_stats(events['dT'].compressed())
+            stats[id][i]['dT'] = mean
+            stats[id][i]['dT_err'] = mean_error
+            stats[id][i]['dT_std'] = std
+            stats[id][i]['dT_std_err'] = std_error
+            mean, mean_error, std, std_error = get_stats(events['dx'].compressed())
+            stats[id][i]['dx'] = mean
+            stats[id][i]['dx_err'] = mean_error
+            stats[id][i]['dx_std'] = std
+            stats[id][i]['dx_std_err'] = std_error
+            mean, mean_error, std, std_error = get_stats(events['dy'].compressed())
+            stats[id][i]['dy'] = mean
+            stats[id][i]['dy_err'] = mean_error
+            stats[id][i]['dy_std'] = std
+            stats[id][i]['dy_std_err'] = std_error
+            mean, mean_error, std, std_error = get_stats(events['dz'].compressed())
+            stats[id][i]['dz'] = mean
+            stats[id][i]['dz_err'] = mean_error
+            stats[id][i]['dz_std'] = std
+            stats[id][i]['dz_std_err'] = std_error
+            mean, mean_error, std, std_error = get_stats(events['theta'].compressed())
+            stats[id][i]['theta'] = mean
+            stats[id][i]['theta_err'] = mean_error
+            stats[id][i]['theta_std'] = std
+            stats[id][i]['theta_std_err'] = std_error
+
+    for id in stats:
+        label = 'Muon' if id == IDP_MU_MINUS else 'Electron'
+
+        T = stats[id]['T']
+        dT = stats[id]['dT']
+        dT_err = stats[id]['dT_err']
+        std_dT = stats[id]['dT_std']
+        std_dT_err = stats[id]['dT_std_err']
+        dx = stats[id]['dx']
+        dx_err = stats[id]['dx_err']
+        std_dx = stats[id]['dx_std']
+        std_dx_err = stats[id]['dx_std_err']
+        dy = stats[id]['dy']
+        dy_err = stats[id]['dy_err']
+        std_dy = stats[id]['dy_std']
+        std_dy_err = stats[id]['dy_std_err']
+        dz = stats[id]['dz']
+        dz_err = stats[id]['dz_err']
+        std_dz = stats[id]['dz_std']
+        std_dz_err = stats[id]['dz_std_err']
+        theta = stats[id]['theta']
+        theta_err = stats[id]['theta_err']
+        std_theta = stats[id]['theta_std']
+        std_theta_err = stats[id]['theta_std_err']
+
+        plt.figure(1)
+        plt.errorbar(T,dT*100/T,yerr=dT_err*100/T,fmt='o',label=label)
+        plt.xlabel("Kinetic Energy (MeV)")
+        plt.ylabel("Energy bias (%)")
+        plt.title("Energy Bias")
+        plt.legend()
+
+        plt.figure(2)
+        plt.errorbar(T,std_dT*100/T,yerr=std_dT_err*100/T,fmt='o',label=label)
+        plt.xlabel("Kinetic Energy (MeV)")
+        plt.ylabel("Energy resolution (%)")
+        plt.title("Energy Resolution")
+        plt.legend()
+
+        plt.figure(3)
+        plt.errorbar(T,dx,yerr=dx_err,fmt='o',label='%s (x)' % label)
+        plt.errorbar(T,dy,yerr=dy_err,fmt='o',label='%s (y)' % label)
+        plt.errorbar(T,dz,yerr=dz_err,fmt='o',label='%s (z)' % label)
+        plt.xlabel("Kinetic Energy (MeV)")
+        plt.ylabel("Position bias (cm)")
+        plt.title("Position Bias")
+        plt.legend()
+
+        plt.figure(4)
+        plt.errorbar(T,std_dx,yerr=std_dx_err,fmt='o',label='%s (x)' % label)
+        plt.errorbar(T,std_dy,yerr=std_dy_err,fmt='o',label='%s (y)' % label)
+        plt.errorbar(T,std_dz,yerr=std_dz_err,fmt='o',label='%s (z)' % label)
+        plt.xlabel("Kinetic Energy (MeV)")
+        plt.ylabel("Position resolution (cm)")
+        plt.title("Position Resolution")
+        plt.ylim((0,plt.gca().get_ylim()[1]))
+        plt.legend()
+
+        plt.figure(5)
+        plt.errorbar(T,std_theta,yerr=std_theta_err,fmt='o',label=label)
+        plt.xlabel("Kinetic Energy (MeV)")
+        plt.ylabel("Angular resolution (deg)")
+        plt.title("Angular Resolution")
+        plt.ylim((0,plt.gca().get_ylim()[1]))
+        plt.legend()
+
+        plt.figure(6)
+        plt.scatter(a[a['id'] == id]['Te'],a[a['id'] == id]['ratio'],label=label)
+        plt.xlabel("Reconstructed Electron Energy (MeV)")
+        plt.ylabel(r"Log Likelihood Ratio (e/$\mu$)")
+        plt.title("Log Likelihood Ratio vs Reconstructed Electron Energy")
+        plt.legend()
+    plt.show()