Как элегантно написать пользовательскую поэлементную функцию в numpy?

Question

Я хочу использовать функцию np.meshgrid для выполнения некоторых действий на 2D-сетке:

import numpy as np

def gauss2d(x, y):
    return np.exp(-(np.power(x-np.arange(10), 2) + np.power(y-np.arange(10), 2)) / 2).sum()

x, y = np.meshgrid(np.arange(5), np.arange(6))
z = gauss2d(x, y)

, но получил ошибку:

ValueError: operands could not be broadcast together with shapes (6,5) (10,) 

Я могу сделать это только правильно простым вложенным для реализации l oop:

z = np.zeros_like(x)
for i in range(x.shape[1]):
    for j in range(x.shape[0]):
        z[i, j] = gauss2d(x[i, j], y[i, j])

Итак, как это элегантно сделать numpy образом?

Answer 1

In [4]: x, y = np.meshgrid(np.arange(5), np.arange(6))                                           
In [5]: x.shape                                                                                  
Out[5]: (6, 5)
In [6]: y.shape                                                                                  
Out[6]: (6, 5)

x не может работать напрямую с массивом размера 10:

In [7]: x-np.arange(10)                                                                          
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-7-cd10c5c1eea4> in <module>
----> 1 x-np.arange(10)

ValueError: operands could not be broadcast together with shapes (6,5) (10,) 

, но, добавив измерение к x, сделав его (6,5,1), он будет 'широковещательный' с (10,):

In [8]: (x[:,:,None]-np.arange(10)).shape                                                        
Out[8]: (6, 5, 10)

Эффективно: (6,5,1) (10,) => (6,5,1) (1,1,10) => (6,5,10)

Тогда мы можем сделать сумму по последнему измерению

In [9]: (x[:,:,None]-np.arange(10)).sum(axis=2).shape                                            
Out[9]: (6, 5)

В качестве альтернативы (x - np.arange(10)[:,None,None]).sum(axis=0), поставив (10,) на первое место.

Answer 2

Сначала наберите reshape, чтобы numpy знал, что вы работаете над двумя числами каждый раз.

In [25]: x = np.reshape(x, [-1, 1])

In [26]: y = np.reshape(y, [-1, 1])

In [28]: def gauss2d(x, y):
    ...:     return np.exp(-(np.power(x-np.arange(10), 2) + np.power(y-np.arange
    ...: (10), 2)) / 2).sum()
    ...:
    ...:

In [29]: gauss2d(x, y)
Out[29]: 26.295857532549885