Cli improvements

Some simple CLI improvements:

• allow saving to multiple files simultaneously
• fix interactive shell

closes #53 (closed)

gwinc/__main__.py

@@ -3,7 +3,6 @@ import os
 import signal
 import argparse
 import numpy as np
-from IPython.terminal.embed import InteractiveShellEmbed
 
 import logging
 logging.basicConfig(
@@ -51,6 +50,7 @@ See documentation for inspiral_range package for details.
 
 IFO = 'aLIGO'
 FREQ = '5:3000:6000'
+DATA_SAVE_FORMATS = ['.hdf5', '.h5']
 
 parser = argparse.ArgumentParser(
     prog='gwinc',
@@ -75,23 +75,23 @@ group.add_argument(
     '--interactive', '-i', action='store_true',
     help="interactive plot with interactive shell")
 group.add_argument(
-    '--save', '-s',
-    help="save budget traces (.hdf5/.h5) or plot (.pdf/.png/.svg) to file")
+    '--save', '-s', action='append',
+    help="save plot (.png/.pdf/.svg) or budget traces (.hdf5/.h5) to file (may be specified multiple times)")
 group.add_argument(
     '--yaml', '-y', action='store_true',
-    help="print IFO as yaml to stdout and exit")
+    help="print IFO as yaml to stdout and exit (budget not calculated)")
 group.add_argument(
     '--text', '-x', action='store_true',
-    help="print IFO as text table to stdout and exit")
+    help="print IFO as text table to stdout and exit (budget not calculated)")
 group.add_argument(
     '--diff', '-d', metavar='IFO',
-    help="show differences table between another IFO description")
+    help="show differences table between another IFO description and exit (budget not calculated)")
 group.add_argument(
     '--no-plot', '-np', action='store_false', dest='plot',
     help="supress plotting")
 parser.add_argument(
     'IFO',
-    help="IFO name, description file path (.yaml, .mat, .m), budget module (.py), or HDF5 data file (.hdf5, .h5)")
+    help="IFO name, or path to budget module (.py), description file (.yaml/.mat/.m), or HDF5 data file (.hdf5/.h5)")
 
 
 def freq_from_spec(spec):
@@ -113,7 +113,7 @@ def main():
     ##########
     # initial arg processing
 
-    if os.path.splitext(os.path.basename(args.IFO))[1] in ['.hdf5', '.h5']:
+    if os.path.splitext(os.path.basename(args.IFO))[1] in DATA_SAVE_FORMATS:
         from .io import load_hdf5
         Budget = None
         freq, traces, attrs = load_hdf5(args.IFO)
@@ -135,6 +135,14 @@ def main():
         plot_style = getattr(Budget, 'plot_style', {})
         traces = None
 
+    out_data_files = set()
+    out_plot_files = set()
+    if args.save:
+        for path in args.save:
+            if os.path.splitext(path)[1] in DATA_SAVE_FORMATS:
+                out_data_files.add(path)
+        out_plot_files = set(args.save) - out_data_files
+
     if args.ifo:
         for paramval in args.ifo:
             param, val = paramval.split('=', 1)
@@ -234,31 +242,12 @@ def main():
         plot_style['title'] += '\n{}'.format(fom_title)
 
     ##########
-    # output
-
-    # save noise traces to HDF5 file
-    if args.save and os.path.splitext(args.save)[1] in ['.hdf5', '.h5']:
-        from .io import save_hdf5
-        logging.info("saving budget traces {}...".format(args.save))
-        if ifo:
-            plot_style['IFO'] = ifo.to_yaml()
-        save_hdf5(
-            path=args.save,
-            freq=freq,
-            traces=traces,
-            **plot_style
-        )
+    # interactive
 
     # interactive shell plotting
-    elif args.interactive:
+    if args.interactive:
+        from IPython.terminal.embed import InteractiveShellEmbed
         ipshell = InteractiveShellEmbed(
-            user_ns={
-                'freq': freq,
-                'traces': traces,
-                'ifo': ifo,
-                'plot_style': plot_style,
-                'plot_noise': plot_noise,
-            },
             banner1='''
 GWINC interactive plotter
 
@@ -266,14 +255,39 @@ You may interact with plot using "plt." methods, e.g.:
 
 >>> plt.title("foo")
 >>> plt.savefig("foo.pdf")
-''')
+''',
+            user_ns={
+                'freq': freq,
+                'traces': traces,
+                'ifo': ifo,
+                'plot_style': plot_style,
+                'plot_noise': plot_noise,
+            },
+        )
         ipshell.enable_pylab()
-        ipshell.run_code("plot_noise(freq, traces, **plot_style)")
-        ipshell.run_code("plt.title('{}')".format(plot_style['title']))
+        ipshell.ex("fig = plot_noise(freq, traces, **plot_style)")
+        ipshell.ex("plt.title('{}')".format(plot_style['title']))
         ipshell()
 
+    ##########
+    # output
+
+    # save noise traces to HDF5 file
+    if out_data_files:
+        from .io import save_hdf5
+        attrs = dict(plot_style)
+        attrs['IFO'] = ifo.to_yaml()
+        for path in out_data_files:
+            logging.info("saving budget traces {}...".format(path))
+            save_hdf5(
+                path=path,
+                freq=freq,
+                traces=traces,
+                **attrs,
+            )
+
     # standard plotting
-    elif args.plot:
+    if args.plot or out_plot_files:
         logging.info("plotting noises...")
         fig = plt.figure()
         ax = fig.add_subplot(1, 1, 1)
@@ -284,10 +298,9 @@ You may interact with plot using "plt." methods, e.g.:
             **plot_style
         )
         fig.tight_layout()
-        if args.save:
-            fig.savefig(
-                args.save,
-            )
+        if out_plot_files:
+            for path in out_plot_files:
+                fig.savefig(path)
         else:
             plt.show()