simplify budget loading with new load_budget() function

This replaces load_ifo(), and returns only the Budget class.

Any ifo.yaml found in a budget package is loaded as ifo attribute.  other
stuff can just be assumed to be provided as class attributes as well.

README.md

@@ -16,10 +16,10 @@ description is loaded, the noise budget can be calculated and plotted:
 >>> import gwinc
 >>> import numpy as np
 >>> freq = np.logspace(1, 3, 1000)
->>> Budget, ifo, freq_, plot_style = gwinc.load_ifo('aLIGO')
->>> ifo = gwinc.precompIFO(freq, ifo)
+>>> Budget = gwinc.load_budget('aLIGO')
+>>> ifo = gwinc.precompIFO(freq, Budget.ifo)
 >>> traces = Budget(freq, ifo=ifo).calc_trace()
->>> fig = gwinc.plot_noise(freq, traces, **plot_style)
+>>> fig = gwinc.plot_noise(freq, traces)
 >>> fig.show()
 ```
 

gwinc/__init__.py

-from .ifo import available_ifos, load_ifo
+import os
+import logging
+import importlib
+
+from .ifo import available_ifos
 from .struct import Struct
 from .precomp import precompIFO
 from .plot import plot_noise
 from .io import load_hdf5, save_hdf5
+
+
+def _load_module(name_or_path):
+    """Load module from name or path.
+
+    Return loaded module and module path.
+
+    """
+    if os.path.exists(name_or_path):
+        path = name_or_path.rstrip('/')
+        modname = os.path.splitext(os.path.basename(path))[0]
+        if os.path.isdir(path):
+            path = os.path.join(path, '__init__.py')
+        spec = importlib.util.spec_from_file_location(modname, path)
+        mod = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(mod)
+    else:
+        mod = importlib.import_module(name_or_path)
+    try:
+        path = mod.__path__[0]
+    except AttributeError:
+        path = mod.__file__
+    return mod, path
+
+
+def load_budget(name_or_path):
+    """Load GWINC IFO Budget by name or from path.
+
+    Named IFOs should correspond to one of the IFOs available in the
+    gwinc package (see gwinc.available_ifos()).  If a path is provided
+    it should either be a budget package (directory) or module (ending
+    in .py), or an IFO struct definition (see gwinc.Struct).
+
+    If a budget package path is provided, and the package includes an
+    'ifo.yaml' file, that file will be loaded into a Struct and
+    assigned as an attribute to the returned Budget class.
+
+    If a bare struct is provided the base aLIGO budget definition will
+    be assumed.
+
+    Returns a Budget class with 'ifo', 'freq', and 'plot_style', ifo
+    structure, frequency array, and plot style dictionary, with the
+    last three being None if they are not defined in the budget.
+
+    """
+    ifo = None
+
+    if os.path.exists(name_or_path):
+        path = name_or_path.rstrip('/')
+        bname, ext = os.path.splitext(os.path.basename(path))
+
+        if ext in Struct.STRUCT_EXT:
+            logging.info("loading struct {}...".format(path))
+            ifo = Struct.from_file(path)
+            bname = 'aLIGO'
+            modname = 'gwinc.ifo.aLIGO'
+            logging.info("loading budget {}...".format(modname))
+
+        else:
+            modname = path
+            logging.info("loading module path {}...".format(modname))
+
+    else:
+        if name_or_path not in available_ifos():
+            raise RuntimeError("Unknonw IFO '{}' (available IFOs: {}).".format(
+                name_or_path,
+                available_ifos(),
+            ))
+        bname = name_or_path
+        modname = 'gwinc.ifo.'+name_or_path
+        logging.info("loading module {}...".format(modname))
+
+    mod, modpath = _load_module(modname)
+
+    Budget = getattr(mod, bname)
+    ifopath = os.path.join(modpath, 'ifo.yaml')
+    if not ifo and ifopath:
+        ifo = Struct.from_file(ifopath)
+    Budget.ifo = ifo
+
+    return Budget

gwinc/__main__.py

@@ -9,9 +9,8 @@ import logging
 logging.basicConfig(format='%(message)s',
                     level=os.getenv('LOG_LEVEL', logging.INFO))
 
-from .ifo import available_ifos, load_ifo
+from . import available_ifos, load_budget, plot_noise
 from .precomp import precompIFO
-from . import plot_noise
 from . import io
 
 ##################################################
@@ -97,10 +96,12 @@ def main():
         plot_style = attrs
 
     else:
-        Budget, ifo, freq, plot_style = load_ifo(args.IFO)
+        Budget = load_budget(args.IFO)
+        ifo = Budget.ifo
         # FIXME: this should be done only if specified, to allow for
-        # using any FREQ specified in budget module by default
+        # using any FREQ specified in the Budget
         freq = np.logspace(np.log10(args.flo), np.log10(args.fhi), args.npoints)
+        plot_style = getattr(Budget, 'plot_style', {})
         traces = None
 
     if args.yaml:
@@ -117,7 +118,7 @@ def main():
         if not ifo:
             parser.exit(2, "no IFO structure available.")
         fmt = '{:30} {:>20} {:>20}'
-        Budget, ifoo, freq, plot_style = load_ifo(args.diff)
+        Budget = load_ifo(args.diff)
         diffs = ifo.diff(ifoo)
         if diffs:
             print(fmt.format('', args.IFO, args.diff))

gwinc/gwinc.py

 from __future__ import division
 import numpy as np
 from numpy import pi, sqrt
-from collections import OrderedDict
 import logging
 
-from . import load_ifo
+from . import load_budget, plot_noise
 from .precomp import precompIFO
-from . import noise
-from .plot import plot_noise
 
 
 def gwinc(freq, ifoin, source=None, plot=False, PRfixed=True):
@@ -29,13 +26,14 @@ def gwinc(freq, ifoin, source=None, plot=False, PRfixed=True):
     # assume generic aLIGO configuration
     # FIXME: how do we allow adding arbitrary addtional noise sources
     # from just ifo description, without having to specify full budget
-    Budget, ifo_, freq_, plot_style = load_ifo('aLIGO')
+    Budget = load_budget('aLIGO')
 
     # add some precomputed info to the ifo struct
     # this implicitly deepcopies and the return value is the copy
     ifo = precompIFO(freq, ifoin, PRfixed)
 
     traces = Budget(freq, ifo=ifo).calc_trace()
+    plot_style = getattr(Budget, 'plot_style', {})
 
     # construct matgwinc-compatible noises structure
     noises = {}

gwinc/ifo/__init__.py

 import os
-import logging
-
-from ..struct import Struct
-from ..util import load_module
 
 
 PLOT_STYLE = dict(
@@ -10,6 +6,20 @@ PLOT_STYLE = dict(
 )
 
 
+def lpath(file0, file1):
+    """Return path of file1 when expressed relative to file0.
+
+    For instance, if file0 is "/path/to/file0" and file1 is
+    "../for/file1" then what is returned is "/path/for/file1".
+
+    This is useful for resolving paths within packages with e.g.:
+
+      rpath = lpath(__file__, '../resource.yaml')
+
+    """
+    return os.path.abspath(os.path.join(os.path.dirname(file0), file1))
+
+
 def available_ifos():
     """List available included pre-defined IFOs
 
@@ -20,57 +30,3 @@ def available_ifos():
         if os.path.isdir(os.path.join(root, f)) and f[0] != '_':
             ifos.append(f)
     return sorted(ifos)
-
-
-def load_ifo(name_or_path):
-    """Load GWINC IFO Budget by name or from file.
-
-    Named IFOs should correspond to one of the IFOs available in the
-    gwinc package (see gwinc.available_ifos()).  If a path is provided
-    it should either be a budget package (directory) or module (ending
-    in .py), or an IFO struct (see gwinc.Struct).  In the latter case
-    the base aLIGO budget definition will be used.
-
-    Returns primary Budget class, ifo structure, frequency array, and
-    plot style dictionary, with the last three being None if they are
-    not defined in the budget.
-
-    """
-    ifo = None
-
-    if os.path.exists(name_or_path):
-        path = name_or_path.rstrip('/')
-        bname, ext = os.path.splitext(os.path.basename(path))
-
-        if ext in Struct.STRUCT_EXT:
-            logging.info("loading struct {}...".format(path))
-            ifo = Struct.from_file(path)
-            bname = 'aLIGO'
-            modname = 'gwinc.ifo.aLIGO'
-            logging.info("loading budget {}...".format(modname))
-
-        else:
-            modname = path
-            logging.info("loading module path {}...".format(modname))
-
-    else:
-        if name_or_path not in available_ifos():
-            raise RuntimeError("Unknonw IFO '{}' (available IFOs: {}).".format(
-                name_or_path,
-                available_ifos(),
-            ))
-        bname = name_or_path
-        modname = 'gwinc.ifo.'+name_or_path
-        logging.info("loading module {}...".format(modname))
-
-    mod, modpath = load_module(modname)
-
-    Budget = getattr(mod, bname)
-    ifo = getattr(mod, 'IFO', ifo)
-    ifopath = os.path.join(modpath, 'ifo.yaml')
-    if not ifo and ifopath:
-        ifo = Struct.from_file(ifopath)
-    freq = getattr(mod, 'FREQ', None)
-    plot_style = getattr(mod, 'PLOT_STYLE', PLOT_STYLE)
-
-    return Budget, ifo, freq, plot_style

gwinc/test/__main__.py

@@ -12,7 +12,7 @@ import logging
 logging.basicConfig(format='%(message)s',
                     level=os.getenv('LOG_LEVEL', logging.INFO))
 
-from .. import load_ifo
+from .. import load_budget
 from ..gwinc import gwinc
 from ..gwinc_matlab import gwinc_matlab
 try:
@@ -58,7 +58,8 @@ def main():
     args = parser.parse_args()
 
     logging.info("loading IFO '{}'...".format(args.IFO))
-    Budget, ifo, freq, plot_style = load_ifo(args.IFO)
+    Budget = load_budget(args.IFO)
+    ifo = Budget.ifo
 
     freq = np.logspace(np.log10(FLO), np.log10(FHI), NPOINTS)
 

gwinc/util.py

-import os
-import importlib
-
-
-def lpath(file0, file1):
-    """Return path of file1 when expressed relative to file0.
-
-    For instance, if file0 is "/path/to/file0" and file1 is
-    "../for/file1" then what is returned is "/path/for/file1".
-
-    This is useful for resolving paths within packages with e.g.:
-
-      rpath = lpath(__file__, '../resource.yaml')
-
-    """
-    return os.path.abspath(os.path.join(os.path.dirname(file0), file1))
-
-
-def load_module(name_or_path):
-    """Load module from name or path.
-
-    Return loaded module and module path.
-
-    """
-    if os.path.exists(name_or_path):
-        path = name_or_path.rstrip('/')
-        modname = os.path.splitext(os.path.basename(path))[0]
-        if os.path.isdir(path):
-            path = os.path.join(path, '__init__.py')
-        spec = importlib.util.spec_from_file_location(modname, path)
-        mod = importlib.util.module_from_spec(spec)
-        spec.loader.exec_module(mod)
-    else:
-        mod = importlib.import_module(name_or_path)
-    try:
-        path = mod.__path__[0]
-    except AttributeError:
-        path = mod.__file__
-    return mod, path