將函數應用於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]

我應該怎么做才能將此功能應用於每個像素？ （或2D數組的每個1D元素）

我嘗試過的

np.apply_along_axis

這個問題和我的一樣。 因此，我首先嘗試了它。

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不是為了提高速度

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筆記本並為每個測試重新啟動了內核）

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數組，將第三個軸交換為第一個軸，然后可以將函數直接應用於它：

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

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

或將sum函數重寫為：

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

Answer 2

從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％（盡管apply_along_axis自1.12以來已大大改變了實現）