高效地对3D numpy数组的每个2D切片应用函数

Question

我想对一个三维数组的每个二维切片应用一个函数，这个函数输入一个二维数组并返回一个同样形状的二维数组。有没有什么高效的方法可以做到这一点？numpy.fromiter 返回的是一维数组，而 numpy.fromfunction 需要对每个坐标单独应用。

目前我在做的是

foo = np.array([func(arg, bar2D) for bar2D in bar3D])

这样做能得到我想要的结果，但使用列表推导式的速度非常慢。而且，func 是一个一维的导数，带有特定的边界条件。numpy.gradient 似乎只处理 N 维的导数，其中 N 是数组的维度，但也许还有其他方法可以帮我完成这个任务？

编辑: 列表推导式可以工作，但我在寻找更快的方法。bar3D 可能很大，最大可以达到 (500,500,1000)。我找到的所有 numpy 方法似乎都假设函数或数组是 1 维的。

Answer 1

假设你有一个数组，叫做a：

>>> a=np.random.random((4,3,2))

array([[[ 0.27252091,  0.78545835],
        [ 0.83604934,  0.48509821],
        [ 0.77828735,  0.26630055]],

       [[ 0.98623474,  0.29839813],
        [ 0.15893604,  0.61870988],
        [ 0.62281607,  0.27193647]],

       [[ 0.47976331,  0.2471835 ],
        [ 0.77323041,  0.30137068],
        [ 0.52906156,  0.53950597]],

       [[ 0.59207654,  0.86355457],
        [ 0.50250812,  0.75688653],
        [ 0.91046136,  0.5785383 ]]])

你可以这样访问二维切片：

>>> for x in range(a.shape[0]):
        print a[x,:,:]

>>> for x in range(a.shape[1]):
        print a[:,x,:]

>>> for x in range(a.shape[2]):
        print a[:,:,x]

Answer 2

我不知道有没有一种通用的方法可以对数组的N维切片应用函数。不过，有两种方法可以解决这个问题。

如果你想对每个二维切片的每一行或每一列应用一维导数，这其实相当于对每个一维切片进行导数计算，你可以使用np.apply_along_axis：

values = np.arange(4)*np.arange(3)[:, None]+np.arange(2)[:, None, None]*2
>>> array([[[0, 0, 0, 0],
            [0, 1, 2, 3],
            [0, 2, 4, 6]],

       [[2, 2, 2, 2],
        [2, 3, 4, 5],
        [2, 4, 6, 8]]])

np.apply_along_axis(np.gradient, 2, values)
>>> array([[[ 0.,  0.,  0.,  0.],
            [ 1.,  1.,  1.,  1.],
            [ 2.,  2.,  2.,  2.]],

           [[ 0.,  0.,  0.,  0.],
            [ 1.,  1.,  1.,  1.],
            [ 2.,  2.,  2.,  2.]]])

这个方法会对每个二维切片的行进行导数计算。如果你想对每一列进行导数计算，可以用np.apply_along_axis(np.gradient, 2, values)。

如果你想做一些需要用到两个维度的操作，通常可以通过广播和轴参数来实现。例如，如果你想要V[i, j] = sqrt((V[i,j]-V[i, j-1])^2 + (V[i, j]-V[i-1, j])^2)在每个切片V上进行计算，你可以这样做：

xdiffs = np.zeros_like(values) 
xdiffs[:, 1:, :]= np.diff(values, axis=1) 

ydiffs = np.zeros_like(values)
ydiffs[:, :, 1:] = np.diff(values, axis=2)

diffnorms = np.linalg.norm(xdiffs, ydiffs)

>>> array(
  [[[ 0.        ,  0.        ,  0.        ,  0.        ],
    [ 0.        ,  1.41421356,  2.23606798,  3.16227766],
    [ 0.        ,  2.23606798,  2.82842712,  3.60555128]],

   [[ 0.        ,  0.        ,  0.        ,  0.        ],
    [ 0.        ,  1.41421356,  2.23606798,  3.16227766],
    [ 0.        ,  2.23606798,  2.82842712,  3.60555128]]])

虽然调整维度有点麻烦，但通常这是最有效的解决方案。

这个例子在边界使用了零值，如果你需要其他值，就需要把normdiff[:, :, 0]和normdiff[:, 0, :]设置为正确的边界值。