Numba中的布尔签名

Question

我在我的一些代码上尝试了Numba，通过添加一些@jit装饰器，我获得了相当大的性能改进，这非常棒

试图压缩更多的内容，我想键入函数输出，因为我只需要一个布尔数组，其中一个函数参数是整数。尽管如此，我还是遇到了一些困难。作为一个最低限度的工作示例，我有一个函数，如

@jit
def start_config(size,  s, b):
    prob = np.exp(-energy(1,s)/b)
    con = np.zeros(size)
    rand_unif = np.random.uniform(0,1,size)
    for i in range(size):
        if rand_unif[i] < prob:
            con[i] = 1
        else:
            con[i] = 0
    return con

它根据某些参数生成1和0的Numpy数组con。 这个数组由Float64数字填充，这是一个相当过分的做法，因为我会避开布尔值

第一件事是，如果我尝试将布尔赋值给con，如中所示

@jit
def start_config(size,  s, a):
    prob = np.exp(-energy(1,s)/a)
    con = np.zeros(size)
    rand_unif = np.random.uniform(0,1,size)
    for i in range(size):
        if rand_unif[i] < prob:
            con[i] = True
        else:
            con[i] = False
    return con

它仍然返回Float64元素，这是我的第一个困惑点

于是我试了一下

@jit(nopython = True)
def start_config(size,  s, a):
    prob = np.exp(-energy(1,s)/a)
    con = np.zeros(size, dtype=bool)
    rand_unif = np.random.uniform(0,1,size)
    for i in range(size):
        if rand_unif[i] < prob:
            con[i] = True
        else:
            con[i] = False
    return con

没有了Numba，数组现在被初始化为布尔值，我可以确认。通过检查类型。但是如果我用Numba的@jit进行装饰，调用函数时会出现错误

No implementation of function Function(<built-in function zeros>) found for signature:

最后，我想使用Numba，甚至添加一个函数签名，因为我认为这将有助于Numba，比如（这个语法正确吗？）

@jit(boolean(int32, float64, float64), nopython=True)

如何做到这一点

谢谢你的帮助

Answer 1

您对np.zeros的调用正在中断，因为numba在使用nopython标志时需要实际的numpy类型。只需将其切换到numpy版本，它就可以正常工作：

con = np.zeros(size, dtype=np.bool_)

关于第二点，您几乎是正确的，您需要将返回类型声明为布尔数组（此外，我发现只传递字符串而不是导入实际类型非常方便）：

@jit("boolean[:](int32, float64, float64)", nopython=True)