找不到用 numba 实现 3d 维数组的工作方法

Question

我正在尝试实现 4 阶 Runge-Kutta 方法，但有一些变化，它应该求解微分方程组，它们也是空间分布的，所以基本上它必须求解[n,n] 矩阵，用于每个元素。鉴于它是一个系统，函数的结果，定义方程右侧的形状为 (2, n, n)。那是代码：

@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)

当运行代码时，我得到的是：

---------------------------------------------------------------------------
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:\Users033\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:\Users033\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:\Users033\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:\Users033\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:\Users033\AppData\Local\Temp/ipykernel_11656/3211464979.py (5)


File "..\..\..\..\Users033\AppData\Local\Temp\ipykernel_1165611464979.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:\Users033\AppData\Local\Temp/ipykernel_11656/1432930034.py (5)

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


File "..\..\..\..\Users033\AppData\Local\Temp\ipykernel_1165632930034.py", line 5:
<source missing, REPL/exec in use?>

这显然意味着我将 returning 数组设置错误并且 numba 不理解它。我试过去掉return np.array(...)里面的方括号，但是没有用，报错是类似的，但是现在是array(array(float64,2d),array(float64,2d))。 return 数组的工作方式是什么？

Answer 1

Numba 错误的发生是因为无法推导return 表达式类型。更准确地说，看起来 Numba 无法推断出 returned 数组 的维数。事实上，即使它可以工作，它也会非常低效，因为在创建一个新的更大的数组之前需要创建两个 sub-arrays 并将其附加到列表中。这会导致必要的昂贵内存副本。您可以手动创建输出数组并将两个临时数组的结果写入输出数组的子部分（使用普通循环）。

还有一个问题：input/output 数组的维数是可变的。实际上，第一次调用的输入是 2D 数组，输出是 3D 数组。然后在表达式 A + h * k0/2 中使用 k0 3D 数组生成另一个 3D 数组。这导致新函数调用从 3D 数组生成 4D 数组，依此类推。最终结果是两个形状为 (2, 2, 2, 1000, 1000) 的 5D 数组的元组。我不认为这是正常的，正如你所说的“方程的形状是 (2, n, n)”。您可以通过简单地删除 @numba.njit Numba 装饰器然后运行函数来测试它。因此，我认为 RungeK1step 函数中存在错误。

另请注意，代码中缺少 delta_t 变量。

另请注意，提供函数签名有助于更早地跟踪异常调用，从而导致更多 user-friendly 错误。它还会导致函数被急切地编译，这通常也非常适合基准测试和早期错误报告（在代码运行中延迟编译的 Numba 函数中间有一个模糊的输入错误，因为几个小时可能会非常令人沮丧；））。这是一个例子：

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

关于签名的更多信息，请阅读the documentation。

找不到用 numba 实现 3d 维数组的工作方法

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

python

arrays

numba