如何解决Numba降低错误? [英] How to Solve Numba Lowering error?

问题描述

我有一个函数，我正在尝试使用 Numba 模块中的 @jit 装饰器来加速它.对我来说，尽可能加快速度是必不可少的，因为我的主代码调用了这个函数数百万次.这是我的功能:

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

该函数运行良好，无需 Jit 装饰器.但是，当我使用@jit 运行它时，出现此错误:

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)

第 82 行对应于以 j 作为迭代器的 For 循环.

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

你能帮我吗?

根据彼得的建议并将其与 Einsum 相结合，我能够删除循环.这使我的函数速度 3 倍.这是新代码:

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?

推荐答案

可能还有其他问题，但一个是目前似乎不支持在模块命名空间中引用数组(下面的简单重现).尝试导入 omega 作为名称.

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)

这篇关于如何解决Numba降低错误?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！