[Numpy-discussion] 2-D function and meshgrid

[Paulo J. S. Silva]

Hello,

I have a function that receives a array of shape (2,) and returns a
number (a function from R^2 -> R). It basically looks like this:

    def weirdDistance2(x):
       return dot(dot(weirdMatrix, x), x)

(weirdMatrix is a "global" (2,2) array)

I want to see its level sets in the box [0, 1] x [0, 1], hence I have to
create a meshgrid and then compute it at each point of the mesh:

    x = linspace(0, 1, 200)
    y = x.copy()
    X, Y = meshgrid(x, y)

My problem is how to actually compute the function at each point of the
mesh. I have come out with two solutions. One very short and clear, but
slow, and another longer and more convoluted (it has a loop, I hate
loops in numpy code), but faster. Does anyone know a "no explicit-loops"
and fast solution?

Solution1:

    def myDistance(a, b):
        return weirdDistance(np.array((a, b)))
    vecDistance = np.vectorize(myDistance)
    return vecDistance(X, Y)

Solution 2:

    nPoints = X.size
    result = np.zeros(nPoints)
    points = np.array( [X.ravel(), Y.ravel()] ).T
    for i in xrange(nPoints):
        result[i] = weirdDistance(points[i])
    result = result.reshape(X.shape)

Of course, the first one is slow because the myDistance function creates
an array at each call. The second one, even with a loop, avoids the
array creations.

Best,

Paulo