Backport PR #11648: FIX: colorbar placement in constrained layout

ImportanceOfBeingErnest · MeeseeksDev[bot] · commit f4baa7d9f666 · 2018-10-02T21:15:12.000Z
diff --git a/.flake8 b/.flake8
@@ -95,7 +95,7 @@ per-file-ignores =
     tutorials/colors/colormaps.py: E501
     tutorials/colors/colors.py: E402
     tutorials/intermediate/artists.py: E402, E501
-    tutorials/intermediate/constrainedlayout_guide.py: E402, E501
+    tutorials/intermediate/constrainedlayout_guide.py: E402
     tutorials/intermediate/gridspec.py: E402, E501
     tutorials/intermediate/legend_guide.py: E402, E501
     tutorials/intermediate/tight_layout_guide.py: E402, E501
diff --git a/examples/subplots_axes_and_figures/colorbar_placement.py b/examples/subplots_axes_and_figures/colorbar_placement.py
@@ -0,0 +1,54 @@
+"""
+=================
+Placing Colorbars
+=================
+
+Colorbars indicate the quantitative extent of image data.  Placing in
+a figure is non-trivial because room needs to be made for them.
+
+The simplest case is just attaching a colorbar to each axes:
+"""
+import matplotlib.pyplot as plt
+import numpy as np
+
+fig, axs = plt.subplots(2, 2)
+cm = ['RdBu_r', 'viridis']
+for col in range(2):
+    for row in range(2):
+        ax = axs[row, col]
+        pcm = ax.pcolormesh(np.random.random((20, 20)) * (col + 1),
+                            cmap=cm[col])
+        fig.colorbar(pcm, ax=ax)
+plt.show()
+
+######################################################################
+# The first column has the same type of data in both rows, so it may
+# be desirable to combine the colorbar which we do by calling
+# `.Figure.colorbar` with a list of axes instead of a single axes.
+
+fig, axs = plt.subplots(2, 2)
+cm = ['RdBu_r', 'viridis']
+for col in range(2):
+    for row in range(2):
+        ax = axs[row, col]
+        pcm = ax.pcolormesh(np.random.random((20, 20)) * (col + 1),
+                            cmap=cm[col])
+    fig.colorbar(pcm, ax=axs[:, col], shrink=0.6)
+plt.show()
+
+######################################################################
+# Relatively complicated colorbar layouts are possible using this
+# paradigm.  Note that this example works far better with
+# ``constrained_layout=True``
+
+fig, axs = plt.subplots(3, 3, constrained_layout=True)
+for ax in axs.flat:
+    pcm = ax.pcolormesh(np.random.random((20, 20)))
+
+fig.colorbar(pcm, ax=axs[0, :2], shrink=0.6, location='bottom')
+fig.colorbar(pcm, ax=[axs[0, 2]], location='bottom')
+fig.colorbar(pcm, ax=axs[1:, :], location='right', shrink=0.6)
+fig.colorbar(pcm, ax=[axs[2, 1]], location='left')
+
+
+plt.show()
diff --git a/lib/matplotlib/_constrained_layout.py b/lib/matplotlib/_constrained_layout.py
@@ -572,6 +572,36 @@ def layoutcolorbarsingle(ax, cax, shrink, aspect, location, pad=0.05):
     return lb, lbpos
 
 
+def _getmaxminrowcolumn(axs):
+    # helper to get the min/max rows and columns of a list of axes.
+    maxrow = -100000
+    minrow = 1000000
+    maxax = None
+    minax = None
+    maxcol = -100000
+    mincol = 1000000
+    maxax_col = None
+    minax_col = None
+
+    for ax in axs:
+        subspec = ax.get_subplotspec()
+        nrows, ncols, row_start, row_stop, col_start, col_stop = \
+            subspec.get_rows_columns()
+        if row_stop > maxrow:
+            maxrow = row_stop
+            maxax = ax
+        if row_start < minrow:
+            minrow = row_start
+            minax = ax
+        if col_stop > maxcol:
+            maxcol = col_stop
+            maxax_col = ax
+        if col_start < mincol:
+            mincol = col_start
+            minax_col = ax
+    return (minrow, maxrow, minax, maxax, mincol, maxcol, minax_col, maxax_col)
+
+
 def layoutcolorbargridspec(parents, cax, shrink, aspect, location, pad=0.05):
     """
     Do the layout for a colorbar, to not oeverly pollute colorbar.py
@@ -586,6 +616,10 @@ def layoutcolorbargridspec(parents, cax, shrink, aspect, location, pad=0.05):
     lb = layoutbox.LayoutBox(parent=gslb.parent,
                              name=gslb.parent.name + '.cbar',
                              artist=cax)
+    # figure out the row and column extent of the parents.
+    (minrow, maxrow, minax_row, maxax_row,
+     mincol, maxcol, minax_col, maxax_col) = _getmaxminrowcolumn(parents)
+
     if location in ('left', 'right'):
         lbpos = layoutbox.LayoutBox(
                 parent=lb,
@@ -594,39 +628,43 @@ def layoutcolorbargridspec(parents, cax, shrink, aspect, location, pad=0.05):
                 pos=True,
                 subplot=False,
                 artist=cax)
-
-        if location == 'right':
-            # arrange to right of the gridpec sibbling
-            layoutbox.hstack([gslb, lb], padding=pad * gslb.width,
-                             strength='strong')
-        else:
-            layoutbox.hstack([lb, gslb], padding=pad * gslb.width)
+        for ax in parents:
+            if location == 'right':
+                order = [ax._layoutbox, lb]
+            else:
+                order = [lb, ax._layoutbox]
+            layoutbox.hstack(order, padding=pad * gslb.width,
+                         strength='strong')
         # constrain the height and center...
         # This isn't quite right.  We'd like the colorbar
         # pos to line up w/ the axes poss, not the size of the
         # gs.
-        maxrow = -100000
-        minrow = 1000000
-        maxax = None
-        minax = None
 
-        for ax in parents:
-            subspec = ax.get_subplotspec()
-            nrows, ncols = subspec.get_gridspec().get_geometry()
-            for num in [subspec.num1, subspec.num2]:
-                rownum1, colnum1 = divmod(subspec.num1, ncols)
-                if rownum1 > maxrow:
-                    maxrow = rownum1
-                    maxax = ax
-                if rownum1 < minrow:
-                    minrow = rownum1
-                    minax = ax
-        # invert the order so these are bottom to top:
-        maxposlb = minax._poslayoutbox
-        minposlb = maxax._poslayoutbox
+        # Horizontal Layout: need to check all the axes in this gridspec
+        for ch in gslb.children:
+            subspec = ch.artist
+            nrows, ncols, row_start, row_stop, col_start, col_stop = \
+                subspec.get_rows_columns()
+            if location == 'right':
+                if col_stop <= maxcol:
+                    order = [subspec._layoutbox, lb]
+                    # arrange to right of the parents
+                if col_start > maxcol:
+                    order = [lb, subspec._layoutbox]
+            elif location == 'left':
+                if col_start >= mincol:
+                    order = [lb, subspec._layoutbox]
+                if col_stop < mincol:
+                    order = [subspec._layoutbox, lb]
+            layoutbox.hstack(order, padding=pad * gslb.width,
+                             strength='strong')
+
+        # Vertical layout:
+        maxposlb = minax_row._poslayoutbox
+        minposlb = maxax_row._poslayoutbox
         # now we want the height of the colorbar pos to be
-        # set by the top and bottom of these poss
-        # bottom              top
+        # set by the top and bottom of the min/max axes...
+        # bottom            top
         #     b             t
         # h = (top-bottom)*shrink
         # b = bottom + (top-bottom - h) / 2.
@@ -650,29 +688,35 @@ def layoutcolorbargridspec(parents, cax, shrink, aspect, location, pad=0.05):
                 subplot=False,
                 artist=cax)
 
-        if location == 'bottom':
-            layoutbox.vstack([gslb, lb], padding=pad * gslb.width)
-        else:
-            layoutbox.vstack([lb, gslb], padding=pad * gslb.width)
-
-        maxcol = -100000
-        mincol = 1000000
-        maxax = None
-        minax = None
-
         for ax in parents:
-            subspec = ax.get_subplotspec()
-            nrows, ncols = subspec.get_gridspec().get_geometry()
-            for num in [subspec.num1, subspec.num2]:
-                rownum1, colnum1 = divmod(subspec.num1, ncols)
-                if colnum1 > maxcol:
-                    maxcol = colnum1
-                    maxax = ax
-                if rownum1 < mincol:
-                    mincol = colnum1
-                    minax = ax
-        maxposlb = maxax._poslayoutbox
-        minposlb = minax._poslayoutbox
+            if location == 'bottom':
+                order = [ax._layoutbox, lb]
+            else:
+                order = [lb, ax._layoutbox]
+            layoutbox.vstack(order, padding=pad * gslb.width,
+                         strength='strong')
+
+        # Vertical Layout: need to check all the axes in this gridspec
+        for ch in gslb.children:
+            subspec = ch.artist
+            nrows, ncols, row_start, row_stop, col_start, col_stop = \
+                subspec.get_rows_columns()
+            if location == 'bottom':
+                if row_stop <= minrow:
+                    order = [subspec._layoutbox, lb]
+                if row_start > maxrow:
+                    order = [lb, subspec._layoutbox]
+            elif location == 'top':
+                if row_stop < minrow:
+                    order = [subspec._layoutbox, lb]
+                if row_start >= maxrow:
+                    order = [lb, subspec._layoutbox]
+            layoutbox.vstack(order, padding=pad * gslb.width,
+                             strength='strong')
+
+        # Do horizontal layout...
+        maxposlb = maxax_col._poslayoutbox
+        minposlb = minax_col._poslayoutbox
         lbpos.constrain_width((maxposlb.right - minposlb.left) *
                               shrink)
         lbpos.constrain_left(
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -1243,6 +1243,7 @@ def make_axes(parents, location=None, orientation=None, fraction=0.15,
     Returns (cax, kw), the child axes and the reduced kw dictionary to be
     passed when creating the colorbar instance.
     '''
+
     locations = ["left", "right", "top", "bottom"]
     if orientation is not None and location is not None:
         raise TypeError('position and orientation are mutually exclusive. '
diff --git a/lib/matplotlib/tests/baseline_images/test_constrainedlayout/test_colorbar_location.png b/lib/matplotlib/tests/baseline_images/test_constrainedlayout/test_colorbar_location.png
diff --git a/lib/matplotlib/tests/test_constrainedlayout.py b/lib/matplotlib/tests/test_constrainedlayout.py
@@ -378,3 +378,24 @@ def test_constrained_layout23():
     for i in range(2):
         fig, ax = plt.subplots(num="123", constrained_layout=True, clear=True)
         fig.suptitle("Suptitle{}".format(i))
+
+
+@image_comparison(baseline_images=['test_colorbar_location'],
+        extensions=['png'], remove_text=True, style='mpl20')
+def test_colorbar_location():
+    """
+    Test that colorbar handling is as expected for various complicated
+    cases...
+    """
+
+    fig, axs = plt.subplots(4, 5, constrained_layout=True)
+    for ax in axs.flatten():
+        pcm = example_pcolor(ax)
+        ax.set_xlabel('')
+        ax.set_ylabel('')
+    fig.colorbar(pcm, ax=axs[:, 1], shrink=0.4)
+    fig.colorbar(pcm, ax=axs[-1, :2], shrink=0.5, location='bottom')
+    fig.colorbar(pcm, ax=axs[0, 2:], shrink=0.5, location='bottom')
+    fig.colorbar(pcm, ax=axs[-2, 3:], shrink=0.5, location='top')
+    fig.colorbar(pcm, ax=axs[0, 0], shrink=0.5, location='left')
+    fig.colorbar(pcm, ax=axs[1:3, 2], shrink=0.5, location='right')
diff --git a/tutorials/intermediate/constrainedlayout_guide.py b/tutorials/intermediate/constrainedlayout_guide.py
@@ -118,10 +118,10 @@ def example_plot(ax, fontsize=12, nodec=False):
 #
 # .. note::
 #
-#   For the `~.axes.Axes.pcolormesh` kwargs (``pc_kwargs``) we use a dictionary.
-#   Below we will assign one colorbar to a number of axes each containing
-#   a `~.cm.ScalarMappable`; specifying the norm and colormap ensures
-#   the colorbar is accurate for all the axes.
+#   For the `~.axes.Axes.pcolormesh` kwargs (``pc_kwargs``) we use a
+#   dictionary. Below we will assign one colorbar to a number of axes each
+#   containing a `~.cm.ScalarMappable`; specifying the norm and colormap
+#   ensures the colorbar is accurate for all the axes.
 
 arr = np.arange(100).reshape((10, 10))
 norm = mcolors.Normalize(vmin=0., vmax=100.)
@@ -133,14 +133,25 @@ def example_plot(ax, fontsize=12, nodec=False):
 
 ############################################################################
 # If you specify a list of axes (or other iterable container) to the
-# ``ax`` argument of ``colorbar``, constrained_layout will take space from all
-# axes that share the same gridspec.
+# ``ax`` argument of ``colorbar``, constrained_layout will take space from
+# the specified axes.
 
 fig, axs = plt.subplots(2, 2, figsize=(4, 4), constrained_layout=True)
 for ax in axs.flatten():
     im = ax.pcolormesh(arr, **pc_kwargs)
 fig.colorbar(im, ax=axs, shrink=0.6)
 
+############################################################################
+# If you specify a list of axes from inside a grid of axes, the colorbar
+# will steal space appropriately, and leave a gap, but all subplots will
+# still be the same size.
+
+fig, axs = plt.subplots(3, 3, figsize=(4, 4), constrained_layout=True)
+for ax in axs.flatten():
+    im = ax.pcolormesh(arr, **pc_kwargs)
+fig.colorbar(im, ax=axs[1:, ][:, 1], shrink=0.8)
+fig.colorbar(im, ax=axs[:, -1], shrink=0.6)
+
 ############################################################################
 # Note that there is a bit of a subtlety when specifying a single axes
 # as the parent.  In the following, it might be desirable and expected
@@ -458,6 +469,21 @@ def docomplicated(suptitle=None):
 ax2 = fig.add_axes(bb_ax2)
 
 ###############################################################################
+# Manually turning off ``constrained_layout``
+# ===========================================
+#
+# ``constrained_layout`` usually adjusts the axes positions on each draw
+# of the figure.  If you want to get the spacing provided by
+# ``constrained_layout`` but not have it update, then do the initial
+# draw and then call ``fig.set_constrained_layout(False)``.
+# This is potentially useful for animations where the tick labels may
+# change length.
+#
+# Note that ``constrained_layout`` is turned off for ``ZOOM`` and ``PAN``
+# GUI events for the backends that use the toolbar.  This prevents the
+# axes from changing position during zooming and panning.
+#
+#
 # Limitations
 # ========================
 #
