Shorten tight_layout code.

anntzer · anntzer · commit 8524caaa2be5 · 2020-04-23T01:14:28.000+02:00
- Turn hspaces/vspaces into 2D arrays, rather than flat lists that
  need to be repeatedly accessed with flattened indices, and directly
  accumulate widths/heights in there.
- Merge the num2 is None case by treating it as a slice of length 1,
  which then allows just inlining _get_{left,right,top,bottom}.
- Minor refactorings.
diff --git a/lib/matplotlib/tight_layout.py b/lib/matplotlib/tight_layout.py
@@ -9,27 +9,13 @@
 for some cases (for example, left or right margin is affected by xlabel).
 """
 
+import numpy as np
+
 from matplotlib import cbook, rcParams
 from matplotlib.font_manager import FontProperties
 from matplotlib.transforms import TransformedBbox, Bbox
 
 
-def _get_left(tight_bbox, axes_bbox):
-    return axes_bbox.xmin - tight_bbox.xmin
-
-
-def _get_right(tight_bbox, axes_bbox):
-    return tight_bbox.xmax - axes_bbox.xmax
-
-
-def _get_bottom(tight_bbox, axes_bbox):
-    return axes_bbox.ymin - tight_bbox.ymin
-
-
-def _get_top(tight_bbox, axes_bbox):
-    return tight_bbox.ymax - axes_bbox.ymax
-
-
 def auto_adjust_subplotpars(
         fig, renderer, nrows_ncols, num1num2_list, subplot_list,
         ax_bbox_list=None, pad=1.08, h_pad=None, w_pad=None, rect=None):
@@ -64,15 +50,8 @@ def auto_adjust_subplotpars(
     font_size_inches = (
         FontProperties(size=rcParams["font.size"]).get_size_in_points() / 72)
     pad_inches = pad * font_size_inches
-    if h_pad is not None:
-        vpad_inches = h_pad * font_size_inches
-    else:
-        vpad_inches = pad_inches
-
-    if w_pad is not None:
-        hpad_inches = w_pad * font_size_inches
-    else:
-        hpad_inches = pad_inches
+    vpad_inches = h_pad * font_size_inches if h_pad is not None else pad_inches
+    hpad_inches = w_pad * font_size_inches if w_pad is not None else pad_inches
 
     if len(num1num2_list) != len(subplot_list) or len(subplot_list) == 0:
         raise ValueError
@@ -81,23 +60,15 @@ def auto_adjust_subplotpars(
         margin_left = margin_bottom = margin_right = margin_top = None
     else:
         margin_left, margin_bottom, _right, _top = rect
-        if _right:
-            margin_right = 1 - _right
-        else:
-            margin_right = None
-        if _top:
-            margin_top = 1 - _top
-        else:
-            margin_top = None
-
-    vspaces = [[] for _ in range((rows + 1) * cols)]
-    hspaces = [[] for _ in range(rows * (cols + 1))]
+        margin_right = 1 - _right if _right else None
+        margin_top = 1 - _top if _top else None
 
-    union = Bbox.union
+    vspaces = np.zeros((rows + 1, cols))
+    hspaces = np.zeros((rows, cols + 1))
 
     if ax_bbox_list is None:
         ax_bbox_list = [
-            union([ax.get_position(original=True) for ax in subplots])
+            Bbox.union([ax.get_position(original=True) for ax in subplots])
             for subplots in subplot_list]
 
     for subplots, ax_bbox, (num1, num2) in zip(subplot_list,
@@ -106,66 +77,39 @@ def auto_adjust_subplotpars(
         if all(not ax.get_visible() for ax in subplots):
             continue
 
-        tight_bbox_raw = union([ax.get_tightbbox(renderer) for ax in subplots
-                                if ax.get_visible()])
+        tight_bbox_raw = Bbox.union([
+            ax.get_tightbbox(renderer) for ax in subplots if ax.get_visible()])
         tight_bbox = TransformedBbox(tight_bbox_raw,
                                      fig.transFigure.inverted())
 
         row1, col1 = divmod(num1, cols)
-
         if num2 is None:
-            # left
-            hspaces[row1 * (cols + 1) + col1].append(
-                                        _get_left(tight_bbox, ax_bbox))
-            # right
-            hspaces[row1 * (cols + 1) + (col1 + 1)].append(
-                                        _get_right(tight_bbox, ax_bbox))
-            # top
-            vspaces[row1 * cols + col1].append(
-                                        _get_top(tight_bbox, ax_bbox))
-            # bottom
-            vspaces[(row1 + 1) * cols + col1].append(
-                                        _get_bottom(tight_bbox, ax_bbox))
+            num2 = num1
+        row2, col2 = divmod(num2, cols)
 
-        else:
-            row2, col2 = divmod(num2, cols)
-
-            for row_i in range(row1, row2 + 1):
-                # left
-                hspaces[row_i * (cols + 1) + col1].append(
-                                    _get_left(tight_bbox, ax_bbox))
-                # right
-                hspaces[row_i * (cols + 1) + (col2 + 1)].append(
-                                    _get_right(tight_bbox, ax_bbox))
-            for col_i in range(col1, col2 + 1):
-                # top
-                vspaces[row1 * cols + col_i].append(
-                                    _get_top(tight_bbox, ax_bbox))
-                # bottom
-                vspaces[(row2 + 1) * cols + col_i].append(
-                                    _get_bottom(tight_bbox, ax_bbox))
+        for row_i in range(row1, row2 + 1):
+            hspaces[row_i, col1] += ax_bbox.xmin - tight_bbox.xmin  # left
+            hspaces[row_i, col2 + 1] += tight_bbox.xmax - ax_bbox.xmax  # right
+        for col_i in range(col1, col2 + 1):
+            vspaces[row1, col_i] += tight_bbox.ymax - ax_bbox.ymax  # top
+            vspaces[row2 + 1, col_i] += ax_bbox.ymin - tight_bbox.ymin  # bot.
 
     fig_width_inch, fig_height_inch = fig.get_size_inches()
 
-    # margins can be negative for axes with aspect applied. And we
-    # append + [0] to make minimum margins 0
-
+    # margins can be negative for axes with aspect applied, so use max(, 0) to
+    # make them nonnegative.
     if not margin_left:
-        margin_left = max([sum(s) for s in hspaces[::cols + 1]] + [0])
-        margin_left += pad_inches / fig_width_inch
-
+        margin_left = (max(hspaces[:, 0].max(), 0)
+                       + pad_inches / fig_width_inch)
     if not margin_right:
-        margin_right = max([sum(s) for s in hspaces[cols::cols + 1]] + [0])
-        margin_right += pad_inches / fig_width_inch
-
+        margin_right = (max(hspaces[:, -1].max(), 0)
+                        + pad_inches / fig_width_inch)
     if not margin_top:
-        margin_top = max([sum(s) for s in vspaces[:cols]] + [0])
-        margin_top += pad_inches / fig_height_inch
-
+        margin_top = (max(vspaces[0, :].max(), 0)
+                      + pad_inches / fig_height_inch)
     if not margin_bottom:
-        margin_bottom = max([sum(s) for s in vspaces[-cols:]] + [0])
-        margin_bottom += pad_inches / fig_height_inch
-
+        margin_bottom = (max(vspaces[-1, :].max(), 0)
+                         + pad_inches / fig_height_inch)
     if margin_left + margin_right >= 1:
         cbook._warn_external('Tight layout not applied. The left and right '
                              'margins cannot be made large enough to '
@@ -181,12 +125,9 @@ def auto_adjust_subplotpars(
                   right=1 - margin_right,
                   bottom=margin_bottom,
                   top=1 - margin_top)
+
     if cols > 1:
-        hspace = (
-            max(sum(s)
-                for i in range(rows)
-                for s in hspaces[i * (cols + 1) + 1:(i + 1) * (cols + 1) - 1])
-            + hpad_inches / fig_width_inch)
+        hspace = hspaces[:, 1:-1].max() + hpad_inches / fig_width_inch
         # axes widths:
         h_axes = (1 - margin_right - margin_left - hspace * (cols - 1)) / cols
         if h_axes < 0:
@@ -196,10 +137,8 @@ def auto_adjust_subplotpars(
             return None
         else:
             kwargs["wspace"] = hspace / h_axes
-
     if rows > 1:
-        vspace = (max(sum(s) for s in vspaces[cols:-cols])
-                  + vpad_inches / fig_height_inch)
+        vspace = vspaces[1:-1, :].max() + vpad_inches / fig_height_inch
         v_axes = (1 - margin_top - margin_bottom - vspace * (rows - 1)) / rows
         if v_axes < 0:
             cbook._warn_external('Tight layout not applied. tight_layout '
