I was in a very similar situation and I found a way to use a Numba-JITed function inside of a class.
The trick is to use a static method, since this kind of methods are not called prepending the object instance to the argument list. The downside of not having access to self
is that you cannot use variables defined outside of the method. So you have to pass them to the static method from a calling method that has access to self
. In my case I did not need to define a wrapper method. I just had to split the method I wanted to JIT compile into two methods.
In the case of your example, the solution would be:
from numba import jit
class MyClass:
def __init__(self):
self.k = 1
def calculation(self):
k = self.k
return self.complicated([1,2,3],k)
@staticmethod
@jit(nopython=True)
def complicated(x,k):
for a in x:
b = a**2 .+ a**3 .+ k