如何应用将向量返回到每个 numpy 数组元素的函数（并获取更高维度的数组）

我们直接用代码写吧

注意：我将映射器（原始示例使用 x -> (x, 2 * x, 3 * x) 仅作为示例）编辑为通用黑盒函数，这导致了麻烦。

import numpy as np

def blackbox_fn(x): #I can't be changed!
    assert np.array(x).shape == (), "I'm a fussy little function!"
    return np.array([x, 2*x, 3*x])

# let's have 2d array
arr2d = np.array(list(range(4)), dtype=np.uint8).reshape(2, 2)

# each element should be mapped to vector
def mapper(x, blackbox_fn):
    # there is some 3rdparty non-trivial function, returning np.array
    # in examples returns np.array((x, 2 * x, 3 * x))
    # but still this 3rdparty function operates only on scalar values
    return vectorized_blackbox_fn(x) 

所以对于输入二维数组

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

我想要 3d 数组

array([[[0, 0, 0],
        [1, 2, 3]],

       [[2, 4, 6],
        [3, 6, 9]]], dtype=uint8)

我可以使用 for 循环编写简单的算法

# result should be 3d array, last dimension is same as mapper result size
arr3d = np.empty(arr2d.shape + (3,), dtype=np.uint8)
for y in range(arr2d.shape[1]):
    for x in xrange(arr2d.shape[0]):
        arr3d[x, y] = mapper(arr2d[x, y])

但对于大型阵列来说似乎相当慢。 我知道有np.vectorize，但是使用

np.vectorize(mapper)(arr2d)

不工作，因为

ValueError: setting an array element with a sequence.

（似乎矢量化不能改变维度） 有没有更好的（numpy 惯用且更快）的解决方案？

np.vectorize使用新的签名选项可以解决这个问题。它并没有提高速度，但使维度记账变得更容易。

In [159]: def blackbox_fn(x): #I can't be changed!
     ...:     assert np.array(x).shape == (), "I'm a fussy little function!"
     ...:     return np.array([x, 2*x, 3*x])
     ...: 

的文档signature有点神秘。我以前用过它，所以做了一个很好的初步猜测：

In [161]: f = np.vectorize(blackbox_fn, signature='()->(n)')
In [162]: f(np.ones((2,2)))
Out[162]: 
array([[[ 1.,  2.,  3.],
        [ 1.,  2.,  3.]],

       [[ 1.,  2.,  3.],
        [ 1.,  2.,  3.]]])

用你的数组：

In [163]: arr2d = np.array(list(range(4)), dtype=np.uint8).reshape(2, 2)
In [164]: f(arr2d)
Out[164]: 
array([[[0, 0, 0],
        [1, 2, 3]],

       [[2, 4, 6],
        [3, 6, 9]]])
In [165]: _.dtype
Out[165]: dtype('int32')

The dtype没有被保留，因为你的blackbox_fn不保留它。作为默认值vectorize对第一个元素进行测试计算，并使用其dtype确定结果的 dtype。可以使用以下命令指定返回数据类型otypes范围。

它可以处理 2d 以外的数组：

In [166]: f(np.arange(3))
Out[166]: 
array([[0, 0, 0],
       [1, 2, 3],
       [2, 4, 6]])
In [167]: f(3)
Out[167]: array([3, 6, 9])

With a signature vectorize正在使用Python级别的迭代。没有签名它使用np.frompyfunc，具有更好的性能。但只要blackbox_fn必须调用输入的元素，我们不能提高太多速度（最多 2 倍）。

np.frompyfunc返回一个对象数据类型数组：

In [168]: fpy = np.frompyfunc(blackbox_fn, 1,1)
In [169]: fpy(1)
Out[169]: array([1, 2, 3])
In [170]: fpy(np.arange(3))
Out[170]: array([array([0, 0, 0]), array([1, 2, 3]), array([2, 4, 6])], dtype=object)
In [171]: np.stack(_)
Out[171]: 
array([[0, 0, 0],
       [1, 2, 3],
       [2, 4, 6]])
In [172]: fpy(arr2d)
Out[172]: 
array([[array([0, 0, 0]), array([1, 2, 3])],
       [array([2, 4, 6]), array([3, 6, 9])]], dtype=object)

stack无法删除这种二维情况下的数组嵌套：

In [173]: np.stack(_)
Out[173]: 
array([[array([0, 0, 0]), array([1, 2, 3])],
       [array([2, 4, 6]), array([3, 6, 9])]], dtype=object)

但我们可以把它拆开，然后堆叠起来。它需要一个reshape:

In [174]: np.stack(__.ravel())
Out[174]: 
array([[0, 0, 0],
       [1, 2, 3],
       [2, 4, 6],
       [3, 6, 9]])

速度测试：

In [175]: timeit f(np.arange(1000))
14.7 ms ± 322 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [176]: timeit fpy(np.arange(1000))
4.57 ms ± 161 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [177]: timeit np.stack(fpy(np.arange(1000).ravel()))
6.71 ms ± 207 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [178]: timeit np.array([blackbox_fn(i) for i in np.arange(1000)])
6.44 ms ± 235 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

让你的函数返回一个列表而不是任何数组可能会使重新组装结果变得更容易，甚至可能更快

def foo(x):
    return [x, 2*x, 3*x]

或玩弄frompyfunc参数;

def foo(x):
    return x, 2*x, 3*x   # return a tuple
In [204]: np.stack(np.frompyfunc(foo, 1,3)(arr2d),2)
Out[204]: 
array([[[0, 0, 0],
        [1, 2, 3]],

       [[2, 4, 6],
        [3, 6, 9]]], dtype=object)

速度提高了 10 倍 - 我很惊讶：

In [212]: foo1 = np.frompyfunc(foo, 1,3)
In [213]: timeit np.stack(foo1(np.arange(1000)),1)
428 µs ± 17.9 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)