add a script to plot fit results

1 files changed, 78 insertions, 0 deletions

diff --git a/src/plot.py b/src/plot.py
new file mode 100755
index 0000000..6a2fe73
--- /dev/null
+++ b/src/plot.py
@@ -0,0 +1,78 @@
+#!/usr/bin/env python
+from __future__ import print_function, division
+import yaml
+
+SNOMAN_MASS = {
+    20: 0.511,
+    21: 0.511,
+    22: 105.658,
+    23: 105.658
+}
+
+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 = []
+        e = []
+        for event in data['data']:
+            # get the particle ID
+            id = event['mctk'][-1]['id']
+            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'][0]['energy']
+            e.append(energy-ke)
+            true_posx = event['mcvx'][0]['posx']
+            posx = event['ev'][0]['fit'][0]['posx']
+            dx.append(posx-true_posx)
+            true_posy = event['mcvx'][0]['posy']
+            posy = event['ev'][0]['fit'][0]['posy']
+            dy.append(posy-true_posy)
+            true_posz = event['mcvx'][0]['posz']
+            posz = event['ev'][0]['fit'][0]['posz']
+            dz.append(posz-true_posz)
+
+        print("dT = %.2g +/- %.2g" % (np.mean(e), np.std(e)))
+        print("dx = %4.2g +/- %.2g" % (np.mean(dx), np.std(dx)))
+        print("dy = %4.2g +/- %.2g" % (np.mean(dy), np.std(dy)))
+        print("dz = %4.2g +/- %.2g" % (np.mean(dz), np.std(dz)))
+
+        plt.figure(1)
+        plt.hist(e, histtype='step', label=filename)
+        plt.xlabel("Kinetic Energy difference (MeV)")
+        plt.figure(2)
+        plt.hist(dx, histtype='step', label=filename)
+        plt.xlabel("X Position difference (cm)")
+        plt.figure(3)
+        plt.hist(dy, histtype='step', label=filename)
+        plt.xlabel("Y Position difference (cm)")
+        plt.figure(4)
+        plt.hist(dz, histtype='step', label=filename)
+        plt.xlabel("Z Position difference (cm)")
+
+    plt.figure(1)
+    plt.legend()
+    plt.figure(2)
+    plt.legend()
+    plt.figure(3)
+    plt.legend()
+    plt.figure(4)
+    plt.legend()
+    plt.show()