Can't find a working way to implement the 3d dimensional array with numba-CodePudding

I am trying to implement a 4-th order Runge-Kutta method, but with several changes, It is supposed to solve a system of differential equations, that are also spacially distributed, so basically it has to solve for a [n,n] matrix, for each of the elements. And given that it's a system, the result of a function, defining the right sides of the equation would be of shape (2, n, n). That's the code:

@numba.njit
def f1(A, B,a1, klim):
    As = (a1-2 np.sqrt(a1**2-16))/(2*(5-a1))
    Bs = (a1**2-a1-8 (a1-1)*np.sqrt(a1**2-16))/(2*(5-a1))
    return np.array([np.multiply(((a1*A**2)/(1 A B) - A),(A < klim*As)),
        np.multiply(((4*A**2)/(1 A) - B),(B < klim*Bs))])


@numba.njit
def RungeK1step(f,A,B,a1,h,klim):
    resultA = np.copy(A)
    resultB = np.copy(B)
    k0 = f(A, B, a1, klim)  
    k1=f(A   h * k0/2, B   h*k0/2, a1, klim) 
    k2=f(A   h * k1/2,B   h * k1/2,a1, klim)
    k3=f(A   h * k2, B   h * k2, a1, klim) 
    k=h * (k0   2.*k1   2.*k2   k3) / 6
    resultA, resultB = np.array((A,B))   k 
    return resultA, resultB


#initial conditions for A and B
L = 200
N = 1000
xs, delta_x= np.linspace(0,L,N,retstep=True)
ys = np.linspace(0,L,N)
x, y = np.meshgrid(xs,ys)
Ainit = Binit = 10 * (1   10*((x-L/2)**2   (y-L/2)**2))**(-1)

A_array = []
Ainit = np.array(10 * (1   10*((x-L/2)**2   (y-L/2)**2))**(-1))
A = B = Ainit
a1 = 4.2

for i in range(100):
    A , B = RungeK1step(f1, A, B,a1,delta_t,2)
    if (i % 10 == 0):
        A_array.append(A)
        print(i)

And when running the code, all I get is this:

---------------------------------------------------------------------------
TypingError                               Traceback (most recent call last)
~\AppData\Local\Temp/ipykernel_11656/1571656124.py in <module>
      5 
      6 for i in range(100):
----> 7     A , B = RungeK1step(f1, A, B,a1,delta_t,2)
      8     if (i % 10 == 0):
      9         A_array.append(A)

c:\Users\79033\AppData\Local\Programs\Python\Python39\lib\site-packages\numba\core\dispatcher.py in _compile_for_args(self, *args, **kws)
    466                 e.patch_message(msg)
    467 
--> 468             error_rewrite(e, 'typing')
    469         except errors.UnsupportedError as e:
    470             # Something unsupported is present in the user code, add help info

c:\Users\79033\AppData\Local\Programs\Python\Python39\lib\site-packages\numba\core\dispatcher.py in error_rewrite(e, issue_type)
    407                 raise e
    408             else:
--> 409                 raise e.with_traceback(None)
    410 
    411         argtypes = []

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<built-in function array>) found for signature:
 
 >>> array(list(array(float64, 2d, C))<iv=None>)
 
There are 4 candidate implementations:
      - Of which 4 did not match due to:
      Overload in function '_OverloadWrapper._build.<locals>.ol_generated': File: numba\core\overload_glue.py: Line 131.
        With argument(s): '(list(array(float64, 2d, C))<iv=None>)':
       Rejected as the implementation raised a specific error:
         TypingError: Failed in nopython mode pipeline (step: nopython frontend)
       No implementation of function Function(<intrinsic stub>) found for signature:
        
        >>> stub(list(array(float64, 2d, C))<iv=None>)
        
       There are 2 candidate implementations:
         - Of which 2 did not match due to:
         Intrinsic in function 'stub': File: numba\core\overload_glue.py: Line 35.
           With argument(s): '(list(array(float64, 2d, C))<iv=None>)':
          Rejected as the implementation raised a specific error:
            TypingError: array(float64, 2d, C) not allowed in a homogeneous sequence
         raised from c:\Users\79033\AppData\Local\Programs\Python\Python39\lib\site-packages\numba\core\typing\npydecl.py:487
       
       During: resolving callee type: Function(<intrinsic stub>)
       During: typing of call at <string> (3)
       
       
       File "<string>", line 3:
       <source missing, REPL/exec in use?>

  raised from c:\Users\79033\AppData\Local\Programs\Python\Python39\lib\site-packages\numba\core\typeinfer.py:1086

During: resolving callee type: Function(<built-in function array>)
During: typing of call at C:\Users\79033\AppData\Local\Temp/ipykernel_11656/3211464979.py (5)


File "..\..\..\..\Users\79033\AppData\Local\Temp\ipykernel_11656\3211464979.py", line 5:
<source missing, REPL/exec in use?>

During: resolving callee type: type(CPUDispatcher(<function f1 at 0x00000130ACCB0670>))
During: typing of call at C:\Users\79033\AppData\Local\Temp/ipykernel_11656/1432930034.py (5)

During: resolving callee type: type(CPUDispatcher(<function f1 at 0x00000130ACCB0670>))
During: typing of call at C:\Users\79033\AppData\Local\Temp/ipykernel_11656/1432930034.py (5)


File "..\..\..\..\Users\79033\AppData\Local\Temp\ipykernel_11656\1432930034.py", line 5:
<source missing, REPL/exec in use?>

Which obviously means that I have set the returning array wrong and numba doesn't understand it. I tried taking away the square brackets inside the return np.array(...), but to no avail, the error is similar but it's array(array(float64,2d),array(float64,2d)) now. What is the working way to return array here?

CodePudding user response：

The Numba error occurs because the return expression type cannot be deduced. More precisely, it looks like Numba fails to deduce the dimension of the returned array. In fact, even if it would work, it would be pretty inefficient since the two sub-arrays need to be created and appended into a list before creating a new bigger array. This cause necessary expensive memory copies. You can create the output array manually and write the result of the two temporary array in the sub part of the output array (using plain loops).

The thing is there is another issue: the dimension of the input/output array is variable. Indeed, the input of the first call are 2D arrays and the output is a 3D array. The k0 3D array is then used in the expression A h * k0/2 producing another 3D array. This cause the new function call to produce a 4D array from the 3D array, and so on. The final result is a tuple of two 5D arrays of shape (2, 2, 2, 1000, 1000). I do not think this is normal as you said "the equation would be of shape (2, n, n)". You can test that by simply removing the @numba.njit Numba decorators and then running the function. Thus, I think there is a bug in the RungeK1step function.

Note also that the delta_t variable is missing in the code.

Also note that providing the signature of functions help to track abnormal calls earlier resulting in much more user-friendly errors. It also cause functions to be compiled eagerly which is generally also great for benchmarks and early error reports (having an obscure typing error in the middle of a Numba function compiled lazily in a code running since few hours can be quite frustrating ;) ). Here is an example:

@numba.njit('float64[:,:,::1](float64[:,::1], float64[:,::1], float64, int_)')
def f1(A, B, a1, klim):
    # [...]

For more information about signature, please read the documentation.