How to Solve Numba Lowering error?

问题

I have a function, which I am trying to speed up using the @jit decorator from Numba module. For me it is essential to speed this up as much as possible, because my main code calls upon this function for millions of times. Here is my function:

from numba import jit, types
import Sweep    #My own module, works fine


@jit(types.Tuple((types.complex128[:], types.float64[:]))(types.complex128[:], types.complex128[:], types.float64[:], types.float64[:], types.float64))
def MultiModeSL(Ef, Ef2, Nf, u, tijd ):
  dEdt= np.zeros(nrModes, dtype=np.complex128)
  dNdt0= np.zeros(nrMoments, dtype=np.complex128)
  Efcon = np.conjugate(Ef)

  for j in range(nrModes):
    for n in range(nrMoments):
      dEdt +=  0.5 * CMx[:,j,n,0] * dg * (1+ A*1j) * Nf[n] * Ef[j] * np.exp( 1j* (Sweep.omega[j]-Sweep.omega) *tijd)
      for k in range(nrModes):
        if n==0:
          dNdt0 += g* CMx[j, k, 0,:] * Efcon[j] * Ef[k] * np.exp( 1j* (Sweep.omega[k]-Sweep.omega[j]) *tijd) 
        dNdt0 += dg*(1+A*1j) * CMx[j,k,n,:] * Nf[n] * Efcon[j] * Ef[k] * np.exp( 1j* (Sweep.omega[k]-Sweep.omega[j]) *tijd) 

  dEdt += - 0.5*(pd-g)*Ef +     fbr*Ef2  + Kinj*EAinj*(1 + np.exp(1j*(u+Vmzm)) ) 
  dNdt = Sweep.Jn - Nf*ed - dNdt0.real
  return dEdt, dNdt

The function works perfectly well, without the Jit decorator. However, when I run it with the @jit, I get this error:

numba.errors.LoweringError: Failed at object (object mode frontend)
Failed at object (object mode backend)
dEdt.1
File "Functions.py", line 82
[1] During: lowering "$237 = call $236(Ef, Ef2, Efcon, Nf, dEdt.1, dNdt0, tijd, u)" at /home/humblebee/MEGA/GUI RC/General_Formula/Functions.py (82)

Line 82 corresponds to the For loop with j as iterator.

Can you help me out?

EDIT: Based on Peter's suggestion and combining it with Einsum, I was able to remove the loops. This made my function 3 times faster. Here is the new code:

def MultiModeSL(Ef, Ef2, Nf, u, tijd ):
  dEdt= np.zeros(nrModes, dtype=np.complex128)
  dNdt0= np.zeros(nrMoments, dtype=np.complex128)
  Efcon = np.conjugate(Ef)
  dEdt = 0.5*  np.einsum("k, jkm, mk, kj -> j",  dg*(1+A*1j), CMx[:, :, :, 0],  (Ef[:] * Nf[:, None] ),  np.exp( 1j* (OMEGA[:, None]-OMEGA) *tijd))
  dEdt += - 0.5*(pd-g)*Ef + fbr*Ef2  + Kinj*EAinj*(1 + np.exp(1j*(u+Vmzm)) )

  dNdt = - np.einsum("j, jkm, jk, kj ", g, CMx[:,:,:,0], (Ef*Efcon[:,None]),  np.exp( 1j* (OMEGA[:, None]-OMEGA) *tijd))
  dNdt += -np.einsum("j, j, jknm, kjm, kj",dg, (1+A*1j), CMx, (Nf[:]*Efcon[:,None]*Ef[:,None,None]), np.exp( 1j* (OMEGA[:, None]-OMEGA) *tijd)  )
  dNdt += JN - Nf*ed
  return dNdt

Can you suggest more techniques to speed this up?

回答1:

I can't see from your code why this isn't vectorizable. Vectorizing can speed up this kind of Python code by around 100x. Not sure of how it does relative to jit.

It looks like you could, for instance, take your dEdt out of the loop, and compute it in one step with something like :

dEdt = 0.5 * (Cmx[:, :, :, 0] * dg * (1+A*1j) * Nf[:] * Ef[:, None] * np.exp( 1j* (Sweep.omega[None, :, None, None]-Sweep.omega) *tijd)).sum(axis=2).sum(axis=1) - 0.5*(pd-g)*Ef + fbr*Ef2  + Kinj*EAinj*(1 + np.exp(1j*(u+Vmzm)) )

(Though I don't really know what the dimensionality of your Sweet.omega is).

回答2:

There may be other issues, but one is that referencing an array in a module namespace seems to currently be unsupported (simple repro below). Try importing omega as a name.

In [14]: %%file Sweep.py
    ...: import numpy as np
    ...: constant_val = 0.5
    ...: constant_arr = np.array([0, 1.5, 2.])
Overwriting Sweep.py

In [15]: Sweep.constant_val
Out[15]: 0.5

In [16]: Sweep.constant_arr
Out[16]: array([ 0. ,  1.5,  2. ])

In [17]: @njit
    ...: def f(value):
    ...:     return value + Sweep.constant_val
    ...: 

In [18]: f(100)
Out[18]: 100.5

In [19]: @njit
    ...: def f(value):
    ...:     return value + Sweep.constant_arr[0]

In [20]: f(100)
LoweringError: Failed at nopython (nopython mode backend)
'NoneType' object has no attribute 'module'
File "<ipython-input-19-0a259ade6b9e>", line 3
[1] During: lowering "$0.3 = getattr(value=$0.2, attr=constant_arr)" at <ipython-input-19-0a259ade6b9e> (3)

来源：https://stackoverflow.com/questions/39379556/how-to-solve-numba-lowering-error

标签

python

numpy

numba