Models Stefanescu-Jirsa 2D and 3D in TVB have more than one mode, and only a python defined dfun.
They could be rewritten with a numba dfun, considering the following aspects:
- other numba gufuncs drop the mode dimension, which isn't ok here, so the gufunc signature should have the extra mode dimension
- some derived parameters are matrices and need to have a (n_mode, n_mode) shape in the gufunc signature instead of (1,)
- these models have a dot product between matrix parameters along the mode dimension, might require tweaking
This could be a unit test:
class TestNumbaDfun: def dfun(self, *args, **kwargs): f1 = self._numpy_dfun(*args, **kwargs) f2 = super().dfun(*args, **kwargs) assert np.allclose(f1, f2) return f2 class TestModel(TestNumbaDfun, ReducedHindmarshRoseNumba): pass sim = simulator.Simulator(model=TestModel(), ...) sim.run()