python - 如何优雅地在 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,) 

我只能通过一个简单的嵌套 for 循环实现来做到这一点：

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 的方式优雅地完成它呢？

标签： pythonnumpy

Answer 1

做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