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