Applying numba to a function whose input is a method of an object

Numba Support: How do I do ...?
1

I have a question about applying numba to a function whose argument is a method of an object. Here’s a simple, hypothetical situation. I am curious if anyone has suggestions on applying numba so that I don’t have to alter the code design as much. In the hypothetical example below, I don’t have loops whatsoever just for simplicity.

Suppose that I have a class of swimming pools. A method of this class is

calc_pressure(self, depth)

where depth is an input.

Now, I have a function (outside of this class), say to compute temperature

def compute_temp(calc_pressure,depth):
return calc_pressure(depth)

Suppose that I initialize an object of this class, say pool1. In my main function, I call compute_temp like this:

temp = compute_temp(pool1.calc_pressure,10)

My problem now is, when I added the decorator @jit(nopython=True) for calc_pressure method and for the compute_temp function, python is complaining to me that in calc_pressure(depth), I am missing an input for depth, presumably because numba doesn’t understand “self”. IF I don’t jit the function, it works fine. How can I apply numba in this situation without changing so much the design of the code?

Suggestions appreciated!

2

Can you post a working minimal reproducer of what you’re trying to do?

3

Hi @nelson2005, thanks for getting back. Here’s a working minimal reproducer of what I’m trying to do without the jit:

class pool():
def init(self):
self.a = 4.3
self.b = 2.7

def helper_pressure(self, c):
    return self.a + self.b + c

def test_main():

endTime = 5

test_obj = pool()

def pressure(c):
    
    return test_obj.helper_pressure(c)

out = calcTemp(endTime, pressure)
    
return out

def calcTemp(endTime, pressure):

temp_sum = 0
for k in range(endTime):
    temp_sum = temp_sum + pressure(1)

return temp_sum

if name == “main”:

out = test_main()

I would like to add @jit(nopython=True) before calcTemp and before the function pressure inside test_main. However, numba is complaining that Failed in nopython mode pipeline (step: nopython frontend)
Untyped global name ‘test_obj’: Cannot determine Numba type of <class ‘main.pool’>.

I’d like to avoid using jitclass since the example above is only a simplification of what I want to do (I plan to have a parent and child class). I just want to jit methods of a class which are not static…

4

Thanks for the code.

I don’t think you can do this as-is… as I understand it, numba njit can’t work well with generic python classes. This might work with and objmode block but once you’re doing that you end up losing much of the benefit of nopython mode.

One of the experts like @stuartarchibald may be able to opine more, but I don’t know how to do what you want to do.