Numba cache and iterative solving

Viewed 48

I am trying to use NumbaLSODA, see here, to solve many systems iteratively. However, it seems that I can't put the function in the numba cache and so I have a big overhead at each iteration. I was wondering if anybody could help.

This is my code. First I put all numba related function inside a file

numba_func.py

NOPYTHON = True
CACHE = True
PARALLEL = True

@nb.cfunc(lsoda_sig, nopython=NOPYTHON, cache=CACHE)
def rhs(x, i, di, p):
    di[0] = ...
    di[1] = ...
    di[2] = ...

funcptr = rhs.address

@nb.njit('(float64)(float64, float64, float64, int16)', nopython=NOPYTHON, parallel=PARALLEL, cache=CACHE)
def solve(a, b, c, funcptr):

    w0 = np.array([a, b, ...], dtype=np.float64)
    p  = np.array([c], dtype=np.float64)

    x = np.linspace(0, 100, 500)

    usol, success = lsoda(funcptr, w0, x, data=p)
    
    return usol[-1][1]

Then I use another file to solve the systems one after the others

from numba_func import solve, funcptr

gs = []
for a, b, c in zip(as, bs, cs):
    gs = np.append(gs, solve(a, b, c, funcptr))

I got the following warning:

NumbaWarning: Cannot cache compiled function "solve" as it uses dynamic globals (such as ctypes pointers and large global arrays)

I guess the idea is to pass the variable funcptr correctly so that numba is happy but so far I failed...

0 Answers
Related