Создание данных из данных meshgrid (Numpy)

Question

Я хотел бы спросить, как сгенерировать соответствующие значения из meshgrid. У меня есть функция foo, которая принимает один массив 1D с длиной 2 и возвращает некоторое действительное число.

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 

Я генерирую решетки X и Y, используя meshgrid.

Затем, как я могу сгенерировать соответствующие значения Z, используя функцию «foo», чтобы построить их в 3D, например, с использованием функции plot_surface с значениями X, Y, Z?

Здесь возникает вопрос, как сгенерировать значения Z, которые имеют одинаковую форму с X и Y, используя функцию «foo». Поскольку моя функция «foo» принимает только 1D-массив, я не знаю, как я могу использовать эту функцию с X и Y для генерации соответствующих значений Z.

Answer 1

Сложите два массива numpy в «глубину» с помощью np.dstack, а затем измените функцию foo, чтобы она работала только по последней оси вашего массива. Это легко сделать с помощью np.sum с параметром axis=-1, вместо того, чтобы с помощью встроенного 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)) 

И теперь, вы должны получить:

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

---

Если вы хотите, чтобы достичь того же эффекта, но без изменение foo, то вы можете использовать np.apply_along_axis, что должно делать именно то, что вам нужно:

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