python 具有复数全局数组的 numba

Question

我正在尝试使用 numba 优化我的代码。

我将代码设计为包含一个 gl.py 文件，其中包含一些数组，这些数组将由 main.py 和 main() 中从 main.py 调用的函数。

auxiliary.py 看起来像：

import numpy as np
from numba import jit, types

from cmath import sqrt, exp, sin


N_timesteps_imag = 100

N_epsilon_divs = 60
N_z_divs = 2000
K = N_z_divs # equal to the variable N_z_divs

delta_epsilon = 0.1
delta_z = 0.1

lambd = 1.5

z_max = (N_z_divs/2) * delta_z
epsilon_range = np.linspace(0.0, N_epsilon_divs*delta_epsilon, N_epsilon_divs+1)
z_range = np.linspace(-z_max, z_max, N_z_divs+1)

psi_ground = np.zeros((N_z_divs+1, N_epsilon_divs+1, N_timesteps_imag+1), dtype=types.complex128)

@jit(nopython=True)
def pop_psiground_t0():
    for c1 in range(1, psi_ground.shape[0]-1):
        for c2 in range(1, psi_ground.shape[1]-1):
            zed = (c1 - N_z_divs/2) * delta_z 
            epsi = c2 * delta_epsilon
            psi_ground[c1, c2, 0] = sqrt(3) * epsi * exp(-sqrt(epsi**(2*lambd) + zed**2))

pop_psiground_t0()

main.py 看起来像 (MWE):

import numpy as np
import auxiliary

def main():
    print(auxiliary.psi_ground[1000, 40, 0]) # shall NOT be 0 + 0j !!!


if __name__ == '__main__':
    main()

不管我在 auxiliary.py 中为 psi_ground 声明的关键字参数 dtype 放置了什么，无论是 numba.types.complex128、np.complex128、np.clongdouble，没有任何效果。特别是，对于 np.complex128，当运行 python3 main.py:

时出现以下错误

No implementation of function Function(<built-in function setitem>) found for signature:
 
 >>> setitem(readonly array(complex128, 3d, C), Tuple(int64, int64, Literal[int](0)), complex128)
 
There are 16 candidate implementations:
  - Of which 14 did not match due to:
  Overload of function 'setitem': File: <numerous>: Line N/A.
    With argument(s): '(readonly array(complex128, 3d, C), UniTuple(int64 x 3), complex128)':
   No match.
  - Of which 2 did not match due to:
  Overload in function 'SetItemBuffer.generic': File: numba/core/typing/arraydecl.py: Line 171.
    With argument(s): '(readonly array(complex128, 3d, C), UniTuple(int64 x 3), complex128)':
   Rejected as the implementation raised a specific error:
     TypeError: Cannot modify value of type readonly array(complex128, 3d, C)
  raised from /home/velenos14/.local/lib/python3.8/site-packages/numba/core/typing/arraydecl.py:177

During: typing of setitem at /mnt/c/Users/iusti/Desktop/test_python/auxiliary.py (45)

File "auxiliary.py", line 45:
def pop_psiground_t0():
    <source elided>
            epsi = c2 * delta_epsilon
            psi_ground[c1, c2, 0] = sqrt(3) * epsi * exp(-sqrt(epsi**(2*lambd) + zed**2))

我该如何处理？我试着按照这里写的：numba TypingError with complex numpy array and native data types

是的，我需要 psi_ground 数组为复杂类型，具有很高的精度，即使最初它由实数填充也是如此。稍后在 main() 中将由复数重新填充。谢谢！

Answer 1

psi_ground = np.zeros((N_z_divs+1, N_epsilon_divs+1, N_timesteps_imag+1), dtype=types.complex128)

必须在 numba 函数中定义。该错误清楚地表明 numba 无法更改 psi_ground 数组的值。

下面是修改后的代码

import numpy as np
from numba import jit, types

from cmath import sqrt, exp, sin


N_timesteps_imag = 100

N_epsilon_divs = 60
N_z_divs = 2000
K = N_z_divs # equal to the variable N_z_divs

delta_epsilon = 0.1
delta_z = 0.1

lambd = 1.5

z_max = (N_z_divs/2) * delta_z
epsilon_range = np.linspace(0.0, N_epsilon_divs*delta_epsilon, N_epsilon_divs+1)
z_range = np.linspace(-z_max, z_max, N_z_divs+1)


@jit(nopython=True)
def pop_psiground_t0():
    psi_ground = np.zeros((N_z_divs+1, N_epsilon_divs+1, N_timesteps_imag+1), dtype=types.complex128)
    for c1 in range(1, psi_ground.shape[0]-1):
        for c2 in range(1, psi_ground.shape[1]-1):
            zed = (c1 - N_z_divs/2) * delta_z 
            epsi = c2 * delta_epsilon
            psi_ground[c1, c2, 0] = sqrt(3) * epsi * exp(-sqrt(epsi**(2*lambd) + zed**2))

pop_psiground_t0()

python 具有复数全局数组的 numba

python numba with complex numbers global array

python

arrays

compiler-errors

numpy

numba