Numbaを解決するにはどうすればよいですか？

Question

私はNumbaモジュールから@jitデコレータを使ってスピードアップしようとしている関数を持っています。私の主なコードは何百万回もこの機能を要求しているので、できるだけ高速化することが不可欠です。ここに私の関数はあります：

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で実行すると、次のエラーが発生します。

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行目はイテレータとしてのForループに対応しています。

お手伝いできますか？

EDIT： ピーターの提案に基づいており、Einsumとそれを組み合わせることで、私はループを削除することができました。これは私の機能を倍速くしました。新しいコードは次のとおりです。

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 

これをスピードアップするための技術をもっと提案できますか？

Answer 1

これはベクトル化できない理由がコードからわかりません。ベクトル化を行うと、この種のPythonコードの速度を約100倍向上させることができます。それがジットに比べてどうなるかはわかりません。あなたは、例えば、ループの外にあなたのDEDTを取る、とのようなものを1つのステップでそれを計算することができように

に見えます：私は本当に知らないけど（

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))) 

何の次元あなたのSweet.omegaは）。

Answer 2

他にも問題があるかもしれませんが、モジュール名前空間の配列を参照することは現在サポートされていないようです（下の簡単な説明）。名前としてomegaをインポートしてみてください。

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)