Ann Descomp Ann Descomp - 2 months ago 14
Python Question

Generating data from meshgrid data (Numpy)

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

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]])