python - 使用 numpy 缩小轴

Question

我有一个 NxMx3 带有 dtype=object 的 numpy 数组。我还有一个函数 f(a,b,c)，它获取此数组最后一个轴中的三个元素并返回一个 np.int32。我的问题是如何将 f 应用到我的 NxMx3 数组以生成一个 NxM 数组 dtype=np.int32?

我目前的解决方案是使用

newarr = np.fromfunction(lambda i,j: f(arr[i,j,0], arr[i,j,1], arr[i,j,2]),
                          arr.shape[:2], dtype=np.int)

尽管这比我希望的要冗长一些。

Answer 1

你可以使用 vectorize :

np.vectorize(f, otypes=[np.int32])(arr[:, :, 0], arr[:, :, 1], arr[:, :, 2])

这可以通过轴滚动和迭代来简化:

np.vectorize(f, otypes=[np.int32])(*np.rollaxis(arr, 2, 0))

或者，您可以使用 dsplit 显式拆分数组:

np.vectorize(f, otypes=[np.int32])(*np.dsplit(arr, 3))[..., 0]

或

np.vectorize(f, otypes=[np.int32])(*np.dsplit(arr, 3)).reshape(arr.shape[:-1])

或

np.vectorize(f, otypes=[np.int32])(*np.dsplit(arr, 3)).squeeze()

但是，apply_along_axis 可能更简单:

np.apply_along_axis(lambda x: f(*x), 2, arr)