Merge pull request matplotlib#29397 from scottshambaugh/masked_array_…

…performance

3D plotting performance improvements

doc/users/next_whats_new/3d_speedups.rst

@@ -0,0 +1,6 @@
+3D performance improvements
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Draw time for 3D plots has been improved, especially for surface and wireframe
+plots. Users should see up to a 10x speedup in some cases. This should make
+interacting with 3D plots much more responsive.

lib/mpl_toolkits/mplot3d/art3d.py

@@ -75,7 +75,7 @@ def get_dir_vector(zdir):
 
 def _viewlim_mask(xs, ys, zs, axes):
     """
-    Return original points with points outside the axes view limits masked.
+    Return the mask of the points outside the axes view limits.
 
     Parameters
     ----------
@@ -86,19 +86,16 @@ def _viewlim_mask(xs, ys, zs, axes):
 
     Returns
     -------
-    xs_masked, ys_masked, zs_masked : np.ma.array
-        The masked points.
+    mask : np.array
+        The mask of the points as a bool array.
     """
     mask = np.logical_or.reduce((xs < axes.xy_viewLim.xmin,
                                  xs > axes.xy_viewLim.xmax,
                                  ys < axes.xy_viewLim.ymin,
                                  ys > axes.xy_viewLim.ymax,
                                  zs < axes.zz_viewLim.xmin,
                                  zs > axes.zz_viewLim.xmax))
-    xs_masked = np.ma.array(xs, mask=mask)
-    ys_masked = np.ma.array(ys, mask=mask)
-    zs_masked = np.ma.array(zs, mask=mask)
-    return xs_masked, ys_masked, zs_masked
+    return mask
 
 
 class Text3D(mtext.Text):
@@ -182,14 +179,13 @@ def set_3d_properties(self, z=0, zdir='z', axlim_clip=False):
     @artist.allow_rasterization
     def draw(self, renderer):
         if self._axlim_clip:
-            xs, ys, zs = _viewlim_mask(self._x, self._y, self._z, self.axes)
-            position3d = np.ma.row_stack((xs, ys, zs)).ravel().filled(np.nan)
+            mask = _viewlim_mask(self._x, self._y, self._z, self.axes)
+            pos3d = np.ma.array([self._x, self._y, self._z],
+                                mask=mask, dtype=float).filled(np.nan)
         else:
-            xs, ys, zs = self._x, self._y, self._z
-            position3d = np.asanyarray([xs, ys, zs])
+            pos3d = np.array([self._x, self._y, self._z], dtype=float)
 
-        proj = proj3d._proj_trans_points(
-            [position3d, position3d + self._dir_vec], self.axes.M)
+        proj = proj3d._proj_trans_points([pos3d, pos3d + self._dir_vec], self.axes.M)
         dx = proj[0][1] - proj[0][0]
         dy = proj[1][1] - proj[1][0]
         angle = math.degrees(math.atan2(dy, dx))
@@ -313,7 +309,12 @@ def get_data_3d(self):
     @artist.allow_rasterization
     def draw(self, renderer):
         if self._axlim_clip:
-            xs3d, ys3d, zs3d = _viewlim_mask(*self._verts3d, self.axes)
+            mask = np.broadcast_to(
+                _viewlim_mask(*self._verts3d, self.axes),
+                (len(self._verts3d), *self._verts3d[0].shape)
+            )
+            xs3d, ys3d, zs3d = np.ma.array(self._verts3d,
+                                           dtype=float, mask=mask).filled(np.nan)
         else:
             xs3d, ys3d, zs3d = self._verts3d
         xs, ys, zs, tis = proj3d._proj_transform_clip(xs3d, ys3d, zs3d,
@@ -404,7 +405,8 @@ def do_3d_projection(self):
         """Project the points according to renderer matrix."""
         vs_list = [vs for vs, _ in self._3dverts_codes]
         if self._axlim_clip:
-            vs_list = [np.ma.row_stack(_viewlim_mask(*vs.T, self.axes)).T
+            vs_list = [np.ma.array(vs, mask=np.broadcast_to(
+                       _viewlim_mask(*vs.T, self.axes), vs.shape))
                        for vs in vs_list]
         xyzs_list = [proj3d.proj_transform(*vs.T, self.axes.M) for vs in vs_list]
         self._paths = [mpath.Path(np.ma.column_stack([xs, ys]), cs)
@@ -450,22 +452,32 @@ def do_3d_projection(self):
         """
         Project the points according to renderer matrix.
         """
-        segments = self._segments3d
+        segments = np.asanyarray(self._segments3d)
+
+        mask = False
+        if np.ma.isMA(segments):
+            mask = segments.mask
+
         if self._axlim_clip:
-            all_points = np.ma.vstack(segments)
-            masked_points = np.ma.column_stack([*_viewlim_mask(*all_points.T,
-                                                               self.axes)])
-            segment_lengths = [np.shape(segment)[0] for segment in segments]
-            segments = np.split(masked_points, np.cumsum(segment_lengths[:-1]))
-        xyslist = [proj3d._proj_trans_points(points, self.axes.M)
-                   for points in segments]
-        segments_2d = [np.ma.column_stack([xs, ys]) for xs, ys, zs in xyslist]
+            viewlim_mask = _viewlim_mask(segments[..., 0],
+                                         segments[..., 1],
+                                         segments[..., 2],
+                                         self.axes)
+            if np.any(viewlim_mask):
+                # broadcast mask to 3D
+                viewlim_mask = np.broadcast_to(viewlim_mask[..., np.newaxis],
+                                               (*viewlim_mask.shape, 3))
+                mask = mask | viewlim_mask
+        xyzs = np.ma.array(proj3d._proj_transform_vectors(segments, self.axes.M),
+                           mask=mask)
+        segments_2d = xyzs[..., 0:2]
         LineCollection.set_segments(self, segments_2d)
 
         # FIXME
-        minz = 1e9
-        for xs, ys, zs in xyslist:
-            minz = min(minz, min(zs))
+        if len(xyzs) > 0:
+            minz = min(xyzs[..., 2].min(), 1e9)
+        else:
+            minz = np.nan
         return minz
 
 
@@ -531,7 +543,9 @@ def get_path(self):
     def do_3d_projection(self):
         s = self._segment3d
         if self._axlim_clip:
-            xs, ys, zs = _viewlim_mask(*zip(*s), self.axes)
+            mask = _viewlim_mask(*zip(*s), self.axes)
+            xs, ys, zs = np.ma.array(zip(*s),
+                                     dtype=float, mask=mask).filled(np.nan)
         else:
             xs, ys, zs = zip(*s)
         vxs, vys, vzs, vis = proj3d._proj_transform_clip(xs, ys, zs,
@@ -587,7 +601,9 @@ def set_3d_properties(self, path, zs=0, zdir='z', axlim_clip=False):
     def do_3d_projection(self):
         s = self._segment3d
         if self._axlim_clip:
-            xs, ys, zs = _viewlim_mask(*zip(*s), self.axes)
+            mask = _viewlim_mask(*zip(*s), self.axes)
+            xs, ys, zs = np.ma.array(zip(*s),
+                                     dtype=float, mask=mask).filled(np.nan)
         else:
             xs, ys, zs = zip(*s)
         vxs, vys, vzs, vis = proj3d._proj_transform_clip(xs, ys, zs,
@@ -701,14 +717,18 @@ def set_3d_properties(self, zs, zdir, axlim_clip=False):
 
     def do_3d_projection(self):
         if self._axlim_clip:
-            xs, ys, zs = _viewlim_mask(*self._offsets3d, self.axes)
+            mask = _viewlim_mask(*self._offsets3d, self.axes)
+            xs, ys, zs = np.ma.array(self._offsets3d, mask=mask)
         else:
             xs, ys, zs = self._offsets3d
         vxs, vys, vzs, vis = proj3d._proj_transform_clip(xs, ys, zs,
                                                          self.axes.M,
                                                          self.axes._focal_length)
         self._vzs = vzs
-        super().set_offsets(np.ma.column_stack([vxs, vys]))
+        if np.ma.isMA(vxs):
+            super().set_offsets(np.ma.column_stack([vxs, vys]))
+        else:
+            super().set_offsets(np.column_stack([vxs, vys]))
 
         if vzs.size > 0:
             return min(vzs)
@@ -851,11 +871,18 @@ def set_depthshade(self, depthshade):
         self.stale = True
 
     def do_3d_projection(self):
+        mask = False
+        for xyz in self._offsets3d:
+            if np.ma.isMA(xyz):
+                mask = mask | xyz.mask
         if self._axlim_clip:
-            xs, ys, zs = _viewlim_mask(*self._offsets3d, self.axes)
+            mask = mask | _viewlim_mask(*self._offsets3d, self.axes)
+            mask = np.broadcast_to(mask,
+                                   (len(self._offsets3d), *self._offsets3d[0].shape))
+            xyzs = np.ma.array(self._offsets3d, mask=mask)
         else:
-            xs, ys, zs = self._offsets3d
-        vxs, vys, vzs, vis = proj3d._proj_transform_clip(xs, ys, zs,
+            xyzs = self._offsets3d
+        vxs, vys, vzs, vis = proj3d._proj_transform_clip(*xyzs,
                                                          self.axes.M,
                                                          self.axes._focal_length)
         # Sort the points based on z coordinates
@@ -1062,16 +1089,37 @@ def get_vector(self, segments3d):
         return self._get_vector(segments3d)
 
     def _get_vector(self, segments3d):
-        """Optimize points for projection."""
-        if len(segments3d):
-            xs, ys, zs = np.vstack(segments3d).T
-        else:  # vstack can't stack zero arrays.
-            xs, ys, zs = [], [], []
-        ones = np.ones(len(xs))
-        self._vec = np.array([xs, ys, zs, ones])
+        """
+        Optimize points for projection.
 
-        indices = [0, *np.cumsum([len(segment) for segment in segments3d])]
-        self._segslices = [*map(slice, indices[:-1], indices[1:])]
+        Parameters
+        ----------
+        segments3d : NumPy array or list of NumPy arrays
+            List of vertices of the boundary of every segment. If all paths are
+            of equal length and this argument is a NumPy array, then it should
+            be of shape (num_faces, num_vertices, 3).
+        """
+        if isinstance(segments3d, np.ndarray):
+            if segments3d.ndim != 3 or segments3d.shape[-1] != 3:
+                raise ValueError("segments3d must be a MxNx3 array, but got "
+                                 f"shape {segments3d.shape}")
+            if isinstance(segments3d, np.ma.MaskedArray):
+                self._faces = segments3d.data
+                self._invalid_vertices = segments3d.mask.any(axis=-1)
+            else:
+                self._faces = segments3d
+                self._invalid_vertices = False
+        else:
+            # Turn the potentially ragged list into a numpy array for later speedups
+            # If it is ragged, set the unused vertices per face as invalid
+            num_faces = len(segments3d)
+            num_verts = np.fromiter(map(len, segments3d), dtype=np.intp)
+            max_verts = num_verts.max(initial=0)
+            segments = np.empty((num_faces, max_verts, 3))
+            for i, face in enumerate(segments3d):
+                segments[i, :len(face)] = face
+            self._faces = segments
+            self._invalid_vertices = np.arange(max_verts) >= num_verts[:, None]
 
     def set_verts(self, verts, closed=True):
         """
@@ -1133,64 +1181,85 @@ def do_3d_projection(self):
                 self._facecolor3d = self._facecolors
             if self._edge_is_mapped:
                 self._edgecolor3d = self._edgecolors
+
+        needs_masking = np.any(self._invalid_vertices)
+        num_faces = len(self._faces)
+        mask = self._invalid_vertices
+
+        # Some faces might contain masked vertices, so we want to ignore any
+        # errors that those might cause
+        with np.errstate(invalid='ignore', divide='ignore'):
+            pfaces = proj3d._proj_transform_vectors(self._faces, self.axes.M)
+
         if self._axlim_clip:
-            xs, ys, zs = _viewlim_mask(*self._vec[0:3], self.axes)
-            if self._vec.shape[0] == 4:  # Will be 3 (xyz) or 4 (xyzw)
-                w_masked = np.ma.masked_where(zs.mask, self._vec[3])
-                vec = np.ma.array([xs, ys, zs, w_masked])
-            else:
-                vec = np.ma.array([xs, ys, zs])
-        else:
-            vec = self._vec
-        txs, tys, tzs = proj3d._proj_transform_vec(vec, self.axes.M)
-        xyzlist = [(txs[sl], tys[sl], tzs[sl]) for sl in self._segslices]
+            viewlim_mask = _viewlim_mask(self._faces[..., 0], self._faces[..., 1],
+                                         self._faces[..., 2], self.axes)
+            if np.any(viewlim_mask):
+                needs_masking = True
+                mask = mask | viewlim_mask
+
+        pzs = pfaces[..., 2]
+        if needs_masking:
+            pzs = np.ma.MaskedArray(pzs, mask=mask)
 
         # This extra fuss is to re-order face / edge colors
         cface = self._facecolor3d
         cedge = self._edgecolor3d
-        if len(cface) != len(xyzlist):
-            cface = cface.repeat(len(xyzlist), axis=0)
-        if len(cedge) != len(xyzlist):
+        if len(cface) != num_faces:
+            cface = cface.repeat(num_faces, axis=0)
+        if len(cedge) != num_faces:
             if len(cedge) == 0:
                 cedge = cface
             else:
-                cedge = cedge.repeat(len(xyzlist), axis=0)
-
-        if xyzlist:
-            # sort by depth (furthest drawn first)
-            z_segments_2d = sorted(
-                ((self._zsortfunc(zs.data), np.ma.column_stack([xs, ys]), fc, ec, idx)
-                 for idx, ((xs, ys, zs), fc, ec)
-                 in enumerate(zip(xyzlist, cface, cedge))),
-                key=lambda x: x[0], reverse=True)
-
-            _, segments_2d, self._facecolors2d, self._edgecolors2d, idxs = \
-                zip(*z_segments_2d)
-        else:
-            segments_2d = []
-            self._facecolors2d = np.empty((0, 4))
-            self._edgecolors2d = np.empty((0, 4))
-            idxs = []
-
-        if self._codes3d is not None:
-            codes = [self._codes3d[idx] for idx in idxs]
-            PolyCollection.set_verts_and_codes(self, segments_2d, codes)
+                cedge = cedge.repeat(num_faces, axis=0)
+
+        if len(pzs) > 0:
+            face_z = self._zsortfunc(pzs, axis=-1)
         else:
-            PolyCollection.set_verts(self, segments_2d, self._closed)
+            face_z = pzs
+        if needs_masking:
+            face_z = face_z.data
+        face_order = np.argsort(face_z, axis=-1)[::-1]
 
-        if len(self._edgecolor3d) != len(cface):
+        if len(pfaces) > 0:
+            faces_2d = pfaces[face_order, :, :2]
+        else:
+            faces_2d = pfaces
+        if self._codes3d is not None and len(self._codes3d) > 0:
+            if needs_masking:
+                segment_mask = ~mask[face_order, :]
+                faces_2d = [face[mask, :] for face, mask
+                               in zip(faces_2d, segment_mask)]
+            codes = [self._codes3d[idx] for idx in face_order]
+            PolyCollection.set_verts_and_codes(self, faces_2d, codes)
+        else:
+            if needs_masking and len(faces_2d) > 0:
+                invalid_vertices_2d = np.broadcast_to(
+                    mask[face_order, :, None],
+                    faces_2d.shape)
+                faces_2d = np.ma.MaskedArray(
+                        faces_2d, mask=invalid_vertices_2d)
+            PolyCollection.set_verts(self, faces_2d, self._closed)
+
+        if len(cface) > 0:
+            self._facecolors2d = cface[face_order]
+        else:
+            self._facecolors2d = cface
+        if len(self._edgecolor3d) == len(cface) and len(cedge) > 0:
+            self._edgecolors2d = cedge[face_order]
+        else:
             self._edgecolors2d = self._edgecolor3d
 
         # Return zorder value
         if self._sort_zpos is not None:
             zvec = np.array([[0], [0], [self._sort_zpos], [1]])
             ztrans = proj3d._proj_transform_vec(zvec, self.axes.M)
             return ztrans[2][0]
-        elif tzs.size > 0:
+        elif pzs.size > 0:
             # FIXME: Some results still don't look quite right.
             #        In particular, examine contourf3d_demo2.py
             #        with az = -54 and elev = -45.
-            return np.min(tzs)
+            return np.min(pzs)
         else:
             return np.nan
 

lib/mpl_toolkits/mplot3d/axes3d.py

@@ -2892,7 +2892,9 @@ def add_collection3d(self, col, zs=0, zdir='z', autolim=True, *,
                 self.auto_scale_xyz(*np.array(col._segments3d).transpose(),
                                     had_data=had_data)
             elif isinstance(col, art3d.Poly3DCollection):
-                self.auto_scale_xyz(*col._vec[:-1], had_data=had_data)
+                self.auto_scale_xyz(col._faces[..., 0],
+                                    col._faces[..., 1],
+                                    col._faces[..., 2], had_data=had_data)
             elif isinstance(col, art3d.Patch3DCollection):
                 pass
                 # FIXME: Implement auto-scaling function for Patch3DCollection