Overhauled module for simpler checking and plotting of data.

tests/test_compare_experimental.py

@@ -16,43 +16,45 @@
 # You should have received a copy of the GNU General Public License
 # along with gprMax.  If not, see <http://www.gnu.org/licenses/>.
 
-import sys, os
+import sys, os, argparse
 import h5py
 import numpy as np
 import matplotlib as mpl
 import matplotlib.pyplot as plt
 
-"""Compare field outputs
-    
-Usage:
-    cd gprMax
-    python -m tests.test_compare_experimental path_to_model_output path_to_real_output
-    
-"""
+from gprMax.exceptions import CmdInputError
 
-modelfile = sys.argv[1]
-realfile = sys.argv[2]
-path = '/rxs/rx1/'
-# Key refers to subplot location
-fields = {0: 'Ex', 1: 'Ey', 2: 'Ez', 3: 'Hx', 4: 'Hy', 5: 'Hz'}
-plotorder = list(fields.keys())
+"""Plots a comparison of fields between given simulation output and experimental data files."""
+
+# Parse command line arguments
+parser = argparse.ArgumentParser(description='Plots a comparison of fields between given simulation output and experimental data files.', usage='cd gprMax; python -m tests.test_compare_experimental modelfile realfile output')
+parser.add_argument('modelfile', help='name of model output file including path')
+parser.add_argument('realfile', help='name of file containing experimental data including path')
+parser.add_argument('output', help='output to be plotted, i.e. Ex Ey Ez', nargs='+')
+args = parser.parse_args()
 
 # Model results
-f = h5py.File(modelfile, 'r')
-floattype = f[path + 'Ex'].dtype
-model = np.zeros((f.attrs['Iterations'], 6), dtype=floattype)
+f = h5py.File(args.modelfile, '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.output[0], ', '.join(availablecomponents)))
+
+floattype = f[path + args.output[0]].dtype
+model = np.zeros((f.attrs['Iterations']), dtype=floattype)
 timemodel = np.zeros((f.attrs['Iterations']), dtype=floattype)
 timemodel = np.arange(0, f.attrs['dt'] * f.attrs['Iterations'], f.attrs['dt']) / 1e-9
-for ID, name in fields.items():
-    model[:,ID] = f[path + str(name)][:] * -1
-    model[:,ID] = model[:,ID] / np.amax(np.abs(model[:,ID]))
+model = f[path + args.output[0]][:] * -1
+model /= np.amax(np.abs(model))
 f.close()
 
 # Select model field of interest and find max
-modelmax = np.where(np.abs(model[:,1]) == 1)[0][0]
+modelmax = np.where(np.abs(model) == 1)[0][0]
 
 # Real results
-with open(realfile, 'r') as f:
+with open(args.realfile, 'r') as f:
     real = np.loadtxt(f)
 real[:,1] = real[:,1] / np.amax(np.abs(real[:,1]))
 realmax = np.where(np.abs(real[:,1]) == 1)[0][0]
@@ -60,18 +62,17 @@ realmax = np.where(np.abs(real[:,1]) == 1)[0][0]
 difftime = - (timemodel[modelmax] - real[realmax,0])
 
 # Plot modelled and real data
-fig, ax = plt.subplots(num=modelfile + ' versus ' + realfile, figsize=(20, 10), facecolor='w', edgecolor='w')
-ax.plot(timemodel + difftime, model[:,1], 'r', lw=2, label='Model')
+fig, ax = plt.subplots(num=args.modelfile + ' versus ' + args.realfile, figsize=(20, 10), facecolor='w', edgecolor='w')
+ax.plot(timemodel + difftime, model, 'r', lw=2, label='Model')
 ax.plot(real[:,0], real[:,1], 'r', ls='--', lw=2, label='Experiment')
 ax.set_xlabel('Time [ns]')
 ax.set_ylabel('Amplitude')
 ax.set_xlim([0, timemodel[-1]])
 ax.set_ylim([-1, 1])
 ax.legend()
-[label.set_bbox(dict(facecolor='white', edgecolor='None', alpha=0.65 )) for label in ax.get_xticklabels() + ax.get_yticklabels()]
 ax.grid()
 
 # Show/print plots
-savename = os.path.abspath(os.path.dirname(modelfile)) + os.sep + os.path.splitext(os.path.split(modelfile)[1])[0] + '_vs_' + os.path.splitext(os.path.split(realfile)[1])[0]
+savename = os.path.abspath(os.path.dirname(args.modelfile)) + os.sep + os.path.splitext(os.path.split(args.modelfile)[1])[0] + '_vs_' + os.path.splitext(os.path.split(args.realfile)[1])[0]
 #fig.savefig(savename + '.pdf', dpi=None, format='pdf', bbox_inches='tight', pad_inches=0.1)
 plt.show()