ENH: reuse gridspec if possible

jklymak · jklymak · commit eeaf3e2e7533 · 2020-05-26T08:29:09.000-07:00
diff --git a/doc/users/next_whats_new/2020-05-26-cl-subplot.rst b/doc/users/next_whats_new/2020-05-26-cl-subplot.rst
@@ -0,0 +1,17 @@
+Constrained layout now supports subplot and subplot2grid
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Previously ``constrained_layout`` depended on a single ``GridSpec``
+for each logical layout on a figure, whereas ``plt.subplot`` added
+a new ``GridSpec`` for each call.  Now ``plt.subplot`` attempts to
+reuse the ``GridSpec`` if the number of rows and columns is the same
+as the top level gridspec already in the figure.  i.e. ``plt.subplot(2, 1, 2)``
+will use the same gridspec as ``plt.subplot(2, 1, 1)`` and the
+``constrained_layout=True`` option to `~.figure.Figure` will work.  Note that
+mixing ``nrows`` and ``ncols`` will *not* work: ``plt.subplot(2, 2, 1)``
+followed by ``plt.subplots(2, 1, 2)`` will produce two gridspecs, and
+``constrained_layout=True`` will give bad results.  In order to get the
+same effect, the second call can specify the cells the second axes is meant
+to cover:  ``plt.subplots(2, 2, (2, 4))``, or the more pythonic
+``plt.subplots2grid((2, 2), (0, 1), rowspan=2)`` can be used.
+
diff --git a/lib/matplotlib/gridspec.py b/lib/matplotlib/gridspec.py
@@ -202,6 +202,29 @@ def get_grid_positions(self, fig, raw=False):
         fig_lefts, fig_rights = (left + cell_ws).reshape((-1, 2)).T
         return fig_bottoms, fig_tops, fig_lefts, fig_rights
 
+    @staticmethod
+    def _check_gridspec_exists(figure, nrows, ncols):
+        """
+        Check if the figure already has a gridspec with these dimensions,
+        or create a new one
+        """
+        for ax in figure.get_axes():
+            if hasattr(ax, 'get_subplotspec'):
+                gs = ax.get_subplotspec().get_gridspec()
+                if hasattr(gs, 'get_topmost_subplotspec'):
+                    # This is needed for colorbar gridspec layouts.
+                    # This is probably OK becase this whole logic tree
+                    # is for when the user is doing simple things with the
+                    # add_subplot command.  For complicated layouts
+                    # like subgridspecs the proper gridspec is passed in...
+                    gs = gs.get_topmost_subplotspec().get_gridspec()
+                (nrow, ncol) = gs.get_geometry()
+                if nrow == nrows and ncol == ncols:
+                    return gs
+        # else gridspec not found:
+        return GridSpec(nrows, ncols, figure=figure)
+
+
     def __getitem__(self, key):
         """Create and return a `.SubplotSpec` instance."""
         nrows, ncols = self.get_geometry()
@@ -665,7 +688,7 @@ def _from_subplot_args(figure, args):
                         f"Single argument to subplot must be a three-digit "
                         f"integer, not {arg}") from None
                 # num - 1 for converting from MATLAB to python indexing
-                return GridSpec(rows, cols, figure=figure)[num - 1]
+                i = j = num
         elif len(args) == 3:
             rows, cols, num = args
             if not (isinstance(rows, Integral) and isinstance(cols, Integral)):
@@ -678,19 +701,24 @@ def _from_subplot_args(figure, args):
                     i, j = map(int, num)
                 else:
                     i, j = num
-                return gs[i-1:j]
             else:
                 if not isinstance(num, Integral):
                     cbook.warn_deprecated("3.3", message=message)
                     num = int(num)
                 if num < 1 or num > rows*cols:
                     raise ValueError(
                         f"num must be 1 <= num <= {rows*cols}, not {num}")
-                return gs[num - 1]   # -1 due to MATLAB indexing.
+                else:
+                    i = j = num
         else:
             raise TypeError(f"subplot() takes 1 or 3 positional arguments but "
                             f"{len(args)} were given")
 
+        gs = GridSpec._check_gridspec_exists(figure, rows, cols)
+        if gs is None:
+            gs = GridSpec(rows, cols, figure=figure)
+        return gs[i-1:j]
+
     # num2 is a property only to handle the case where it is None and someone
     # mutates num1.
 
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -1311,10 +1311,10 @@ def subplot2grid(shape, loc, rowspan=1, colspan=1, fig=None, **kwargs):
     if fig is None:
         fig = gcf()
 
-    s1, s2 = shape
-    subplotspec = GridSpec(s1, s2).new_subplotspec(loc,
-                                                   rowspan=rowspan,
-                                                   colspan=colspan)
+    rows, cols = shape
+    gs = GridSpec._check_gridspec_exists(fig, rows, cols)
+
+    subplotspec = gs.new_subplotspec(loc, rowspan=rowspan, colspan=colspan)
     ax = fig.add_subplot(subplotspec, **kwargs)
     bbox = ax.bbox
     axes_to_delete = []
diff --git a/tutorials/intermediate/constrainedlayout_guide.py b/tutorials/intermediate/constrainedlayout_guide.py
@@ -492,40 +492,43 @@ def docomplicated(suptitle=None):
 # Incompatible functions
 # ----------------------
 #
-# ``constrained_layout`` will not work on subplots created via
-# `.pyplot.subplot`.  The reason is that each call to `.pyplot.subplot` creates
-# a separate `.GridSpec` instance and ``constrained_layout`` uses (nested)
-# gridspecs to carry out the layout.  So the following fails to yield a nice
-# layout:
+# ``constrained_layout`` will work with `.pyplot.subplot`, but only if the
+# number of rows and columns is the same for each call.
+# The reason is that each call to `.pyplot.subplot` will create a new
+# `.GridSpec` instance if the geometry is not the same, and
+# ``constrained_layout``.  So the following works fine:
+
 
 fig = plt.figure()
 
-ax1 = plt.subplot(221)
-ax2 = plt.subplot(223)
-ax3 = plt.subplot(122)
+ax1 = plt.subplot(2, 2, 1)
+ax2 = plt.subplot(2, 2, 3)
+# third axes that spans both rows in second column:
+ax3 = plt.subplot(2, 2, (2, 4))
 
 example_plot(ax1)
 example_plot(ax2)
 example_plot(ax3)
+plt.suptitle('Homogenous nrows, ncols')
 
 ###############################################################################
-# Of course that layout is possible using a gridspec:
+# but the following leads to a poor layout:
 
 fig = plt.figure()
-gs = fig.add_gridspec(2, 2)
 
-ax1 = fig.add_subplot(gs[0, 0])
-ax2 = fig.add_subplot(gs[1, 0])
-ax3 = fig.add_subplot(gs[:, 1])
+ax1 = plt.subplot(2, 2, 1)
+ax2 = plt.subplot(2, 2, 3)
+ax3 = plt.subplot(1, 2, 2)
 
 example_plot(ax1)
 example_plot(ax2)
 example_plot(ax3)
+plt.suptitle('Mixed nrows, ncols')
 
 ###############################################################################
 # Similarly,
-# :func:`~matplotlib.pyplot.subplot2grid` doesn't work for the same reason:
-# each call creates a different parent gridspec.
+# :func:`~matplotlib.pyplot.subplot2grid` works with the same limitation
+# that nrows and ncols cannot change for the layout to look good.
 
 fig = plt.figure()
 
@@ -538,23 +541,7 @@ def docomplicated(suptitle=None):
 example_plot(ax2)
 example_plot(ax3)
 example_plot(ax4)
-
-###############################################################################
-# The way to make this plot compatible with ``constrained_layout`` is again
-# to use ``gridspec`` directly
-
-fig = plt.figure()
-gs = fig.add_gridspec(3, 3)
-
-ax1 = fig.add_subplot(gs[0, 0])
-ax2 = fig.add_subplot(gs[0, 1:])
-ax3 = fig.add_subplot(gs[1:, 0:2])
-ax4 = fig.add_subplot(gs[1:, -1])
-
-example_plot(ax1)
-example_plot(ax2)
-example_plot(ax3)
-example_plot(ax4)
+fig.suptitle('subplot2grid')
 
 ###############################################################################
 # Other Caveats
