API change: load_budget function returns instantiated Budget object

This addresses issues about confusion over the IFO being set as a class
attribute of the Budget, and is required to close #69.

README.md

@@ -106,7 +106,7 @@ $ python3 -m gwinc -h
 ```
 
 
-### python library
+### library interface
 
 For custom plotting, parameter optimization, etc. all functionality can be
 accessed directly through the `gwinc` library interface:
@@ -114,17 +114,18 @@ accessed directly through the `gwinc` library interface:
 >>> import gwinc
 >>> import numpy as np
 >>> freq = np.logspace(1, 3, 1000)
->>> Budget = gwinc.load_budget('aLIGO')
->>> traces = Budget(freq).run()
->>> fig = gwinc.plot_noise(freq, traces)
+>>> budget = gwinc.load_budget('aLIGO', freq)
+>>> trace = budget.run()
+>>> fig = gwinc.plot_noise(freq, trace)
 >>> fig.show()
 ```
 
 The `load_budget()` function takes most of the same inputs as the
 command line interface (e.g. IFO names, budget module paths, YAML
-parameter files), and returns the un-instantiated `Budget` class
-defined in the specified budget module (see [budget
-interface](#budget-interface) below).
+parameter files), and returns the `Budget` object defined in the
+specified budget module (see [budget interface](#budget-interface)
+below).  The budget `ifo` `gwinc.Struct` is assigned to the
+`budget.ifo` attribute.
 
 The budget `run()` method is a convenience method that calculates all
 budget noises and the noise total and returns a (possibly) nested
@@ -315,8 +316,7 @@ style for the noise.
 This budget can be loaded with the `gwinc.load_budget()` function, and
 processed with the `Budget.run()` method:
 ```python
-Budget = load_budget('/path/to/MyBudget')
-budget = Budget(freq)
+budget = load_budget('/path/to/MyBudget', freq)
 traces = budget.run()
 ```
 
@@ -333,18 +333,14 @@ suspension Struct is extracted from the `self.ifo` Struct at
 
 If a budget module defined as a package includes an `ifo.yaml`
 [parameter file](#parameter-files) in the package directory, the
-`load_budget()` function will automatically load the YAML data into a
-`gwinc.Struct` and include it as an `Budget.ifo` attribute in the
-returned `Budget` class.  This would provide the `self.ifo` needed in
-the `SuspensionThermal` Noise class above and is therefore a
-convenient way to provide parameter structures in budget packages.
-Otherwise it would need to be created/loaded in some other way and
-passed to the budget at instantiation, e.g.:
+`load_budget()` function will automatically load the YAML data into an
+`ifo` `gwinc.Struct` and assign it to the `budget.ifo` attribute.
+Alternate ifos can be specified at run time:
 ```python
-Budget = load_budget('/path/to/MyBudget')
+budget = load_budget('/path/to/MyBudget', freq)
 ifo = Struct.from_file('/path/to/MyBudget.ifo')
-budget = Budget(freq, ifo=ifo)
-traces = budget.run()
+traces = budget.run(ifo=ifo)
+...
 ```
 
 The IFOs included in `gwinc.ifo` provide examples of the use of the

gwinc/__init__.py

@@ -37,24 +37,24 @@ def _load_module(name_or_path):
     return mod, path
 
 
-def load_budget(name_or_path):
-    """Load GWINC IFO Budget by name or from path.
+def load_budget(name_or_path, freq=None):
+    """Load GWINC Budget
 
-    Named IFOs should correspond to one of the IFOs available in the
-    gwinc package (see gwinc.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).
+    Accepts either the name of a built-in canonical budget (see
+    gwinc.IFOS), the path to a budget package (directory) or module
+    (ending in .py), or the path to an IFO struct definition file (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 the budget is a package directory which includes an 'ifo.yaml'
+    file the ifo Struct will be loaded from that file and assigned to
+    the budget.ifo attribute.  If a struct definition file is provided
+    the base aLIGO budget definition and ifo will be assumed.
 
-    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.
+    Returns a Budget object instantiated with the provided frequency
+    array (if specified), and with any included ifo assigned as an
+    object attribute.  If a frequency array is not provided and the
+    budget class does not define it's own, the frequency array must be
+    provided at budget update() or run() time.
 
     """
     ifo = None
@@ -83,14 +83,11 @@ def load_budget(name_or_path):
 
     logger.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 os.path.exists(ifopath):
         ifo = Struct.from_file(ifopath)
-    Budget.ifo = ifo
-
-    return Budget
+    return Budget(freq=freq, ifo=ifo)
 
 
 def gwinc(freq, ifo, source=None, plot=False, PRfixed=True):
@@ -112,9 +109,9 @@ def gwinc(freq, ifo, 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 = load_budget('aLIGO')
-    traces = Budget(freq, ifo=ifo).run()
-    plot_style = getattr(Budget, 'plot_style', {})
+    budget = load_budget('aLIGO', freq)
+    traces = budget.run()
+    plot_style = getattr(budget, 'plot_style', {})
 
     # construct matgwinc-compatible noises structure
     noises = {}

gwinc/__main__.py

@@ -122,7 +122,7 @@ def main():
         if args.ifo:
             parser.exit(2, "IFO parameter specification not allowed when loading traces from file.\n")
         from .io import load_hdf5
-        Budget = None
+        budget = None
         freq, traces, attrs = load_hdf5(args.IFO)
         ifo = attrs.get('ifo')
         # FIXME: deprecate 'IFO'
@@ -130,16 +130,16 @@ def main():
         plot_style = attrs
 
     else:
-        Budget = load_budget(args.IFO)
-        ifo = Budget.ifo
+        budget = load_budget(args.IFO)
+        ifo = budget.ifo
         if args.freq:
             try:
                 freq = freq_from_spec(args.freq)
             except IndexError:
                 parser.exit(2, "Improper frequency specification: {}\n".format(args.freq))
         else:
-            freq = getattr(Budget, 'freq', freq_from_spec(FREQ))
-        plot_style = getattr(Budget, 'plot_style', {})
+            freq = getattr(budget, 'freq', freq_from_spec(FREQ))
+        plot_style = getattr(budget, 'plot_style', {})
         traces = None
 
     if args.ifo:
@@ -160,9 +160,8 @@ def main():
     if args.diff:
         if not ifo:
             parser.exit(2, "IFO structure not provided.\n")
-        _, difo = load_budget(args.diff)
-        Budget = load_budget(args.diff)
-        diffs = ifo.diff(Budget.ifo)
+        dbudget = load_budget(args.diff)
+        diffs = ifo.diff(dbudget.ifo)
         if diffs:
             w = max([len(d[0]) for d in diffs])
             fmt = '{{:{}}} {{:>20}} {{:>20}}'.format(w)
@@ -236,14 +235,14 @@ def main():
 
     if not traces:
         logger.info("calculating budget...")
-        traces = Budget(freq=freq, ifo=ifo).run()
+        traces = budget.run(freq=freq)
 
     if args.title:
         plot_style['title'] = args.title
     elif 'title' in plot_style:
         pass
-    elif Budget:
-        plot_style['title'] = "GWINC Noise Budget: {}".format(Budget.name)
+    elif budget:
+        plot_style['title'] = "GWINC Noise Budget: {}".format(budget.name)
     else:
         plot_style['title'] = "GWINC Noise Budget: {}".format(args.IFO)
 

gwinc/ifo/__main__.py

@@ -29,7 +29,7 @@ def main():
     budgets = {}
     range_pad = 0
     for ifo in IFOS:
-        budget = load_budget(ifo)(freq)
+        budget = load_budget(ifo, freq)
         name = budget.name
         budgets[name] = budget
         range_pad = max(len(name), range_pad)

gwinc/nb.py

@@ -66,8 +66,6 @@ class BudgetItem:
         if freq is not None:
             assert isinstance(freq, np.ndarray)
             self.freq = freq
-        elif not hasattr(self, 'freq'):
-            raise AttributeError("Frequency array not provided or defined.")
         for key, val in kwargs.items():
             setattr(self, key, val)
 

gwinc/test/__main__.py

@@ -336,8 +336,8 @@ gwinc/test/cache/<SHA1>.  Old caches are automatically pruned.""",
 
         freq, traces_ref, attrs = load_hdf5(path)
 
-        Budget = load_budget(name)
-        traces_cur = Budget(freq).run()
+        budget = load_budget(name, freq)
+        traces_cur = budget.run()
 
         if inspiral_range:
             total_ref = traces_ref['Total'][0]