Restructured so this module can be imported as a package as well as used with command line arguments.

tools/plot_builtin_wave.py

@@ -25,69 +25,76 @@ from gprMax.exceptions import CmdInputError
 from gprMax.utilities import round_value
 from gprMax.waveforms import Waveform
 
+"""Plot built-in waveforms that can be used with sources."""
 
-"""Plot built-in waveforms that can be used for sources."""
 
-# Parse command line arguments
+def check_timewindow(timewindow, dt):
-parser = argparse.ArgumentParser(description='Plot built-in waveforms that can be used for sources.', usage='cd gprMax; python -m tools.plot_builtin_wave type amp freq timewindow dt')
+    """Checks and sets time window and number of iterations.
-parser.add_argument('type', help='type of waveform, e.g. gaussian, ricker etc...')
-parser.add_argument('amp', type=float, help='amplitude of waveform')
-parser.add_argument('freq', type=float, help='centre frequency of waveform')
-parser.add_argument('timewindow', help='time window to view waveform')
-parser.add_argument('dt', type=float, help='time step to view waveform')
-parser.add_argument('-fft', action='store_true', default=False, help='plot FFT')
-args = parser.parse_args()
             
-# Check waveform parameters
+    Args:
-if args.type.lower() not in Waveform.types:
+        timewindow (float): Time window.
-    raise CmdInputError('The waveform must have one of the following types {}'.format(', '.join(Waveform.types)))
+        dt (float): Time discretisation.
-if args.freq <= 0:
-    raise CmdInputError('The waveform requires an excitation frequency value of greater than zero')
             
-w = Waveform()
+    Returns:
-w.type = args.type
+        timewindow (float): Time window.
-w.amp = args.amp
+        iterations (int): Number of interations.
-w.freq = args.freq
+    """
-dt = args.dt
     
-# Check time window
+    # Time window could be a string, float or int, so convert to string then check
-if '.' in args.timewindow or 'e' in args.timewindow:
+    timewindow = str(timewindow)
-    if float(args.timewindow) > 0:
+    
-        timewindow = float(args.timewindow)
+    try:
-        iterations = round_value((float(args.timewindow) / dt)) + 1
+        timewindow = int(timewindow)
+        iterations = timewindow
+        timewindow = (timewindow - 1) * dt
+
+    except:
+        timewindow = float(timewindow)
+        if timewindow > 0:
+            iterations = round_value((timewindow / dt)) + 1
         else:
             raise CmdInputError('Time window must have a value greater than zero')
-    # If number of iterations given
-else:
-    timewindow = (int(args.timewindow) - 1) * dt
-    iterations = int(args.timewindow)
 
-time = np.linspace(0, 1, iterations)
+    return timewindow, iterations
-time *= (iterations * dt)
-waveform = np.zeros(len(time))
-timeiter = np.nditer(time, flags=['c_index'])
 
-while not timeiter.finished:
+
+def plot_waveform(w, timewindow, dt, iterations, fft=False):
+    """Plots waveform and prints useful information about its properties.
+            
+    Args:
+        w (class): Waveform class instance.
+        timewindow (float): Time window.
+        dt (float): Time discretisation.
+        iterations (int): Number of iterations.
+        fft (boolean): Plot FFT of waveform.
+    """
+    
+    time = np.linspace(0, 1, iterations)
+    time *= (iterations * dt)
+    waveform = np.zeros(len(time))
+    timeiter = np.nditer(time, flags=['c_index'])
+
+    while not timeiter.finished:
         waveform[timeiter.index] = w.calculate_value(timeiter[0], dt)
         timeiter.iternext()
 
-print('Waveform characteristics...')
+    print('Waveform characteristics...')
-print('Type: {}'.format(w.type))
+    print('Type: {}'.format(w.type))
-print('Amplitude: {:g}'.format(w.amp))
+    print('Amplitude: {:g}'.format(w.amp))
-print('Centre frequency: {:g} Hz'.format(w.freq))
+    print('Centre frequency: {:g} Hz'.format(w.freq))
-print('Time to centre of pulse: {:g} s'.format(1 / w.freq))
+    print('Time to centre of pulse: {:g} s'.format(1 / w.freq))
 
-# Calculate pulse width for gaussian
+    # Calculate pulse width for gaussian
-if w.type == 'gaussian':
+    if w.type == 'gaussian':
         powerdrop = -3 #dB
         start = np.where((10 * np.log10(waveform / np.amax(waveform))) > powerdrop)[0][0]
         stop = np.where((10 * np.log10(waveform[start:] / np.amax(waveform))) < powerdrop)[0][0] + start
         print('Pulse width at {:d}dB, i.e. FWHM: {:g} s'.format(powerdrop, time[stop] - time[start]))
 
-print('Time window: {:g} s ({} iterations)'.format(timewindow, iterations))
+    print('Time window: {:g} s ({} iterations)'.format(timewindow, iterations))
-print('Time step: {:g} s'.format(dt))
+    print('Time step: {:g} s'.format(dt))
 
-if args.fft:
+    if fft:
         # Calculate magnitude of frequency spectra of waveform
         power = 10 * np.log10(np.abs(np.fft.fft(waveform))**2)
         freqs = np.fft.fftfreq(power.size, d=dt)
@@ -115,7 +122,7 @@ if args.fft:
         ax2.set_xlabel('Frequency [Hz]')
         ax2.set_ylabel('Power [dB]')
 
-else:
+    else:
         fig, ax1 = plt.subplots(num=w.type, figsize=(20, 10), facecolor='w', edgecolor='w')
 
         # Plot waveform
@@ -123,11 +130,42 @@ else:
         ax1.set_xlabel('Time [s]')
         ax1.set_ylabel('Amplitude')
 
-[ax.grid() for ax in fig.axes] # Turn on grid
+    [ax.grid() for ax in fig.axes] # Turn on grid
+
+    # Save a PDF/PNG of the figure
+    #fig.savefig(os.path.dirname(os.path.abspath(__file__)) + os.sep + w.type + '.pdf', dpi=None, format='pdf', bbox_inches='tight', pad_inches=0.1)
+    #fig.savefig(os.path.dirname(os.path.abspath(__file__)) + os.sep + w.type + '.png', dpi=150, format='png', bbox_inches='tight', pad_inches=0.1)
+
+    plt.show()
+
+
+if __name__ == "__main__":
+    
+    # Parse command line arguments
+    parser = argparse.ArgumentParser(description='Plot built-in waveforms that can be used for sources.', usage='cd gprMax; python -m tools.plot_builtin_wave type amp freq timewindow dt')
+    parser.add_argument('type', help='type of waveform', choices=Waveform.types)
+    parser.add_argument('amp', type=float, help='amplitude of waveform')
+    parser.add_argument('freq', type=float, help='centre frequency of waveform')
+    parser.add_argument('timewindow', help='time window to view waveform')
+    parser.add_argument('dt', type=float, help='time step to view waveform')
+    parser.add_argument('-fft', action='store_true', default=False, help='plot FFT of waveform')
+    args = parser.parse_args()
+
+    # Check waveform parameters
+    if args.type.lower() not in Waveform.types:
+        raise CmdInputError('The waveform must have one of the following types {}'.format(', '.join(Waveform.types)))
+    if args.freq <= 0:
+        raise CmdInputError('The waveform requires an excitation frequency value of greater than zero')
+
+    # Create waveform instance
+    w = Waveform()
+    w.type = args.type
+    w.amp = args.amp
+    w.freq = args.freq
+
+    timewindow, iterations = check_timewindow(args.timewindow, args.dt)
+    plot_waveform(w, timewindow, args.dt, iterations, args.fft)
+
 
-# Save a PDF/PNG of the figure
-#fig.savefig(os.path.dirname(os.path.abspath(__file__)) + os.sep + w.type + '.pdf', dpi=None, format='pdf', bbox_inches='tight', pad_inches=0.1)
-#fig.savefig(os.path.dirname(os.path.abspath(__file__)) + os.sep + w.type + '.png', dpi=150, format='png', bbox_inches='tight', pad_inches=0.1)
 
-plt.show()