Modified and included arguments for spatial plotting functions and filled doc strings with missing variables (#136)

pabloitu · web-flow · commit 4a789d9ebe1d · 2021-08-05T07:16:18.000-07:00
- Basemap plotting: Added fontsize arg for grid labels. Added figsize arg. Added manual Zoom Tile level, in case a greater or lower resolution of the map is desired.
- Catalog plotting: Included modified basemap options. Added legend title, framealpha, label fontsize, title fontsize as arguments. Check if given mag_ticks args is list, array or tuple.
- Basemap plotting: Added title_size into correct function args. Added colormap title and ticks fontsize as argument
diff --git a/csep/utils/plots.py b/csep/utils/plots.py
@@ -553,24 +553,31 @@ def get_hist(x, mws, normed=True):
     if show:
         pyplot.show()
 
-def plot_basemap(basemap, extent, ax=None,  coastline=True, borders=False, linecolor='black', linewidth=True,
-                 grid=False, grid_labels=False, set_global=False, show=False, projection=ccrs.PlateCarree(), apprx=False,
-                 central_latitude=0.0):
+def plot_basemap(basemap, extent, ax=None, figsize=None, coastline=True, borders=False, tile_scaling='auto',
+                 set_global=False, projection=ccrs.PlateCarree(), apprx=False, central_latitude=0.0,
+                 linecolor='black', linewidth=True,
+                 grid=False, grid_labels=False, grid_fontsize=None,
+                 show=False):
     """ Wrapper function for multiple cartopy base plots, including access to standard raster webservices
 
      Args:
          basemap (str): Possible values are: stock_img, stamen_terrain, stamen_terrain-background, google-satellite, ESRI_terrain, ESRI_imagery, ESRI_relief, ESRI_topo, ESRI_terrain, or webservice link (see examples in :func:`csep.utils.plots._get_basemap`. Default is None
          extent (list):  [lon_min, lon_max, lat_min, lat_max]
-         show (bool): Flag if the figure is displayed
-         set_global (bool): Display the complete globe as basemap
+         ax (:class:`matplotlib.pyplot.ax`): Previously defined ax object
+         figsize (tuple): If no ax is provided, a tuple of floats can be provided to define figure size
          coastline (str): Flag to plot coastline. default True,
          borders (bool): Flag to plot country borders. default False,
-         linewidth (float): Line width of borders and coast lines. default 1.5,
+         tile_scaling (str/int): Zoom level (1-12) of the basemap tiles. If 'auto', is automatically derived from extent
+         set_global (bool): Display the complete globe as basemap
+         projection (:class:`cartopy.crs.Projection`): Projection to be used in the basemap
+         apprx (bool): If true, approximates transformation by setting aspect ratio of axes based on middle latitude
+         central_latitude (float): average latitude from plotting region
          linecolor (str): Color of borders and coast lines. default 'black',
+         linewidth (float): Line width of borders and coast lines. default 1.5,
          grid (bool): Draws a grid in the basemap
          grid_labels (bool): Annotate grid values
-         apprx (bool): If true, approximates transformation by setting aspect ratio of axes based on middle latitude
-         central_latitude (float): average latitude from plotting region
+         grid_fontsize (float): Font size of the grid x and y labels
+         show (bool): Flag if the figure is displayed
 
      Returns:
          :class:`matplotlib.pyplot.ax` object
@@ -579,14 +586,14 @@ def plot_basemap(basemap, extent, ax=None,  coastline=True, borders=False, linec
     if ax is None:
         if apprx:
             projection = ccrs.PlateCarree()
-            fig = pyplot.figure()
+            fig = pyplot.figure(figsize=figsize)
             ax = fig.add_subplot(111, projection=projection)
             # Set plot aspect according to local longitude-latitude ratio in metric units
             # (only compatible with plain PlateCarree "projection")
             LATKM = 110.574  # length of a ° of latitude [km]; constant --> ignores Earth's flattening
             ax.set_aspect(LATKM / (111.320 * numpy.cos(numpy.deg2rad(central_latitude))))
         else:
-            fig = pyplot.figure()
+            fig = pyplot.figure(figsize=figsize)
             ax = fig.add_subplot(111, projection=projection)
 
     if set_global:
@@ -600,7 +607,10 @@ def plot_basemap(basemap, extent, ax=None,  coastline=True, borders=False, linec
         tile_autoscaler = cartopy.feature.AdaptiveScaler(5, ((6, 50), (7, 15)))
         tiles = None
         # Set tile depth
-        tile_depth = 4 if set_global else tile_autoscaler.scale_from_extent(extent)
+        if tile_scaling == 'auto':
+            tile_depth = 4 if set_global else tile_autoscaler.scale_from_extent(extent)
+        else:
+            tile_depth = tile_scaling
         if coastline:
             ax.coastlines(color=linecolor, linewidth=linewidth)
         if borders:
@@ -621,6 +631,8 @@ def plot_basemap(basemap, extent, ax=None,  coastline=True, borders=False, linec
         gl = ax.gridlines(draw_labels=grid_labels, alpha=0.5)
         gl.right_labels = False
         gl.top_labels = False
+        gl.xlabel_style['fontsize'] = grid_fontsize
+        gl.ylabel_style['fontsize'] = grid_fontsize
         gl.xformatter = LONGITUDE_FORMATTER
         gl.yformatter = LATITUDE_FORMATTER
 
@@ -646,22 +658,25 @@ def plot_catalog(catalog, ax=None, show=False, extent=None, set_global=False, pl
            - :filename: :class:`str` - File to save figure. default None
            - :projection: :class:`cartopy.crs.Projection` - default :class:`cartopy.crs.PlateCarree`. Note: this can be
                 'fast' to apply an approximate transformation of axes.
+           - :basemap:  :class:`str`/:class:`None`. Possible values are: stock_img, stamen_terrain, stamen_terrain-background, google-satellite, ESRI_terrain, ESRI_imagery, ESRI_relief, ESRI_topo, ESRI_terrain, or webservice link. Default is None
+           - :coastline: :class:`bool` - Flag to plot coastline. default True,
            - :grid: :class:`bool` - default True
            - :grid_labels: :class:`bool` - default True
+           - :grid_fontsize: :class:`float` - default 10.0
            - :marker: :class:`str` - Marker type
            - :markersize: :class:`float` - Constant size for all earthquakes
            - :markercolor: :class:`str` - Color for all earthquakes
-           - :basemap:  :class:`str`/:class:`None`. Possible values are: stock_img, stamen_terrain, stamen_terrain-background, google-satellite, ESRI_terrain, ESRI_imagery, ESRI_relief, ESRI_topo, ESRI_terrain, or webservice link. Default is None
-           - :coastline: :class:`bool` - Flag to plot coastline. default True,
            - :borders: :class:`bool` - Flag to plot country borders. default False,
-           - :linewidth: :class:`float` - Line width of borders and coast lines. default 1.5,
-           - :linecolor: :class:`str` - Color of borders and coast lines. default 'black',
+           - :region_border: :class:`bool` - Flag to plot the catalog region border. default True,
            - :alpha: :class:`float` - Transparency for the earthquakes scatter
            - :mag_scale: :class:`float` - Scaling of the scatter
            - :legend: :class:`bool` - Flag to display the legend box
            - :legend_loc: :class:`int`/:class:`str` - Position of the legend
            - :mag_ticks: :class:`list` - Ticks to display in the legend
            - :labelspacing: :class:`int` - Separation between legend ticks
+           - :tile_scaling: :class:`str`/:class:`int`. Zoom level (1-12) of the basemap tiles. If 'auto', is automatically derived from extent
+           - :linewidth: :class:`float` - Line width of borders and coast lines. default 1.5,
+           - :linecolor: :class:`str` - Color of borders and coast lines. default 'black',
 
     Returns:
         :class:`matplotlib.pyplot.ax` object
@@ -682,17 +697,23 @@ def plot_catalog(catalog, ax=None, show=False, extent=None, set_global=False, pl
     alpha = plot_args.get('alpha', 1)
     mag_scale = plot_args.get('mag_scale', 1)
     legend = plot_args.get('legend', False)
+    legend_title = plot_args.get('legend_title', r"Magnitudes")
     legend_loc = plot_args.get('legend_loc', 1)
+    legend_framealpha = plot_args.get('legend_framealpha', None)
+    legend_fontsize = plot_args.get('legend_fontsize', None)
+    legend_titlesize = plot_args.get('legend_titlesize', None)
     mag_ticks = plot_args.get('mag_ticks', False)
     labelspacing = plot_args.get('labelspacing', 1)
     region_border = plot_args.get('region_border', True)
     # cartopy properties
     projection = plot_args.get('projection', ccrs.PlateCarree(central_longitude=0.0))
-    grid = plot_args.get('grid', True)
-    grid_labels = plot_args.get('grid_labels', False)
     basemap = plot_args.get('basemap', None)
     coastline = plot_args.get('coastline', True)
+    grid = plot_args.get('grid', True)
+    grid_labels = plot_args.get('grid_labels', False)
+    grid_fontsize = plot_args.get('grid_fontsize', False)
     borders = plot_args.get('borders', False)
+    tile_scaling = plot_args.get('tile_scaling', 'auto')
     linewidth = plot_args.get('linewidth', True)
     linecolor = plot_args.get('linecolor', 'black')
 
@@ -728,7 +749,7 @@ def plot_catalog(catalog, ax=None, show=False, extent=None, set_global=False, pl
         ax.set_extent(extents=extent, crs=ccrs.PlateCarree())  # Defined extent always in lat/lon
 
     # Basemap plotting
-    ax = plot_basemap(basemap, extent, ax=ax, coastline=coastline, borders=borders,
+    ax = plot_basemap(basemap, extent, ax=ax, coastline=coastline, borders=borders, tile_scaling=tile_scaling,
                       linecolor=linecolor, linewidth=linewidth, projection=projection, apprx=apprx,
                       central_latitude=central_latitude)
 
@@ -751,18 +772,17 @@ def size_map(markersize, values, scale):
 
     # Legend
     if legend:
-        if not mag_ticks:
-            mag_ticks = numpy.linspace(mw_range[0], mw_range[1], 4)
-        else:
+        if isinstance(mag_ticks, (tuple, list, numpy.ndarray)):
             if not numpy.all([ i >= mw_range[0] and i <= mw_range[1] for i in mag_ticks]):
                 print("Magnitude ticks do not lie within the catalog magnitude range")
-
+        elif mag_ticks is False:
+            mag_ticks = numpy.linspace(mw_range[0], mw_range[1], 4)
         handles, labels = scatter.legend_elements(prop="sizes",
                                                   num=list(size_map(markersize, mag_ticks, mag_scale)),
                                                   alpha=0.3)
         ax.legend(handles, numpy.round(mag_ticks, 1),
-                  loc=legend_loc, title=r"Magnitudes", title_fontsize=16,
-                  labelspacing=labelspacing, handletextpad=5, framealpha=False)
+                  loc=legend_loc, title=legend_title, fontsize=legend_fontsize, title_fontsize=legend_titlesize,
+                  labelspacing=labelspacing, handletextpad=5, framealpha=legend_framealpha)
 
     if region_border:
         try:
@@ -777,6 +797,8 @@ def size_map(markersize, values, scale):
         gl = ax.gridlines(draw_labels=grid_labels, alpha=0.5)
         gl.right_labels = False
         gl.top_labels = False
+        gl.xlabel_style['fontsize'] = grid_fontsize
+        gl.ylabel_style['fontsize'] = grid_fontsize
         gl.xformatter = LONGITUDE_FORMATTER
         gl.yformatter = LATITUDE_FORMATTER
 
@@ -808,14 +830,19 @@ def plot_spatial_dataset(gridded, region, ax=None, show=False, extent=None, set_
            - :projection: :class:`cartopy.crs.Projection` - default :class:`cartopy.crs.PlateCarree`
            - :grid: :class:`bool` - default True
            - :grid_labels: :class:`bool` - default True
+           - :grid_fontsize: :class:`float` - default 10.0
            - :basemap:  :class:`str`. Possible  values are: stock_img, stamen_terrain, stamen_terrain-background, google-satellite, ESRI_terrain, ESRI_imagery, ESRI_relief, ESRI_topo, ESRI_terrain, or webservice link. Default is None
            - :coastline: :class:`bool` - Flag to plot coastline. default True,
            - :borders: :class:`bool` - Flag to plot country borders. default False,
+           - :region_border: :class:`bool` - Flag to plot the dataset region border. default True,
+           - :tile_scaling: :class:`str`/:class:`int`. Zoom level (1-12) of the basemap tiles. If 'auto', is automatically derived from extent
            - :linewidth: :class:`float` - Line width of borders and coast lines. default 1.5,
            - :linecolor: :class:`str` - Color of borders and coast lines. default 'black',
            - :cmap: :class:`str`/:class:`pyplot.colors.Colormap` -  default 'viridis'
-           - :clabel: :class:`str` - default None
-           - :clim: :class:`list` - default None
+           - :clim: :class:`list` - Range of the colorbar. default None
+           - :clabel: :class:`str` - Label of the colorbar. default None
+           - :clabel_fontsize: :class:`float` - default None
+           - :cticks_fontsize: :class:`float` - default None
            - :alpha: :class:`float` - default 1
            - :alpha_exp: :class:`float` - Exponent for the alpha func (recommended between 0.4 and 1). default 0
 
@@ -841,16 +868,20 @@ def plot_spatial_dataset(gridded, region, ax=None, show=False, extent=None, set_
     projection = plot_args.get('projection', ccrs.PlateCarree(central_longitude=0.0))
     grid = plot_args.get('grid', True)
     grid_labels = plot_args.get('grid_labels', False)
+    grid_fontsize = plot_args.get('grid_fontsize', False)
     basemap = plot_args.get('basemap', None)
     coastline = plot_args.get('coastline', True)
     borders = plot_args.get('borders', False)
+    tile_scaling = plot_args.get('tile_scaling', 'auto')
     linewidth = plot_args.get('linewidth', True)
     linecolor = plot_args.get('linecolor', 'black')
     region_border = plot_args.get('region_border', True)
     # color bar properties
     cmap = plot_args.get('cmap', None)
-    clabel = plot_args.get('clabel', '')
     clim = plot_args.get('clim', None)
+    clabel = plot_args.get('clabel', None)
+    clabel_fontsize = plot_args.get('clabel_fontsize', None)
+    cticks_fontsize = plot_args.get('cticks_fontsize', None)
     alpha = plot_args.get('alpha', 1)
     alpha_exp = plot_args.get('alpha_exp', 0)
 
@@ -862,7 +893,7 @@ def plot_spatial_dataset(gridded, region, ax=None, show=False, extent=None, set_
         n_lats = len(region.ys) // 2
         central_latitude = region.ys[n_lats]
 
-    # Instantiage GeoAxes object
+    # Instantiate GeoAxes object
     if ax is None:
         fig = pyplot.figure(figsize=figsize)
         ax = fig.add_subplot(111, projection=projection)
@@ -877,7 +908,7 @@ def plot_spatial_dataset(gridded, region, ax=None, show=False, extent=None, set_
     # Basemap plotting
     ax = plot_basemap(basemap, extent, ax=ax, coastline=coastline, borders=borders,
                       linecolor=linecolor, linewidth=linewidth, projection=projection, apprx=apprx,
-                      central_latitude=central_latitude)
+                      central_latitude=central_latitude, tile_scaling=tile_scaling)
 
     ## Define colormap and transparency function
     if isinstance(cmap, str) or not cmap:
@@ -902,22 +933,25 @@ def plot_spatial_dataset(gridded, region, ax=None, show=False, extent=None, set_
     cax = fig.add_axes([ax.get_position().x1 + 0.01, ax.get_position().y0, 0.025, ax.get_position().height],
                        label='Colorbar')
     cbar = fig.colorbar(im, ax=ax, cax=cax)
-    cbar.set_label(clabel)
+    cbar.set_label(clabel, fontsize=clabel_fontsize)
+    cbar.ax.tick_params(labelsize=cticks_fontsize)
 
     # Gridline options
     if grid:
         gl = ax.gridlines(draw_labels=grid_labels, alpha=0.5)
         gl.right_labels = False
         gl.top_labels = False
+        gl.xlabel_style['fontsize'] = grid_fontsize
+        gl.ylabel_style['fontsize'] = grid_fontsize
         gl.xformatter = LONGITUDE_FORMATTER
         gl.yformatter = LATITUDE_FORMATTER
 
     if region_border:
         pts = region.tight_bbox()
-        ax.plot(pts[:,0], pts[:,1], lw=1, color='black', transform=ccrs.PlateCarree())
+        ax.plot(pts[:, 0], pts[:, 1], lw=1, color='black', transform=ccrs.PlateCarree())
 
     # matplotlib figure options
-    ax.set_title(title, y=1.06)
+    ax.set_title(title, y=1.06, fontsize=title_size)
     if filename is not None:
         ax.get_figure().savefig(filename + '.pdf')
         ax.get_figure().savefig(filename + '.png', dpi=300)
