python - 如何向量化参数为numpy中向量的函数

Question

我有一个维度为 (n, n, n, 3) 的 numpy 数组 R 和一个函数 f ，它将一维向量转换为标量。我需要一个新数组 A，它与 R 的关系是

A[i, j, k] = f(R[i, j, k, :])

如果没有三个 for 语句，我如何在 numpy 中执行此操作。

Answer 1

理想情况下，您可以通过更改 f 的实现来使用适当处理高维输入的技术来做到这一点。例如，您可以将 np.sum(whatever) 更改为 np.sum(whatever, axis=-1) 以获得最后一个轴而不是整个轴的总和大批。这将产生最有效的结果，但这可能很困难或不可能，具体取决于 f。

较慢但更容易的答案是 np.apply_along_axis :

A = np.apply_along_axis(f, -1, R)

这比 3 个 for 循环更漂亮，但它可能不会更有效。