在返回向量的函数上使用 Numpy 向量化

Question

numpy.vectorize 是一个可以把一个函数 f:a->b 转换成 g:a[]->b[].

当 a 和 b 是单个数值时，这个方法很好用，但我想不出为什么它在 b 是一个 ndarray 或列表时就不行，也就是说 f:a->b[] 和 g:a[]->b[][] 也应该可以。

举个例子：

import numpy as np
def f(x):
    return x * np.array([1,1,1,1,1], dtype=np.float32)
g = np.vectorize(f, otypes=[np.ndarray])
a = np.arange(4)
print(g(a))

这样做会得到：

array([[ 0.  0.  0.  0.  0.],
       [ 1.  1.  1.  1.  1.],
       [ 2.  2.  2.  2.  2.],
       [ 3.  3.  3.  3.  3.]], dtype=object)

好的，这样得到了正确的值，但数据类型不对。而且更糟糕的是：

g(a).shape

会得到：

(4,)

所以这个数组几乎没什么用。我知道我可以通过以下方式转换它：

np.array(map(list, a), dtype=np.float32)

这样可以得到我想要的结果：

array([[ 0.,  0.,  0.,  0.,  0.],
       [ 1.,  1.,  1.,  1.,  1.],
       [ 2.,  2.,  2.,  2.,  2.],
       [ 3.,  3.,  3.,  3.,  3.]], dtype=float32)

但这样做既不高效，也不符合 Python 的风格。有没有人能找到更简洁的方法来解决这个问题？

Answer 1

import numpy as np
def f(x):
    return x * np.array([1,1,1,1,1], dtype=np.float32)
g = np.vectorize(f, otypes=[np.ndarray])
a = np.arange(4)
b = g(a)
b = np.array(b.tolist())
print(b)#b.shape = (4,5)
c = np.ones((2,3,4))
d = g(c)
d = np.array(d.tolist())
print(d)#d.shape = (2,3,4,5)

这样做应该能解决问题，而且不管你的输入有多大都能正常工作。“map”只适用于一维的输入。使用“.tolist()”并创建一个新的ndarray可以更彻底、更好地解决这个问题（我觉得）。希望这对你有帮助。

Answer 2

在1.12.0版本中新增了一个参数signature，它正好满足你的需求。

def f(x):
    return x * np.array([1,1,1,1,1], dtype=np.float32)

g = np.vectorize(f, signature='()->(n)')

然后使用g(np.arange(4)).shape会得到(4L, 5L)。

这里指定了f的签名。(n)表示返回值的形状，而()表示参数的形状，这里是标量（单个数值）。参数也可以是数组。如果你想了解更复杂的签名，可以查看通用函数API。

Answer 3

np.vectorize 只是一个方便的工具。它并不会让代码运行得更快。如果使用 np.vectorize 不方便，那就自己写一个符合你需求的函数。

np.vectorize 的目的是把那些不懂 numpy 的函数（比如输入和输出都是浮点数的函数）转换成可以处理 numpy 数组的函数。

你的函数 f 已经是懂 numpy 的了——它在定义中使用了 numpy 数组，并且返回的是一个 numpy 数组。所以 np.vectorize 并不适合你的情况。

因此，解决办法就是自己写一个符合你需求的函数 f。