Merge pull request #11438 from jklymak/enh-get-gridspec

tacaswell · web-flow · commit 21b102ec4767 · 2018-07-07T16:27:30.000-04:00
ENH: add get_gridspec convenience method to subplots
diff --git a/doc/users/next_whats_new/subplot_get_gridspec.rst b/doc/users/next_whats_new/subplot_get_gridspec.rst
@@ -0,0 +1,20 @@
+Add ``ax.get_gridspec`` to `.SubplotBase`
+-----------------------------------------
+
+New method `.SubplotBase.get_gridspec` is added so that users can
+easily get the gridspec that went into making an axes:
+
+  .. code::
+
+    import matplotlib.pyplot as plt
+
+    fig, axs = plt.subplots(3, 2)
+    gs = axs[0, -1].get_gridspec()
+
+    # remove the last column
+    for ax in axs[:,-1].flatten():
+      ax.remove()
+
+    # make a subplot in last column that spans rows.
+    ax = fig.add_subplot(gs[:, -1])
+    plt.show()
diff --git a/examples/subplots_axes_and_figures/gridspec_and_subplots.py b/examples/subplots_axes_and_figures/gridspec_and_subplots.py
@@ -0,0 +1,24 @@
+"""
+==================================================
+Combining two subplots using subplots and GridSpec
+==================================================
+
+Sometimes we want to combine two subplots in an axes layout created with
+`~.Figure.subplots`.  We can get the `~.gridspec.GridSpec` from the axes
+and then remove the covered axes and fill the gap with a new bigger axes.
+Here we create a layout with the bottom two axes in the last column combined.
+
+See: :doc:`/tutorials/intermediate/gridspec`
+"""
+import matplotlib.pyplot as plt
+
+fig, axs = plt.subplots(ncols=3, nrows=3)
+gs = axs[1, 2].get_gridspec()
+# remove the underlying axes
+for ax in axs[1:, -1]:
+    ax.remove()
+axbig = fig.add_subplot(gs[1:, -1])
+axbig.annotate('Big Axes \nGridSpec[1:, -1]', (0.1, 0.5),
+               xycoords='axes fraction', va='center')
+
+fig.tight_layout()
diff --git a/lib/matplotlib/axes/_subplots.py b/lib/matplotlib/axes/_subplots.py
@@ -116,6 +116,10 @@ def set_subplotspec(self, subplotspec):
         """set the SubplotSpec instance associated with the subplot"""
         self._subplotspec = subplotspec
 
+    def get_gridspec(self):
+        """get the GridSpec instance associated with the subplot"""
+        return self._subplotspec.get_gridspec()
+
     def update_params(self):
         """update the subplot position from fig.subplotpars"""
 
diff --git a/lib/matplotlib/tests/test_subplots.py b/lib/matplotlib/tests/test_subplots.py
@@ -143,3 +143,9 @@ def test_subplots_offsettext():
     axes[1, 0].plot(x, x)
     axes[0, 1].plot(y, x)
     axes[1, 1].plot(y, x)
+
+
+def test_get_gridspec():
+    # ahem, pretty trivial, but...
+    fig, ax = plt.subplots()
+    assert ax.get_subplotspec().get_gridspec() == ax.get_gridspec()
diff --git a/tutorials/intermediate/gridspec.py b/tutorials/intermediate/gridspec.py
@@ -122,6 +122,22 @@
 
 fig5.tight_layout()
 
+############################################################################
+# The ``subplots`` and ``gridspec`` methods can be combined since it is
+# sometimes more convenient to make most of the subplots using ``subplots``
+# and then remove some and combine them.  Here we create a layout with
+# the bottom two axes in the last column combined.
+
+fig, axs = plt.subplots(ncols=3, nrows=3)
+gs = axs[1, 2].get_gridspec()
+# remove the underlying axes
+for ax in axs[1:, -1]:
+    ax.remove()
+axbig = fig.add_subplot(gs[1:, -1])
+axbig.annotate('Big Axes \nGridSpec[1:, -1]', (0.1, 0.5),
+               xycoords='axes fraction', va='center')
+
+fig.tight_layout()
 
 ###############################################################################
 # Fine Adjustments to a Gridspec Layout
