How do I apply some function to a python meshgrid?

Question

Say I want to calculate a value for every point on a grid. I would define some function func that takes two values x and y as parameters and returns a third value. In the example below, calculating this value requires a look-up in an external dictionary. I would then generate a grid of points and evaluate func on each of them to get my desired result.

The code below does precisely this, but in a somewhat roundabout way. First I reshape both the X and Y coordinate matrices into one-dimensional arrays, calculate all the values, and then reshape the result back into a matrix. My questions is, can this be done in a more elegant manner?

import collections as c  # some arbitrary lookup table a = c.defaultdict(int) a[1] = 2 a[2] = 3 a[3] = 2 a[4] = 3  def func(x,y):     # some arbitrary function     return a[x] + a[y]  X,Y = np.mgrid[1:3, 1:4] X = X.T Y = Y.T  Z = np.array([func(x,y) for (x,y) in zip(X.ravel(), Y.ravel())]).reshape(X.shape) print Z 

The purpose of this code is to generate a set of values that I can use with pcolor in matplotlib to create a heatmap-type plot.

Answer 1

I'd use numpy.vectorize to "vectorize" your function. Note that despite the name, vectorize is not intended to make your code run faster -- Just simplify it a bit.

Here's some examples:

>>> import numpy as np >>> @np.vectorize ... def foo(a, b): ...    return a + b ...  >>> foo([1,3,5], [2,4,6]) array([ 3,  7, 11]) >>> foo(np.arange(9).reshape(3,3), np.arange(9).reshape(3,3)) array([[ 0,  2,  4],        [ 6,  8, 10],        [12, 14, 16]]) 

With your code, it should be enough to decorate func with np.vectorize and then you can probably just call it as func(X, Y) -- No raveling or reshapeing necessary:

import numpy as np import collections as c  # some arbitrary lookup table a = c.defaultdict(int) a[1] = 2 a[2] = 3 a[3] = 2 a[4] = 3  @np.vectorize def func(x,y):     # some arbitrary function     return a[x] + a[y]  X,Y = np.mgrid[1:3, 1:4] X = X.T Y = Y.T  Z = func(X, Y)