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.
Answers:
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 ravel
ing or reshape
ing 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)
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.
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 ravel
ing or reshape
ing 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)