New option --plot to save plots to files

Author: claudiodsf

CHANGELOG.md

@@ -24,6 +24,8 @@ to the latest version.
   frequency band specified in the config file
 - New option `--detailed` for `print_families`, to print more detailed
   information about each family, including the list of events
+- New option `--output` to save plots to files instead of showing them on
+  screen
 - Config option `waveform_data_path` renamed to `sds_data_path`
 - New config option `event_data_path` to specify the path to a local directory
   with waveform files organized per event

requake/config/parse_arguments.py

@@ -204,10 +204,23 @@ def parse_arguments(progname='requake'):
              'config file. Specify FREQ_MIN and FREQ_MAX in Hz.'
     )
     # ---
+    # --- parent parser for saving plots to file instead of showing
+    #     them on screen
+    output = argparse.ArgumentParser(add_help=False)
+    output.add_argument(
+        '-o', '--output', type=str, default='NO_OUTPUT', nargs='?',
+        metavar='FILE_OR_FORMAT',
+        help='Save the plot instead of displaying it. Provide a filename '
+             '(e.g., "plot.png") to specify the output name and format, or '
+             'just a format (e.g., "png", "pdf", "svg") to automatically '
+             'generate a default filename. If omitted, the plot is saved as '
+             'a PNG file with a default name.'
+    )
+    # ---
     # --- plot_pair
     plot_pair = subparser.add_parser(
         'plot_pair',
-        parents=[traceid, freq_band],
+        parents=[traceid, freq_band, output],
         help='plot traces for a given event pair'
     )
     plot_pair.add_argument('evid1')
@@ -309,37 +322,37 @@ def parse_arguments(progname='requake'):
     )
     # ---
     # --- print_families
-    printfamilies = subparser.add_parser(
+    print_families = subparser.add_parser(
         'print_families',
         parents=[
             longerthan, shorterthan, minevents, family_numbers, print_format
         ],
         help='print families to screen'
     )
-    printfamilies.add_argument(
+    print_families.add_argument(
         '-d', '--detailed', action='store_true',
         help='print detailed information for each family, including a list '
              'of events'
     )
     # ---
     # --- plot_families
-    plotfamilies = subparser.add_parser(
+    plot_families = subparser.add_parser(
         'plot_families',
         parents=[
             longerthan, shorterthan, minevents, family_numbers,
-            traceid, freq_band
+            traceid, freq_band, output
         ],
         help='plot traces for one ore more event families'
     )
-    plotfamilies.add_argument(
+    plot_families.add_argument(
         '-s', '--starttime', type=float, default=None,
         help='start time, in seconds relative to trace start, for the plot'
     )
-    plotfamilies.add_argument(
+    plot_families.add_argument(
         '-e', '--endtime', type=float, default=None,
         help='end time, in seconds relative to trace start, for the plot'
     )
-    plotfamilies.add_argument(
+    plot_families.add_argument(
         '-T', '--template', action='store_true',
         help='plot family members found with template scan'
     )
@@ -349,7 +362,7 @@ def parse_arguments(progname='requake'):
         'plot_timespans',
         parents=[
             longerthan, shorterthan, minevents, family_numbers, colorby,
-            colormap, color_range
+            colormap, color_range, output
         ],
         help='plot family timespans'
     )
@@ -369,7 +382,7 @@ def parse_arguments(progname='requake'):
         'plot_cumulative',
         parents=[
             longerthan, shorterthan, minevents, family_numbers, colorby,
-            colormap, color_range
+            colormap, color_range, output
         ],
         help='cumulative plot for one or more families'
     )
@@ -390,7 +403,7 @@ def parse_arguments(progname='requake'):
         'map_families',
         parents=[
             longerthan, shorterthan, minevents, family_numbers, colorby,
-            colormap, color_range
+            colormap, color_range, output
         ],
         help='plot families on a map'
     )

requake/plot/map_families.py

@@ -19,7 +19,8 @@
 from ..config import config, rq_exit
 from ..families import FamilyNotFoundError, read_selected_families
 from .plot_utils import (
-    plot_title, hover_annotation, duration_string, family_colors, plot_colorbar
+    plot_title, hover_annotation, duration_string, family_colors,
+    plot_colorbar, save_or_show_plot
 )
 from .cached_tiler import CachedTiler
 from .map_tiles import (
@@ -155,4 +156,4 @@ def map_families():
     annot.set_visible(False)
     annot.hover_annotation = True
     fig.canvas.mpl_connect('motion_notify_event', hover_annotation)
-    plt.show()
+    save_or_show_plot(fig, 'family_map')

requake/plot/plot_cumulative.py

@@ -19,7 +19,7 @@
 from ..formulas import mag_to_slip_in_cm, mag_to_moment
 from .plot_utils import (
     format_time_axis, plot_title, hover_annotation, duration_string,
-    family_colors, plot_colorbar
+    family_colors, plot_colorbar, save_or_show_plot
 )
 logger = logging.getLogger(__name__.rsplit('.', maxsplit=1)[-1])
 # Reduce logging level for Matplotlib to avoid DEBUG messages
@@ -174,4 +174,4 @@ def plot_cumulative():
     annot.set_visible(False)
     annot.hover_annotation = True
     fig.canvas.mpl_connect('motion_notify_event', hover_annotation)
-    plt.show()
+    save_or_show_plot(fig, 'family_cumulative')

requake/plot/plot_families.py

@@ -23,6 +23,7 @@
     get_family_aligned_waveforms_and_template
 )
 from ..waveforms import process_waveforms, NoWaveformError
+from .plot_utils import save_or_show_plot
 logger = logging.getLogger(__name__.rsplit('.', maxsplit=1)[-1])
 # Reduce logging level for Matplotlib to avoid DEBUG messages
 mpl_logger = logging.getLogger('matplotlib')
@@ -139,6 +140,10 @@ def _plot_family(family):
         f'{tr0.stats.freq_min:.1f}-{tr0.stats.freq_max:.1f} Hz'
     )
     ax.set_title(title, loc='right')
+    figure_name = (
+        f'family_{family.number}_{tr0.id}_'
+        f'{tr0.stats.freq_min:.1f}-{tr0.stats.freq_max:.1f}Hz'
+    )
 
     def _zoom_lines(zoom_level):
         for line in tracelines:
@@ -206,6 +211,7 @@ def _keypress(event):
     _keypress.time_zoom_level = 1
     _keypress.pan_amount = 0
     fig.canvas.mpl_connect('key_press_event', _keypress)
+    return fig, figure_name
 
 
 def plot_families():
@@ -226,13 +232,15 @@ def plot_families():
             'how to select specific families.')
         families = families[:20]
     for family in families:
-        _plot_family(family)
-    print('''
-    Use left/right arrow keys to scroll backwards/forward in time.
-    Use shift+left/shift+right to increase/decrease the time window.
-    Use up/down arrow keys to increase/decrease trace amplitude.
-    Press '0' to reset the view.
-    Press 'a' to show/hide theoretical arrivals.
-    Press 'q' to close a plot.
-    ''')
-    plt.show()
+        fig, figure_name = _plot_family(family)
+        save_or_show_plot(fig, figure_name, show=False)
+    if config.args.output == 'NO_OUTPUT':
+        print('''
+        Use left/right arrow keys to scroll backwards/forward in time.
+        Use shift+left/shift+right to increase/decrease the time window.
+        Use up/down arrow keys to increase/decrease trace amplitude.
+        Press '0' to reset the view.
+        Press 'a' to show/hide theoretical arrivals.
+        Press 'q' to close a plot.
+        ''')
+        plt.show()

requake/plot/plot_pair.py

@@ -18,6 +18,7 @@
     WaveformPair, process_waveforms, align_pair,
     NoWaveformError
 )
+from .plot_utils import save_or_show_plot
 logger = logging.getLogger(__name__.rsplit('.', maxsplit=1)[-1])
 # Reduce logging level for Matplotlib to avoid DEBUG messages
 mpl_logger = logging.getLogger('matplotlib')
@@ -124,4 +125,7 @@ def plot_pair():
         _ax.yaxis.set_tick_params(which='minor', bottom=False)
         _ax.grid(True)
         _ax.legend(loc='upper right')
-    plt.show()
+    plot_file_basename = (
+        f'{evid1}_{evid2}_{tr_id}_{freq_min:.1f}-{freq_max:.1f}Hz'
+    )
+    save_or_show_plot(fig, plot_file_basename)

requake/plot/plot_timespans.py

@@ -17,7 +17,7 @@
 from ..families import FamilyNotFoundError, read_selected_families
 from .plot_utils import (
     format_time_axis, plot_title, hover_annotation, duration_string,
-    family_colors, plot_colorbar
+    family_colors, plot_colorbar, save_or_show_plot
 )
 logger = logging.getLogger(__name__.rsplit('.', maxsplit=1)[-1])
 # Reduce logging level for Matplotlib to avoid DEBUG messages
@@ -138,4 +138,4 @@ def plot_timespans():
     annot.set_visible(False)
     annot.hover_annotation = True
     fig.canvas.mpl_connect('motion_notify_event', hover_annotation)
-    plt.show()
+    save_or_show_plot(fig, 'family_timespans')

requake/plot/plot_utils.py

@@ -11,8 +11,9 @@
 """
 import logging
 import contextlib
-import matplotlib.dates as mdates
 import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.dates as mdates
 from matplotlib import cm, colors
 from ..config import config
 from .colormaps import cmaps
@@ -349,3 +350,43 @@ def plot_colorbar(fig, ax, cmap, norm):
         format_time_axis(cbar.ax, which='yaxis', grid=False)
     with contextlib.suppress(AttributeError):
         cbar.ax.set_ylabel(cmap.label)
+
+
+def save_or_show_plot(fig, plot_file_basename, show=True):
+    """
+    Save or show the plot.
+
+    :param fig: The Matplotlib figure to save or show.
+    :param plot_file_basename: The base filename for the plot.
+        Only used if the 'output' command line argument contains just the
+        file format (e.g. 'png', 'pdf') or nothing is specified.
+    :param show: Whether to show the plot or not.
+    :type show: bool
+    """
+    output = config.args.output
+    if output == 'NO_OUTPUT':
+        if show:
+            plt.show()
+        return
+    if output is None:
+        plot_format = 'png'
+    elif output.lower() in ('pdf', 'svg', 'png', 'jpg', 'jpeg'):
+        plot_format = output.lower()
+    else:
+        plot_format = None
+    if plot_format is None:
+        plot_file_name = output
+    else:
+        plot_file_name = f'{plot_file_basename}.{plot_format}'
+        # sanitize plot filename
+        plot_file_name = plot_file_name\
+            .replace(' ', '_')\
+            .replace(':', '-')\
+            .replace('/', '-')
+    fig.savefig(
+        plot_file_name,
+        dpi=300,
+        bbox_inches='tight',
+    )
+    print(f'Plot saved as {plot_file_name}')
+    plt.close(fig)