Keep autoscale speedup but return original collection

Author: scottshambaugh

lib/mpl_toolkits/mplot3d/axes3d.py

@@ -2446,22 +2446,19 @@ def plot_wireframe(self, X, Y, Z, *, axlim_clip=False, **kwargs):
         row_lines = np.stack([X[rii], Y[rii], Z[rii]], axis=-1)
         col_lines = np.stack([tX[cii], tY[cii], tZ[cii]], axis=-1)
 
-        nr, nc = len(rii), len(cii)
-        if nr == nc:
-            lines = np.concatenate([row_lines, col_lines])
-        elif nr == 0:
-            lines = col_lines
-        elif nc == 0:
-            lines = row_lines
-        else:
-            lines = np.full((nr + nc, max(row_lines.shape[1], col_lines.shape[1]), 3),
-                            np.nan)
-            lines[:nr, :row_lines.shape[1], :] = row_lines
-            lines[nr:, :col_lines.shape[1], :] = col_lines
-
+        # We autoscale twice because autoscaling is much faster with vectorized numpy
+        # arrays, but row_lines and col_lines might not be the same shape, so we can't
+        # stack them to check them in a single pass.
+        # Note that while the column and row grid points are the same, the lines
+        # between them may expand the view limits, so we have to check both.
+        self.auto_scale_xyz(row_lines[..., 0], row_lines[..., 1], row_lines[..., 2],
+                            had_data)
+        self.auto_scale_xyz(col_lines[..., 0], col_lines[..., 1], col_lines[..., 2],
+                            had_data=True)
+
+        lines = list(row_lines) + list(col_lines)
         linec = art3d.Line3DCollection(lines, axlim_clip=axlim_clip, **kwargs)
         self.add_collection(linec)
-        self.auto_scale_xyz(lines[..., 0], lines[..., 1], lines[..., 2], had_data)
 
         return linec