python numba with complex numbers global array

Question

I am trying to optimize my code with numba.

I designed the code to contain a gl.py file which contains some arrays which will be used by main.py and functions called inside main() from main.py.

The auxiliary.py looks like:

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

The main.py looks like (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()

Irrespective to what I put for the keyword argument dtype for the declaration of psi_ground inside auxiliary.py, be it numba.types.complex128, np.complex128, np.clongdouble, nothing works. In particular, for np.complex128, I get the following error when running 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))

How can I proceed with this? I tried to follow what it's written here: numba TypingError with complex numpy array and native data types

And yes, I need that psi_ground array to be of complex type, with a lot of precision, even if initially it's populated by real numbers. Later in the main() will get re-populated by complex numbers. Thank you!

Answer 1

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

must be defined inside numba function. The error clearly states that numba is unable to change the values of the psi_ground array.

Below is the modified code

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