skip to Main Content

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

2

Answers


  1. Chosen as BEST ANSWER

    For sure the best and most complete answer was from @Rafnus, but to report another simple work around to someone that face the problem in the future is just addind the null imaginary part to other elements of the array:

    from numba import jit
    from numpy import array, vectorize
    
    @jit(nopython=True)
    def func(x):
        a = 0+1j
        v = array([[1+0j,a],
                   [2+0j,3+0j]])
        return x*v
    func_vec = vectorize(func)
    
    print(func_vec(10))
    

  2. 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

    Login or Signup to reply.
Please signup or login to give your own answer.
Back To Top
Search