update plot-muons script so it works outside the utils/ directory

1 files changed, 21 insertions, 1 deletions

diff --git a/utils/sddm/plot_energy.py b/utils/sddm/plot_energy.py
index a4b7005..0ba4076 100755
--- a/utils/sddm/plot_energy.py
+++ b/utils/sddm/plot_energy.py
@@ -22,6 +22,8 @@ from .dc import DC_MUON, DC_JUNK, DC_CRATE_ISOTROPY, DC_QVNHIT, DC_NECK, DC_FLAS
 from . import grouper, print_warning, AV_RADIUS, PSUP_RADIUS, read_hdf
 import pandas as pd
 from scipy.misc import logsumexp
+from os.path import exists, join
+from os import environ
 
 EPSILON = 1e-10
 
@@ -383,9 +385,27 @@ def get_dx(row):
     else:
         return 0
 
+def find_file(filename):
+    """
+    Returns the full path to `filename` by searching the current working
+    directory and the path specified by the environment variable SDDM_DATA.
+
+    If the file isn't found in either place, just returns the filename.
+    """
+    if exists(filename):
+        return filename
+
+    if 'SDDM_DATA' in environ:
+        path = join(environ['SDDM_DATA'],filename)
+
+        if exists(path):
+            return path
+
+    return filename
+
 def dx_to_energy(dx):
     lines = []
-    with open("../src/muE_water_liquid.txt") as f:
+    with open(find_file("muE_water_liquid.txt")) as f:
         for i, line in enumerate(f):
             if i < 10:
                 continue