ENH: constrained_layout simple compress axes

jklymak · jklymak · commit 6fa3fa93921a · 2021-04-18T17:45:02.000-07:00
diff --git a/lib/matplotlib/_constrained_layout.py b/lib/matplotlib/_constrained_layout.py
@@ -61,7 +61,7 @@
 
 ######################################################
 def do_constrained_layout(fig, renderer, h_pad, w_pad,
-                          hspace=None, wspace=None):
+                          hspace=None, wspace=None, compress=False):
     """
     Do the constrained_layout.  Called at draw time in
      ``figure.constrained_layout()``
@@ -83,6 +83,11 @@ def do_constrained_layout(fig, renderer, h_pad, w_pad,
        A value of 0.2 for a three-column layout would have a space
        of 0.1 of the figure width between each column.
        If h/wspace < h/w_pad, then the pads are used instead.
+
+    compress : boolean, False
+        Whether to try and push axes together if their aspect ratios
+        make it so that the they will have lots of extra white space
+        between them.  Useful for grids of images or maps.
     """
 
     # list of unique gridspecs that contain child axes:
@@ -98,7 +103,7 @@ def do_constrained_layout(fig, renderer, h_pad, w_pad,
                            'Possibly did not call parent GridSpec with the'
                            ' "figure" keyword')
 
-    for _ in range(2):
+    for nn in range(2):
         # do the algorithm twice.  This has to be done because decorations
         # change size after the first re-position (i.e. x/yticklabels get
         # larger/smaller).  This second reposition tends to be much milder,
@@ -118,16 +123,56 @@ def do_constrained_layout(fig, renderer, h_pad, w_pad,
         # update all the variables in the layout.
         fig._layoutgrid.update_variables()
 
+        warn_collapsed = ('constrained_layout not applied because '
+                          'axes sizes collapsed to zero.  Try making '
+                          'figure larger or axes decorations smaller.')
         if _check_no_collapsed_axes(fig):
             _reposition_axes(fig, renderer, h_pad=h_pad, w_pad=w_pad,
                              hspace=hspace, wspace=wspace)
+            if compress:
+                _compress_fixed_aspect(fig)
+                # update all the variables in the layout.
+                fig._layoutgrid.update_variables()
+                if _check_no_collapsed_axes(fig):
+                    _reposition_axes(fig, renderer, h_pad=h_pad, w_pad=w_pad,
+                                     hspace=hspace, wspace=wspace)
+                else:
+                    _api.warn_external(warn_collapsed)
         else:
-            _api.warn_external('constrained_layout not applied because '
-                               'axes sizes collapsed to zero.  Try making '
-                               'figure larger or axes decorations smaller.')
+            _api.warn_external(warn_collapsed)
         _reset_margins(fig)
 
 
+def _compress_fixed_aspect(fig):
+    extraw = dict()
+    extrah = dict()
+    for ax in fig.axes:
+        if hasattr(ax, 'get_subplotspec'):
+            actual = ax.get_position(original=False)
+            ax.apply_aspect()
+            sub = ax.get_subplotspec()
+            gs = sub.get_gridspec()
+            if gs not in extraw.keys():
+                extraw[gs] = np.zeros(gs.ncols)
+                extrah[gs] = np.zeros(gs.nrows)
+            orig = ax.get_position(original=True)
+            actual = ax.get_position(original=False)
+            dw = orig.width - actual.width
+            if dw > 0:
+                for i in sub.colspan:
+                    extraw[gs][i] = max(extraw[gs][i], dw)
+            dh = orig.height - actual.height
+            if dh > 0:
+                for i in sub.rowspan:
+                    extrah[gs][i] = max(extrah[gs][i], dh)
+
+    fig._layoutgrid.edit_margin_min('left', np.sum(extraw[gs]) / 2)
+    fig._layoutgrid.edit_margin_min('right', np.sum(extraw[gs]) / 2)
+
+    fig._layoutgrid.edit_margin_min('top', np.sum(extrah[gs]) / 2)
+    fig._layoutgrid.edit_margin_min('bottom', np.sum(extrah[gs]) / 2)
+
+
 def _check_no_collapsed_axes(fig):
     """
     Check that no axes have collapsed to zero size.
diff --git a/lib/matplotlib/_layoutgrid.py b/lib/matplotlib/_layoutgrid.py
@@ -120,7 +120,9 @@ def __repr__(self):
                        f'innerW{self.inner_widths[j].value():1.3f}, ' \
                        f'innerH{self.inner_heights[i].value():1.3f}, ' \
                        f'ML{self.margins["left"][j].value():1.3f}, ' \
-                       f'MR{self.margins["right"][j].value():1.3f}, \n'
+                       f'MR{self.margins["right"][j].value():1.3f}, ' \
+                       f'MB{self.margins["bottom"][j].value():1.3f}, ' \
+                       f'MT{self.margins["top"][j].value():1.3f}, \n'
         return str
 
     def reset_margins(self):
diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
@@ -2449,6 +2449,10 @@ def set_constrained_layout_pads(self, **kwargs):
             Height padding between subplots, expressed as a fraction of the
             subplot width. The total padding ends up being h_pad + hspace.
 
+        compress : boolean
+            Try to compress axes in constrained layout.  Useful when
+            axes aspect ratios make it so that there is substantial
+            white space between them.
         """
 
         todo = ['w_pad', 'h_pad', 'wspace', 'hspace']
@@ -2458,6 +2462,8 @@ def set_constrained_layout_pads(self, **kwargs):
             else:
                 self._constrained_layout_pads[td] = (
                     mpl.rcParams['figure.constrained_layout.' + td])
+        self._constrained_layout_pads['compress'] = (
+            kwargs.get('compress', False))
 
     def get_constrained_layout_pads(self, relative=False):
         """
@@ -2477,14 +2483,15 @@ def get_constrained_layout_pads(self, relative=False):
         h_pad = self._constrained_layout_pads['h_pad']
         wspace = self._constrained_layout_pads['wspace']
         hspace = self._constrained_layout_pads['hspace']
+        compress = self._constrained_layout_pads['compress']
 
         if relative and (w_pad is not None or h_pad is not None):
             renderer0 = layoutgrid.get_renderer(self)
             dpi = renderer0.dpi
             w_pad = w_pad * dpi / renderer0.width
             h_pad = h_pad * dpi / renderer0.height
 
-        return w_pad, h_pad, wspace, hspace
+        return w_pad, h_pad, wspace, hspace, compress
 
     def set_canvas(self, canvas):
         """
@@ -3073,15 +3080,16 @@ def execute_constrained_layout(self, renderer=None):
                                "or you need to call figure or subplots "
                                "with the constrained_layout=True kwarg.")
             return
-        w_pad, h_pad, wspace, hspace = self.get_constrained_layout_pads()
+        w_pad, h_pad, wspace, hspace, comp = self.get_constrained_layout_pads()
         # convert to unit-relative lengths
         fig = self
         width, height = fig.get_size_inches()
         w_pad = w_pad / width
         h_pad = h_pad / height
         if renderer is None:
             renderer = _get_renderer(fig)
-        do_constrained_layout(fig, renderer, h_pad, w_pad, hspace, wspace)
+        do_constrained_layout(fig, renderer, h_pad, w_pad, hspace, wspace,
+                              compress=comp)
 
     def tight_layout(self, *, pad=1.08, h_pad=None, w_pad=None, rect=None):
         """
