python-如何应用将向量返回到每个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惯用且更快)吗？

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])
     ...: 

签名文档有点含糊.我以前使用过它,所以做了一个很好的第一个猜测：

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')

不会保留dtype,因为您的blackbox_fn不会保留它.默认情况下,vectorize使用第一个元素进行测试计算,并使用其dtype确定结果的dtype.可以使用otypes参数指定return dtype.

它可以处理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])

具有签名的矢量化使用Python级别的迭代.如果没有签名,它将使用np.frompyfunc,性能会更好一些.但是,只要必须为输入的元素调用blackbox_fn,我们就无法将速度提高很多(最多2倍).

np.frompyfunc返回一个对象dtype数组：

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)

在这种2d情况下,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)

但我们可以将其拆散并堆叠.它需要重塑：

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)