Numpy - Efficiently compute a function (1-d array) over a grid

Question

So I have a function f(x) where the argument x is a row array, with dimension k. The function can be given an array with a number of rows >1 and is optimized to operate row-wise on arrays, clearly much faster than iterating over array rows and calling the function. Now I would like to apply the function on a grid covering a k-dimensional space. Let us say the k-dimension = 3. Then

N_DIV = 2
x0 = np.linspace(0,1,N_DIV)
x1 = np.linspace(0,1,N_DIV)
x2 = np.linspace(0,1,N_DIV)

and I would like to compute the function for all combinations such as

x0  x1  x2
0   0   0
0   0   0.5
0   0.5 0
0   0.5 0.5

etc.

I thought about using np.meshgrid so

xx, yy, zz = np.meshgrid(x0,x1,x2)

but what next? The brutal approach

prev_array=no.array([0,0,0])
for i in range(N_DIV):
    for j in range(N_DIV):
        for k in range(N_DIV):
            prev_array = np.vstack((prev_array, 
                                    np.array([xx[i,j,k],yy[i,j,k],zz[i,j,k]])))

cannot be right, any suggestions please? I would like to efficiently compute the functionf over a grid covering the k-dimensional space, thanks *** EDIT

The post Evaluate function on a grid of points has been suggested as a solution, but I fail to see how it could answer my question. They have a f(x,y) of two scalar variables, and I see how the idea result = func(xaxis[:,None], yaxis[None,:]). But my function takes a row vector as input, so {x,y}, and hence the idea above seems not directly applicable, to me at least, thanks again

BACKGROUND - What I am trying to achieve

Say I have a function of 3 variables

def func(x,y,z):
    return x**3 - 3*y + z**2

and I want to plot it as a function of say(x,y), for a fixed value of z.

I could do

N_DIV = 30
x =np.linspace(0,5,N_DIV)
y =np.linspace(0,5,N_DIV)
z =np.linspace(0,5,N_DIV)
xx , yy, zz = np.meshgrid(x,y,z)
W = func(xx,yy,zz)
import matplotlib.pyplot as plt
# Plot the surface
fig, ax = plt.subplots(subplot_kw={"projection": "3d"})
surf = ax.plot_surface(xx[:,:,1],yy[:,:,1],W[:,:,1], cmap="Spectral",
                       linewidth=0)

plt.show()

worls fine and the repeated function evaluations are fast. But, if my function is defined as

def func_vect(x):
    return x[0]**3 - 3*x[1] + x[2]**2

how to achieve the same result? That is, creating an array W of output results, ready to plot using as before?

The brute force approach would be to create it by looping, but I am also confused by the following

def func2(x,y):
    return x**3 - 3*y 
xx2 , yy2 = np.meshgrid(x,y)
W2 = func2(xx2,yy2)
W2_loop = np.zeros((N_DIV, N_DIV))
for i in range(N_DIV):
    for j in range(N_DIV):
            W2_loop[i,j] = func2(x[j],y[i])
np.isclose(W2,W2_loop)

returns all True, but I cannot figure out how to make it work in three dimensions, as (SECOND FUNDAMENTAL ISSUE)

W_loop = np.zeros((N_DIV, N_DIV,N_DIV))
for i in range(N_DIV):
    for j in range(N_DIV):
        for k in range(N_DIV):
            W_loop[i,j,k] = func(x[k],y[j],z[i])

is different from the W created above.

Thanks

Does this answer your question? [numpy - evaluate function on a grid of points](https://stackoverflow.com/questions/22774726/numpy-evaluate-function-on-a-grid-of-points) — Matt Hall, Jul 19 '23 at 10:11
Thanks for this I did read that question but I fail to see how it could answer my question. They have a ```f(x,y)``` of two scalar variables, and I see how the idea ```result = func(xaxis[:,None], yaxis[None,:])```. But my function take a row vector as input, so {x,y}, and hence the idea above seems not directly applicable, to me at least, thanks again — user37292, Jul 19 '23 at 10:31
So you want a 2d, (n,3) array with all combinations of the 3 arrays? — hpaulj, Jul 19 '23 at 15:24
@hpaulj thanks for your answer, which I am going through in detail. Yes ultimately I would like a 2d, (n,3) array with all combinations of the 3 arrays, but it has also to be in a certain order consistent with the xx, yy, zz arrays for plotting purposes. I added an edit to the question explaining what my aim is, I hope it will make it clearer, thanks a lot again — user37292, Jul 19 '23 at 18:31

hpaulj · Accepted Answer · 2023-07-19T18:37:40.897

First that repeated use of vstack is very bad. If you must build arrays iteratively, use list and list.append.

In [134]: N_DIV = 2
     ...: x0 = np.linspace(0,1,N_DIV)
     ...: x1 = np.linspace(0,1,N_DIV)
     ...: x2 = np.linspace(0,1,N_DIV)

Your 3 meshgrid arrays:

In [136]: a,b,c = np.meshgrid(x0,x1,x2, indexing='ij')
In [137]: a,b,c
Out[137]: 
(array([[[0., 0.],
         [0., 0.]],
 
        [[1., 1.],
         [1., 1.]]]),
 array([[[0., 0.],
         [1., 1.]],
 
        [[0., 0.],
         [1., 1.]]]),
 array([[[0., 1.],
         [0., 1.]],
 
        [[0., 1.],
         [0., 1.]]]))

np.stack, like np.array can combine them on a new leading axis.

In [138]: np.stack((a,b,c)).shape
Out[138]: (3, 2, 2, 2)

But you want 3 to be the last axis:

In [139]: np.stack((a,b,c),axis=-1).shape
Out[139]: (2, 2, 2, 3)

So let's reshape that to a (n,3) array:

In [140]: np.stack((a,b,c),axis=-1).reshape(-1,3)
Out[140]: 
array([[0., 0., 0.],
       [0., 0., 1.],
       [0., 1., 0.],
       [0., 1., 1.],
       [1., 0., 0.],
       [1., 0., 1.],
       [1., 1., 0.],
       [1., 1., 1.]])

another tool is

In [142]: from itertools import product

In [144]: list(product(x0,x1,x2))
Out[144]: 
[(0.0, 0.0, 0.0),
 (0.0, 0.0, 1.0),
 (0.0, 1.0, 0.0),
 (0.0, 1.0, 1.0),
 (1.0, 0.0, 0.0),
 (1.0, 0.0, 1.0),
 (1.0, 1.0, 0.0),
 (1.0, 1.0, 1.0)]

A variant on meshgrid is mgrid:

In [153]: idx = 1j*N_DIV
In [156]: np.mgrid[0:1:idx,0:1:idx,0:1:idx].reshape(3,-1).T
Out[156]: 
array([[0., 0., 0.],
       [0., 0., 1.],
       [0., 1., 0.],
       [0., 1., 1.],
       [1., 0., 0.],
       [1., 0., 1.],
       [1., 1., 0.],
       [1., 1., 1.]])

Have you considered taking advantage of broadcasting. meshgrid can produce "sparse" arrays:

In [162]: a,b,c =np.meshgrid(x0,x1,x2,indexing='ij',sparse=True)
In [163]: a,b,c
Out[163]: 
(array([[[0.]],
 
        [[1.]]]),
 array([[[0.],
         [1.]]]),
 array([[[0., 1.]]]))

ix_ can produce the same:

In [164]: a,b,c =np.ix_(x0,x1,x2)

They can be used in exactly the same way as the "regular" meshgrid arrays, for example adding them to make a (2,2,2) array:

In [165]: a*100+b*10+c
Out[165]: 
array([[[  0.,   1.],
        [ 10.,  11.]],

       [[100., 101.],
        [110., 111.]]])

edit

Your latest example, with different sizes in each dimension:

In [222]: N_DIV = 30
     ...: x =np.linspace(0,5,N_DIV//2)
     ...: y =np.linspace(0,5,N_DIV)
     ...: z =np.linspace(0,5,N_DIV*2)
In [223]: a,b,c =np.ix_(x,y,z)
In [224]: V = a**3 - 3*b + c**2
In [225]: V.shape
Out[225]: (15, 30, 60)

thanks for this, I am going through it but I also added an edit to the original question explaining what is I want to achieve, as I need a specific row-order too — user37292, Jul 19 '23 at 18:33
I added the calculation for your `x[0]**3 - 3*x[1] + x[2]**2`, producing a 3d result (via broadcasting). You can tweak that shape how ever you want, with transpose and reshapes. — hpaulj, Jul 19 '23 at 18:41
inasmuch I appreciate your support a lot, and is very instructive, I fail to see how this answers it. Ultimately I want to plot with ```plot_surface(xx[:,:,1],yy[:,:,1],W[:,:,1]```, ```W ```has then to be ```(N_DIV, N_DIV, N_DIV),``` and how do I get it given `func_vect` accepting a row vector. Basically, how to plot ```func_vect``` over ```x``` and ````y``` for a fixed ```z```. But ok I am sure the answer is there, will go through it carefully thanks a lot — user37292, Jul 19 '23 at 18:55
for one, ```func_vect(V).shape``` is (30,30) , for ```N_DIV=30```, while I expected it to be (30,30,30), as a scalar function over 30x30X30 evaluations. But very very instructive thanks — user37292, Jul 19 '23 at 19:02
I tweaked the shapes to more readily identify which array goes with which dimension. I much prefer illustrating code with (2,3,4) array than a (3,3,3). — hpaulj, Jul 19 '23 at 19:37

Numpy - Efficiently compute a function (1-d array) over a grid

1 Answers1

edit