Ann Descomp - 1 year ago 127

Python 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 Source

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