python - 用于数组索引的 Numba 列表参数
问题描述
我希望利用 numba 迭代一个大型二维数组,其中对于迭代,数组的一个子集将由 [x, y] 位置(作为参数传递)选择。我在构建这种与 numba 配合得很好的正确方法时遇到了麻烦,特别是在将列表列表作为参数传递给方法时。任何指针?
x_y_list = [[1, 2], [3, 4], [5, 6]]
array = ([[1, 2, 3, 4, 5, 6],
[1, 2, 3, 4, 5, 6],
[1, 2, 3, 4, 5, 6]])
@jit
def arrIndexing(array, x_y_list):
for index in x_y_list:
subset = array[index[0]-1:index[0]+1, index[1]-1:index[1]+1]
# do some other stuff
解决方案
像这样的东西?应该用 Numba 做得很好,但我没有测试过(在我不支持 Numba 的手机上做了这个)
import numpy as np
def xy():
x_y_list = np.array([[1, 2], [2, 4], [0, 5]])
array = np.array([[1, 2, 3, 4, 5, 6],[1, 2, 3, 4, 5, 6],[1, 2, 3, 4, 5, 6]])
for i,j in x_y_list:
print(array[np.ix_((i-1, i), (j-1, j))])
>>> xy()
[[2 3]
[2 3]]
[[4 5]
[4 5]]
[[5 6]
[5 6]]