Python 将函数应用于numpy中的三维阵列

我有一个3D numpy.ndarray（想象一下带有RGB的图像）像

以及处理列表输入的函数

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

如何将此功能应用于每个像素？（或2D阵列的每个1D元素）

我试过的 np.沿_轴应用_ 和我的差不多。所以，我首先尝试了一下

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

起初，我认为这是解决方案，但这太慢了，因为

矢量化 我将np.vetorize应用于我的函数

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.

---

一个选项是转置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]])

---

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

不幸的是，这是一个比正常的

矢量化

慢得多的代码路径，在1.12中，它至少在C中运行for循环

---

在我的机器上，使用numpy

master

，这只比

apply\u沿轴

快10%左右（尽管

apply\u沿轴

自1.12以来的实现发生了巨大变化）

感谢您的回答。他们都工作得很好，速度也很快。我在我的问题中添加了我的基准测试结果。更清楚地说，你可以做

np.moveaxis（a，-1，0）

，读为

moveaxis-1到位置0

，不客气，很高兴这有帮助。谢谢分享计时结果。谢谢@Eric。我添加了基准测试结果。虽然它与转置具有相同的速度，但它似乎是expicit并使其更好（在我看来）

矢量化

在这里对您没有帮助，很遗憾，因为它模拟了

ufunc

，不是一个

gufunc

这三个选项的测试数据有多大？我添加了我在基准测试中实际做的事情。对于那些在这篇评论之前看到这个基准的人，基准设计得很糟糕，没有做我希望他们做的事情。现在，我相信它已经修好了。谢谢你的指点。基准现在不包括创建我的金额。

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

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)