использование Numba для ускорения - PullRequest
Новая
       111

использование Numba для ускорения

2 голосов
/ 21 февраля 2020

Я использую метод линий с 'solve_ivp` для решения нелинейного PDE : 

@njit(fastmath=True,error_model="numpy",cache=True)
def thinFilmEq(t,h,dx,Ma,phiFun,tempFun):
    phi = phiFun(h)
    temperature = tempFun(h)
    hxx = (np.roll(h,1) - 2*h + np.roll(h,-1))/dx**2  # use np.roll as I'm implementing periodic BC
    p = phi - hxx
    px = (np.roll(p,-1) - np.roll(p,1))/(2*dx)
    Tx = (np.roll(temperature,-1) - np.roll(temperature,1))/(2*dx)

    flux = h**3*px/3 + Ma*h**2*Tx/2
    dhdt = (np.roll(flux,-1) - np.roll(flux,1))/(2*dx)

    return dhdt

Я получаю следующую ошибку: TypingError: non-precise type pyobject [1] During: typing of argument at C:/Users/yhcha/method_of_lines/test_01_thinFilmEq.py (28) Я подозреваю, что это связано с phiFun и tempFun. Это функции, которые я предоставляю во время вызова. Я делаю функции аргументами функции dhdt только для того, чтобы сделать вещи более общими. Когда я пытаюсь удалить phiFun и tempFun и явно дать форму функции внутри thinFilmEq, ошибка исчезает.

Затем я вижу следующую ошибку TypingError: Use of unsupported NumPy function 'numpy.roll' or unsupported use of the function. Я подумал, что возможно np.roll не поддерживается, хотя он включен в официальный веб-сайт . Я пытался «увеличить» массив, чтобы как-то вручную применить то же самое, что и np.roll, когда имел дело с конечной разницей для periodi c B C: 

    def augment(x):
        x2 = np.empty(len(x)+2)
        x2[1:-1] = x
        x2[0] = x[-1]
        x2[-1] = x[0]
        return x2

    H = augment(x)
    hx = (H[2:]-[H:-2])/dx   # use this instead of hx=(roll(h,-1)-roll(h,1))/dx

Мои вопросы:

  1. Кажется, что я могу заставить numba работать за счет того, чтобы сделать коды менее общими (не могу предоставить произвольную функцию, такую ​​как phiFun и элегантную (например, не могу использовать одно- лайнер с np.roll). Есть ли способы обойти это или это просто цена, которую я должен заплатить при использовании numba для «компиляции» кода?

  2. Оригинал версия без numba примерно в 10 раз медленнее, чем версия Matlab, которую я кодировал, а версия numba все еще примерно в 3-4 раза медленнее, чем Matlab. Я не ожидаю, что scipy превзойдет Matlab, но есть другие способы ускорить код для преодоления разрыва?

...