多个数组的Python apply_along_axis

Question

如果我有一个函数 f(x)，它接受一个一维数组作为参数并生成一个一维数组作为输出，我可以使用 numpy.apply_along_axis将函数应用于二维数组 X 的每一行，其行是 f 的有效参数。

现在我想用一个带有两个参数的函数来做类似的事情。例如。我有一个函数 f(x, y)，它将两个一维数组作为参数，我还有两个二维数组 X、Y，它们都有 n 行。我想将 f 应用于每一对行，生成一个又包含 n 行的数组。

如何以高效的方式实现这一目标？

我也对变体感兴趣，其中 f 接受更多参数或涉及更高维数组:

例如 f 可以取 3 个形状为 (2,2) 的数组 x, y, z； (3,); (5,)并产生形状 (4,4) 的结果。

我有 X、Y、Z 形状 (50, 100, 2, 2)； (50, 100, 3); (50, 100, 5)并想要形状 (50, 100, 4, 4) 的结果

Answer 1

查看 numpy.apply_along_axis 的代码，我看到它只是迭代其他维度，将您的函数应用于每个“行”。有额外的代码允许大约 2 维。但是对于 2d X 它归结为:

result = np.empty_like(X)
for i, x in enumerate(X):
    result[i] = func1d(x)

还有一些代码可以推断结果 应该具有什么形状。例如，如果 func1d 是 np.sum，那么 result 将是 1d，而不是像输入那样的 2d。

所以这个函数没有特别的“效率”。多输入的扩展可以是普通的 Python zip:

 result = np.empty_like(X)
 for i,(x,y) in enumerate(zip(X,Y)):
     result[i] = func1d(x,y)

---

np.ndindex 是生成索引的便捷工具。值得一看它的代码。它使用通用的 numpy 迭代器 np.nditer，见:http://docs.scipy.org/doc/numpy/reference/arrays.nditer.html

For example f could take 3 arrays x, y, z of shape (2,2); (3,); (5,) and produce a result of shape (4,4).

I have X, Y, Z of shapes (50, 100, 2, 2); (50, 100, 3); (50, 100, 5) and want a result of shape (50, 100, 4, 4)

for i,j in np.ndindex(50,100):
    result[i,j,:,:] = f(X[i,j,:,:], Y[i,j,:,:], Z[i,j,:,:])

':' 不是必需的，但要明确表示我们在两个维度上建立索引，并对其余维度进行切片。如果您想迭代第 1 维和第 3 维，并对第 2 维进行切片，则需要它们。