[Python] python class and numba jitclass for codes with numba functions

At some point in my code, I call a Numba function and all subsequent computations are made with Numba jitted functions until post-processing steps.

Over the past days, I have been looking for an efficient way to send to the Numba part of the code all the variables (booleans, integers, floats, and float arrays mostly) while trying to keep the code readable and clear. In my case, that implies limiting the number of arguments and, if possible, regrouping some variables depending on the system they refer to.

I identified four ways to do this:

1. brut force: sending all the variables one by one as arguments of the first called Numba function. I find this solution is not acceptable as it makes the code barely readable (very large list of arguments) and inconsistent with my wish to my wish of regrouping variables,
2. Numba typed dictionaries (see this post for instance): I did not find this solution acceptable as I understand a given dictionary can only contain variables of similar types (dictionary of floats64 only for instance) while a given system may have related variables of different types. In addition, I observed a significant loss of performance (~ +10% computation times) with this option,
3. Numba namedtuples: fairly easy to implement and use, but it is my understanding that these can only be efficiently used if defined within a Numba function and thus cannot be sent from a non jitted function/code to a jitted function without making it impossible to benefit from the cache=True option. This is a dealbreaker for me as compilation times may exceed the execution of the code itself.
4. Numba @jitclass: I was initially reluctant to use classes for my code but it turns out it is very practical... but same as for the namedtuples, if an object from a @jitclass is initialized within a non jitted function, I observed that it becomes impossible to benefit from the cache=True option (see this post).

In the end, none of these four alternatives allow me to do what I wanted to. I am probably missing something...

Here is what I did in the end : I combined the use of regular python classes and Numba @jitclass in order to maintain the possibility to benefit from the cache=True option.

Here is my mwe:

import numba as nb
from numba import jit
from numba.experimental import jitclass

spec_cls = [
('a', nb.types.float64),
('b', nb.types.float64),
]

# python class
class ClsWear_py(object):
def __init__(self, a, b):
self.a = a
self.b = b

# mirror Numba class
@jitclass(spec_cls)
class ClsWear(object):
def __init__(self, a, b):
self.a = a
self.b = b

def function_python(obj):
print('from the python class :', obj.a)
# call of a Numba function => this is where I must list explicitly all the keys of the python class object
oa, ob = function_numba(obj.a, obj.b)
return obj, oa, ob

@jit(nopython=True)
def function_numba(oa, ob):
# at the beginning of the Numba function, the arguments are used to define the @jitclass object
obj_nb = ClsWear(oa, ob)
print('from the numba class :', obj_nb.a)
return obj_nb.a, obj_nb.b

# main code :
obj_py = ClsWear_py(11,22)

obj_rt, a, b = function_python(obj_py)


The output of this code is :

$ python mwe.py
from the python class : 11
from the numba class : 11.0


On the plus side :

• I have a clean data structure in python and Numba (use of classes)
• I keep a fast running code and cache=True is working

But on the down side :

• I must define python classes and their mirror in Numba
• there remains one barely readable part of the code: the first call of a jitted function where all the content of my objects must be listed explicitly

Am I missing something ? Is there a more obvious way to do this ?

Continue reading...