python - 将函数应用于 3D numpy 数组的每个 2D 切片的有效方法

Question

我想对 3D 数组的每个 2D 切片应用一个接受 2D 数组(并返回相同形状之一)的函数。这样做的有效方法是什么？ numpy.fromiter 返回一维数组，numpy.fromfunction 需要分别应用于每个坐标。

目前我正在做

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

这给了我想要的，但列表理解速度很慢。此外，func 是具有特定边界条件的一维导数。 numpy.gradient 似乎只对数组的 N 维进行 N 维导数，但也许有另一个例程可以为我完成所有事情？

编辑:列表理解有效，但我正在寻找一种更快的方法。 bar3D 可以很大，最大可达 (500,500,1000)。我发现的所有用于将函数应用于数组的 numpy 例程似乎都假设函数或数组是一维的。

Answer 1

我不知道有什么通用方法可以将函数应用于数组的 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.]]])

这会区分每个 2D 切片的行。要区分每一列，请执行 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, :] 到正确的边界值。