Added save (to file) option and rx gather option

toolboxes/Plotting/plot_Ascan.py

@@ -30,18 +30,19 @@ from gprMax.utilities.utilities import fft_power
 logger = logging.getLogger(__name__)
 
 
-def mpl_plot(filename, outputs=Rx.defaultoutputs, fft=False):
+def mpl_plot(filename, outputs=Rx.defaultoutputs, fft=False, save=False):
     """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.
 
     Args:
-        filename (string): Filename (including path) of output file.
+        filename: string of filename (including path) of output file.
-        outputs (list): List of field/current components to plot.
+        outputs: list of field/current components to plot.
-        fft (boolean): Plot FFT switch.
+        fft: boolean flag to plot FFT.
+        save: boolean flag to save plot to file.
 
     Returns:
-        plt (object): matplotlib plot object.
+        plt: matplotlib plot object.
     """
 
     file = Path(filename)
@@ -80,7 +81,8 @@ def mpl_plot(filename, outputs=Rx.defaultoutputs, fft=False):
         # Check for single output component when doing a FFT
         if fft:
             if not len(outputs) == 1:
-                logger.exception('A single output must be specified when using the -fft option')
+                logger.exception('A single output must be specified when using ' +
+                                 'the -fft option')
                 raise ValueError
 
         # New plot for each receiver
@@ -102,7 +104,9 @@ def mpl_plot(filename, outputs=Rx.defaultoutputs, fft=False):
                     output = outputs[0]
 
                 if output not in availableoutputs:
-                    logger.exception(f"{output} output requested to plot, but the available output for receiver 1 is {', '.join(availableoutputs)}")
+                    logger.exception(f"{output} output requested to plot, but " +
+                                     f"the available output for receiver 1 is " +
+                                     f"{', '.join(availableoutputs)}")
                     raise ValueError
 
                 outputdata = f[rxpath + output][:] * polarity
@@ -145,19 +149,22 @@ def mpl_plot(filename, outputs=Rx.defaultoutputs, fft=False):
                     ax2.set_ylabel('Power [dB]')
                     ax2.grid(which='both', axis='both', linestyle='-.')
 
-                    # Change colours and labels for magnetic field components or currents
+                    # Change colours and labels for magnetic field components 
+                    # or currents
                     if 'H' in outputs[0]:
                         plt.setp(line1, color='g')
                         plt.setp(line2, color='g')
                         plt.setp(ax1, ylabel=outputtext + ' field strength [A/m]')
                         plt.setp(stemlines, 'color', 'g')
-                        plt.setp(markerline, 'markerfacecolor', 'g', 'markeredgecolor', 'g')
+                        plt.setp(markerline, 'markerfacecolor', 'g', 
+                                 'markeredgecolor', 'g')
                     elif 'I' in outputs[0]:
                         plt.setp(line1, color='b')
                         plt.setp(line2, color='b')
                         plt.setp(ax1, ylabel=outputtext + ' current [A]')
                         plt.setp(stemlines, 'color', 'b')
-                        plt.setp(markerline, 'markerfacecolor', 'b', 'markeredgecolor', 'b')
+                        plt.setp(markerline, 'markerfacecolor', 'b', 
+                                 'markeredgecolor', 'b')
 
                     plt.show()
 
@@ -179,7 +186,8 @@ def mpl_plot(filename, outputs=Rx.defaultoutputs, fft=False):
                         plt.setp(line, color='b')
                         plt.setp(ax, ylabel=outputtext + ', current [A]')
 
-            # If multiple outputs required, create all nine subplots and populate only the specified ones
+            # If multiple outputs required, create all nine subplots and 
+            # populate only the specified ones
             else:
                 fig, ax = plt.subplots(subplot_kw=dict(xlabel='Time [s]'),
                                        num=rxpath + ' - ' + f[rxpath].attrs['Name'],
@@ -190,7 +198,8 @@ def mpl_plot(filename, outputs=Rx.defaultoutputs, fft=False):
                     gs = gridspec.GridSpec(3, 2, hspace=0.3, wspace=0.3)
 
                 for output in outputs:
-                    # Check for polarity of output and if requested output is in file
+                    # Check for polarity of output and if requested output 
+                    # is in file
                     if output[-1] == 'm':
                         polarity = -1
                         outputtext = '-' + output[0:-1]
@@ -201,7 +210,10 @@ def mpl_plot(filename, outputs=Rx.defaultoutputs, fft=False):
 
                     # Check if requested output is in file
                     if output not in availableoutputs:
-                        logger.exception(f"Output(s) requested to plot: {', '.join(outputs)}, but available output(s) for receiver {rx} in the file: {', '.join(availableoutputs)}")
+                        logger.exception(f"Output(s) requested to plot: " +
+                                         f"{', '.join(outputs)}, but available output(s) " +
+                                         f"for receiver {rx} in the file: " +
+                                         f"{', '.join(availableoutputs)}")
                         raise ValueError
 
                     outputdata = f[rxpath + output][:] * polarity
@@ -252,31 +264,39 @@ def mpl_plot(filename, outputs=Rx.defaultoutputs, fft=False):
                     ax.set_xlim([0, np.amax(time)])
                     ax.grid(which='both', axis='both', linestyle='-.')
 
-            # Save a PDF/PNG of the figure
-            savename = file.stem + '_rx' + str(rx)
-            savename = file.parent / savename
-            # fig.savefig(savename.with_suffix('.pdf'), dpi=None, format='pdf',
-            #             bbox_inches='tight', pad_inches=0.1)
-            # fig.savefig(savename.with_suffix('.png'), dpi=150, format='png',
-            #             bbox_inches='tight', pad_inches=0.1)
-
     f.close()
 
+    if save:
+        # Save a PDF of the figure
+        fig.savefig(filename[:-3] + '.pdf', dpi=None, format='pdf', 
+                    bbox_inches='tight', pad_inches=0.1)
+        # Save a PNG of the figure
+        # fig.savefig(filename[:-3] + '.png', dpi=150, format='png', 
+        #             bbox_inches='tight', pad_inches=0.1)
+
     return plt
 
 
 if __name__ == "__main__":
 
     # 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 = 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 toolboxes.Plotting.plot_Ascan outputfile')
     parser.add_argument('outputfile', help='name of output file including path')
     parser.add_argument('--outputs', help='outputs to be plotted',
                         default=Rx.defaultoutputs,
-                        choices=['Ex', 'Ey', 'Ez', 'Hx', 'Hy', 'Hz', 'Ix', 'Iy', 'Iz', 'Ex-', 'Ey-', 'Ez-', 'Hx-', 'Hy-', 'Hz-', 'Ix-', 'Iy-', 'Iz-'],
+                        choices=['Ex', 'Ey', 'Ez', 'Hx', 'Hy', 'Hz', 
+                                 'Ix', 'Iy', 'Iz', 'Ex-', 'Ey-', 'Ez-', 
+                                 'Hx-', 'Hy-', 'Hz-', 'Ix-', 'Iy-', 'Iz-'],
                         nargs='+')
-    parser.add_argument('-fft', action='store_true', help='plot FFT (single output must be specified)',
+    parser.add_argument('-fft', action='store_true', default=False,
-                        default=False)
+                        help='plot FFT (single output must be specified)')
+    parser.add_argument('-save', action='store_true', default=False,
+                        help='save plot directly to file, i.e. do not display')
     args = parser.parse_args()
 
-    plthandle = mpl_plot(args.outputfile, args.outputs, fft=args.fft)
+    plthandle = mpl_plot(args.outputfile, args.outputs, fft=args.fft, save=args.save)
-    plthandle.show()
+    
+    plthandle.show()

toolboxes/Plotting/plot_Bscan.py

@@ -24,23 +24,24 @@ import h5py
 import matplotlib.pyplot as plt
 import numpy as np
 
-from .outputfiles_merge import get_output_data
+from ..Utilities.outputfiles_merge import get_output_data
 
 logger = logging.getLogger(__name__)
 
 
-def mpl_plot(filename, outputdata, dt, rxnumber, rxcomponent):
+def mpl_plot(filename, outputdata, dt, rxnumber, rxcomponent, save=False):
     """Creates a plot (with matplotlib) of the B-scan.
 
     Args:
-        filename (string): Filename (including path) of output file.
+        filename: string of filename (including path) of output file.
-        outputdata (array): Array of A-scans, i.e. B-scan data.
+        outputdata: array of A-scans, i.e. B-scan data.
-        dt (float): Temporal resolution of the model.
+        dt: float of temporal resolution of the model.
-        rxnumber (int): Receiver output number.
+        rxnumber: int of receiver output number.
-        rxcomponent (str): Receiver output field/current component.
+        rxcomponent: string of receiver output field/current component.
+        save: boolean flag to save plot to file.
 
     Returns:
-        plt (object): matplotlib plot object.
+        plt: matplotlib plot object.
     """
 
     file = Path(filename)
@@ -65,11 +66,13 @@ def mpl_plot(filename, outputdata, dt, rxnumber, rxcomponent):
     elif 'I' in rxcomponent:
         cb.set_label('Current [A]')
 
-    # Save a PDF/PNG of the figure
+    if save:
-    # fig.savefig(file.with_suffix('.pdf'), dpi=None, format='pdf',
+        # Save a PDF of the figure
-    #             bbox_inches='tight', pad_inches=0.1)
+        fig.savefig(filename[:-3] + '.pdf', dpi=None, format='pdf', 
-    # fig.savefig(file.with_suffix('.png'), dpi=150, format='png',
+                    bbox_inches='tight', pad_inches=0.1)
-    #             bbox_inches='tight', pad_inches=0.1)
+        # Save a PNG of the figure
+        # fig.savefig(filename[:-3] + '.png', dpi=150, format='png', 
+        #             bbox_inches='tight', pad_inches=0.1)
 
     return plt
 
@@ -78,10 +81,14 @@ if __name__ == "__main__":
 
     # Parse command line arguments
     parser = argparse.ArgumentParser(description='Plots a B-scan image.',
-                                     usage='cd gprMax; python -m tools.plot_Bscan outputfile output')
+                                     usage='cd gprMax; python -m toolboxes.Plotting.plot_Bscan outputfile output')
     parser.add_argument('outputfile', help='name of output file including path')
     parser.add_argument('rx_component', help='name of output component to be plotted',
                         choices=['Ex', 'Ey', 'Ez', 'Hx', 'Hy', 'Hz', 'Ix', 'Iy', 'Iz'])
+    parser.add_argument('-gather', action='store_true', default=False,
+                        help='gather together all receiver outputs in file')
+    parser.add_argument('-save', action='store_true', default=False,
+                        help='save plot directly to file, i.e. do not display')
     args = parser.parse_args()
 
     # Open output file and read number of outputs (receivers)
@@ -94,8 +101,28 @@ if __name__ == "__main__":
         logger.exception(f'No receivers found in {args.outputfile}')
         raise ValueError
 
+    rxsgather = []
     for rx in range(1, nrx + 1):
         outputdata, dt = get_output_data(args.outputfile, rx, args.rx_component)
-        plthandle = mpl_plot(args.outputfile, outputdata, dt, rx, args.rx_component)
+        print(outputdata.shape)
+        if args.gather:
+            rxsgather.append(outputdata)
+        else:
+            plthandle = mpl_plot(args.outputfile, outputdata, dt, rx, 
+                                 args.rx_component, save=args.save)
 
-    plthandle.show()
+    # Plot all receivers from single output file together if required
+    if args.gather:
+        rxsgather = np.array(rxsgather).transpose()
+        plthandle = mpl_plot(args.outputfile, rxsgather, dt, rx, 
+                             args.rx_component, save=args.save)
+
+    if args.save:
+        # Save a PDF of the figure
+        plthandle.savefig(args.outputfile[:-3] + '.pdf', dpi=None, 
+                          format='pdf', bbox_inches='tight', pad_inches=0.1)
+        # # Save a PNG of the figure
+        # plthandle.savefig(args.outputfile[:-3] + '.png', dpi=150, 
+        #                   format='png', bbox_inches='tight', pad_inches=0.1)
+    else:
+        plthandle.show()