Add a fast path for NumPy arrays to Collection.set_verts

apaszke · apaszke · commit 0abe3670522f · 2020-03-05T23:07:27.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,28 +1090,41 @@ def set_verts(self, verts, closed=True):
             Whether the polygon should be closed by adding a CLOSEPOLY
             connection at the end.
         """
+        self.stale = True
+        # This is much faster than having Path do it one at a time.
         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[0:1]])
-                    else:
-                        xy = np.asarray(xy)
-                        xy = np.concatenate([xy, xy[0:1]])
-                    codes = np.empty(xy.shape[0], dtype=mpath.Path.code_type)
-                    codes[:] = mpath.Path.LINETO
-                    codes[0] = mpath.Path.MOVETO
-                    codes[-1] = mpath.Path.CLOSEPOLY
-                    self._paths.append(mpath.Path(xy, codes))
-                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)
+            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[0:1]])
+                else:
+                    xy = np.asarray(xy)
+                    xy = np.concatenate([xy, xy[0:1]])
+                codes = np.empty(xy.shape[0], dtype=mpath.Path.code_type)
+                codes[:] = mpath.Path.LINETO
+                codes[0] = mpath.Path.MOVETO
+                codes[-1] = mpath.Path.CLOSEPOLY
+                self._paths.append(mpath.Path(xy, codes))
+            else:
+                self._paths.append(mpath.Path(xy))
 
     set_paths = set_verts
 
