jitclass with class composition and class members

[J Ello]

Hi!

I'm currently building a simulation environment in python (modularity is much more important than pure speed, thus python) and now I got to a point, where speeding it up would be a nice thing.
So I thought I could try it using numba as I was quite impressed with the performance boost compared to Cython (well, getting Cython to the same speed just takes alot of time in longer code).

My simulation environment consists of a parent class `ParentCls` which contains the iterating and solving method, as well as a list with the instances of the created subclasses and the connection methods for the subclasses. Thus the subclasses, here `ClsA` and `ClsB`, inherit from the parent class, to access the parent class' methods.
Here is a minimum working example:

import numba
import numpy as np


class ParentCls():
    cls_list = list()
    test = 55
    count_cls = 0
    def __init__(self):
        return
    def add_cls(self, cls_name, **kwargs):
        ParentCls.cls_list.append(cls_name(**kwargs))
        ParentCls.count_cls += 1
    def solver(self, step):
        for classes in ParentCls.cls_list:
            classes.results[step] = classes.calc()
    def iterate(self, num):
        for classes in ParentCls.cls_list:  #preallocate results grid
            classes.results = np.empty(num)
        step = 0
        for step in range(num):  # iterate over solver
            self.solver(step)
            step += 1

spec = [
    ('cls_id', numba.int32),
    ('a', numba.float64),
    ('b', numba.float64),
    ('results', numba.float64[:]),
]


#@numba.jitclass(spec)
class ClsA(ParentCls):
    def __init__(self, a, b):
        self.cls_id = ParentCls.count_cls
        self.a = a
        self.b = b
        self.results = np.array((20))
    def calc(self):
        return (self.a + self.b) * self.test


#@numba.jitclass(spec)
class ClsB(ParentCls):
    def __init__(self, a, b):
        self.cls_id = ParentCls.count_cls
        self.a = a
        self.b = b
        self.results = np.array((20))
    def calc(self):
        return self.a * self.b + self.test


env = ParentCls()

env.add_cls(ClsA, a=5, b=33.7)
env.add_cls(ClsB, a=3.8, b=9)
env.add_cls(ClsB, a=5, b=33.7)
env.add_cls(ClsA, a=0, b=13)

env.iterate(35)

So iterate iterates over the solver which itself iterates over the subclasses saved in `cls_list`.
Since the iteration is numerically quite expensive with a few million steps and inside the solver I have to iterate over, depending on the environment to simulate, at least 100 subclasses which itself do not only have simple `calc` operations like `a+b` but rather array operations over thousands of cells each, I wanted to jit at least the `solver` and `ìterate` methods.
Is there any way to do that? I already "outsourced the `calc` methods to speed up the array operations, now the function overhead of calling these subfunctions is quite big. My ClsA for example now looks like this:

class ClsA(ParentCls):
    def __init__(self, a, b):
        self.cls_id = ParentCls.count_cls
        self.a = a
        self.b = b
        self.results = np.array((20))
    def calc(self):
        return calc_a(self.a, self.b, self.test)

@numba.jit(float64(float64, float64, float64),
           nopython=True, cache=True, nogil=True)
def calc_a(a, b, test):
    return (a + b) * test

I guess I could with some effort also get rid of the class inheritance by outsourcing all methods of the `ParentCls` and jitting them.
But is there any way to have class members or equivalently stored variables which can be accessed from within all classes (parent class and subclasses)? As soon as I turn class members of `ParentCls` in instance variables, I can only access them in subclasses by knowing the instance name of `ParentCls`...
And is there any workaround on how and where to store instances of subclasses in a parent class instance (or if class members are somehow with a workaround supported: how to store them in class members)?
It would also be ok to outsource the `solver` and `iterate` methods to jitted functions to which I pass the (also outsourced and jitted) class methods `calc` of the subclasses as arguments (well, I also know that passing functions as arguments is not allowed, but perhaps somehow...).
Is anything of this possible and reasonably doable? And is there any release date for when class members, inheritance and python dictionairies will be supported?
Sorry for my lengthy question and thanks for help in advance!

Sincerely,
Jo

[J Ello]

Anyone? Any help, also only for small parts of my problem, is appreciated!