res_reg_lmnt_awikner.lorenzrungekutta_numba 

 1import numpy as np
 2from numba import jit
 3@jit(nopython = True, fastmath = True)
 4def dxdt_lorenz(x, sigma = 10., beta = 8/3, rho = 28.):
 5    # Evaluates derivative of Lorenz '63 system with a time-dependent
 6    # rho value. For constant rho, input the r_t_const function.
 7    return np.array([sigma*(- x[0] + x[1]),\
 8                     rho*x[0] - x[1] - x[0]*x[2],\
 9                     x[0]*x[1]-beta*x[2]])
10
11@jit(nopython = True, fastmath = True)
12def dxdt_lorenz_array(x, sigma = 10., beta = 8/3, rho = 28.):
13    # Evaluates derivative of Lorenz '63 system with a time-dependent
14    # rho value. For constant rho, input the r_t_const function.
15    return np.stack((sigma*(- x[0] + x[1]),\
16                      rho*x[0] - x[1] - x[0]*x[2],\
17                      x[0]*x[1]-beta*x[2]))
18
19@jit(nopython = True, fastmath = True)
20def rk4(x, tau, dxdt):
21    # Fourth order Runge-Kutta integrator
22
23    k1 = dxdt(x)
24    k2 = dxdt(x + k1/2*tau)
25    k3 = dxdt(x + k2/2*tau)
26    k4 = dxdt(x + tau*k3)
27
28    xnext = x + 1/6*tau*(k1+2*k2+2*k3+k4)
29    return xnext
30
31@jit(nopython = True, fastmath = True)
32def lorenzrungekutta(u0 = np.array([1,1,30]), T = 100, tau = 0.1, int_step = 10):
33    steps = T*int_step
34    output = np.zeros((3,steps+int_step))
35    output[:,0] = u0
36
37    #loops from t = 0 to T
38    for i in range(0, steps+int_step):
39        output[:,i+1] = rk4(output[:,i],tau/int_step,dxdt_lorenz)
40
41
42    return output[:,::int_step]
43
44@jit(nopython = True, fastmath = True)
45def lorenzrungekutta_pred(u0_array, tau = 0.1, int_step = 10):
46
47    #loops from t = 0 to T
48    for i in range(0, int_step):
49        u0_array = rk4(u0_array,tau/int_step,dxdt_lorenz_array)
50
51    output = u0_array
52    return output
@jit(nopython=True, fastmath=True)
def dxdt_lorenz(x, sigma=10.0, beta=2.6666666666666665, rho=28.0): 
 4@jit(nopython = True, fastmath = True)
 5def dxdt_lorenz(x, sigma = 10., beta = 8/3, rho = 28.):
 6    # Evaluates derivative of Lorenz '63 system with a time-dependent
 7    # rho value. For constant rho, input the r_t_const function.
 8    return np.array([sigma*(- x[0] + x[1]),\
 9                     rho*x[0] - x[1] - x[0]*x[2],\
10                     x[0]*x[1]-beta*x[2]])
@jit(nopython=True, fastmath=True)
def dxdt_lorenz_array(x, sigma=10.0, beta=2.6666666666666665, rho=28.0): 
12@jit(nopython = True, fastmath = True)
13def dxdt_lorenz_array(x, sigma = 10., beta = 8/3, rho = 28.):
14    # Evaluates derivative of Lorenz '63 system with a time-dependent
15    # rho value. For constant rho, input the r_t_const function.
16    return np.stack((sigma*(- x[0] + x[1]),\
17                      rho*x[0] - x[1] - x[0]*x[2],\
18                      x[0]*x[1]-beta*x[2]))
@jit(nopython=True, fastmath=True)
def rk4(x, tau, dxdt): 
20@jit(nopython = True, fastmath = True)
21def rk4(x, tau, dxdt):
22    # Fourth order Runge-Kutta integrator
23
24    k1 = dxdt(x)
25    k2 = dxdt(x + k1/2*tau)
26    k3 = dxdt(x + k2/2*tau)
27    k4 = dxdt(x + tau*k3)
28
29    xnext = x + 1/6*tau*(k1+2*k2+2*k3+k4)
30    return xnext
@jit(nopython=True, fastmath=True)
def lorenzrungekutta(u0=array([ 1, 1, 30]), T=100, tau=0.1, int_step=10): 
32@jit(nopython = True, fastmath = True)
33def lorenzrungekutta(u0 = np.array([1,1,30]), T = 100, tau = 0.1, int_step = 10):
34    steps = T*int_step
35    output = np.zeros((3,steps+int_step))
36    output[:,0] = u0
37
38    #loops from t = 0 to T
39    for i in range(0, steps+int_step):
40        output[:,i+1] = rk4(output[:,i],tau/int_step,dxdt_lorenz)
41
42
43    return output[:,::int_step]
@jit(nopython=True, fastmath=True)
def lorenzrungekutta_pred(u0_array, tau=0.1, int_step=10): 
45@jit(nopython = True, fastmath = True)
46def lorenzrungekutta_pred(u0_array, tau = 0.1, int_step = 10):
47
48    #loops from t = 0 to T
49    for i in range(0, int_step):
50        u0_array = rk4(u0_array,tau/int_step,dxdt_lorenz_array)
51
52    output = u0_array
53    return output