try to import CLoader if possible since it's *much* faster

1 files changed, 8 insertions, 4 deletions

diff --git a/utils/plot-energy b/utils/plot-energy
index 5322ff3..a49fcb0 100755
--- a/utils/plot-energy
+++ b/utils/plot-energy
@@ -16,6 +16,10 @@
 
 from __future__ import print_function, division
 import yaml
+try:
+    from yaml import CLoader as Loader
+except ImportError:
+    from yaml.loader import SafeLoader as Loader
 import numpy as np
 from scipy.stats import iqr
 from matplotlib.lines import Line2D
@@ -26,9 +30,6 @@ from matplotlib.lines import Line2D
 import matplotlib
 matplotlib.use("Qt5Agg")
 
-IDP_E_MINUS  =    20
-IDP_MU_MINUS =    22
-
 SNOMAN_MASS = {
     20: 0.511,
     21: 0.511,
@@ -64,13 +65,16 @@ if __name__ == '__main__':
     for filename in args.filenames:
         print(filename)
         with open(filename) as f:
-            data = yaml.load(f.read())
+            data = yaml.load(f.read(),Loader=Loader)
 
         for i, event in enumerate(data['data']):
             for ev in event['ev']:
                 if 'fit' not in ev:
                     continue
                 for id, fit_result in ev['fit'].iteritems():
+                    # FIXME: Should I just store the particle ids in the YAML
+                    # output as a list of particle ids instead of a single
+                    # integer?
                     ids = map(int,chunks(str(id),2))
                     energy = 0.0
                     for i, ke in zip(ids,np.atleast_1d(fit_result['energy'])):