数组中的 Python Numba 值答案

【问题标题】：Python Numba Value in Array数组中的 Python Numba 值
【发布时间】：2023-03-29 14:43:01
【问题描述】：

我正在尝试检查一个数字是否在 int8s 的 NumPy 数组中。我试过这个，但它不起作用。

from numba import njit
import numpy as np

@njit
def c(b):
    return 9 in b

a = np.array((9, 10, 11), 'int8')
print(c(a))

我得到的错误是

Invalid use of Function(<built-in function contains>) with argument(s) of type(s): (array(int8, 1d, C), Literal[int](9))
 * parameterized
In definition 0:
    All templates rejected with literals.
In definition 1:
    All templates rejected without literals.
In definition 2:
    All templates rejected with literals.
In definition 3:
    All templates rejected without literals.
In definition 4:
    All templates rejected with literals.
In definition 5:
    All templates rejected without literals.
This error is usually caused by passing an argument of a type that is unsupported by the named function.
[1] During: typing of intrinsic-call at .\emptyList.py (6)

如何在保持性能的同时解决此问题？将检查数组的两个值，1 和 -1，长度为 32 项。它们没有排序。

【问题讨论】：

Numba 似乎不支持 in 用于 numpy-arrays，但它支持用于集合。所以你唯一要做的就是把 b 改成 set(b)
@max9111 如果我理解正确，那不会需要一点时间吗？我会多次这样做，所以速度很重要。
最快的算法取决于数据。（数组是否已排序？如果它在数组中，是否要检查多个值？）
@max9111 数组未排序，我正在检查两个值，1 和 -1。该数组将有 32 个项目长。我会添加这个。

标签： python arrays numpy numba

【解决方案1】：

检查两个值是否在一个数组中

为了仅检查数组中是否出现两个值，我建议使用简单的蛮力算法。

代码

import numba as nb
import numpy as np

@nb.njit(fastmath=True)
def isin(b):
  for i in range(b.shape[0]):
    res=False
    if (b[i]==-1):
      res=True
    if (b[i]==1):
      res=True
  return res

#Parallelized call to isin if the data is an array of shape (n,m)
@nb.njit(fastmath=True,parallel=True)
def isin_arr(b):
  res=np.empty(b.shape[0],dtype=nb.boolean)
  for i in nb.prange(b.shape[0]):
    res[i]=isin(b[i,:])

  return res

性能

#Create some data (320MB)
A=(np.random.randn(10000000,32)-0.5)*5
A=A.astype(np.int8)
res=isin_arr(A) 11ms per call

因此，通过这种方法，我可以获得大约 29GB/s 的吞吐量，这与内存带宽相差不远。您还可以尝试减小 Testdatasize 以使其适合 L3-cache 以避免内存带宽限制。使用 3.2 MB 的测试数据，我获得了 100 GB/s 的吞吐量（远远超出了我的内存带宽），这清楚地表明此实现受到内存带宽的限制。

【讨论】：

谢谢！我不知道代码可以达到它的速度。
在 for 循环中包含 res=False 似乎是无意的。

【解决方案2】：

这只是 max9111 答案的轻微变化。稍微短一点，并将要搜索的元素作为参数。一旦找到元素，它也会退出循环。

from numba import njit

@njit
def isin(val, arr):
    for i in range(len(arr)):
        if arr[i] == val:
            return True
    return False

【讨论】：