def Model(dt=0.25,DL=32,Nx=128):
"Define step(), x0, etc. Alternative schemes (step_XXX) also implemented."
h = dt # alias -- prevents o/w in step()
# Fourier stuff
kk = np.append(arange(0,Nx/2),0)*2/DL # wave nums for rfft
# kk = ccat([arange(0,Nx/2),[0], arange(-Nx/2+1,0)])*2/DL # wave nums for fft
# kk = np.fft.fftfreq(Nx, DL/Nx/2) # (altern. method)
# Operators
D = 1j*kk # Differentiation to compute: F[ u_x ]
L = kk**2 - kk**4 # Linear operator for K-S eqn: F[ - u_xx - u_xxxx]
# NonLinear term (-u*u_x) in Fourier domain via time domain
def NL(v):
return -0.5 * D * fft( ifft(v).real ** 2 )
# Evolution equation
@byFourier
def dxdt(v):
return NL(v) + L*v
# Jacobian of dxdt(u)
def d2x_dtdx(u):
assert is1d(u)
dL = ifft ( L * fft(np.eye(Nx)) ) . T
dNL = - ifft ( D * fft(np.diag(u)) ) . T
return dL + dNL
# dstep_dx = FD_Jac(step)
def dstep_dx(x,t,dt):
return integrate_TLM(d2x_dtdx(x),dt,method='analytic')
# Runge-Kutta -- Requries dt<1e-2:
# ------------------------------------------------
step_RK4 = with_rk4(dxdt,autonom=True) # Bad, not recommended.
step_RK1 = with_rk4(dxdt,autonom=True,order=1) # Truly terrible.
# "Semi-implicit RK3": explicit RK3 for nonlinear term,
# ------------------------------------------------
# implicit trapezoidal adjustment for linear term.
# Based on github.com/jswhit/pyks (Jeff Whitaker),
# who got it from doi.org/10.1175/MWR3214.1.
@byFourier
def step_SI_RK3(v,t,dt):
v0 = v.copy()
for n in range(3):
dt3 = h/(3-n)
v = v0 + dt3*NL(v)
v = (v + 0.5*L*dt3*v0)/(1 - 0.5*L*dt3) #
return v
# ETD-RK4 -- Integration factor (IF) technique, mixed with RK4.
# ------------------------------------------------
# Based on kursiv.m of Kassam and Trefethen, 2002,
# doi.org/10.1137/S1064827502410633.
#
# Precompute ETDRK4 scalar quantities
E = exp(h*L) # Integrating factor, eval at dt
E2 = exp(h*L/2) # Integrating factor, eval at dt/2
# Roots of unity are used to discretize a circular countour...
nRoots = 16
roots = exp( 1j * pi * (0.5+arange(nRoots))/nRoots )
# ... the associated integral then reduces to the mean,
# g(CL).mean(axis=-1) ~= g(L), whose computation is more stable.
CL = h * L[:,None] + roots # Contour for (each element of) L
# E * exact_integral of integrating factor:
Q = h * ( (exp(CL/2)-1) / CL ).mean(axis=-1).real
# RK4 coefficients (modified by Cox-Matthews):
f1 = h * ( (-4-CL+exp(CL)*(4-3*CL+CL**2)) / CL**3 ).mean(axis=-1).real
f2 = h * ( (2+CL+exp(CL)*(-2+CL)) / CL**3 ).mean(axis=-1).real
f3 = h * ( (-4-3*CL-CL**2+exp(CL)*(4-CL)) / CL**3 ).mean(axis=-1).real
#
@byFourier
def step_ETD_RK4(v,t,dt):
assert dt == h, \
"Model is instantiated with a pre-set dt, "+\
"which differs from the requested value"
N1 = NL(v); v1 = E2*v + Q*N1
N2a = NL(v1); v2a = E2*v + Q*N2a
N2b = NL(v2a); v2b = E2*v1 + Q*(2*N2b-N1)
N3 = NL(v2b); v = E *v + N1*f1 + 2*(N2a+N2b)*f2 + N3*f3
return v
# Select the "official" step method
step=step_ETD_RK4
# Generate IC as end-point of ex. from Kassam and Trefethen.
# x0_Kassam isn't convenient, coz prefer {x0 ∈ attractor} to {x0 ∈ basin}.
grid = DL*pi*linspace(0,1,Nx+1)[1:]
x0_Kassam = cos(grid/16) * (1 + sin(grid/16))
x0 = x0_Kassam.copy()
for k in range(int(150/h)):
x0 = step(x0, np.nan, h)
# Return dict
return Bunch.magic_init(
dt, DL, Nx, x0, x0_Kassam, grid, step,
step_ETD_RK4, step_SI_RK3, step_RK4, step_RK1,
dxdt, d2x_dtdx, dstep_dx)
###########################
PRMS = 'Lorenz'
if PRMS == 'Wilks':
from dapper.mods.LorenzUV.wilks05 import LUV
else:
from dapper.mods.LorenzUV.lorenz96 import LUV
nU = LUV.nU
K = 400
dt = 0.005
t0 = np.nan
dpr.set_seed(30) # 3 5 7 13 15 30
x0 = np.random.randn(LUV.M)
true_step = with_rk4(LUV.dxdt, autonom=True)
model_step = with_rk4(LUV.dxdt_trunc, autonom=True)
true_K = with_recursion(true_step, prog=1)
###########################
# Compute truth trajectory
###########################
x0 = true_K(x0, int(2 / dt), t0, dt)[-1] # BurnIn
xx = true_K(x0, K, t0, dt)
###########################
# Compute unresovled scales
###########################
gg = np.zeros((K, nU)) # "Unresolved tendency"
for k, x in enumerate(xx[:-1]):
X = x[:nU]
beta = 8.0 / 3
@ens_compatible
def dxdt(x):
"Evolution equation (coupled ODEs) specifying the dynamics."
d = np.zeros_like(x)
x, y, z = x
d[0] = sig * (y - x)
d[1] = rho * x - y - x * z
d[2] = x * y - beta * z
return d
# Time-step integration.
step = with_rk4(dxdt, autonom=True)
# Time span for plotting. Typically: ≈10 * "system time scale".
Tplot = 4.0
# Example initial state -- usually not important,
# coz system is chaotic, and we average stats after a BurnIn time.
# But it's often convenient to have a point on the attractor, or basin,
# or at least ensure "physicality".
x0 = np.array([1.509, -1.531, 25.46])
################################################
# OPTIONAL (not necessary for EnKF or PartFilt):
################################################
def d2x_dtdx(x):
# In unitless time (as used here), this means LyapExps ∈ ( −4.87, 1.66 ) .
if mod == "L96":
from dapper.mods.Lorenz96 import step
Nx = 40 # State size (flexible). Usually 36 or 40
T = 1e3 # Length of experiment (unitless time).
dt = 0.1 # Step length
# dt = 0.0083 # Any dt<0.1 yield "almost correct" Lyapunov expos.
x0 = randn(Nx) # Init condition.
eps = 0.0002 # Ens rescaling factor.
N = Nx # Num of perturbations used.
# Lyapunov exponents with F=10: [9.47 9.3 8.72 ..., -33.02 -33.61 -34.79] => n0:64
if mod == "LUV":
from dapper.mods.LorenzUV.lorenz96 import LUV
ii = np.arange(LUV.nU)
step = with_rk4(LUV.dxdt, autonom=True)
Nx = LUV.M
T = 1e2
dt = 0.005
LUV.F = 10
x0 = 0.01 * randn(LUV.M)
eps = 0.001
N = 66 # Don't need all Nx for a good approximation of upper spectrum.
# Lyapunov exponents: [ 0.08 0.07 0.06 ... -37.9 -39.09 -41.55]
if mod == "KS":
from dapper.mods.KS import Model
KS = Model()
step = KS.step
x0 = KS.x0
dt = KS.dt