diff --git a/README.md b/README.md index 3f261c18eb862cadd48d225eb57bdc85a26143ee..717c09b3b157977e2bfd46c0bf325e9bd7ba62a1 100644 --- a/README.md +++ b/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 diff --git a/gwinc/__init__.py b/gwinc/__init__.py index 2643cf8d04672f7b9cb5dfe8c384ceb41b36d584..299379dda23fe462048ab48e1022a43a2ea65f84 100644 --- a/gwinc/__init__.py +++ b/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 diff --git a/gwinc/__main__.py b/gwinc/__main__.py index 017b8b25464e4213787fd3d89a83c502464a7e1c..3fafdf033a637d09ce770db5dc90191f657721c6 100644 --- a/gwinc/__main__.py +++ b/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 ) diff --git a/gwinc/plot.py b/gwinc/plot.py index d850ea523cd2a8fc57effed85140061a230c066c..45b005262a7c553b49c2e15f1dbd1fa4535b4f0d 100644 --- a/gwinc/plot.py +++ b/gwinc/plot.py @@ -1,5 +1,5 @@ -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 diff --git a/gwinc/trace.py b/gwinc/trace.py index 175d5e4628d89de2f2379dd948588cff7d3c5946..e78b3a456ef60da5cce54376a24b49f464c9a977 100644 --- a/gwinc/trace.py +++ b/gwinc/trace.py @@ -1,5 +1,7 @@ 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)