From 8032851b9bfe5e478ddd19462dd18bf72ce851e9 Mon Sep 17 00:00:00 2001
From: Craig Warren <craig_warren@me.com>
Date: Wed, 13 Jan 2016 15:48:33 +0000
Subject: [PATCH] Added better checking of output(s) to plot.

---
 tools/plot_Ascan.py | 12 ++----------
 tools/plot_Bscan.py | 15 +++++++--------
 2 files changed, 9 insertions(+), 18 deletions(-)

diff --git a/tools/plot_Ascan.py b/tools/plot_Ascan.py
index 933f4796..7d0d2c37 100644
--- a/tools/plot_Ascan.py
+++ b/tools/plot_Ascan.py
@@ -26,13 +26,10 @@ from gprMax.exceptions import CmdInputError
 
 """Plots electric and magnetic fields and currents from all receiver points in the given output file. Each receiver point is plotted in a new figure window."""
 
-# Outputs that can be plotted
-outputslist = ['Ex', 'Ey', 'Ez', 'Hx', 'Hy', 'Hz', 'Ix', 'Iy', 'Iz']
-
 # Parse command line arguments
 parser = argparse.ArgumentParser(description='Plots electric and magnetic fields and currents from all receiver points in the given output file. Each receiver point is plotted in a new figure window.', usage='cd gprMax; python -m tools.plot_Ascan outputfile')
 parser.add_argument('outputfile', help='name of output file including path')
-parser.add_argument('--outputs', help='list of outputs to be plotted, i.e. Ex Ey Ez', default=outputslist, nargs='+')
+parser.add_argument('--outputs', help='list of outputs to be plotted, i.e. Ex Ey Ez', default=Rx.availableoutputs, nargs='+')
 parser.add_argument('-fft', action='store_true', default=False, help='plot FFT (single output must be specified)')
 args = parser.parse_args()
 
@@ -45,11 +42,6 @@ iterations = f.attrs['Iterations']
 time = np.arange(0, dt * iterations, dt)
 time = time / 1e-9
 
-# Check for valid output names
-for output in args.outputs:
-    if output not in outputslist:
-        raise CmdInputError('{} not allowed. Options are: Ex Ey Ez Hx Hy Hz'.format(output))
-
 # Check for single output component when doing a FFT
 if args.fft:
     if not len(args.outputs) == 1:
@@ -65,7 +57,7 @@ for rx in range(1, nrx + 1):
         
         # Check if requested output is in file
         if args.outputs[0] not in availablecomponents:
-            raise CmdInputError('Output request to plot {}, but the available output for receiver 1 in the file: {}'.format(args.outputs[0], ', '.join(availablecomponents)))
+            raise CmdInputError('{} output requested to plot, but the available output for receiver 1 is {}'.format(args.outputs[0], ', '.join(availablecomponents)))
         
         outputdata = f[path + args.outputs[0]][:]
         
diff --git a/tools/plot_Bscan.py b/tools/plot_Bscan.py
index 1d35e8be..82db3440 100644
--- a/tools/plot_Bscan.py
+++ b/tools/plot_Bscan.py
@@ -25,22 +25,21 @@ from gprMax.exceptions import CmdInputError
 
 """Plots a B-scan image."""
 
-# Outputs that can be plotted
-outputslist = ['Ex', 'Ey', 'Ez', 'Hx', 'Hy', 'Hz', 'Ix', 'Iy', 'Iz']
-
 # Parse command line arguments
-parser = argparse.ArgumentParser(description='Plots B-scan.', usage='cd gprMax; python -m tools.plot_Bscan outputfile --field fieldcomponent')
+parser = argparse.ArgumentParser(description='Plots a B-scan image.', usage='cd gprMax; python -m tools.plot_Bscan outputfile --field fieldcomponent')
 parser.add_argument('outputfile', help='name of output file including path')
 parser.add_argument('--output', help='name of output to be plotted, i.e. Ex Ey Ez')
 args = parser.parse_args()
 
-# Check for valid output name
-if args.output not in outputslist:
-    raise CmdInputError('{} not allowed. Options are: Ex Ey Ez Hx Hy Hz Ix Iy Iz'.format(args.output))
-
 # Open output file and read some attributes
 f = h5py.File(args.outputfile, 'r')
 path = '/rxs/rx1'
+availablecomponents = list(f[path].keys())
+    
+# Check if requested output is in file
+if args.output[0] not in availablecomponents:
+    raise CmdInputError('{} output requested to plot, but the available output for receiver 1 is {}'.format(args.outputs[0], ', '.join(availablecomponents)))
+
 outputdata = f[path + '/' + args.output]
 f.close()