From 487ac8041ef63c3385c9fa80fb1943a96034ce7b Mon Sep 17 00:00:00 2001 From: Jody Klymak Date: Tue, 26 May 2020 08:29:09 -0700 Subject: [PATCH] ENH: reuse gridspec if possible --- .../next_whats_new/2020-05-26-cl-subplot.rst | 21 ++++++++ lib/matplotlib/gridspec.py | 33 ++++++++++-- lib/matplotlib/pyplot.py | 8 +-- lib/matplotlib/tests/test_figure.py | 15 ++++++ .../intermediate/constrainedlayout_guide.py | 51 +++++++------------ 5 files changed, 88 insertions(+), 40 deletions(-) create mode 100644 doc/users/next_whats_new/2020-05-26-cl-subplot.rst 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 new file mode 100644 index 000000000000..0031cc6656b6 --- /dev/null +++ b/doc/users/next_whats_new/2020-05-26-cl-subplot.rst @@ -0,0 +1,21 @@ +Subplot and subplot2grid can now work with constrained layout +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +``constrained_layout`` depends on a single ``GridSpec`` +for each logical layout on a figure. Previously, ``plt.subplot`` and +``plt.subplot2grid`` added a new ``GridSpec`` each time they were called and +were therefore incompatible with ``constrained_layout``. + +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. + +In contrast, mixing ``nrows`` and ``ncols`` will *not* work with +``constrained_lyaout``: ``plt.subplot(2, 2, 1)`` followed by +``plt.subplots(2, 1, 2)`` will still produce two gridspecs, and +``constrained_layout=True`` will give bad results. In order to get the +desired 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.subplot2grid((2, 2), (0, 1), rowspan=2)`` can be used. diff --git a/lib/matplotlib/gridspec.py b/lib/matplotlib/gridspec.py index 577a50b4db1c..cd07136beb80 100644 --- a/lib/matplotlib/gridspec.py +++ b/lib/matplotlib/gridspec.py @@ -204,6 +204,28 @@ 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() + if gs.get_geometry() == (nrows, 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() @@ -666,8 +688,7 @@ def _from_subplot_args(figure, args): raise ValueError( 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)): @@ -680,7 +701,6 @@ 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) @@ -688,11 +708,16 @@ def _from_subplot_args(figure, args): 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. + 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 index b7e21cb6a78a..ce8f88bd4624 100644 --- a/lib/matplotlib/pyplot.py +++ b/lib/matplotlib/pyplot.py @@ -1399,10 +1399,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/lib/matplotlib/tests/test_figure.py b/lib/matplotlib/tests/test_figure.py index 84e65dfa90fa..22ecafc4ce8e 100644 --- a/lib/matplotlib/tests/test_figure.py +++ b/lib/matplotlib/tests/test_figure.py @@ -778,3 +778,18 @@ def test_fail(self, x, match): def test_hashable_keys(self, fig_test, fig_ref): fig_test.subplot_mosaic([[object(), object()]]) fig_ref.subplot_mosaic([["A", "B"]]) + + +def test_reused_gridspec(): + """Test that these all use the same gridspec""" + fig = plt.figure() + ax1 = fig.add_subplot(3, 2, (3, 5)) + ax2 = fig.add_subplot(3, 2, 4) + ax3 = plt.subplot2grid((3, 2), (2, 1), colspan=2, fig=fig) + + gs1 = ax1.get_subplotspec().get_gridspec() + gs2 = ax2.get_subplotspec().get_gridspec() + gs3 = ax3.get_subplotspec().get_gridspec() + + assert gs1 == gs2 + assert gs1 == gs3 diff --git a/tutorials/intermediate/constrainedlayout_guide.py b/tutorials/intermediate/constrainedlayout_guide.py index 1975153db181..c60ef59be710 100644 --- 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