README.md

@@ -115,7 +115,7 @@ accessed directly through the `gwinc` library interface:
 >>> import gwinc
 >>> budget = gwinc.load_budget('aLIGO')
 >>> trace = budget.run()
->>> fig = gwinc.plot_budget(trace)
+>>> fig = trace.plot()
 >>> fig.show()
 ```
 A default frequency array is used, but alternative frequencies can be

gwinc/__init__.py

@@ -17,6 +17,7 @@ except ModuleNotFoundError:
         __version__ = '?.?.?'
 from .ifo import IFOS
 from .struct import Struct
+from .plot import plot_trace
 from .plot import plot_budget
 from .plot import plot_noise
 
@@ -209,6 +210,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_budget(traces, **plot_style)
+        traces.plot(**plot_style)
 
     return score, noises, ifo

gwinc/__main__.py

@@ -10,7 +10,6 @@ from . import (
     DEFAULT_FREQ,
     freq_from_spec,
     load_budget,
-    plot_budget,
     logger,
 )
 from . import io
@@ -261,14 +260,14 @@ def main():
     if args.interactive:
         banner = """GWINC interactive shell
 
-The 'ifo' Struct and 'budget' trace data objects are available for
+The 'ifo' Struct, 'budget', and 'trace' 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_budget()' function:
+You may plot the budget using the 'trace.plot()' method:
 
-In [.]: plot_budget(budget, **plot_style)
+In [.]: trace.plot(**plot_style)
 """
         banner += """
 You may interact with the plot using the 'plt' functions, e.g.:
@@ -277,19 +276,20 @@ In [.]: plt.title("foo")
 In [.]: plt.savefig("foo.pdf")
 """
         from IPython.terminal.embed import InteractiveShellEmbed
+        if subtitle:
+            plot_style['title'] += '\n' + subtitle
         ipshell = InteractiveShellEmbed(
             banner1=banner,
             user_ns={
-                'budget': trace,
                 'ifo': ifo,
+                'budget': budget,
+                'trace': trace,
                 'plot_style': plot_style,
-                'plot_budget': plot_budget,
             },
         )
-        ipshell.enable_pylab()
+        ipshell.enable_pylab(import_all=False)
         if args.plot:
-            ipshell.ex("fig = plot_budget(budget, **plot_style)")
-            ipshell.ex("plt.title(plot_style['title'])")
+            ipshell.ex("fig = trace.plot(**plot_style)")
         ipshell()
 
     ##########
@@ -313,8 +313,7 @@ In [.]: plt.savefig("foo.pdf")
         ax = fig.add_subplot(1, 1, 1)
         if subtitle:
             plot_style['title'] += '\n' + subtitle
-        plot_budget(
-            trace,
+        trace.plot(
             ax=ax,
             **plot_style
         )

gwinc/ifo/noises.py

@@ -14,7 +14,6 @@ from .. import suspension
 # helper functions
 ############################################################
 
-
 def mirror_struct(ifo, tm):
     """Create a "mirror" Struct for a LIGO core optic
 
@@ -126,7 +125,6 @@ def ifo_power(ifo, PRfixed=True):
 
 ##################################################
 
-
 def precomp_suspension(f, ifo):
     pc = Struct()
     pc.VHCoupling = Struct()
@@ -142,7 +140,6 @@ def precomp_suspension(f, ifo):
     return pc
 
 
-
 def precomp_quantum(f, ifo):
     pc = Struct()
     mirror_mass = mirror_struct(ifo, 'ETM').MirrorMass
@@ -345,35 +342,58 @@ class StandardQuantumLimit(nb.Noise):
         ETM = mirror_struct(self.ifo, 'ETM')
         return 8 * const.hbar / (ETM.MirrorMass * (2 * np.pi * self.freq) ** 2)
 
+
 #########################
 # seismic
 #########################
 
-class Seismic(nb.Noise):
-    """Seismic
+class SeismicHorizontal(nb.Noise):
+    """Horizontal seismic noise
 
     """
     style = dict(
-        label='Seismic',
-        color='#855700',
+        label='Horizontal',
+        color='xkcd:muted blue',
+    )
+
+    @nb.precomp(sustf=precomp_suspension)
+    def calc(self, sustf):
+        nt, nr = noise.seismic.platform_motion(self.freq, self.ifo)
+        n = noise.seismic.seismic_suspension_filtered(sustf, nt, 'horiz')
+        return n * 4
+
+
+class SeismicVertical(nb.Noise):
+    """Vertical seismic noise
+
+    """
+    style = dict(
+        label='Vertical',
+        color='xkcd:brick red',
     )
 
     @nb.precomp(sustf=precomp_suspension)
     def calc(self, sustf):
-        if 'PlatformMotion' in self.ifo.Seismic:
-            if self.ifo.Seismic.PlatformMotion == 'BSC':
-                nt, nr = noise.seismic.seismic_BSC_ISI(self.freq)
-            elif self.ifo.Seismic.PlatformMotion == '6D':
-                nt, nr = noise.seismic.seismic_BSC_ISI_6D(self.freq)
-            else:
-                nt, nr = noise.seismic.seismic_BSC_ISI(self.freq)
-        else:
-            nt, nr = noise.seismic.seismic_BSC_ISI(self.freq)
-        n, nh, nv = noise.seismic.seismic_suspension_fitered(
-            self.ifo.Suspension, sustf, nt)
+        nt, nr = noise.seismic.platform_motion(self.freq, self.ifo)
+        n = noise.seismic.seismic_suspension_filtered(sustf, nt, 'vert')
         return n * 4
 
 
+class Seismic(nb.Budget):
+    """Seismic
+
+    """
+    style = dict(
+        label='Seismic',
+        color='#855700',
+    )
+
+    noises = [
+        SeismicHorizontal,
+        SeismicVertical,
+    ]
+
+
 #########################
 # Newtonian
 #########################
@@ -572,7 +592,6 @@ class SubstrateThermoElastic(nb.Noise):
 # residual gas
 #########################
 
-
 class ExcessGas(nb.Noise):
     """Excess Gas
 

gwinc/noise/seismic.py

@@ -6,32 +6,50 @@ import numpy as np
 from scipy.interpolate import PchipInterpolator as interp1d
 
 
-def seismic_suspension_fitered(sus, sustf, in_trans):
+def seismic_suspension_filtered(sustf, in_trans, direction):
     """Seismic displacement noise for single suspended test mass.
 
-    :sus: suspension Struct
     :sustf: sus transfer function Struct
     :in_trans: input translational displacement spectrum
+    :direction: 'horiz' for horizontal or 'vert' for vertical
 
-    :returns: tuple of displacement noise power spectrum at :f:, and
-    horizontal and vertical components.
+    :returns: tuple of displacement noise power spectrum at :f:
 
     """
-    hTable = sustf.hTable
-    vTable = sustf.vTable
+    if direction == 'horiz':
+        # horizontal noise total
+        n = (abs(sustf.hTable)**2) * in_trans**2
 
-    theta = sustf.VHCoupling.theta
+    elif direction == 'vert':
+        # vertical to horizontal coupling
+        theta = sustf.VHCoupling.theta
 
-    # horizontal noise total
-    nh = (abs(hTable)**2) * in_trans**2
+        # vertical noise total
+        n = (abs(theta * sustf.vTable)**2) * in_trans**2
 
-    # vertical noise total
-    nv = (abs(theta * vTable)**2) * in_trans**2
+    return n
 
-    # new total noise
-    n = nv + nh
 
-    return n, nh, nv
+def platform_motion(f, ifo):
+    """Compute the platform motion
+
+    :f: frequency array in Hz
+    :ifo: the IFO struct
+
+    :returns: tuple of displacement noise power spectrum at :f: for
+    translational and rotational DOFs.
+    """
+    if 'PlatformMotion' in ifo.Seismic:
+        if ifo.Seismic.PlatformMotion == 'BSC':
+            nt, nr = seismic_BSC_ISI(f)
+        elif ifo.Seismic.PlatformMotion == '6D':
+            nt, nr = seismic_BSC_ISI_6D(f)
+        else:
+            nt, nr = seismic_BSC_ISI(f)
+    else:
+        nt, nr = seismic_BSC_ISI(f)
+
+    return nt, nr
 
 
 def seismic_BSC_ISI(f):

gwinc/plot.py

-def plot_budget(
-        budget,
+def plot_trace(
+        trace,
         ax=None,
         **kwargs
 ):
@@ -17,29 +17,29 @@ def plot_budget(
     else:
         fig = ax.figure
 
-    total = budget.asd
+    total = trace.asd
     ylim = [min(total)/10, max(total)]
     style = dict(
         color='#000000',
         alpha=0.6,
         lw=4,
     )
-    style.update(getattr(budget, 'style', {}))
+    style.update(getattr(trace, 'style', {}))
     if 'label' in style:
         style['label'] = 'Total ' + style['label']
     else:
         style['label'] = 'Total'
-    ax.loglog(budget.freq, total, **style)
+    ax.loglog(trace.freq, total, **style)
 
-    for name, trace in budget.items():
-        style = trace.style
+    for name, strace in trace.items():
+        style = strace.style
         if 'label' not in style:
             style['label'] = name
         if 'linewidth' in style:
             style['lw'] = style['linewidth']
         elif 'lw' not in style:
             style['lw'] = 3
-        ax.loglog(budget.freq, trace.asd, **style)
+        ax.loglog(trace.freq, strace.asd, **style)
 
     ax.grid(
         True,
@@ -56,7 +56,7 @@ def plot_budget(
 
     ax.autoscale(enable=True, axis='y', tight=True)
     ax.set_ylim(kwargs.get('ylim', ylim))
-    ax.set_xlim(budget.freq[0], budget.freq[-1])
+    ax.set_xlim(trace.freq[0], trace.freq[-1])
     ax.set_xlabel('Frequency [Hz]')
     if 'ylabel' in kwargs:
         ax.set_ylabel(kwargs['ylabel'])
@@ -67,4 +67,5 @@ def plot_budget(
 
 
 # FIXME: deprecate
-plot_noise = plot_budget
+plot_noise = plot_trace
+plot_budget = plot_trace

gwinc/trace.py

 import numpy as np
 
+from .plot import plot_trace
+
 
 class BudgetTrace:
     """Budget trace data calculated from a Noise or Budget
@@ -53,7 +55,10 @@ class BudgetTrace:
         return bdict
 
     def __getattr__(self, name):
-        return self._bdict[name]
+        try:
+            return self._bdict[name]
+        except KeyError:
+            raise AttributeError
 
     def __getitem__(self, name):
         """get budget trace by name"""
@@ -77,3 +82,16 @@ class BudgetTrace:
         for trace in self:
             for t in trace.walk():
                 yield t
+
+
+    def plot(self, ax=None, **kwargs):
+        """Plot the trace budget
+
+        If an axis handle `ax` is provided it will be used for the
+        plot.  All remaining keyword arguments are assumed to define
+        various matplotlib plot style attributes.
+
+        Returns the figure handle.
+
+        """
+        return plot_trace(self, ax=ax, **kwargs)