HDF5_SCHEMATA → HDF5_SCHEMATA.md

 # HDF5 Schema for GWINC noise trace storage
 
-
 This file describes a schemata for HDF5 storage of noise trace data
 and plot styling GWINC noise budgets.  Python functions for writing
 budget data to and reading budget data from this format are included
@@ -19,53 +18,22 @@ pairs.
 Bindings are available for most platforms including Python [1] and
 Matlab.
 
-[0] https://en.wikipedia.org/wiki/Hierarchical_Data_Format
-[1] http://www.h5py.org/
-
-
-## version history
-
-v1
-- first versioned schema release
+ * [0] https://en.wikipedia.org/wiki/Hierarchical_Data_Format
+ * [1] http://www.h5py.org/
 
 
 ## schema
 
 The following describes the noise budget schema.  Specific strings are
 enclosed in single quotes (''), and variables are described in
-brackets (<>).  Group objects are indicated by a closing '/'
-separator, data set are indicated by a closing ':' followed by a
-specification of their length and type (e.g. "(N),float"), and
-attributes are specified in the .attrs[] dictionary format.  Optional
-elements are enclosed in parentheses.
-
-A single trace is a length N array (with optional plot style specified
-in attributes:
-```
-/<trace>: (N),float
-(/<trace>.attrs['label'] = <label>)
-(/<trace>.attrs['color] = <color>)
-...
-```
-
-A budget item, i.e. a collection of noises is structured as follows:
-```
-/<budget>/
-    /'Total': (N),float
-    /<trace_0>: (N),float
-    (/<trace_1>: (N),float)
-```
-<!-- ``` -->
-<!--     /'noises'/ -->
-<!--         /*<noise_0>: (N),float -->
-<!--         ... -->
-<!--     /\*'references'/ -->
-<!--         /*<ref_0>: (N),float -->
-<!--         ... -->
-<!-- ``` -->
+brackets (<>).  Group objects are indicated by a closing '/', data
+sets are indicated by a closing ':' followed by a specification of
+their length and type (e.g. "(N),float"), and attributes are specified
+in the .attrs[] dictionary format.  Optional elements are enclosed in
+parentheses.
 
 
-## Top-level Budget
+## top-level attributes
 
 The following two root attributes expected: a string describing the schema,
 and an int schema version number:
@@ -74,34 +42,81 @@ and an int schema version number:
 /.attrs['SCHEMA_VERSION'] = 1
 ```
 
-Top-level attributes can be used for storing any relevant metadata,
-such as "ifo" parameters in YAML format or overall plot styling.
-Typically at least the following root attributes are defined:
+The following root root attributes are defined:
 ```
 /.attrs['title'] = <experiment description string (e.g. 'H1 Strain Budget')>
 /.attrs['date'] = <ISO-formatted string (e.g. '2015-10-24T20:30:00.000000Z')>
+/.attrs['ifo'] = <IFO Struct object, YAML serialized>
+```
+
+The remaining root level attributes usually pertain to the plot style.
+
+The budget frequency array is defined in a top level 'Freq' dataset:
+
+```
+/'Freq': (N),float
+```
+
+## version history
+
+### v1
+
+A single trace is a length N array (with optional plot style specified
+in attributes:
+```
+/<trace>: (N),float
+/<trace>.attrs['label'] = <label>
+/<trace>.attrs['color] = <color>
+...
 ```
 
-The budget frequency array is defined in a 'Freq' dataset:
+A budget item, i.e. a collection of noises is structured as follows:
 ```
-/'Freq': (N), float
+/<budget>/
+/<budget>/Total': (N),float
+/<budget>/<trace_0>: (N),float
+/<budget>/...
 ```
 
+
 The budget traces are defined a traces group.  The overall structure
 looks something like this:
 ```
-/.attrs['SCHEMA'] = 'GWINC Noise Budget'
-/.attrs['SCHEMA_VERSION'] = 1
-/.attrs['title'] = <experiment description string (e.g. 'H1 Strain Budget')>
-/.attrs['date'] = <ISO-formatted string (e.g. '2015-10-24T20:30:00.000000Z')>
-/'Freq': (N), float
 /traces/
-    /'Total': (N),float
-    /<noise_0>: (N),float
-    /<noise_1>: (N),float
-    /<noise_2>/
-        /'Total': (N),float
-        /<noise_3>: (N),float
-        /<noise_4>: (N),float
-    ...
+/traces/'Total': (N),float
+/traces/<noise_0>: (N),float
+/traces/<noise_1>: (N),float
+/traces/<noise_2>/
+/traces/<noise_2>/'Total': (N),float
+/traces/<noise_2>/<noise_3>: (N),float
+/traces/<noise_2>/...
+/traces/...
+```
+
+
+### v2
+
+Each trace is given the following structure:
+```
+/<trace>/
+/<trace>.attrs['style'] = <YAML trace style dict>
+/<trace>/'PSD': (N),float
+/<trace>/budget/
+/<trace>/budget/<subtrace_0>/
+/<trace>/budget/<subtrace_1>/
+/<trace>/budget/...
+```
+
+The overall structure is:
+```
+/<budget>/
+/<budget>/.attrs['plot_style'] = <YAML plot style dict>
+/<budget>/.attrs['style'] = <YAML "Total" trace style dict>
+/<budget>/.attrs[*] = <arbitrary data>
+/<budget>/'Freq': (N),float
+/<budget>/'PSD': (N),float
+/<budget>/budget/
+/<budget>/budget/<trace_0>/...
+/<budget>/budget/<trace_1>/...
+/<budget>/budget/...
 ```

README.md

@@ -39,8 +39,8 @@ data.
 
 The [`inspiral_range`](https://git.ligo.org/gwinc/inspiral-range)
 package can be used to calculate various common "inspiral range"
-figures of merit for gravitational wave detector budgets.  See
-[figures of merit](#figures-of-merit) section below.
+figures of merit for gravitational wave detector budgets.  See the
+[inspiral range](#inspiral-range) section below.
 
 
 ## usage
@@ -48,12 +48,12 @@ figures of merit for gravitational wave detector budgets.  See
 ### command line interface
 
 `pygwinc` provides a command line interface that can be used to
-calculate and plot noise budgets for generic noise budgets or the
-various canonical IFOs described above, save/plot hdf5 trace data, and
-dump budget IFO parameters:
+calculate and plot the various canonical IFO noise budgets described
+above, as well as custom noise budgets (see below):
 ```shell
 $ python3 -m gwinc aLIGO
 ```
+Budget plots save/plot hdf5 trace data, and dump budget IFO parameters:
 
 You can play with IFO parameters and see the effects on the budget by
 dumping the pre-defined parameters to a [YAML-formatted parameter
@@ -84,18 +84,19 @@ $ python3 -m gwinc path/to/mybudget
 ```
 
 The command line interface also includes an "interactive" mode which
-provides an [IPython](https://ipython.org/) shell for interacting with a processed budget:
+provides an [IPython](https://ipython.org/) shell for interacting with
+a processed budget:
 ```shell
 $ python3 -m gwinc -i Aplus
 GWINC interactive shell
 
-The 'ifo' Struct and 'traces' data are available for inspection.
+The 'ifo' Struct and 'trace' data are available for inspection.
 Use the 'whos' command to view the workspace.
 
 You may interact with the plot using the 'plt' functions, e.g.:
 
-In [.]: plt.title("foo")
-In [.]: plt.savefig("foo.pdf")
+In [.]: plt.title("My Special Budget")
+In [.]: plt.savefig("mybudget.pdf")
 
 In [1]: 
 ```
@@ -116,30 +117,30 @@ accessed directly through the `gwinc` library interface:
 >>> freq = np.logspace(1, 3, 1000)
 >>> budget = gwinc.load_budget('aLIGO', freq)
 >>> trace = budget.run()
->>> fig = gwinc.plot_noise(freq, trace)
+>>> fig = gwinc.plot_budget(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 `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.
+parameter files), and returns the instantiated `Budget` object defined
+in the specified budget module (see [budget
+interface](#budget-interface) below).  The budget `ifo` `gwinc.Struct`
+is available in 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
-dictionary of a noise data, in the form of a `(data, style)` tuple
-where 'data' is the PSD data and 'style' is a plot style dictionary
-for the trace.  The dictionary will be nested if the budget includes
-any sub-budgets.
+The budget `run()` method calculates all budget noises and the noise
+total and returns a `BudgetTrace` object with `freq`, `psd`, and `asd`
+properties.  The budget sub-traces are available through a dictionary
+(`trace['QuantumVacuum']`) interface and via attributes
+(`trace.QuantumVacumm`).
 
 
 ## noise functions
 
-`pygwinc` noise functions are available in the `gwinc.noise` package.
-This package includes multiple sub-modules for the different types of
-noises, e.g. `suspensionthermal`, `coatingthermal`, `quantum`, etc.)
+The `pygwinc` analytical noise functions are available in the
+`gwinc.noise` package.  This package includes multiple sub-modules for
+the different types of noises, e.g. `suspensionthermal`,
+`coatingthermal`, `quantum`, etc.)
 
 The various noise functions need many different parameters to
 calculate their noise outputs.  Many parameters are expected to be in
@@ -162,7 +163,7 @@ def coating_brownian(f, materials, wavelength, wBeam, dOpt):
 ```
 
 
-### `gwinc.Struct` objects
+## `gwinc.Struct` objects
 
 `pygwinc` provides a `Struct` class that can hold parameters in
 attributes and additionally acts like a dictionary, for passing to the
@@ -233,6 +234,11 @@ are:
 * `update(**kwargs)`: update data/attributes
 * `calc()`: return final data array
 
+Generally these methods are not called directly.  Instead, the `Noise`
+and `Budget` classes include a `run` method that smartly executes then
+all in sequence.
+
+
 See the built-in documentation for more info (e.g. `pydoc3
 gwinc.nb.BudgetItem`)
 
@@ -314,10 +320,10 @@ The `style` attributes of the various `Noise` classes define plot
 style for the noise.
 
 This budget can be loaded with the `gwinc.load_budget()` function, and
-processed with the `Budget.run()` method:
+processed as usual with the `Budget.run()` method:
 ```python
 budget = load_budget('/path/to/MyBudget', freq)
-traces = budget.run()
+trace = budget.run()
 ```
 
 Other than the necessary `freq` initialization argument that defines
@@ -339,8 +345,7 @@ Alternate ifos can be specified at run time:
 ```python
 budget = load_budget('/path/to/MyBudget', freq)
 ifo = Struct.from_file('/path/to/MyBudget.ifo')
-traces = budget.run(ifo=ifo)
-...
+trace = budget.run(ifo=ifo)
 ```
 
 The IFOs included in `gwinc.ifo` provide examples of the use of the
@@ -354,10 +359,9 @@ interface.  The most straightforward way is to run the full budget,
 and extract the noise data the output traces dictionary:
 
 ```python
-Budget = load_budget('/path/to/MyBudget')
-budget = Budget(freq)
-traces = budget.calc_traces()
-data, plot_style = traces['QuantumVacuum']
+budget = load_budget('/path/to/MyBudget', freq)
+trace = budget.run()
+quantum_trace = trace['QuantumVacuum']
 ```
 
 You can also calculate the final calibrated output noise for just a
@@ -375,7 +379,7 @@ data = budget['QuantumVacuum'].calc()
 ```
 
 
-# figures of merit
+# inspiral range
 
 The [`inspiral_range`](https://git.ligo.org/gwinc/inspiral-range)
 package can be used to calculate various common "inspiral range"
@@ -388,18 +392,15 @@ import inspiral_range
 import numpy as np
 
 freq = np.logspace(1, 3, 1000)
-Budget = gwinc.load_budget('Aplus')
-traces = Budget(freq).run()
+budget = gwinc.load_budget('Aplus', freq)
+trace = budget.run()
 
 range = inspiral_range.range(
-    freq, traces['Total'][0],
+    freq, trace.psd,
     m1=30, m2=30,
 )
 ```
 
-Note you need to extract the zeroth element of the `traces['Total']`
-tuple, which is the actual PSD data.
-
 See the [`inspiral_range`](https://git.ligo.org/gwinc/inspiral-range)
 package for more details.
 

gwinc/__init__.py

@@ -7,6 +7,7 @@ import numpy as np
 
 from .ifo import IFOS
 from .struct import Struct
+from .plot import plot_budget
 from .plot import plot_noise
 
 
@@ -115,9 +116,10 @@ def gwinc(freq, ifo, source=None, plot=False, PRfixed=True):
 
     # construct matgwinc-compatible noises structure
     noises = {}
-    for name, (data, style) in traces.items():
-        noises[style.get('label', name)] = data
-    noises['Freq'] = freq
+    for name, trace in traces.items():
+        noises[name] = trace.psd
+    noises['Total'] = traces.psd
+    noises['Freq'] = traces.freq
 
     pbs = ifo.gwinc.pbs
     parm = ifo.gwinc.parm
@@ -164,6 +166,6 @@ def gwinc(freq, ifo, source=None, plot=False, PRfixed=True):
             logger.info('BBH Inspiral Range:     ' + str(score.effr0bh) + ' Mpc/ z = ' + str(score.zHorizonBH))
             logger.info('Stochastic Omega: %4.1g Universes' % score.Omega)
 
-        plot_noise(ifo, traces, **plot_style)
+        plot_budget(traces, **plot_style)
 
     return score, noises, ifo

gwinc/__main__.py

@@ -5,7 +5,7 @@ import logging
 import argparse
 import numpy as np
 
-from . import IFOS, load_budget, plot_noise, logger
+from . import IFOS, load_budget, plot_budget, logger
 
 logger.setLevel(os.getenv('LOG_LEVEL', 'WARNING').upper())
 formatter = logging.Formatter('%(message)s')
@@ -123,11 +123,10 @@ def main():
             parser.exit(2, "IFO parameter specification not allowed when loading traces from file.\n")
         from .io import load_hdf5
         budget = None
-        freq, traces, attrs = load_hdf5(args.IFO)
-        ifo = attrs.get('ifo')
-        # FIXME: deprecate 'IFO'
-        ifo = attrs.get('IFO', ifo)
-        plot_style = attrs
+        trace = load_hdf5(args.IFO)
+        freq = trace.freq
+        ifo = trace.ifo
+        plot_style = trace.plot_style
 
     else:
         budget = load_budget(args.IFO)
@@ -140,7 +139,7 @@ def main():
         else:
             freq = getattr(budget, 'freq', freq_from_spec(FREQ))
         plot_style = getattr(budget, 'plot_style', {})
-        traces = None
+        trace = None
 
     if args.ifo:
         for paramval in args.ifo:
@@ -233,9 +232,9 @@ def main():
     ##########
     # main calculations
 
-    if not traces:
+    if not trace:
         logger.info("calculating budget...")
-        traces = budget.run(freq=freq)
+        trace = budget.run(freq=freq)
 
     if args.title:
         plot_style['title'] = args.title
@@ -250,15 +249,16 @@ def main():
         logger.info("calculating inspiral {}...".format(range_func))
         H = inspiral_range.CBCWaveform(freq, **range_params)
         logger.debug("waveform params: {}".format(H.params))
-        fom = eval('inspiral_range.{}'.format(range_func))(freq, traces['Total'][0], H=H)
+        fom = eval('inspiral_range.{}'.format(range_func))(freq, trace.psd, H=H)
         logger.info("{}({}) = {:.2f} Mpc".format(range_func, range_func_args, fom))
-        fom_title = 'inspiral {func} {m1}/{m2} $\mathrm{{M}}_\odot$: {fom:.0f} Mpc'.format(
+        subtitle = 'inspiral {func} {m1}/{m2} $\mathrm{{M}}_\odot$: {fom:.0f} Mpc'.format(
             func=range_func,
             m1=H.params['m1'],
             m2=H.params['m2'],
             fom=fom,
             )
-        plot_style['title'] += '\n{}'.format(fom_title)
+    else:
+        subtitle = None
 
     ##########
     # interactive
@@ -267,14 +267,14 @@ def main():
     if args.interactive:
         banner = """GWINC interactive shell
 
-The 'ifo' Struct and 'traces' data objects are available for
+The 'ifo' Struct and 'budget' trace data objects are available for
 inspection.  Use the 'whos' command to view the workspace.
 """
         if not args.plot:
             banner += """
-You may plot the budget using the 'plot_noise()' function:
+You may plot the budget using the 'plot_budget()' function:
 
-In [.]: plot_noise(freq, traces, **plot_style)
+In [.]: plot_budget(budget, **plot_style)
 """
         banner += """
 You may interact with the plot using the 'plt' functions, e.g.:
@@ -286,34 +286,31 @@ In [.]: plt.savefig("foo.pdf")
         ipshell = InteractiveShellEmbed(
             banner1=banner,
             user_ns={
-                'freq': freq,
-                'traces': traces,
+                'budget': trace,
                 'ifo': ifo,
                 'plot_style': plot_style,
-                'plot_noise': plot_noise,
+                'plot_budget': plot_budget,
             },
         )
         ipshell.enable_pylab()
         if args.plot:
-            ipshell.ex("fig = plot_noise(freq, traces, **plot_style)")
+            ipshell.ex("fig = plot_budget(budget, **plot_style)")
             ipshell.ex("plt.title(plot_style['title'])")
         ipshell()
 
     ##########
     # output
 
-    # save noise traces to HDF5 file
+    # save noise trace to HDF5 file
     if out_data_files:
         from .io import save_hdf5
-        attrs = dict(plot_style)
         for path in out_data_files:
-            logger.info("saving budget traces: {}".format(path))
+            logger.info("saving budget trace: {}".format(path))
             save_hdf5(
                 path=path,
-                freq=freq,
-                traces=traces,
+                trace=trace,
                 ifo=ifo,
-                **attrs,
+                plot_style=plot_style,
             )
 
     # standard plotting
@@ -321,10 +318,10 @@ In [.]: plt.savefig("foo.pdf")
         logger.debug("plotting noises...")
         fig = plt.figure()
         ax = fig.add_subplot(1, 1, 1)
-        plot_noise(
-            freq,
-            traces,
+        plot_budget(
+            trace,
             ax=ax,
+            subtitle=subtitle,
             **plot_style
         )
         fig.tight_layout()

gwinc/ifo/CE2/ifo.yaml

@@ -248,7 +248,6 @@ Optics:
   # zz = loadmat('CryogenicLIGO/Sensitivity/GWINC/optRuns/' + qopt_mat)
   ITM:
     SubstrateAbsorption: 1e-3     # 1/m; 10 ppm/cm for MCZ Si
-    BeamRadius: 0.141             # m; 1/e^2 power radius w1
     CoatingAbsorption: 0.5e-6     # absorption of ITM
     Transmittance: 0.014          # zz['x'][0][3]
     #CoatingThicknessLown: 0.308
@@ -269,7 +268,6 @@ Optics:
       - 0.3867382
       - 0.08814237
   ETM:
-    BeamRadius: 0.141                  # m; 1/e^2 power radius w2
     Transmittance: 5e-6                # Transmittance of ETM
     #CoatingThicknessLown: 0.27
     #CoatingThicknessCap: 0.5

gwinc/ifo/Voyager/ifo.yaml

@@ -273,7 +273,6 @@ Optics:
   # zz = loadmat('CryogenicLIGO/Sensitivity/GWINC/optRuns/' + qopt_mat)
   ITM:
     SubstrateAbsorption: 1e-3   # 1/m; 10 ppm/cm for MCZ Si
-    BeamRadius: 0.0585          # m; 1/e^2 power radius w1
     CoatingAbsorption: 1e-6     # absorption of ITM
     Transmittance: 2e-3         # zz['x'][0][3]
     #CoatingThicknessLown: 0.308
@@ -294,7 +293,6 @@ Optics:
       - 0.386738199718736
       - 0.088142365969070
   ETM:
-    BeamRadius: 0.0835                 # m; 1/e^2 power radius w2
     Transmittance: 5e-6                # Transmittance of ETM
     #CoatingThicknessLown: 0.27
     #CoatingThicknessCap: 0.5

gwinc/ifo/__main__.py

@@ -35,11 +35,11 @@ def main():
         range_pad = max(len(name), range_pad)
 
     for name, budget in budgets.items():
-        data = budget.calc()
+        trace = budget.run()
         try:
             import inspiral_range
             label_range = ' {:>6.0f} Mpc'.format(
-                inspiral_range.range(freq, data),
+                inspiral_range.range(freq, trace.psd),
             )
         except ModuleNotFoundError:
             label_range = ''
@@ -49,7 +49,7 @@ def main():
             pad=range_pad,
             range=label_range,
         )
-        ax.loglog(freq, np.sqrt(data), label=label, lw=2)
+        ax.loglog(freq, trace.asd, label=label, lw=2)
 
     ax.grid(
         True,

gwinc/ifo/aLIGO/ifo.yaml

@@ -188,7 +188,7 @@ Suspension:
 
 ## Optic Material -------------------------------------------------------
 Materials:
-  MassRadius: 0.17                # m; 
+  MassRadius: 0.17                # m;
   MassThickness: 0.200            # m; Peter F 8/11/2005
 
   ## Dielectric coating material parameters----------------------------------
@@ -251,13 +251,11 @@ Optics:
     dc: 1.5707963                 # pi/2 # demod/detection/homodyne phase
 
   ITM:
-    # BeamRadius: 0.055             # m, 1/e^2 power radius, now in precompIFO
     Transmittance: 0.014
     CoatingThicknessLown: 0.308
     CoatingThicknessCap: 0.5
     CoatingAbsorption: 0.5e-6
   ETM:
-    # BeamRadius: 0.062             # m, 1/e^2 power radius, now in precompIFO
     Transmittance: 5e-6
     CoatingThicknessLown: 0.27
     CoatingThicknessCap: 0.5

gwinc/ifo/noises.py

@@ -296,10 +296,21 @@ class CoatingBrownian(nb.Noise):
         w0, wBeam_ITM, wBeam_ETM = arm_cavity(self.ifo)
         dOpt_ITM = coating_thickness(self.ifo, 'ITM')
         dOpt_ETM = coating_thickness(self.ifo, 'ETM')
+        mTi_ITM = None
+        mTi_ETM = None
+        Ic = None
+        if 'IncCoatBrAmpNoise' in materials.Coating:
+            if materials.Coating.IncCoatBrAmpNoise.lower() == 'yes':
+                mTi_ITM = self.ifo.Optics.ITM.Transmittance
+                mTi_ETM = self.ifo.Optics.ETM.Transmittance
+                Ic = ifo_power(self.ifo)[1]   # Circulating Arm Power
+                precomp_mirror(self.freq, self.ifo)
         nITM = noise.coatingthermal.coating_brownian(
-            self.freq, materials, wavelength, wBeam_ITM, dOpt_ITM)
+            self.freq, materials, wavelength, wBeam_ITM,
+            dOpt_ITM, Ic, mTi_ITM)
         nETM = noise.coatingthermal.coating_brownian(
-            self.freq, materials, wavelength, wBeam_ETM, dOpt_ETM)
+            self.freq, materials, wavelength, wBeam_ETM,
+            dOpt_ETM, Ic, mTi_ETM)
         return (nITM + nETM) * 2
 
 

gwinc/io.py

@@ -4,25 +4,23 @@ import datetime
 
 from . import logger
 from . import Struct
+from .trace import BudgetTrace
 
 
 SCHEMA = 'GWINC noise budget'
-SCHEMA_VERSION = 1
 
 
-def _write_trace_recursive(grp, traces):
-    for name, trace in traces.items():
-        if isinstance(trace, dict):
-            tgrp = grp.create_group(name)
-            _write_trace_recursive(tgrp, trace)
-        else:
-            data, style = trace
-            dset = grp.create_dataset(name, data=data)
-            for key, val in style.items():
-                dset.attrs[key] = val
+def _write_trace_recursive(f, trace):
+    f.attrs['style'] = yaml.dump(trace.style)
+    f.create_dataset('PSD', data=trace.psd)
+    if trace.budget:
+        bgrp = f.create_group('budget')
+    for t in trace:
+        tgrp = bgrp.create_group(t.name)
+        _write_trace_recursive(tgrp, t)
 
 
-def save_hdf5(path, freq, traces, **kwargs):
+def save_hdf5(path, trace, ifo=None, plot_style=None, **kwargs):
     """Save GWINC budget data to an HDF5 file.
 
     The `freq` argument should be the frequency array, and `traces`
@@ -36,46 +34,107 @@ def save_hdf5(path, freq, traces, **kwargs):
     """
     with h5py.File(path, 'w') as f:
         f.attrs['SCHEMA'] = SCHEMA
-        f.attrs['SCHEMA_VERSION'] = SCHEMA_VERSION
+        f.attrs['SCHEMA_VERSION'] = 2
         # FIXME: add budget code hash or something
         f.attrs['date'] = datetime.datetime.now().isoformat()
+        if ifo:
+            f.attrs['ifo'] = ifo.to_yaml()
+        if plot_style:
+            f.attrs['plot_style'] = yaml.dump(plot_style)
         for key, val in kwargs.items():
-            if key == 'ifo':
-                f.attrs['ifo'] = val.to_yaml()
+            f.attrs[key] = val
+        f.create_dataset('Freq', data=trace.freq)
+        tgrp = f.create_group(trace.name)
+        _write_trace_recursive(tgrp, trace)
+
+
+def _load_hdf5_v1(f):
+    def read_recursive(element):
+        budget = []
+        for name, item in element.items():
+            if name == 'Total':
+                total = item[:]
+                continue
+            if isinstance(item, h5py.Group):
+                trace = read_recursive(item)
+                trace.name = name
             else:
-                f.attrs[key] = val
-        f.create_dataset('Freq', data=freq)
-        tgrp = f.create_group('traces')
-        _write_trace_recursive(tgrp, traces)
+                trace = BudgetTrace(
+                    name=name,
+                    style=dict(item.attrs),
+                    psd=item[:],
+                    budget=[],
+                )
+            budget.append(trace)
+        return BudgetTrace(
+            psd=total,
+            budget=budget,
+        )
+    trace = read_recursive(f['/traces'])
+    trace.name = 'Total'
+    trace._freq = f['Freq'][:]
+    attrs = dict(f.attrs)
+    if 'ifo' in attrs:
+        try:
+            ifo = Struct.from_yaml(attrs['ifo'])
+            del attrs['ifo']
+        except yaml.constructor.ConstructorError:
+            logger.warning("HDF5 load warning: Could not de-serialize 'ifo' YAML attribute.")
+            ifo = None
+    trace.style = attrs
+    trace.ifo = ifo
+    return trace
 
 
-def _read_trace_recursive(element):
-    trace = {}
-    for name, item in element.items():
-        if isinstance(item, h5py.Group):
-            trace[name] = _read_trace_recursive(item)
+def _load_hdf5_v2(f):
+    def read_recursive(element):
+        psd = element['PSD'][:]
+        style = yaml.safe_load(element.attrs['style'])
+        budget = []
+        if 'budget' in element:
+            for name, item in element['budget'].items():
+                trace = read_recursive(item)
+                trace.name = name
+                budget.append(trace)
+        return BudgetTrace(
+            psd=psd,
+            budget=budget,
+            style=style,
+        )
+    for name, item in f.items():
+        if name == 'Freq':
+            freq = item[:]
         else:
-            trace[name] = item[:], dict(item.attrs.items())
+            trace = read_recursive(item)
+            trace.name = name
+    trace._freq = freq
+    trace.ifo = None
+    attrs = dict(f.attrs)
+    ifo = attrs.get('ifo')
+    if ifo:
+        try:
+            trace.ifo = Struct.from_yaml(ifo)
+            del attrs['ifo']
+        except yaml.constructor.ConstructorError:
+            logger.warning("HDF5 load warning: Could not de-serialize 'ifo' YAML attribute.")
+    trace.plot_style = yaml.safe_load(attrs['plot_style'])
     return trace
 
 
 def load_hdf5(path):
     """Load GWINC budget data from an HDF5 file.
 
-    Returns (freq, traces, attrs).  An 'ifo' attr will be
-    de-serialized from YAML into a Struct object.
+    Returns BudgetTrace.  An 'ifo' attr will be de-serialized from
+    YAML into a Struct object and assigned as an attribute to the
+    BudgetTrace.
 
     See HDF5_SCHEMA.
 
     """
+    loaders = {
+        1: _load_hdf5_v1,
+        2: _load_hdf5_v2,
+    }
     with h5py.File(path, 'r') as f:
-        # FIXME: check SCHEMA name/version
-        freq = f['Freq'][:]
-        traces = _read_trace_recursive(f['/traces'])
-        attrs = dict(f.attrs)
-        if 'ifo' in attrs:
-            try:
-                attrs['ifo'] = Struct.from_yaml(attrs['ifo'])
-            except yaml.constructor.ConstructorError:
-                logger.warning("HDF5 load warning: Could not de-serialize 'ifo' YAML attribute.")
-        return freq, traces, attrs
+        version = f.attrs.get('SCHEMA_VERSION', 1)
+        return loaders[version](f)

gwinc/nb.py

-import operator
 import itertools
 import collections
 import numpy as np
 import scipy.interpolate
-from functools import reduce
 
 from . import logger
+from .trace import BudgetTrace
 
 
 def quadsum(data):
@@ -20,7 +19,6 @@ def quadsum(data):
     return np.nansum(data, 0)
 
 
-
 class BudgetItem:
     """GWINC BudgetItem class
 
@@ -37,6 +35,11 @@ class BudgetItem:
         By default any keyword arguments provided are written directly
         as attribute variables (as with __init__).
 
+        When overloading this method it is recommended to execute the
+        method from the base class with e.g.:
+
+          super().update(**kwargs)
+
         """
         for key, val in kwargs.items():
             setattr(self, key, val)
@@ -132,27 +135,31 @@ class Noise(BudgetItem):
     style = {}
     """Trace plot style dictionary"""
 
-    def calc_trace(self, calibrations=None, calc=True):
-        """Returns noise (PSD, style) tuple.
+    def _make_trace(self, psd=None, budget=None):
+        return BudgetTrace(
+            name=self.name,
+            style=self.style,
+            freq=self.freq,
+            psd=psd,
+            budget=budget,
+        )
+
+    def calc_trace(self, calibration=1, calc=True):
+        """Calculate noise and return BudgetTrace object
 
-        If `calibrations` is not None it is assumed to be a list of
-        len(self.freq) arrays that will be multiplied to the output
-        PSD.
+        `calibration` should either be a scalar or a len(self.freq)
+        array that will be multiplied to the output PSD of the budget
+        and all sub noises.
 
-        If calc=False, the noise will not be calculated and the PSD
-        will be None.  This is useful for just getting the style.
+        If `calc` is False, the noise will not be calculated and the
+        trace PSD will be None.  This is useful for just getting the
+        trace style info.
 
         """
+        total = None
         if calc:
-            data = self.calc()
-            if calibrations:
-                data *= reduce(
-                    operator.mul,
-                    calibrations,
-                )
-        else:
-            data = None
-        return data, self.style
+            total = self.calc() * calibration
+        return self._make_trace(psd=total)
 
     def run(self, **kwargs):
         """Convenience method to load, update, and return calc traces.
@@ -361,95 +368,67 @@ class Budget(Noise):
 
     def update(self, **kwargs):
         """Update all noise and cal objects with supplied kwargs."""
+        super().update(**kwargs)
         for name, item in itertools.chain(
                 self._cal_objs.items(),
                 self._noise_objs.items()):
             logger.debug("update {}".format(item))
             item.update(**kwargs)
 
-    def calc_noise(self, name):
-        """Return calibrated individual noise term.
+    update.__doc__ = Noise.update.__doc__
 
-        The noise PSD and calibration transfer functions are
-        calculated, and the calibrated noise array is returned.
+    def calc_noise(self, name, calibration=1, calc=True, _cals=None):
+        """Return calibrated individual noise BudgetTrace.
 
-        """
-        noise = self._noise_objs[name]
-        cals = [self._cal_objs[cal].calc() for cal in self._noise_cals[name]]
-        data = noise.calc_trace(cals)
-        if isinstance(data, dict):
-            return data['Total'][0]
-        else:
-            return data[0]
-
-    def calc(self):
-        """Calculate sum of all noises.
+        The noise and calibration transfer functions are calculated,
+        and the calibrated noise BudgetTrace is returned.
+        `calibration` is an overall calculated calibration to apply to
+        the noise.
 
         """
-        data = [self.calc_noise(name) for name in self._noise_objs.keys() if name in self._budget_noises]
-        return quadsum(data)
-
-    def calc_trace(self, calibrations=None, calc=True):
-        """Returns a dictionary of noises traces, keyed by noise names.
+        for cal in self._noise_cals[name]:
+            if _cals:
+                calibration *= _cals[cal]
+            else:
+                obj = self._cal_objs[cal]
+                logger.debug("calc {}".format(obj))
+                calibration *= obj.calc()
+        noise = self._noise_objs[name]
+        logger.debug("calc {}".format(noise))
+        return noise.calc_trace(
+            calibration=calibration,
+            calc=calc,
+        )
 
-        Values are (data, style) trace tuples (see Noise.calc_trace).
-        The key of the budget sum total is 'Total'.  The values of sub
-        budgets are themselves dictionaries returned from
-        calc_trace() of the sub budget.
+    def calc_trace(self, calibration=1, calc=True):
+        """Calculate all budget noises and return BudgetTrace object
 
-        If `calibrations` is not None it is assumed to be a list of
-        len(self.freq) array that will be multiplied to the output PSD
-        of the budget and all sub noises.
+        `calibration` should either be a scalar or a len(self.freq)
+        array that will be multiplied to the output PSD of the budget
+        and all sub noises.
 
-        If calc=False, the noise will not be calculated and the PSD
-        will be None.  This is useful for just getting style the
-        style.
+        If `calc` is False, the noise will not be calculated and the
+        trace PSD will be None.  This is useful for just getting the
+        trace style info.
 
         """
-        # start by creating an empty OrderedDict used for outputing trace data
-        # or style info with the following order:
-        #   references
-        #   total
-        #   constituents
-        d = collections.OrderedDict()
-        # allocate references
-        for name, noise in self._noise_objs.items():
-            if name in self._budget_noises:
-                continue
-            d[name] = noise.calc_trace(calc=False)
-        # allocate total
-        if self._budget_noises:
-            d['Total'] = None, self.style
-        # allocate constituent
-        for name, noise in self._noise_objs.items():
-            if name not in self._budget_noises:
-                continue
-            d[name] = noise.calc_trace(calc=False)
-        # if we're not calc'ing, just return the dict now
-        if not calc:
-            return d
-
-        # calc all calibrations
-        c = {}
-        for name, cal in self._cal_objs.items():
-            logger.debug("calc {}".format(name))
-            c[name] = cal.calc()
-        # calc all noises
-        for name, noise in self._noise_objs.items():
-            cals = [c[cal] for cal in self._noise_cals[name]]
-            if calibrations:
-                cals += calibrations
-            logger.debug("calc {}".format(name))
-            d[name] = noise.calc_trace(
-                calibrations=cals,
+        _cals = {}
+        if calc:
+            for name, cal in self._cal_objs.items():
+                logger.debug("calc {}".format(cal))
+                _cals[name] = cal.calc()
+
+        budget = []
+        for name in self._noise_objs:
+            trace = self.calc_noise(
+                name,
+                calibration=calibration,
+                calc=calc,
+                _cals=_cals,
             )
-        # calc budget total
-        constituent_data = []
-        for name in self._budget_noises:
-            if isinstance(d[name], dict):
-                data = d[name]['Total'][0]
-            else:
-                data = d[name][0]
-            constituent_data.append(data)
-        d['Total'] = quadsum(constituent_data), self.style
-        return d
+            budget.append(trace)
+
+        total = quadsum([trace.psd for trace in budget])
+        return self._make_trace(
+            psd=total, budget=budget
+        )

gwinc/plot.py

-from numpy import sqrt
-from collections.abc import Mapping
-
-
-def plot_noise(
-        freq,
-        traces,
+def plot_budget(
+        budget,
         ax=None,
         **kwargs
 ):
-    """Plot a GWINC noise budget from calculated noises
+    """Plot a GWINC BudgetTrace noise budget from calculated noises
 
     If an axis handle is provided it will be used for the plot.
 
@@ -22,36 +17,29 @@ def plot_noise(
     else:
         fig = ax.figure
 
-    ylim = kwargs.get('ylim')
-
-    for name, trace in traces.items():
-        if isinstance(trace, Mapping):
-            trace = trace['Total']
+    total = budget.asd
+    ylim = [min(total)/10, max(total)]
+    style = dict(
+        color='#000000',
+        alpha=0.6,
+        lw=4,
+    )
+    style.update(getattr(budget, 'style', {}))
+    if 'label' in style:
+        style['label'] = 'Total ' + style['label']
+    else:
+        style['label'] = 'Total'
+    ax.loglog(budget.freq, total, **style)
 
-        try:
-            data = trace[0]
-            style = dict(**trace[1])
-        except TypeError:
-            data = trace
-            style = {}
-        # assuming all data is PSD
-        data = sqrt(data)
-        if name == 'Total' and not style:
-            style = dict(
-                color='#000000',
-                alpha=0.6,
-                lw=4,
-            )
-            if ylim is None:
-                ylim = [min(data)/10, max(data)]
+    for name, trace in budget.items():
+        style = trace.style
         if 'label' not in style:
-            style['label'] = name
+            style['label'] = budget.name
         if 'linewidth' in style:
             style['lw'] = style['linewidth']
         elif 'lw' not in style:
             style['lw'] = 3
-
-        ax.loglog(freq, data, **style)
+        ax.loglog(budget.freq, trace.asd, **style)
 
     ax.grid(
         True,
@@ -67,14 +55,19 @@ def plot_noise(
     )
 
     ax.autoscale(enable=True, axis='y', tight=True)
-    if ylim:
-        ax.set_ylim(ylim)
-    ax.set_xlim(freq[0], freq[-1])
-
+    ax.set_ylim(kwargs.get('ylim', ylim))
+    ax.set_xlim(budget.freq[0], budget.freq[-1])
     ax.set_xlabel('Frequency [Hz]')
     if 'ylabel' in kwargs:
         ax.set_ylabel(kwargs['ylabel'])
     if 'title' in kwargs:
-        ax.set_title(kwargs['title'])
+        title = kwargs['title']
+        if 'subtitle' in kwargs:
+            title += '\n' + kwargs['subtitle']
+        ax.set_title(title)
 
     return fig
+
+
+# FIXME: deprecate
+plot_noise = plot_budget

gwinc/test/__main__.py

@@ -6,10 +6,8 @@ import logging
 import tempfile
 import argparse
 import subprocess
-import numpy as np
 import matplotlib.pyplot as plt
 from collections import OrderedDict
-from collections.abc import Mapping
 from PyPDF2 import PdfFileReader, PdfFileWriter
 
 from .. import IFOS, load_budget
@@ -120,81 +118,45 @@ def load_cache(path):
     return cache
 
 
-def walk_traces(traces, root=()):
-    """recursively walk through traces
-
-    yields (key_tuple, value), where key_tuple is a tuple of nested
-    dict keys of the traces dict locating the given value.
-
-    """
-    for key, val in traces.items():
-        r = root+(key,)
-        if isinstance(val, Mapping):
-            for k, v in walk_traces(val, root=r):
-                yield k, v
-            continue
-        else:
-            yield r, val
-
-
-def zip_noises(tracesA, tracesB, skip):
+def zip_noises(traceA, traceB, skip):
     """zip matching noises from traces"""
-    for keys, (noiseA, _) in walk_traces(tracesA):
-        # keys is a tuple of nested keys for noise
-        name = keys[-1]
-        # skip keys we'll add at the end
-        if name in ['Total', 'int73']:
-            continue
+    B_dict = dict(traceB.walk())
+    for name, tA in traceA.walk():
         if skip and name in skip:
             logging.warning("SKIPPING TEST: '{}'".format(name))
             continue
-
         try:
-            t = tracesB
-            for key in keys:
-                t = t[key]
-            noiseB, style = t
+            tB = B_dict[name]
         except (KeyError, TypeError):
             logging.warning("MISSING B TRACE: '{}'".format(name))
             continue
+        yield name, tA, tB
 
-        yield name, noiseA, noiseB
 
-    yield 'Total', tracesA['Total'][0], tracesB['Total'][0]
-
-    if 'int73' in tracesA:
-        yield 'int73', tracesA['int73'][0], tracesB['int73'][0]
-
-
-def compare_traces(tracesA, tracesB, tolerance=TOLERANCE, skip=None):
-    """Compare two gwinc trace dictionaries
+def compare_traces(traceA, traceB, tolerance=TOLERANCE, skip=None):
+    """Compare two gwinc traces
 
     Noises listed in `skip` will not be compared.
 
     Returns a dictionary of noises that differ fractionally (on a
-    point-by-point basis) by more that `tolerance` between `tracesA`
-    and `tracesB`.
+    point-by-point basis) by more than `tolerance` between the traces
+    in `traceA` and `traceB`.
 
     """
     #name_width = max([len(n[0][-1]) for n in walk_traces(tracesA)])
     name_width = 15
     diffs = OrderedDict()
-    for name, noiseA, noiseB in zip_noises(tracesA, tracesB, skip):
+    for name, tA, tB in zip_noises(traceA, traceB, skip):
         logging.debug("comparing {}...".format(name))
-
-        ampA = np.sqrt(noiseA)
-        ampB = np.sqrt(noiseB)
+        ampA = tA.asd
+        ampB = tB.asd
         diff = ampB - ampA
         frac = abs(diff / ampA)
-
         if max(frac) < tolerance:
             continue
-
         logging.warning("EXCESSIVE DIFFERENCE: {:{w}} {:6.1f} ppm".format(
             name, max(frac)*1e6, w=name_width))
-
-        diffs[name] = (noiseA, noiseB, frac)
-
+        diffs[name] = (ampA, ampB, frac)
     return diffs
 
 
@@ -202,11 +164,11 @@ def plot_diffs(freq, diffs, styleA, styleB):
     spec = (len(diffs)+1, 2)
     sharex = None
     for i, nname in enumerate(diffs):
-        noiseA, noiseB, frac = diffs[nname]
+        ampA, ampB, frac = diffs[nname]
 
         axl = plt.subplot2grid(spec, (i, 0), sharex=None)
-        axl.loglog(freq, np.sqrt(noiseA), **styleA)
-        axl.loglog(freq, np.sqrt(noiseB), **styleB)
+        axl.loglog(freq, ampA, **styleA)
+        axl.loglog(freq, ampB, **styleB)
         axl.grid()
         axl.legend(loc='upper right')
         axl.set_ylabel(nname)
@@ -334,14 +296,17 @@ gwinc/test/cache/<SHA1>.  Old caches are automatically pruned.""",
             fail |= True
             continue
 
-        freq, traces_ref, attrs = load_hdf5(path)
-
+        traces_ref = load_hdf5(path)
+        traces_ref.name = 'Total'
+        freq = traces_ref.freq
         budget = load_budget(name, freq)
         traces_cur = budget.run()
+        traces_cur.name = 'Total'
 
-        if inspiral_range:
-            total_ref = traces_ref['Total'][0]
-            total_cur = traces_cur['Total'][0]
+        # FIXME: add this back
+        if False: #inspiral_range:
+            total_ref = traces_ref.psd
+            total_cur = traces_cur.psd
             range_func = inspiral_range.range
             H = inspiral_range.waveform.CBCWaveform(freq)
             fom_ref = range_func(freq, total_ref, H=H)

gwinc/trace.py

+import numpy as np
+
+
+class BudgetTrace:
+    """Budget trace data calculated from a Noise or Budget
+
+
+    """
+    def __init__(self, name=None, style=None, freq=None, psd=None, budget=None):
+        self.name = name
+        self.style = style or {}
+        self._freq = freq
+        self._psd = psd
+        self.budget = budget or []
+
+    def __repr__(self):
+        if self.budget:
+            bs = ' [{}]'.format(', '.join([str(b.name) for b in self]))
+        else:
+            bs = ''
+        return '<{} {}{}>'.format(
+            self.__class__.__name__,
+            self.name,
+            bs,
+        )
+
+    @property
+    def freq(self):
+        """trace frequency array, in Hz"""
+        return self._freq
+
+    @property
+    def psd(self):
+        """trace power spectral density array"""
+        return self._psd
+
+    @property
+    def asd(self):
+        """trace amplitude spectral density array"""
+        return np.sqrt(self._psd)
+
+    def len(self):
+        """Length of data"""
+        return len(self._freq)
+
+    def __iter__(self):
+        """iterator of budget traces"""
+        return iter(self.budget)
+
+    @property
+    def _bdict(self):
+        bdict = {trace.name: trace for trace in self.budget}
+        return bdict
+
+    def __getattr__(self, name):
+        return self._bdict[name]
+
+    def __getitem__(self, name):
+        """get budget trace by name"""
+        return self._bdict[name]
+
+    def items(self):
+        """iterator of budget (name, trace) tuples"""
+        return self._bdict.items()
+
+    def keys(self):
+        """iterator of budget trace names"""
+        return self._bdict.keys()
+
+    def values(self):
+        """iterator of budget traces"""
+        return self._bdict.values()
+
+    def walk(self):
+        """walk recursively through all traces"""
+        yield self.name, self
+        for trace in self:
+            for t in trace.walk():
+                yield t