NumPy:将向量值函数应用于网格

Question

我正在尝试在NumPy中做以下事情:

import numpy as np
def f(x):
return x[0] + x[1]
X1 = np.array([0, 1, 2])
X2 = np.array([0, 1, 2])
X = np.meshgrid(X1, X2)
result = np.vectorize(f)(X)

，预期结果是array([[0, 1, 2], [1, 2, 3], [2, 3, 4]])，但它返回以下错误:

2 
3 def f(x):
----> 4   return x[0] + x[1]
5 
6 X1 = np.array([0, 1, 2])
IndexError: invalid index to scalar variable

这是因为它试图将f应用于网格的所有18个标量元素，而我希望它应用于9对2标量。正确的做法是什么?

注意:我知道这段代码将工作，如果我不向量化f，但这是重要的，因为f可以是任何函数，例如，它可以包含一个if语句抛出值错误没有向量化。

Answer 1

如果您坚持使用numpy.vectorize，则需要在定义向量化函数时定义signature。

import numpy as np
def f(x):
return x[0] + x[1]
# Or
# return np.add.reduce(x, axis=0)

X1 = np.array([0, 1, 2])
X2 = np.array([0, 1, 2])
X = np.meshgrid(X1, X2)
# np.asarray(X).shape -> (2, 3, 3)
# shape of the desired result is (3, 3)
f_vec = np.vectorize(f, signature='(n,m,m)->(m,m)')
result = f_vec(X)
print(result)

输出:

[[0 1 2]
[1 2 3]
[2 3 4]]

Answer 2

对于您在评论中提到的函数:

f = lambda x: x[0] + x[1] if x[0] > 0 else 0

您可以使用np.where:

def f(x):
return np.where(x > 0, x[0] + x[1], 0)
# np.where(some_condition, value_if_true, value_if_false)

Numpy在设计时考虑了向量化——除非你有一些疯狂的边缘情况，否则几乎总有一种方法可以利用Numpy的广播和向量化。我强烈建议在轻易放弃和求助于for循环之前寻找矢量化的解决方案。

Answer 3

如果你太懒，或者太无知，去做都是"正确的";'矢量化'，您可以使用np.vectorize。但你需要花时间真正阅读它的文档。这不是魔法。它可能是有用的，特别是当你需要利用广播，而函数，由于某种原因，只接受标量。

重写您的函数以处理标量输入(尽管在这种情况下它也可以很好地处理数组):

In [91]: def foo(x,y): return x+y
...: f = np.vectorize(foo)

标量输入:

In [92]: f(1,2)
Out[92]: array(3)

2数组((2,1)和(3)),返回一个(2、3):

In [93]: f(np.array([1,2])[:,None], np.arange(1,4))
Out[93]: 
array([[2, 3, 4],
[3, 4, 5]])

与meshgrid:

相同

In [94]: I,J = np.meshgrid(np.array([1,2]), np.arange(1,4),indexing='ij')
In [95]: I
Out[95]: 
array([[1, 1, 1],
[2, 2, 2]])
In [96]: J
Out[96]: 
array([[1, 2, 3],
[1, 2, 3]])
In [97]: f(I,J)
Out[97]: 
array([[2, 3, 4],
[3, 4, 5]])

或定义在[93]中的网格数组:

In [98]: I,J = np.meshgrid(np.array([1,2]), np.arange(1,4),indexing='ij', sparse=True)
In [99]: I,J
Out[99]: 
(array([[1],
[2]]),
array([[1, 2, 3]]))

但是在真正的向量化意义上，你可以把这两个数组相加:

In [100]: I+J
Out[100]: 
array([[2, 3, 4],
[3, 4, 5]])

np.vectorize文档的第一段(我的重点):

定义一个向量化函数，接受嵌套的对象序列或numpy数组作为输入，并返回单个numpy数组或numpy元组数组。向量化函数对连续元组计算pyfunc的输入数组，除了它使用<<strong>广播规则/strong>numpy。

编辑

从一个需要两个元素元组的函数开始，我们可以添加一个将其分成两个的cover，并对其应用vectorize:

In [103]: def foo1(x): return x[0]+x[1]
...: def foo2(x,y): return foo1((x,y))
...: f = np.vectorize(foo2)
In [104]: f(1,2)
Out[104]: array(3)

X是一个2d元素元组:

In [105]: X = np.meshgrid(np.array([1,2]), np.arange(1,4),indexing='ij')
In [106]: X
Out[106]: 
[array([[1, 1, 1],
[2, 2, 2]]),
array([[1, 2, 3],
[1, 2, 3]])]

可以传递给f:

In [107]: f(X[0],X[1])
Out[107]: 
array([[2, 3, 4],
[3, 4, 5]])

但是没有必要放慢迭代的速度。只需将元组传递给foo1:

In [108]: foo1(X)
Out[108]: 
array([[2, 3, 4],
[3, 4, 5]])

在f = lambda x: x[0] + x[1] if x[0] > 0 else 0中，你会得到' ambiguous ' valueerror，因为if只适用于标量。但是有很多更快的方法来代替这样的if步骤。