diff --git a/doc/users/next_whats_new/3d_speedups.rst b/doc/users/next_whats_new/3d_speedups.rst new file mode 100644 index 000000000000..70e80bfbdccb --- /dev/null +++ b/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. diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py index deb0ca34302c..88c2ae625c00 100644 --- a/lib/mpl_toolkits/mplot3d/art3d.py +++ b/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,8 +86,8 @@ 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, @@ -95,10 +95,7 @@ def _viewlim_mask(xs, ys, zs, axes): 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,52 +1181,73 @@ 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 @@ -1186,11 +1255,11 @@ def do_3d_projection(self): 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 diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py index 45d38a346a69..8ec63aba3349 100644 --- a/lib/mpl_toolkits/mplot3d/axes3d.py +++ b/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 diff --git a/lib/mpl_toolkits/mplot3d/proj3d.py b/lib/mpl_toolkits/mplot3d/proj3d.py index 923bd32c9ce0..34a03969c961 100644 --- a/lib/mpl_toolkits/mplot3d/proj3d.py +++ b/lib/mpl_toolkits/mplot3d/proj3d.py @@ -133,15 +133,33 @@ def _ortho_transformation(zfront, zback): def _proj_transform_vec(vec, M): vecw = np.dot(M, vec.data) - w = vecw[3] - txs, tys, tzs = vecw[0]/w, vecw[1]/w, vecw[2]/w - if np.ma.isMA(vec[0]): # we check each to protect for scalars - txs = np.ma.array(txs, mask=vec[0].mask) - if np.ma.isMA(vec[1]): - tys = np.ma.array(tys, mask=vec[1].mask) - if np.ma.isMA(vec[2]): - tzs = np.ma.array(tzs, mask=vec[2].mask) - return txs, tys, tzs + ts = vecw[0:3]/vecw[3] + if np.ma.isMA(vec): + ts = np.ma.array(ts, mask=vec.mask) + return ts[0], ts[1], ts[2] + + +def _proj_transform_vectors(vecs, M): + """ + Vectorized version of ``_proj_transform_vec``. + + Parameters + ---------- + vecs : ... x 3 np.ndarray + Input vectors + M : 4 x 4 np.ndarray + Projection matrix + """ + vecs_shape = vecs.shape + vecs = vecs.reshape(-1, 3).T + + vecs_pad = np.empty((vecs.shape[0] + 1,) + vecs.shape[1:]) + vecs_pad[:-1] = vecs + vecs_pad[-1] = 1 + product = np.dot(M, vecs_pad) + tvecs = product[:3] / product[3] + + return tvecs.T.reshape(vecs_shape) def _proj_transform_vec_clip(vec, M, focal_length): diff --git a/lib/mpl_toolkits/mplot3d/tests/baseline_images/test_axes3d/surface3d_masked.png b/lib/mpl_toolkits/mplot3d/tests/baseline_images/test_axes3d/surface3d_masked.png index df893f9c843f..9e5af36ffbfc 100644 Binary files a/lib/mpl_toolkits/mplot3d/tests/baseline_images/test_axes3d/surface3d_masked.png and b/lib/mpl_toolkits/mplot3d/tests/baseline_images/test_axes3d/surface3d_masked.png differ diff --git a/lib/mpl_toolkits/mplot3d/tests/test_axes3d.py b/lib/mpl_toolkits/mplot3d/tests/test_axes3d.py index 0852a797c946..7067576e358a 100644 --- a/lib/mpl_toolkits/mplot3d/tests/test_axes3d.py +++ b/lib/mpl_toolkits/mplot3d/tests/test_axes3d.py @@ -996,7 +996,7 @@ def test_poly3dcollection_verts_validation(): art3d.Poly3DCollection(poly) # should be Poly3DCollection([poly]) poly = np.array(poly, dtype=float) - with pytest.raises(ValueError, match=r'list of \(N, 3\) array-like'): + with pytest.raises(ValueError, match=r'MxNx3 array'): art3d.Poly3DCollection(poly) # should be Poly3DCollection([poly])