Make Collection.set_paths even faster

apaszke · apaszke · commit 230325de501a · 2020-03-23T12:49:37.000+01:00
diff --git a/lib/matplotlib/cbook/__init__.py b/lib/matplotlib/cbook/__init__.py
@@ -16,6 +16,7 @@
 import os
 from pathlib import Path
 import re
+import gc
 import shlex
 import subprocess
 import sys
@@ -2343,3 +2344,17 @@ def _backend_module_name(name):
     """
     return (name[9:] if name.startswith("module://")
             else "matplotlib.backends.backend_{}".format(name.lower()))
+
+
+def _without_gc(f):
+    """A decorator to disable garbage collection in the decorated function."""
+    @functools.wraps(f)
+    def wrapper(*args, **kwargs):
+        was_enabled = gc.isenabled()
+        gc.disable()
+        try:
+            return f(*args, **kwargs)
+        finally:
+            if was_enabled:
+                gc.enable()
+    return wrapper
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -1070,6 +1070,10 @@ def __init__(self, verts, sizes=None, closed=True, **kwargs):
         self.set_verts(verts, closed)
         self.stale = True
 
+    # This function creates a lot of Path object, which is pretty much
+    # guaranteed to trigger the GC multiple times. None of them will be garbage
+    # just yet, so all those runs are completely unnecessary.
+    @cbook._without_gc
     def set_verts(self, verts, closed=True):
         """
         Set the vertices of the polygons.
@@ -1094,15 +1098,19 @@ def set_verts(self, verts, closed=True):
             return
 
         # Fast path for arrays
-        if isinstance(verts, np.ndarray):
-            verts_pad = np.concatenate((verts, verts[:, :1]), axis=1)
+        if isinstance(verts, np.ndarray) and len(verts):
+            verts_pad = (np.concatenate((verts, verts[:, :1]), axis=1)
+                           .astype(mpath.Path.verts_type))
             # Creating the codes once is much faster than having Path do it
             # separately each time by passing closed=True.
-            codes = np.empty(verts_pad.shape[1], dtype=mpath.Path.code_type)
-            codes[:] = mpath.Path.LINETO
-            codes[0] = mpath.Path.MOVETO
-            codes[-1] = mpath.Path.CLOSEPOLY
-            self._paths = [mpath.Path(xy, codes) for xy in verts_pad]
+            example_path = mpath.Path(verts_pad[0], closed=True)
+            # Looking up the values once speeds up the iteration a bit
+            _make_path = mpath.Path._fast_from_codes_and_verts
+            codes = example_path.codes
+            self._paths = [_make_path(xy, codes,
+                                      internals_from=example_path,
+                                      unmask_verts=False)
+                           for xy in verts_pad]
             return
 
         self._paths = []
diff --git a/lib/matplotlib/path.py b/lib/matplotlib/path.py
@@ -74,8 +74,12 @@ class Path:
         made up front in the constructor that will not change when the
         data changes.
     """
+    __slots__ = ('_vertices', '_codes', '_readonly',
+                 '_should_simplify', '_simplify_threshold',
+                 '_interpolation_steps', '__weakref__')
 
     code_type = np.uint8
+    verts_type = float
 
     # Path codes
     STOP = code_type(0)         # 1 vertex
@@ -160,7 +164,8 @@ def __init__(self, vertices, codes=None, _interpolation_steps=1,
             self._readonly = False
 
     @classmethod
-    def _fast_from_codes_and_verts(cls, verts, codes, internals_from=None):
+    def _fast_from_codes_and_verts(cls, verts, codes, internals_from=None,
+                                   unmask_verts=True):
         """
         Creates a Path instance without the expense of calling the constructor.
 
@@ -173,9 +178,17 @@ def _fast_from_codes_and_verts(cls, verts, codes, internals_from=None):
             ``should_simplify``, ``simplify_threshold``, and
             ``interpolation_steps`` will be copied.  Note that ``readonly`` is
             never copied, and always set to ``False`` by this constructor.
+        unmask_verts : bool
+            If False, vertices should be an unmasked NumPy array with dtype
+            set to Path.verts_type.
         """
         pth = cls.__new__(cls)
-        pth._vertices = _to_unmasked_float_array(verts)
+        # NOTE: _to_unmasked_float_array has non-trivial overhead in tight
+        #       loops, so we allow skipping it if the caller wants to
+        if unmask_verts:
+            pth._vertices = _to_unmasked_float_array(verts)
+        else:
+            pth._vertices = verts
         pth._codes = codes
         pth._readonly = False
         if internals_from is not None:
@@ -269,16 +282,6 @@ def readonly(self):
         """
         return self._readonly
 
-    def __copy__(self):
-        """
-        Returns a shallow copy of the `Path`, which will share the
-        vertices and codes with the source `Path`.
-        """
-        import copy
-        return copy.copy(self)
-
-    copy = __copy__
-
     def __deepcopy__(self, memo=None):
         """
         Returns a deepcopy of the `Path`.  The `Path` will not be
diff --git a/lib/matplotlib/tests/test_transforms.py b/lib/matplotlib/tests/test_transforms.py
@@ -621,28 +621,6 @@ def test_invalid_arguments():
         t.transform([[1, 2, 3]])
 
 
-def test_transformed_path():
-    points = [(0, 0), (1, 0), (1, 1), (0, 1)]
-    path = Path(points, closed=True)
-
-    trans = mtransforms.Affine2D()
-    trans_path = mtransforms.TransformedPath(path, trans)
-    assert_allclose(trans_path.get_fully_transformed_path().vertices, points)
-
-    # Changing the transform should change the result.
-    r2 = 1 / np.sqrt(2)
-    trans.rotate(np.pi / 4)
-    assert_allclose(trans_path.get_fully_transformed_path().vertices,
-                    [(0, 0), (r2, r2), (0, 2 * r2), (-r2, r2)],
-                    atol=1e-15)
-
-    # Changing the path does not change the result (it's cached).
-    path.points = [(0, 0)] * 4
-    assert_allclose(trans_path.get_fully_transformed_path().vertices,
-                    [(0, 0), (r2, r2), (0, 2 * r2), (-r2, r2)],
-                    atol=1e-15)
-
-
 def test_transformed_patch_path():
     trans = mtransforms.Affine2D()
     patch = mpatches.Wedge((0, 0), 1, 45, 135, transform=trans)
