numpy.ndarray.astype not working in numba.jit

Question

Hello fellow programmers!

I am trying to free a two dimensional array from complex numbers in a function decorated with numba.jit:

Function without numba.jit (works):

import numpy as np
import numba

def main():
    KAM = np.array([[ -106.66666667-0.j,    42.66666667+0.j,    42.66666667+0.j,     21.33333333+0.j,     0.        +0.j],
                   [   42.66666667+0.j,   -42.66666667-0.j,     0.        +0.j,      0.        +0.j,     0.        +0.j],
                   [   42.66666667+0.j,     0.        +0.j,   -42.66666667-0.j,      0.        +0.j,     0.        +0.j],
                   [   21.33333333+0.j,     0.        +0.j,     0.        +0.j,  -1088.        -0.j,  1066.66666667+0.j],
                   [    0.        +0.j,     0.        +0.j,     0.        +0.j,   1066.66666667+0.j, -1066.66666667-0.j]])

    KAM = KAM.astype(float)
    print(KAM)

if __name__ == '__main__':
    main()

Output:

C:\Users\Artur\Anaconda\python.exe C:/Users/Artur/Desktop/RL_framework/help_functions/test2.py
C:/Users/Artur/Desktop/RL_framework/help_functions/test2.py:11: ComplexWarning: Casting complex values to real discards the imaginary part
  KAM = KAM.astype(float)
[[ -106.66666667    42.66666667    42.66666667    21.33333333
      0.        ]
 [   42.66666667   -42.66666667     0.             0.
      0.        ]
 [   42.66666667     0.           -42.66666667     0.
      0.        ]
 [   21.33333333     0.             0.         -1088.
   1066.66666667]
 [    0.             0.             0.          1066.66666667
  -1066.66666667]]

Process finished with exit code 0

Function decorated with numba.jit:

import numpy as np
import numba

@numba.jit(nopython=True)
def main():
    KAM = np.array([[ -106.66666667-0.j,    42.66666667+0.j,    42.66666667+0.j,     21.33333333+0.j,     0.        +0.j],
                   [   42.66666667+0.j,   -42.66666667-0.j,     0.        +0.j,      0.        +0.j,     0.        +0.j],
                   [   42.66666667+0.j,     0.        +0.j,   -42.66666667-0.j,      0.        +0.j,     0.        +0.j],
                   [   21.33333333+0.j,     0.        +0.j,     0.        +0.j,  -1088.        -0.j,  1066.66666667+0.j],
                   [    0.        +0.j,     0.        +0.j,     0.        +0.j,   1066.66666667+0.j, -1066.66666667-0.j]])

    KAM = KAM.astype(float)
    print(KAM)

if __name__ == '__main__':
    main()

Output:

C:\Users\Artur\Anaconda\python.exe C:/Users/Artur/Desktop/RL_framework/help_functions/test2.py
Traceback (most recent call last):
  File "C:/Users/Artur/Desktop/RL_framework/help_functions/test2.py", line 16, in <module>
    main()
  File "C:\Users\Artur\Anaconda\lib\site-packages\numba\core\dispatcher.py", line 401, in _compile_for_args
    error_rewrite(e, 'typing')
  File "C:\Users\Artur\Anaconda\lib\site-packages\numba\core\dispatcher.py", line 344, in error_rewrite
    reraise(type(e), e, None)
  File "C:\Users\Artur\Anaconda\lib\site-packages\numba\core\utils.py", line 80, in reraise
    raise value.with_traceback(tb)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Invalid use of BoundFunction(array.astype for array(complex128, 2d, C)) with parameters (Function(<class 'float'>))
 * parameterized
[1] During: resolving callee type: BoundFunction(array.astype for array(complex128, 2d, C))
[2] During: typing of call at C:/Users/Artur/Desktop/RL_framework/help_functions/test2.py (12)


File "test2.py", line 12:
def main():
    <source elided>

    KAM = KAM.astype(float)
    ^


Process finished with exit code 1

This problem seems to not come from the complex values although they are mentioned in the error message. Running the same code without complex values in KAM results to the same error:

import numpy as np
import numba

@numba.jit(nopython=True)
def main():
    KAM = np.array([[ -106.66666667,    42.66666667,    42.66666667,     21.33333333,     0.        ],
                   [   42.66666667,   -42.66666667,     0.        ,      0.        ,     0.        ],
                   [   42.66666667,     0.        ,   -42.66666667,      0.        ,     0.        ],
                   [   21.33333333,     0.        ,     0.        ,  -1088.        ,  1066.66666667],
                   [    0.        ,     0.        ,     0.        ,   1066.66666667, -1066.66666667]])

    KAM = KAM.astype(float)
    print(KAM)

if __name__ == '__main__':
    main()

Output:

C:\Users\Artur\Anaconda\python.exe C:/Users/Artur/Desktop/RL_framework/help_functions/test2.py
Traceback (most recent call last):
  File "C:/Users/Artur/Desktop/RL_framework/help_functions/test2.py", line 22, in <module>
    main()
  File "C:\Users\Artur\Anaconda\lib\site-packages\numba\core\dispatcher.py", line 401, in _compile_for_args
    error_rewrite(e, 'typing')
  File "C:\Users\Artur\Anaconda\lib\site-packages\numba\core\dispatcher.py", line 344, in error_rewrite
    reraise(type(e), e, None)
  File "C:\Users\Artur\Anaconda\lib\site-packages\numba\core\utils.py", line 80, in reraise
    raise value.with_traceback(tb)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Invalid use of BoundFunction(array.astype for array(float64, 2d, C)) with parameters (Function(<class 'float'>))
 * parameterized
[1] During: resolving callee type: BoundFunction(array.astype for array(float64, 2d, C))
[2] During: typing of call at C:/Users/Artur/Desktop/RL_framework/help_functions/test2.py (18)


File "test2.py", line 18:
def main():
    <source elided>

    KAM = KAM.astype(float)
    ^


Process finished with exit code 1

Answer 1

Don't use .astype(float) to convert complex numbers to real ones.

Use np.real(arr) to get the real parts and if desired np.imag(arr) to get the imaginary parts. If you want the magnitudes of the complex numbers use np.abs(arr). All work with numba.