Make Collection.set_paths even faster

apaszke · apaszke · commit 7a1bae755e66 · 2020-03-16T21:35:28.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
@@ -2337,3 +2338,16 @@ def _backend_module_name(name):
     """
     return (name[9:] if name.startswith("module://")
             else "matplotlib.backends.backend_{}".format(name.lower()))
+
+
+def _without_gc(f):
+    @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
@@ -1076,6 +1076,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 necessary.
+    @cbook._without_gc
     def set_verts(self, verts, closed=True):
         """
         Set the vertices of the polygons.
@@ -1101,14 +1105,21 @@ def set_verts(self, verts, closed=True):
 
         # Fast path for arrays
         if isinstance(verts, np.ndarray):
-            verts_pad = np.concatenate((verts, verts[:, :1]), axis=1)
+            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], codes)
+            # Looking up the function once before the iteration gives a nice speedup
+            _make_path = mpath.Path._fast_from_codes_and_verts
+            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')
 
     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,15 @@ 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)
+        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 +280,14 @@ def readonly(self):
         """
         return self._readonly
 
-    def __copy__(self):
+    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
