I’ve been fiddling with numba for a while. I’m quite happy with the way it feels–but I am worried that some things I am doing are unnecessarily clunky.
One of the things I’m trying to do is develop a scientific library to be hopefully used with relative ease by someone who doesn’t have much knowledge of numba (or even python).
The functionality I’m trying to achieve is the folloqing:
- The user first defines a class by passing a series of environmental parameters, which, for commodity and also ease of generalization, are given as a dictionary. The class may process these parameters somehow.
- The user can then call the class methods to perform operations on external variables, the result of which depends on the environmental parameters stored in the class. Naturally, these operations are very expensive which is where numba comes in.
Doing this in jit has proven a bit tricky.
- It is not possible to use a jitclass, because there are pretty strong restrictions on the kind of parameters that can be passed through a dictionary (I tried asking about this before here).
- Using a standard python class has its own issues, however, because if we decorate a method with @njit, we cannot use self and access the class attributes.
- The solution I’ve arrived at (similar to this other issue) is to define a first class method which acts as a “builder” for a jitted function, which then is stored as an instance attribute of the class and can be called by the user. An example is below:
Class Summer():
def __init__(self, pars):
self.a = pars["a"]
self.b = pars["b"]
self.build_sum()
def build_sum(self):
a = self.a
b = self.b
@jit
def jit_sum(x,y):
return a*x + b*y
self.jit_sum = jit_sum
def sum(self, x, y):
return self.jit_sum(x,y)
# this second method could be skipped altogether, in this simple case, by self.sum = jit_sum in build_sum.
I’ve actually been quite happy with this method in another project, where I’ve managed to achieve the intended result while reaching a significant speedup. In this case, the environmental parameters are actually large matrices produced by the class itself, but once created they never change, which is why the approach is successful.
However, there are two issues: first, as the code grows more complicated, keeping tracks of the builder functions and the corresponding jitted functions becomes more cumbersome. This is especially true if we want to split some funcionality across multiple classes. Second, I’m uncertain on how to document the jitted functions; a docstring describing build_sum or sum is not very useful, but a docstring on the jitted function would not be recognized as a class method. Documenting jit_sum in the docstring for build_sum, however, would be confusing.
I am interested in understanding if I’m doing the “right” thing–maybe there are more straightforward ways of achieving the same functionality in numba that I missed or didn’t understand properly. If not, are there any tricks to keep in mind when extending this approach to more complex problems, and is there an established way to write the documentation?