Speedup get_tightbbox() for non-rectilinear axes

lukelbd · lukelbd · commit fc1a38611589 · 2021-12-16T17:17:23.000-07:00
diff --git a/lib/matplotlib/artist.py b/lib/matplotlib/artist.py
@@ -320,23 +320,6 @@ def get_window_extent(self, renderer):
         """
         return Bbox([[0, 0], [0, 0]])
 
-    def _get_clipping_extent_bbox(self):
-        """
-        Return a bbox with the extents of the intersection of the clip_path
-        and clip_box for this artist, or None if both of these are
-        None, or ``get_clip_on`` is False.
-        """
-        bbox = None
-        if self.get_clip_on():
-            clip_box = self.get_clip_box()
-            if clip_box is not None:
-                bbox = clip_box
-            clip_path = self.get_clip_path()
-            if clip_path is not None and bbox is not None:
-                clip_path = clip_path.get_fully_transformed_path()
-                bbox = Bbox.intersection(bbox, clip_path.get_extents())
-        return bbox
-
     def get_tightbbox(self, renderer):
         """
         Like `.Artist.get_window_extent`, but includes any clipping.
@@ -358,7 +341,7 @@ def get_tightbbox(self, renderer):
             if clip_box is not None:
                 bbox = Bbox.intersection(bbox, clip_box)
             clip_path = self.get_clip_path()
-            if clip_path is not None and bbox is not None:
+            if clip_path is not None:
                 clip_path = clip_path.get_fully_transformed_path()
                 bbox = Bbox.intersection(bbox, clip_path.get_extents())
         return bbox
@@ -844,6 +827,30 @@ def get_in_layout(self):
         """
         return self._in_layout
 
+    def _fully_clipped_to_axes(self):
+        """
+        Return a boolean flag, ``True`` if the artist can be safely skipped
+        in axes layout calculations. Requires ``get_clip_on`` is True, one of
+        `clip_box` or `clip_path` is set, ``clip_box.extents`` is equivalent
+        to ``axes.bbox.extents`` (if set), and ``clip_path._patch`` is
+        equivalent to ``axes.patch`` (if set).
+        """
+        # Note that ``clip_path.get_fully_transformed_path().get_extents()``
+        # cannot be directly compared to ``axes.bbox.extents`` because the
+        # extents may be undefined (i.e. equivalent to ``Bbox.null()``)
+        # before the associated artist is drawn, and this method is meant
+        # to determine whether ``axes.get_tightbbox()`` may bypass drawing
+        clip_box = self.get_clip_box()
+        clip_path = self.get_clip_path()
+        return (self.axes is not None
+                and self.get_clip_on()
+                and (clip_box is not None or clip_path is not None)
+                and (clip_box is None
+                     or np.all(clip_box.extents == self.axes.bbox.extents))
+                and (clip_path is None
+                     or isinstance(clip_path, TransformedPatchPath)
+                     and clip_path._patch is self.axes.patch))
+
     def get_clip_on(self):
         """Return whether the artist uses clipping."""
         return self._clipon
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -21,6 +21,7 @@
 import matplotlib.font_manager as font_manager
 import matplotlib.image as mimage
 import matplotlib.lines as mlines
+from matplotlib.offsetbox import OffsetBox
 import matplotlib.patches as mpatches
 import matplotlib.path as mpath
 from matplotlib.rcsetup import cycler, validate_axisbelow
@@ -4541,21 +4542,25 @@ def get_default_bbox_extra_artists(self):
 
         artists = self.get_children()
 
+        for _axis in self._get_axis_list():
+            # axis tight bboxes are calculated separately inside
+            # Axes.get_tightbbox() using for_layout_only=True
+            artists.remove(_axis)
         if not (self.axison and self._frameon):
             # don't do bbox on spines if frame not on.
             for spine in self.spines.values():
                 artists.remove(spine)
 
-        if not self.axison:
-            for _axis in self._get_axis_list():
-                artists.remove(_axis)
-
         artists.remove(self.title)
         artists.remove(self._left_title)
         artists.remove(self._right_title)
 
-        return [artist for artist in artists
-                if (artist.get_visible() and artist.get_in_layout())]
+        # always include types that do not internally implement clipping
+        # to axes. may have clip_on set to True and clip_box equivalent
+        # to ax.bbox but then ignore these properties during draw.
+        noclip = (_AxesBase, maxis.Axis, OffsetBox, mtext._AnnotationBase)
+        return [a for a in artists if a.get_visible() and a.get_in_layout()
+                and (isinstance(a, noclip) or not a._fully_clipped_to_axes())]
 
     def get_tightbbox(self, renderer, call_axes_locator=True,
                       bbox_extra_artists=None, *, for_layout_only=False):
@@ -4612,17 +4617,11 @@ def get_tightbbox(self, renderer, call_axes_locator=True,
         else:
             self.apply_aspect()
 
-        if self.axison:
-            if self.xaxis.get_visible():
-                bb_xaxis = martist._get_tightbbox_for_layout_only(
-                    self.xaxis, renderer)
-                if bb_xaxis:
-                    bb.append(bb_xaxis)
-            if self.yaxis.get_visible():
-                bb_yaxis = martist._get_tightbbox_for_layout_only(
-                    self.yaxis, renderer)
-                if bb_yaxis:
-                    bb.append(bb_yaxis)
+        for axis in self._get_axis_list():
+            if self.axison and axis.get_visible():
+                ba = martist._get_tightbbox_for_layout_only(axis, renderer)
+                if ba:
+                    bb.append(ba)
         self._update_title_position(renderer)
         axbbox = self.get_window_extent(renderer)
         bb.append(axbbox)
@@ -4643,17 +4642,6 @@ def get_tightbbox(self, renderer, call_axes_locator=True,
             bbox_artists = self.get_default_bbox_extra_artists()
 
         for a in bbox_artists:
-            # Extra check here to quickly see if clipping is on and
-            # contained in the Axes.  If it is, don't get the tightbbox for
-            # this artist because this can be expensive:
-            clip_extent = a._get_clipping_extent_bbox()
-            if clip_extent is not None:
-                clip_extent = mtransforms.Bbox.intersection(
-                    clip_extent, axbbox)
-                if np.all(clip_extent.extents == axbbox.extents):
-                    # clip extent is inside the Axes bbox so don't check
-                    # this artist
-                    continue
             bbox = a.get_tightbbox(renderer)
             if (bbox is not None
                     and 0 < bbox.width < np.inf
diff --git a/lib/matplotlib/tests/test_tightlayout.py b/lib/matplotlib/tests/test_tightlayout.py
@@ -342,3 +342,27 @@ def test_manual_colorbar():
     fig.colorbar(pts, cax=cax)
     with pytest.warns(UserWarning, match="This figure includes Axes"):
         fig.tight_layout()
+
+
+def test_skip_layout():
+    # Ensure that _fully_clipped_to_axes() returns True under default
+    # conditions for all projection types. Axes.get_tightbbox()
+    # uses this to skip artists in layout calculations.
+    arr = np.arange(100).reshape((10, 10))
+    fig = plt.figure(figsize=(6, 2))
+    ax1 = fig.add_subplot(131, projection='rectilinear')
+    ax2 = fig.add_subplot(132, projection='mollweide')
+    ax3 = fig.add_subplot(133, projection='polar')
+    for ax in (ax1, ax2, ax3):
+        # Default conditions (clipped by ax.bbox or ax.patch)
+        ax.grid(False)
+        h, = ax.plot(arr[:, 0])
+        m = ax.pcolor(arr)
+        assert h._fully_clipped_to_axes()
+        assert m._fully_clipped_to_axes()
+        # Non-default conditions (not clipped by ax.patch)
+        rect = Rectangle((0, 0), 0.5, 0.5, transform=ax.transAxes)
+        h.set_clip_path(rect)
+        m.set_clip_path(rect.get_path(), rect.get_transform())
+        assert not h._fully_clipped_to_axes()
+        assert not m._fully_clipped_to_axes()
