Skip to content

3D plotting performance improvements #29397

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

Merged
timhoffm merged 14 commits into from
Jan 28, 2025
Merged

3D plotting performance improvements #29397

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
9f66076
Poly3DCollection masking speedups
scottshambaugh Jan 2, 2025
94c6ca7
3d text masking
scottshambaugh Jan 2, 2025
606e76f
Single 3d object masking
scottshambaugh Jan 2, 2025
b48cdcf
Line3D collection and Path3DCollection masking
scottshambaugh Jan 2, 2025
0bf8b7b
Patch3DCollection masking
scottshambaugh Jan 3, 2025
fdcc0d8
simplify projection
scottshambaugh Jan 3, 2025
e0ff208
Fix Poly3Dcollection autoscaling
scottshambaugh Jan 3, 2025
80a5a95
Collection3D masking
scottshambaugh Jan 3, 2025
d18c52b
fix Line3D masking
scottshambaugh Jan 3, 2025
1782309
Fix tests
scottshambaugh Jan 3, 2025
a700fe9
3d speedups whats new
scottshambaugh Jan 6, 2025
52a06a6
Code review updates
scottshambaugh Jan 7, 2025
68b8a6c
Ignore div0 errors on masked vertices
scottshambaugh Jan 8, 2025
4fc3745
code review updates
scottshambaugh Jan 27, 2025
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
6 changes: 6 additions & 0 deletions doc/users/next_whats_new/3d_speedups.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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.
231 changes: 150 additions & 81 deletions lib/mpl_toolkits/mplot3d/art3d.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,7 @@ def get_dir_vector(zdir):

def _viewlim_mask(xs, ys, zs, axes):
Comment thread
scottshambaugh marked this conversation as resolved.
"""
Return original points with points outside the axes view limits masked.
Return the mask of the points outside the axes view limits.

Parameters
----------
Expand All @@ -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):
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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)
Comment thread
scottshambaugh marked this conversation as resolved.
else:
xs3d, ys3d, zs3d = self._verts3d
xs, ys, zs, tis = proj3d._proj_transform_clip(xs3d, ys3d, zs3d,
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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


Expand Down Expand Up @@ -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),
Comment thread
scottshambaugh marked this conversation as resolved.
dtype=float, mask=mask).filled(np.nan)
else:
xs, ys, zs = zip(*s)
vxs, vys, vzs, vis = proj3d._proj_transform_clip(xs, ys, zs,
Expand Down Expand Up @@ -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),
Comment thread
scottshambaugh marked this conversation as resolved.
dtype=float, mask=mask).filled(np.nan)
else:
xs, ys, zs = zip(*s)
vxs, vys, vzs, vis = proj3d._proj_transform_clip(xs, ys, zs,
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Comment thread
scottshambaugh marked this conversation as resolved.
# 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):
"""
Expand Down Expand Up @@ -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
Comment on lines +1189 to +1190
Copy link
Copy Markdown
Contributor

@greglucas greglucas Jan 7, 2025

Choose a reason for hiding this comment

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

Should these errors be handled within the _proj_transform_vectors case directly closer to where these errors would actually arise? I know there has been some work in numpy to not raise invalid/division warnings on masked elements, so is this still an issue now?

Copy link
Copy Markdown
Contributor Author

@scottshambaugh scottshambaugh Jan 7, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

I removed the error ignore line and did not see any issues. The tests cover partially masked Poly3DCollections, so that case is exercised, and it looks like the improved masked handling covers any previous errors here.

Copy link
Copy Markdown
Contributor Author

@scottshambaugh scottshambaugh Jan 8, 2025

Choose a reason for hiding this comment

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

Looks like the error does crop up in the AppVeyor build, so adding the original error suppression back in. I don't want to put it in _proj_transform_vectors since it might not always be the case that we know that we have masked out the problematic vertices outside of that function.

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

Expand Down
4 changes: 3 additions & 1 deletion lib/mpl_toolkits/mplot3d/axes3d.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading