我如何使用 numba 更改数组中的索引

Question

我有一个函数，我在其中执行一些操作并想用 numba 加快它的速度。在我的代码中，使用高级索引更改数组中的值不起作用。我认为他们确实在 numba 文档中这么说了。但是像 numpy.put()?

这样的解决方法是什么？

这是我想做的一个简短例子：

#example array
array([[ 0,  1,  2],
       [ 0,  2, -1],
       [ 0,  3, -1]])

使用在 numba 中工作的任何方法更改给定索引处的值...以获得：更改值位于：[0,0]、[1,2]、[2,1]

#changed example array by given indexes with one given value (10)
array([[ 10,  1,  2],
       [ 0,  2, 10],
       [ 0,  10, -1]])

这是我在 python 中所做的，但没有使用 numba：

indexList 是一个元组，与 numpy.take()

一起使用

这是工作示例 python 代码，数组中的值更改为 100。

x = np.zeros((151,151))
print(x.ndim)
indexList=np.array([[0,1,3],[0,1,2]])
indexList=tuple(indexList)

def change(xx,filter_list):
    xx[filter_list] = 100
    return xx

Z = change(x,indexList)

现在在函数上使用@jit:

@jit
def change(xx,filter_list):
    xx[filter_list] = 100
    return xx

Z = change(x,indexList)

编译正在回退到对象模式并启用循环提升，因为函数“更改”失败类型推断，原因是：未找到用于签名的函数 Function() 的实现：setitem(array (float64, 2d, C), UniTuple(array(int32, 1d, C) x 2), Literalint)

出现这个错误。所以我需要一个解决方法。 numba 不支持 numpy.put()。

如有任何想法，我将不胜感激。

谢谢

Answer 1

如果将 indexList 保留为数组没有问题，您可以将它与 change 函数中的 for 循环结合使用，以使其与 numba 兼容:

indexList = np.array([[0,1,3],[0,1,2]]).T

@njit()
def change(xx, filter_list):
    for y, x in filter_list:
        xx[y, x] = 100
    return xx
change(x, indexList)

请注意，indexList 必须转置才能使 y、x 坐标沿第一轴。换句话说，要更改 n 点，它必须具有 (n, 2) 而不是 (2, n) 的形状。实际上它现在是一个坐标列表：[[0, 0],[1, 1],[3, 2]]

Answer 2

@mandulaj 发布了要走的路。在 mandulaj 给出他的回答之前，我采用了一种不同的方式。

使用此功能我会收到弃用警告...所以最好使用@mandulaj 并且不要忘记转置 indexList。

@jit
def change_arr(arr,idx,val): # change values in array by np index array to value
    for i,item in enumerate(idx):
        arr[item[0],item[1]]= val
    return arr

我如何使用 numba 更改数组中的索引

how do i change indexes in an array using numba

python

arrays

indexing

numpy

numba