[Numpy-discussion] some vectorization help

[Eric Carlson]

Hello All,

I have a bit of code that nicely accomplishes what I need for a course I 
am teaching. I'd like to extend this for larger 3D grids, and I suspect 
that the looping will be a brutal performance hit. Even if my suspicions 
are not confirmed, I still would like to know if it's possible to 
vectorize the following code:


from scipy import shape, prod, array, zeros,ravel, reshape,sin, mgrid
from scipy.misc import derivative
def gradient2D_vect(func,x,y):
     the_shape = shape(x)
     x1=ravel(x)
     y1=ravel(y)
     N = prod(the_shape)
     the_result = zeros([N,2])
     for k in range(N):
         func_x=lambda x: func(x,y1[k])
         func_y=lambda y: func(x1[k],y)
         the_result[k,:]= array([derivative(func_x,x1[k],dx=.01, order=5),
                                 derivative(func_y,y1[k],dx=.01, order=5)])
     if prod(shape(the_shape))==1:
         return the_result
     else:
         return reshape(the_result,[the_shape[0],the_shape[1],2])

fxy = lambda x,y: sin(x*y) #just a little test
x,y=mgrid[0:5,0:4]
the_gradient = gradient2D_vect(fxy, x,y)



Cheers,
Eric Carlson