Skip to content

plot_wireframe plotting speedup #29399

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.

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

Conversation

scottshambaugh
Copy link
Contributor

@scottshambaugh scottshambaugh commented Jan 4, 2025

PR summary

For large wireframe plots (in my testing, 8000x8000), a large chunk of time is spent doing autoscaling of the axes. There is some opportunity to vectorize these calculations and get speedups.

Changes here:

  1. Vectorize bounding box limits updates
  2. Faster array creation for transforms.update_from_data_x and transforms.update_from_data_y
  3. Switch from list comprehensions to vectorized approach for axes3d.plot_wireframe. This also stores the data as a numpy array, which gives speedups when drawing.

The compiled function update_path_extents is actually just about equal with a vectorized numpy approach speedwise when updating both x and y limits. But when doing single axis limits, it needs to check for NaNs for each pair of x or y, so the numpy approach lets you skip the extra x or y calculations and comes out ahead. This bit is a benefit for all 3D functions, not just plot_wireframe. I spent some time trying to refactor the compiled function to operate on just a single axis, but it's tricky to disentangle x and y. The numpy approach was easier and just as fast.

For my test case I'm seeing an average over 5 runs on main of 6.82 seconds, and on this PR 4.00 seconds - a 1.7x speedup.

import time
import numpy as np
import matplotlib.pyplot as plt

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

x = y = np.linspace(-1, 1, 8000)
X, Y = np.meshgrid(x, y)
Z = X ** 2 + Y ** 2

print("Timing...")

start_time = time.perf_counter()
ax.plot_wireframe(X, Y, Z, axlim_clip=False)
ax.set_xlim(0, 1)
fig.canvas.draw()
end_time = time.perf_counter()

plt.close()

print(f"Time taken: {end_time - start_time:.4f} seconds")

Plotting and drawing on main (avg 6.82 sec):
image

Plotting and drawing on this PR (avg 4.00 sec):
image

PR checklist

@scottshambaugh scottshambaugh added this to the v3.11.0 milestone Jan 4, 2025
@scottshambaugh scottshambaugh marked this pull request as draft January 4, 2025 01:19
@scottshambaugh scottshambaugh force-pushed the plot_wireframe_speedup branch 7 times, most recently from 3b832d3 to 8cc938d Compare January 5, 2025 19:22
@scottshambaugh scottshambaugh force-pushed the plot_wireframe_speedup branch 7 times, most recently from ea4a0dc to dfe72ed Compare January 5, 2025 22:23
@scottshambaugh scottshambaugh marked this pull request as ready for review January 5, 2025 22:34
@scottshambaugh scottshambaugh force-pushed the plot_wireframe_speedup branch from dfe72ed to 703ef23 Compare January 5, 2025 22:43
@scottshambaugh scottshambaugh force-pushed the plot_wireframe_speedup branch from 703ef23 to cb0a07c Compare January 5, 2025 22:45
@scottshambaugh scottshambaugh force-pushed the plot_wireframe_speedup branch from ff7a2db to 28d03bf Compare January 7, 2025 18:24
@scottshambaugh scottshambaugh force-pushed the plot_wireframe_speedup branch from 74543da to 00062a6 Compare January 8, 2025 16:07
@scottshambaugh
Copy link
Contributor Author

CI failures are unrelated, merging

@scottshambaugh scottshambaugh merged commit d08fd8f into matplotlib:main Jan 8, 2025
37 of 40 checks passed
@anntzer anntzer mentioned this pull request Mar 29, 2025
5 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants