Vectorize patch extraction in Axes3D.plot_surface by apaszke · Pull Request #16675 · matplotlib/matplotlib · GitHub

PR Summary

First of the improvements suggested in #16659.

When rstride and cstride divide the numbers of rows and columns minus 1,
all polygons have the same number of vertices which can be recovered
using some stride tricks on NumPy arrays. On a toy benchmark of a
800x800 grid this allows to reduce the time it takes to produce this
list of polygons from ~10s to ~60ms.

The formatting gets a bit awkward in some places due to the 79 column limit 😕 Also, Flake 8 seems to report many errors in the files I've modified, but I've only tried to make sure the lines I've added don't show up to minimize the diff.

PR Checklist

Has Pytest style unit tests
Code is Flake 8 compliant
New features are documented, with examples if plot related
Documentation is sphinx and numpydoc compliant
Added an entry to doc/users/next_whats_new/ if major new feature (follow instructions in README.rst there)
Documented in doc/api/api_changes.rst if API changed in a backward-incompatible way

I just did a quick benchmark, and this PR reduces the runtime of the following script from 16s to 10s on my old MBP (speedup might be higher on systems with more cores):

import numpy as np
import matplotlib.pyplot as plt

fig = plt.figure()
ax = fig.gca(projection='3d')

x = y = np.linspace(-1, 1, 400)
x, y = np.meshgrid(x, y)
z = x ** 2 + y ** 2

ax.plot_surface(x, y, z, rstride=1, cstride=1)
fig.savefig('dump.png')
plt.close()

Some minor comments which need to be addressed but are not actually really blocking, this basically looks great.

Also note that I haven't added any tests because I've checked that existing plot_surface tests trigger both code paths.

@anntzer I added the padding in #16699.

let's keep that discussion in the corresponding PR then.

@anntzer Please see if the _array_patch_perimeters comment is satisfactory now. I think all your concerns should be addressed now.

looks great, just needs a second positive review.
you may also consider squashing your commits.

Is there any reason to prefer squashing on my side to merge + squash on GitHub? I've usually done the latter for most projects. What's the policy here?

I think it's mostly habit, though perhaps @tacaswell will have an opinion there :)

Personally, I dislike GitHub squash merge because it breaks contributors' git prune + git branch -d.

Small correction, parameter names in docstrings should be *var*, not **var**.

@QuLogic I've applied the suggested changes and added an example for _unfold.

Thanks, the example is quite helpful.

Do you want to squash this?

When rstride and cstride divide the numbers of rows and columns minus 1, all polygons have the same number of vertices which can be recovered using some stride tricks on NumPy arrays. On a toy benchmark of a 800x800 grid this allows to reduce the time it takes to produce this list of polygons from ~10s to ~60ms.

@QuLogic done!

Failure looks like disk / filesystem related problems.

Thanks @apaszke ! Congratulations on your first Matplotlib PR 🎉 hopefully we will hear from you again.

Failure was #16735.

@tacaswell Thanks! I already have a few followup PRs posted, so I'd appreciate reviews on those!

Also thanks to reviewers for helping to get this in!

anntzer added the topic: mplot3d label Mar 5, 2020

apaszke mentioned this pull request Mar 5, 2020

Store vertices as (masked) arrays in Poly3DCollection #16688

Closed

6 tasks done