NumPy итерация в течение двух 2d массивов

Question

Скажем, у меня есть две матрицы:

X, Y = np.meshgrid(np.arange(0, 2, 0.1), np.arange(3, 5, 0.1)) 

и функцию, что-то вроде:

def func(x) : 
    return x[0]**2 + x[1]**2 

Как я могу заполнить матрицу Z (от размера np.shape (Х)), где каждый элемент формируется путем вызова func на двух соответствующих значений X и Y, то есть:

Z[i, j] = func([X[i, j], Y[i, j]]) 

Есть ли способ без использования двойного вложенного цикла?

Answer 1

Это также работает в векторной форме оценки функции:

import numpy as np 
X, Y = np.meshgrid(np.arange(0, 2, 0.1), np.arange(3, 5, 0.1)) 
def func(x) : 
    return x[0]**2 + x[1]**2 

Z = func([X,Y]) 

Answer 2

Для заданных Numpy массивов X и Y, вы можете просто сделать -

Zout = X**2 + Y**2 

Если вы на самом деле построения X и Y подобное, есть прямой путь, чтобы получить Z с broadcasting и таким образом избежать np.meshgrid, подобный так -

Zout = np.arange(0, 2, 0.1)**2 + np.arange(3, 5, 0.1)[:,None]**2