plot_surface final patch (for now)

Hi again,

I'll try to keep my messages in plain text from now on. I kept working
on some more speed improvements for plot_surface functions and have
included the diff that has those extra changes.

Here are the main extra changes:
1. vectorized the code in _shade_colors when inputting a grid of
colors by moving to an already vectorized to_rgba_array() function,
using the vectorized Normalize class, and using numpy to multiply the
two. Also, making the result be a numpy array stops set_face_color()
function from redundantly calling to_rgba on each individual item
again. This improvement speeds up the mplot3d example surface3d_demo3
by about 3x on my machine.
2. I preallocated the space for the two vectors that will be crossed
for each area of points to make the normals and do a vectorized
np.cross() afterwards instead. Shaves off a bit more time.
3. Made the functions that select points not include multiples of the
corners of the rectangle of points (if that makes any sense). Not a
huge savings, but still worthwhile I feel.

I think that I'm done with working on these functions now. I hope that
these changes are amenable. Let me know if there are concerns.

Justin Peel

plot_surface_final.diff (4.24 KB)


---------- Forwarded message ----------
From: J P <jpscipy@...149...>
Date: Tue, Nov 16, 2010 at 4:20 PM
Subject: plot_surface patch

Hi all, here's my first patch for matplotlib. Someone noticed at Stack
Overflow that the plot_surface function in mplot3d wasn't especially
fast for a lot of points (and small rstrides/cstrides) and using
shading and a single color. I found some parts of the code that
weren't vectorized. These are my changes so far.

Summary of changes:
1. Changed from double looping over aranges to using xrange
2. Made the normalization of the normals and their dot product with
the vector [-1,-1,0.5] to find the shading a vectorized operation.
3. Changed a list comprehension which calculated the colors using an
iterative approach to using the already built-in vectorization of the
Normalization class and using the np.outer function. The result is a
numpy array rather than a list which actually speeds up things down
the line.
4. removed the corners array from plot_surface which wasn't ever used
or returned. It didn't really slow things down, but I'm thinking that
it is cruft.

For change number two, I made a separate function that generates the
shades, but feel free to move that around if you prefer.. or maybe it
should be a function that begins with a _ because it shouldn't be used
externally. These changes give varying levels of speed improvement
depending on the number of points and the rstrides/cstrides arguments.
With larger numbers of points and small rstrides/cstrides, these
changes can more than halve the running time. I have found no
difference in output after my changes.

I know there is more work to be done within the plot_surface function
and I'll submit more changes soon.