python - 用 numpy 减少轴

Question

我有一个NxMx3numpy 数组dtype=object。我还有一个函数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)