Skip to content

Store vertices as (masked) arrays in Poly3DCollection #16688

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 4 commits into from
Closed

Store vertices as (masked) arrays in Poly3DCollection #16688

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
8633bbd
Store vertices as (masked) arrays in Poly3DCollection
apaszke Mar 5, 2020
06105a1
Don't store faces as masked arrays
apaszke Mar 23, 2020
076bfcb
Update lib/mpl_toolkits/mplot3d/proj3d.py
apaszke Mar 23, 2020
301b6fe
Ignore floating point errors
apaszke Mar 23, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
97 changes: 67 additions & 30 deletions lib/mpl_toolkits/mplot3d/art3d.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down Expand Up @@ -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)
Comment on lines +680 to +682
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tempting to move this down to line 692, where it's used for the first time.

Copy link
Contributor Author

@apaszke apaszke Mar 23, 2020 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I felt like it's logically connected to pfaces and it's used on lines 692, 725 and 729, so I didn't want to make the first use special in any way.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I suggested this because it might be worth adding an optimization today along the lines of:

if numpy_version >= LooseVersion(1.17.0):
   # use the where argument, skip the masked array entirely
else:
   ...

Happy to have that declared as out of scope

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm happy to do that, but I'm still a bit afraid that it will make the code even less readable. I'd have to insert the same if on lines 692 and 729. WDYT?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe try it and see if it actually helps with performance locally in any significant way? If not, then I think it's not worth it.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've tried this, and the where= reductions are indeed faster!

In [1]: import numpy as np                                            
                                                  
In [2]: maskf = np.zeros((399 * 399, 4), dtype=bool)                  
                                                  
In [3]: maskt = np.ones((399 * 399, 4), dtype=bool)                   
                         
In [4]: x = np.ones((399 * 399, 4), dtype=float)                      
                         
In [5]: mx = np.ma.MaskedArray(x, maskf)          
                                                  
In [6]: %timeit np.max(mx, axis=-1)                                   
10.3 ms ± 441 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [7]: %timeit np.max(x, axis=-1, where=maskt, initial=0)            
7.09 ms ± 139 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

On the other hand, it looks like np.average doesn't support the where= argument, so I don't think we can do this...

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh, that's a good catch. You can almost get away with passing the mask into the weights argument of np.average, but that behaves poorly on NaNs.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hah, good point! And it looks like it still beats the masked array implementation, which is very surprising to me:

In [9]: %timeit np.average(x, axis=-1, weights=maskt)                                                
8.69 ms ± 137 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [10]: %timeit np.average(mx, axis=-1)          
12.9 ms ± 370 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

But yes, there's an issue with invalid floating point values and there's no default argument we can pass in. I guess the latter problem could be fixed by using a wrapper like lambda *args, **kwargs: np.average(*args, weight=kwargs.pop('where', 1.), **kwargs) in the dict defining zsort functions.

So what's the final decision?


# 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)
Copy link
Contributor

@eric-wieser eric-wieser Mar 23, 2020 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Perhaps a marker for future cleanup would be handy here.

Suggested change
face_z = self._zsortfunc(pzs, axis=-1)
# TODO: When we drop numpy < 1.17, investigate using the `where` argument to reductions
face_z = self._zsortfunc(pzs, axis=-1)

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since matplotlib 3.5 we require numpy >= 1.17, so this can be updated immediately.

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

Expand All @@ -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

Expand Down
27 changes: 24 additions & 3 deletions lib/mpl_toolkits/mplot3d/proj3d.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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
Expand Down