speed up plot_wireframe

scottshambaugh · scottshambaugh · commit 7f99935c27ea · 2025-01-03T17:32:04.000-07:00
diff --git a/lib/matplotlib/transforms.py b/lib/matplotlib/transforms.py
@@ -874,11 +874,13 @@ def update_from_path(self, path, ignore=None, updatex=True, updatey=True):
         if updatex:
             points[0, 0] = min(points[0, 0], np.min(path.vertices[..., 0]))
             points[1, 0] = max(points[1, 0], np.max(path.vertices[..., 0]))
-            minpos[0] = min(minpos[0], np.min(np.where(path.vertices[..., 0] > 0, path.vertices[..., 0], np.inf)))
+            minpos[0] = min(minpos[0], np.min(np.where(path.vertices[..., 0] > 0,
+                                                       path.vertices[..., 0], np.inf)))
         if updatey:
             points[0, 1] = min(points[0, 1], np.min(path.vertices[..., 1]))
             points[1, 1] = max(points[1, 1], np.max(path.vertices[..., 1]))
-            minpos[1] = min(minpos[1], np.min(np.where(path.vertices[..., 1] > 0, path.vertices[..., 1], np.inf)))
+            minpos[1] = min(minpos[1], np.min(np.where(path.vertices[..., 1] > 0,
+                                                       path.vertices[..., 1], np.inf)))
         if np.any(points != self._points) or np.any(minpos != self._minpos):
             changed = True
 
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -2419,41 +2419,34 @@ def plot_wireframe(self, X, Y, Z, *, axlim_clip=False, **kwargs):
         tX, tY, tZ = np.transpose(X), np.transpose(Y), np.transpose(Z)
 
         if rstride:
-            rii = list(range(0, rows, rstride))
+            rii = np.arange(0, rows, rstride)
             # Add the last index only if needed
             if rows > 0 and rii[-1] != (rows - 1):
-                rii += [rows-1]
+                rii = np.append(rii, rows-1)
         else:
-            rii = []
+            rii = np.array([], dtype=int)
+
         if cstride:
-            cii = list(range(0, cols, cstride))
+            cii = np.arange(0, cols, cstride)
             # Add the last index only if needed
             if cols > 0 and cii[-1] != (cols - 1):
-                cii += [cols-1]
+                cii = np.append(cii, cols-1)
         else:
-            cii = []
+            cii = np.array([], dtype=int)
 
         if rstride == 0 and cstride == 0:
             raise ValueError("Either rstride or cstride must be non zero")
 
         # If the inputs were empty, then just
         # reset everything.
         if Z.size == 0:
-            rii = []
-            cii = []
-
-        xlines = [X[i] for i in rii]
-        ylines = [Y[i] for i in rii]
-        zlines = [Z[i] for i in rii]
+            rii = np.array([], dtype=int)
+            cii = np.array([], dtype=int)
 
-        txlines = [tX[i] for i in cii]
-        tylines = [tY[i] for i in cii]
-        tzlines = [tZ[i] for i in cii]
+        row_lines = np.stack([X[rii], Y[rii], Z[rii]], axis=-1)
+        col_lines = np.stack([tX[cii], tY[cii], tZ[cii]], axis=-1)
 
-        lines = ([list(zip(xl, yl, zl))
-                 for xl, yl, zl in zip(xlines, ylines, zlines)]
-                 + [list(zip(xl, yl, zl))
-                 for xl, yl, zl in zip(txlines, tylines, tzlines)])
+        lines = np.concatenate([row_lines, col_lines])
 
         linec = art3d.Line3DCollection(lines, axlim_clip=axlim_clip, **kwargs)
         self.add_collection(linec)
