res_reg_lmnt_awikner.ks_etdrk4

 1import numpy as np
 2from numpy.fft import fft, ifft
 3
 4
 5from numba import jit, prange, objmode
 6
 7@jit(nopython = True, fastmath = True, parallel = True)
 8def mean_numba_axis1(mat):
 9
10    res = np.zeros(mat.shape[0])
11    for i in prange(mat.shape[0]):
12        res[i] = np.mean(mat[i])
13
14    return res
15
16@jit(nopython = True, fastmath = True)
17def precompute_KS_params(N, d, tau, M = 16, const = 0):
18    k = np.concatenate((np.arange(int(N/2)), np.arange(-int(N/2), 0)))*2*np.pi/d
19    L = (1+const)*k**2.0 - k**4.0
20    E = np.exp(tau*L)
21    E2 = np.exp(tau/2*L)
22    r = np.exp(1j * np.pi * (np.arange(1, M+1)-0.5)/M)
23    LR = tau*(np.zeros((1,M)) + L.reshape(-1,1)) + (np.zeros((N,1)) + r)
24    Q  = tau*mean_numba_axis1(np.real((np.exp(LR/2)-1)/LR))
25    f1 = tau*mean_numba_axis1(np.real((-4-LR+np.exp(LR)*(4-3*LR+LR**2.0))/(LR**3.0)))
26    f2 = tau*mean_numba_axis1(np.real((2+LR+np.exp(LR)*(-2+LR))/(LR**3.0)))
27    f3 = tau*mean_numba_axis1(np.real((-4-3*LR-LR**2.0+np.exp(LR)*(4-LR))/(LR**3.0)))
28    g  = -0.5*1j*k
29    params = np.zeros((7,N), dtype = np.complex128)
30    params[0] = E
31    params[1] = E2
32    params[2] = Q
33    params[3] = f1
34    params[4] = f2
35    params[5] = f3
36    params[6] = g
37
38    return params
39
40
41@jit(nopython = True, fastmath = True)
42def kursiv_forecast(u, params):
43
44    with objmode(unext = 'double[:]'):
45        v  = fft(u)
46        Nv = params[6]*fft(np.real(ifft(v))**2.0)
47        a  = params[1]*v + params[2]*Nv
48        Na = params[6]*fft(np.real(ifft(a))**2.0)
49        b  = params[1]*v + params[2]*Na
50        Nb = params[6]*fft(np.real(ifft(b))**2.0)
51        c  = params[1]*a + params[2]*(2*Nb - Nv)
52        Nc = params[6]*fft(np.real(ifft(c))**2.0)
53        vnext  = params[0]*v + Nv*params[3] + 2*(Na+Nb)*params[4] + Nc*params[5]
54        unext = np.real(ifft(vnext))
55    return unext
56
57@jit(nopython = True, fastmath = True)
58def kursiv_forecast_pred(u, params):
59    u = np.ascontiguousarray(u.T)
60    with objmode(unext = 'double[:,:]'):
61        v  = fft(u,axis = 1)
62        Nv = params[6]*fft(np.real(ifft(v, axis = 1))**2.0, axis = 1)
63        a  = params[1]*v + params[2]*Nv
64        Na = params[6]*fft(np.real(ifft(a, axis = 1))**2.0, axis = 1)
65        b  = params[1]*v + params[2]*Na
66        Nb = params[6]*fft(np.real(ifft(b, axis = 1))**2.0, axis = 1)
67        c  = params[1]*a + params[2]*(2*Nb - Nv)
68        Nc = params[6]*fft(np.real(ifft(c, axis = 1))**2.0, axis = 1)
69        v  = params[0]*v + Nv*params[3] + 2*(Na+Nb)*params[4] + Nc*params[5]
70        unext = np.real(ifft(v, axis = 1))
71    return unext.T
72
73
74@jit(nopython = True, fastmath = True)
75def kursiv_predict(u0, tau = 0.25, N = 64, d = 22, T = 100, params = np.array([[],[]], dtype = np.complex128), int_steps = 1):
76    if params.size == 0:
77        new_params = precompute_KS_params(N, d, tau)
78    else:
79        new_params = params
80    steps = T*int_steps
81
82
83    u_arr = np.zeros((N, steps+int_steps))
84    u_arr[:,0] = u0
85    for i in range(steps+int_steps-1):
86        u_arr[:,i+1] = kursiv_forecast(u_arr[:,i], new_params)
87    return np.ascontiguousarray(u_arr[:,::int_steps]), new_params
88
89@jit(nopython = True, fastmath = True)
90def kursiv_predict_pred(u0_array, tau = 0.25, N = 64, d = 22, T = 100, params = np.array([[],[]], dtype = np.complex128)):
91    if params.size == 0:
92        new_params = precompute_KS_params(N, d, tau)
93    else:
94        new_params = params
95    steps = T
96
97    u_arr = kursiv_forecast_pred(u0_array, new_params)
98    return u_arr, new_params