diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py index 206641e5814f..41a26f0ef8ff 100644 --- a/lib/mpl_toolkits/mplot3d/art3d.py +++ b/lib/mpl_toolkits/mplot3d/art3d.py @@ -598,16 +598,37 @@ def set_zsort(self, zsort): self.stale = True def get_vector(self, segments3d): - """Optimize points for projection.""" - if len(segments3d): - xs, ys, zs = np.row_stack(segments3d).T - else: # row_stack 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 arrray, 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 " + + "shape {}".format(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: + 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] + assert self._invalid_vertices is False or \ + self._invalid_vertices.shape == self._faces.shape[:-1] def set_verts(self, verts, closed=True): """Set 3D vertices.""" @@ -649,37 +670,53 @@ def do_3d_projection(self, renderer): self.update_scalarmappable() self._facecolors3d = self._facecolors - txs, tys, tzs = proj3d._proj_transform_vec(self._vec, renderer.M) - xyzlist = [(txs[sl], tys[sl], tzs[sl]) for sl in self._segslices] + needs_masking = self._invalid_vertices is not False + num_faces = len(self._faces) + + # 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, renderer.M) + pzs = pfaces[..., 2] + if needs_masking: + pzs = np.ma.MaskedArray(pzs, mask=self._invalid_vertices) # This extra fuss is to re-order face / edge colors cface = self._facecolors3d cedge = self._edgecolors3d - 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) + cedge = cedge.repeat(num_faces, axis=0) - # sort by depth (furthest drawn first) - z_segments_2d = sorted( - ((self._zsortfunc(zs), np.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) + face_z = self._zsortfunc(pzs, axis=-1) + if needs_masking: + face_z = face_z.data + face_order = np.argsort(face_z, axis=-1)[::-1] - segments_2d = [s for z, s, fc, ec, idx in z_segments_2d] + faces_2d = pfaces[face_order, :, :2] if self._codes3d is not None: - codes = [self._codes3d[idx] for z, s, fc, ec, idx in z_segments_2d] - PolyCollection.set_verts_and_codes(self, segments_2d, codes) + if needs_masking: + segment_mask = ~self._invalid_vertices[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: - PolyCollection.set_verts(self, segments_2d, self._closed) - - self._facecolors2d = [fc for z, s, fc, ec, idx in z_segments_2d] + if needs_masking: + invalid_vertices_2d = np.broadcast_to( + self._invalid_vertices[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) + + self._facecolors2d = cface[face_order] if len(self._edgecolors3d) == len(cface): - self._edgecolors2d = [ec for z, s, fc, ec, idx in z_segments_2d] + self._edgecolors2d = cedge[face_order] else: self._edgecolors2d = self._edgecolors3d @@ -688,11 +725,11 @@ def do_3d_projection(self, renderer): zvec = np.array([[0], [0], [self._sort_zpos], [1]]) ztrans = proj3d._proj_transform_vec(zvec, renderer.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/proj3d.py b/lib/mpl_toolkits/mplot3d/proj3d.py index e2bb1cf69571..46ee44dec098 100644 --- a/lib/mpl_toolkits/mplot3d/proj3d.py +++ b/lib/mpl_toolkits/mplot3d/proj3d.py @@ -103,10 +103,8 @@ def ortho_transformation(zfront, zback): def _proj_transform_vec(vec, M): vecw = np.dot(M, vec) - w = vecw[3] # clip here.. - txs, tys, tzs = vecw[0]/w, vecw[1]/w, vecw[2]/w - return txs, tys, tzs + return vecw[:3] / vecw[3] def _proj_transform_vec_clip(vec, M): @@ -146,6 +144,29 @@ def proj_transform(xs, ys, zs, M): transform = proj_transform +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_clip(xs, ys, zs, M): """ Transform the points by the projection matrix