对3D numpy数组中的向量应用函数

Question

我有一个关于如何在3D numpy数组中对向量应用函数的问题。我的问题是这样的：假设我有一个像这样的数组：

a = np.arange(24)
a = a.reshape([4,3,2])

我想对所有后面的向量应用一个函数来修改它们：

[0 6], [1 7], [2 8], [4 10], [3 9] ...

使用什么方法最好呢？因为我的数组比较大，在三个维度中的两个维度上循环会很慢……

提前谢谢你们！

Answer 1

你可以使用函数 np.apply_along_axis。在文档中这样描述：

这个函数可以对给定轴上的一维切片应用一个函数。

举个例子：

>>> import numpy as np
>>> a = np.arange(24)
>>> a = a.reshape([4,3,2])
>>> 
>>> def my_func(a):
...   print "vector: " + str(a)
...   return sum(a) / len(a)
... 
>>> np.apply_along_axis(my_func, 0, a)
vector: [ 0  6 12 18]
vector: [ 1  7 13 19]
vector: [ 2  8 14 20]
vector: [ 3  9 15 21]
vector: [ 4 10 16 22]
vector: [ 5 11 17 23]
array([[ 9, 10],
       [11, 12],
       [13, 14]])

在上面的例子中，我使用的是第0个轴。如果你需要处理 n 个轴，你可以执行这个函数 n 次。