Как решить ошибку опускания Numba?

Question

У меня есть функция, которую я пытаюсь ускорить, используя декоратор @jit из модуля Numba. Для меня важно максимально ускорить это, потому что мой основной код вызывает эту функцию миллионы раз. Вот моя функция:

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 соответствует Для цикла с J в качестве итератора.

Можете ли вы мне помочь?

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 примерно на 100x. Не уверен, как это относится к джитту.

Похоже, что вы могли бы, например, принять ваше dEdt из петли, и вычислить его в одном шаге что-то вроде:

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)