时间戳

首页 > 解决方案 > Numba 中的布尔签名

python - Numba 中的布尔签名

问题描述

我在我的一些代码上使用了 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

con它根据某些参数生成一个 1 和 0 的 Numpy 数组。该数组由 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)

如何做到这一点？

谢谢你的帮助

标签： pythonnumpynumba

解决方案

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

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

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

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

推荐阅读