如何用numpy有效地将函数应用于三维数组？

Question

我想把任意函数应用到3d-ndarray作为元素，它使用(3维)数组作为参数，并返回scalar.As一个结果，我们应该得到2d矩阵。

例如，伪码

A = [[[1,2,3],[4,5,6]],
     [[7,8,9],[10,11,12]]]
A.apply_3d_array(sum) ## or apply_3d_array(A,sum) is Okey.
>> [[6,15],[24,33]]

我理解使用ndarray.shape函数的循环是可能的，但正如官方文档所述，直接索引访问效率很低。有没有比使用循环更有效的方法？

def chromaticity(pixel):
    geo_mean = math.pow(sum(pixel),1/3)
    return map(lambda x: math.log(x/geo_mean),pixel ) 

Answer 1

考虑到函数实现，我们可以使用NumPy ufuncs将其矢量化，这将对整个输入数组A进行一次操作，从而避免了不支持数组向量化的math库函数。在此过程中，我们还将引入非常高效的矢量化工具：NumPy broadcasting。所以，我们会有这样的实现-

np.log(A/np.power(np.sum(A,2,keepdims=True),1/3))

示例运行与验证

不使用lamdba构造和引入NumPy函数而不是math库函数的函数实现如下所示-

def chromaticity(pixel): 
    geo_mean = np.power(np.sum(pixel),1/3) 
    return np.log(pixel/geo_mean)

使用迭代实现运行的示例-

In [67]: chromaticity(A[0,0,:])
Out[67]: array([-0.59725316,  0.09589402,  0.50135913])

In [68]: chromaticity(A[0,1,:])
Out[68]: array([ 0.48361096,  0.70675451,  0.88907607])

In [69]: chromaticity(A[1,0,:])
Out[69]: array([ 0.88655887,  1.02009026,  1.1378733 ])

In [70]: chromaticity(A[1,1,:])
Out[70]: array([ 1.13708257,  1.23239275,  1.31940413])    

用拟议的矢量化实现进行示例运行-

In [72]: np.log(A/np.power(np.sum(A,2,keepdims=True),1/3))
Out[72]: 
array([[[-0.59725316,  0.09589402,  0.50135913],
        [ 0.48361096,  0.70675451,  0.88907607]],

       [[ 0.88655887,  1.02009026,  1.1378733 ],
        [ 1.13708257,  1.23239275,  1.31940413]]])

运行时测试

In [131]: A = np.random.randint(0,255,(512,512,3)) # 512x512 colored image

In [132]: def org_app(A):
     ...:     out = np.zeros(A.shape)     
     ...:     for i in range(A.shape[0]):
     ...:         for j in range(A.shape[1]):
     ...:             out[i,j] = chromaticity(A[i,j])
     ...:     return out
     ...: 

In [133]: %timeit org_app(A)
1 loop, best of 3: 5.99 s per loop

In [134]: %timeit np.apply_along_axis(chromaticity, 2, A) #@hpaulj's soln
1 loop, best of 3: 9.68 s per loop

In [135]: %timeit np.log(A/np.power(np.sum(A,2,keepdims=True),1/3))
10 loops, best of 3: 90.8 ms per loop

这就是为什么在使用数组矢量化并在一次工作中处理尽可能多的元素时，总是尝试使用NumPy funcs！