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

4 files changed, 23 insertions, 7 deletions

diff --git a/utils/cat-grid-jobs b/utils/cat-grid-jobs
index 43fef50..84623b8 100755
--- a/utils/cat-grid-jobs
+++ b/utils/cat-grid-jobs
@@ -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
 
 if __name__ == '__main__':
     import argparse
@@ -30,7 +34,7 @@ if __name__ == '__main__':
 
     for filename in args.filenames:
         with open(filename) as f:
-            data = yaml.load(f.read(),Loader=yaml.loader.SafeLoader)
+            data = yaml.load(f.read(),Loader=Loader)
 
         if data is None:
             continue
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'])):
diff --git a/utils/plot-fit-results b/utils/plot-fit-results
index e1f7987..773a0dc 100755
--- a/utils/plot-fit-results
+++ b/utils/plot-fit-results
@@ -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
@@ -81,7 +85,7 @@ 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)
 
         a = np.ma.empty(len(data['data']),
                         dtype=[('id',np.int),       # particle id
diff --git a/utils/submit-grid-jobs b/utils/submit-grid-jobs
index 13bddfe..37b3303 100755
--- a/utils/submit-grid-jobs
+++ b/utils/submit-grid-jobs
@@ -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 string
 from os.path import split, splitext, join, abspath
 import uuid
@@ -127,7 +131,7 @@ if __name__ == '__main__':
         else:
             output = check_output([join(args.dir,"zdab-cat"),filename])
 
-        data = yaml.load(output,Loader=yaml.loader.SafeLoader)
+        data = yaml.load(output,Loader=Loader)
 
         filename = abspath(filename)