What's wrong with numba compilation with using numpy?

Question

I cannot compile this code:

import numpy as np
import numba
from numba import jit, float64, complex128
import math

@jit(complex128[:](float64,float64[:],float64))
def GaborWavelet(omega, t, Gabor_coef):
    
    c1 = 0.3251520240633*math.sqrt(omega)
    c2 = -0.5*Gabor_coef
    c3 = omega*0.187390625129278
    
    res = np.array(c2*(t * c3)**2, dtype = np.complex128)
    
    res.imag = omega*t
    
    return c1*np.exp(res)

It raises:

Compilation is falling back to object mode WITH looplifting enabled because Function "GaborWavelet" failed type inference due to: No implementation of function Function() found for signature:

array(array(float64, 1d, C), dtype=class(complex128))

There are 2 candidate implementations: - Of which 2 did not match due to: Overload in function 'array': File: numba\core\typing\npydecl.py: Line 504. With argument(s): '(array(float64, 1d, C), dtype=class(complex128))': Rejected as the implementation raised a specific error: TypingError: array(float64, 1d, C) not allowed in a homogeneous sequence

res = np.array(c2*(t * c3)**2, dtype = np.complex128)
^

What do I wrong?

How to compile this code (with numpy methods inside)?

Niklas Mertsch Niklas Mertsch · Accepted Answer · 2020-09-11T10:48:56

Numba does not support two of the things you use, but equivalent options:

Type conversion via np.array(arr, dtype=type). Instead use arr.astype(type).
Setting arr.imag=values for complex datatypes. Instead use arr += values*1j.

The following code works on my machine and should yield equivalent results:

import numpy as np
import numba
from numba import jit, float64, complex128
import math

@jit(complex128[:](float64,float64[:],float64))
def GaborWavelet(omega, t, Gabor_coef):

    c1 = 0.3251520240633*math.sqrt(omega)
    c2 = -0.5*Gabor_coef
    c3 = omega*0.187390625129278

    res = (c2*(t * c3)**2).astype(np.complex128)
    res += omega*t*1j

    return c1*np.exp(res)

What's wrong with numba compilation with using numpy?

1 Answers