Generating data from meshgrid data (Numpy)

Question

I'd like to ask how to generate corresponding values from a meshgrid. I have a function "foo" that takes one 1D array with the length of 2, and returns some real number.

import numpy as np

def foo(X):  
    #this function takes a vector, e.g., np.array([2,3]), and returns a real number.
    return sum(X)**np.sin( sum(X) );

x = np.arange(-2, 1, 1)          # points in the x axis
y = np.arange( 3, 8, 1)          # points in the y axis
X, Y = np.meshgrid(x, y)         # X, Y : grid

I generate X and Y grids using meshgrid.

Then, how can I generate corresponding Z values using "foo" function, in order to plot them in 3D, e.g., plotting using plot_surface function with X,Y,Z values?

Here the question is how to generate Z values, which has the same shape to X and Y, using "foo" function. Since my "foo" function only takes an 1D array, I do not know how I can uses this function with X and Y to generate corresponding Z values.

Answer 1

Stack your two numpy arrays in "depth" using np.dstack, and then modify your foo function, so that it operates on only the last axis of your stacked array. This is easily done using np.sum with parameter axis=-1, instead of using the builtin sum:

import numpy as np

def foo(xy):
    return np.sum(xy, axis=-1) ** np.sin(np.sum(xy, axis=-1))

x = np.arange(-2, 1, 1)          # points in the x axis
y = np.arange( 3, 8, 1)          # points in the y axis
X, Y = np.meshgrid(x, y)         # X, Y : grid
XY = np.dstack((X, Y))

And now, you should get:

>>> XY.shape
(5, 3, 2)
>>> foo(XY)
array([[ 1.        ,  1.87813065,  1.1677002 ],
       [ 1.87813065,  1.1677002 ,  0.35023496],
       [ 1.1677002 ,  0.35023496,  0.2136686 ],
       [ 0.35023496,  0.2136686 ,  0.60613935],
       [ 0.2136686 ,  0.60613935,  3.59102217]])

---

If you want to achieve the same effect, but without modifying foo, then you can use np.apply_along_axis, which should do exactly what you need:

>>> np.apply_along_axis(foo, -1, XY)
array([[ 1.        ,  1.87813065,  1.1677002 ],
       [ 1.87813065,  1.1677002 ,  0.35023496],
       [ 1.1677002 ,  0.35023496,  0.2136686 ],
       [ 0.35023496,  0.2136686 ,  0.60613935],
       [ 0.2136686 ,  0.60613935,  3.59102217]])