5

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.

Ann Descomp
  • 95
  • 1
  • 7

1 Answers1

3

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]])
Praveen
  • 6,872
  • 3
  • 43
  • 62
  • Thanks for the wonderful solution. This is exactly what I wanted to do. However, I have some problems to apply this approach to my code bock. I actually posted another question regarding this matter [link](http://stackoverflow.com/questions/39715227/making-a-function-that-can-take-arguments-in-various-shapes). My actual foo function is actually something like `return X[0] + X[1];`, so it cannot perform an iterative operation on a nested array, and can only take the shape of (2,). Thus, if I put XY into my actual foo function, it spits our error. In this case, how can generate Z as you did? – Ann Descomp Sep 27 '16 at 02:51
  • Oops, sorry, I forgot it! – Ann Descomp Sep 27 '16 at 03:21
  • By the way, can you let me know how I should change your ode if I just want to use my current foo function, which can only take (2,) shape? – Ann Descomp Sep 27 '16 at 03:28
  • @AnnDescomp I've posted an answer on your other question. See if that works for you... – Praveen Sep 27 '16 at 03:30
  • I am now looking at "..." indexing. In addition, I'd want to know if there is a way to generate Z with my current foo function, which can explicitly take only one (2,) shape Numpy array as an argument. – Ann Descomp Sep 27 '16 at 03:45
  • @AnnDescomp You _don't_ want to modify your foo function, but you still want to achieve the same result? – Praveen Sep 27 '16 at 03:51
  • @AnnDescomp See the edit I've made to my answer. Hopefully that should solve your issue. Again, don't forget to upvote! :-) – Praveen Sep 27 '16 at 04:09