Experimental subplot implementation that mimics matplotlib's style

Experimenting with a more matplotlib-like syntax for managing subplots, namely the `fig, axs = plt.subplots()` style of defining subplots! A high-level `subplots` function is defined in pygmtplots.py, that wraps around the new class SubPlot (inherited from the main Figure class to include subplot functionality). Re-aliased F to figsize (previously called dimensions). Also updated tests and tutorials to reference this new syntax.

Author: weiji14

doc/api/index.rst

@@ -16,6 +16,7 @@ All plotting is handled through the :class:`pygmt.Figure` class and its methods.
     :toctree: generated
 
     Figure
+    subplots
 
 Plotting data and laying out the map:
 
@@ -34,7 +35,6 @@ Plotting data and laying out the map:
     Figure.logo
     Figure.image
     Figure.shift_origin
-    Figure.subplot
     Figure.text
     Figure.meca
 

examples/tutorials/subplots.py

@@ -6,59 +6,58 @@
 you'll need to put many individual plots into one large figure, and label them
 'abcd'. These individual plots are called subplots.
 
-There are two main ways to handle subplots in GMT:
+There are two main ways to create subplots in GMT:
 
 - Use :meth:`pygmt.Figure.shift_origin` to manually move each individual plot
   to the right position.
-- Use :meth:`pygmt.Figure.subplot` to define the layout of the subplots.
+- Use :meth:`pygmt.subplots` to define the layout of the subplots.
 
 The first method is easier to use and should handle simple cases involving a
 couple of subplots. For more advanced subplot layouts however, we recommend the
-use of :meth:`pygmt.Figure.subplot` which offers finer grained control, and
-this is what the tutorial below will cover.
+use of :meth:`pygmt.subplots` which offers finer grained control, and this is
+what the tutorial below will cover.
 """
 
 ###############################################################################
-# Let's start by importing the PyGMT library and initiating a figure.
+# Let's start by importing the PyGMT library
 
 import pygmt
 
-fig = pygmt.Figure()
-
 ###############################################################################
 # Define subplot layout
 # ---------------------
 #
-# The ``fig.subplot(directive="begin")`` command is used to setup the layout,
-# size, and other attributes of the figure. It divides the whole canvas into
-# regular grid areas with n rows and m columns. Each grid area can contain an
-# individual subplot. For example:
+# The ``pygmt.subplots`` command is used to setup the layout, size, and other
+# attributes of the figure. It divides the whole canvas into regular grid areas
+# with n rows and m columns. Each grid area can contain an individual subplot.
+# For example:
 
-fig.subplot(directive="begin", row=2, col=3, dimensions="s5c/3c", frame="lrtb")
+fig, axs = pygmt.subplots(nrows=2, ncols=3, figsize=("15c", "6c"), frame="lrtb")
 
 ###############################################################################
 # will define our figure to have a 2 row and 3 column grid layout.
-# ``dimensions="s5c/3c"`` specifies that each 's'ubplot will have a width of
-# 5cm and height of 3cm. Alternatively, you can set ``dimensions="f15c/6c"`` to
-# define the overall size of the 'f'igure to be 15cm wide by 6cm high. Using
-# ``frame="lrtb"`` allows us to customize the map frame for all subplots. The
-# figure layout will look like the following:
-
-for index in range(2 * 3):
-    i = index // 3  # row
-    j = index % 3  # column
-    fig.subplot(directive="set", row=i, col=j)
+# ``figsize=("15c", "6c")`` defines the overall size of the figure to be 15cm
+# wide by 6cm high. Using ``frame="lrtb"`` allows us to customize the map frame
+# for all subplots instead of setting them individually. The figure layout will
+# look like the following:
+
+for index in axs.flatten():
+    i = index // axs.shape[1]  # row
+    j = index % axs.shape[1]  # column
+    fig.sca(ax=axs[i, j])  # sets the current Axes
     fig.text(
         x=0.5, y=0.5, text=f"index: {index}, row: {i}, col: {j}", region=[0, 1, 0, 1],
     )
-fig.subplot(directive="end")
+fig.end_subplot()
 fig.show()
 
 ###############################################################################
-# The ``fig.subplot(directive="set")`` command activates a specified subplot,
-# and all subsequent plotting commands will take place in that subplot. In
-# order to specify a subplot, you will need to know the identifier for each
-# subplot. This can be done by setting the ``row`` and ``col`` arguments.
+# The ``fig.sca`` command activates a specified subplot, and all subsequent
+# plotting commands will take place in that subplot. This is similar to
+# matplotlib's ``plt.sca`` method. In order to specify a subplot, you will need
+# to provide the identifier for that subplot via the ``ax`` argument. This can
+# be found in the ``axs`` variable referenced by the ``row`` and ``col``
+# number.
 
 ###############################################################################
 # .. note::
@@ -68,20 +67,19 @@
 #     numbers will go from 0 to M-1.
 
 ###############################################################################
-# For example, to activate the subplot on the top right corner (index: 2) so
-# that all subsequent plotting commands happen there, you can use the following
-# command:
+# For example, to activate the subplot on the top right corner (index: 2) at
+# ``row=0`` and ``col=2``, so that all subsequent plotting commands happen
+# there, you can use the following command:
 
 ###############################################################################
 # .. code-block:: default
 #
-#     fig.subplot(directive="set", row=0, col=2)
+#     fig.sca(ax=axs[0, 2])
 
 ###############################################################################
-# Finally, remember to use ``fig.subplot(directive="end")`` to exit the subplot
-# mode.
+# Finally, remember to use ``fig.end_subplot()`` to exit the subplot mode.
 
 ###############################################################################
 # .. code-block:: default
 #
-#     fig.subplot(directive="end")
+#     fig.end_subplot()

pygmt/__init__.py

@@ -13,12 +13,13 @@
 
 # Import modules to make the high-level GMT Python API
 from .session_management import begin as _begin, end as _end
-from .figure import Figure
+from .figure import Figure, SubPlot
 from .filtering import blockmedian
 from .gridding import surface
 from .sampling import grdtrack
 from .mathops import makecpt
 from .modules import GMTDataArrayAccessor, config, info, grdinfo, which
+from .pygmtplot import subplots
 from .gridops import grdcut
 from .x2sys import x2sys_init, x2sys_cross
 from . import datasets

pygmt/base_plotting.py

@@ -883,59 +883,6 @@ def legend(self, spec=None, position="JTR+jTR+o0.2c", box="+gwhite+p1p", **kwarg
             arg_str = " ".join([specfile, build_arg_string(kwargs)])
             lib.call_module("legend", arg_str)
 
-    @fmt_docstring
-    @use_alias(F="dimensions", B="frame")
-    def subplot(self, directive: str, row: int = None, col: int = None, **kwargs):
-        """
-        Manage modern mode figure subplot configuration and selection.
-
-        The subplot module is used to split the current figure into a
-        rectangular layout of subplots that each may contain a single
-        self-contained figure. A subplot setup is started with the begin
-        directive that defines the layout of the subplots, while positioning to
-        a particular subplot for plotting is done via the set directive. The
-        subplot process is completed via the end directive.
-
-        Full option list at :gmt-docs:`subplot.html`
-
-        {aliases}
-
-        Parameters
-        ----------
-        directive : str
-            Either 'begin', 'set' or 'end'.
-        row : int
-            The number of rows if using the 'begin' directive, or the row
-            number if using the 'set' directive. First row is 0, not 1.
-        col : int
-            The number of columns if using the 'begin' directive, or the column
-            number if using the 'set' directive. First column is 0, not 1.
-        dimensions : str
-            ``[f|s]width(s)/height(s)[+fwfracs/hfracs][+cdx/dy][+gfill][+ppen]
-            [+wpen]``
-            Specify the dimensions of the figure when using the 'begin'
-            directive. There are two different ways to do this: (f) Specify
-            overall figure dimensions or (s) specify the dimensions of a single
-            subplot.
-        """
-        if directive not in ("begin", "set", "end"):
-            raise GMTInvalidInput(
-                f"Unrecognized subplot directive '{directive}',\
-                should be either 'begin', 'set', or 'end'"
-            )
-
-        with Session() as lib:
-            rowcol = ""  # default is blank, e.g. when directive == "end"
-            if row is not None and col is not None:
-                if directive == "begin":
-                    rowcol = f"{row}x{col}"
-                elif directive == "set":
-                    rowcol = f"{row},{col}"
-            arg_str = " ".join(
-                a for a in [directive, rowcol, build_arg_string(kwargs)] if a
-            )
-            lib.call_module(module="subplot", args=arg_str)
-
     @fmt_docstring
     @use_alias(R="region", J="projection", B="frame")
     @use_alias(

pygmt/figure.py

@@ -374,3 +374,78 @@ def _repr_html_(self):
         base64_png = base64.encodebytes(raw_png)
         html = '<img src="data:image/png;base64,{image}" width="{width}px">'
         return html.format(image=base64_png.decode("utf-8"), width=500)
+
+
+class SubPlot(Figure):
+    """
+    Manage modern mode figure subplot configuration and selection.
+
+    The subplot module is used to split the current figure into a
+    rectangular layout of subplots that each may contain a single
+    self-contained figure. A subplot setup is started with the begin
+    directive that defines the layout of the subplots, while positioning to
+    a particular subplot for plotting is done via the set directive. The
+    subplot process is completed via the end directive.
+
+    Full option list at :gmt-docs:`subplot.html`
+    """
+
+    def __init__(self, nrows, ncols, figsize, **kwargs):
+        super().__init__()
+        # Activate main Figure, and initiate subplot
+        self._activate_figure()
+        self.begin_subplot(row=nrows, col=ncols, figsize=figsize, **kwargs)
+
+    @fmt_docstring
+    @use_alias(Ff="figsize", B="frame")
+    @kwargs_to_strings(Ff="sequence")
+    def begin_subplot(self, row=None, col=None, **kwargs):
+        """
+        The begin directive of subplot defines the layout of the entire
+        multi-panel illustration. Several options are available to specify
+        the systematic layout, labeling, dimensions, and more for the
+        subplots.
+
+        {aliases}
+        """
+        arg_str = " ".join(["begin", f"{row}x{col}", build_arg_string(kwargs)])
+        with Session() as lib:
+            lib.call_module(module="subplot", args=arg_str)
+
+    @fmt_docstring
+    @use_alias(F="dimensions")
+    def sca(self, ax=None, **kwargs):
+        """
+        Set the current Axes instance to *ax*.
+
+        Before you start plotting you must first select the active subplot.
+        Note: If any projection (J) option is passed with ? as scale or
+        width when plotting subplots, then the dimensions of the map are
+        automatically determined by the subplot size and your region. For
+        Cartesian plots: If you want the scale to apply equally to both
+        dimensions then you must specify ``projection="x"`` [The default
+        ``projection="X"`` will fill the subplot by using unequal scales].
+
+        {aliases}
+        """
+        arg_str = " ".join(["set", f"{ax}", build_arg_string(kwargs)])
+        with Session() as lib:
+            lib.call_module(module=f"subplot", args=arg_str)
+
+    @fmt_docstring
+    @use_alias(V="verbose")
+    def end_subplot(self, **kwargs):
+        """
+        This command finalizes the current subplot, including any placement
+        of tags, and updates the gmt.history to reflect the dimensions and
+        linear projection required to draw the entire figure outline. This
+        allows subsequent commands, such as colorbar, to use
+        ``position="J"`` to place bars with reference to the complete
+        figure dimensions. We also reset the current plot location to where
+        it was prior to the subplot.
+
+        {aliases}
+        """
+        arg_str = " ".join(["end", build_arg_string(kwargs)])
+        with Session() as lib:
+            lib.call_module(module="subplot", args=arg_str)

pygmt/pygmtplot.py

@@ -0,0 +1,39 @@
+import numpy as np
+
+from .figure import SubPlot
+
+
+def subplots(nrows=1, ncols=1, figsize=(6.4, 4.8), **kwargs):
+    """
+    Create a figure with a set of subplots.
+
+    Parameters
+    ----------
+    nrows : int
+        Number of rows of the subplot grid.
+
+    ncols : int
+        Number of columns of the subplot grid.
+
+    figsize : tuple
+        Figure dimensions as ``(width, height)``.
+
+    Returns
+    -------
+    fig : :class:`pygmt.Figure`
+        A PyGMT Figure instance.
+
+    axs : numpy.ndarray
+        Array of Axes objects.
+    """
+    # Get PyGMT Figure with SubPlot initiated
+    fig = SubPlot(nrows=nrows, ncols=ncols, figsize=figsize, **kwargs)
+
+    # Setup matplotlib-like Axes
+    axs = np.empty(shape=(nrows, ncols), dtype=object)
+    for index in range(nrows * ncols):
+        i = index // ncols  # row
+        j = index % ncols  # column
+        axs[i, j] = index
+
+    return fig, axs

pygmt/tests/test_subplot.py

@@ -3,22 +3,20 @@
 """
 import pytest
 
-from .. import Figure
-from ..exceptions import GMTInvalidInput
+from ..pygmtplot import subplots
 
 
 @pytest.mark.mpl_image_compare
 def test_subplot_basic():
     """
     Create a subplot figure with 1 row and 2 columns.
     """
-    fig = Figure()
-    fig.subplot(directive="begin", row=1, col=2, dimensions="f6c/3c")
-    fig.subplot(directive="set", row=0, col=0)
+    fig, axs = subplots(nrows=1, ncols=2, figsize=("6c", "3c"))
+    fig.sca(ax=axs[0, 0])
     fig.basemap(region=[0, 3, 0, 3], frame=True)
-    fig.subplot(directive="set", row=0, col=1)
+    fig.sca(ax=axs[0, 1])
     fig.basemap(region=[0, 3, 0, 3], frame=True)
-    fig.subplot(directive="end")
+    fig.end_subplot()
     return fig
 
 
@@ -27,20 +25,10 @@ def test_subplot_frame():
     """
     Check that map frame setting is applied to all subplot figures
     """
-    fig = Figure()
-    fig.subplot(directive="begin", row=1, col=2, dimensions="f6c/3c", frame="WSne")
-    fig.subplot(directive="set", row=0, col=0)
+    fig, axs = subplots(nrows=1, ncols=2, figsize=("6c", "3c"), frame="WSne")
+    fig.sca(ax=axs[0, 0])
     fig.basemap(region=[0, 3, 0, 3], frame="+tplot0")
-    fig.subplot(directive="set", row=0, col=1)
+    fig.sca(ax=axs[0, 1])
     fig.basemap(region=[0, 3, 0, 3], frame="+tplot1")
-    fig.subplot(directive="end")
+    fig.end_subplot()
     return fig
-
-
-def test_subplot_incorrect_directive():
-    """
-    Check that subplot fails when an incorrect directive is used
-    """
-    fig = Figure()
-    with pytest.raises(GMTInvalidInput):
-        fig.subplot(directive="start")