Numba "LoweringError" for complex numbers in numpy array-CodePudding

I have to make a calculation using complex arrays, however when using numba to speed up the process I get an error numba.core.errors.LoweringError: Failed in nopython mode pipeline (step: nopython mode backend). Here it is a simplified version of my code:

import numpy as np
from numba import jit
from numpy import array

@jit(nopython=True)
def func(x):
   a = 1j
   v = x*array([[1.,a],
                [2.,3.]])
   return v
func_vec = np.vectorize(func)

print(func_vec(10.))

It is important to note that if a is real, everything works well. I have already tested a dtype=np.complex128 for v, but the issue remain.

Numba version: 0.51.0

Numpy version: 1.22.3

Python version: 3.8.10

System: Ubuntu 20.4

CodePudding user response：

The reason your code crashes is because of a limitation with numba. Numba seemingly does not like it when you use the numpy array function to create an (multidimensional!) array filled with complex values. You can work around the issue as follows:

import numpy as np
import numba as nb

inp_arr = np.array([[1, 0],
                    [2, 3j]], dtype =np.complex128)

@nb.njit(nb.complex128[:, ::1](nb.complex128, nb.complex128[:, ::1]))
def func(x, arr):
    v = x*inp_arr
    return v

func_vec = np.vectorize(func)

print(func_vec(10))

Alternatively, you can just simply assume the array is compile time constant:

import numpy as np
import numba as nb

def enclosingFunc(x):
    inp_arr = np.array([[1, 0],
                        [2, 3j]], dtype =np.complex128)

    @nb.njit(nb.complex128[:, ::1](nb.complex128))
    def func(x):
        v = x*inp_arr
        return v

    return func(x)

enclosingFunc = np.vectorize(enclosingFunc)
print(enclosingFunc(10))

Note that the numpy vectorize function is not made for good perfomance, according to the numpy documentation. You could look into using the numba vectorize decorater instead, but that is outside the scope of your question.

Interestingly, giving the numba decorater the type hint locals={'arr':nb.complex128[:,::1]} did not work as an alternative answer