From 8633bbd22d2be5581baec6ce221c892f26a1572b Mon Sep 17 00:00:00 2001 From: Adam Paszke Date: Thu, 5 Mar 2020 21:53:18 +0100 Subject: [PATCH 1/4] Store vertices as (masked) arrays in Poly3DCollection This removes a bunch of very expensive list comprehensions, dropping rendering times by half for some 3D surface plots. --- lib/mpl_toolkits/mplot3d/art3d.py | 79 ++++++++++++++++++++---------- lib/mpl_toolkits/mplot3d/proj3d.py | 33 +++++++++++-- 2 files changed, 82 insertions(+), 30 deletions(-) diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py index 206641e5814f..6ef0a6a40634 100644 --- a/lib/mpl_toolkits/mplot3d/art3d.py +++ b/lib/mpl_toolkits/mplot3d/art3d.py @@ -598,16 +598,33 @@ 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)) + self._segments = segments3d + else: + num_faces = len(segments3d) + num_verts = np.fromiter(map(len, segments3d), dtype=np.intp) + max_verts = num_verts.max(initial=0) + padded = np.empty((num_faces, max_verts, 3)) + for i, face in enumerate(segments3d): + padded[i, :len(face)] = face + mask = np.arange(max_verts) >= num_verts[:, None] + mask = mask[..., None] # add a component axis + # ma.array does not broadcast the mask for us + mask = np.broadcast_to(mask, padded.shape) + self._segments = np.ma.array(padded, mask=mask) def set_verts(self, verts, closed=True): """Set 3D vertices.""" @@ -649,37 +666,45 @@ 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] + psegments = proj3d._proj_transform_vectors(self._segments, renderer.M) + is_masked = isinstance(psegments, np.ma.MaskedArray) + num_faces = len(psegments) # 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(psegments[..., 2], axis=-1) + if is_masked: + # NOTE: Unpacking .data is safe here, because every face has to + # contain a valid vertex. + 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] + segments_2d = psegments[face_order, :, :2] if self._codes3d is not None: - codes = [self._codes3d[idx] for z, s, fc, ec, idx in z_segments_2d] + if is_masked: + # NOTE: We cannot assert the same on segments_2d, as it is a + # result of advanced indexing, so its mask is a newly + # allocated tensor. However, both of those asserts are + # equivalent. + assert psegments.mask.strides[-1] == 0 + segments_2d = [s.compressed().reshape(-1, 2) for s in segments_2d] + codes = [self._codes3d[idx] for idx in face_order] PolyCollection.set_verts_and_codes(self, segments_2d, codes) else: PolyCollection.set_verts(self, segments_2d, self._closed) - self._facecolors2d = [fc for z, s, fc, ec, idx in z_segments_2d] + 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 +713,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 psegments.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(psegments[..., 2]) else: return np.nan diff --git a/lib/mpl_toolkits/mplot3d/proj3d.py b/lib/mpl_toolkits/mplot3d/proj3d.py index e2bb1cf69571..988876d5d389 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,35 @@ def proj_transform(xs, ys, zs, M): transform = proj_transform +def _proj_transform_vectors(vecs, M): + """Vector version of ``project_transform`` able to handle MaskedArrays. + + Parameters + ---------- + vecs : ... x 3 np.ndarray or np.ma.MaskedArray + Input vectors + + M : 4 x 4 np.ndarray + Projection matrix + """ + vecs_shape = vecs.shape + vecs = vecs.reshape(-1, 3).T + + is_masked = isinstance(vecs, np.ma.MaskedArray) + if is_masked: + mask = vecs.mask + + 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] + + if is_masked: + tvecs = np.ma.array(tvecs, mask=mask) + return tvecs.T.reshape(vecs_shape) + + def proj_transform_clip(xs, ys, zs, M): """ Transform the points by the projection matrix From 06105a177458bc02e1467266a30cdd2d565633d4 Mon Sep 17 00:00:00 2001 From: Adam Paszke Date: Mon, 23 Mar 2020 13:42:29 +0100 Subject: [PATCH 2/4] Don't store faces as masked arrays --- lib/mpl_toolkits/mplot3d/art3d.py | 63 +++++++++++++++++------------- lib/mpl_toolkits/mplot3d/proj3d.py | 8 +--- 2 files changed, 37 insertions(+), 34 deletions(-) diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py index 6ef0a6a40634..0abad0eddffe 100644 --- a/lib/mpl_toolkits/mplot3d/art3d.py +++ b/lib/mpl_toolkits/mplot3d/art3d.py @@ -612,19 +612,23 @@ def get_vector(self, segments3d): if segments3d.ndim != 3 or segments3d.shape[-1] != 3: raise ValueError("segments3d must be a MxNx3 array, but got " + "shape {}".format(segments3d.shape)) - self._segments = segments3d + 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) - padded = np.empty((num_faces, max_verts, 3)) + segments = np.empty((num_faces, max_verts, 3)) for i, face in enumerate(segments3d): - padded[i, :len(face)] = face - mask = np.arange(max_verts) >= num_verts[:, None] - mask = mask[..., None] # add a component axis - # ma.array does not broadcast the mask for us - mask = np.broadcast_to(mask, padded.shape) - self._segments = np.ma.array(padded, mask=mask) + 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.""" @@ -666,9 +670,13 @@ def do_3d_projection(self, renderer): self.update_scalarmappable() self._facecolors3d = self._facecolors - psegments = proj3d._proj_transform_vectors(self._segments, renderer.M) - is_masked = isinstance(psegments, np.ma.MaskedArray) - num_faces = len(psegments) + needs_masking = self._invalid_vertices is not False + num_faces = len(self._faces) + + 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 @@ -681,26 +689,27 @@ def do_3d_projection(self, renderer): else: cedge = cedge.repeat(num_faces, axis=0) - face_z = self._zsortfunc(psegments[..., 2], axis=-1) - if is_masked: - # NOTE: Unpacking .data is safe here, because every face has to - # contain a valid vertex. + 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 = psegments[face_order, :, :2] + faces_2d = pfaces[face_order, :, :2] if self._codes3d is not None: - if is_masked: - # NOTE: We cannot assert the same on segments_2d, as it is a - # result of advanced indexing, so its mask is a newly - # allocated tensor. However, both of those asserts are - # equivalent. - assert psegments.mask.strides[-1] == 0 - segments_2d = [s.compressed().reshape(-1, 2) for s in segments_2d] + 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, segments_2d, codes) + PolyCollection.set_verts_and_codes(self, faces_2d, codes) else: - PolyCollection.set_verts(self, segments_2d, self._closed) + 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): @@ -713,11 +722,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 psegments.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(psegments[..., 2]) + 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 988876d5d389..039f05db3c0f 100644 --- a/lib/mpl_toolkits/mplot3d/proj3d.py +++ b/lib/mpl_toolkits/mplot3d/proj3d.py @@ -145,7 +145,7 @@ def proj_transform(xs, ys, zs, M): def _proj_transform_vectors(vecs, M): - """Vector version of ``project_transform`` able to handle MaskedArrays. + """Vectorized version of ``_proj_transform_vec``. Parameters ---------- @@ -158,18 +158,12 @@ def _proj_transform_vectors(vecs, M): vecs_shape = vecs.shape vecs = vecs.reshape(-1, 3).T - is_masked = isinstance(vecs, np.ma.MaskedArray) - if is_masked: - mask = vecs.mask - 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] - if is_masked: - tvecs = np.ma.array(tvecs, mask=mask) return tvecs.T.reshape(vecs_shape) From 076bfcbef20276d7cdd745de565d3b1208504dd9 Mon Sep 17 00:00:00 2001 From: Adam Paszke Date: Mon, 23 Mar 2020 14:05:22 +0100 Subject: [PATCH 3/4] Update lib/mpl_toolkits/mplot3d/proj3d.py Co-Authored-By: Eric Wieser --- lib/mpl_toolkits/mplot3d/proj3d.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/mpl_toolkits/mplot3d/proj3d.py b/lib/mpl_toolkits/mplot3d/proj3d.py index 039f05db3c0f..46ee44dec098 100644 --- a/lib/mpl_toolkits/mplot3d/proj3d.py +++ b/lib/mpl_toolkits/mplot3d/proj3d.py @@ -149,7 +149,7 @@ def _proj_transform_vectors(vecs, M): Parameters ---------- - vecs : ... x 3 np.ndarray or np.ma.MaskedArray + vecs : ... x 3 np.ndarray Input vectors M : 4 x 4 np.ndarray From 301b6fe59d311a942a4fdb5da55dbacf5a4d70a6 Mon Sep 17 00:00:00 2001 From: Adam Paszke Date: Mon, 23 Mar 2020 15:56:31 +0100 Subject: [PATCH 4/4] Ignore floating point errors --- lib/mpl_toolkits/mplot3d/art3d.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py index 0abad0eddffe..41a26f0ef8ff 100644 --- a/lib/mpl_toolkits/mplot3d/art3d.py +++ b/lib/mpl_toolkits/mplot3d/art3d.py @@ -673,7 +673,10 @@ def do_3d_projection(self, renderer): needs_masking = self._invalid_vertices is not False num_faces = len(self._faces) - pfaces = proj3d._proj_transform_vectors(self._faces, renderer.M) + # 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)