Сделать Рунге Кутта Быстрее

Question

Попытка ускорить расчет лагранжевых траекторий с учетом двумерного флюидного поля.

Пока что, используя 4-й порядок РК, имеем:

    for n in range(tau):

        t=2*n
        [x1,z1]=[x+0.5*U_xz[t](x0,z0),z+0.5*W_xz[t](x0,z0)]

        t+=1
        [x2,z2]=[x+0.5*U_xz[t](x1,z1),z+0.5*W_xz[t](x1,z1)]
        [x3,z3]=[x+U_xz[t](x2,z2),z+W_xz[t](x2,z2)]

        t+=1
        [x4,z4]=[x+U_xz[t](x3,z3),z+W_xz[t](x3,z3)]
        [x0,z0]=[(x1+2*x2+2*x3+x4)/6.0,(z1+2*z2+2*z3+z4)/6.0]

        x=x%nx

Последние несколько строк не так важны. U_xz и W_xz являются функциями (линейно интерполируются из данных с дискретными временными шагами).

Хотели бы немного ускорить процесс, какие-нибудь предложения, может быть, @jit?