如何直接向量化以索引作为输入并返回 2x2 np.array 的函数

Question

假设我有一个函数

f(i,j,a)

，它从 2D-numPy 数组

a

获取行和列，并返回另一个 2 x 2 numPy 数组，其中

i,j

位于左上角

b

i.e

def f(i,j,a):
 b = a[np.arange(2)[:, np.newaxis]  + i, np.arange(2) + j]
 return b

（上面的代码可以通过用其他数字替换 2 来推广）。我想通过本质上迭代

f(i_array,j_array,a)

来定义

i,j

来向量化这个过程。这对于列表理解来说很简单：

b_array = np.array([ f(i,j,a) for i in i_array, for j in j_array])

.

是否有更有效的方法可以直接使用 numPy 函数/索引来执行此操作，或者列表理解是我能做的最好的方法？我知道

np.vectorize

可以做到这一点，但我不相信它比简单的列表理解更快。

Answer 1

我想通了（但我的解决方案有点笨拙，所以任何修改）将不胜感激）。我的代码如下所示：

import numpy as np

def f(i_array,j_array, a):
 i_transposed, j_transposed = i_array.T, j_array.T
 i_matrix = np.arange(2)[:, np.newaxis][:, np.newaxis]+ i_transposed
 j_matrix = np.arange(2)[:, np.newaxis]+ j_transposed
 b_grids = np.swapaxes(grid[i_matrix, j_matrix].T, 1,2)

 return b_grids

在这里，我们利用广播来生成先前矩阵的

np.array

。

i_matrix

项采用

np.arange(2)

，将其维度增加 2，并使用广播创建一个

(2, 1, 2)

数组，其中所有可能的

i_transposed

添加到 2。

j_matrix

也会发生同样的情况，只不过它低了 1 维。然后我们使用numPy索引在

grid[i_matrix, j_matrix]

中进行广播。由于我们广播的方式，尺寸是错误的，所以我们必须使用

np.swapaxes.