二维numpy阵列的映射函数

Question

我有一个函数foo，它以NxM numpy数组作为参数并返回标量值。我有一个AxNxM numpy数组data，我想在这个数组上映射foo，从而得到长度为a的numpy数组

当然，我正在做的是：

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

它是有效的，但似乎我没有利用核弹魔法（和速度）。有更好的办法吗？

我已经看过numpy.vectorize和numpy.apply_along_axis，但是这两个都不适用于2D数组的函数。

编辑：我正在对24x24图像补丁进行增强回归，所以我的AxNxM类似于1000x24x24。我在上面所说的foo将类似Haar的特性应用到一个补丁（所以，计算量不是很高）。

Answer 1

如果NxM很大（比如100），它们在A上的迭代成本将基本上分摊为零。

假设数组是1000 X 100 X 100。

迭代是O（1000），但是内部函数的累计开销是O（1000 X 100 X 100）-慢10000倍。（注意，我的术语有点古怪，但我知道我在说什么）

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

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

你可以在建立列表时节省大量的内存分配。。。但是循环开销会更大。

或者可以编写一个并行版本的循环，并将其拆分为多个进程。这可能要快得多，这取决于foo的密集程度（因为它必须抵消数据处理）。

Answer 2

您可以通过将3D数组重塑为具有相同前导维度的2D数组来实现这一点，并根据foo的要求将函数foo包装为可用于1D数组的函数。例如（使用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]