ENH: Add gridspec method to figure, and subplotspecs

jklymak · jklymak · commit 2977f37d46bf · 2018-07-16T09:10:49.000-07:00
diff --git a/doc/users/whats_new.rst b/doc/users/whats_new.rst
@@ -15,7 +15,7 @@ revision, see the :ref:`github-stats`.
    For a release, add a new section after this, then comment out the include
    and toctree below by indenting them. Uncomment them after the release.
 
-   .. include:: next_whats_new/README.rst
+  .. include:: next_whats_new/README.rst
    .. toctree::
       :glob:
       :maxdepth: 1
@@ -190,6 +190,24 @@ specify a number that is close (i.e. ``ax.title.set_position(0.5, 1.01)``)
 and the title will not be moved via this algorithm.
 
 
+New convenience methods for GridSpec
+------------------------------------
+
+There are new convenience methods for `.gridspec.GridSpec` and
+`.gridspec.GridSpecFromSubplotSpec`.  Instead of the former we can
+now call `.Figure.add_gridspec` and for the latter `.SubplotSpec.subgridspec`.
+
+.. code-block:: python
+
+    import matplotlib.pyplot as plt
+
+    fig = plt.figure()
+    gs0 = fig.add_gridspec(3, 1)
+    ax1 = fig.add_subplot(gs0[0])
+    ax2 = fig.add_subplot(gs0[1])
+    gssub = gs0[2].subgridspec(1, 3)
+    for i in range(3):
+        fig.add_subplot(gssub[0, i])
 
 
 
diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
@@ -394,6 +394,9 @@ def __init__(self,
         self._align_xlabel_grp = cbook.Grouper()
         self._align_ylabel_grp = cbook.Grouper()
 
+        # list of child gridspecs for this figure
+        self._gridspecs = []
+
     # TODO: I'd like to dynamically add the _repr_html_ method
     # to the figure in the right context, but then IPython doesn't
     # use it, for some reason.
@@ -1480,6 +1483,7 @@ def subplots(self, nrows=1, ncols=1, sharex=False, sharey=False,
         else:
             # this should turn constrained_layout off if we don't want it
             gs = GridSpec(nrows, ncols, figure=None, **gridspec_kw)
+        self._gridspecs.append(gs)
 
         # Create array to hold all axes.
         axarr = np.empty((nrows, ncols), dtype=object)
@@ -2474,6 +2478,49 @@ def align_labels(self, axs=None):
         self.align_xlabels(axs=axs)
         self.align_ylabels(axs=axs)
 
+    def add_gridspec(self, nrows, ncols, **kwargs):
+        """
+        Return a `.GridSpec` that has this figure as a parent.  This allows
+        complex layout of axes in the figure.
+
+        Parameters
+        ----------
+        nrows : int
+            Number of rows in grid.
+
+        ncols : int
+            Number or columns in grid.
+
+        Returns
+        -------
+        gridspec : `.GridSpec`
+
+        Other Parameters
+        ----------------
+        *kwargs* are passed to `.GridSpec`.
+
+        See Also
+        --------
+        matplotlib.pyplot.subplots
+
+        Examples
+        --------
+        Adding a subplot that spans two rows::
+
+            fig = plt.figure()
+            gs = fig.add_gridspec(2, 2)
+            ax1 = fig.add_subplot(gs[0, 0])
+            ax2 = fig.add_subplot(gs[1, 0])
+            # spans two rows:
+            ax3 = fig.add_subplot(gs[:, 1])
+
+        """
+
+        _ = kwargs.pop('figure', None)  # pop in case user has added this...
+        gs = GridSpec(nrows=nrows, ncols=ncols, figure=self, **kwargs)
+        self._gridspecs.append(gs)
+        return gs
+
 
 def figaspect(arg):
     """
diff --git a/lib/matplotlib/gridspec.py b/lib/matplotlib/gridspec.py
@@ -498,3 +498,44 @@ def __eq__(self, other):
 
     def __hash__(self):
         return hash((self._gridspec, self.num1, self.num2))
+
+    def subgridspec(self, nrows, ncols, **kwargs):
+        """
+        Return a `.GridSpecFromSubplotSpec` that has this subplotspec as
+        a parent.
+
+        Parameters
+        ----------
+        nrows : int
+            Number of rows in grid.
+
+        ncols : int
+            Number or columns in grid.
+
+        Returns
+        -------
+        gridspec : `.GridSpec`
+
+        Other Parameters
+        ----------------
+        **kwargs
+            All other parameters are passed to `.GridSpec`.
+
+        See Also
+        --------
+        matplotlib.pyplot.subplots
+
+        Examples
+        --------
+        Adding three subplots in the space occupied by a single subplot::
+
+            fig = plt.figure()
+            gs0 = fig.add_gridspec(3, 1)
+            ax1 = fig.add_subplot(gs0[0])
+            ax2 = fig.add_subplot(gs0[1])
+            gssub = gs0[2].subgridspec(1, 3)
+            for i in range(3):
+                fig.add_subplot(gssub[0, i])
+        """
+
+        return GridSpecFromSubplotSpec(nrows, ncols, self, **kwargs)
diff --git a/lib/matplotlib/tests/test_constrainedlayout.py b/lib/matplotlib/tests/test_constrainedlayout.py
@@ -94,9 +94,9 @@ def test_constrained_layout5():
 def test_constrained_layout6():
     'Test constrained_layout for nested gridspecs'
     fig = plt.figure(constrained_layout=True)
-    gs = gridspec.GridSpec(1, 2, figure=fig)
-    gsl = gridspec.GridSpecFromSubplotSpec(2, 2, gs[0])
-    gsr = gridspec.GridSpecFromSubplotSpec(1, 2, gs[1])
+    gs = fig.add_gridspec(1, 2, figure=fig)
+    gsl = gs[0].subgridspec(2, 2)
+    gsr = gs[1].subgridspec(1, 2)
     axsl = []
     for gs in gsl:
         ax = fig.add_subplot(gs)
diff --git a/tutorials/intermediate/constrainedlayout_guide.py b/tutorials/intermediate/constrainedlayout_guide.py
@@ -78,7 +78,7 @@ def example_plot(ax, fontsize=12, nodec=False):
         ax.set_yticklabels('')
 
 
-fig, ax = plt.subplots()
+fig, ax = plt.subplots(constrained_layout=False)
 example_plot(ax, fontsize=24)
 
 ###############################################################################
@@ -334,8 +334,10 @@ def example_plot(ax, fontsize=12, nodec=False):
 # with :func:`~matplotlib.figure.Figure.subplots` or
 # :func:`~matplotlib.gridspec.GridSpec` and
 # :func:`~matplotlib.figure.Figure.add_subplot`.
+#
+# Note that in what follows ``constrained_layout=True``
 
-fig = plt.figure(constrained_layout=True)
+fig = plt.figure()
 
 gs1 = gridspec.GridSpec(2, 1, figure=fig)
 ax1 = fig.add_subplot(gs1[0])
@@ -345,20 +347,21 @@ def example_plot(ax, fontsize=12, nodec=False):
 example_plot(ax2)
 
 ###############################################################################
-# More complicated gridspec layouts are possible.
+# More complicated gridspec layouts are possible.  Note here we use the
+# convenenience functions ``add_gridspec`` and ``subgridspec``
 
-fig = plt.figure(constrained_layout=True)
+fig = plt.figure()
 
-gs0 = gridspec.GridSpec(1, 2, figure=fig)
+gs0 = fig.add_gridspec(1, 2)
 
-gs1 = gridspec.GridSpecFromSubplotSpec(2, 1, gs0[0])
+gs1 = gs0[0].subgridspec(2, 1)
 ax1 = fig.add_subplot(gs1[0])
 ax2 = fig.add_subplot(gs1[1])
 
 example_plot(ax1)
 example_plot(ax2)
 
-gs2 = gridspec.GridSpecFromSubplotSpec(3, 1, gs0[1])
+gs2 = gs0[1].subgridspec(3, 1)
 
 for ss in gs2:
     ax = fig.add_subplot(ss)
@@ -373,9 +376,9 @@ def example_plot(ax, fontsize=12, nodec=False):
 # extent.  If we want the top and bottom of the two grids to line up then
 # they need to be in the same gridspec:
 
-fig = plt.figure(constrained_layout=True)
+fig = plt.figure()
 
-gs0 = gridspec.GridSpec(6, 2, figure=fig)
+gs0 = fig.add_gridspec(6, 2)
 
 ax1 = fig.add_subplot(gs0[:3, 0])
 ax2 = fig.add_subplot(gs0[3:, 0])
@@ -398,10 +401,10 @@ def example_plot(ax, fontsize=12, nodec=False):
 
 
 def docomplicated(suptitle=None):
-    fig = plt.figure(constrained_layout=True)
-    gs0 = gridspec.GridSpec(1, 2, figure=fig, width_ratios=[1., 2.])
-    gsl = gridspec.GridSpecFromSubplotSpec(2, 1, gs0[0])
-    gsr = gridspec.GridSpecFromSubplotSpec(2, 2, gs0[1])
+    fig = plt.figure()
+    gs0 = fig.add_gridspec(1, 2, figure=fig, width_ratios=[1., 2.])
+    gsl = gs0[0].subgridspec(2, 1)
+    gsr = gs0[1].subgridspec(2, 2)
 
     for gs in gsl:
         ax = fig.add_subplot(gs)
@@ -430,7 +433,7 @@ def docomplicated(suptitle=None):
 # no effect on it anymore. (Note that constrained_layout still leaves the
 # space for the axes that is moved).
 
-fig, axs = plt.subplots(1, 2, constrained_layout=True)
+fig, axs = plt.subplots(1, 2)
 example_plot(axs[0], fontsize=12)
 axs[1].set_position([0.2, 0.2, 0.4, 0.4])
 
@@ -444,7 +447,7 @@ def docomplicated(suptitle=None):
 
 from matplotlib.transforms import Bbox
 
-fig, axs = plt.subplots(1, 2, constrained_layout=True)
+fig, axs = plt.subplots(1, 2)
 example_plot(axs[0], fontsize=12)
 fig.execute_constrained_layout()
 # put into data-space:
@@ -468,7 +471,7 @@ def docomplicated(suptitle=None):
 # to yield a nice layout:
 
 
-fig = plt.figure(constrained_layout=True)
+fig = plt.figure()
 
 ax1 = plt.subplot(221)
 ax2 = plt.subplot(223)
@@ -481,8 +484,8 @@ def docomplicated(suptitle=None):
 ###############################################################################
 # Of course that layout is possible using a gridspec:
 
-fig = plt.figure(constrained_layout=True)
-gs = gridspec.GridSpec(2, 2, figure=fig)
+fig = plt.figure()
+gs = fig.add_gridspec(2, 2)
 
 ax1 = fig.add_subplot(gs[0, 0])
 ax2 = fig.add_subplot(gs[1, 0])
@@ -497,7 +500,7 @@ def docomplicated(suptitle=None):
 # :func:`~matplotlib.pyplot.subplot2grid` doesn't work for the same reason:
 # each call creates a different parent gridspec.
 
-fig = plt.figure(constrained_layout=True)
+fig = plt.figure()
 
 ax1 = plt.subplot2grid((3, 3), (0, 0))
 ax2 = plt.subplot2grid((3, 3), (0, 1), colspan=2)
@@ -513,8 +516,8 @@ def docomplicated(suptitle=None):
 # The way to make this plot compatible with ``constrained_layout`` is again
 # to use ``gridspec`` directly
 
-fig = plt.figure(constrained_layout=True)
-gs = gridspec.GridSpec(3, 3, figure=fig)
+fig = plt.figure()
+gs = fig.add_gridspec(3, 3)
 
 ax1 = fig.add_subplot(gs[0, 0])
 ax2 = fig.add_subplot(gs[0, 1:])
diff --git a/tutorials/intermediate/gridspec.py b/tutorials/intermediate/gridspec.py
