add plot method to BudgetTrace object

just calls plot.plot_trace

closes #75

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/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
@@ -77,3 +79,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)