在numpy中将函数应用于3D数组

Question

我有一个 3D numpy.ndarray（想想一个带有 RGB 的图像），比如

a = np.arange(12).reshape(2,2,3)
'''array(
  [[[ 0,  1,  2],　[ 3,  4,  5]],
   [[ 6,  7,  8],　[ 9, 10, 11]]])'''

还有一个处理列表输入的函数；

my_sum = lambda x: x[0] + x[1] + x[2]

如何将此函数应用于每个像素？（或二维数组的每个一维元素）

我尝试过的

np.apply_along_axis

This question 和我的一样。所以，我先试了一下。

np.apply_along_axis(my_sum, 0, a.T).T #EDIT np.apply_along_axis(my_sum, -1, a) is better

起初，我以为这是解决方案，但这太慢了，因为np.apply_along_axis is not for speed

np.vectorize

我将 np.vetorize 应用于 my_func。

vector_my_func = np.vectorize(my_sum)

但是，我什至不知道如何调用这个矢量化函数。

vector_my_func(0,1,2) 
#=> TypeError: <lambda>() takes 1 positional argument but 3 were given

vector_my_func(np.arange(3)) 
#=> IndexError: invalid index to scalar variable.

vector_my_func(np.arange(12).reshape(4,3)) 
#=> IndexError: invalid index to scalar variable.

vector_my_func(np.arange(12).reshape(2,2,3)) 
#=> IndexError: invalid index to scalar variable.

我完全不知道该怎么做。

编辑

建议方法的基准测试结果。 （每次测试使用 jupyter notebook 并重启内核）

a = np.ones((1000,1000,3))
my_sum = lambda x: x[0] + x[1] + x[2]
my_sum_ellipsis = lambda x: x[..., 0] + x[..., 1] + x[..., 2]
vector_my_sum = np.vectorize(my_sum, signature='(i)->()')

%timeit np.apply_along_axis(my_sum, -1, a)
#1 loop, best of 3: 3.72 s per loop

%timeit vector_my_sum(a)
#1 loop, best of 3: 2.78 s per loop

%timeit my_sum(a.transpose(2,0,1))
#100 loops, best of 3: 12 ms per loop

%timeit my_sum_ellipsis(a)
#100 loops, best of 3: 12.2 ms per loop

%timeit my_sum(np.moveaxis(a, -1, 0))
#100 loops, best of 3: 12.2 ms per loop

Answer 1

从 numpy 1.12 开始，vectorize 获得了 signature 参数。所以你可以把它当作：

my_sum = lambda x: x[0] + x[1] + x[2]
vector_my_sum = np.vectorize(my_sum, signature='(i)->()')  # vector to scalar
vector_my_sum(a)

不幸的是，这是一个比普通vectorize 慢得多的代码路径，在 1.12 中至少运行 C 中的 for 循环。

在我的机器上，使用 numpy master，这仅比 apply_along_axis 快 10% 左右（尽管自 1.12 以来 apply_along_axis 改变了实现方式）

Answer 2

一种选择是将numpy数组转置，将第三个轴交换为第一个，然后您可以直接将函数应用于它：

my_sum(a.transpose(2,0,1))

#array([[ 3, 12],
#       [21, 30]])

或者将求和函数改写为：

my_sum = lambda x: x[..., 0] + x[..., 1] + x[..., 2]
my_sum(a)
#array([[ 3, 12],
#       [21, 30]])