2D numpy数组的映射函数

Question

我有一个叫做 foo 的函数，它接收一个 NxM 的 numpy 数组作为输入，然后返回一个标量值。现在我有一个 AxNxM 的 numpy 数组 data，我想把 foo 应用到这个数组上，最终得到一个长度为 A 的 numpy 数组。

目前，我是这样做的：

result = numpy.array([foo(x) for x in data])

这样做是有效的，但我觉得我没有充分利用 numpy 的强大功能和速度。有没有更好的方法呢？

我查过 numpy.vectorize 和 numpy.apply_along_axis，但这两个方法都不适合处理二维数组的函数。

补充说明：我正在对 24x24 的图像块进行增强回归，所以我的 AxNxM 大约是 1000x24x24。我上面提到的 foo 函数是用来对一个图像块应用类似 Haar 的特征（所以计算量不是很大）。

Answer 1

你可以通过把你的三维数组变成二维数组来实现这个目标，保持它的第一维不变。然后，你可以用一个函数把你的函数foo包裹起来，这样它就能处理一维数组了，前提是把这些一维数组调整成foo所需要的格式。下面是一个例子（这里用trace代替foo）：

from numpy import *

def apply2d_along_first(func2d, arr3d):
    a, n, m = arr3d.shape
    def func1d(arr1d):
        return func2d(arr1d.reshape((n,m)))
    arr2d = arr3d.reshape((a,n*m))
    return apply_along_axis(func1d, -1, arr2d)

A, N, M = 3, 4, 5
data = arange(A*N*M).reshape((A,N,M))

print data
print apply2d_along_first(trace, data)

输出结果：

[[[ 0  1  2  3  4]
  [ 5  6  7  8  9]
  [10 11 12 13 14]
  [15 16 17 18 19]]

 [[20 21 22 23 24]
  [25 26 27 28 29]
  [30 31 32 33 34]
  [35 36 37 38 39]]

 [[40 41 42 43 44]
  [45 46 47 48 49]
  [50 51 52 53 54]
  [55 56 57 58 59]]]
[ 36 116 196]

Answer 2

如果 NxM 的值很大（比如说，100），那么遍历数组 A 的成本几乎可以忽略不计。

假设这个数组的大小是 1000 X 100 X 100。

遍历这个数组的时间复杂度是 O(1000)，但是里面函数的总成本是 O(1000 X 100 X 100)，也就是慢了 10,000 倍。（注意，我的术语可能有点不太准确，但我知道我在说什么）

我不太确定，但你可以试试这个：

result = numpy.empty(data.shape[0])
for i in range(len(data)):
    result[i] = foo(data[i])

这样在构建列表的时候可以节省一些内存分配的开销……不过循环的开销可能会更大。

或者你可以写一个并行版本的循环，把它分散到多个进程中去。这样可能会快很多，具体要看 foo 的复杂程度（因为这需要处理数据的开销）。