vectorize - python generate grid of points
numpy-evaluate function on a grid of points (4)
What is a good way to produce a numpy array containing the values of a function evaluated on an n-dimensional grid of points?
For example, suppose I want to evaluate the function defined by
def func(x, y): return <some function of x and y>
Suppose I want to evaluate it on a two dimensional array of points with the x values going from 0 to 4 in ten steps, and the y values going from -1 to 1 in twenty steps. What's a good way to do this in numpy?
P.S. This has been asked in various forms on StackOverflow many times, but I couldn't find a concisely stated question and answer. I posted this to provide a concise simple solution (below).
I use this function to get X, Y, Z values ready for plotting:
def npmap2d(fun, x_spec, y_spec, doPrint=False): xs = np.linspace(*x_spec) ys = np.linspace(*y_spec) Z = np.empty(len(xs) * len(ys)) i = 0 for y in ys: for x in xs: Z[i] = fun(x, y) if doPrint: print([i, x, y, Z[i]]) i += 1 X, Y = np.meshgrid(xs, ys) Z.shape = X.shape return X, Y, Z
def f(x, y): # ...some function that can't handle numpy arrays X, Y, Z = npmap2d(f, (0, 0.5, 21), (0.6, 0.4, 41)) fig = plt.figure() ax = fig.add_subplot(111, projection='3d') ax.plot_wireframe(X, Y, Z)
The same result can be achieved using map:
xs = np.linspace(0, 4, 10) ys = np.linspace(-1, 1, 20) X, Y = np.meshgrid(xs, ys) Z = np.fromiter(map(f, X.ravel(), Y.ravel()), X.dtype).reshape(X.shape)
In the case your function actually takes a tuple of
d elements, i.e.
f((x1,x2,x3,...xd)) (for example the scipy.stats.multivariate_normal function), and you want to evaluate
f on N^d combinations/grid of N variables, you could also do the following (2D case):
x=np.arange(-1,1,0.2) # each variable is instantiated N=10 times y=np.arange(-1,1,0.2) Z=f(np.dstack(np.meshgrid(x,y))) # result is an NxN (10x10) matrix, whose entries are f((xi,yj))
np.dstack(np.meshgrid(x,y)) creates an 10x10 "matrix" (technically a 10x10x2 numpy array) whose entries are the 2-dimensional tuples to be evaluated by
shorter, faster and clearer answer, avoiding meshgrid:
import numpy as np def func(x, y): return np.sin(y * x) xaxis = np.linspace(0, 4, 10) yaxis = np.linspace(-1, 1, 20) result = func(x[:,None], y[None,:])
This will be faster in memory if you get something like x^2+y as function, since than x^2 is done on a 1D array (instead of a 2D one), and the increase in dimension only happens when you do the "+". For meshgrid, x^2 will be done on a 2D array, in which essentially every row is the same, causing massive time increases.
Edit: the "x[:,None]", makes x to a 2D array, but with an empty second dimension. This "None" is the same as using "x[:,numpy.newaxis]". The same thing is done with Y, but with making an empty first dimension.
Edit: in 3 dimensions:
def func2(x, y, z): return np.sin(y * x)+z xaxis = np.linspace(0, 4, 10) yaxis = np.linspace(-1, 1, 20) zaxis = np.linspace(0, 1, 20) result2 = func2(xaxis[:,None,None], yaxis[None,:,None],zaxis[None,None,:])
This way you can easily extend to n dimensions if you wish, using as many
: as you have dimensions. Each
: makes a dimension, and each
None makes an "empty" dimension. The next example shows a bit more how these empty dimensions work. As you can see, the shape changes if you use
None, showing that it is a 3D object in the next example, but the empty dimensions only get filled up whenever you multiply with an object that actually has something in those dimensions (sounds complicated, but the next example shows what i mean)
In : import numpy In : a = numpy.linspace(-1,1,20) In : a.shape Out: (20,) In : a[None,:,None].shape Out: (1, 20, 1) In : b = a[None,:,None] # this is a 3D array, but with the first and third dimension being "empty" In : c = a[:,None,None] # same, but last two dimensions are "empty" here In : d=b*c In : d.shape # only the last dimension is "empty" here Out: (20, 20, 1)
edit: without needing to type the None yourself
def ndm(*args): return [x[(None,)*i+(slice(None),)+(None,)*(len(args)-i-1)] for i, x in enumerate(args)] x2,y2,z2 = ndm(xaxis,yaxis,zaxis) result3 = func2(x2,y2,z2)
This way, you make the
None-slicing to create the extra empty dimensions, by making the first argument you give to ndm as the first full dimension, the second as second full dimension etc- it does the same as the 'hardcoded' None-typed syntax used before.
Short explanation: doing
x2, y2, z2 = ndm(xaxis, yaxis, zaxis) is the same as doing
x2 = xaxis[:,None,None] y2 = yaxis[None,:,None] z2 = zaxis[None,None,:]
but the ndm method should also work for more dimensions, without needing to hardcode the
None-slices in multiple lines like just shown. This will also work in numpy versions before 1.8, while numpy.meshgrid only works for higher than 2 dimensions if you have numpy 1.8 or higher.
import numpy as np def func(x, y): return np.sin(y * x) xaxis = np.linspace(0, 4, 10) yaxis = np.linspace(-1, 1, 20) x, y = np.meshgrid(xaxis, yaxis) result = func(x, y)