How to specify numba jitclass when the class's attribute contains another class instance?

Question

I'm trying to use numba to boost the python performance of scipy.integrate.odeint. To this end, I have to use @nb.jit(nopython=True) for the function defining the ODE system. However, this function has to take another python-class instance as an argument in my program. I had to jit the class also with @nb.jitclass(spec) with appropriate specs. This worked fine, until I found a serious issue when the specs of the class includes another type of class instance as its method. My code is following.

import numba as nb
from scipy.integrate import odeint


spec1=[("hi", nb.i4)]
@nb.jitclass(spec1)
class Hi(object):
    def __init__(self):
        self.hi = 0

spec2=[("dummy", nb.i4), ("dummy1", nb.i4)]
@nb.jitclass(spec2)
class Dummy(object):
    def __init__(self, anotherClassInstance):
        self.dummy = 0
        self.dummy1 = anotherClassInstance

class A:
    def __init__(self, someClassInstance):
        self.a1 = someClassInstance

    def odeSystem(self, x, t):
        return _odeSystem(x, t, self.a1)

    def odeSolve(self, iValues, ts):
        sol = odeint(self.odeSystem, iValues, ts)
        return sol

@nb.jit(nopython=True)
def _odeSystem(x, t, someClassInstance):
    return 1-x



if __name__ == "__main__":
    c = Hi()
    b = Dummy(c)
    a = A(b)
    print a.odeSolve(0.5, range(0, 10))

In summary: So here "class A" is my ode solver.

1. To compile the method "odeSystem" with numba, it must not a class method. So I made another function outside of the class "_odeSystem".

2. Unfortunately however, my odeSystem has to have a class instance as an argument. Hence I used @jitclass to compile the class instance argument properly.

3. I again encountered another problem, where this class "Dummy" also takes another type of class instance as one of its attribute. I have no idea how to set "spec" for this class. I tried the type for "dummy1" with "nb.typeof(Hi)" but it didn't work.

Please help me. Thanks in advance.

Answer 1

You can use .class_type.instance_type in your spec definition for holding an instance of another type. See example in the numba source tree here

spec2=[("dummy", nb.i4), ("dummy1", Hi.class_type.instance_type)]
@nb.jitclass(spec2)
class Dummy(object):
    def __init__(self, anotherClassInstance):
        self.dummy = 0
        self.dummy1 = anotherClassInstance