Boolean signature in Numba

Question

I had a go with Numba on some code of mine, and I got quite a performance improvement just by adding some @jit decorators, which was great.

Trying to squeeze something more I would like to type the function output, as I need an array booleans only, and one of the function arguments is an integer. Nevertheless I ran into some difficulties. As a minimum working example, I have a function such as

@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

which generates a Numpy array con of 1 and 0s, based on certain parameters. The array is populated by Float64 numbers, which is quite an overkill, as I would get away with booleans.

The first thing is, if I try to assign booleans to con as in

@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

it still returns Float64 elements, which is my first point of confusion.

I hence tried

@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

Without Numba, the array is now initialised to booleans, as I can confirm. by checking type. But if I decorate with Numba's @jit, I get an error upon calling the function,

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

Ultimately, I would like to use Numba and even add a function signature, as I assume this would help Numba, such as (is this syntax correct at all?)

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

how to achieve this?

Thanks for any help

Answer 1

Your call to np.zeros is breaking because numba requires actual numpy-like types when using the nopython flag. Just switch it to the numpy version and it should work just fine:

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

On your second point, you are almost correct, you need to declare the return type as an array of booleans (besides I find it handy to just pass strings instead of importing the actual type):

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