3

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

fsl
  • 3,250
  • 1
  • 10
  • 20
user37292
  • 248
  • 2
  • 12
  • 2
    The correct datatype would be `np.bool_` . You don't need a signature at all if you don't want to compile ahead of time. It will just limit the function to the datatypes. You also don't need to zero the array before writing. `con = np.empty(size, dtype=np.bool_)` is enough and should be a bit faster (You are overwriting the whole array anyway) – max9111 Mar 21 '21 at 20:00
  • @max9111, thanks for your comment, very useful. Just one thing: are you saying that declaring the types, regardless of compiling ahead of time, will not in general produce more efficient code? I would have guessed that a priori knowledge of the types would facilitate compiler's optimisations. Thanks a lot – user37292 Mar 22 '21 at 08:15
  • No generally it isn't more efficient. The first step of jit-compilation is to determine the datatypes (also if a array is non-contiguous, c-contiguous or fortran contiguous). That is exactly what you do if you pass a signature. BTW: If you declare a array explicitly as non-contiguous eg. `boolean[:]` this could worsen the performance `boolean[::1]` would be the recommended signature. https://stackoverflow.com/a/57316221/4045774 – max9111 Mar 22 '21 at 08:37

1 Answers1

2

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)
fsl
  • 3,250
  • 1
  • 10
  • 20
  • Great thanks a lot. Un underscore in ````np.bool_````seems to be missing (are they aliases? on my machine the non-underscore version does not work), but it works great thanks a lot – user37292 Mar 22 '21 at 08:13
  • My bad, it's really supposed to be `np.bool_` as `np.bool` is deprecated (it was an alias for `bool`). – fsl Mar 22 '21 at 23:07