Add a fast path for NumPy arrays to Collection.set_verts

apaszke · apaszke · commit 95977d6443a2 · 2020-03-16T15:27:39.000+01:00
This reduces the run time for larger 3D plots by over a second on some
toy examples.
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -1090,23 +1090,37 @@ def set_verts(self, verts, closed=True):
             Whether the polygon should be closed by adding a CLOSEPOLY
             connection at the end.
         """
+        self.stale = True
         if isinstance(verts, np.ma.MaskedArray):
             verts = verts.astype(float).filled(np.nan)
-            # This is much faster than having Path do it one at a time.
-        if closed:
-            self._paths = []
-            for xy in verts:
-                if len(xy):
-                    if isinstance(xy, np.ma.MaskedArray):
-                        xy = np.ma.concatenate([xy, xy[:1]])
-                    else:
-                        xy = np.concatenate([xy, xy[:1]])
-                    self._paths.append(mpath.Path(xy, closed=True))
-                else:
-                    self._paths.append(mpath.Path(xy))
-        else:
+
+        # No need to do anything fancy if the path isn't closed.
+        if not closed:
             self._paths = [mpath.Path(xy) for xy in verts]
-        self.stale = True
+            return
+
+        # Fast path for arrays
+        if isinstance(verts, np.ndarray):
+            verts_pad = np.concatenate((verts, verts[:, :1]), axis=1)
+            # 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]
+            return
+
+        self._paths = []
+        for xy in verts:
+            if len(xy):
+                if isinstance(xy, np.ma.MaskedArray):
+                    xy = np.ma.concatenate([xy, xy[:1]])
+                else:
+                    xy = np.concatenate([xy, xy[:1]])
+                self._paths.append(mpath.Path(xy, closed=True))
+            else:
+                self._paths.append(mpath.Path(xy))
 
     set_paths = set_verts
 
diff --git a/lib/matplotlib/tests/test_collections.py b/lib/matplotlib/tests/test_collections.py
@@ -12,7 +12,8 @@
 import matplotlib.pyplot as plt
 import matplotlib.collections as mcollections
 import matplotlib.transforms as mtransforms
-from matplotlib.collections import Collection, LineCollection, EventCollection
+from matplotlib.collections import (Collection, LineCollection,
+                                    EventCollection, PolyCollection)
 from matplotlib.testing.decorators import image_comparison
 
 
@@ -612,3 +613,13 @@ def test_EventCollection_nosort():
     arr = np.array([3, 2, 1, 10])
     coll = EventCollection(arr)
     np.testing.assert_array_equal(arr, np.array([3, 2, 1, 10]))
+
+
+def test_collection_set_verts_array():
+    verts = np.arange(80, dtype=np.double).reshape(10, 4, 2)
+    col_arr = PolyCollection(verts)
+    col_list = PolyCollection(list(verts))
+    assert len(col_arr._paths) == len(col_list._paths)
+    for ap, lp in zip(col_arr._paths, col_list._paths):
+        assert np.array_equal(ap._vertices, lp._vertices)
+        assert np.array_equal(ap._codes, lp._codes)
